[DT][CPU]: RVV scalable encoding materialization (#24600)
This PR introduces scalable data-tiled layout selection and
materialization for RVV. Though since there exist microkernels for the
RV64 backend with the V extension, previous mechanism to select these
data-tiled layouts _statically_ based on the `+zvl*b` flags is also
preserved and the scalable data-tiled layout materialization is hidden
behind the flag scalable vectorization flag.
The exactl layout selection (the tile sizes) though is based on the same
principles as the static one - fully utilize all 32 of the V registers.
F16 tile sizes that underutilize the V registers in the static case (see
[this discord
thread](https://discord.com/channels/689900678990135345/1456679894812590326/1461764521403224209))).
**NOTE**: the `+zvl*b` flags DO NOT give us the _exact_ values of VLEN
that the hardware implements but rather the minimum one, meaning that
the previous static mechanism can underutilize the vector registers and
is not vector-length agnostic.
---------
Signed-off-by: Ege Beysel <beyselege@gmail.com>
diff --git a/compiler/src/iree/compiler/Codegen/Common/test/materialize_encoding_aarch64.mlir b/compiler/src/iree/compiler/Codegen/Common/test/materialize_encoding_aarch64.mlir
index 9c422c4..3ddee79 100644
--- a/compiler/src/iree/compiler/Codegen/Common/test/materialize_encoding_aarch64.mlir
+++ b/compiler/src/iree/compiler/Codegen/Common/test/materialize_encoding_aarch64.mlir
@@ -153,7 +153,7 @@
%arg1: tensor<16x1xf32, #encoding_rhs>,
%arg2: tensor<16x1xf32, #encoding_result>
) -> tensor<16x1xf32, #encoding_result> attributes {
- hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz", {target_triple="aarch64-xyz-xyz", iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}>
+ hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz", {cpu_features="+sve", target_triple="aarch64-xyz-xyz", iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}>
} {
%0 = linalg.matmul ins(%arg0, %arg1 : tensor<16x16xf32, #encoding_lhs>, tensor<16x1xf32, #encoding_rhs>)
outs(%arg2 : tensor<16x1xf32, #encoding_result>) -> tensor<16x1xf32, #encoding_result>
@@ -161,11 +161,17 @@
}
// CHECK-LABEL: func @matvec_shaped_matmul_lowering_f32f32f32_aarch64
-// CHECK-SAME: %[[LHS:[a-zA-Z0-9]+]]: tensor<2x16x8x1xf32>
-// CHECK-SAME: %[[RHS:[a-zA-Z0-9]+]]: tensor<1x16x1x1xf32>
-// CHECK-SAME: %[[OUT:[a-zA-Z0-9]+]]: tensor<1x2x1x8xf32>
+// NO-SVE-SAME: %[[LHS:[a-zA-Z0-9]+]]: tensor<2x16x8x1xf32>
+// NO-SVE-SAME: %[[RHS:[a-zA-Z0-9]+]]: tensor<1x16x1x1xf32>
+// NO-SVE-SAME: %[[OUT:[a-zA-Z0-9]+]]: tensor<1x2x1x8xf32>
+
+// WITH-SVE-SAME: %[[LHS:[a-zA-Z0-9]+]]: tensor<2x16x8x1xf32>
+// WITH-SVE-SAME: %[[RHS:[a-zA-Z0-9]+]]: tensor<?x16x?x1xf32>
+// WITH-SVE-SAME: %[[OUT:[a-zA-Z0-9]+]]: tensor<2x?x8x?xf32>
+
// CHECK: %[[MMT4D:.+]] = linalg.mmt4d
-// CHECK-SAME: ins(%[[RHS]], %[[LHS]]
+// NO-SVE-SAME: ins(%[[RHS]], %[[LHS]]
+// WITH-SVE-SAME: ins(%[[LHS]], %[[RHS]]
// CHECK-SAME: outs(%[[OUT]]
// CHECK: return %[[MMT4D]]
@@ -221,7 +227,7 @@
}
// CHECK-DAG: #[[$MAP0:.+]] = affine_map<()[s0] -> (s0 ceildiv 8)>
// WITH-SVE-DAG: #[[$MAP1:.+]] = affine_map<()[s0, s1] -> (s0 ceildiv s1)>
-// CHECK-LABEL: func @matmul_lowering_f32f32f32_aarch64()
+// CHECK: func @matmul_lowering_f32f32f32_aarch64()
// WITH-SVE-DAG: %[[C8:.+]] = arith.constant 8 : index
// CHECK-DAG: %[[M:.+]] = hal.interface.constant.load layout({{.+}}) ordinal(0)
// CHECK-DAG: %[[N:.+]] = hal.interface.constant.load layout({{.+}}) ordinal(1)
@@ -310,7 +316,7 @@
%rhs: tensor<16x1xf32, #encoding_rhs>,
%result: tensor<16x1xf32, #encoding_result>
) -> tensor<16x1xf32, #encoding_result> attributes {
- hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz", {target_triple="aarch64-xyz-xyz", iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}>
+ hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz", {cpu_features="+sve", target_triple="aarch64-xyz-xyz", iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}>
} {
%matmul = linalg.matmul
ins(%lhs, %rhs : tensor<16x16xf32, #encoding_lhs>, tensor<16x1xf32, #encoding_rhs>)
@@ -319,13 +325,20 @@
return %matmul : tensor<16x1xf32, #encoding_result>
}
// CHECK-LABEL: func @matvec_lowering_f32f32f32_aarch64(
-// CHECK-SAME: %[[LHS:[a-zA-Z0-9]+]]: tensor<2x16x8x1xf32>
-// CHECK-SAME: %[[RHS:[a-zA-Z0-9]+]]: tensor<1x16x1x1xf32>
-// CHECK-SAME: %[[OUTS:[a-zA-Z0-9]+]]: tensor<1x2x1x8xf32>
-// CHECK: %[[MMT4D:.+]] = linalg.mmt4d
-// CHECK-SAME: ins(%[[RHS]], %[[LHS]] :
-// CHECK-SAME: outs(%[[OUTS]] :
-// CHECK: return %[[MMT4D]]
+// NO-SVE-SAME: %[[LHS:[a-zA-Z0-9]+]]: tensor<2x16x8x1xf32>
+// NO-SVE-SAME: %[[RHS:[a-zA-Z0-9]+]]: tensor<1x16x1x1xf32>
+// NO-SVE-SAME: %[[OUT:[a-zA-Z0-9]+]]: tensor<1x2x1x8xf32>
+
+
+// WITH-SVE-SAME: %[[LHS:[a-zA-Z0-9]+]]: tensor<2x16x8x1xf32>
+// WITH-SVE-SAME: %[[RHS:[a-zA-Z0-9]+]]: tensor<?x16x?x1xf32>
+// WITH-SVE-SAME: %[[OUT:[a-zA-Z0-9]+]]: tensor<2x?x8x?xf32>
+
+// CHECK: %[[MMT4D:.+]] = linalg.mmt4d
+// NO-SVE-SAME: ins(%[[RHS]], %[[LHS]] :
+// WITH-SVE-SAME: ins(%[[LHS]], %[[RHS]] :
+// CHECK-SAME: outs(%[[OUT]] :
+// CHECK: return %[[MMT4D]]
// -----
diff --git a/compiler/src/iree/compiler/Codegen/Common/test/materialize_encoding_riscv.mlir b/compiler/src/iree/compiler/Codegen/Common/test/materialize_encoding_riscv.mlir
index b1a54e0..a7ed1af 100644
--- a/compiler/src/iree/compiler/Codegen/Common/test/materialize_encoding_riscv.mlir
+++ b/compiler/src/iree/compiler/Codegen/Common/test/materialize_encoding_riscv.mlir
@@ -1,12 +1,46 @@
-// RUN: iree-opt --pass-pipeline="builtin.module(func.func(iree-codegen-materialize-device-encoding))" --split-input-file %s | FileCheck %s
+// Tests for RISC-V targets with static VLEN-based tiles and scalable vectorization.
+//
+// STATIC mode (default): Expected to use VLEN-based tile sizes computed from zvl* features.
+// SCALABLE mode (opt-in): Expected to use smaller base tile sizes that scale with vscale at runtime.
+//
+// RUN: iree-opt --pass-pipeline="builtin.module(func.func(iree-codegen-materialize-device-encoding))" --split-input-file %s | FileCheck %s --check-prefixes=CHECK,STATIC
+// RUN: iree-opt --pass-pipeline="builtin.module(func.func(iree-codegen-materialize-device-encoding))" --iree-llvmcpu-enable-scalable-vectorization=true --split-input-file %s | FileCheck %s --check-prefixes=CHECK,SCALABLE
+
+// RISC-V f32 matmul encoding tests.
#map = affine_map<(d0, d1, d2) -> (d0, d2)>
#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
-#encoding_lhs = #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [i8, i8, i32], user_indexing_maps = [#map, #map1, #map2], iteration_sizes = [?, ?, ?]>
-#encoding_rhs = #iree_encoding.encoding<operand_index = 1, op_type = matmul, element_types = [i8, i8, i32], user_indexing_maps = [#map, #map1, #map2], iteration_sizes = [?, ?, ?]>
-#encoding_result = #iree_encoding.encoding<operand_index = 2, op_type = matmul, element_types = [i8, i8, i32], user_indexing_maps = [#map, #map1, #map2], iteration_sizes = [?, ?, ?]>
-func.func @matmul_lowering_f32f32f32_riscv(%lhs: tensor<?x?xf32>, %rhs: tensor<?x?xf32>, %acc: tensor<?x?xf32>) -> tensor<?x?xf32> attributes {
+#encoding_lhs = #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map, #map1, #map2], iteration_sizes = [?, ?, ?]>
+#encoding_rhs = #iree_encoding.encoding<operand_index = 1, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map, #map1, #map2], iteration_sizes = [?, ?, ?]>
+#encoding_res = #iree_encoding.encoding<operand_index = 2, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map, #map1, #map2], iteration_sizes = [?, ?, ?]>
+
+// RISC-V64 without V extension - no data tiling (fallback)
+func.func @negative_matmul_lowering_f32f32f32_riscv64_no_v_ext(%lhs: tensor<?x?xf32>, %rhs: tensor<?x?xf32>, %acc: tensor<?x?xf32>, %m : index, %n : index, %k : index) -> tensor<?x?xf32> attributes {
+ hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz", {target_triple="riscv64-xyz-xyz", iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}>
+} {
+ %0 = iree_encoding.set_encoding %lhs encoding_dims{%m, %n, %k} : tensor<?x?xf32> -> tensor<?x?xf32, #encoding_lhs>
+ %1 = iree_encoding.set_encoding %rhs encoding_dims{%m, %n, %k} : tensor<?x?xf32> -> tensor<?x?xf32, #encoding_rhs>
+ %2 = iree_encoding.set_encoding %acc encoding_dims{%m, %n, %k} : tensor<?x?xf32> -> tensor<?x?xf32, #encoding_res>
+ %3 = linalg.matmul
+ ins(%0, %1 : tensor<?x?xf32, #encoding_lhs>,
+ tensor<?x?xf32, #encoding_rhs>)
+ outs(%2 : tensor<?x?xf32, #encoding_res>)
+ -> tensor<?x?xf32, #encoding_res>
+ %4 = iree_encoding.unset_encoding %3 encoding_dims{%m, %n, %k} : tensor<?x?xf32, #encoding_res> -> tensor<?x?xf32>{%m, %n}
+ return %4 : tensor<?x?xf32>
+}
+
+// RISC-V64 without V extension does not implement data-tiling.
+// CHECK-LABEL: func @negative_matmul_lowering_f32f32f32_riscv64_no_v_ext
+// CHECK-NOT: linalg.pack
+// CHECK-NOT: linalg.mmt4d
+// CHECK: %[[RES:.+]] = linalg.matmul
+// CHECK-NOT: linalg.unpack
+// CHECK: return %[[RES]]
+
+// RISC-V32 without V extension - no data tiling
+func.func @negative_matmul_lowering_f32f32f32_riscv32_no_v_ext(%lhs: tensor<?x?xf32>, %rhs: tensor<?x?xf32>, %acc: tensor<?x?xf32>) -> tensor<?x?xf32> attributes {
hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz", {target_triple="riscv32-xyz-xyz", iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}>
} {
%c0 = arith.constant 0 : index
@@ -16,22 +50,167 @@
%K = tensor.dim %lhs, %c1 : tensor<?x?xf32>
%0 = iree_encoding.set_encoding %lhs encoding_dims{%M, %N, %K} : tensor<?x?xf32> -> tensor<?x?xf32, #encoding_lhs>
%1 = iree_encoding.set_encoding %rhs encoding_dims{%M, %N, %K} : tensor<?x?xf32> -> tensor<?x?xf32, #encoding_rhs>
- %2 = iree_encoding.set_encoding %acc encoding_dims{%M, %N, %K} : tensor<?x?xf32> -> tensor<?x?xf32, #encoding_result>
+ %2 = iree_encoding.set_encoding %acc encoding_dims{%M, %N, %K} : tensor<?x?xf32> -> tensor<?x?xf32, #encoding_res>
%3 = linalg.matmul
ins(%0, %1 : tensor<?x?xf32, #encoding_lhs>,
tensor<?x?xf32, #encoding_rhs>)
- outs(%2 : tensor<?x?xf32, #encoding_result>)
- -> tensor<?x?xf32, #encoding_result>
- %4 = iree_encoding.unset_encoding %3 encoding_dims{%M, %N, %K} : tensor<?x?xf32, #encoding_result> -> tensor<?x?xf32>{%M, %N}
+ outs(%2 : tensor<?x?xf32, #encoding_res>)
+ -> tensor<?x?xf32, #encoding_res>
+ %4 = iree_encoding.unset_encoding %3 encoding_dims{%M, %N, %K} : tensor<?x?xf32, #encoding_res> -> tensor<?x?xf32>{%M, %N}
return %4 : tensor<?x?xf32>
}
-// RISC-V targets does not implement data-tiling yet.
-// CHECK-LABEL: func @matmul_lowering_f32f32f32_riscv
+// RISC-V32 without V extension does not implement data-tiling.
+// CHECK-LABEL: func @negative_matmul_lowering_f32f32f32_riscv32_no_v_ext
+// CHECK-NOT: linalg.pack
+// CHECK-NOT: linalg.mmt4d
// CHECK: %[[RES:.+]] = linalg.matmul
+// CHECK-NOT: linalg.unpack
// CHECK: return %[[RES]]
+// RISC-V64 with V extension - set_encoding for LHS (f32)
+// The LHS operand corresponds to the M dimension, which is NOT scalable.
+// Therefore, STATIC and SCALABLE produce the same result for LHS.
+func.func @matmul_set_encoding_LHS_f32_riscv64(%arg0: tensor<?x?xf32>, %m: index, %n: index, %k: index) -> tensor<?x?xf32, #encoding_lhs> attributes {
+ hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz", {target_triple="riscv64-xyz-xyz", cpu_features="+v", iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}>
+} {
+ %0 = iree_encoding.set_encoding %arg0 encoding_dims{%m, %n, %k} : tensor<?x?xf32> -> tensor<?x?xf32, #encoding_lhs>
+ return %0 : tensor<?x?xf32, #encoding_lhs>
+}
+// CHECK-LABEL: func.func @matmul_set_encoding_LHS_f32_riscv64
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<?x?xf32>
+// CHECK: %[[PACK:.+]] = linalg.pack %[[ARG0]]
+// CHECK-SAME: outer_dims_perm = [0, 1]
+// CHECK-SAME: inner_dims_pos = [0, 1]
+// CHECK-SAME: inner_tiles = [7, 1]
+// CHECK-SAME: -> tensor<?x?x7x1xf32>
+// CHECK: return %[[PACK]]
+
+// RISC-V64 with V extension - set_encoding for RHS (f32)
+// The RHS operand corresponds to the N dimension, which IS scalable in when scalable vectorization is enabled.
+// In static mode: N tile is fixed based on VLEN (16 for VLEN=128)
+// In scalable mode: N tile is scalable (8 * vscale)
+func.func @matmul_set_encoding_RHS_f32_riscv64(%arg0: tensor<?x?xf32>, %m: index, %n: index, %k: index) -> tensor<?x?xf32, #encoding_rhs> attributes {
+ hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz", {target_triple="riscv64-xyz-xyz", cpu_features="+v", iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}>
+} {
+ %0 = iree_encoding.set_encoding %arg0 encoding_dims{%m, %n, %k} : tensor<?x?xf32> -> tensor<?x?xf32, #encoding_rhs>
+ return %0 : tensor<?x?xf32, #encoding_rhs>
+}
+// CHECK-LABEL: func.func @matmul_set_encoding_RHS_f32_riscv64
+// CHECK-SAME: %[[SRC:[a-zA-Z0-9]+]]: tensor<?x?xf32>
+// CHECK-DAG: %[[PAD:.+]] = arith.constant 0.000000e+00 : f32
+
+// SCALABLE-DAG: %[[C8:.*]] = arith.constant 8 : index
+// SCALABLE-DAG: %[[VSCALE:.*]] = vector.vscale
+// SCALABLE: %[[C8_VSCALE:.*]] = arith.muli %[[VSCALE]], %[[C8]] : index
+
+/// With dynamic dimensions the N dimension is padded up to a multiple of the
+/// (possibly scalable) inner tile, so both modes emit a padding value.
+// CHECK: %[[PACK:.+]] = linalg.pack %[[SRC]]
+// CHECK-SAME: padding_value(%[[PAD]] : f32)
+// CHECK-SAME: outer_dims_perm = [1, 0]
+// CHECK-SAME: inner_dims_pos = [1, 0]
+
+// STATIC-SAME: inner_tiles = [16, 1]
+// STATIC-SAME: -> tensor<?x?x16x1xf32>
+
+// SCALABLE-SAME: inner_tiles = [%[[C8_VSCALE]], 1]
+// SCALABLE-SAME: -> tensor<?x?x?x1xf32>
+
+// CHECK: return %[[PACK]]
+
+// RISC-V64 with V extension - unset_encoding for RESULT (f32)
+// The RESULT operand has both M and N dimensions. M is fixed, N differs by mode.
+func.func @matmul_unset_encoding_RESULT_f32_riscv64(%arg0: tensor<?x?xf32, #encoding_res>, %m: index, %n: index, %k: index) -> tensor<?x?xf32> attributes {
+ hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz", {target_triple="riscv64-xyz-xyz", cpu_features="+v", iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}>
+} {
+ %0 = iree_encoding.unset_encoding %arg0 encoding_dims{%m, %n, %k} : tensor<?x?xf32, #encoding_res> -> tensor<?x?xf32>{%m, %n}
+ return %0 : tensor<?x?xf32>
+}
+// CHECK-LABEL: func.func @matmul_unset_encoding_RESULT_f32_riscv64
+// STATIC-SAME: %[[INPUT:[a-zA-Z0-9]+]]: tensor<?x?x7x16xf32>
+// SCALABLE-SAME: %[[INPUT:[a-zA-Z0-9]+]]: tensor<?x?x7x?xf32>
+
+// CHECK-DAG: %[[EMPTY:.+]] = tensor.empty{{.*}} : tensor<?x?xf32>
+
+// SCALABLE-DAG: %[[C8:.+]] = arith.constant 8 : index
+// SCALABLE: %[[VSCALE:.+]] = vector.vscale
+// SCALABLE: %[[C8_VSCALE:.+]] = arith.muli %[[VSCALE]], %[[C8]] : index
+
+// CHECK: %[[UNPACK:.+]] = linalg.unpack %[[INPUT]]
+// CHECK-SAME: outer_dims_perm = [0, 1] inner_dims_pos = [0, 1]
+// STATIC-SAME: inner_tiles = [7, 16]
+// SCALABLE-SAME: inner_tiles = [7, %[[C8_VSCALE]]]
+
+// CHECK: return %[[UNPACK]] : tensor<?x?xf32>
+
+// RISC-V64 with V extension - matmul lowering (f32)
+func.func @matmul_lowering_f32f32f32_riscv64(
+ %lhs: tensor<?x?xf32, #encoding_lhs>,
+ %rhs: tensor<?x?xf32, #encoding_rhs>,
+ %result: tensor<?x?xf32, #encoding_res>
+) -> tensor<?x?xf32, #encoding_res> attributes {
+ hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz", {target_triple="riscv64-xyz-xyz", cpu_features="+v", iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}>
+} {
+ %matmul = linalg.matmul
+ ins(%lhs, %rhs : tensor<?x?xf32, #encoding_lhs>, tensor<?x?xf32, #encoding_rhs>)
+ outs(%result : tensor<?x?xf32, #encoding_res>)
+ -> tensor<?x?xf32, #encoding_res>
+ return %matmul : tensor<?x?xf32, #encoding_res>
+}
+
+// CHECK-LABEL: func @matmul_lowering_f32f32f32_riscv64(
+// STATIC-SAME: %[[LHS:[a-zA-Z0-9]+]]: tensor<?x?x7x1xf32>
+// STATIC-SAME: %[[RHS:[a-zA-Z0-9]+]]: tensor<?x?x16x1xf32>
+// STATIC-SAME: %[[OUTS:[a-zA-Z0-9]+]]: tensor<?x?x7x16xf32>
+
+// SCALABLE-SAME: %[[LHS:[a-zA-Z0-9]+]]: tensor<?x?x7x1xf32>
+// SCALABLE-SAME: %[[RHS:[a-zA-Z0-9]+]]: tensor<?x?x?x1xf32>
+// SCALABLE-SAME: %[[OUTS:[a-zA-Z0-9]+]]: tensor<?x?x7x?xf32>
+
+// CHECK: %[[MMT4D:.+]] = linalg.mmt4d
+// CHECK-SAME: ins(%[[LHS]], %[[RHS]] :
+// CHECK-SAME: outs(%[[OUTS]] :
+// CHECK: return %[[MMT4D]]
+
+// Checks that the +zvl512b flag indeed reflects on the data-tiled layout for the static case
+// and does not make a difference on the scalable case.
+
+func.func @matmul_lowering_f32f32f32_riscv64_zvl512b(
+ %lhs: tensor<?x?xf32, #encoding_lhs>,
+ %rhs: tensor<?x?xf32, #encoding_rhs>,
+ %result: tensor<?x?xf32, #encoding_res>
+) -> tensor<?x?xf32, #encoding_res> attributes {
+ hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz", {target_triple="riscv64-xyz-xyz", cpu_features="+v,+zvl512b", iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}>
+} {
+ %matmul = linalg.matmul
+ ins(%lhs, %rhs : tensor<?x?xf32, #encoding_lhs>, tensor<?x?xf32, #encoding_rhs>)
+ outs(%result : tensor<?x?xf32, #encoding_res>)
+ -> tensor<?x?xf32, #encoding_res>
+ return %matmul : tensor<?x?xf32, #encoding_res>
+}
+
+// CHECK-LABEL: func @matmul_lowering_f32f32f32_riscv64_zvl512b(
+// STATIC-SAME: %[[LHS:[a-zA-Z0-9]+]]: tensor<?x?x7x1xf32>
+
+// +v implies +zvl128b, so the data-tiled layout without the flag had 16 for the N dim.
+// +zvl512b therefore has 16 * (512/128) = 64
+
+// STATIC-SAME: %[[RHS:[a-zA-Z0-9]+]]: tensor<?x?x64x1xf32>
+// STATIC-SAME: %[[OUTS:[a-zA-Z0-9]+]]: tensor<?x?x7x64xf32>
+
+// SCALABLE-SAME: %[[LHS:[a-zA-Z0-9]+]]: tensor<?x?x7x1xf32>
+// SCALABLE-SAME: %[[RHS:[a-zA-Z0-9]+]]: tensor<?x?x?x1xf32>
+// SCALABLE-SAME: %[[OUTS:[a-zA-Z0-9]+]]: tensor<?x?x7x?xf32>
+
+// CHECK: %[[MMT4D:.+]] = linalg.mmt4d
+// CHECK-SAME: ins(%[[LHS]], %[[RHS]] :
+// CHECK-SAME: outs(%[[OUTS]] :
+// CHECK: return %[[MMT4D]]
+
// -----
+// RISC-V32 with ukernels - uses mmt4d with default tiles
+
#map = affine_map<(d0, d1, d2) -> (d0, d2)>
#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
@@ -60,3 +239,125 @@
// CHECK-SAME: ins(%[[LHS]], %[[RHS]]
// CHECK-SAME: outs(%[[ACC]]
// CHECK: return %[[MMT4D]]
+
+// -----
+
+// Checks that we don't transpose M/N for the narrow-N case when scalable inner tiles are present.
+// For scalable vectors, we want to keep the N dimension scalable.
+
+#map = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+#encoding_lhs = #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map, #map1, #map2], iteration_sizes = [100, 4, 500]>
+#encoding_rhs = #iree_encoding.encoding<operand_index = 1, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map, #map1, #map2], iteration_sizes = [100, 4, 500]>
+#encoding_result = #iree_encoding.encoding<operand_index = 2, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map, #map1, #map2], iteration_sizes = [100, 4, 500]>
+func.func @matmul_lowering_narrow_n_f32f32f32_riscv64(
+ %lhs: tensor<100x500xf32, #encoding_lhs>,
+ %rhs: tensor<500x4xf32, #encoding_rhs>,
+ %result: tensor<100x4xf32, #encoding_result>
+) -> tensor<100x4xf32, #encoding_result> attributes {
+ hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz", {target_triple="riscv64-xyz-xyz", cpu_features="+v", iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}>
+} {
+ %matmul = linalg.matmul
+ ins(%lhs, %rhs : tensor<100x500xf32, #encoding_lhs>, tensor<500x4xf32, #encoding_rhs>)
+ outs(%result : tensor<100x4xf32, #encoding_result>)
+ -> tensor<100x4xf32, #encoding_result>
+ return %matmul : tensor<100x4xf32, #encoding_result>
+}
+
+// CHECK-LABEL: func @matmul_lowering_narrow_n_f32f32f32_riscv64(
+// STATIC-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<7x500x16x1xf32>
+// STATIC-SAME: %[[ARG1:[a-zA-Z0-9]+]]: tensor<1x500x4x1xf32>
+// STATIC-SAME: %[[ARG2:[a-zA-Z0-9]+]]: tensor<1x7x4x16xf32>
+
+// SCALABLE-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<15x500x7x1xf32>
+// SCALABLE-SAME: %[[ARG1:[a-zA-Z0-9]+]]: tensor<?x500x?x1xf32>
+// SCALABLE-SAME: %[[ARG2:[a-zA-Z0-9]+]]: tensor<15x?x7x?xf32>
+
+// STATIC: %[[MMT4D:.+]] = linalg.mmt4d
+// STATIC-SAME: ins(%[[ARG1]], %[[ARG0]] :
+// STATIC-SAME: outs(%[[ARG2]] :
+
+// SCALABLE: %[[MMT4D:.+]] = linalg.mmt4d
+// SCALABLE-SAME: ins(%[[ARG0]], %[[ARG1]] :
+// SCALABLE-SAME: outs(%[[ARG2]] :
+
+// CHECK: return %[[MMT4D]]
+
+// -----
+
+// RISC-V64 with V extension and zvfh - full matmul lowering (f16)
+
+#map = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+#encoding_lhs = #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f16, f16, f16], user_indexing_maps = [#map, #map1, #map2], iteration_sizes = [?, ?, ?]>
+#encoding_rhs = #iree_encoding.encoding<operand_index = 1, op_type = matmul, element_types = [f16, f16, f16], user_indexing_maps = [#map, #map1, #map2], iteration_sizes = [?, ?, ?]>
+#encoding_result = #iree_encoding.encoding<operand_index = 2, op_type = matmul, element_types = [f16, f16, f16], user_indexing_maps = [#map, #map1, #map2], iteration_sizes = [?, ?, ?]>
+func.func @full_matmul_lowering_f16f16f16_riscv64(
+ %lhs: tensor<?x?xf16>,
+ %rhs: tensor<?x?xf16>,
+ %acc: tensor<?x?xf16>,
+ %m : index, %n : index, %k : index
+) -> tensor<?x?xf16> attributes {
+ hal.executable.target = #hal.executable.target<"llvm-cpu", "xyz", {target_triple="riscv64-xyz-xyz", cpu_features="+v,+zvfh", iree.encoding.resolver = #iree_cpu.cpu_encoding_resolver<>}>
+} {
+ %0 = iree_encoding.set_encoding %lhs encoding_dims{%m, %n, %k} : tensor<?x?xf16> -> tensor<?x?xf16, #encoding_lhs>
+ %1 = iree_encoding.set_encoding %rhs encoding_dims{%m, %n, %k} : tensor<?x?xf16> -> tensor<?x?xf16, #encoding_rhs>
+ %2 = iree_encoding.set_encoding %acc encoding_dims{%m, %n, %k} : tensor<?x?xf16> -> tensor<?x?xf16, #encoding_result>
+ %3 = linalg.matmul
+ ins(%0, %1 : tensor<?x?xf16, #encoding_lhs>,
+ tensor<?x?xf16, #encoding_rhs>)
+ outs(%2 : tensor<?x?xf16, #encoding_result>)
+ -> tensor<?x?xf16, #encoding_result>
+ %4 = iree_encoding.unset_encoding %3 encoding_dims{%m, %n, %k} : tensor<?x?xf16, #encoding_result> -> tensor<?x?xf16>{%m, %n}
+ return %4 : tensor<?x?xf16>
+}
+
+/// F16 on RISC-V64 with V+zvfh extension:
+/// For the static case: M=7, N=16 (vlen=128), K=1
+/// For the scalable case: M=7, N=8*vscale (vlen=64*vscale), K=1
+
+// CHECK-LABEL: func @full_matmul_lowering_f16f16f16_riscv64(
+// CHECK-SAME: %[[LHS:[a-zA-Z0-9]+]]: tensor<?x?xf16>
+// CHECK-SAME: %[[RHS:[a-zA-Z0-9]+]]: tensor<?x?xf16>
+// CHECK-SAME: %[[ACC:[a-zA-Z0-9]+]]: tensor<?x?xf16>
+
+// SCALABLE: %[[C8:.+]] = arith.constant 8
+
+/// Pack LHS: [M, K] -> [?, ?, 7, 1] (this happens before vscale computation)
+// CHECK: %[[PACK_LHS:.+]] = linalg.pack %[[LHS]]
+// CHECK-SAME: outer_dims_perm = [0, 1]
+// CHECK-SAME: inner_dims_pos = [0, 1]
+// CHECK-SAME: inner_tiles = [7, 1]
+
+// SCALABLE: %[[VSCALE:.+]] = vector.vscale
+// SCALABLE: %[[C8_VSCALE:.+]] = arith.muli %[[VSCALE]], %[[C8]]
+
+/// Pack RHS: [K, N] -> [?, ?, N_tile, 1] with outer_dims_perm = [1, 0]
+// CHECK: %[[PACK_RHS:.+]] = linalg.pack %[[RHS]]
+// CHECK-SAME: outer_dims_perm = [1, 0]
+// CHECK-SAME: inner_dims_pos = [1, 0]
+// STATIC-SAME: inner_tiles = [16, 1]
+// SCALABLE-SAME: inner_tiles = [%[[C8_VSCALE]], 1]
+
+/// Pack ACC: [M, N] -> [?, ?, 7, N_tile]
+// CHECK: %[[PACK_ACC:.+]] = linalg.pack %[[ACC]]
+// CHECK-SAME: outer_dims_perm = [0, 1]
+// CHECK-SAME: inner_dims_pos = [0, 1]
+// STATIC-SAME: inner_tiles = [7, 16]
+// SCALABLE-SAME: inner_tiles = [7, %[[C8_VSCALE]]]
+
+/// The mmt4d operation
+// CHECK: %[[MMT4D:.+]] = linalg.mmt4d
+// CHECK-SAME: ins(%[[PACK_LHS]], %[[PACK_RHS]] :
+// CHECK-SAME: outs(%[[PACK_ACC]] :
+
+/// Unpack result: [?, ?, 7, N_tile] -> [M, N]
+// CHECK: %[[UNPACK:.+]] = linalg.unpack %[[MMT4D]]
+// CHECK-SAME: outer_dims_perm = [0, 1]
+// CHECK-SAME: inner_dims_pos = [0, 1]
+// STATIC-SAME: inner_tiles = [7, 16]
+// SCALABLE-SAME: inner_tiles = [7, %[[C8_VSCALE]]]
+
+// CHECK: return %[[UNPACK]]
diff --git a/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CPUEncodingExternalModels.cpp b/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CPUEncodingExternalModels.cpp
index a9e0670..cd14497 100644
--- a/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CPUEncodingExternalModels.cpp
+++ b/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CPUEncodingExternalModels.cpp
@@ -71,14 +71,30 @@
// Utilities.
//===----------------------------------------------------------------------===//
+/// Determines which tile dimensions should be scalable for a given target.
+/// For AArch64 SVE/SVE2 and RISC-V with the V extension (when not using static
+/// vectors), the N dimension is made scalable to take advantage of scalable
+/// vector registers.
+/// Returns failure if scalable tiling is not applicable for the target.
static FailureOr<IREE::Codegen::ScalableTileFlags>
getScalableTileFlags(linalg::ContractionDimensions cDims,
IREE::Encoding::EncodingAttr encoding,
DictionaryAttr config) {
// TODO(egebeysel): I think this isScalable*Enabled flag should be temporary
// and the temporary SME flag should probably come next to it.
- if (!isAArch64(config) || !isScalableVectorizationEnabled()) {
- LDBG() << "Pre-conditions to enable scalable tiling are not met!";
+ if (!isScalableVectorizationEnabled()) {
+ LDBG() << "Scalable vectorization is not enabled!";
+ return failure();
+ }
+
+ // Check if target supports scalable vectors.
+ bool hasAArch64ScalableSupport =
+ isAArch64(config) &&
+ (hasFeature(config, "+sve") || hasFeature(config, "+sve2"));
+ bool hasRISCVScalableSupport = isRISCV64(config) && hasFeature(config, "+v");
+
+ if (!hasAArch64ScalableSupport && !hasRISCVScalableSupport) {
+ LDBG() << "Target does not support scalable vectors!";
return failure();
}
@@ -90,7 +106,9 @@
: encoding.mapDimToOperandIndex(cDims.n[0]);
std::optional<unsigned> kDim = encoding.mapDimToOperandIndex(cDims.k[0]);
IREE::Codegen::ScalableTileFlags scalableTiles;
- // TODO(egebeysel): Add logic for SME.
+
+ // M dimension: fixed for both AArch64 SVE and RISC-V V extension.
+ // TODO(egebeysel): Add logic for SME which may have scalable M.
if (mDim.has_value()) {
if (hasFeature(config, "+sme")) {
LDBG() << "SME with data-tiling is not supported yet!";
@@ -98,10 +116,17 @@
}
scalableTiles.push_back(false);
}
+
+ // N dimension: scalable for both AArch64 SVE/SVE2 and RISC-V V extension
+ // (when using scalable vectors). This allows the tile to grow with the
+ // vector register length.
if (nDim.has_value()) {
- scalableTiles.push_back(hasFeature(config, "+sve") ||
- hasFeature(config, "+sve2"));
+ scalableTiles.push_back(hasAArch64ScalableSupport ||
+ hasRISCVScalableSupport);
}
+
+ // K dimension: fixed for both targets. The reduction dimension does not
+ // benefit from scalability in the same way as the N dimension.
if (kDim.has_value()) {
scalableTiles.push_back(false);
}
@@ -219,7 +244,7 @@
IREE::Encoding::MatmulNarrowDim narrowDim) {
// Handle narrow-N by transposing to reduce to narrow-M. Note: the
// enumeratedTiles currently only enumerate narrow-M cases.
- if (narrowDim.isN()) {
+ if (narrowDim.isN() && !isScalableVectorizationEnabled()) {
narrowDim.dim = IREE::Encoding::MatmulNarrowDim::Dim::M;
TileMxNxK tile = chooseMatmulTile(enumeratedTiles, narrowDim);
std::swap(tile.M, tile.N);
@@ -994,18 +1019,42 @@
// Enumerate tile sizes to choose from on riscv64.
// For narrow-{M,N} cases, this only enumerates on narrow M. The narrow-N cases
// are handled by transposition in chooseMatmulTile.
+//
+// When using static VLEN-based tile sizes, the N dimension is computed from the
+// VLEN using zvl* features.
+//
+// When using scalable vectors
+// (iree-llvmcpu-enable-scalable-vectorization=true), we use smaller base tile
+// sizes that scale with vscale (VLEN/64). The base tile sizes assume VLEN=64
+// bits, meaning vscale=1 at the minimum. When scalable vectorization is
+// enabled, these base sizes will be multiplied by vscale at runtime to utilize
+// the full vector register width.
+//
+// Base tile calculation for VLEN=64:
+// - f32 (32 bits): 64/32 = 2 elements per vscale
+// - f16 (16 bits): 64/16 = 4 elements per vscale
static SmallVector<TileMxNxK>
enumerateMatmulTileRiscv64(TypeRange elementTypes, DictionaryAttr config) {
-
- // Data-Tiling is only implemented for the V extension
+ // Data-tiling is only implemented for the V extension.
if (!hasFeature(config, "+v")) {
return {};
}
- size_t vlen = getRISCVVVlenFromCPUFeatures(config);
+
assert(elementTypes.size() == 3);
Type lhs = elementTypes[0];
Type rhs = elementTypes[1];
Type out = elementTypes[2];
+
+ // RVV has both static and scalable vector modes.
+ // Static mode: use VLEN-based tile sizes from zvl* features.
+ // Scalable mode: use base VLEN=64 bits as the base size, the vscale
+ // multiplier and the correct scalable data-tiled layout will be emitted
+ // accordingly in `getScalableTileFlags`.
+ constexpr size_t kBaseVlen = 64;
+ size_t vlen = isScalableVectorizationEnabled()
+ ? kBaseVlen
+ : getRISCVVVlenFromCPUFeatures(config);
+
if (lhs.isF32() && rhs.isF32() && out.isF32()) {
// VLEN-aware Tile size selection
// One concern that needs to be addressed here is that
@@ -1020,10 +1069,13 @@
};
}
if (lhs.isF16() && rhs.isF16()) {
+ // TODO(egebeysel): currently, the F16 path uses LMUL=2 and therefore 14
+ // registers for accumulators. This has to be investigated and the
+ // corresponding microkernels have to be adjusted if need be.
int N0 = vlen / 8;
if (hasFeature(config, "+zvfh")) {
return {
- TileMxNxK{7, N0, 1},
+ TileMxNxK{7, N0, 1}, // Aim to use vfmacc.
TileMxNxK{4, N0, 1}, // Truncation of the above.
TileMxNxK{2, N0, 1}, // Truncation of the above.
TileMxNxK{1, N0, 1}, // Truncation of the above.
@@ -1031,7 +1083,7 @@
}
if (hasFeature(config, "+zvfhmin")) {
return {
- TileMxNxK{6, N0, 1},
+ TileMxNxK{6, N0, 1}, // Aim to use vfmacc.
TileMxNxK{4, N0, 1}, // Truncation of the above.
TileMxNxK{2, N0, 1}, // Truncation of the above.
TileMxNxK{1, N0, 1}, // Truncation of the above.
@@ -1326,6 +1378,27 @@
return info;
}
auto narrowDim = IREE::Encoding::getPo2MatmulNarrowDim(encoding);
+ // In case of scalable vectors, we explicitly use the narrow M dimension
+ // config to use the M dimension for calculating padding penalties while we
+ // choose the layout.
+ // TODO(egebeysel): reconcile this with the static case. Currently,
+ // this results in vector types with non-trailing scalable dimensions, e.g.
+ // vector<[8]x8xf32>, which cannot be legalized or lowered to LLVM properly.
+ // LLVM only handles vector types with only-trailing scalable dimensions
+ // properly, e.g. vector<8x[8]xf32> is lowered to !llvm.array<8 x
+ // vector<[8]xf32>>. The former have to be legalized, so that they can be
+ // lowered to an array of scalable vectors, but the patterns are missing.
+ if (isScalableVectorizationEnabled() && narrowDim.isN() &&
+ !cDims.m.empty()) {
+ int64_t m = getEncodingContractionLikeSizes(encoding).value().M;
+ if (ShapedType::isDynamic(m)) {
+ narrowDim = {};
+ } else {
+ narrowDim = {IREE::Encoding::MatmulNarrowDim::Dim::M, m};
+ }
+ }
+ // Choose a final matmul TileMxNxK from the above-enumerated tile shapes,
+ // taking narrow dimensions into account.
TileMxNxK chosenTileMxNxK =
chooseMatmulTile(enumeratedTileMxNxK, narrowDim);
FailureOr<MaterializeEncodingInfo> maybeEncodingInfo =
