[CPU] Scalarize vector loads used in a scalar fashion (#14150)
This PR makes sure that we scalarize the vector load that are used in a
scalar fashion (i.e., vector.transfer_read + vector.extract) in matmuls.
The LLVM backend is sometimes able to scalarize this vector load but not
always. For cases where it's not scalarized, register pressure increases
as the liveness of this vector load spans across all the unrolled
iterations of the matmul. There are also some shuffles/permutations
generated that don't make too much sense from the performance point of
view.
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUVectorLowering.cpp b/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUVectorLowering.cpp
index abc7d8e..5aec874 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUVectorLowering.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUVectorLowering.cpp
@@ -84,7 +84,12 @@
vector::populateVectorContractLoweringPatterns(
patterns, vectorTransformOptions,
/*benefit=*/1,
- /*disableOuterProductLowering=*/true);
+ /*disableOuterProductLowering=*/false);
+ // This pattern will transform vector loads whose elements are used in a
+ // scalar fashion into scalar loads. This will let scalar loads to be folded
+ // into broadcast/arithmetic operations and reduce register pressure.
+ vector::populateScalarVectorTransferLoweringPatterns(
+ patterns, /*benefit=*/1, /*allowMultipleUses=*/true);
vector::populateVectorTransferPermutationMapLoweringPatterns(patterns);
vector::populateVectorMultiReductionLoweringPatterns(
patterns, vectorMultiReductionLowering);
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/pad_conv_pipeline_tests.mlir b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/pad_conv_pipeline_tests.mlir
index 98af3d0..11e6ea1 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/pad_conv_pipeline_tests.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/pad_conv_pipeline_tests.mlir
@@ -59,5 +59,5 @@
// Check that the stack buffer is bounded by tiling sizes.
//
// CHECK: memref.alloca() {alignment = 64 : i64} : memref<1x8x1x8xf32>
-// CHECK: vector.outerproduct %{{.+}}, %{{.+}}, %{{.+}} {kind = #vector.kind<add>}
+// CHECK: vector.fma
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/pad_pipeline_tests.mlir b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/pad_pipeline_tests.mlir
index 3a90248..182021c 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/pad_pipeline_tests.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/pad_pipeline_tests.mlir
@@ -213,7 +213,7 @@
// CHECK: scf.for
// CHECK-COUNT-7: vector.load %[[INPUT]]
// CHECK-COUNT-8: vector.load %[[FILTER_SUBVIEW]]
-// CHECK-COUNT-8: vector.outerproduct
+// CHECK-COUNT-8: vector.fma
// CHECK: scf.yield
// CHECK: scf.yield
// CHECK: scf.yield
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/peel_and_vectorize.mlir b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/peel_and_vectorize.mlir
index 4ea2909..4ba2af9 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/peel_and_vectorize.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/peel_and_vectorize.mlir
@@ -53,7 +53,7 @@
// CHECK: scf.for
// CHECK: scf.for
// CHECK: scf.for
-// CHECK: vector.outerproduct
+// CHECK: vector.fma
// CHECK-NOT: scf.for
// -----
@@ -109,7 +109,7 @@
// CHECK: scf.for
// CHECK: scf.for
// CHECK: scf.for
-// CHECK: vector.outerproduct
+// CHECK: vector.fma
// 2nd dim peeling:
// CHECK: scf.for
@@ -187,7 +187,7 @@
// CHECK: scf.for
// CHECK: scf.for
// CHECK: scf.for
-// CHECK: vector.outerproduct
+// CHECK: vector.fma
// 1nd dim peeling:
// CHECK: scf.for
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/pipeline_tests.mlir b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/pipeline_tests.mlir
index 566cea6..adff8f4 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/pipeline_tests.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/pipeline_tests.mlir
@@ -106,7 +106,7 @@
}
}
// CHECK-LABEL: func.func @preset_pad_config_matmul
-// CHECK: vector.outerproduct
+// CHECK: vector.fma
// -----
@@ -166,7 +166,7 @@
// CHECK-DAG: memref.alloca() {{.+}} memref<8x16xf32>
// CHECK-DAG: memref.alloca() {{.+}} memref<16x32xf32>
// CHECK-DAG: memref.alloca() {{.+}} memref<8x32xf32>
-// CHECK: vector.outerproduct
+// CHECK: vector.fma
// -----
@@ -226,7 +226,7 @@
// CHECK-DAG: memref.alloca() {{.+}} memref<1x32xf32>
// CHECK-DAG: memref.alloca() {{.+}} memref<1x32xf32>
// CHECK-NOT: memref.alloca
-// CHECK: vector.outerproduct
+// CHECK: vector.fma
// CHECK: arith.addf {{.*}} : vector<
// CHECK: arith.maxf {{.*}} : vector<
@@ -281,7 +281,7 @@
}
}
// CHECK-LABEL: func.func @batch_matmul_dynamic
-// CHECK: vector.outerproduct
+// CHECK: vector.fma
// -----
@@ -392,7 +392,7 @@
// CHECK-LABEL: func.func @vectorize_fill_conv2d_generic
// CHECK-NOT: memref.alloca
// CHECK-NOT: linalg.fill
-// CHECK: vector.outerproduct %{{.+}}, %{{.+}}, %{{.+}} {kind = #vector.kind<add>}
+// CHECK: vector.fma
// CHECK-NOT: linalg.generic
// CHECK: arith.cmpf olt, %{{.+}}, %{{.+}} : vector<4x4xf32>
@@ -441,7 +441,7 @@
// CHECK: scf.for
// CHECK: scf.for
// CHECK: scf.for
-// CHECK-COUNT-16: vector.outerproduct
+// CHECK-COUNT-16: vector.fma
// CHECK: arith.addf %{{.+}}, %{{.+}} : vector<8x32xf32>
// -----
@@ -501,7 +501,22 @@
// CHECK: scf.for
// CHECK: scf.for
// CHECK: scf.for
-// CHECK: vector.outerproduct
+// CHECK: arith.muli
+// CHECK-NEXT: arith.addi
+// CHECK: arith.muli
+// CHECK-NEXT: arith.addi
+// CHECK: arith.muli
+// CHECK-NEXT: arith.addi
+// CHECK: arith.muli
+// CHECK-NEXT: arith.addi
+// CHECK: arith.muli
+// CHECK-NEXT: arith.addi
+// CHECK: arith.muli
+// CHECK-NEXT: arith.addi
+// CHECK: arith.muli
+// CHECK-NEXT: arith.addi
+// CHECK: arith.muli
+// CHECK-NEXT: arith.addi
// CHECK: %{{.+}} = "tosa.apply_scale"({{.+}}) <{double_round = true}> : (vector<8x32xi32>, vector<8x32xi32>, vector<8x32xi8>) -> vector<8x32xi32>
// -----
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/vector_lowering.mlir b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/vector_lowering.mlir
index 758ff8c..7ef532a 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/vector_lowering.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/vector_lowering.mlir
@@ -52,19 +52,76 @@
}
return
}
-// CHECK-LABEL: func.func @matmul_391x384x384_f32
-// CHECK: scf.for
-// CHECK: scf.for
-// CHECK: %{{.*}} scf.if
-// CHECK: scf.yield %{{.*}}
-// CHECK: } else {
-// CHECK: linalg.fill
-// CHECK: memref.copy
-// CHECK: scf.yield %{{.*}}
-// CHECK: }
-// CHECK: scf.for
-// CHECK: vector.outerproduct
-// CHECK: %{{.*}} = scf.if
-// CHECK: scf.if %{{.*}} {
-// CHECK: memref.copy
-// CHECK: }
+
+// CHECK-LABEL: func.func @matmul_391x384x384_f32
+// CHECK: scf.for
+// CHECK: scf.for
+// CHECK: %{{.*}} scf.if
+// CHECK: scf.yield %{{.*}}
+// CHECK: } else {
+// CHECK: linalg.fill
+// CHECK: memref.copy
+// CHECK: scf.yield %{{.*}}
+// CHECK: }
+// CHECK: scf.for
+// CHECK-COUNT-128: vector.fma
+// CHECK: %{{.*}} = scf.if
+// CHECK: scf.if %{{.*}} {
+// CHECK: memref.copy
+// CHECK: }
+
+// -----
+
+// Check that vector.loads whose elements are extracted and
+// consumed in a scalar fashion are scalarized.
+
+func.func @matmul_scalar_loads() {
+ %cst = arith.constant 0.000000e+00 : f32
+ %c0 = arith.constant 0 : index
+ %c384 = arith.constant 384 : index
+ %c32 = arith.constant 32 : index
+ %c8 = arith.constant 8 : index
+ %c128 = arith.constant 128 : index
+ %c16 = arith.constant 16 : index
+ %cst_0 = arith.constant dense<0.000000e+00> : vector<8x32xf32>
+ %cst_1 = arith.constant dense<6.000000e+00> : vector<8x32xf32>
+ %alloca = memref.alloca() {alignment = 64 : i64} : memref<8x32xf32>
+ %0 = hal.interface.binding.subspan set(0) binding(0) type(storage_buffer) alignment(64) offset(%c0) flags(ReadOnly) : memref<391x384xf32>
+ memref.assume_alignment %0, 64 : memref<391x384xf32>
+ %1 = hal.interface.binding.subspan set(0) binding(1) type(storage_buffer) alignment(64) offset(%c0) flags(ReadOnly) : memref<384x384xf32>
+ memref.assume_alignment %1, 64 : memref<384x384xf32>
+ %2 = hal.interface.binding.subspan set(0) binding(2) type(storage_buffer) alignment(64) offset(%c0) flags(ReadOnly) : memref<384xf32>
+ memref.assume_alignment %2, 64 : memref<384xf32>
+ %3 = hal.interface.binding.subspan set(0) binding(3) type(storage_buffer) alignment(64) offset(%c0) : memref<391x384xf32>
+ memref.assume_alignment %3, 64 : memref<391x384xf32>
+ %workgroup_id_x = hal.interface.workgroup.id[0] : index
+ %workgroup_id_y = hal.interface.workgroup.id[1] : index
+ %4 = affine.apply affine_map<()[s0] -> (s0 * 128)>()[%workgroup_id_y]
+ %5 = affine.apply affine_map<()[s0] -> (s0 * 128)>()[%workgroup_id_x]
+ %6 = affine.min affine_map<(d0) -> (-d0 + 391, 128)>(%4)
+ %subview = memref.subview %0[%4, 0] [%6, 384] [1, 1] : memref<391x384xf32> to memref<?x384xf32, strided<[384, 1], offset: ?>>
+ %subview_2 = memref.subview %3[%4, %5] [%6, 128] [1, 1] : memref<391x384xf32> to memref<?x128xf32, strided<[384, 1], offset: ?>>
+ %subview_3 = memref.subview %2[%5] [128] [1] : memref<384xf32> to memref<128xf32, strided<[1], offset: ?>>
+ %subview_4 = memref.subview %1[0, %5] [384, 128] [1, 1] : memref<384x384xf32> to memref<384x128xf32, strided<[384, 1], offset: ?>>
+ scf.for %arg0 = %c0 to %6 step %c8 {
+ %7 = affine.min affine_map<(d0)[s0] -> (-d0 + s0, 8)>(%arg0)[%6]
+ %subview_5 = memref.subview %subview[%arg0, 0] [%7, 384] [1, 1] : memref<?x384xf32, strided<[384, 1], offset: ?>> to memref<?x384xf32, strided<[384, 1], offset: ?>>
+ scf.for %arg1 = %c0 to %c128 step %c32 {
+ vector.transfer_write %cst_0, %alloca[%c0, %c0] {in_bounds = [true, true]} : vector<8x32xf32>, memref<8x32xf32>
+ %subview_6 = memref.subview %alloca[0, 0] [%7, 32] [1, 1] : memref<8x32xf32> to memref<?x32xf32, strided<[32, 1]>>
+ %8 = vector.transfer_read %subview_6[%c0, %c0], %cst {in_bounds = [false, true]} : memref<?x32xf32, strided<[32, 1]>>, vector<8x32xf32>
+ %9 = scf.for %arg2 = %c0 to %c384 step %c16 iter_args(%arg3 = %8) -> (vector<8x32xf32>) {
+ %15 = vector.transfer_read %subview_5[%c0, %arg2], %cst : memref<?x384xf32, strided<[384, 1], offset: ?>>, vector<8x16xf32>
+ %16 = vector.transfer_read %subview_4[%arg2, %arg1], %cst {in_bounds = [true, true]} : memref<384x128xf32, strided<[384, 1], offset: ?>>, vector<16x32xf32>
+ %17 = vector.contract {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d2)>, affine_map<(d0, d1, d2) -> (d2, d1)>, affine_map<(d0, d1, d2) -> (d0, d1)>], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %15, %16, %arg3 : vector<8x16xf32>, vector<16x32xf32> into vector<8x32xf32>
+ scf.yield %17 : vector<8x32xf32>
+ }
+ %subview_7 = memref.subview %subview_2[%arg0, %arg1] [%7, 32] [1, 1] : memref<?x128xf32, strided<[384, 1], offset: ?>> to memref<?x32xf32, strided<[384, 1], offset: ?>>
+ vector.transfer_write %9, %subview_7[%c0, %c0] {in_bounds = [false, true]} : vector<8x32xf32>, memref<?x32xf32, strided<[384, 1], offset: ?>>
+ }
+ }
+ return
+}
+
+// CHECK-LABEL: func.func @matmul_scalar_load
+// CHECK-COUNT-128: memref.load
