[Codegen][ROCDL] Add support for nhwc convolution with mfma (#16579)
This supports convolutions using the matmul pipeline by simply tiling
the kernel dimensions to 1 and then folding them away, leaving a
contraction along the remaining image/input channel dimensions. The
config selection logic is kept the same as matmul for now, ignoring the
additional dimensions.
This only support nhwc convolutions, however supporting nchw
convolutions is simply a matter of supporting matmul_transpose_a
variants. This also does not support strided convolutions because this
approach will fail to vectorize in such cases.
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/KernelConfig.cpp b/compiler/src/iree/compiler/Codegen/LLVMGPU/KernelConfig.cpp
index 7e46f1e..48fa797 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/KernelConfig.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/KernelConfig.cpp
@@ -84,6 +84,10 @@
} // namespace
+//====---------------------------------------------------------------------===//
+// Matmul Configuration Helpers
+//====---------------------------------------------------------------------===//
+
/// Return the best combination of tile size and wg size. It will then used to
/// pick the best size aligned with the shape dimension.
static void getMatmulConfig(SmallVectorImpl<TileWorkgroupSizePair> &tileSizes) {
@@ -286,13 +290,162 @@
return codegenPipeline;
}
+//====---------------------------------------------------------------------===//
+// Vector Distribution Contraction/Convolution Pipeline Configuration
+//====---------------------------------------------------------------------===//
+
static LogicalResult
-setVectorDistributionConfig(mlir::FunctionOpInterface entryPoint,
- linalg::LinalgOp op, const TargetInfo &targetInfo) {
- if (!isMatmulOrBatchMatmul(op)) {
+setConvolutionVectorDistributionConfig(mlir::FunctionOpInterface entryPoint,
+ linalg::LinalgOp op,
+ const TargetInfo &targetInfo) {
+ FailureOr<ArrayAttr> mmaKinds = getSupportedMmaTypes(entryPoint);
+ if (failed(mmaKinds)) {
return failure();
}
+ // This pipeline needs to know the subgroup size for distributing to virtual
+ // lane IDs.
+ if (targetInfo.supportedSubgroupSizes.empty()) {
+ return failure();
+ }
+ const int64_t subgroupSize = targetInfo.supportedSubgroupSizes.front();
+
+ SmallVector<int64_t, 4> bounds = op.getStaticLoopRanges();
+ FailureOr<mlir::linalg::ConvolutionDimensions> convolutionDims =
+ mlir::linalg::inferConvolutionDims(op);
+ assert(succeeded(convolutionDims) && "Could not infer contraction dims");
+
+ // This strategy turns non-strided/dilated convolution problems into matmul
+ // problems by tiling certain dimensions to 1:
+ // - Batch dimensions (parallel shared by the image and output)
+ // - Filter dimensions (reduction on the filter, and convolved on the image)
+ // - All output image dimensions except the outermost one
+ //
+ // After this, the remaining non-unit dimensions are:
+ // - One output image dimension corresponding to the M dimension of a matmul.
+ // - The output channel dimension, corresponding to the N dimension.
+ // - The input channel dimension, corresponding to the K dimension.
+
+ // TODO: Relax this condition to strictly alignment requirements.
+ if (convolutionDims->outputChannel.size() != 1 ||
+ convolutionDims->inputChannel.size() != 1 ||
+ convolutionDims->filterLoop.size() < 1 ||
+ convolutionDims->outputImage.size() < 1 ||
+ convolutionDims->depth.size() != 0) {
+ return failure();
+ }
+
+ auto isAllOnesList = [](ArrayRef<int64_t> list) {
+ return llvm::all_of(list, [](int64_t i) { return i == 1; });
+ };
+
+ // TODO: Support non-unit strides/dilations.
+ if (!isAllOnesList(convolutionDims->strides) ||
+ !isAllOnesList(convolutionDims->dilations)) {
+ return failure();
+ }
+
+ int64_t mDim = convolutionDims->outputImage.back();
+ int64_t nDim = convolutionDims->outputChannel.front();
+ // TODO: Support NCHW convolutions. This is just a matmul_transpose_a, however
+ // the distribution patterns currently do not support that variant.
+ if (mDim > nDim) {
+ return failure();
+ }
+ int64_t kDim = convolutionDims->inputChannel.front();
+
+ Value lhs = op.getDpsInputOperand(0)->get();
+ Value rhs = op.getDpsInputOperand(1)->get();
+ Value init = op.getDpsInitOperand(0)->get();
+
+ Type lhsElemType = getElementTypeOrSelf(lhs);
+ Type rhsElemType = getElementTypeOrSelf(rhs);
+ Type initElemType = getElementTypeOrSelf(init);
+
+ GPUMatmulShapeType problem{bounds[mDim], bounds[nDim], bounds[kDim],
+ lhsElemType, rhsElemType, initElemType};
+
+ auto mmaAttrs = llvm::to_vector(mmaKinds->getAsRange<IREE::GPU::MmaAttr>());
+ SmallVector<GPUMatmulShapeType> intrinsics;
+ intrinsics.reserve(mmaKinds->size());
+ for (auto mma : mmaAttrs) {
+ auto [mSize, nSize, kSize] = mma.getMNKShape();
+ auto [aType, bType, cType] = mma.getABCElementTypes();
+ intrinsics.emplace_back(mSize, nSize, kSize, aType, bType, cType);
+ }
+
+ // Note that the following heuristic seeds are just placeholder values.
+ // We need to clean it up and make it adjusting to different targets.
+ // See https://github.com/openxla/iree/issues/16341 for details.
+ GPUMMAHeuristicSeeds seeds{/*bestSubgroupCountPerWorkgroup=*/4,
+ /*bestMNTileCountPerSubgroup=*/8,
+ /*bestKTileCountPerSubgroup=*/2};
+
+ // First try to find a schedule with an exactly matching intrinsic.
+ std::optional<GPUMMASchedule> schedule =
+ deduceMMASchedule(problem, intrinsics, seeds);
+ if (!schedule) {
+ // Then try again by allowing upcasting accumulator.
+ schedule =
+ deduceMMASchedule(problem, intrinsics, seeds, /*canUpcastAcc=*/true);
+ }
+ if (!schedule) {
+ return failure();
+ }
+
+ std::array<int64_t, 3> workgroupSize{schedule->nWarpCount * subgroupSize,
+ schedule->mWarpCount, 1};
+
+ SmallVector<int64_t> workgroupTileSizes(op.getNumLoops(), 0);
+ // Tile all batch dimensions with unit size.
+ for (int64_t batch : convolutionDims->batch) {
+ workgroupTileSizes[batch] = 1;
+ }
+ // Tile all output image dimensions with unit size except the last one.
+ for (int64_t oi :
+ ArrayRef<unsigned int>(convolutionDims->outputImage).drop_back()) {
+ workgroupTileSizes[oi] = 1;
+ }
+ // Compute the M/N dimension tile size by multiply subgroup information.
+ workgroupTileSizes[mDim] =
+ schedule->mWarpCount * schedule->mTileCount * schedule->mSize;
+ workgroupTileSizes[nDim] =
+ schedule->nWarpCount * schedule->nTileCount * schedule->nSize;
+
+ // Follow the LLVMGPU convention of keeping all of the tile sizes in one list.
+ workgroupTileSizes[kDim] = schedule->kTileCount * schedule->kSize;
+
+ // Tile all filter loop dimensions to 1.
+ for (int64_t filterDim : convolutionDims->filterLoop) {
+ workgroupTileSizes[filterDim] = 1;
+ }
+
+ TileSizesListType tileSizes;
+ tileSizes.push_back(workgroupTileSizes);
+
+ // Attach the MMA schedule as an attribute to the entry point export function
+ // for later access in the pipeline.
+ MLIRContext *context = op.getContext();
+ auto scheduleAttr = IREE::GPU::MMAScheduleAttr::get(
+ context, mmaAttrs[schedule->index], schedule->mWarpCount,
+ schedule->nWarpCount, schedule->mTileCount, schedule->nTileCount,
+ schedule->kTileCount);
+ SmallVector<NamedAttribute, 1> attrs;
+ attrs.emplace_back(
+ StringAttr::get(context, IREE::GPU::MMAScheduleAttr::getMnemonic()),
+ scheduleAttr);
+ auto configDict = DictionaryAttr::get(context, attrs);
+
+ return setOpConfigAndEntryPointFnTranslation(
+ entryPoint, op, tileSizes,
+ IREE::Codegen::DispatchLoweringPassPipeline::LLVMGPUVectorDistribute,
+ workgroupSize, subgroupSize, configDict);
+}
+
+static LogicalResult
+setMatmulVectorDistributionConfig(mlir::FunctionOpInterface entryPoint,
+ linalg::LinalgOp op,
+ const TargetInfo &targetInfo) {
FailureOr<ArrayAttr> mmaKinds = getSupportedMmaTypes(entryPoint);
if (failed(mmaKinds)) {
return failure();
@@ -398,6 +551,24 @@
workgroupSize, subgroupSize, configDict);
}
+static LogicalResult
+setVectorDistributionConfig(mlir::FunctionOpInterface entryPoint,
+ linalg::LinalgOp linalgOp,
+ const TargetInfo &targetInfo) {
+ if (isMatmulOrBatchMatmul(linalgOp)) {
+ return setMatmulVectorDistributionConfig(entryPoint, linalgOp, targetInfo);
+ }
+ if (isa<linalg::ConvolutionOpInterface>(*linalgOp)) {
+ return setConvolutionVectorDistributionConfig(entryPoint, linalgOp,
+ targetInfo);
+ }
+ return failure();
+}
+
+//====---------------------------------------------------------------------===//
+// Contraction Pipeline Configuration
+//====---------------------------------------------------------------------===//
+
static LogicalResult setContractConfig(mlir::FunctionOpInterface entryPoint,
linalg::LinalgOp op,
const TargetInfo &targetInfo) {
@@ -587,6 +758,10 @@
IREE::Codegen::DispatchLoweringPassPipeline::LLVMGPUMatmulSimt);
}
+//====---------------------------------------------------------------------===//
+// FFT Pipeline Configuration
+//====---------------------------------------------------------------------===//
+
static LogicalResult setFftConfig(mlir::FunctionOpInterface entryPoint,
IREE::LinalgExt::FftOp op,
const TargetInfo &targetInfo) {
@@ -619,6 +794,10 @@
workgroupSize);
}
+//====---------------------------------------------------------------------===//
+// Sort Pipeline Configuration
+//====---------------------------------------------------------------------===//
+
static LogicalResult setSortConfig(mlir::FunctionOpInterface entryPoint,
Operation *op,
const TargetInfo &targetInfo) {
@@ -661,6 +840,10 @@
workgroupSize);
}
+//====---------------------------------------------------------------------===//
+// Pack/Unpack Pipeline Configuration
+//====---------------------------------------------------------------------===//
+
static SmallVector<int64_t>
getDefaultWorkgroupTileSizesForPackUnPack(TilingInterface op,
int64_t defaultSize) {
@@ -707,6 +890,10 @@
workgroupSizes);
}
+//====---------------------------------------------------------------------===//
+// Default Pipeline Configuration
+//====---------------------------------------------------------------------===//
+
// Basic default properties for linalg ops that haven't been tuned.
static LogicalResult setRootDefaultConfig(mlir::FunctionOpInterface entryPoint,
Operation *op,
@@ -839,6 +1026,10 @@
targetInfo.supportedSubgroupSizes.front());
}
+//====---------------------------------------------------------------------===//
+// Transform Dialect Pipeline Configuration
+//====---------------------------------------------------------------------===//
+
/// Set configuration for transform dialect based strategies.
static LogicalResult
setTransformDialectConfig(mlir::FunctionOpInterface entryPoint, Operation *op,
@@ -918,6 +1109,10 @@
return true;
}
+//====---------------------------------------------------------------------===//
+// Warp Reduction Pipeline Configuration
+//====---------------------------------------------------------------------===//
+
/// Set the configuration for reductions that can be mapped to warp reductions.
static LogicalResult
setWarpReductionConfig(mlir::FunctionOpInterface entryPoint,
@@ -1131,6 +1326,10 @@
linalgOp.getNumParallelLoops() <= 3;
}
+//====---------------------------------------------------------------------===//
+// Transpose Pipeline Configuration
+//====---------------------------------------------------------------------===//
+
static LogicalResult setTransposeConfig(mlir::FunctionOpInterface entryPoint,
linalg::LinalgOp linalgOp) {
LinalgOpInfo opInfo(linalgOp, sharedMemTransposeFilter);
@@ -1186,6 +1385,10 @@
workgroupSize);
}
+//====---------------------------------------------------------------------===//
+// UKernel Pipeline Configuration
+//====---------------------------------------------------------------------===//
+
/// Set the configuration for argmax that can be mapped to argmax uKernel.
/// Distribute all parallel dim across different workgroups, and only use single
/// subgroup per workgroup.
@@ -1342,6 +1545,10 @@
return false;
}
+//====---------------------------------------------------------------------===//
+// Convolution Pipeline Configuration
+//====---------------------------------------------------------------------===//
+
static LogicalResult setConvolutionConfig(linalg::LinalgOp linalgOp,
const int64_t subgroupSize,
const int64_t bestTilingFactor) {
@@ -1435,6 +1642,10 @@
pipeline, workgroupSize);
}
+//====---------------------------------------------------------------------===//
+// Pipeline Configuration
+//====---------------------------------------------------------------------===//
+
static LogicalResult setRootConfig(mlir::FunctionOpInterface entryPointFn,
Operation *computeOp) {
TargetInfo targetInfo = getTargetInfo(entryPointFn);
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/Passes.cpp b/compiler/src/iree/compiler/Codegen/LLVMGPU/Passes.cpp
index 57a936c..7920f21 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/Passes.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/Passes.cpp
@@ -496,6 +496,20 @@
nestedModulePM.addNestedPass<func::FuncOp>(
createGPUTensorTileToSerialLoops());
+ // Generalize convolutions and fold away unit extent dims. All convolutions
+ // are expected to have the kernel dimensions tiled to 1 by this point, so
+ // folding unit dims like this directly maps it to a matrix multiplication.
+ // After vectorization we expect to get a pure matmul (or a transposed
+ // variant) as a `vector.contract`.
+ nestedModulePM.addNestedPass<func::FuncOp>(
+ createGPUGeneralizeNamedConvolutionOpsPass());
+ LinalgFoldUnitExtentDimsPassOptions options;
+ options.useRankReducingSlices = true;
+ nestedModulePM.addNestedPass<func::FuncOp>(
+ mlir::createLinalgFoldUnitExtentDimsPass(options));
+ nestedModulePM.addPass(createCanonicalizerPass());
+ nestedModulePM.addPass(createCSEPass());
+
nestedModulePM.addNestedPass<func::FuncOp>(
createOptimizeTensorInsertExtractSlicesPass());
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/config_vector_distribute.mlir b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/config_vector_distribute.mlir
index 4f739e1..169e7c1 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/config_vector_distribute.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/config_vector_distribute.mlir
@@ -5,6 +5,54 @@
// to be migrated to the rocdl heuristics, but for now is just physically
// located here.
+// CHECK: #[[$TILE_SIZES:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[1, 1, 64, 128, 1, 1, 32]{{\]}}
+// CHECK: #iree_codegen.translation_info<LLVMGPUVectorDistribute
+// CHECK-SAME: mma_schedule = #iree_gpu.mma_schedule
+// CHECK-SAME: intrinsic = #iree_gpu.mfma_layout<F16_16x16x16_F32>
+// CHECK-SAME: subgroup_m_count = 2, subgroup_n_count = 2,
+// CHECK-SAME: subgroup_m_tile_count = 2, subgroup_n_tile_count = 4, subgroup_k_tile_count = 2
+
+#pipeline_layout = #hal.pipeline.layout<push_constants = 0, sets = [
+ #hal.descriptor_set.layout<0, bindings = [
+ #hal.descriptor_set.binding<0, storage_buffer>,
+ #hal.descriptor_set.binding<1, storage_buffer>
+ ]>
+]>
+hal.executable @conv_nhwc_dispatch_0 {
+hal.executable.variant @rocm target(<"rocm", "rocm-hsaco-fb", {
+ target_arch = "gfx940",
+ mma_intrinsics = [#iree_gpu.mfma_layout<F16_16x16x16_F32>,
+ #iree_gpu.mfma_layout<F16_32x32x8_F32>]
+ }>) {
+ hal.executable.export @conv_nhwc layout(#pipeline_layout) {
+ ^bb0(%arg0: !hal.device, %arg1: index, %arg2 : index):
+ %x, %y, %z = flow.dispatch.workgroup_count_from_dag_root %arg1, %arg2
+ hal.return %x, %y, %z : index, index, index
+ }
+ builtin.module {
+ func.func @conv_nhwc() {
+ %c0 = arith.constant 0 : index
+ %cst = arith.constant 0.000000e+00 : f32
+ %0 = hal.interface.binding.subspan set(0) binding(0) type(storage_buffer) alignment(64) offset(%c0) : !flow.dispatch.tensor<readonly:tensor<2x258x514x768xf16>>
+ %1 = hal.interface.binding.subspan set(0) binding(1) type(storage_buffer) alignment(64) offset(%c0) : !flow.dispatch.tensor<readonly:tensor<3x3x768x256xf16>>
+ %2 = hal.interface.binding.subspan set(0) binding(2) type(storage_buffer) alignment(64) offset(%c0) : !flow.dispatch.tensor<writeonly:tensor<2x256x512x256xf32>>
+ %3 = flow.dispatch.tensor.load %0, offsets = [0, 0, 0, 0], sizes = [2, 258, 514, 768], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2x258x514x768xf16>> -> tensor<2x258x514x768xf16>
+ %4 = flow.dispatch.tensor.load %1, offsets = [0, 0, 0, 0], sizes = [3, 3, 768, 256], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<3x3x768x256xf16>> -> tensor<3x3x768x256xf16>
+ %5 = tensor.empty() : tensor<2x256x512x256xf32>
+ %6 = linalg.fill ins(%cst : f32) outs(%5 : tensor<2x256x512x256xf32>) -> tensor<2x256x512x256xf32>
+ %7 = linalg.conv_2d_nhwc_hwcf {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} ins(%3, %4 : tensor<2x258x514x768xf16>, tensor<3x3x768x256xf16>) outs(%6 : tensor<2x256x512x256xf32>) -> tensor<2x256x512x256xf32>
+ flow.dispatch.tensor.store %7, %2, offsets = [0, 0, 0, 0], sizes = [2, 256, 512, 256], strides = [1, 1, 1, 1] : tensor<2x256x512x256xf32> -> !flow.dispatch.tensor<writeonly:tensor<2x256x512x256xf32>>
+ return
+ }
+ }
+ }
+}
+
+// CHECK-LABEL: hal.executable public @conv_nhwc_dispatch_0
+// CHECK: linalg.conv_2d_nhwc_hwcf {{.*}} lowering_config = #[[$TILE_SIZES]]
+
+// -----
+
#pipeline_layout = #hal.pipeline.layout<push_constants = 0, sets = [
#hal.descriptor_set.layout<0, bindings = [
#hal.descriptor_set.binding<0, storage_buffer>,
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/pipeline_vector_distribute.mlir b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/pipeline_vector_distribute.mlir
index 5abc043..31bdc44 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/pipeline_vector_distribute.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/pipeline_vector_distribute.mlir
@@ -114,3 +114,48 @@
// CHECK: %[[TRUNC:.+]] = arith.truncf %157 : vector<2x4x1x1x1x4xf32> to vector<2x4x1x1x1x4xf16>
// CHECK: scf.yield %[[TRUNC]] : vector<2x4x1x1x1x4xf16>
// CHECK-COUNT-8: vector.transfer_write {{.+}} {in_bounds = [true, true]} : vector<4x1xf16>, memref<256x256xf16, #hal.descriptor_type<storage_buffer>>
+
+// -----
+
+#pipeline_layout = #hal.pipeline.layout<push_constants = 0, sets = [
+ #hal.descriptor_set.layout<0, bindings = [
+ #hal.descriptor_set.binding<0, storage_buffer>,
+ #hal.descriptor_set.binding<1, storage_buffer>
+ ]>
+]>
+hal.executable @conv_nhwc_dispatch_0 {
+hal.executable.variant @rocm target(<"rocm", "rocm-hsaco-fb", {
+ target_arch = "gfx940",
+ mma_intrinsics = [#iree_gpu.mfma_layout<F16_16x16x16_F32>,
+ #iree_gpu.mfma_layout<F16_32x32x8_F32>]
+ }>) {
+ hal.executable.export @conv_nhwc layout(#pipeline_layout) {
+ ^bb0(%arg0: !hal.device):
+ %x, %y, %z = flow.dispatch.workgroup_count_from_slice
+ hal.return %x, %y, %z : index, index, index
+ }
+ builtin.module {
+ func.func @conv_nhwc() {
+ %c0 = arith.constant 0 : index
+ %cst = arith.constant 0.000000e+00 : f32
+ %0 = hal.interface.binding.subspan set(0) binding(0) type(storage_buffer) alignment(64) offset(%c0) : !flow.dispatch.tensor<readonly:tensor<2x258x514x768xf16>>
+ %1 = hal.interface.binding.subspan set(0) binding(1) type(storage_buffer) alignment(64) offset(%c0) : !flow.dispatch.tensor<readonly:tensor<3x3x768x256xf16>>
+ %2 = hal.interface.binding.subspan set(0) binding(2) type(storage_buffer) alignment(64) offset(%c0) : !flow.dispatch.tensor<writeonly:tensor<2x256x512x256xf32>>
+ %3 = flow.dispatch.tensor.load %0, offsets = [0, 0, 0, 0], sizes = [2, 258, 514, 768], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2x258x514x768xf16>> -> tensor<2x258x514x768xf16>
+ %4 = flow.dispatch.tensor.load %1, offsets = [0, 0, 0, 0], sizes = [3, 3, 768, 256], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<3x3x768x256xf16>> -> tensor<3x3x768x256xf16>
+ %5 = tensor.empty() : tensor<2x256x512x256xf32>
+ %6 = linalg.fill ins(%cst : f32) outs(%5 : tensor<2x256x512x256xf32>) -> tensor<2x256x512x256xf32>
+ %7 = linalg.conv_2d_nhwc_hwcf {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} ins(%3, %4 : tensor<2x258x514x768xf16>, tensor<3x3x768x256xf16>) outs(%6 : tensor<2x256x512x256xf32>) -> tensor<2x256x512x256xf32>
+ flow.dispatch.tensor.store %7, %2, offsets = [0, 0, 0, 0], sizes = [2, 256, 512, 256], strides = [1, 1, 1, 1] : tensor<2x256x512x256xf32> -> !flow.dispatch.tensor<writeonly:tensor<2x256x512x256xf32>>
+ return
+ }
+ }
+ }
+}
+
+// CHECK-LABEL: func.func @conv_nhwc
+// CHECK: scf.for {{.*}} = %c0 to %c3
+// CHECK: scf.for {{.*}} = %c0 to %c3
+// CHECK: scf.for {{.*}} = %c0 to %c768 step %c32 iter_args(%[[ARG:.+]] = {{.*}}) -> (vector<2x4x1x1x1x4xf32>)
+// CHECK-COUNT-16: amdgpu.mfma {{.*}} {blocks = 1 : i32, k = 16 : i32, m = 16 : i32, n = 16 : i32} blgp = none : vector<4xf16>, vector<4xf16>, vector<4xf32>
+// CHECK-COUNT-8: vector.transfer_write {{.+}} : vector<4x1xf32>, memref<2x256x512x256xf32, #hal.descriptor_type<storage_buffer>>
