[ROCm] Triton in XLA for ROCm - ir_emitter_triton related changes.#10649
[ROCm] Triton in XLA for ROCm - ir_emitter_triton related changes.#10649zoranjovanovic-ns wants to merge 1 commit intoopenxla:mainfrom
Conversation
a7ea1f4 to
2080824
Compare
2080824 to
34b6c28
Compare
34b6c28 to
06fe001
Compare
06fe001 to
5ca24f4
Compare
xla/service/gpu/ir_emitter_triton.cc
Outdated
| #include "tsl/platform/path.h" | ||
| #include "tsl/platform/status.h" | ||
| #include "tsl/platform/statusor.h" | ||
| #ifdef TENSORFLOW_USE_ROCM |
There was a problem hiding this comment.
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.
5ca24f4 to
3d43cfc
Compare
3d43cfc to
15dd1a8
Compare
xla/service/gpu/ir_emitter_triton.cc
Outdated
| #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 |
There was a problem hiding this comment.
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?
There was a problem hiding this comment.
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'.
There was a problem hiding this comment.
Thank you for the proposal and clarification. I will change the implementation (based on proposal) and update PR.
15dd1a8 to
5be1280
Compare
5be1280 to
d7a2b5b
Compare
23d442f to
a7b1560
Compare
…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
…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
…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
a7b1560 to
36b991a
Compare
|
Hi @xla-rotation, I have rebased and pushed, do I need to do something more? |
…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
|
@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. :) |
…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
…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
…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
Second commit of the series for enabling Triton in XLA for ROCm.