[ROCM] Clean up XLA unit tests for warp size 64

Author: amd-jianli12

third_party/xla/xla/service/gpu/tests/gpu_codegen_test.cc

@@ -71,7 +71,7 @@ std::string GpuCodegenTest::MakePlatformSpecificLlvm(absl::string_view input) {
         is_built_with_rocm_ ? "amdgpu_kernel void" : "void"},
        {"BARRIER",
         is_built_with_rocm_ ? "@llvm.amdgcn.s.barrier" : "@llvm.nvvm.barrier0"},
-       {"SHUFFLE", is_built_with_rocm_ ? "i32 @llvm.amdgcn.ds.swizzle"
+       {"SHUFFLE", is_built_with_rocm_ ? "i32 @llvm.amdgcn.ds.bpermute"
                                        : "float @llvm.nvvm.shfl.sync.down.f32"},
        {"TIDX", is_built_with_rocm_ ? "@llvm.amdgcn.workitem.id.x"
                                     : "@llvm.nvvm.read.ptx.sreg.tid.x"},

third_party/xla/xla/service/gpu/tests/gpu_kernel_tiling_test.cc

@@ -649,7 +649,7 @@ TEST_F(GpuKernelTilingTest, RowReductionCorrectShmemUsage) {
   auto &debug_options = hlo_module->mutable_config().mutable_debug_options();
   debug_options.set_xla_gpu_mlir_emitter_level(3);
   auto expected_ir = is_built_with_rocm_ ? R"(
-; CHECK: %llvm.amdgcn.kernel.input_reduce_fusion.lds.t = type { [4 x [2 x float]] }
+; CHECK: %llvm.amdgcn.kernel.input_reduce_fusion.lds.t = type { [4 x [1 x float]] }
 ; CHECK: @llvm.amdgcn.kernel.input_reduce_fusion.lds = internal addrspace(3) global %llvm.amdgcn.kernel.input_reduce_fusion.lds.t poison
   )"
                                          : R"(

third_party/xla/xla/service/gpu/tests/reduction_vectorization_test.cc

@@ -143,7 +143,7 @@ CHECK-NOT: SHUFFLE
       expected_optimized_llvm_ir,
       {{"X_THREAD", is_built_with_rocm_ ? "@llvm.amdgcn.workitem.id.x"
                                         : "@llvm.nvvm.read.ptx.sreg.tid.x"},
-       {"SHUFFLE", is_built_with_rocm_ ? "@llvm.amdgcn.ds.swizzle"
+       {"SHUFFLE", is_built_with_rocm_ ? "@llvm.amdgcn.ds.bpermute"
                                        : "llvm.nvvm.shfl.sync.down.f32"}});
 
   CompileAndVerifyIr(hlo_text, expected_optimized_llvm_ir, true);

third_party/xla/xla/tests/multioutput_fusion_test.cc

@@ -383,19 +383,19 @@ XLA_TEST_F(MultiOutputFusionTest, MultiOutputReduceFusionMinorWithExtraOutput) {
 XLA_TEST_F(MultiOutputFusionTest, MultiOutputReduceFusionMajorWithExtraOutput) {
   const std::string testcase = absl::StrCat(kScalarOps, R"(
     fused_reduce {
-      p0 = f32[32,32,2]{2,1,0} parameter(0)
+      p0 = f32[64,64,2]{2,1,0} parameter(0)
       c0 = f32[] constant(0)
-      r1 = f32[32,2]{1,0} reduce(p0, c0), dimensions={0}, to_apply=Add
-      mul = f32[32,32,2]{2,1,0} multiply(p0, p0)
+      r1 = f32[64,2]{1,0} reduce(p0, c0), dimensions={0}, to_apply=Add
+      mul = f32[64,64,2]{2,1,0} multiply(p0, p0)
       c1 = f32[] constant(5)
-      r2 = f32[32,2]{1,0} reduce(mul, c1), dimensions={0}, to_apply=Max
-      ROOT tuple = (f32[32,2]{1,0}, f32[32,32,2]{2,1,0}, f32[32,2]{1,0})
+      r2 = f32[64,2]{1,0} reduce(mul, c1), dimensions={0}, to_apply=Max
+      ROOT tuple = (f32[64,2]{1,0}, f32[64,64,2]{2,1,0}, f32[64,2]{1,0})
                      tuple(r1, mul, r2)
     }
 
     ENTRY reduce {
-      p = f32[32,32,2]{2,1,0} parameter(0)
-      ROOT fusion = (f32[32,2]{1,0}, f32[32,32,2]{2,1,0}, f32[32,2]{1,0})
+      p = f32[64,64,2]{2,1,0} parameter(0)
+      ROOT fusion = (f32[64,2]{1,0}, f32[64,64,2]{2,1,0}, f32[64,2]{1,0})
         fusion(p), kind=kInput, calls=fused_reduce
     })");
   auto module = ParseAndReturnVerifiedModule(testcase).value();