Merge pull request #30 from ROCm/ci_rocm-jaxlib-v0.4.28-gemm_alg_pick…

…er-fix Fixing gemm_algorithm_picker_test due to missing commit
ROCm · Jul 19, 2024 · 3483ccf · 3483ccf
2 parents 1a7f339 + 8d4bc3e
commit 3483ccf
Showing 1 changed file with 6 additions and 7 deletions.
diff --git a/xla/service/gpu/gemm_algorithm_picker_test.cc b/xla/service/gpu/gemm_algorithm_picker_test.cc
@@ -106,7 +106,7 @@ TF_ASSERT_OK_AND_ASSIGN(auto module,
     TF_ASSERT_OK_AND_ASSIGN(
       bool changed,
       RunHloPass(
-          GemmRewriter(gpu_comp(), /*toolkit_version=*/12040),
+          GemmRewriter(gpu_comp()),
           module.get()));
 
     AutotuneConfig cfg{DeviceConfig{stream_exec(), nullptr}, debug_opts};
@@ -135,7 +135,7 @@ TF_ASSERT_OK_AND_ASSIGN(auto module,
     TF_ASSERT_OK_AND_ASSIGN(
       bool changed,
       RunHloPass(
-          GemmRewriter(gpu_comp(), /*toolkit_version=*/12040),
+          GemmRewriter(gpu_comp()),
           module.get()));
 
     AutotuneConfig cfg{DeviceConfig{stream_exec(), nullptr}, debug_opts};
@@ -165,7 +165,7 @@ ENTRY main {
   TF_ASSERT_OK_AND_ASSIGN(
       changed,
       RunHloPass(
-          GemmRewriter(gpu_comp(), /*toolkit_version=*/12040),
+          GemmRewriter(gpu_comp()),
           m.get()));
   changed = false;
   DebugOptions opts;
@@ -192,7 +192,7 @@ ENTRY main {
   TF_ASSERT_OK_AND_ASSIGN(
       changed,
       RunHloPass(
-          GemmRewriter(gpu_comp(), /*toolkit_version=*/12040),
+          GemmRewriter(gpu_comp()),
           m.get()));
   changed = false;
   TF_ASSERT_OK_AND_ASSIGN(changed,
@@ -226,7 +226,7 @@ ENTRY main {
   TF_ASSERT_OK_AND_ASSIGN(
       changed,
       RunHloPass(
-          GemmRewriter(gpu_comp(), /*toolkit_version=*/12040),
+          GemmRewriter(gpu_comp()),
           m.get()));
   changed = false;
 
@@ -260,8 +260,7 @@ ENTRY main {
   TF_ASSERT_OK_AND_ASSIGN(
       changed,
       RunHloPass(
-          GemmRewriter(gpu_comp(),
-              /*toolkit_version=*/12040),
+          GemmRewriter(gpu_comp()),
           m.get()));
   changed = false;
   TF_ASSERT_OK_AND_ASSIGN(