[ROCm] Triton in XLA for ROCm - changed ir_emitter_triton_test by zoranjovanovic-ns · Pull Request #11352 · openxla/xla

zoranjovanovic-ns · 2024-04-09T12:44:22Z

Modified the test case to pass on rocm and cuda.
There will be at least one more XLA related PR - to switch on Triton usage for ROCm, but only after PR for adding Triton build files for ROCm on openxla/triton.

bchetioui · 2024-04-11T06:47:33Z

xla/service/gpu/ir_emitter_triton_test.cc

+  const se::GpuComputeCapability& GpuComputeComp() {
+    return device_desc().gpu_compute_capability();
+  }
+  enum class Switch : uint32_t {


Could we just use a bool directly instead of this Switch enum? I don't think we gain more insights by using this type, especially given that the constructors are called False and True. Am I missing something?

Upon looking at all the CudaOrRocmChecks below, it also seems like this could all be deleted. Could you explain what is the motivation behind going through these APIs for testing?

It seems that reason for this type of checks is gone with removal of Cuda related checks (e.g. cc.IsAtLeast(e::CudaComputeCapability::AMPERE)).
I will modify my changes and update PR.

bchetioui · 2024-04-11T06:51:37Z

xla/service/gpu/ir_emitter_triton_test.cc

 }

 TEST_F(TritonGemmTestAny, DoAddConstantToScalarAndBroadcastThat) {
+  if (CudaOrRocmCheck(Switch::False, Switch::True)) {


Could this just be if (std::holds_alternative<se::RocmComputeCapability>(GpuComputeComp())) {?

bchetioui · 2024-04-11T06:51:43Z

xla/service/gpu/ir_emitter_triton_test.cc

 }

 TEST_F(TritonGemmTest, SingleElementTileIsHandled) {
+  if (CudaOrRocmCheck(Switch::False, Switch::True)) {


Can this just be if (std::holds_alternative<se::RocmComputeCapability>(GpuComputeComp())) {?

bchetioui · 2024-04-11T06:52:01Z

xla/service/gpu/ir_emitter_triton_test.cc


 TEST_F(TritonGemmTest, SingleElementTileIsHandled) {
+  if (CudaOrRocmCheck(Switch::False, Switch::True)) {
+    GTEST_SKIP() << "Not using autotuner on ROCM yet..";


bchetioui · 2024-04-11T06:52:16Z

xla/service/gpu/ir_emitter_triton_test.cc

 }

 TEST_F(TritonGemmTestWithoutTritonGemmAny, SkipF32F32) {
+  if (CudaOrRocmCheck(Switch::False, Switch::True)) {


Can this just be if (std::holds_alternative<se::RocmComputeCapability>(GpuComputeComp())) {?

bchetioui · 2024-04-11T06:52:22Z

xla/service/gpu/ir_emitter_triton_test.cc

 }

 TEST_F(TritonGemmTestWithoutTritonGemmAny, SkipU8) {
+  if (CudaOrRocmCheck(Switch::False, Switch::True)) {


Can this just be if (std::holds_alternative<se::RocmComputeCapability>(GpuComputeComp())) {?

bchetioui · 2024-04-11T06:52:41Z

xla/service/gpu/ir_emitter_triton_test.cc

 }

 TEST_F(TritonGemmTest, FailIfTooMuchShmem) {
+  if (CudaOrRocmCheck(Switch::False, Switch::True)) {


Can this just be if (std::holds_alternative<se::RocmComputeCapability>(GpuComputeComp())) {?

bchetioui · 2024-04-11T10:56:15Z

xla/service/gpu/ir_emitter_triton_test.cc

+        GpuComputeComp());
+  }
+  bool SkipBF16Tests() {
+    return CudaOrRocmCheck(


Can the check logic just be inlined? Overall, it does seem like we have a lot of CudaOrRocmCheck implementations that don't really need to exist---we can just implement SkipBF16Tests directly and check for alternatives in tests, I believe?

…s on rocm and cuda

zoranjovanovic-ns · 2024-04-15T14:11:23Z

@xla-rotation modified according to comments.

bchetioui

Thank you!

bchetioui · 2024-04-15T18:44:11Z

xla/service/gpu/ir_emitter_triton_test.cc

 }

 TEST_F(TritonGemmLevel2Test, DoubleBroadcastOfScalarConstantIsHandled) {
+  if (SkipBF16Tests()) {


Note: I think that many of these tests could be rewritten to not use bf16, thereby allowing the removal of this filter and allowing more coverage on rocm. Be that as it may, this is fine for now.

bchetioui · 2024-04-15T18:45:32Z

xla/service/gpu/ir_emitter_triton_test.cc

+  const se::GpuComputeCapability& GpuComputeComp() {
+    return device_desc().gpu_compute_capability();
+  }
+  bool SkipBF16Tests() {


nit: there should be one empty line between definition. I'll make the edit on my end before submitting, just wanted to highlight it.

…test Imported from GitHub PR openxla/xla#11352 Modified the test case to pass on rocm and cuda. There will be at least one more XLA related PR - to switch on Triton usage for ROCm, but only after PR for adding Triton build files for ROCm on openxla/triton. Copybara import of the project: -- dd80a6f44295642efd3ae6af6ffd3e2a3302d36e by Zoran Jovanovic <zjovanov@amd.com>: [ROCm] Triton in XLA for ROCm - changed ir_emitter_triton_test to pass on rocm and cuda Merging this change closes #11352 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#11352 from ROCm:rocm_triton_backend_5 dd80a6f44295642efd3ae6af6ffd3e2a3302d36e PiperOrigin-RevId: 625033187

…test Imported from GitHub PR #11352 Modified the test case to pass on rocm and cuda. There will be at least one more XLA related PR - to switch on Triton usage for ROCm, but only after PR for adding Triton build files for ROCm on openxla/triton. Copybara import of the project: -- dd80a6f by Zoran Jovanovic <zjovanov@amd.com>: [ROCm] Triton in XLA for ROCm - changed ir_emitter_triton_test to pass on rocm and cuda Merging this change closes #11352 FUTURE_COPYBARA_INTEGRATE_REVIEW=#11352 from ROCm:rocm_triton_backend_5 dd80a6f PiperOrigin-RevId: 625043989

…test Imported from GitHub PR openxla/xla#11352 Modified the test case to pass on rocm and cuda. There will be at least one more XLA related PR - to switch on Triton usage for ROCm, but only after PR for adding Triton build files for ROCm on openxla/triton. Copybara import of the project: -- dd80a6f44295642efd3ae6af6ffd3e2a3302d36e by Zoran Jovanovic <zjovanov@amd.com>: [ROCm] Triton in XLA for ROCm - changed ir_emitter_triton_test to pass on rocm and cuda Merging this change closes #11352 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#11352 from ROCm:rocm_triton_backend_5 dd80a6f44295642efd3ae6af6ffd3e2a3302d36e PiperOrigin-RevId: 625043989

…test Imported from GitHub PR #11352 Modified the test case to pass on rocm and cuda. There will be at least one more XLA related PR - to switch on Triton usage for ROCm, but only after PR for adding Triton build files for ROCm on openxla/triton. Copybara import of the project: -- dd80a6f by Zoran Jovanovic <zjovanov@amd.com>: [ROCm] Triton in XLA for ROCm - changed ir_emitter_triton_test to pass on rocm and cuda Merging this change closes #11352 FUTURE_COPYBARA_INTEGRATE_REVIEW=#11352 from ROCm:rocm_triton_backend_5 dd80a6f PiperOrigin-RevId: 625043989

…test Imported from GitHub PR openxla/xla#11352 Modified the test case to pass on rocm and cuda. There will be at least one more XLA related PR - to switch on Triton usage for ROCm, but only after PR for adding Triton build files for ROCm on openxla/triton. Copybara import of the project: -- dd80a6f44295642efd3ae6af6ffd3e2a3302d36e by Zoran Jovanovic <zjovanov@amd.com>: [ROCm] Triton in XLA for ROCm - changed ir_emitter_triton_test to pass on rocm and cuda Merging this change closes #11352 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#11352 from ROCm:rocm_triton_backend_5 dd80a6f44295642efd3ae6af6ffd3e2a3302d36e PiperOrigin-RevId: 625043989

…test Imported from GitHub PR #11352 Modified the test case to pass on rocm and cuda. There will be at least one more XLA related PR - to switch on Triton usage for ROCm, but only after PR for adding Triton build files for ROCm on openxla/triton. Copybara import of the project: -- dd80a6f by Zoran Jovanovic <zjovanov@amd.com>: [ROCm] Triton in XLA for ROCm - changed ir_emitter_triton_test to pass on rocm and cuda Merging this change closes #11352 FUTURE_COPYBARA_INTEGRATE_REVIEW=#11352 from ROCm:rocm_triton_backend_5 dd80a6f PiperOrigin-RevId: 625043989

…test Imported from GitHub PR openxla/xla#11352 Modified the test case to pass on rocm and cuda. There will be at least one more XLA related PR - to switch on Triton usage for ROCm, but only after PR for adding Triton build files for ROCm on openxla/triton. Copybara import of the project: -- dd80a6f44295642efd3ae6af6ffd3e2a3302d36e by Zoran Jovanovic <zjovanov@amd.com>: [ROCm] Triton in XLA for ROCm - changed ir_emitter_triton_test to pass on rocm and cuda Merging this change closes #11352 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#11352 from ROCm:rocm_triton_backend_5 dd80a6f44295642efd3ae6af6ffd3e2a3302d36e PiperOrigin-RevId: 625043989

FUTURE_COPYBARA_INTEGRATE_REVIEW=#11352 from ROCm:rocm_triton_backend_5 dd80a6f PiperOrigin-RevId: 625365616

…test Imported from GitHub PR openxla/xla#11352 Modified the test case to pass on rocm and cuda. There will be at least one more XLA related PR - to switch on Triton usage for ROCm, but only after PR for adding Triton build files for ROCm on openxla/triton. Copybara import of the project: -- dd80a6f44295642efd3ae6af6ffd3e2a3302d36e by Zoran Jovanovic <zjovanov@amd.com>: [ROCm] Triton in XLA for ROCm - changed ir_emitter_triton_test to pass on rocm and cuda Merging this change closes #11352 PiperOrigin-RevId: 625361694

FUTURE_COPYBARA_INTEGRATE_REVIEW=#11352 from ROCm:rocm_triton_backend_5 dd80a6f PiperOrigin-RevId: 625367225

github-actions bot added the kokoro:force-run Forces CI to rerun label Apr 9, 2024

github-actions bot assigned kamaljeeti and xla-rotation Apr 9, 2024

kokoro-team removed the kokoro:force-run Forces CI to rerun label Apr 9, 2024

kamaljeeti requested review from ddunl and golechwierowicz April 10, 2024 07:00

golechwierowicz requested a review from bchetioui April 10, 2024 09:32

bchetioui requested changes Apr 11, 2024

View reviewed changes

[ROCm] Triton in XLA for ROCm - changed ir_emitter_triton_test to pas…

dd80a6f

…s on rocm and cuda

zoranjovanovic-ns force-pushed the rocm_triton_backend_5 branch from bb5a5fd to dd80a6f Compare April 12, 2024 13:45

github-actions bot added the kokoro:force-run Forces CI to rerun label Apr 12, 2024

kokoro-team removed the kokoro:force-run Forces CI to rerun label Apr 12, 2024

bchetioui approved these changes Apr 15, 2024

View reviewed changes

copybara-service bot mentioned this pull request Apr 15, 2024

PR #11352: [ROCm] Triton in XLA for ROCm - changed ir_emitter_triton_test tensorflow/tensorflow#65746

Closed

copybara-service bot mentioned this pull request Apr 15, 2024

PR #11352: [ROCm] Triton in XLA for ROCm - changed ir_emitter_triton_test #11533

Merged

copybara-service bot mentioned this pull request Apr 15, 2024

PR #11352: [ROCm] Triton in XLA for ROCm - changed ir_emitter_triton_test tensorflow/tensorflow#65750

Merged

copybara-service bot pushed a commit that referenced this pull request Apr 16, 2024

Remove unnecessary mentions of tf_runtime in visibilities

6103370

FUTURE_COPYBARA_INTEGRATE_REVIEW=#11352 from ROCm:rocm_triton_backend_5 dd80a6f PiperOrigin-RevId: 625365616

copybara-service bot mentioned this pull request Apr 16, 2024

Remove unnecessary mentions of tf_runtime in visibilities #11581

Closed

copybara-service bot pushed a commit that referenced this pull request Apr 16, 2024

Integrate StableHLO at openxla/stablehlo@e81411ef

556edc6

FUTURE_COPYBARA_INTEGRATE_REVIEW=#11352 from ROCm:rocm_triton_backend_5 dd80a6f PiperOrigin-RevId: 625367225

copybara-service bot mentioned this pull request Apr 16, 2024

Integrate StableHLO at openxla/stablehlo@e81411ef #11582

Merged

copybara-service bot closed this in b98180d Apr 16, 2024

i-chaochen mentioned this pull request Apr 19, 2024

[ROCM] Include rocm Triton build in build system. openxla/triton#7

Closed

alekstheod deleted the rocm_triton_backend_5 branch July 3, 2025 08:11

alekstheod restored the rocm_triton_backend_5 branch July 15, 2025 08:23

Conversation

zoranjovanovic-ns commented Apr 9, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

bchetioui Apr 11, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zoranjovanovic-ns commented Apr 15, 2024

Uh oh!

bchetioui left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

zoranjovanovic-ns commented Apr 9, 2024 •

edited

Loading

bchetioui Apr 11, 2024 •

edited

Loading