Fix dot_bf16 using feature conditions

ROCm · Jun 20, 2024 · 59ada7f · 59ada7f
1 parent c09263e
commit 59ada7f
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/xla/lit.cfg.py b/xla/lit.cfg.py
@@ -43,8 +43,13 @@
     ("%PYTHON", os.getenv("PYTHON", sys.executable)),
 ])
 
+if lit_config.params.get('PTX', 'GCN'):
+    config.available_features.add("IS_ROCM")
+
+
 # Include additional substitutions that may be defined via params
 config.substitutions.extend(
     ("%%{%s}" % key, val)
     for key, val in lit_config.params.items()
 )
+
diff --git a/xla/service/gpu/tests/dot_bf16.hlo b/xla/service/gpu/tests/dot_bf16.hlo
@@ -1,5 +1,6 @@
-// RUN: hlo-opt %s --platform=gpu --stage=hlo --xla_gpu_target_config_filename=%S/../../../tools/hlo_opt/gpu_specs/v100.txtpb --split-input-file | FileCheck %s --check-prefixes=CHECK-SM70
-// RUN: hlo-opt %s --platform=gpu --stage=hlo --xla_gpu_target_config_filename=%S/../../../tools/hlo_opt/gpu_specs/a100_80.txtpb --split-input-file --xla_gpu_autotune_level=0 --xla_gpu_enable_triton_gemm=false | FileCheck %s --check-prefixes=CHECK-SM80
+// RUN: %if !IS_ROCM %{ hlo-opt %s --platform=gpu --stage=hlo --xla_gpu_target_config_filename=%S/../../../tools/hlo_opt/gpu_specs/v100.txtpb --split-input-file | FileCheck %s --check-prefixes=CHECK-SM70 %}
+// RUN: %if !IS_ROCM %{ hlo-opt %s --platform=gpu --stage=hlo --xla_gpu_target_config_filename=%S/../../../tools/hlo_opt/gpu_specs/a100_80.txtpb --split-input-file --xla_gpu_autotune_level=0 --xla_gpu_enable_triton_gemm=false | FileCheck %s --check-prefixes=CHECK-SM80 %}
+// RUN: %if IS_ROCM %{ hlo-opt %s --platform=gpu --stage=hlo --xla_gpu_target_config_filename=%S/../../../tools/hlo_opt/gpu_specs/mi200.txtpb --split-input-file --xla_gpu_autotune_level=0 --xla_gpu_enable_triton_gemm=false | FileCheck %s --check-prefixes=CHECK-SM80 %}
 
 
 // CHECK-SM70: custom-call(f32