Rename `unroll_{m,n,k}` to `intrinsics_{m,n,k}` (#19652)
In `DataTiledMMAAttr`, we had those `unroll_{m,n,k}` parameters which
were questionably named as they control the tile size, counter in
"number of intrinsics, stacked along each dimension".
They were passed down, still under that same name, to the ukernel, where
they sounded awkward too. They sounded like unrolling factors when in
that context whey really were loop trip counts.
---------
Signed-off-by: Benoit Jacob <jacob.benoit.1@gmail.com>diff --git a/compiler/plugins/target/ROCM/builtins/ukernel/iree_uk_amdgpu_multi_mma_mfma_i32_16x16x32_i8.c b/compiler/plugins/target/ROCM/builtins/ukernel/iree_uk_amdgpu_multi_mma_mfma_i32_16x16x32_i8.c
index 064bedb..c329270 100644
--- a/compiler/plugins/target/ROCM/builtins/ukernel/iree_uk_amdgpu_multi_mma_mfma_i32_16x16x32_i8.c
+++ b/compiler/plugins/target/ROCM/builtins/ukernel/iree_uk_amdgpu_multi_mma_mfma_i32_16x16x32_i8.c
@@ -14,14 +14,14 @@
[[clang::always_inline]] void iree_uk_amdgpu_multi_mma_mfma_i32_16x16x32_i8(
const int8_t *a_buffer, int64_t a_offset, const int8_t *b_buffer,
int64_t b_offset, int32_t *c_buffer, int64_t c_offset, int32_t k_size,
- int32_t unroll_m, int32_t subgroups_m, int32_t unroll_n,
- int32_t subgroups_n, int32_t unroll_k) {
+ int32_t intrinsics_m, int32_t subgroups_m, int32_t intrinsics_n,
+ int32_t subgroups_n, int32_t intrinsics_k) {
// Load existing accumulators. The VLA becomes a normal array after inlining.
- int32x4_t c[unroll_m][unroll_n];
+ int32x4_t c[intrinsics_m][intrinsics_n];
int32x4_t *c_global = (int32x4_t *)(c_buffer + c_offset);
- for (int m = 0; m < unroll_m; ++m) {
- for (int n = 0; n < unroll_n; ++n) {
- c[m][n] = c_global[64 * (m * unroll_n + n)];
+ for (int m = 0; m < intrinsics_m; ++m) {
+ for (int n = 0; n < intrinsics_n; ++n) {
+ c[m][n] = c_global[64 * (m * intrinsics_n + n)];
}
}
@@ -29,23 +29,23 @@
const int64_t *a_global = (const int64_t *)(a_buffer + a_offset);
const int64_t *b_global = (const int64_t *)(b_buffer + b_offset);
for (int k_outer = 0; k_outer < k_size; ++k_outer) {
- for (int m = 0; m < unroll_m; ++m) {
- for (int n = 0; n < unroll_n; ++n) {
- for (int k = 0; k < unroll_k; ++k) {
+ for (int m = 0; m < intrinsics_m; ++m) {
+ for (int n = 0; n < intrinsics_n; ++n) {
+ for (int k = 0; k < intrinsics_k; ++k) {
c[m][n] = __builtin_amdgcn_mfma_i32_16x16x32_i8(
- a_global[64 * unroll_k * m + k], b_global[64 * unroll_k * n + k],
- c[m][n], 0, 0, 0);
+ a_global[64 * intrinsics_k * m + k],
+ b_global[64 * intrinsics_k * n + k], c[m][n], 0, 0, 0);
}
}
}
- a_global += 64 * unroll_m * subgroups_m * unroll_k;
- b_global += 64 * unroll_n * subgroups_n * unroll_k;
+ a_global += 64 * intrinsics_m * subgroups_m * intrinsics_k;
+ b_global += 64 * intrinsics_n * subgroups_n * intrinsics_k;
}
// Store accumulators.
- for (int m = 0; m < unroll_m; ++m) {
- for (int n = 0; n < unroll_n; ++n) {
- c_global[64 * (m * unroll_n + n)] = c[m][n];
+ for (int m = 0; m < intrinsics_m; ++m) {
+ for (int n = 0; n < intrinsics_n; ++n) {
+ c_global[64 * (m * intrinsics_n + n)] = c[m][n];
}
}
}
diff --git a/compiler/plugins/target/ROCM/test/config_ukernel_multi_mma_gfx942.mlir b/compiler/plugins/target/ROCM/test/config_ukernel_multi_mma_gfx942.mlir
index 2fd7813..654dc54 100644
--- a/compiler/plugins/target/ROCM/test/config_ukernel_multi_mma_gfx942.mlir
+++ b/compiler/plugins/target/ROCM/test/config_ukernel_multi_mma_gfx942.mlir
@@ -16,7 +16,7 @@
#iree_gpu.iterator_type<reduction>
], kind = #iree_gpu.data_tiled_mma_layout<
intrinsic = MFMA_I32_16x16x32_I8,
- unroll_m = 8, unroll_n = 2, subgroups_n = 4, unroll_k = 2
+ intrinsics_m = 8, intrinsics_n = 2, subgroups_n = 4, intrinsics_k = 2
>} : tensor<1x2x8x4x16x2x8xi8>, tensor<1x2x4x2x4x16x2x8xi8> into tensor<1x1x8x4x2x4x16x4xi32>
return %d : tensor<1x1x8x4x2x4x16x4xi32>
}
diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/GPULowerToUKernels.cpp b/compiler/src/iree/compiler/Codegen/Common/GPU/GPULowerToUKernels.cpp
index 86df9f7..9cf3bd6 100644
--- a/compiler/src/iree/compiler/Codegen/Common/GPU/GPULowerToUKernels.cpp
+++ b/compiler/src/iree/compiler/Codegen/Common/GPU/GPULowerToUKernels.cpp
@@ -110,15 +110,16 @@
};
Value k = castIndexToI32(
rewriter.create<tensor::DimOp>(op.getLoc(), op.getLhs(), 1));
- Value unrollM = constI32(mma.getUnrollM());
+ Value intrinsicsM = constI32(mma.getIntrinsicsM());
Value subgroupsM = constI32(mma.getSubgroupsM());
- Value unrollN = constI32(mma.getUnrollN());
+ Value intrinsicsN = constI32(mma.getIntrinsicsN());
Value subgroupsN = constI32(mma.getSubgroupsN());
- Value unrollK = constI32(mma.getUnrollK());
+ Value intrinsicsK = constI32(mma.getIntrinsicsK());
rewriter.replaceOpWithNewOp<IREE::Codegen::UKernelGenericOp>(
op, TypeRange{op.getAccType()}, ukernelAttr.getName(),
ValueRange{op.getLhs(), op.getRhs()}, op.getAcc(),
- ValueRange{k, unrollM, subgroupsM, unrollN, subgroupsN, unrollK},
+ ValueRange{k, intrinsicsM, subgroupsM, intrinsicsN, subgroupsN,
+ intrinsicsK},
ukernelAttr.getDefAttrs(),
/*strided_outer_dims=*/rewriter.getIndexAttr(0));
return success();
diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_lower_to_ukernels.mlir b/compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_lower_to_ukernels.mlir
index bc9331f..2095393 100644
--- a/compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_lower_to_ukernels.mlir
+++ b/compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_lower_to_ukernels.mlir
@@ -73,7 +73,7 @@
%d = iree_gpu.multi_mma %a, %b, %c {
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 = [#iree_gpu.iterator_type<parallel>, #iree_gpu.iterator_type<parallel>, #iree_gpu.iterator_type<reduction>],
- kind = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_I32_16x16x32_I8, unroll_m = 8, unroll_n = 2, subgroups_n = 4, unroll_k = 2>,
+ kind = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_I32_16x16x32_I8, intrinsics_m = 8, intrinsics_n = 2, subgroups_n = 4, intrinsics_k = 2>,
lowering_config = #iree_gpu.lowering_config<{
reduction = [0, 0, 0],
ukernel = #iree_gpu.ukernel_config<name = "iree_uk_amdgpu_multi_mma_mfma_i32_16x16x32_i8", def_attrs = {vm.import.module = "rocm"}>,
diff --git a/compiler/src/iree/compiler/Codegen/Common/test/gpu_materialize_encoding_gfx1100.mlir b/compiler/src/iree/compiler/Codegen/Common/test/gpu_materialize_encoding_gfx1100.mlir
index 645fd71..a942751 100644
--- a/compiler/src/iree/compiler/Codegen/Common/test/gpu_materialize_encoding_gfx1100.mlir
+++ b/compiler/src/iree/compiler/Codegen/Common/test/gpu_materialize_encoding_gfx1100.mlir
@@ -56,5 +56,5 @@
// CHECK: %[[MMA:.+]] = iree_gpu.multi_mma %[[LHS]], %[[RHS]], %[[ACC]]
// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]],
// CHECK-SAME: iterator_types = [#iree_gpu.iterator_type<parallel>, #iree_gpu.iterator_type<parallel>, #iree_gpu.iterator_type<reduction>]
-// CHECK-SAME: kind = #iree_gpu.data_tiled_mma_layout<intrinsic = WMMA_F32_16x16x16_F16, unroll_m = 4, subgroups_n = 4>
+// CHECK-SAME: kind = #iree_gpu.data_tiled_mma_layout<intrinsic = WMMA_F32_16x16x16_F16, intrinsics_m = 4, subgroups_n = 4>
// CHECK: flow.dispatch.tensor.store %[[MMA]], %[[ACC_BINDING]]
diff --git a/compiler/src/iree/compiler/Codegen/Common/test/gpu_materialize_encoding_gfx908.mlir b/compiler/src/iree/compiler/Codegen/Common/test/gpu_materialize_encoding_gfx908.mlir
index a9fc2bc..cf54195 100644
--- a/compiler/src/iree/compiler/Codegen/Common/test/gpu_materialize_encoding_gfx908.mlir
+++ b/compiler/src/iree/compiler/Codegen/Common/test/gpu_materialize_encoding_gfx908.mlir
@@ -56,5 +56,5 @@
// CHECK: %[[MMA:.+]] = iree_gpu.multi_mma %[[LHS]], %[[RHS]], %[[ACC]]
// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]],
// CHECK-SAME: iterator_types = [#iree_gpu.iterator_type<parallel>, #iree_gpu.iterator_type<parallel>, #iree_gpu.iterator_type<reduction>]
-// CHECK-SAME: kind = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_I32_16x16x16_I8, unroll_m = 4, unroll_n = 2, subgroups_n = 4, unroll_k = 4>
+// CHECK-SAME: kind = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_I32_16x16x16_I8, intrinsics_m = 4, intrinsics_n = 2, subgroups_n = 4, intrinsics_k = 4>
// CHECK: flow.dispatch.tensor.store %[[MMA]], %[[ACC_BINDING]]
diff --git a/compiler/src/iree/compiler/Codegen/Common/test/gpu_materialize_encoding_gfx90a.mlir b/compiler/src/iree/compiler/Codegen/Common/test/gpu_materialize_encoding_gfx90a.mlir
index 89fe357..86388f1 100644
--- a/compiler/src/iree/compiler/Codegen/Common/test/gpu_materialize_encoding_gfx90a.mlir
+++ b/compiler/src/iree/compiler/Codegen/Common/test/gpu_materialize_encoding_gfx90a.mlir
@@ -56,7 +56,7 @@
// CHECK: %[[MMA:.+]] = iree_gpu.multi_mma %[[LHS]], %[[RHS]], %[[ACC]]
// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]],
// CHECK-SAME: iterator_types = [#iree_gpu.iterator_type<parallel>, #iree_gpu.iterator_type<parallel>, #iree_gpu.iterator_type<reduction>]
-// CHECK-SAME: kind = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_F32_16x16x8_BF16, unroll_m = 4, unroll_n = 2, subgroups_n = 4, unroll_k = 4>
+// CHECK-SAME: kind = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_F32_16x16x8_BF16, intrinsics_m = 4, intrinsics_n = 2, subgroups_n = 4, intrinsics_k = 4>
// CHECK: flow.dispatch.tensor.store %[[MMA]], %[[ACC_BINDING]]
// -----
@@ -115,5 +115,5 @@
// CHECK: %[[MMA:.+]] = iree_gpu.multi_mma %[[LHS]], %[[RHS]], %[[ACC]]
// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]],
// CHECK-SAME: iterator_types = [#iree_gpu.iterator_type<parallel>, #iree_gpu.iterator_type<parallel>, #iree_gpu.iterator_type<reduction>]
-// CHECK-SAME: kind = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_F64_16x16x4_F64, unroll_m = 4, subgroups_n = 4, unroll_k = 2>
+// CHECK-SAME: kind = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_F64_16x16x4_F64, intrinsics_m = 4, subgroups_n = 4, intrinsics_k = 2>
// CHECK: flow.dispatch.tensor.store %[[MMA]], %[[ACC_BINDING]]
diff --git a/compiler/src/iree/compiler/Codegen/Common/test/gpu_materialize_encoding_gfx942.mlir b/compiler/src/iree/compiler/Codegen/Common/test/gpu_materialize_encoding_gfx942.mlir
index 392ed29..bd70cad 100644
--- a/compiler/src/iree/compiler/Codegen/Common/test/gpu_materialize_encoding_gfx942.mlir
+++ b/compiler/src/iree/compiler/Codegen/Common/test/gpu_materialize_encoding_gfx942.mlir
@@ -366,7 +366,7 @@
// CHECK: %[[MMA:.+]] = iree_gpu.multi_mma %[[LHS]], %[[RHS]], %[[ACC]]
// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]],
// CHECK-SAME: iterator_types = [#iree_gpu.iterator_type<parallel>, #iree_gpu.iterator_type<parallel>, #iree_gpu.iterator_type<reduction>]
-// CHECK-SAME: kind = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_F32_16x16x4_F32, unroll_m = 8, unroll_n = 2, subgroups_n = 4, unroll_k = 4>
+// CHECK-SAME: kind = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_F32_16x16x4_F32, intrinsics_m = 8, intrinsics_n = 2, subgroups_n = 4, intrinsics_k = 4>
// CHECK: flow.dispatch.tensor.store %[[MMA]], %[[ACC_BINDING]]
// -----
@@ -426,7 +426,7 @@
// CHECK: %[[MMA:.+]] = iree_gpu.multi_mma %[[LHS]], %[[RHS]], %[[ACC]]
// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]],
// CHECK-SAME: iterator_types = [#iree_gpu.iterator_type<parallel>, #iree_gpu.iterator_type<parallel>, #iree_gpu.iterator_type<parallel>, #iree_gpu.iterator_type<reduction>]
-// CHECK-SAME: kind = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_F32_16x16x4_F32, unroll_m = 8, unroll_n = 2, subgroups_n = 4, unroll_k = 4>
+// CHECK-SAME: kind = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_F32_16x16x4_F32, intrinsics_m = 8, intrinsics_n = 2, subgroups_n = 4, intrinsics_k = 4>
// CHECK: flow.dispatch.tensor.store %[[MMA]], %[[ACC_BINDING]]
// -----
@@ -622,7 +622,7 @@
// CHECK: %[[MMA:.+]] = iree_gpu.multi_mma %[[LHS]], %[[RHS]], %[[ACC]]
// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]],
// CHECK-SAME: iterator_types = [#iree_gpu.iterator_type<parallel>, #iree_gpu.iterator_type<parallel>, #iree_gpu.iterator_type<reduction>]
-// CHECK-SAME: kind = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_I32_16x16x32_I8, unroll_m = 8, unroll_n = 2, subgroups_n = 4, unroll_k = 2>
+// CHECK-SAME: kind = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_I32_16x16x32_I8, intrinsics_m = 8, intrinsics_n = 2, subgroups_n = 4, intrinsics_k = 2>
// CHECK: flow.dispatch.tensor.store %[[MMA]], %[[ACC_BINDING]]
// -----
@@ -631,7 +631,7 @@
// 3. Custom target parameters to test more MaterializeEncoding heuristics.
//-------------------------------------------------------------------------
-// Custom {max_load_instruction_bits = 64} => implied default {unroll_k = 1} (omitted in output) instead of {unroll_k = 2}.
+// Custom {max_load_instruction_bits = 64} => implied default {intrinsics_k = 1} (omitted in output) instead of {intrinsics_k = 2}.
#target_gfx942_except_max_load_instruction_bits_64 = #hal.executable.target<"rocm", "rocm-hsaco-fb", {
iree.gpu.target = #iree_gpu.target<
@@ -700,11 +700,11 @@
// CHECK: func.func @matmul_lowering_MFMA_I32_16x16x32_I8_custom_max_load_instruction_bits_64
// CHECK: iree_gpu.multi_mma %[[LHS]], %[[RHS]], %[[ACC]]
// CHECK-SAME: iterator_types = [#iree_gpu.iterator_type<parallel>, #iree_gpu.iterator_type<parallel>, #iree_gpu.iterator_type<reduction>]
-// CHECK-SAME: kind = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_I32_16x16x32_I8, unroll_m = 8, unroll_n = 2, subgroups_n = 4>
+// CHECK-SAME: kind = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_I32_16x16x32_I8, intrinsics_m = 8, intrinsics_n = 2, subgroups_n = 4>
// -----
-// Custom {max_load_instruction_bits = 256} => {unroll_k = 4} instead of {unroll_k = 2}.
+// Custom {max_load_instruction_bits = 256} => {intrinsics_k = 4} instead of {intrinsics_k = 2}.
#target_gfx942_except_max_load_instruction_bits_256 = #hal.executable.target<"rocm", "rocm-hsaco-fb", {
iree.gpu.target = #iree_gpu.target<
@@ -773,11 +773,11 @@
// CHECK: func.func @matmul_lowering_MFMA_I32_16x16x32_I8_custom_max_load_instruction_bits_64
// CHECK: iree_gpu.multi_mma %[[LHS]], %[[RHS]], %[[ACC]]
// CHECK-SAME: iterator_types = [#iree_gpu.iterator_type<parallel>, #iree_gpu.iterator_type<parallel>, #iree_gpu.iterator_type<reduction>]
-// CHECK-SAME: kind = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_I32_16x16x32_I8, unroll_m = 8, unroll_n = 2, subgroups_n = 4, unroll_k = 4>
+// CHECK-SAME: kind = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_I32_16x16x32_I8, intrinsics_m = 8, intrinsics_n = 2, subgroups_n = 4, intrinsics_k = 4>
// -----
-// Custom {simds_per_wgp = 1} => implied default {subgroups_n = 1} (omitted in output) and {unroll_n = 8} instead of {subgroups_n = 4}.
+// Custom {simds_per_wgp = 1} => implied default {subgroups_n = 1} (omitted in output) and {intrinsics_n = 8} instead of {subgroups_n = 4}.
#target_gfx942_except_simds_per_wgp_1 = #hal.executable.target<"rocm", "rocm-hsaco-fb", {
iree.gpu.target = #iree_gpu.target<
@@ -846,11 +846,11 @@
// CHECK: func.func @matmul_lowering_MFMA_I32_16x16x32_I8_custom_simds_per_wgp_1
// CHECK: iree_gpu.multi_mma %[[LHS]], %[[RHS]], %[[ACC]]
// CHECK-SAME: iterator_types = [#iree_gpu.iterator_type<parallel>, #iree_gpu.iterator_type<parallel>, #iree_gpu.iterator_type<reduction>]
-// CHECK-SAME: kind = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_I32_16x16x32_I8, unroll_m = 8, unroll_n = 8, unroll_k = 2>
+// CHECK-SAME: kind = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_I32_16x16x32_I8, intrinsics_m = 8, intrinsics_n = 8, intrinsics_k = 2>
// -----
-// Custom 2x smaller {vgpr_space_bits = 8192} => smaller unroll_m and unroll_n
+// Custom 2x smaller {vgpr_space_bits = 8192} => smaller intrinsics_m and intrinsics_n
#target_gfx942_except_vgpr_space_bits_8192 = #hal.executable.target<"rocm", "rocm-hsaco-fb", {
iree.gpu.target = #iree_gpu.target<
@@ -919,11 +919,11 @@
// CHECK: func.func @matmul_lowering_MFMA_I32_16x16x32_I8_custom_vgpr_space_bits_8192
// CHECK: iree_gpu.multi_mma %[[LHS]], %[[RHS]], %[[ACC]]
// CHECK-SAME: iterator_types = [#iree_gpu.iterator_type<parallel>, #iree_gpu.iterator_type<parallel>, #iree_gpu.iterator_type<reduction>]
-// CHECK-SAME: kind = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_I32_16x16x32_I8, unroll_m = 4, unroll_n = 2, subgroups_n = 4, unroll_k = 2>
+// CHECK-SAME: kind = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_I32_16x16x32_I8, intrinsics_m = 4, intrinsics_n = 2, subgroups_n = 4, intrinsics_k = 2>
// -----
-// Custom 4x smaller {vgpr_space_bits = 4096} => smaller unroll_m and unroll_n
+// Custom 4x smaller {vgpr_space_bits = 4096} => smaller intrinsics_m and intrinsics_n
#target_gfx942_except_vgpr_space_bits_4096 = #hal.executable.target<"rocm", "rocm-hsaco-fb", {
iree.gpu.target = #iree_gpu.target<
@@ -992,11 +992,11 @@
// CHECK: func.func @matmul_lowering_MFMA_I32_16x16x32_I8_custom_vgpr_space_bits_4096
// CHECK: iree_gpu.multi_mma %[[LHS]], %[[RHS]], %[[ACC]]
// CHECK-SAME: iterator_types = [#iree_gpu.iterator_type<parallel>, #iree_gpu.iterator_type<parallel>, #iree_gpu.iterator_type<reduction>]
-// CHECK-SAME: kind = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_I32_16x16x32_I8, unroll_m = 4, subgroups_n = 4, unroll_k = 2>
+// CHECK-SAME: kind = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_I32_16x16x32_I8, intrinsics_m = 4, subgroups_n = 4, intrinsics_k = 2>
// -----
-// Custom smaller {vgpr_space_bits = 32768} => larger unroll_m and/or unroll_n
+// Custom smaller {vgpr_space_bits = 32768} => larger intrinsics_m and/or intrinsics_n
#target_gfx942_except_vgpr_space_bits_32768 = #hal.executable.target<"rocm", "rocm-hsaco-fb", {
iree.gpu.target = #iree_gpu.target<
@@ -1065,7 +1065,7 @@
// CHECK: func.func @matmul_lowering_MFMA_I32_16x16x32_I8_custom_vgpr_space_bits_32768
// CHECK: iree_gpu.multi_mma %[[LHS]], %[[RHS]], %[[ACC]]
// CHECK-SAME: iterator_types = [#iree_gpu.iterator_type<parallel>, #iree_gpu.iterator_type<parallel>, #iree_gpu.iterator_type<reduction>]
-// CHECK-SAME: kind = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_I32_16x16x32_I8, unroll_m = 8, unroll_n = 4, subgroups_n = 4, unroll_k = 2>
+// CHECK-SAME: kind = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_I32_16x16x32_I8, intrinsics_m = 8, intrinsics_n = 4, subgroups_n = 4, intrinsics_k = 2>
// -----
@@ -1128,7 +1128,7 @@
// CHECK: %[[MMA:.+]] = iree_gpu.multi_mma %[[LHS]], %[[RHS]], %[[ACC]]
// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]],
// CHECK-SAME: iterator_types = [#iree_gpu.iterator_type<parallel>, #iree_gpu.iterator_type<parallel>, #iree_gpu.iterator_type<parallel>, #iree_gpu.iterator_type<reduction>]
-// CHECK-SAME: kind = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_F32_16x16x32_F8E4M3FNUZ, unroll_m = 8, unroll_n = 2, subgroups_n = 4, unroll_k = 2>
+// CHECK-SAME: kind = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_F32_16x16x32_F8E4M3FNUZ, intrinsics_m = 8, intrinsics_n = 2, subgroups_n = 4, intrinsics_k = 2>
// CHECK: flow.dispatch.tensor.store %[[MMA]], %[[ACC_BINDING]]
// -----
@@ -1188,5 +1188,5 @@
// CHECK: %[[MMA:.+]] = iree_gpu.multi_mma %[[LHS]], %[[RHS]], %[[ACC]]
// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]],
// CHECK-SAME: iterator_types = [#iree_gpu.iterator_type<parallel>, #iree_gpu.iterator_type<parallel>, #iree_gpu.iterator_type<parallel>, #iree_gpu.iterator_type<reduction>]
-// CHECK-SAME: kind = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_F32_16x16x16_BF16, unroll_m = 8, unroll_n = 2, subgroups_n = 4, unroll_k = 2>
+// CHECK-SAME: kind = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_F32_16x16x16_BF16, intrinsics_m = 8, intrinsics_n = 2, subgroups_n = 4, intrinsics_k = 2>
// CHECK: flow.dispatch.tensor.store %[[MMA]], %[[ACC_BINDING]]
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/GPUTileSwizzleUtils.cpp b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/GPUTileSwizzleUtils.cpp
index cbff06a..2829922 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/GPUTileSwizzleUtils.cpp
+++ b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/GPUTileSwizzleUtils.cpp
@@ -147,14 +147,14 @@
case IREE::GPU::MMAFragment::Lhs:
// A-matrix (LHS). Source dimensions are M (index 0) and K (index 1).
// Unroll on K with interleaving, then on M.
- if (mma.getUnrollK() > 1) {
- expand(swizzle, 1, {Kind::CrossIntrinsic, mma.getUnrollK()});
+ if (mma.getIntrinsicsK() > 1) {
+ expand(swizzle, 1, {Kind::CrossIntrinsic, mma.getIntrinsicsK()});
int interleavingIdx =
getInnermostNonInternalDimIdx(swizzle.expandShape[1]);
interleave(swizzle, 1, interleavingIdx);
}
- if (mma.getUnrollM() > 1) {
- expand(swizzle, 0, {Kind::CrossIntrinsic, mma.getUnrollM()});
+ if (mma.getIntrinsicsM() > 1) {
+ expand(swizzle, 0, {Kind::CrossIntrinsic, mma.getIntrinsicsM()});
}
if (mma.getSubgroupsM() > 1) {
expand(swizzle, 0, {Kind::CrossThread, mma.getSubgroupsM()});
@@ -164,14 +164,14 @@
// B-matrix (RHS). Since the pack ops already took care of transposing B,
// source dimensions are N (index 0) and K (index 1).
// Unroll on K with interleaving, then on N.
- if (mma.getUnrollK() > 1) {
- expand(swizzle, 1, {Kind::CrossIntrinsic, mma.getUnrollK()});
+ if (mma.getIntrinsicsK() > 1) {
+ expand(swizzle, 1, {Kind::CrossIntrinsic, mma.getIntrinsicsK()});
int interleavingIdx =
getInnermostNonInternalDimIdx(swizzle.expandShape[1]);
interleave(swizzle, 1, interleavingIdx);
}
- if (mma.getUnrollN() > 1) {
- expand(swizzle, 0, {Kind::CrossIntrinsic, mma.getUnrollN()});
+ if (mma.getIntrinsicsN() > 1) {
+ expand(swizzle, 0, {Kind::CrossIntrinsic, mma.getIntrinsicsN()});
}
if (mma.getSubgroupsN() > 1) {
expand(swizzle, 0, {Kind::CrossThread, mma.getSubgroupsN()});
@@ -180,11 +180,11 @@
case IREE::GPU::MMAFragment::Acc:
// C-matrix (accumulator). Source dimensions are M (index 0) and N (index
// 1). Unroll on N, then on M.
- if (mma.getUnrollN() > 1) {
- expand(swizzle, 1, {Kind::CrossIntrinsic, mma.getUnrollN()});
+ if (mma.getIntrinsicsN() > 1) {
+ expand(swizzle, 1, {Kind::CrossIntrinsic, mma.getIntrinsicsN()});
}
- if (mma.getUnrollM() > 1) {
- expand(swizzle, 0, {Kind::CrossIntrinsic, mma.getUnrollM()});
+ if (mma.getIntrinsicsM() > 1) {
+ expand(swizzle, 0, {Kind::CrossIntrinsic, mma.getIntrinsicsM()});
}
if (mma.getSubgroupsN() > 1) {
expand(swizzle, 1, {Kind::CrossThread, mma.getSubgroupsN()});
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.cpp b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.cpp
index 11b2f96..74f3853 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.cpp
+++ b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.cpp
@@ -565,9 +565,9 @@
std::tuple<int64_t, int64_t, int64_t> DataTiledMMAAttr::getMNKShape() const {
MLIRContext *ctx = getContext();
auto opaqueLayout = getOpaqueMMALayout(ctx, getIntrinsic().getValue());
- return {opaqueLayout.mSize * getUnrollM() * getSubgroupsM(),
- opaqueLayout.nSize * getUnrollN() * getSubgroupsN(),
- opaqueLayout.kSize * getUnrollK()};
+ return {opaqueLayout.mSize * getIntrinsicsM() * getSubgroupsM(),
+ opaqueLayout.nSize * getIntrinsicsN() * getSubgroupsN(),
+ opaqueLayout.kSize * getIntrinsicsK()};
}
std::tuple<VectorType, VectorType, VectorType>
@@ -791,12 +791,12 @@
getVectorType(builder.getContext(), intrinsic, MMAFragment::Acc);
// Loop over the 3 unroll_{m,n,k} dimensions to create the intrinsics.
- for (int mu = 0; mu < getUnrollM(); ++mu) {
- for (int nu = 0; nu < getUnrollN(); ++nu) {
- for (int ku = 0; ku < getUnrollK(); ++ku) {
- Value lhs = intrinsicsLhs[mu * getUnrollK() + ku];
- Value rhs = intrinsicsRhs[nu * getUnrollK() + ku];
- Value &acc = intrinsicsAcc[mu * getUnrollN() + nu];
+ for (int mu = 0; mu < getIntrinsicsM(); ++mu) {
+ for (int nu = 0; nu < getIntrinsicsN(); ++nu) {
+ for (int ku = 0; ku < getIntrinsicsK(); ++ku) {
+ Value lhs = intrinsicsLhs[mu * getIntrinsicsK() + ku];
+ Value rhs = intrinsicsRhs[nu * getIntrinsicsK() + ku];
+ Value &acc = intrinsicsAcc[mu * getIntrinsicsN() + nu];
acc = createMmaOp(builder, loc, intrinsic, intrinCType, lhs, rhs, acc);
}
}
@@ -926,7 +926,7 @@
return success();
}
-int64_t VirtualMMAAttr::getUnrollK() const {
+int64_t VirtualMMAAttr::getIntrinsicsK() const {
switch (getIntrinsic().getValue()) {
case VirtualMMAIntrinsic::VMFMA_F32_16x16x32_F16:
case VirtualMMAIntrinsic::VMFMA_F32_32x32x16_F16: {
@@ -963,7 +963,7 @@
case VirtualMMAIntrinsic::VMFMA_F32_32x32x16_F8E4M3FNUZ:
case VirtualMMAIntrinsic::VMFMA_F32_32x32x16_F16: {
// Generate mfma's for K with unrolled kernels.
- const int64_t unrollKFactor = getUnrollK();
+ const int64_t unrollKFactor = getIntrinsicsK();
auto [m, n, k] = getMNKShape();
// Compute actual/native intrinsic's K size.
int64_t nativeKSize = k / unrollKFactor;
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.td b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.td
index e4b66bf..c4dbc00 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.td
+++ b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.td
@@ -232,11 +232,11 @@
let parameters = (ins
"::mlir::iree_compiler::IREE::GPU::MMAIntrinsicAttr":$intrinsic,
- DefaultValuedParameter<"int64_t", "1", "Unrolling along the M dimension, on the same thread.">:$unroll_m,
- DefaultValuedParameter<"int64_t", "1", "Unrolling along the M dimension, distributed across this many more threads.">:$subgroups_m,
- DefaultValuedParameter<"int64_t", "1", "Unrolling along the N dimension, on the same thread.">:$unroll_n,
- DefaultValuedParameter<"int64_t", "1", "Unrolling along the N dimension, distributed across this many more threads.">:$subgroups_n,
- DefaultValuedParameter<"int64_t", "1", "Unrolling along the K dimension, on the same thread, with interleaved layout.">:$unroll_k
+ DefaultValuedParameter<"int64_t", "1", "Intrinsic count along the M dimension.">:$intrinsics_m,
+ DefaultValuedParameter<"int64_t", "1", "Subgroup count along the M dimension.">:$subgroups_m,
+ DefaultValuedParameter<"int64_t", "1", "Intrinsic count along the N dimension.">:$intrinsics_n,
+ DefaultValuedParameter<"int64_t", "1", "Subgroup count along the N dimension.">:$subgroups_n,
+ DefaultValuedParameter<"int64_t", "1", "Intrinsic count along the K dimension, with interleaved layout.">:$intrinsics_k
);
}
@@ -260,7 +260,7 @@
let description = [{
This mma variant represents "virtual" MMA ops that has modification to
- its native layouts by unrollK and/or interleave reads. The |intrinsic|
+ its native layouts by intrinsicsK and/or interleave reads. The |intrinsic|
field represents different kinds of "Virtual" MMA Ops we found helpful.
These interleaving and/or unrolling changes in the layout is especially
@@ -282,8 +282,8 @@
// Factor to unroll K from native MMA/intrinsic size to virtual size.
// e.g MFMA_F32_16x16x16 has K of 16, while VMFMA_F32_16x16x32 has K of 32
- // in this example, unrollK = 32/16 = 2.
- int64_t getUnrollK() const;
+ // in this example, intrinsicsK = 32/16 = 2.
+ int64_t getIntrinsicsK() const;
}];
}
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/test/iree_gpu_attrs.mlir b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/test/iree_gpu_attrs.mlir
index 2a611ea..d506140 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/test/iree_gpu_attrs.mlir
+++ b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/test/iree_gpu_attrs.mlir
@@ -29,25 +29,25 @@
module {
func.func @test_data_tiled_mfma_f32_16x16x4_f32() attributes {
- mma_types = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_F32_16x16x4_F32, unroll_m = 4, subgroups_m = 2, unroll_k = 1>} {
+ mma_types = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_F32_16x16x4_F32, intrinsics_m = 4, subgroups_m = 2, intrinsics_k = 1>} {
return
}
}
// CHECK-LABEL: func @test_data_tiled_mfma_f32_16x16x4_f32
-// CHECK-SAME: mma_types = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_F32_16x16x4_F32, unroll_m = 4, subgroups_m = 2>
+// CHECK-SAME: mma_types = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_F32_16x16x4_F32, intrinsics_m = 4, subgroups_m = 2>
module {
func.func @test_data_tiled_mfma_f32_16x16x16_f16() attributes {
- mma_types = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_F32_16x16x16_F16, unroll_m = 1, subgroups_n = 2, unroll_k = 2>} {
+ mma_types = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_F32_16x16x16_F16, intrinsics_m = 1, subgroups_n = 2, intrinsics_k = 2>} {
return
}
}
// CHECK-LABEL: func @test_data_tiled_mfma_f32_16x16x16_f16
-// CHECK-SAME: mma_types = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_F32_16x16x16_F16, subgroups_n = 2, unroll_k = 2>
+// CHECK-SAME: mma_types = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_F32_16x16x16_F16, subgroups_n = 2, intrinsics_k = 2>
module {
func.func @test_data_tiled_mfma_i32_16x16x32_i8() attributes {
- mma_types = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_I32_16x16x32_I8, unroll_m = 1, unroll_n = 1, unroll_k = 1>} {
+ mma_types = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_I32_16x16x32_I8, intrinsics_m = 1, intrinsics_n = 1, intrinsics_k = 1>} {
return
}
}
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/test/iree_gpu_ops.mlir b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/test/iree_gpu_ops.mlir
index d64493c..38ad149 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/test/iree_gpu_ops.mlir
+++ b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/test/iree_gpu_ops.mlir
@@ -254,7 +254,7 @@
%0 = iree_gpu.multi_mma %lhs, %rhs, %acc {
indexing_maps = #contraction_accesses,
iterator_types = [#iree_gpu.iterator_type<parallel>, #iree_gpu.iterator_type<parallel>, #iree_gpu.iterator_type<reduction>],
- kind = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_F32_16x16x4_F32, unroll_m = 2, unroll_n = 2, unroll_k = 4>
+ kind = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_F32_16x16x4_F32, intrinsics_m = 2, intrinsics_n = 2, intrinsics_k = 4>
} : tensor<?x?x2x4x16x1x4xf32>, tensor<?x?x2x4x16x1x4xf32> into tensor<?x?x2x2x4x16x4x1xf32>
return %0 : tensor<?x?x2x2x4x16x4x1xf32>
}
@@ -267,7 +267,7 @@
// CHECK: iree_gpu.multi_mma %arg0, %arg1, %arg2
// CHECK-SAME: indexing_maps = [#[[$MAP]], #[[$MAP1]], #[[$MAP2]]]
// CHECK-SAME: iterator_types = [#iree_gpu.iterator_type<parallel>, #iree_gpu.iterator_type<parallel>, #iree_gpu.iterator_type<reduction>]
-// CHECK-SAME: kind = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_F32_16x16x4_F32, unroll_m = 2, unroll_n = 2, unroll_k = 4>
+// CHECK-SAME: kind = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_F32_16x16x4_F32, intrinsics_m = 2, intrinsics_n = 2, intrinsics_k = 4>
// CHECK-SAME: : tensor<?x?x2x4x16x1x4xf32>, tensor<?x?x2x4x16x1x4xf32> into tensor<?x?x2x2x4x16x4x1xf32>
// -----
@@ -281,7 +281,7 @@
%0 = iree_gpu.multi_mma %lhs, %rhs, %acc {
indexing_maps = #contraction_accesses,
iterator_types = [#iree_gpu.iterator_type<parallel>, #iree_gpu.iterator_type<parallel>, #iree_gpu.iterator_type<reduction>],
- kind = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_F32_16x16x4_F32, subgroups_m = 2, subgroups_n = 2, unroll_k = 4>
+ kind = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_F32_16x16x4_F32, subgroups_m = 2, subgroups_n = 2, intrinsics_k = 4>
} : tensor<?x?x2x4x16x1x4xf32>, tensor<?x?x2x4x16x1x4xf32> into tensor<?x?x2x2x4x16x4x1xf32>
return %0 : tensor<?x?x2x2x4x16x4x1xf32>
}
@@ -294,7 +294,7 @@
// CHECK: iree_gpu.multi_mma %arg0, %arg1, %arg2
// CHECK-SAME: indexing_maps = [#[[$MAP]], #[[$MAP1]], #[[$MAP2]]]
// CHECK-SAME: iterator_types = [#iree_gpu.iterator_type<parallel>, #iree_gpu.iterator_type<parallel>, #iree_gpu.iterator_type<reduction>]
-// CHECK-SAME: kind = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_F32_16x16x4_F32, subgroups_m = 2, subgroups_n = 2, unroll_k = 4>
+// CHECK-SAME: kind = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_F32_16x16x4_F32, subgroups_m = 2, subgroups_n = 2, intrinsics_k = 4>
// CHECK-SAME: : tensor<?x?x2x4x16x1x4xf32>, tensor<?x?x2x4x16x1x4xf32> into tensor<?x?x2x2x4x16x4x1xf32>
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/GPU/Transforms/test/distribute_mma_to_lanes.mlir b/compiler/src/iree/compiler/Codegen/Dialect/GPU/Transforms/test/distribute_mma_to_lanes.mlir
index a5a0ff1..f4ba3ae 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/GPU/Transforms/test/distribute_mma_to_lanes.mlir
+++ b/compiler/src/iree/compiler/Codegen/Dialect/GPU/Transforms/test/distribute_mma_to_lanes.mlir
@@ -416,7 +416,7 @@
%0 = iree_gpu.multi_mma %lhs, %rhs, %acc {
indexing_maps = #contraction_accesses,
iterator_types = [#iree_gpu.iterator_type<parallel>, #iree_gpu.iterator_type<parallel>, #iree_gpu.iterator_type<reduction>],
- kind = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_F32_16x16x4_F32, unroll_m = 2, unroll_n = 2, unroll_k = 4>
+ kind = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_F32_16x16x4_F32, intrinsics_m = 2, intrinsics_n = 2, intrinsics_k = 4>
} : tensor<1x1x2x4x16x4xf32>, tensor<1x1x2x4x16x4xf32> into tensor<1x1x2x2x4x16x4xf32>
return %0 : tensor<1x1x2x2x4x16x4xf32>
}
@@ -434,7 +434,7 @@
// CHECK-DAG: %[[ACC_SLICE:.+]] = tensor.extract_slice %[[ACC_ARG]]
// CHECK-SAME: [0, 0, 0, 0, %[[IN_IDS]]#1, %[[IN_IDS]]#2, 0] [1, 1, 2, 2, 1, 1, 4] [1, 1, 1, 1, 1, 1, 1]
// CHECK: %[[MMA:.+]] = iree_gpu.multi_mma %[[LHS_SLICE]], %[[RHS_SLICE]], %[[ACC_SLICE]]
-// CHECK-SAME: kind = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_F32_16x16x4_F32, unroll_m = 2, unroll_n = 2, unroll_k = 4>
+// CHECK-SAME: kind = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_F32_16x16x4_F32, intrinsics_m = 2, intrinsics_n = 2, intrinsics_k = 4>
// CHECK-SAME: : tensor<1x1x2x1x1x4xf32>, tensor<1x1x2x1x1x4xf32> into tensor<1x1x2x2x1x1x4xf32>
// CHECK: tensor.parallel_insert_slice %[[MMA]] into %[[ACC_ARG]]
// CHECK-SAME: [0, 0, 0, 0, %[[IN_IDS]]#1, %[[IN_IDS]]#2, 0] [1, 1, 2, 2, 1, 1, 4] [1, 1, 1, 1, 1, 1, 1]
@@ -452,7 +452,7 @@
%0 = iree_gpu.multi_mma %lhs, %rhs, %acc {
indexing_maps = #contraction_accesses,
iterator_types = [#iree_gpu.iterator_type<parallel>, #iree_gpu.iterator_type<parallel>, #iree_gpu.iterator_type<reduction>],
- kind = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_F32_16x16x4_F32, subgroups_m = 2, subgroups_n = 2, unroll_k = 4>
+ kind = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_F32_16x16x4_F32, subgroups_m = 2, subgroups_n = 2, intrinsics_k = 4>
} : tensor<1x1x2x4x16x4xf32>, tensor<1x1x2x4x16x4xf32> into tensor<1x1x2x2x4x16x4xf32>
return %0 : tensor<1x1x2x2x4x16x4xf32>
}
@@ -471,7 +471,7 @@
// CHECK-DAG: %[[ACC_SLICE:.+]] = tensor.extract_slice %[[ACC_ARG]]
// CHECK-SAME: [0, 0, %[[ACC_IDS]]#1, %[[ACC_IDS]]#2, %[[ACC_IDS]]#3, %[[ACC_IDS]]#4, 0] [1, 1, 1, 1, 1, 1, 4] [1, 1, 1, 1, 1, 1, 1]
// CHECK: %[[MMA:.+]] = iree_gpu.multi_mma %[[LHS_SLICE]], %[[RHS_SLICE]], %[[ACC_SLICE]]
-// CHECK-SAME: kind = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_F32_16x16x4_F32, subgroups_m = 2, subgroups_n = 2, unroll_k = 4>}
+// CHECK-SAME: kind = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_F32_16x16x4_F32, subgroups_m = 2, subgroups_n = 2, intrinsics_k = 4>}
// CHECK-SAME: : tensor<1x1x1x1x1x4xf32>, tensor<1x1x1x1x1x4xf32> into tensor<1x1x1x1x1x1x4xf32>
// CHECK: tensor.parallel_insert_slice %[[MMA]] into %[[ACC_ARG]]
// CHECK-SAME: [0, 0, %[[ACC_IDS]]#1, %[[ACC_IDS]]#2, %[[ACC_IDS]]#3, %[[ACC_IDS]]#4, 0] [1, 1, 1, 1, 1, 1, 4] [1, 1, 1, 1, 1, 1, 1]
diff --git a/compiler/src/iree/compiler/Codegen/ExternalInterfaces/GPUEncodingExternalModels.cpp b/compiler/src/iree/compiler/Codegen/ExternalInterfaces/GPUEncodingExternalModels.cpp
index b3bd093..8920e4c 100644
--- a/compiler/src/iree/compiler/Codegen/ExternalInterfaces/GPUEncodingExternalModels.cpp
+++ b/compiler/src/iree/compiler/Codegen/ExternalInterfaces/GPUEncodingExternalModels.cpp
@@ -96,13 +96,13 @@
};
auto [intrinsicA, intrinsicB, intrinsicC] = intrinsicMma.getABCVectorTypes();
- // The unrollK factor serves to allow loads from the A and B matrices to use
- // the target ISA's vector loads. For instance, if the ISA has 128-bit loads
- // and each intrinsic consumes only 32 bits from A and B, then we want to set
- // unrollK=4 to turn 4 separate 32-bit loads into one 128-bit load.
+ // The intrinsicsK factor serves to allow loads from the A and B matrices to
+ // use the target ISA's vector loads. For instance, if the ISA has 128-bit
+ // loads and each intrinsic consumes only 32 bits from A and B, then we want
+ // to set intrinsicsK=4 to turn 4 separate 32-bit loads into one 128-bit load.
int intrinsicLoadBits =
std::min(sizeInBits(intrinsicA), sizeInBits(intrinsicB));
- const int unrollK =
+ const int intrinsicsK =
std::max(1, *wgp.getMaxLoadInstructionBits() / intrinsicLoadBits);
// The total amount of unrolling along the M and N dimensions is normally
@@ -113,25 +113,25 @@
// correspondingly divides the available register space between this many
// subgroups, making it cancel out of the equation here.
//
- // We need to solve for two variables here, unroll_m and unroll_n, constrained
- // by one quadratic equation expressing that the A, B and C tiles must fit in
- // VGPR space. Since we have only 1 constraint for two variables, we
- // self-impose a second constraint for now: that the unrolling shape should be
- // square, i.e. unrollM == unrollN.
+ // We need to solve for two variables here, intrinsics_m and intrinsics_n,
+ // constrained by one quadratic equation expressing that the A, B and C tiles
+ // must fit in VGPR space. Since we have only 1 constraint for two variables,
+ // we self-impose a second constraint for now: that the unrolling shape should
+ // be square, i.e. intrinsicsM == intrinsicsN.
// TODO(#18850): that is suboptimal for narrow cases.
//
// Now we have only one variable, call it x, to solve for.
// The register space taken is:
- // A-tile: x * unrollK * sizeInBits(intrinsicA)
- // B-tile: x * unrollK * sizeInBits(intrinsicB)
+ // A-tile: x * intrinsicsK * sizeInBits(intrinsicA)
+ // B-tile: x * intrinsicsK * sizeInBits(intrinsicB)
// C-tile: x^2 * sizeInBits(intrinsicC)
// So the equation to solve is:
// x^2 * sizeInBits(intrinsicC)
- // + x * unrollK * (sizeInBits(intrinsicA) + sizeInBits(intrinsicB))
+ // + x * intrinsicsK * (sizeInBits(intrinsicA) + sizeInBits(intrinsicB))
// == wgp.getVgprSpaceBits()
float c2 = sizeInBits(intrinsicC);
- float c1 = unrollK * (sizeInBits(intrinsicA) + sizeInBits(intrinsicB));
+ float c1 = intrinsicsK * (sizeInBits(intrinsicA) + sizeInBits(intrinsicB));
float c0 = -*wgp.getVgprSpaceBits(); // negative by construction.
// Now the equation to solve is: c2 * x^2 + c1 * x + c0 == 0.
float discriminant = c1 * c1 - 4 * c0 * c2; // positive, because c0 < 0.
@@ -185,8 +185,8 @@
// to think about unroll-to-subgroups when making the narrowing adjustment.
int subgroupsM = 1;
int subgroupsN = *wgp.getSimdsPerWgp();
- int unrollM = totalUnrollM / subgroupsM;
- int unrollN = totalUnrollN / subgroupsN;
+ int intrinsicsM = totalUnrollM / subgroupsM;
+ int intrinsicsN = totalUnrollN / subgroupsN;
//
// Step 3: Adjust the unrolling factors when there is a narrow dimension.
@@ -195,19 +195,22 @@
IREE::Encoding::MatmulNarrowDim narrowDim =
IREE::Encoding::getMatmulNarrowDim(encoding);
if (narrowDim.isM()) {
- unrollM = std::min(unrollM, static_cast<int>(llvm::divideCeil(
- narrowDim.size, intrinsicMma.getMSize())));
+ intrinsicsM =
+ std::min(intrinsicsM, static_cast<int>(llvm::divideCeil(
+ narrowDim.size, intrinsicMma.getMSize())));
}
if (narrowDim.isN()) {
- std::swap(unrollM, unrollN);
+ std::swap(intrinsicsM, intrinsicsN);
std::swap(subgroupsM, subgroupsN);
assert(subgroupsN == 1);
- unrollN = std::min(unrollN, static_cast<int>(llvm::divideCeil(
- narrowDim.size, intrinsicMma.getNSize())));
+ intrinsicsN =
+ std::min(intrinsicsN, static_cast<int>(llvm::divideCeil(
+ narrowDim.size, intrinsicMma.getNSize())));
}
- return DataTiledMMAAttr::get(ctx, intrinsicMma.getIntrinsic(), unrollM,
- subgroupsM, unrollN, subgroupsN, unrollK);
+ return DataTiledMMAAttr::get(ctx, intrinsicMma.getIntrinsic(), intrinsicsM,
+ subgroupsM, intrinsicsN, subgroupsN,
+ intrinsicsK);
}
static Operation *lowerContractionOpToMultiMmaOp(OpBuilder &builder,
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/config_tile_and_fuse.mlir b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/config_tile_and_fuse.mlir
index 125ad8b..47ccb62 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/config_tile_and_fuse.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/config_tile_and_fuse.mlir
@@ -242,9 +242,9 @@
#iree_gpu.iterator_type<reduction>],
kind = #iree_gpu.data_tiled_mma_layout<
intrinsic = MFMA_F32_16x16x4_F32,
- unroll_m = 8, unroll_n = 2,
+ intrinsics_m = 8, intrinsics_n = 2,
subgroups_n = 4,
- unroll_k = 4>}
+ intrinsics_k = 4>}
: tensor<1x8x8x4x16x4xf32>, tensor<1x8x4x2x4x16x4xf32> into tensor<1x1x8x4x2x4x16x4xf32>
return %6 : tensor<1x1x8x4x2x4x16x4xf32>
}
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/pipeline_tile_and_fuse.mlir b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/pipeline_tile_and_fuse.mlir
index 886b39b..1a521e6 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/pipeline_tile_and_fuse.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/pipeline_tile_and_fuse.mlir
@@ -729,10 +729,10 @@
#iree_gpu.iterator_type<reduction>],
kind = #iree_gpu.data_tiled_mma_layout<
intrinsic = MFMA_F32_16x16x4_F32,
- unroll_m = 8,
- unroll_n = 2,
+ intrinsics_m = 8,
+ intrinsics_n = 2,
subgroups_n = 4,
- unroll_k = 4>}
+ intrinsics_k = 4>}
: tensor<4x1x8x4x16x4xf32>, tensor<4x1x4x2x4x16x4xf32> into tensor<4x4x8x4x2x4x16x4xf32>
flow.dispatch.tensor.store %6, %2, offsets = [0, 0, 0, 0, 0, 0, 0, 0], sizes = [4, 4, 8, 4, 2, 4, 16, 4], strides = [1, 1, 1, 1, 1, 1, 1, 1] : tensor<4x4x8x4x2x4x16x4xf32> -> !flow.dispatch.tensor<readwrite:tensor<4x4x8x4x2x4x16x4xf32>>
return
diff --git a/compiler/src/iree/compiler/DispatchCreation/test/convert_region_to_workgroups.mlir b/compiler/src/iree/compiler/DispatchCreation/test/convert_region_to_workgroups.mlir
index d69af93..d6cba05 100644
--- a/compiler/src/iree/compiler/DispatchCreation/test/convert_region_to_workgroups.mlir
+++ b/compiler/src/iree/compiler/DispatchCreation/test/convert_region_to_workgroups.mlir
@@ -62,7 +62,7 @@
affine_map<(d0, d1, d2) -> (d1, d2)>,
affine_map<(d0, d1, d2) -> (d0, d1)>],
iterator_types = [#iree_gpu.iterator_type<parallel>, #iree_gpu.iterator_type<parallel>, #iree_gpu.iterator_type<reduction>],
- kind = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_F32_16x16x16_F16, unroll_m = 8, unroll_n = 2, subgroups_n = 4, unroll_k = 2>}
+ kind = #iree_gpu.data_tiled_mma_layout<intrinsic = MFMA_F32_16x16x16_F16, intrinsics_m = 8, intrinsics_n = 2, subgroups_n = 4, intrinsics_k = 2>}
: tensor<4x16x8x4x16x2x4xf16>, tensor<4x16x4x2x4x16x2x4xf16> into tensor<4x4x8x4x2x4x16x4xf32>
flow.return %13 : tensor<4x4x8x4x2x4x16x4xf32>
}
