Skip to content

[ROCm] Triton in XLA for ROCm - ir_emitter_triton related changes.#10649

Closed
zoranjovanovic-ns wants to merge 1 commit intoopenxla:mainfrom
ROCm:rocm_triton_backend_3
Closed

[ROCm] Triton in XLA for ROCm - ir_emitter_triton related changes.#10649
zoranjovanovic-ns wants to merge 1 commit intoopenxla:mainfrom
ROCm:rocm_triton_backend_3

Conversation

@zoranjovanovic-ns
Copy link
Contributor

Second commit of the series for enabling Triton in XLA for ROCm.

@github-actions github-actions bot added the kokoro:force-run Forces CI to rerun label Mar 18, 2024
@kokoro-team kokoro-team removed the kokoro:force-run Forces CI to rerun label Mar 18, 2024
@github-actions github-actions bot added the kokoro:force-run Forces CI to rerun label Mar 18, 2024
@kokoro-team kokoro-team removed the kokoro:force-run Forces CI to rerun label Mar 18, 2024
@github-actions github-actions bot added the kokoro:force-run Forces CI to rerun label Mar 18, 2024
@kokoro-team kokoro-team removed the kokoro:force-run Forces CI to rerun label Mar 18, 2024
@github-actions github-actions bot added the kokoro:force-run Forces CI to rerun label Mar 18, 2024
@kokoro-team kokoro-team removed the kokoro:force-run Forces CI to rerun label Mar 18, 2024
@github-actions github-actions bot added the kokoro:force-run Forces CI to rerun label Mar 18, 2024
@kokoro-team kokoro-team removed the kokoro:force-run Forces CI to rerun label Mar 18, 2024
@kamaljeeti kamaljeeti requested review from ddunl and tdanyluk March 19, 2024 04:10
#include "tsl/platform/path.h"
#include "tsl/platform/status.h"
#include "tsl/platform/statusor.h"
#ifdef TENSORFLOW_USE_ROCM
Copy link
Member

@ddunl ddunl Mar 19, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think it'd be good if possible to minimize the number of ifdefs, so maybe ifdefs used everywhere at the top, then

#ifdef GOOGLE_CUDA
<includes> 
#endif
#ifdef TENSORFLOW_USE_ROCM
<includes>
#endif

or similar at the bottom of the includes.

@github-actions github-actions bot added the kokoro:force-run Forces CI to rerun label Mar 20, 2024
@kokoro-team kokoro-team removed the kokoro:force-run Forces CI to rerun label Mar 20, 2024
@github-actions github-actions bot added the kokoro:force-run Forces CI to rerun label Mar 20, 2024
@kokoro-team kokoro-team removed the kokoro:force-run Forces CI to rerun label Mar 20, 2024
Comment on lines 761 to 770
#ifndef TENSORFLOW_USE_ROCM
absl::Status CreateTritonPipeline(
mlir::OpPassManager& pm, const se::CudaComputeCapability& cc,
mlir::OpPassManager& pm, const se::GpuComputeCapability& cc,
const TritonGemmConfig& config,
mt::nvidia_gpu::ClusterInfo& out_cluster_info) {
const int ccAsInt = cc.major * 10 + cc.minor;
#else
absl::Status CreateTritonPipeline(
mlir::OpPassManager& pm, const se::GpuComputeCapability& cc,
const TritonGemmConfig& config) {
#endif
Copy link
Contributor

@tdanyluk tdanyluk Mar 20, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In general I would recommend to reduce the number of ifdefs inside the code, to close to 0.
Could we have some kind of abstraction, like have a
Something class, and a derived CudaSomething and RocmSomething or just use a runtime if based on whether the se::GpuComputeCapability is cuda or rocm?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I talked with my colleagues and a class is not really needed.
We could use something like this:

something.h:

int something();

something_cuda.cc:

int something() {
 // implemented with cuda
}

something_rocm.cc:

int something() {
 // implemented with rocm
}

And then build either something_rocm.cc or something_cuda.cc with a select.
So ir_emitter_triton.cc would remain mostly device vendor independent and just include the shared "something.h'.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you for the proposal and clarification. I will change the implementation (based on proposal) and update PR.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

thank you!

@github-actions github-actions bot added the kokoro:force-run Forces CI to rerun label Mar 25, 2024
@kokoro-team kokoro-team removed the kokoro:force-run Forces CI to rerun label Mar 25, 2024
@github-actions github-actions bot added the kokoro:force-run Forces CI to rerun label Apr 4, 2024
@kokoro-team kokoro-team removed the kokoro:force-run Forces CI to rerun label Apr 4, 2024
copybara-service bot pushed a commit to tensorflow/tensorflow that referenced this pull request Apr 4, 2024
…changes.

Imported from GitHub PR openxla/xla#10649

Second commit of the series for enabling Triton in XLA for ROCm.

Copybara import of the project:

--
23d442f83c731cd86131bcd1d91c4e3d7cc42468 by Zoran Jovanovic <zjovanov@amd.com>:

[ROCm] Triton in XLA for ROCm - ir_emitter_triton related changes.

Merging this change closes #10649

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#10649 from ROCm:rocm_triton_backend_3 23d442f83c731cd86131bcd1d91c4e3d7cc42468
PiperOrigin-RevId: 621830985
copybara-service bot pushed a commit to tensorflow/tensorflow that referenced this pull request Apr 5, 2024
…changes.

Imported from GitHub PR openxla/xla#10649

Second commit of the series for enabling Triton in XLA for ROCm.

Copybara import of the project:

--
23d442f83c731cd86131bcd1d91c4e3d7cc42468 by Zoran Jovanovic <zjovanov@amd.com>:

[ROCm] Triton in XLA for ROCm - ir_emitter_triton related changes.

Merging this change closes #10649

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#10649 from ROCm:rocm_triton_backend_3 23d442f83c731cd86131bcd1d91c4e3d7cc42468
PiperOrigin-RevId: 621830985
copybara-service bot pushed a commit to tensorflow/tensorflow that referenced this pull request Apr 5, 2024
…changes.

Imported from GitHub PR openxla/xla#10649

Second commit of the series for enabling Triton in XLA for ROCm.

Copybara import of the project:

--
23d442f83c731cd86131bcd1d91c4e3d7cc42468 by Zoran Jovanovic <zjovanov@amd.com>:

[ROCm] Triton in XLA for ROCm - ir_emitter_triton related changes.

Merging this change closes #10649

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#10649 from ROCm:rocm_triton_backend_3 23d442f83c731cd86131bcd1d91c4e3d7cc42468
PiperOrigin-RevId: 621830985
@github-actions github-actions bot added the kokoro:force-run Forces CI to rerun label Apr 5, 2024
@kokoro-team kokoro-team removed the kokoro:force-run Forces CI to rerun label Apr 5, 2024
@zoranjovanovic-ns
Copy link
Contributor Author

Hi @xla-rotation, I have rebased and pushed, do I need to do something more?

copybara-service bot pushed a commit to tensorflow/tensorflow that referenced this pull request Apr 5, 2024
…changes.

Imported from GitHub PR openxla/xla#10649

Second commit of the series for enabling Triton in XLA for ROCm.

Copybara import of the project:

--
23d442f83c731cd86131bcd1d91c4e3d7cc42468 by Zoran Jovanovic <zjovanov@amd.com>:

[ROCm] Triton in XLA for ROCm - ir_emitter_triton related changes.

Merging this change closes #10649

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#10649 from ROCm:rocm_triton_backend_3 23d442f83c731cd86131bcd1d91c4e3d7cc42468
PiperOrigin-RevId: 621830985
@penpornk
Copy link
Member

penpornk commented Apr 5, 2024

@zoranjovanovic-ns Thank you for checking and for merging conflicts! We are applying some more fixes to pass internal tests. I'll let you know if we need more help. :)

copybara-service bot pushed a commit to tensorflow/tensorflow that referenced this pull request Apr 5, 2024
…changes.

Imported from GitHub PR openxla/xla#10649

Second commit of the series for enabling Triton in XLA for ROCm.

Copybara import of the project:

--
23d442f83c731cd86131bcd1d91c4e3d7cc42468 by Zoran Jovanovic <zjovanov@amd.com>:

[ROCm] Triton in XLA for ROCm - ir_emitter_triton related changes.

Merging this change closes #10649

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#10649 from ROCm:rocm_triton_backend_3 23d442f83c731cd86131bcd1d91c4e3d7cc42468
PiperOrigin-RevId: 621830985
copybara-service bot pushed a commit to tensorflow/tensorflow that referenced this pull request Apr 5, 2024
…changes.

Imported from GitHub PR openxla/xla#10649

Second commit of the series for enabling Triton in XLA for ROCm.

Copybara import of the project:

--
23d442f83c731cd86131bcd1d91c4e3d7cc42468 by Zoran Jovanovic <zjovanov@amd.com>:

[ROCm] Triton in XLA for ROCm - ir_emitter_triton related changes.

Merging this change closes #10649

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#10649 from ROCm:rocm_triton_backend_3 23d442f83c731cd86131bcd1d91c4e3d7cc42468
PiperOrigin-RevId: 621830985
copybara-service bot pushed a commit to tensorflow/tensorflow that referenced this pull request Apr 5, 2024
…changes.

Imported from GitHub PR openxla/xla#10649

Second commit of the series for enabling Triton in XLA for ROCm.

Copybara import of the project:

--
23d442f83c731cd86131bcd1d91c4e3d7cc42468 by Zoran Jovanovic <zjovanov@amd.com>:

[ROCm] Triton in XLA for ROCm - ir_emitter_triton related changes.

Merging this change closes #10649

PiperOrigin-RevId: 622202797
copybara-service bot pushed a commit that referenced this pull request Apr 5, 2024
FUTURE_COPYBARA_INTEGRATE_REVIEW=#10649 from ROCm:rocm_triton_backend_3 23d442f
PiperOrigin-RevId: 622213543
copybara-service bot pushed a commit that referenced this pull request Apr 5, 2024
Reverts 3d6326c

FUTURE_COPYBARA_INTEGRATE_REVIEW=#10649 from ROCm:rocm_triton_backend_3 23d442f
PiperOrigin-RevId: 622186078
copybara-service bot pushed a commit that referenced this pull request Apr 5, 2024
Ubuntu image used in TF SIG Build Dockerfile upgraded from 20.04 to 22.04 (LTS).

Reverts 3d6326c

FUTURE_COPYBARA_INTEGRATE_REVIEW=#10649 from ROCm:rocm_triton_backend_3 23d442f
PiperOrigin-RevId: 615889505
copybara-service bot pushed a commit that referenced this pull request Apr 5, 2024
…eters.

Use TiledHloInstructions in the Triton emitter.

Reverts 3d6326c

FUTURE_COPYBARA_INTEGRATE_REVIEW=#10649 from ROCm:rocm_triton_backend_3 23d442f
PiperOrigin-RevId: 621553019
@copybara-service copybara-service bot closed this in aa08925 Apr 5, 2024
@alekstheod alekstheod deleted the rocm_triton_backend_3 branch July 3, 2025 08:11
@alekstheod alekstheod restored the rocm_triton_backend_3 branch July 15, 2025 08:23
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

8 participants