[mlir][linalg] Support minus reduction

[guraypp]

Summary

This revision binds the minus reduction to add to the reduction. It also adds a test with transform dialect that tests the vectorization of the sub reduction.

[springerm]

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

108

typo

mlir/test/Dialect/Linalg/transform-op-vectorize.mlir

171

should this be a subtraction?

[guraypp]

address typos

[guraypp]

We found more issues with that revision. It needs more work.

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

enum VectorizationStatus {
   NoReplace,
   /// Op vectorized into a new Op whose results will replace original Op's
   /// results.
   NewOp
   // TODO: support values if Op vectorized to Many-Ops whose results we need to
   // aggregate for replacement.
 };
 struct VectorizationResult {
   /// Return status from vectorizing the current op.
   enum VectorizationStatus status = VectorizationStatus::Failure;
   /// New vectorized operation to replace the current op.
   /// Replacement behavior is specified by `status`.
   Operation *newOp;
 };
 
 llvm::Optional<vector::CombiningKind>
 mlir::linalg::getCombinerOpKind(Operation *combinerOp) {
   using ::mlir::vector::CombiningKind;
 
   if (!combinerOp)
     return llvm::None;
   return llvm::TypeSwitch<Operation *, llvm::Optional<CombiningKind>>(
              combinerOp)
       .Case<arith::AddIOp, arith::AddFOp>(
           [&](auto op) { return CombiningKind::ADD; })
+      .Case<arith::SubIOp>([&](auto op) { return CombiningKind::ADD; })
+      .Case<arith::SubFOp>([&](auto op) { return CombiningKind::ADD; })
       .Case<arith::AndIOp>([&](auto op) { return CombiningKind::AND; })
       .Case<arith::MaxSIOp>([&](auto op) { return CombiningKind::MAXSI; })
       .Case<arith::MaxFOp>([&](auto op) { return CombiningKind::MAXF; })
       .Case<arith::MinSIOp>([&](auto op) { return CombiningKind::MINSI; })
       .Case<arith::MinFOp>([&](auto op) { return CombiningKind::MINF; })
       .Case<arith::MulIOp, arith::MulFOp>(
           [&](auto op) { return CombiningKind::MUL; })
       .Case<arith::OrIOp>([&](auto op) { return CombiningKind::OR; })

mlir/test/Dialect/Linalg/transform-op-vectorize.mlir

 ^bb0(%arg0: !pdl.operation):
   transform.sequence %arg0 failures(propagate) {
   ^bb1(%arg1: !pdl.operation):
     %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1
     // expected-error @below {{op requires isolated-from-above targets}}
     %2 = transform.structured.vectorize %0
   }
 }
+
+// -----
+
+// CHECK-LABEL: @vectorize_subred
+func.func @vectorize_subred(%arg0: tensor<32xf32>, %out: tensor<f32>) -> tensor<f32> {
+  // CHECK: %[[vA:.+]] = vector.transfer_read 
+  // CHECK: %[[vB:.+]] = vector.transfer_read 
+  // CHECK: %[[vC:.+]] = vector.extractelement
+  // CHECK: %[[vR:.+]] = vector.multi_reduction <add>, %[[vA]], %[[vC]]
+    %0 = linalg.generic {
+        indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> ()>],
+        iterator_types = ["reduction"]} 
+        ins(%arg0 : tensor<32xf32>)
+        outs(%out : tensor<f32>) {
+        ^bb0(%arg7: f32, %arg9: f32):
+            %1 = arith.subf %arg7, %arg9 : f32            
+            linalg.yield %1 : f32
+    } -> tensor<f32>
+    return %0 : tensor<f32>
+}
+
+transform.with_pdl_patterns {
+  ^bb0(%IN_MAT1: !pdl.operation):
+  transform.sequence %IN_MAT1 failures(propagate) {
+    ^bb1(%IN_MAT2: !pdl.operation):
+    %0 = transform.structured.match ops{["linalg.generic"]} in %IN_MAT2
+    %1 = get_closest_isolated_parent %0
+    transform.structured.vectorize %1
+  }
+}