PR #19484: [ROCm] Fix //xla/tests:complex_unary_op_test and //xla/ser…

…vice/gpu/te… Imported from GitHub PR #19484 …sts:gpu_input_fusible_slice_test Copybara import of the project: -- 0d30738 by Dragan Mladjenovic <[email protected]>: [ROCm] Fix //xla/tests:complex_unary_op_test and //xla/service/gpu/tests:gpu_input_fusible_slice_test Merging this change closes #19484 COPYBARA_INTEGRATE_REVIEW=#19484 from ROCm:mlir_tests_new 0d30738 PiperOrigin-RevId: 698374588
openxla · Nov 20, 2024 · 66073ea · 66073ea
1 parent ffd2c11
commit 66073ea
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/xla/service/gpu/fusions/transforms/lower_to_llvm.cc b/xla/service/gpu/fusions/transforms/lower_to_llvm.cc
@@ -68,6 +68,8 @@ class LowerToLLVMPass : public impl::LowerToLLVMPassBase<LowerToLLVMPass> {
     mlir::arith::populateArithExpandOpsPatterns(patterns);
     mlir::arith::populateArithToLLVMConversionPatterns(type_converter,
                                                        patterns);
+    mlir::populateMathToLLVMConversionPatterns(type_converter, patterns,
+                                               /* approximateLog1p */ false);
     if (!this->is_amd_gpu_) {
       mlir::populateGpuToNVVMConversionPatterns(type_converter, patterns);
     } else {
@@ -79,7 +81,6 @@ class LowerToLLVMPass : public impl::LowerToLLVMPassBase<LowerToLLVMPass> {
     mlir::cf::populateControlFlowToLLVMConversionPatterns(type_converter,
                                                           patterns);
     mlir::populateComplexToLLVMConversionPatterns(type_converter, patterns);
-    mlir::populateMathToLLVMConversionPatterns(type_converter, patterns);
 
     //  Setup target.
     if (!this->is_amd_gpu_) {

diff --git a/xla/service/gpu/tests/gpu_input_fusible_slice_test.cc b/xla/service/gpu/tests/gpu_input_fusible_slice_test.cc
@@ -68,7 +68,8 @@ TEST_F(GpuSliceInputFusionTest, InputFusionWithATupleOfSlices) {
           .value();
   auto expected_ir = is_built_with_rocm_ ? R"(
 ; CHECK-LABEL: define amdgpu_kernel void @{{[a-z_]*}}fusion
-; CHECK: slice2
+; CHECK: store half %{{.*}}, ptr %{{.*}}, align 2
+; CHECK: store half %{{.*}}, ptr %{{.*}}, align 2
 ; CHECK: }
 )"
                                          : R"(
@@ -116,7 +117,8 @@ TEST_F(GpuSliceInputFusionTest, ConcatThenSplit) {
           .value();
   auto expected_ir = is_built_with_rocm_ ? R"(
 ; CHECK-LABEL: define amdgpu_kernel void @{{[a-z_]*}}fusion
-; CHECK: slice2
+; CHECK: store half %{{.*}}, ptr %{{.*}}, align 2
+; CHECK: store half %{{.*}}, ptr %{{.*}}, align 2
 ; CHECK: }
 )"
                                          : R"(