add lit tests for swizzle config

Made-with: Cursor
Signed-off-by: Yu-Zhewen <zhewenyu@amd.com>

Author: Yu-Zhewen

compiler/src/iree/compiler/Codegen/Dialect/GPU/TargetUtils/ConfigUtils.cpp

@@ -916,22 +916,27 @@ getMatmulOrIGEMMLoweringConfigAndWorkgroupSize(
   if (!scaled && useDirectLoad) {
     Attribute lhsAttr = useGlobalDma;
     Attribute rhsAttr = useGlobalDma;
-    // Apply XOR swizzle for bank conflict avoidance. Only swizzle operands
-    // whose reduction dim is innermost (contiguous reads).
-    if (!transposedLhs) {
-      FailureOr<Attribute> lhsSwizzleAttr =
-          getXorShuffleAttr(context, useGlobalDma, target, kind,
-                            schedule->kTileSizes, kMMAOperandLhs);
-      if (succeeded(lhsSwizzleAttr)) {
-        lhsAttr = *lhsSwizzleAttr;
+    // Apply XOR swizzle for BF16 DMA operands whose reduction dim is
+    // innermost (contiguous reads) to avoid LDS bank conflicts.
+    SmallVector<Type> elemTypes;
+    kind.getElementTypes(elemTypes);
+    bool isBF16 = !elemTypes.empty() && elemTypes[0].isBF16();
+    if (isBF16) {
+      if (!transposedLhs) {
+        FailureOr<Attribute> lhsSwizzleAttr =
+            getXorShuffleAttr(context, useGlobalDma, target, kind,
+                              schedule->kTileSizes, kMMAOperandLhs);
+        if (succeeded(lhsSwizzleAttr)) {
+          lhsAttr = *lhsSwizzleAttr;
+        }
       }
-    }
-    if (transposedRhs) {
-      FailureOr<Attribute> rhsSwizzleAttr =
-          getXorShuffleAttr(context, useGlobalDma, target, kind,
-                            schedule->kTileSizes, kMMAOperandRhs);
-      if (succeeded(rhsSwizzleAttr)) {
-        rhsAttr = *rhsSwizzleAttr;
+      if (transposedRhs) {
+        FailureOr<Attribute> rhsSwizzleAttr =
+            getXorShuffleAttr(context, useGlobalDma, target, kind,
+                              schedule->kTileSizes, kMMAOperandRhs);
+        if (succeeded(rhsSwizzleAttr)) {
+          rhsAttr = *rhsSwizzleAttr;
+        }
       }
     }
     promotionArray = {lhsAttr, rhsAttr};

compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/config_tile_and_fuse_gfx950.mlir

@@ -386,6 +386,38 @@ func.func @scaled_matmul_accumulate(
 
 // -----
 
+// BF16 matmul with direct-load DMA gets LHS XOR swizzle for bank conflict
+// avoidance. Only LHS is swizzled (reduction dim is innermost for A[M,K]).
+func.func @matmul_bf16(
+    %arg0: tensor<4096x4096xbf16>,
+    %arg1: tensor<4096x4096xbf16>,
+    %arg2: tensor<4096x4096xf32>) -> tensor<4096x4096xf32> {
+  %0 = linalg.matmul ins(%arg0, %arg1 : tensor<4096x4096xbf16>, tensor<4096x4096xbf16>)
+                      outs(%arg2 : tensor<4096x4096xf32>) -> tensor<4096x4096xf32>
+  return %0 : tensor<4096x4096xf32>
+}
+// CHECK-LABEL: func.func @matmul_bf16
+// CHECK:   lowering_config = #iree_gpu.lowering_config
+// CHECK-SAME: mma_kind = #iree_gpu.mma_layout<MFMA_F32_16x16x32_BF16>
+
+// CHECK-REMARKS: [Analysis] SharedMemoryUsage
+// CHECK-REMARKS-SAME: Category:deduceMMASchedule
+// CHECK-REMARKS-SAME: Remark=16384
+
+// CHECK-REMARKS-DIRECT-LOAD-2: [Analysis] SharedMemoryUsage
+// CHECK-REMARKS-DIRECT-LOAD-2-SAME: Category:deduceMMASchedule
+// CHECK-REMARKS-DIRECT-LOAD-2-SAME: Remark=32768
+
+// CHECK-REMARKS-DIRECT-LOAD-3: [Analysis] SharedMemoryUsage
+// CHECK-REMARKS-DIRECT-LOAD-3-SAME: Category:deduceMMASchedule
+// CHECK-REMARKS-DIRECT-LOAD-3-SAME: Remark=49152
+
+// CHECK-DIRECT-LOAD-LABEL: func.func @matmul_bf16
+// CHECK-DIRECT-LOAD:       linalg.matmul {lowering_config = #iree_gpu.lowering_config
+// CHECK-DIRECT-LOAD-SAME:    promotion_types = [#iree_gpu.swizzle_operand<copy_config = #iree_gpu.use_global_load_dma, swizzle = #iree_codegen.xor_shuffle<128, 8>>, #iree_gpu.use_global_load_dma]
+
+// -----
+
 // Very large f16 matmul — compute-bound, so picks 32x32x16 (higher compute per
 // instruction, lower VGPR pressure than 16x16x32).
 func.func @matmul_f16_compute_bound(