[mlir][vector] Add helper that builds a scalar reduction according to CombiningKind

mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h

 //===- VectorUtils.h - Vector Utilities -------------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_VECTOR_UTILS_VECTORUTILS_H_
 #define MLIR_DIALECT_VECTOR_UTILS_VECTORUTILS_H_
 
+#include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/Support/LLVM.h"
 
 #include "llvm/ADT/DenseMap.h"
 
 namespace mlir {
 
 // Forward declarations.
 class AffineApplyOp;
 class AffineForOp;
 class AffineMap;
 class Block;
 class Location;
 class OpBuilder;
 class Operation;
 class ShapedType;
 class Value;
 class VectorType;
 class VectorTransferOpInterface;
 
 namespace vector {
-class TransferWriteOp;
-class TransferReadOp;
-
 /// Helper function that creates a memref::DimOp or tensor::DimOp depending on
 /// the type of `source`.
 Value createOrFoldDimOp(OpBuilder &b, Location loc, Value source, int64_t dim);
+
+/// Return the result value of reducing two scalar/vector values with the
+/// corresponding arith operation.
+Value makeArithReduction(OpBuilder &b, Location loc, CombiningKind kind,
+                         Value v1, Value v2);
 } // namespace vector
 
 /// Return the number of elements of basis, `0` if empty.
 int64_t computeMaxLinearIndex(ArrayRef<int64_t> basis);
 
 /// Given the shape and sizes of a vector, returns the corresponding
 /// strides for each dimension.
 /// TODO: needs better doc of how it is used.

mlir/lib/Dialect/Vector/Transforms/VectorMultiDimReductionTransforms.cpp

if (!elementType.isIntOrIndexOrFloat())
       return failure();
 
     Value result =
         rewriter.create<vector::ExtractOp>(loc, multiReductionOp.source(), 0)
             .getResult();
     for (int64_t i = 1; i < srcShape[0]; i++) {
       auto operand =
           rewriter.create<vector::ExtractOp>(loc, multiReductionOp.source(), i);
-      switch (multiReductionOp.kind()) {
-      case vector::CombiningKind::ADD:
-        if (elementType.isIntOrIndex())
-          result = rewriter.create<arith::AddIOp>(loc, operand, result);
-        else
-          result = rewriter.create<arith::AddFOp>(loc, operand, result);
-        break;
-      case vector::CombiningKind::MUL:
-        if (elementType.isIntOrIndex())
-          result = rewriter.create<arith::MulIOp>(loc, operand, result);
-        else
-          result = rewriter.create<arith::MulFOp>(loc, operand, result);
-        break;
-      case vector::CombiningKind::MINUI:
-        result = rewriter.create<arith::MinUIOp>(loc, operand, result);
-        break;
-      case vector::CombiningKind::MINSI:
-        result = rewriter.create<arith::MinSIOp>(loc, operand, result);
-        break;
-      case vector::CombiningKind::MINF:
-        result = rewriter.create<arith::MinFOp>(loc, operand, result);
-        break;
-      case vector::CombiningKind::MAXUI:
-        result = rewriter.create<arith::MaxUIOp>(loc, operand, result);
-        break;
-      case vector::CombiningKind::MAXSI:
-        result = rewriter.create<arith::MaxSIOp>(loc, operand, result);
-        break;
-      case vector::CombiningKind::MAXF:
-        result = rewriter.create<arith::MaxFOp>(loc, operand, result);
-        break;
-      case vector::CombiningKind::AND:
-        result = rewriter.create<arith::AndIOp>(loc, operand, result);
-        break;
-      case vector::CombiningKind::OR:
-        result = rewriter.create<arith::OrIOp>(loc, operand, result);
-        break;
-      case vector::CombiningKind::XOR:
-        result = rewriter.create<arith::XOrIOp>(loc, operand, result);
-        break;
-      }
+      result = makeArithReduction(rewriter, loc, multiReductionOp.kind(),
+                                  operand, result);
     }
 
     rewriter.replaceOp(multiReductionOp, result);
     return success();
   }
 };
 
 /// Converts 2d vector.multi_reduction with inner most reduction dimension into

mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp

 //===- VectorTransforms.cpp - Conversion within the Vector dialect --------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements target-independent rewrites as 1->N patterns.
 //
 //===----------------------------------------------------------------------===//
 
+#include "mlir/Dialect/Vector/Transforms/VectorTransforms.h"
+
 #include <type_traits>
 
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/SCF.h"
 #include "mlir/Dialect/Utils/StructuredOpsUtils.h"
-
-#include "mlir/Dialect/Vector/Transforms/VectorTransforms.h"
+#include "mlir/Dialect/Vector/Utils/VectorUtils.h"
 #include "mlir/IR/ImplicitLocOpBuilder.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Interfaces/VectorInterfaces.h"
 
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/STLExtras.h"
static Optional<Value> genMultI(Location loc, Value x, Value y, Value acc,
                                   vector::CombiningKind kind,
                                   PatternRewriter &rewriter) {
     using vector::CombiningKind;
 
     auto mul = rewriter.create<arith::MulIOp>(loc, x, y);
     if (!acc)
       return Optional<Value>(mul);
 
-    Value combinedResult;
-    switch (kind) {
-    case CombiningKind::ADD:
-      combinedResult = rewriter.create<arith::AddIOp>(loc, mul, acc);
-      break;
-    case CombiningKind::MUL:
-      combinedResult = rewriter.create<arith::MulIOp>(loc, mul, acc);
-      break;
-    case CombiningKind::MINUI:
-      combinedResult = rewriter.create<arith::MinUIOp>(loc, mul, acc);
-      break;
-    case CombiningKind::MINSI:
-      combinedResult = rewriter.create<arith::MinSIOp>(loc, mul, acc);
-      break;
-    case CombiningKind::MAXUI:
-      combinedResult = rewriter.create<arith::MaxUIOp>(loc, mul, acc);
-      break;
-    case CombiningKind::MAXSI:
-      combinedResult = rewriter.create<arith::MaxSIOp>(loc, mul, acc);
-      break;
-    case CombiningKind::AND:
-      combinedResult = rewriter.create<arith::AndIOp>(loc, mul, acc);
-      break;
-    case CombiningKind::OR:
-      combinedResult = rewriter.create<arith::OrIOp>(loc, mul, acc);
-      break;
-    case CombiningKind::XOR:
-      combinedResult = rewriter.create<arith::XOrIOp>(loc, mul, acc);
-      break;
-    case CombiningKind::MINF: // Only valid for floating point types.
-    case CombiningKind::MAXF: // Only valid for floating point types.
+    if (kind == CombiningKind::MINF || kind == CombiningKind::MAXF)
+      // Only valid for floating point types.
       return Optional<Value>();
-    }
-    return Optional<Value>(combinedResult);
+
+    return makeArithReduction(rewriter, loc, kind, mul, acc);
   }
 
   static Optional<Value> genMultF(Location loc, Value x, Value y, Value acc,
                                   vector::CombiningKind kind,
                                   PatternRewriter &rewriter) {
     using vector::CombiningKind;
 
     // Special case for fused multiply-add.
     if (acc && kind == CombiningKind::ADD) {
       return Optional<Value>(rewriter.create<vector::FMAOp>(loc, x, y, acc));
     }
 
     auto mul = rewriter.create<arith::MulFOp>(loc, x, y);
 
     if (!acc)
       return Optional<Value>(mul);
 
-    Value combinedResult;
-    switch (kind) {
-    case CombiningKind::MUL:
-      combinedResult = rewriter.create<arith::MulFOp>(loc, mul, acc);
-      break;
-    case CombiningKind::MINF:
-      combinedResult = rewriter.create<arith::MinFOp>(loc, mul, acc);
-      break;
-    case CombiningKind::MAXF:
-      combinedResult = rewriter.create<arith::MaxFOp>(loc, mul, acc);
-      break;
-    case CombiningKind::ADD:   // Already handled this special case above.
-    case CombiningKind::AND:   // Only valid for integer types.
-    case CombiningKind::MINUI: // Only valid for integer types.
-    case CombiningKind::MINSI: // Only valid for integer types.
-    case CombiningKind::MAXUI: // Only valid for integer types.
-    case CombiningKind::MAXSI: // Only valid for integer types.
-    case CombiningKind::OR:    // Only valid for integer types.
-    case CombiningKind::XOR:   // Only valid for integer types.
+    if (kind == CombiningKind::ADD || kind == CombiningKind::AND ||
+        kind == CombiningKind::MINUI || kind == CombiningKind::MINSI ||
+        kind == CombiningKind::MAXUI || kind == CombiningKind::MAXSI ||
+        kind == CombiningKind::OR || kind == CombiningKind::XOR)
+      // Already handled or only valid for integer types.
       return Optional<Value>();
-    }
-    return Optional<Value>(combinedResult);
+
+    return makeArithReduction(rewriter, loc, kind, mul, acc);
   }
 };
 
 /// Progressive lowering of ConstantMaskOp.
 /// One:
 ///   %x = vector.constant_mask [a,b]
 /// is replaced by:
 ///   %z = zero-result

mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp

 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/IR/Operation.h"
+#include "mlir/IR/TypeUtilities.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Support/MathExtras.h"
 #include <numeric>
 
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/SetVector.h"
 
 using namespace mlir;
 
 /// Helper function that creates a memref::DimOp or tensor::DimOp depending on
 /// the type of `source`.
 Value mlir::vector::createOrFoldDimOp(OpBuilder &b, Location loc, Value source,
                                       int64_t dim) {
   if (source.getType().isa<UnrankedMemRefType, MemRefType>())
     return b.createOrFold<memref::DimOp>(loc, source, dim);
   if (source.getType().isa<UnrankedTensorType, RankedTensorType>())
     return b.createOrFold<tensor::DimOp>(loc, source, dim);
   llvm_unreachable("Expected MemRefType or TensorType");
 }
 
+Value mlir::vector::makeArithReduction(OpBuilder &b, Location loc,
+                                       CombiningKind kind, Value v1, Value v2) {
+  Type t1 = getElementTypeOrSelf(v1.getType());
+  Type t2 = getElementTypeOrSelf(v2.getType());
+  switch (kind) {
+  case CombiningKind::ADD:
+    if (t1.isIntOrIndex() && t2.isIntOrIndex())
+      return b.createOrFold<arith::AddIOp>(loc, v1, v2);
+    else if (t1.isa<FloatType>() && t2.isa<FloatType>())
+      return b.createOrFold<arith::AddFOp>(loc, v1, v2);
+    llvm_unreachable("invalid value types for ADD reduction");
+  case CombiningKind::AND:
+    assert(t1.isIntOrIndex() && t2.isIntOrIndex() && "expected int values");
+    return b.createOrFold<arith::AndIOp>(loc, v1, v2);
+  case CombiningKind::MAXF:
+    assert(t1.isa<FloatType>() && t2.isa<FloatType>() &&
+           "expected float values");
+    return b.createOrFold<arith::MaxFOp>(loc, v1, v2);
+  case CombiningKind::MINF:
+    assert(t1.isa<FloatType>() && t2.isa<FloatType>() &&
+           "expected float values");
+    return b.createOrFold<arith::MinFOp>(loc, v1, v2);
+  case CombiningKind::MAXSI:
+    assert(t1.isIntOrIndex() && t2.isIntOrIndex() && "expected int values");
+    return b.createOrFold<arith::MaxSIOp>(loc, v1, v2);
+  case CombiningKind::MINSI:
+    assert(t1.isIntOrIndex() && t2.isIntOrIndex() && "expected int values");
+    return b.createOrFold<arith::MinSIOp>(loc, v1, v2);
+  case CombiningKind::MAXUI:
+    assert(t1.isIntOrIndex() && t2.isIntOrIndex() && "expected int values");
+    return b.createOrFold<arith::MaxUIOp>(loc, v1, v2);
+  case CombiningKind::MINUI:
+    assert(t1.isIntOrIndex() && t2.isIntOrIndex() && "expected int values");
+    return b.createOrFold<arith::MinUIOp>(loc, v1, v2);
+  case CombiningKind::MUL:
+    if (t1.isIntOrIndex() && t2.isIntOrIndex())
+      return b.createOrFold<arith::MulIOp>(loc, v1, v2);
+    else if (t1.isa<FloatType>() && t2.isa<FloatType>())
+      return b.createOrFold<arith::MulFOp>(loc, v1, v2);
+    llvm_unreachable("invalid value types for MUL reduction");
+  case CombiningKind::OR:
+    assert(t1.isIntOrIndex() && t2.isIntOrIndex() && "expected int values");
+    return b.createOrFold<arith::OrIOp>(loc, v1, v2);
+  case CombiningKind::XOR:
+    assert(t1.isIntOrIndex() && t2.isIntOrIndex() && "expected int values");
+    return b.createOrFold<arith::XOrIOp>(loc, v1, v2);
+  };
+  llvm_unreachable("unknown CombiningKind");
+}
+
 /// Return the number of elements of basis, `0` if empty.
 int64_t mlir::computeMaxLinearIndex(ArrayRef<int64_t> basis) {
   if (basis.empty())
     return 0;
   return std::accumulate(basis.begin(), basis.end(), 1,
                          std::multiplies<int64_t>());
 }