[mlir][spirv] Add conversion from vector.reduction

[antiagainst]

Summary

Only supports addition and multiplication for now; other cases
to be implemented.

[hanchung]

mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp

265

nit: we can use getVectorType() method. (and we don't have to check if dyn_cast succeeds or not.

[antiagainst]

mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp

265

I don't think there is a getVectorType() API right now.

mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp

 
 #include "../PassDetail.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVTypes.h"
 #include "mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinTypes.h"
 #include "mlir/Transforms/DialectConversion.h"
+#include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/STLExtras.h"
 #include <numeric>
 
 using namespace mlir;
 
 /// Gets the first integer value from `attr`, assuming it is an integer array
 /// attribute.
 static uint64_t getFirstIntValue(ArrayAttr attr) {
matchAndRewrite(vector::InsertStridedSliceOp insertOp, OpAdaptor adaptor,
     rewriter.replaceOpWithNewOp<spirv::VectorShuffleOp>(
         insertOp, dstVector.getType(), dstVector, srcVector,
         rewriter.getI32ArrayAttr(indices));
 
     return success();
   }
 };
 
+struct VectorReductionPattern final
+    : public OpConversionPattern<vector::ReductionOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(vector::ReductionOp reduceOp, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    Type resultType = typeConverter->convertType(reduceOp.getType());
+    if (!resultType)
+      return failure();
+
+    auto srcVectorType = adaptor.getVector().getType().dyn_cast<VectorType>();
+    if (!srcVectorType || srcVectorType.getRank() != 1)
+      return rewriter.notifyMatchFailure(reduceOp, "not 1-D vector source");
+
+    // Extract all elements.
+    int numElements = srcVectorType.getDimSize(0);
+    SmallVector<Value, 4> values;
+    values.reserve(numElements + (adaptor.getAcc() != nullptr));
+    Location loc = reduceOp.getLoc();
+    for (int i = 0; i < numElements; ++i) {
+      values.push_back(rewriter.create<spirv::CompositeExtractOp>(
+          loc, srcVectorType.getElementType(), adaptor.getVector(),
+          rewriter.getI32ArrayAttr({i})));
+    }
+    if (Value acc = adaptor.getAcc())
+      values.push_back(acc);
+
+    // Reduce them.
+    Value result = values.front();
+    for (Value next : llvm::makeArrayRef(values).drop_front()) {
+      switch (reduceOp.getKind()) {
+#define INT_FLOAT_CASE(kind, iop, fop)                                         \
+  case vector::CombiningKind::kind:                                            \
+    if (resultType.isa<IntegerType>()) {                                       \
+      result = rewriter.create<spirv::iop>(loc, resultType, result, next);     \
+    } else {                                                                   \
+      assert(resultType.isa<FloatType>());                                     \
+      result = rewriter.create<spirv::fop>(loc, resultType, result, next);     \
+    }                                                                          \
+    break
+
+        INT_FLOAT_CASE(ADD, IAddOp, FAddOp);
+        INT_FLOAT_CASE(MUL, IMulOp, FMulOp);
+
+      case vector::CombiningKind::MINUI:
+      case vector::CombiningKind::MINSI:
+      case vector::CombiningKind::MINF:
+      case vector::CombiningKind::MAXUI:
+      case vector::CombiningKind::MAXSI:
+      case vector::CombiningKind::MAXF:
+      case vector::CombiningKind::AND:
+      case vector::CombiningKind::OR:
+      case vector::CombiningKind::XOR:
+        return rewriter.notifyMatchFailure(reduceOp, "unimplemented");
+      }
+    }
+
+    rewriter.replaceOp(reduceOp, result);
+    return success();
+  }
+};
+
 class VectorSplatPattern final : public OpConversionPattern<vector::SplatOp> {
 public:
   using OpConversionPattern<vector::SplatOp>::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(vector::SplatOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     VectorType dstVecType = op.getType();

 } // namespace
 
 void mlir::populateVectorToSPIRVPatterns(SPIRVTypeConverter &typeConverter,
                                          RewritePatternSet &patterns) {
   patterns.add<VectorBitcastConvert, VectorBroadcastConvert,
                VectorExtractElementOpConvert, VectorExtractOpConvert,
                VectorExtractStridedSliceOpConvert, VectorFmaOpConvert,
                VectorInsertElementOpConvert, VectorInsertOpConvert,
-               VectorInsertStridedSliceOpConvert, VectorShuffleOpConvert,
-               VectorSplatPattern>(typeConverter, patterns.getContext());
+               VectorReductionPattern, VectorInsertStridedSliceOpConvert,
+               VectorShuffleOpConvert, VectorSplatPattern>(
+      typeConverter, patterns.getContext());
 }

mlir/test/Conversion/VectorToSPIRV/vector-to-spirv.mlir

 // -----
 
 // CHECK-LABEL: @extract_size1_vector
 //  CHECK-SAME: %[[ARG0:.+]]: vector<1xf32>
 //       CHECK:   %[[R:.+]] = builtin.unrealized_conversion_cast %[[ARG0]]
 //       CHECK:   return %[[R]]
 func.func @extract_size1_vector(%arg0 : vector<1xf32>) -> f32 {
   %0 = vector.extract %arg0[0] : vector<1xf32>
-  return %0: f32
+  return %0: f32
 }
 
 // -----
 
 // CHECK-LABEL: @insert
 //  CHECK-SAME: %[[V:.*]]: vector<4xf32>, %[[S:.*]]: f32
 //       CHECK:   spv.CompositeInsert %[[S]], %[[V]][2 : i32] : f32 into vector<4xf32>
 func.func @insert(%arg0 : vector<4xf32>, %arg1: f32) -> vector<4xf32> {
   %1 = vector.insert %arg1, %arg0[2] : f32 into vector<4xf32>
-  return %1: vector<4xf32>
+  return %1: vector<4xf32>
 }
 
 // -----
 
 // CHECK-LABEL: @insert_size1_vector
 //  CHECK-SAME: %[[V:.*]]: vector<1xf32>, %[[S:.*]]: f32
 //       CHECK:   %[[R:.+]] = builtin.unrealized_conversion_cast %[[S]]
 //       CHECK:   return %[[R]]
 func.func @insert_size1_vector(%arg0 : vector<1xf32>, %arg1: f32) -> vector<1xf32> {
   %1 = vector.insert %arg1, %arg0[0] : f32 into vector<1xf32>
-  return %1 : vector<1xf32>
+  return %1 : vector<1xf32>
 }
 
 // -----
 
 // CHECK-LABEL: @extract_element
 //  CHECK-SAME: %[[V:.*]]: vector<4xf32>, %[[ID:.*]]: i32
 //       CHECK:   spv.VectorExtractDynamic %[[V]][%[[ID]]] : vector<4xf32>, i32
 func.func @extract_element(%arg0 : vector<4xf32>, %id : i32) -> f32 {
   %0 = vector.extractelement %arg0[%id : i32] : vector<4xf32>
   return %0: f32
 }
 
 // -----
 
 // CHECK-LABEL: @extract_element_index
 func.func @extract_element_index(%arg0 : vector<4xf32>, %id : index) -> f32 {
-  // CHECK: vector.extractelement
+  // CHECK: vector.extractelement
   %0 = vector.extractelement %arg0[%id : index] : vector<4xf32>
   return %0: f32
 }
 
 // -----
 
 // CHECK-LABEL: @extract_element_size5_vector
 func.func @extract_element_size5_vector(%arg0 : vector<5xf32>, %id : i32) -> f32 {
-  // CHECK: vector.extractelement
+  // CHECK: vector.extractelement
   %0 = vector.extractelement %arg0[%id : i32] : vector<5xf32>
   return %0: f32
 }
 
 // -----
 
 // CHECK-LABEL: @extract_strided_slice
 //  CHECK-SAME: %[[ARG:.+]]: vector<4xf32>
func.func @insert_element(%val: f32, %arg0 : vector<4xf32>, %id : i32) -> vector<4xf32> {
   %0 = vector.insertelement %val, %arg0[%id : i32] : vector<4xf32>
   return %0: vector<4xf32>
 }
 
 // -----
 
 // CHECK-LABEL: @insert_element_index
 func.func @insert_element_index(%val: f32, %arg0 : vector<4xf32>, %id : index) -> vector<4xf32> {
-  // CHECK: vector.insertelement
+  // CHECK: vector.insertelement
   %0 = vector.insertelement %val, %arg0[%id : index] : vector<4xf32>
   return %0: vector<4xf32>
 }
 
 // -----
 
 // CHECK-LABEL: @insert_element_size5_vector
 func.func @insert_element_size5_vector(%val: f32, %arg0 : vector<5xf32>, %id : i32) -> vector<5xf32> {
-  // CHECK: vector.insertelement
+  // CHECK: vector.insertelement
   %0 = vector.insertelement %val, %arg0[%id : i32] : vector<5xf32>
-  return %0 : vector<5xf32>
+  return %0 : vector<5xf32>
 }
 
 // -----
 
 // CHECK-LABEL: @insert_strided_slice
 //  CHECK-SAME: %[[PART:.+]]: vector<2xf32>, %[[ALL:.+]]: vector<4xf32>
-//       CHECK:   spv.VectorShuffle [0 : i32, 4 : i32, 5 : i32, 3 : i32] %[[ALL]] : vector<4xf32>, %[[PART]] : vector<2xf32> -> vector<4xf32>
+//       CHECK:   spv.VectorShuffle [0 : i32, 4 : i32, 5 : i32, 3 : i32] %[[ALL]] : vector<4xf32>, %[[PART]] : vector<2xf32> -> vector<4xf32>
 func.func @insert_strided_slice(%arg0: vector<2xf32>, %arg1: vector<4xf32>) -> vector<4xf32> {
   %0 = vector.insert_strided_slice %arg0, %arg1 {offsets = [1], strides = [1]} : vector<2xf32> into vector<4xf32>
-  return %0 : vector<4xf32>
+  return %0 : vector<4xf32>
 }
 
 // -----
 
 // CHECK-LABEL: @insert_size1_vector
 //  CHECK-SAME: %[[SUB:.*]]: vector<1xf32>, %[[FULL:.*]]: vector<3xf32>
 //       CHECK:   %[[S:.+]] = builtin.unrealized_conversion_cast %[[SUB]]
 //       CHECK:   spv.CompositeInsert %[[S]], %[[FULL]][2 : i32] : f32 into vector<3xf32>
 func.func @insert_size1_vector(%arg0 : vector<1xf32>, %arg1: vector<3xf32>) -> vector<3xf32> {
   %1 = vector.insert_strided_slice %arg0, %arg1 {offsets = [2], strides = [1]} : vector<1xf32> into vector<3xf32>
-  return %1 : vector<3xf32>
+  return %1 : vector<3xf32>
 }
 
 // -----
 
 // CHECK-LABEL: @fma
 //  CHECK-SAME: %[[A:.*]]: vector<4xf32>, %[[B:.*]]: vector<4xf32>, %[[C:.*]]: vector<4xf32>
 //       CHECK:   spv.GLSL.Fma %[[A]], %[[B]], %[[C]] : vector<4xf32>
 func.func @fma(%a: vector<4xf32>, %b: vector<4xf32>, %c: vector<4xf32>) -> vector<4xf32> {
   %0 = vector.fma %a, %b, %c: vector<4xf32>
-  return %0 : vector<4xf32>
+  return %0 : vector<4xf32>
 }
 
 // -----
 
 // CHECK-LABEL: func @splat
 //  CHECK-SAME: (%[[A:.+]]: f32)
 //       CHECK:   %[[VAL:.+]] = spv.CompositeConstruct %[[A]], %[[A]], %[[A]], %[[A]] : vector<4xf32>
 //       CHECK:   return %[[VAL]]

 // -----
 
 // CHECK-LABEL:  func @shuffle
 func.func @shuffle(%v0 : vector<2x16xf32>, %v1: vector<1x16xf32>) -> vector<3x16xf32> {
   // CHECK: vector.shuffle
   %shuffle = vector.shuffle %v0, %v1 [0, 1, 2] : vector<2x16xf32>, vector<1x16xf32>
   return %shuffle : vector<3x16xf32>
 }
+
+// -----
+
+// CHECK-LABEL: func @reduction
+//  CHECK-SAME: (%[[V:.+]]: vector<4xi32>)
+//       CHECK:   %[[S0:.+]] = spv.CompositeExtract %[[V]][0 : i32] : vector<4xi32>
+//       CHECK:   %[[S1:.+]] = spv.CompositeExtract %[[V]][1 : i32] : vector<4xi32>
+//       CHECK:   %[[S2:.+]] = spv.CompositeExtract %[[V]][2 : i32] : vector<4xi32>
+//       CHECK:   %[[S3:.+]] = spv.CompositeExtract %[[V]][3 : i32] : vector<4xi32>
+//       CHECK:   %[[ADD0:.+]] = spv.IAdd %[[S0]], %[[S1]]
+//       CHECK:   %[[ADD1:.+]] = spv.IAdd %[[ADD0]], %[[S2]]
+//       CHECK:   %[[ADD2:.+]] = spv.IAdd %[[ADD1]], %[[S3]]
+//       CHECK:   return %[[ADD2]]
+func.func @reduction(%v : vector<4xi32>) -> i32 {
+  %reduce = vector.reduction <add>, %v : vector<4xi32> into i32
+  return %reduce : i32
+}
+
+// -----
+
+// CHECK-LABEL: func @reduction
+//  CHECK-SAME: (%[[V:.+]]: vector<3xf32>, %[[S:.+]]: f32)
+//       CHECK:   %[[S0:.+]] = spv.CompositeExtract %[[V]][0 : i32] : vector<3xf32>
+//       CHECK:   %[[S1:.+]] = spv.CompositeExtract %[[V]][1 : i32] : vector<3xf32>
+//       CHECK:   %[[S2:.+]] = spv.CompositeExtract %[[V]][2 : i32] : vector<3xf32>
+//       CHECK:   %[[ADD0:.+]] = spv.FMul %[[S0]], %[[S1]]
+//       CHECK:   %[[ADD1:.+]] = spv.FMul %[[ADD0]], %[[S2]]
+//       CHECK:   %[[ADD2:.+]] = spv.FMul %[[ADD1]], %[[S]]
+//       CHECK:   return %[[ADD2]]
+func.func @reduction(%v : vector<3xf32>, %s: f32) -> f32 {
+  %reduce = vector.reduction <mul>, %v, %s : vector<3xf32> into f32
+  return %reduce : f32
+}