Apply pattern for populateVectorMultiReductionLoweringPatterns from… (#1312)

… mlir-opt in iree-opt

For issue #1306, converting vector.multi_reduction to vector.reduction
is the first step as `vector.reduction` can be converted to Peano
instructions which can vectorize. `mlir-opt` supports it with the flag
`lower-vector-multi-reduction` and the
`populateVectorMultiReductionLoweringPatterns` pattern, reuse it in
IREE, as a last pass to kick in `AMDAIEVectorizationPass` pipeline.

---------

Co-authored-by: James Newling <[email protected]>
Co-authored-by: Vivian Zhang <[email protected]>

Author: 3 people

compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIEVectorization.cpp

@@ -115,6 +115,9 @@ void AMDAIEVectorizationPass::runOnOperation() {
   vector::populateVectorTransferPermutationMapLoweringPatterns(
       vectorizationPatterns);
 
+  vector::populateVectorMultiReductionLoweringPatterns(
+      vectorizationPatterns,
+      vector::VectorMultiReductionLowering::InnerReduction);
   (void)applyPatternsGreedily(funcOp, std::move(vectorizationPatterns));
 }
 }  // namespace

compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/test/CMakeLists.txt

@@ -73,6 +73,7 @@ iree_lit_test_suite(
     "lowering_strategy_objectfifo_npu4.mlir"
     "lowering_strategy_softmax.mlir"
     "map_forall_to_cores.mlir"
+    "multi_reduction_to_reduction.mlir"
     "none_access_to_temporary_buffer.mlir"
     "normalize_loop_bounds.mlir"
     "npu_dma_to_half_dma_cpy_nd.mlir"

compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/test/multi_reduction_to_reduction.mlir

@@ -0,0 +1,42 @@
+// RUN: iree-opt --pass-pipeline='builtin.module(func.func(iree-amdaie-vectorization))' %s | FileCheck %s
+
+// Make sure it's not falling on the InnerParallel pattern
+// CHECK-LABEL: func.func @multi_reduction_innerparallel
+func.func @multi_reduction_innerparallel(%v : vector<4xf16>, %acc: f16) -> f16 {
+  // CHECK-NOT: vector.extract %{{.*}}[0] : f16 from vector<4xf16>
+  // CHECK-NOT: arith.addf %{{.*}}, %{{.*}} : f16
+  %0 = vector.multi_reduction <add>, %v, %acc[0] : vector<4xf16> to f16
+  return %0 : f16
+}
+
+// CHECK-LABEL: func.func @multi_reduction_2d
+func.func @multi_reduction_2d(%v : vector<4x6xf32>, %acc: vector<4xf32>) -> vector<4xf32> {
+  // CHECK: vector.reduction <add>,
+  // CHECK: vector.reduction <add>,
+  // CHECK: vector.reduction <add>,
+  // CHECK: vector.reduction <add>,
+  // CHECK-NOT: vector.reduction <add>,
+  %0 = vector.multi_reduction #vector.kind<add>, %v, %acc [1] : vector<4x6xf32> to vector<4xf32>
+  return %0 : vector<4xf32>
+}
+
+// CHECK-LABEL: func.func @multi_reduction_1d
+func.func @multi_reduction_1d(%v : vector<4xf32>, %acc: f32) -> f32 {
+  // CHECK: vector.reduction <add>, %{{.*}}, %{{.*}} : vector<4xf32> into f32
+  %0 = vector.multi_reduction <add>, %v, %acc[0] : vector<4xf32> to f32
+  return %0 : f32
+}
+
+// CHECK-LABEL: func.func @multi_reduction_1d_mul
+func.func @multi_reduction_1d_mul(%v : vector<4xf32>, %acc: f32) -> f32 {
+  // CHECK: vector.reduction <mul>, %{{.*}}, %{{.*}} : vector<4xf32> into f32
+  %0 = vector.multi_reduction <mul>, %v, %acc[0] : vector<4xf32> to f32
+  return %0 : f32
+}
+
+// CHECK-LABEL: func.func @multi_reduction_1d_bf16
+func.func @multi_reduction_1d_bf16(%v : vector<32xbf16>, %acc: bf16) -> bf16 {
+  // CHECK: vector.reduction <add>, %{{.*}}, %{{.*}} : vector<32xbf16> into bf16
+  %0 = vector.multi_reduction <add>, %v, %acc[0] : vector<32xbf16> to bf16
+  return %0 : bf16
+}