[mlir] Allow polynomial approximations for N-d vectors.

[pifon2a]

Summary

Polynomial approximation can be extented to support N-d vectors.
N-dimensional vectors are useful when vectorizing operations on N-dimensional
tiles. Before lowering to LLVM these vectors are usually unrolled or flattened
to 1-dimensional vectors.

[bkramer]

looks good

[asaadaldien]

This diff broke the build, should be fixed in https://reviews.llvm.org/D112570

mlir/lib/Dialect/Math/Transforms/PolynomialApproximation.cpp

 #include <cstddef>
 
 using namespace mlir;
 using namespace mlir::math;
 using namespace mlir::vector;
 
 using TypePredicate = llvm::function_ref<bool(Type)>;
 
-// Returns vector width if the element type is matching the predicate (scalars
-// that do match the predicate have width equal to `1`).
-static Optional<int> vectorWidth(Type type, TypePredicate pred) {
-  // If the type matches the predicate then its width is `1`.
+// Returns vector shape if the element type is matching the predicate (scalars
+// that do match the predicate have shape equal to `{1}`).
+static Optional<SmallVector<int64_t, 2>> vectorShape(Type type,
+                                                     TypePredicate pred) {
+  // If the type matches the predicate then its shape is `{1}`.
   if (pred(type))
-    return 1;
+    return SmallVector<int64_t, 2>{1};
 
   // Otherwise check if the type is a vector type.
   auto vectorType = type.dyn_cast<VectorType>();
   if (vectorType && pred(vectorType.getElementType())) {
-    assert(vectorType.getRank() == 1 && "only 1d vectors are supported");
-    return vectorType.getDimSize(0);
+    return llvm::to_vector<2>(vectorType.getShape());
   }
 
   return llvm::None;
 }
 
-// Returns vector width of the type. If the type is a scalar returns `1`.
-static int vectorWidth(Type type) {
+// Returns vector shape of the type. If the type is a scalar returns `1`.
+static SmallVector<int64_t, 2> vectorShape(Type type) {
   auto vectorType = type.dyn_cast<VectorType>();
-  return vectorType ? vectorType.getDimSize(0) : 1;
+  return vectorType ? llvm::to_vector<2>(vectorType.getShape())
+                    : SmallVector<int64_t, 2>{1};
 }
 
 // Returns vector element type. If the type is a scalar returns the argument.
 LLVM_ATTRIBUTE_UNUSED static Type elementType(Type type) {
   auto vectorType = type.dyn_cast<VectorType>();
   return vectorType ? vectorType.getElementType() : type;
 }
 
 LLVM_ATTRIBUTE_UNUSED static bool isF32(Type type) { return type.isF32(); }
 
 LLVM_ATTRIBUTE_UNUSED static bool isI32(Type type) {
   return type.isInteger(32);
 }
 
 //----------------------------------------------------------------------------//
 // Broadcast scalar types and values into vector types and values.
 //----------------------------------------------------------------------------//
 
-// Broadcasts scalar type into vector type (iff width is greater then 1).
-static Type broadcast(Type type, int width) {
+// Returns true if shape != {1}.
+static bool isNonScalarShape(ArrayRef<int64_t> shape) {
+  return shape.size() > 1 || shape[0] > 1;
+}
+
+// Broadcasts scalar type into vector type (iff shape is non-scalar).
+static Type broadcast(Type type, ArrayRef<int64_t> shape) {
   assert(!type.isa<VectorType>() && "must be scalar type");
-  return width > 1 ? VectorType::get({width}, type) : type;
+  return isNonScalarShape(shape) ? VectorType::get(shape, type) : type;
 }
 
-// Broadcasts scalar value into vector (iff width is greater then 1).
-static Value broadcast(ImplicitLocOpBuilder &builder, Value value, int width) {
+// Broadcasts scalar value into vector (iff shape is non-scalar).
+static Value broadcast(ImplicitLocOpBuilder &builder, Value value,
+                       ArrayRef<int64_t> shape) {
   assert(!value.getType().isa<VectorType>() && "must be scalar value");
-  auto type = broadcast(value.getType(), width);
-  return width > 1 ? builder.create<BroadcastOp>(type, value) : value;
+  auto type = broadcast(value.getType(), shape);
+  return isNonScalarShape(shape) ? builder.create<BroadcastOp>(type, value)
+                                 : value;
 }
 
 //----------------------------------------------------------------------------//
 // Helper functions to create constants.
 //----------------------------------------------------------------------------//
 
 static Value f32Cst(ImplicitLocOpBuilder &builder, float value) {
   return builder.create<arith::ConstantOp>(builder.getF32FloatAttr(value));

 }
 
 // Decomposes given floating point value `arg` into a normalized fraction and
 // an integral power of two (see std::frexp). Returned values have float type.
 static std::pair<Value, Value> frexp(ImplicitLocOpBuilder &builder, Value arg,
                                      bool is_positive = false) {
   assert(isF32(elementType(arg.getType())) && "argument must be f32 type");
 
-  int width = vectorWidth(arg.getType());
+  auto shape = vectorShape(arg.getType());
 
   auto bcast = [&](Value value) -> Value {
-    return broadcast(builder, value, width);
+    return broadcast(builder, value, shape);
   };
 
   auto i32 = builder.getIntegerType(32);
-  auto i32Vec = broadcast(i32, width);
-  auto f32Vec = broadcast(builder.getF32Type(), width);
+  auto i32Vec = broadcast(i32, shape);
+  auto f32Vec = broadcast(builder.getF32Type(), shape);
 
   Value cst126f = f32Cst(builder, 126.0f);
   Value cstHalf = f32Cst(builder, 0.5f);
   Value cstInvMantMask = f32FromBits(builder, ~0x7f800000u);
 
   // Bitcast to i32 for bitwise operations.
   Value i32Half = builder.create<arith::BitcastOp>(i32, cstHalf);
   Value i32InvMantMask = builder.create<arith::BitcastOp>(i32, cstInvMantMask);
static std::pair<Value, Value> frexp(ImplicitLocOpBuilder &builder, Value arg,
 
   return {normalizedFraction, exponent};
 }
 
 // Computes exp2 for an i32 argument.
 static Value exp2I32(ImplicitLocOpBuilder &builder, Value arg) {
   assert(isI32(elementType(arg.getType())) && "argument must be i32 type");
 
-  int width = vectorWidth(arg.getType());
+  auto shape = vectorShape(arg.getType());
 
   auto bcast = [&](Value value) -> Value {
-    return broadcast(builder, value, width);
+    return broadcast(builder, value, shape);
   };
 
-  auto f32Vec = broadcast(builder.getF32Type(), width);
+  auto f32Vec = broadcast(builder.getF32Type(), shape);
   // The exponent of f32 located at 23-bit.
   auto exponetBitLocation = bcast(i32Cst(builder, 23));
   // Set the exponent bias to zero.
   auto bias = bcast(i32Cst(builder, 127));
 
   Value biasedArg = builder.create<arith::AddIOp>(arg, bias);
   Value exp2ValueInt =
       builder.create<arith::ShLIOp>(biasedArg, exponetBitLocation);
public:
   LogicalResult matchAndRewrite(math::TanhOp op,
                                 PatternRewriter &rewriter) const final;
 };
 } // namespace
 
 LogicalResult
 TanhApproximation::matchAndRewrite(math::TanhOp op,
                                    PatternRewriter &rewriter) const {
-  auto width = vectorWidth(op.operand().getType(), isF32);
-  if (!width.hasValue())
+  auto shape = vectorShape(op.operand().getType(), isF32);
+  if (!shape.hasValue())
     return rewriter.notifyMatchFailure(op, "unsupported operand type");
 
   ImplicitLocOpBuilder builder(op->getLoc(), rewriter);
   auto bcast = [&](Value value) -> Value {
-    return broadcast(builder, value, *width);
+    return broadcast(builder, value, *shape);
   };
 
   // Clamp operand into [plusClamp, minusClamp] range.
   Value minusClamp = bcast(f32Cst(builder, -7.99881172180175781f));
   Value plusClamp = bcast(f32Cst(builder, 7.99881172180175781f));
   Value x = clamp(builder, op.operand(), minusClamp, plusClamp);
 
   // Mask for tiny values that are approximated with `operand`.

 } // namespace
 
 // This approximation comes from Julien Pommier's SSE math library.
 // Link: http://gruntthepeon.free.fr/ssemath
 template <typename Op>
 LogicalResult
 LogApproximationBase<Op>::logMatchAndRewrite(Op op, PatternRewriter &rewriter,
                                              bool base2) const {
-  auto width = vectorWidth(op.operand().getType(), isF32);
-  if (!width.hasValue())
+  auto shape = vectorShape(op.operand().getType(), isF32);
+  if (!shape.hasValue())
     return rewriter.notifyMatchFailure(op, "unsupported operand type");
 
   ImplicitLocOpBuilder builder(op->getLoc(), rewriter);
   auto bcast = [&](Value value) -> Value {
-    return broadcast(builder, value, *width);
+    return broadcast(builder, value, *shape);
   };
 
   Value cstZero = bcast(f32Cst(builder, 0.0f));
   Value cstOne = bcast(f32Cst(builder, 1.0f));
   Value cstNegHalf = bcast(f32Cst(builder, -0.5f));
 
   // The smallest non denormalized float number.
   Value cstMinNormPos = bcast(f32FromBits(builder, 0x00800000u));
LogicalResult matchAndRewrite(math::Log1pOp op,
                                 PatternRewriter &rewriter) const final;
 };
 } // namespace
 
 // Approximate log(1+x).
 LogicalResult
 Log1pApproximation::matchAndRewrite(math::Log1pOp op,
                                     PatternRewriter &rewriter) const {
-  auto width = vectorWidth(op.operand().getType(), isF32);
-  if (!width.hasValue())
+  auto shape = vectorShape(op.operand().getType(), isF32);
+  if (!shape.hasValue())
     return rewriter.notifyMatchFailure(op, "unsupported operand type");
 
   ImplicitLocOpBuilder builder(op->getLoc(), rewriter);
   auto bcast = [&](Value value) -> Value {
-    return broadcast(builder, value, *width);
+    return broadcast(builder, value, *shape);
   };
 
   // Approximate log(1+x) using the following, due to W. Kahan:
   //   u = x + 1.0;
   //   if (u == 1.0 || u == inf) return x;
   //   return x * log(u) / (u - 1.0);
   //          ^^^^^^^^^^^^^^^^^^^^^^
   //             "logLarge" below.

 } // namespace
 
 // Approximate exp(x) using its reduced range exp(y) where y is in the range
 // [0, ln(2)], let y = x - floor(x / ln(2)) * ln(2) = x - k * ln(2), exp(x)
 // = exp(y) * 2^k. exp(y).
 LogicalResult
 ExpApproximation::matchAndRewrite(math::ExpOp op,
                                   PatternRewriter &rewriter) const {
-  auto width = vectorWidth(op.operand().getType(), isF32);
-  if (!width.hasValue())
+  auto shape = vectorShape(op.operand().getType(), isF32);
+  if (!shape.hasValue())
     return rewriter.notifyMatchFailure(op, "unsupported operand type");
   ImplicitLocOpBuilder builder(op->getLoc(), rewriter);
 
   // TODO: Consider a common pattern rewriter with all methods below to
   // write the approximations.
   auto bcast = [&](Value value) -> Value {
-    return broadcast(builder, value, *width);
+    return broadcast(builder, value, *shape);
   };
   auto fmla = [&](Value a, Value b, Value c) {
     return builder.create<math::FmaOp>(a, b, c);
   };
   auto mul = [&](Value a, Value b) -> Value {
     return builder.create<arith::MulFOp>(a, b);
   };
   auto sub = [&](Value a, Value b) -> Value {
ExpApproximation::matchAndRewrite(math::ExpOp op,
   Value y4 = mul(y2, y2);
 
   Value q0 = fmla(cstCephesExpP1, y, cstCephesExpP0);
   Value q1 = fmla(cstCephesExpP3, y, cstCephesExpP2);
   Value q2 = fmla(cstCephesExpP5, y, cstCephesExpP4);
   Value expY = fmla(q1, y2, q0);
   expY = fmla(q2, y4, expY);
 
-  auto i32Vec = broadcast(builder.getI32Type(), *width);
+  auto i32Vec = broadcast(builder.getI32Type(), *shape);
 
   // exp2(k)
   Value k = builder.create<arith::FPToSIOp>(kF32, i32Vec);
   Value exp2KValue = exp2I32(builder, k);
 
   // exp(x) = exp(y) * exp2(k)
   expY = mul(expY, exp2KValue);
 
public:
   LogicalResult matchAndRewrite(math::ExpM1Op op,
                                 PatternRewriter &rewriter) const final;
 };
 } // namespace
 
 LogicalResult
 ExpM1Approximation::matchAndRewrite(math::ExpM1Op op,
                                     PatternRewriter &rewriter) const {
-  auto width = vectorWidth(op.operand().getType(), isF32);
-  if (!width.hasValue())
+  auto shape = vectorShape(op.operand().getType(), isF32);
+  if (!shape.hasValue())
     return rewriter.notifyMatchFailure(op, "unsupported operand type");
 
   ImplicitLocOpBuilder builder(op->getLoc(), rewriter);
   auto bcast = [&](Value value) -> Value {
-    return broadcast(builder, value, *width);
+    return broadcast(builder, value, *shape);
   };
 
   // expm1(x) = exp(x) - 1 = u - 1.
   // We have to handle it carefully when x is near 0, i.e. u ~= 1,
   // and when the input is ~= -inf, i.e. u - 1 ~= -1.
   Value cstOne = bcast(f32Cst(builder, 1.0f));
   Value cstNegOne = bcast(f32Cst(builder, -1.0f));
   Value x = op.operand();

 // the reduced range [0, pi/2] for both sin(x) and cos(x). Then given y in the
 // reduced range sin(x) will be computed as sin(y), -sin(y), cos(y) or -cos(y).
 template <bool isSine, typename OpTy>
 LogicalResult SinAndCosApproximation<isSine, OpTy>::matchAndRewrite(
     OpTy op, PatternRewriter &rewriter) const {
   static_assert(
       llvm::is_one_of<OpTy, math::SinOp, math::CosOp>::value,
       "SinAndCosApproximation pattern expects math::SinOp or math::CosOp");
-  auto width = vectorWidth(op.operand().getType(), isF32);
-  if (!width.hasValue())
+  auto shape = vectorShape(op.operand().getType(), isF32);
+  if (!shape.hasValue())
     return rewriter.notifyMatchFailure(op, "unsupported operand type");
 
   ImplicitLocOpBuilder builder(op->getLoc(), rewriter);
   auto bcast = [&](Value value) -> Value {
-    return broadcast(builder, value, *width);
+    return broadcast(builder, value, *shape);
   };
   auto mul = [&](Value a, Value b) -> Value {
     return builder.create<arith::MulFOp>(a, b);
   };
   auto sub = [&](Value a, Value b) -> Value {
     return builder.create<arith::SubFOp>(a, b);
   };
   auto floor = [&](Value a) { return builder.create<math::FloorOp>(a); };
 
-  auto i32Vec = broadcast(builder.getI32Type(), *width);
+  auto i32Vec = broadcast(builder.getI32Type(), *shape);
   auto fPToSingedInteger = [&](Value a) -> Value {
     return builder.create<arith::FPToSIOp>(a, i32Vec);
   };
 
   auto modulo4 = [&](Value a) -> Value {
     return builder.create<arith::AndIOp>(a, bcast(i32Cst(builder, 3)));
   };
 
struct RsqrtApproximation : public OpRewritePattern<math::RsqrtOp> {
   LogicalResult matchAndRewrite(math::RsqrtOp op,
                                 PatternRewriter &rewriter) const final;
 };
 } // namespace
 
 LogicalResult
 RsqrtApproximation::matchAndRewrite(math::RsqrtOp op,
                                     PatternRewriter &rewriter) const {
-  auto width = vectorWidth(op.operand().getType(), isF32);
+  auto shape = vectorShape(op.operand().getType(), isF32);
   // Only support already-vectorized rsqrt's.
-  if (!width.hasValue() || *width != 8)
+  if (!shape.hasValue() || (*shape)[0] != 8)
     return rewriter.notifyMatchFailure(op, "unsupported operand type");
 
   ImplicitLocOpBuilder builder(op->getLoc(), rewriter);
   auto bcast = [&](Value value) -> Value {
-    return broadcast(builder, value, *width);
+    return broadcast(builder, value, *shape);
   };
 
   Value cstPosInf = bcast(f32FromBits(builder, 0x7f800000u));
   Value cstOnePointFive = bcast(f32Cst(builder, 1.5f));
   Value cstNegHalf = bcast(f32Cst(builder, -0.5f));
   Value cstMinNormPos = bcast(f32FromBits(builder, 0x00800000u));
 
   Value negHalf = builder.create<arith::MulFOp>(op.operand(), cstNegHalf);

mlir/test/Dialect/Math/polynomial-approximation.mlir

 // CHECK:           return %[[VAL_109]] : f32
 // CHECK:         }
 func @expm1_scalar(%arg0: f32) -> f32 {
   %0 = math.expm1 %arg0 : f32
   return %0 : f32
 }
 
 // CHECK-LABEL:   func @expm1_vector(
-// CHECK-SAME:                       %[[VAL_0:.*]]: vector<8xf32>) -> vector<8xf32> {
-// CHECK:           %[[VAL_1:.*]] = arith.constant dense<-1.000000e+00> : vector<8xf32>
+// CHECK-SAME:                       %[[VAL_0:.*]]: vector<8x8xf32>) -> vector<8x8xf32> {
+// CHECK:           %[[VAL_1:.*]] = arith.constant dense<-1.000000e+00> : vector<8x8xf32>
 // CHECK-NOT:       exp
 // CHECK-COUNT-4:   select
 // CHECK-NOT:       log
 // CHECK-COUNT-5:   select
 // CHECK-NOT:       expm1
 // CHECK-COUNT-3:   select
 // CHECK:           %[[VAL_115:.*]] = select
-// CHECK:           return %[[VAL_115]] : vector<8xf32>
+// CHECK:           return %[[VAL_115]] : vector<8x8xf32>
 // CHECK:         }
-func @expm1_vector(%arg0: vector<8xf32>) -> vector<8xf32> {
-  %0 = math.expm1 %arg0 : vector<8xf32>
-  return %0 : vector<8xf32>
+func @expm1_vector(%arg0: vector<8x8xf32>) -> vector<8x8xf32> {
+  %0 = math.expm1 %arg0 : vector<8x8xf32>
+  return %0 : vector<8x8xf32>
 }
 
 // CHECK-LABEL:   func @log_scalar(
 // CHECK-SAME:                             %[[X:.*]]: f32) -> f32 {
 // CHECK:           %[[VAL_1:.*]] = arith.constant 0.000000e+00 : f32
 // CHECK:           %[[VAL_2:.*]] = arith.constant 1.000000e+00 : f32
 // CHECK:           %[[VAL_3:.*]] = arith.constant -5.000000e-01 : f32
 // CHECK:           %[[VAL_4:.*]] = arith.constant 1.17549435E-38 : f32