[CUDA][LLVMGPU] Fix sm_120 WGP params and add BF16 mma.sync coverage (#24648)
## Description
Fixes a correctness issue in `getSM120WgpDetails()` carried over from
the Ampere placeholder, consolidates BF16 test coverage into existing
architecture-scoped files, and extends sm_120 tests with a
mixed-precision case.
- **BF16 mma.sync fix**: `NV_MMA_SYNC_F32_16x8x16_BF16` was missing from
sm_120's `mmaOps` despite being supported from Ampere onwards. sm_121
already advertises this intrinsic. Without this fix, bf16 × bf16 → f32
matmuls silently fall back to SIMT tiling on sm_120.
- **Test consolidation**: Per the reviewer's suggestion, tests no longer
have their own per-architecture files. Instead they are folded into the
existing architecture-scoped test files.
## Testcases
1. config_tile_and_fuse_sm80.mlir / config_tile_and_fuse_sm120.mlir:
each gets a new BF16 section verifying that bf16 × bf16 → f32 selects
`NV_MMA_SYNC_F32_16x8x16_BF16` at the config stage.
1. config_tile_and_fuse_sm120.mlir: adds a mixed-precision f32 × bf16 →
f32 case to verify the SIMT fallback path when no MMA intrinsic matches
the operand types.
1. pipeline_tile_and_fuse_mma_sync.mlir: adds a shared BF16 pipeline
lowering section that covers both sm_80 and sm_120, verifying
`nvgpu.mma.sync` is generated with the correct shape.
1. target_device_features.mlir: expands the SM120 metadata check to a
full wgp attribute check (mma list, workgroup sizes, memory limits),
matching the granularity of the existing SM89 check.
1. bf16_mma_sm80.mlir: deleted — coverage moved into the files above.
---------
Signed-off-by: weimin023 <tnwilly@gmail.com>
diff --git a/compiler/plugins/target/CUDA/test/target_device_features.mlir b/compiler/plugins/target/CUDA/test/target_device_features.mlir
index 91d4f40..158affd 100644
--- a/compiler/plugins/target/CUDA/test/target_device_features.mlir
+++ b/compiler/plugins/target/CUDA/test/target_device_features.mlir
@@ -22,6 +22,12 @@
// PTX80: target_info = #iree_gpu.target<arch = "sm_89", features = "+ptx80",
// SM120: target_info = #iree_gpu.target<arch = "sm_120", features = "+ptx87",
+// SM120-SAME: wgp = <compute = fp64|fp32|fp16|int64|int32|int16|int8, storage = b64|b32|b16|b8,
+// SM120-SAME: subgroup = shuffle|arithmetic, dot = dp4xi8toi32,
+// SM120-SAME: mma = [<NV_MMA_SYNC_F32_16x8x16_F16>, <NV_MMA_SYNC_F16_16x8x16_F16>, <NV_MMA_SYNC_F32_16x8x16_BF16>, <NV_WMMA_F32_16x16x16_F16>, <NV_WMMA_F16_16x16x16_F16>],
+// SM120-SAME: subgroup_size_choices = [32], max_workgroup_sizes = [1024, 1024, 64],
+// SM120-SAME: max_thread_count_per_workgroup = 1024, max_workgroup_memory_bytes = 101376,
+// SM120-SAME: max_workgroup_counts = [2147483647, 65535, 65535]>
// SM121: target_info = #iree_gpu.target<arch = "sm_121", features = "+ptx88",
stream.executable public @target_device_features {
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/GPU/TargetUtils/KnownTargets.cpp b/compiler/src/iree/compiler/Codegen/Dialect/GPU/TargetUtils/KnownTargets.cpp
index 64ccd85..39724bf 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/GPU/TargetUtils/KnownTargets.cpp
+++ b/compiler/src/iree/compiler/Codegen/Dialect/GPU/TargetUtils/KnownTargets.cpp
@@ -1076,14 +1076,10 @@
}
const WgpDetails *getSM120WgpDetails() {
- // TODO: Model Blackwell-specific capabilities once IREE has target details
- // for the architecture. Reuse the existing NVIDIA capability model for now
- // so canonical CUDA targets such as sm_120 can be represented.
- // See:
- // https://docs.nvidia.com/cuda/blackwell-tuning-guide/index.html#nvidia-blackwell-tuning
static const MMAIntrinsic mmaOps[] = {
MMAIntrinsic::NV_MMA_SYNC_F32_16x8x16_F16,
MMAIntrinsic::NV_MMA_SYNC_F16_16x8x16_F16,
+ MMAIntrinsic::NV_MMA_SYNC_F32_16x8x16_BF16,
MMAIntrinsic::NV_WMMA_F32_16x16x16_F16,
MMAIntrinsic::NV_WMMA_F16_16x16x16_F16,
};
@@ -1096,9 +1092,9 @@
0,
nullptr,
{32, 32},
- {1024, 1024, 1024},
+ {1024, 1024, 64},
1024,
- 163 * 1024,
+ 99 * 1024,
{0x7fffffff, 0xffff, 0xffff}};
return &sm120Wgp;
}
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/BUILD.bazel b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/BUILD.bazel
index e5ded30..ae991c2 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/BUILD.bazel
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/BUILD.bazel
@@ -16,7 +16,6 @@
name = "lit",
srcs = enforce_glob(
[
- "bf16_mma_sm80.mlir",
"config_tile_and_fuse_sm120.mlir",
"config_tile_and_fuse_sm80.mlir",
"config_vector_distribute_sm80.mlir",
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/CMakeLists.txt
index 5d58dde..b293f5d 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/CMakeLists.txt
@@ -14,7 +14,6 @@
NAME
lit
SRCS
- "bf16_mma_sm80.mlir"
"config_tile_and_fuse_sm120.mlir"
"config_tile_and_fuse_sm80.mlir"
"config_vector_distribute_sm80.mlir"
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/bf16_mma_sm80.mlir b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/bf16_mma_sm80.mlir
deleted file mode 100644
index 6873631..0000000
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/bf16_mma_sm80.mlir
+++ /dev/null
@@ -1,39 +0,0 @@
-// RUN: iree-opt --mlir-print-local-scope --iree-gpu-test-target=sm_80 \
-// RUN: --pass-pipeline="builtin.module(iree-llvmgpu-select-lowering-strategy)" %s | FileCheck %s --check-prefix=CONFIG
-// RUN: iree-opt --split-input-file --iree-gpu-test-target=sm_80 \
-// RUN: --iree-codegen-llvmgpu-configuration-pipeline --iree-codegen-llvmgpu-nvvm-lowering-pipeline %s | FileCheck %s --check-prefix=NVVM
-
-// Test that Ampere selects and lowers BF16 matmuls through NVIDIA mma.sync.
-
-#pipeline_layout = #hal.pipeline.layout<bindings = [
- #hal.pipeline.binding<storage_buffer>,
- #hal.pipeline.binding<storage_buffer>,
- #hal.pipeline.binding<storage_buffer>
-]>
-// Drives BF16xbf16xf32 matmul selection through the NVIDIA mma.sync path.
-func.func @matmul_256x256x256_bf16_f32() {
- %cst = arith.constant 0.000000e+00 : f32
- %c0 = arith.constant 0 : index
- %0 = hal.interface.binding.subspan layout(#pipeline_layout) binding(0) alignment(64) offset(%c0) flags(ReadOnly) : !iree_tensor_ext.dispatch.tensor<readonly:tensor<256x256xbf16>>
- %1 = hal.interface.binding.subspan layout(#pipeline_layout) binding(1) alignment(64) offset(%c0) flags(ReadOnly) : !iree_tensor_ext.dispatch.tensor<readonly:tensor<256x256xbf16>>
- %2 = hal.interface.binding.subspan layout(#pipeline_layout) binding(2) alignment(64) offset(%c0) : !iree_tensor_ext.dispatch.tensor<writeonly:tensor<256x256xf32>>
- %3 = iree_tensor_ext.dispatch.tensor.load %0, offsets = [0, 0], sizes = [256, 256], strides = [1, 1] : !iree_tensor_ext.dispatch.tensor<readonly:tensor<256x256xbf16>> -> tensor<256x256xbf16>
- %4 = iree_tensor_ext.dispatch.tensor.load %1, offsets = [0, 0], sizes = [256, 256], strides = [1, 1] : !iree_tensor_ext.dispatch.tensor<readonly:tensor<256x256xbf16>> -> tensor<256x256xbf16>
- %5 = tensor.empty() : tensor<256x256xf32>
- %6 = linalg.fill ins(%cst : f32) outs(%5 : tensor<256x256xf32>) -> tensor<256x256xf32>
- %7 = linalg.matmul ins(%3, %4 : tensor<256x256xbf16>, tensor<256x256xbf16>) outs(%6 : tensor<256x256xf32>) -> tensor<256x256xf32>
- iree_tensor_ext.dispatch.tensor.store %7, %2, offsets = [0, 0], sizes = [256, 256], strides = [1, 1] : tensor<256x256xf32> -> !iree_tensor_ext.dispatch.tensor<writeonly:tensor<256x256xf32>>
- return
-}
-
-// CONFIG-LABEL: func.func @matmul_256x256x256_bf16_f32(
-// CONFIG: linalg.matmul {{.*}}lowering_config = #iree_gpu.lowering_config
-// CONFIG-SAME: mma_kind = #iree_gpu.mma_layout<NV_MMA_SYNC_F32_16x8x16_BF16>
-// CONFIG-SAME: promote_operands = [0, 1]
-
-// NVVM-LABEL: llvm.func @matmul_256x256x256_bf16_f32(
-// NVVM: nvvm.mma.sync
-// NVVM-SAME: multiplicandAPtxType = #nvvm.mma_type<bf16>
-// NVVM-SAME: multiplicandBPtxType = #nvvm.mma_type<bf16>
-// NVVM-SAME: shape = #nvvm.shape<m = 16, n = 8, k = 16>
-// NVVM-SAME: : (i32, i32, f32) -> !llvm.struct<(f32, f32, f32, f32)>
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/config_tile_and_fuse_sm120.mlir b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/config_tile_and_fuse_sm120.mlir
index 89cba20..684c23a 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/config_tile_and_fuse_sm120.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/config_tile_and_fuse_sm120.mlir
@@ -92,3 +92,73 @@
// CHECK: linalg.matmul {{.*}}lowering_config = #iree_gpu.lowering_config
// CHECK-SAME: convert_acc_gemm
// CHECK-SAME: mma_kind = #iree_gpu.mma_layout<NV_MMA_SYNC_F32_16x8x16_F16>, promote_operands = [0, 1], reduction = [0, 0, 8], subgroup = [2, 4, 0], workgroup = [64, 64, 0]
+
+// -----
+
+// Mixed precision f32xbf16xf32 matmuls cannot use NVIDIA MMA intrinsics, but
+// should still avoid degenerate SIMT tiling that only computes one M element
+// per workgroup when M has enough work to tile. This is the same fallback
+// behavior as sm_80 since sm_120 has not yet received architecture-specific
+// heuristic tuning.
+func.func @mixed_precision_matmul_36x151936x896_f32xbf16xf32(
+ %lhs: tensor<36x896xf32>,
+ %rhs: tensor<151936x896xbf16>) -> tensor<36x151936xf32> {
+ %cst = arith.constant 0.000000e+00 : f32
+ %empty = tensor.empty() : tensor<36x151936xf32>
+ %fill = linalg.fill ins(%cst : f32) outs(%empty : tensor<36x151936xf32>) -> tensor<36x151936xf32>
+ %result = linalg.generic {
+ indexing_maps = [
+ affine_map<(d0, d1, d2) -> (d0, d2)>,
+ affine_map<(d0, d1, d2) -> (d1, d2)>,
+ affine_map<(d0, d1, d2) -> (d0, d1)>],
+ iterator_types = ["parallel", "parallel", "reduction"]}
+ ins(%lhs, %rhs : tensor<36x896xf32>, tensor<151936x896xbf16>)
+ outs(%fill : tensor<36x151936xf32>) {
+ ^bb0(%in: f32, %in_0: bf16, %out: f32):
+ %0 = arith.extf %in_0 : bf16 to f32
+ %1 = arith.mulf %in, %0 : f32
+ %2 = arith.addf %out, %1 : f32
+ linalg.yield %2 : f32
+ } -> tensor<36x151936xf32>
+ return %result : tensor<36x151936xf32>
+}
+
+// CHECK-LABEL: func.func @mixed_precision_matmul_36x151936x896_f32xbf16xf32(
+// CHECK-SAME: #iree_codegen.translation_info<pipeline = #iree_gpu.pipeline<TileAndFuse> workgroup_size = [32, 8, 1] subgroup_size = 32
+// CHECK: linalg.generic {{.*}}lowering_config = #iree_gpu.lowering_config<{{[{]}}reduction = [0, 0, 32]
+// CHECK-SAME: thread = [1, 16, 0]
+// CHECK-SAME: workgroup = [32, 128, 1]
+
+// -----
+
+// Test that BF16 matmul selects NV_MMA_SYNC_F32_16x8x16_BF16 with TileAndFuse,
+// mirroring the sm_80 BF16 coverage in bf16_mma_sm80.mlir.
+
+#pipeline_layout_bf16 = #hal.pipeline.layout<bindings = [
+ #hal.pipeline.binding<storage_buffer>,
+ #hal.pipeline.binding<storage_buffer>,
+ #hal.pipeline.binding<storage_buffer>
+]>
+func.func @matmul_256x256x256_bf16_f32() {
+ %cst = arith.constant 0.000000e+00 : f32
+ %c0 = arith.constant 0 : index
+ %0 = hal.interface.binding.subspan layout(#pipeline_layout_bf16) binding(0) alignment(64) offset(%c0) flags(ReadOnly) : !iree_tensor_ext.dispatch.tensor<readonly:tensor<256x256xbf16>>
+ %1 = hal.interface.binding.subspan layout(#pipeline_layout_bf16) binding(1) alignment(64) offset(%c0) flags(ReadOnly) : !iree_tensor_ext.dispatch.tensor<readonly:tensor<256x256xbf16>>
+ %2 = hal.interface.binding.subspan layout(#pipeline_layout_bf16) binding(2) alignment(64) offset(%c0) : !iree_tensor_ext.dispatch.tensor<writeonly:tensor<256x256xf32>>
+ %3 = iree_tensor_ext.dispatch.tensor.load %0, offsets = [0, 0], sizes = [256, 256], strides = [1, 1] : !iree_tensor_ext.dispatch.tensor<readonly:tensor<256x256xbf16>> -> tensor<256x256xbf16>
+ %4 = iree_tensor_ext.dispatch.tensor.load %1, offsets = [0, 0], sizes = [256, 256], strides = [1, 1] : !iree_tensor_ext.dispatch.tensor<readonly:tensor<256x256xbf16>> -> tensor<256x256xbf16>
+ %5 = tensor.empty() : tensor<256x256xf32>
+ %6 = linalg.fill ins(%cst : f32) outs(%5 : tensor<256x256xf32>) -> tensor<256x256xf32>
+ %7 = linalg.matmul ins(%3, %4 : tensor<256x256xbf16>, tensor<256x256xbf16>) outs(%6 : tensor<256x256xf32>) -> tensor<256x256xf32>
+ iree_tensor_ext.dispatch.tensor.store %7, %2, offsets = [0, 0], sizes = [256, 256], strides = [1, 1] : tensor<256x256xf32> -> !iree_tensor_ext.dispatch.tensor<writeonly:tensor<256x256xf32>>
+ return
+}
+
+// CHECK-LABEL: func.func @matmul_256x256x256_bf16_f32(
+// CHECK-SAME: #iree_codegen.translation_info<pipeline = #iree_gpu.pipeline<TileAndFuse> workgroup_size = [128, 1, 1] subgroup_size = 32
+// CHECK: linalg.matmul {{.*}}lowering_config = #iree_gpu.lowering_config
+// CHECK-SAME: mma_kind = #iree_gpu.mma_layout<NV_MMA_SYNC_F32_16x8x16_BF16>
+// CHECK-SAME: promote_operands = [0, 1]
+// CHECK-SAME: reduction = [0, 0, 8]
+// CHECK-SAME: subgroup = [2, 4, 0]
+// CHECK-SAME: workgroup = [64, 64, 0]
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/config_tile_and_fuse_sm80.mlir b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/config_tile_and_fuse_sm80.mlir
index 69517d0..6bd2f94 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/config_tile_and_fuse_sm80.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/config_tile_and_fuse_sm80.mlir
@@ -133,3 +133,36 @@
// CHECK: linalg.generic {{.*}}lowering_config = #iree_gpu.lowering_config<{{[{]}}reduction = [0, 0, 32]
// CHECK-SAME: thread = [1, 16, 0]
// CHECK-SAME: workgroup = [32, 128, 1]
+
+// -----
+
+// Test that BF16 matmul selects NV_MMA_SYNC_F32_16x8x16_BF16 with TileAndFuse.
+
+#pipeline_layout_bf16 = #hal.pipeline.layout<bindings = [
+ #hal.pipeline.binding<storage_buffer>,
+ #hal.pipeline.binding<storage_buffer>,
+ #hal.pipeline.binding<storage_buffer>
+]>
+func.func @matmul_256x256x256_bf16_f32() {
+ %cst = arith.constant 0.000000e+00 : f32
+ %c0 = arith.constant 0 : index
+ %0 = hal.interface.binding.subspan layout(#pipeline_layout_bf16) binding(0) alignment(64) offset(%c0) flags(ReadOnly) : !iree_tensor_ext.dispatch.tensor<readonly:tensor<256x256xbf16>>
+ %1 = hal.interface.binding.subspan layout(#pipeline_layout_bf16) binding(1) alignment(64) offset(%c0) flags(ReadOnly) : !iree_tensor_ext.dispatch.tensor<readonly:tensor<256x256xbf16>>
+ %2 = hal.interface.binding.subspan layout(#pipeline_layout_bf16) binding(2) alignment(64) offset(%c0) : !iree_tensor_ext.dispatch.tensor<writeonly:tensor<256x256xf32>>
+ %3 = iree_tensor_ext.dispatch.tensor.load %0, offsets = [0, 0], sizes = [256, 256], strides = [1, 1] : !iree_tensor_ext.dispatch.tensor<readonly:tensor<256x256xbf16>> -> tensor<256x256xbf16>
+ %4 = iree_tensor_ext.dispatch.tensor.load %1, offsets = [0, 0], sizes = [256, 256], strides = [1, 1] : !iree_tensor_ext.dispatch.tensor<readonly:tensor<256x256xbf16>> -> tensor<256x256xbf16>
+ %5 = tensor.empty() : tensor<256x256xf32>
+ %6 = linalg.fill ins(%cst : f32) outs(%5 : tensor<256x256xf32>) -> tensor<256x256xf32>
+ %7 = linalg.matmul ins(%3, %4 : tensor<256x256xbf16>, tensor<256x256xbf16>) outs(%6 : tensor<256x256xf32>) -> tensor<256x256xf32>
+ iree_tensor_ext.dispatch.tensor.store %7, %2, offsets = [0, 0], sizes = [256, 256], strides = [1, 1] : tensor<256x256xf32> -> !iree_tensor_ext.dispatch.tensor<writeonly:tensor<256x256xf32>>
+ return
+}
+
+// CHECK-LABEL: func.func @matmul_256x256x256_bf16_f32(
+// CHECK-SAME: #iree_codegen.translation_info<pipeline = #iree_gpu.pipeline<TileAndFuse> workgroup_size = [128, 1, 1] subgroup_size = 32
+// CHECK: linalg.matmul {{.*}}lowering_config = #iree_gpu.lowering_config
+// CHECK-SAME: mma_kind = #iree_gpu.mma_layout<NV_MMA_SYNC_F32_16x8x16_BF16>
+// CHECK-SAME: promote_operands = [0, 1]
+// CHECK-SAME: reduction = [0, 0, 8]
+// CHECK-SAME: subgroup = [2, 4, 0]
+// CHECK-SAME: workgroup = [64, 64, 0]
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/pipeline_tile_and_fuse_mma_sync.mlir b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/pipeline_tile_and_fuse_mma_sync.mlir
index 5d99b25..748597b 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/pipeline_tile_and_fuse_mma_sync.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/NVVM/pipeline_tile_and_fuse_mma_sync.mlir
@@ -96,3 +96,46 @@
// CHECK: vector.shape_cast {{.*}} : vector<2x2x2xf16> to vector<4x2xf16>
// Verify nvgpu.mma.sync is generated with f16 output type
// CHECK-COUNT-8: nvgpu.mma.sync({{.*}}) {mmaShape = [16, 8, 16]} : ({{.*}}) -> vector<2x2xf16>
+
+// -----
+
+// Test BF16 matmul lowering through TileAndFuse pipeline with NV_MMA_SYNC_F32_16x8x16_BF16.
+
+#executable_target_cuda_bf16 = #hal.executable.target<"cuda", "cuda-nvptx-fb">
+#pipeline_layout_bf16 = #hal.pipeline.layout<bindings = [
+ #hal.pipeline.binding<storage_buffer>,
+ #hal.pipeline.binding<storage_buffer>,
+ #hal.pipeline.binding<storage_buffer>
+]>
+#config_bf16 = #iree_gpu.lowering_config<{
+ workgroup = [64, 64, 0],
+ reduction = [0, 0, 8],
+ subgroup = [2, 4],
+ mma_kind = #iree_gpu.mma_layout<NV_MMA_SYNC_F32_16x8x16_BF16>,
+ promote_operands = [0, 1]
+}>
+func.func @matmul_tile_and_fuse_mma_sync_bf16()
+ attributes {
+ hal.executable.target = #executable_target_cuda_bf16,
+ translation_info = #iree_codegen.translation_info<pipeline = #iree_gpu.pipeline<TileAndFuse> workgroup_size = [128, 1, 1] subgroup_size = 32>} {
+ %cst = arith.constant 0.000000e+00 : f32
+ %c0 = arith.constant 0 : index
+ %0 = hal.interface.binding.subspan layout(#pipeline_layout_bf16) binding(0) alignment(64) offset(%c0) flags(ReadOnly) : !iree_tensor_ext.dispatch.tensor<readonly:tensor<2048x1280xbf16>>
+ %1 = hal.interface.binding.subspan layout(#pipeline_layout_bf16) binding(1) alignment(64) offset(%c0) flags(ReadOnly) : !iree_tensor_ext.dispatch.tensor<readonly:tensor<1280x10240xbf16>>
+ %2 = hal.interface.binding.subspan layout(#pipeline_layout_bf16) binding(2) alignment(64) offset(%c0) : !iree_tensor_ext.dispatch.tensor<writeonly:tensor<2048x10240xf32>>
+ %3 = iree_tensor_ext.dispatch.tensor.load %0, offsets = [0, 0], sizes = [2048, 1280], strides = [1, 1] : !iree_tensor_ext.dispatch.tensor<readonly:tensor<2048x1280xbf16>> -> tensor<2048x1280xbf16>
+ %4 = iree_tensor_ext.dispatch.tensor.load %1, offsets = [0, 0], sizes = [1280, 10240], strides = [1, 1] : !iree_tensor_ext.dispatch.tensor<readonly:tensor<1280x10240xbf16>> -> tensor<1280x10240xbf16>
+ %5 = tensor.empty() : tensor<2048x10240xf32>
+ %6 = linalg.fill ins(%cst : f32) outs(%5 : tensor<2048x10240xf32>) -> tensor<2048x10240xf32>
+ %7 = linalg.matmul {lowering_config = #config_bf16} ins(%3, %4 : tensor<2048x1280xbf16>, tensor<1280x10240xbf16>) outs(%6 : tensor<2048x10240xf32>) -> tensor<2048x10240xf32>
+ iree_tensor_ext.dispatch.tensor.store %7, %2, offsets = [0, 0], sizes = [2048, 10240], strides = [1, 1] : tensor<2048x10240xf32> -> !iree_tensor_ext.dispatch.tensor<writeonly:tensor<2048x10240xf32>>
+ return
+}
+
+// CHECK-LABEL: func @matmul_tile_and_fuse_mma_sync_bf16
+// CHECK-DAG: memref.alloc() : memref<{{.*}}xbf16, #gpu.address_space<workgroup>>
+// CHECK-DAG: memref.alloc() : memref<{{.*}}xbf16, #gpu.address_space<workgroup>>
+// CHECK: scf.for
+// Verify nvgpu.mma.sync is generated with correct shape for BF16
+// CHECK-COUNT-8: nvgpu.mma.sync({{.*}}) {mmaShape = [16, 8, 16]} : ({{.*}}) -> vector<2x2xf32>
+// CHECK: scf.yield
