[mlir] MathApproximations: scalars shape must be 0-rank

[ezhulenev]

Summary

Using [1] for representing shape of a scalar is incorrect, and will break with vectors of size 1.

• remove redundant helper functions
• fix couple of style warnings

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

 //===- PolynomialApproximation.cpp - Approximate math operations ----------===//
 //
 // 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 expansion of math operations to fast approximations
 // that do not rely on any of the library functions.
 //
 //===----------------------------------------------------------------------===//
 
+#include <climits>
+#include <cstddef>
+
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Dialect/Math/IR/Math.h"
 #include "mlir/Dialect/Math/Transforms/Approximation.h"
 #include "mlir/Dialect/Math/Transforms/Passes.h"
 #include "mlir/Dialect/Vector/VectorOps.h"
 #include "mlir/Dialect/Vector/VectorUtils.h"
 #include "mlir/Dialect/X86Vector/X86VectorDialect.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/ImplicitLocOpBuilder.h"
+#include "mlir/IR/TypeUtilities.h"
 #include "mlir/Transforms/Bufferize.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "llvm/ADT/ArrayRef.h"
-#include <climits>
-#include <cstddef>
 
 using namespace mlir;
 using namespace mlir::math;
 using namespace mlir::vector;
 
-using TypePredicate = llvm::function_ref<bool(Type)>;
-
-// 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 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())) {
-    return llvm::to_vector<2>(vectorType.getShape());
-  }
-
-  return llvm::None;
-}
-
-// 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 ? 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) {
+// Returns vector shape if the type is a vector. Returns an empty shape if it is
+// not a vector.
+static ArrayRef<int64_t> vectorShape(Type type) {
   auto vectorType = type.dyn_cast<VectorType>();
-  return vectorType ? vectorType.getElementType() : type;
+  return vectorType ? vectorType.getShape() : ArrayRef<int64_t>();
 }
 
-LLVM_ATTRIBUTE_UNUSED static bool isF32(Type type) { return type.isF32(); }
-
-LLVM_ATTRIBUTE_UNUSED static bool isI32(Type type) {
-  return type.isInteger(32);
+static ArrayRef<int64_t> vectorShape(Value value) {
+  return vectorShape(value.getType());
 }
 
 //----------------------------------------------------------------------------//
 // Broadcast scalar types and values into vector types and values.
 //----------------------------------------------------------------------------//
 
-// 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 isNonScalarShape(shape) ? VectorType::get(shape, type) : type;
+  return !shape.empty() ? VectorType::get(shape, type) : type;
 }
 
 // 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(), shape);
-  return isNonScalarShape(shape) ? builder.create<BroadcastOp>(type, value)
-                                 : value;
+  return !shape.empty() ? builder.create<BroadcastOp>(type, value) : value;
 }
 
 //----------------------------------------------------------------------------//
 // Helper function to handle n-D vectors with 1-D operations.
 //----------------------------------------------------------------------------//
 
 // Expands and unrolls n-D vector operands into multiple fixed size 1-D vectors
 // and calls the compute function with 1-D vector operands. Stitches back all
static Value clamp(ImplicitLocOpBuilder &builder, Value value, Value lowerBound,
                    Value upperBound) {
   return max(builder, min(builder, value, upperBound), lowerBound);
 }
 
 // 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");
-
-  auto shape = vectorShape(arg.getType());
+  assert(getElementTypeOrSelf(arg).isF32() && "arg must be f32 type");
+  ArrayRef<int64_t> shape = vectorShape(arg);
 
   auto bcast = [&](Value value) -> Value {
     return broadcast(builder, value, shape);
   };
 
   auto i32 = builder.getIntegerType(32);
   auto i32Vec = broadcast(i32, shape);
   auto f32Vec = broadcast(builder.getF32Type(), shape);
static std::pair<Value, Value> frexp(ImplicitLocOpBuilder &builder, Value arg,
   Value exponent =
       builder.create<arith::SubFOp>(biasedExponent, bcast(cst126f));
 
   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");
-
-  auto shape = vectorShape(arg.getType());
+  assert(getElementTypeOrSelf(arg).isInteger(32) && "arg must be i32 type");
+  ArrayRef<int64_t> shape = vectorShape(arg);
 
   auto bcast = [&](Value value) -> Value {
     return broadcast(builder, value, shape);
   };
 
   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);
   Value exp2ValueF32 = builder.create<arith::BitcastOp>(f32Vec, exp2ValueInt);
 
   return exp2ValueF32;
 }
 
 namespace {
 Value makePolynomialCalculation(ImplicitLocOpBuilder &builder,
                                 llvm::ArrayRef<Value> coeffs, Value x) {
-  auto shape = vectorShape(x.getType(), isF32);
-  if (coeffs.size() == 0) {
-    return broadcast(builder, f32Cst(builder, 0.0f), *shape);
-  } else if (coeffs.size() == 1) {
+  assert(getElementTypeOrSelf(x).isF32() && "x must be f32 type");
+  ArrayRef<int64_t> shape = vectorShape(x);
+
+  if (coeffs.empty())
+    return broadcast(builder, f32Cst(builder, 0.0f), shape);
+
+  if (coeffs.size() == 1)
     return coeffs[0];
-  }
+
   Value res = builder.create<math::FmaOp>(x, coeffs[coeffs.size() - 1],
                                           coeffs[coeffs.size() - 2]);
   for (auto i = ptrdiff_t(coeffs.size()) - 3; i >= 0; --i) {
     res = builder.create<math::FmaOp>(x, res, coeffs[i]);
   }
   return res;
 }
 } // namespace
public:
   LogicalResult matchAndRewrite(math::TanhOp op,
                                 PatternRewriter &rewriter) const final;
 };
 } // namespace
 
 LogicalResult
 TanhApproximation::matchAndRewrite(math::TanhOp op,
                                    PatternRewriter &rewriter) const {
-  auto shape = vectorShape(op.operand().getType(), isF32);
-  if (!shape.hasValue())
+  if (!getElementTypeOrSelf(op.operand()).isF32())
     return rewriter.notifyMatchFailure(op, "unsupported operand type");
 
+  ArrayRef<int64_t> shape = vectorShape(op.operand());
+
   ImplicitLocOpBuilder builder(op->getLoc(), rewriter);
   auto bcast = [&](Value value) -> Value {
-    return broadcast(builder, value, *shape);
+    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 shape = vectorShape(op.operand().getType(), isF32);
-  if (!shape.hasValue())
+  if (!getElementTypeOrSelf(op.operand()).isF32())
     return rewriter.notifyMatchFailure(op, "unsupported operand type");
 
+  ArrayRef<int64_t> shape = vectorShape(op.operand());
+
   ImplicitLocOpBuilder builder(op->getLoc(), rewriter);
   auto bcast = [&](Value value) -> Value {
-    return broadcast(builder, value, *shape);
+    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 shape = vectorShape(op.operand().getType(), isF32);
-  if (!shape.hasValue())
+  if (!getElementTypeOrSelf(op.operand()).isF32())
     return rewriter.notifyMatchFailure(op, "unsupported operand type");
 
+  ArrayRef<int64_t> shape = vectorShape(op.operand());
+
   ImplicitLocOpBuilder builder(op->getLoc(), rewriter);
   auto bcast = [&](Value value) -> Value {
-    return broadcast(builder, value, *shape);
+    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.

 // where P and Q are polynomials of degree 4.
 // Different coefficients are chosen based on the value of x.
 // The approximation error is ~2.5e-07.
 // Boost's minimax tool that utilizes the Remez method was used to find the
 // coefficients.
 LogicalResult
 ErfPolynomialApproximation::matchAndRewrite(math::ErfOp op,
                                             PatternRewriter &rewriter) const {
-  auto shape = vectorShape(op.operand().getType(), isF32);
-  if (!shape.hasValue())
+  if (!getElementTypeOrSelf(op.operand()).isF32())
     return rewriter.notifyMatchFailure(op, "unsupported operand type");
 
+  ArrayRef<int64_t> shape = vectorShape(op.operand());
+
   ImplicitLocOpBuilder builder(op->getLoc(), rewriter);
   auto bcast = [&](Value value) -> Value {
-    return broadcast(builder, value, *shape);
+    return broadcast(builder, value, shape);
   };
 
   const int intervalsCount = 3;
   const int polyDegree = 4;
 
   Value zero = bcast(f32Cst(builder, 0));
   Value one = bcast(f32Cst(builder, 1));
   Value pp[intervalsCount][polyDegree + 1];

 } // 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 shape = vectorShape(op.operand().getType(), isF32);
-  if (!shape.hasValue())
+  if (!getElementTypeOrSelf(op.operand()).isF32())
     return rewriter.notifyMatchFailure(op, "unsupported operand type");
+
+  ArrayRef<int64_t> shape = vectorShape(op.operand());
+
   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, *shape);
+    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(), *shape);
+  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 shape = vectorShape(op.operand().getType(), isF32);
-  if (!shape.hasValue())
+  if (!getElementTypeOrSelf(op.operand()).isF32())
     return rewriter.notifyMatchFailure(op, "unsupported operand type");
 
+  ArrayRef<int64_t> shape = vectorShape(op.operand());
+
   ImplicitLocOpBuilder builder(op->getLoc(), rewriter);
   auto bcast = [&](Value value) -> Value {
-    return broadcast(builder, value, *shape);
+    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 shape = vectorShape(op.operand().getType(), isF32);
-  if (!shape.hasValue())
+
+  if (!getElementTypeOrSelf(op.operand()).isF32())
     return rewriter.notifyMatchFailure(op, "unsupported operand type");
 
+  ArrayRef<int64_t> shape = vectorShape(op.operand());
+
   ImplicitLocOpBuilder builder(op->getLoc(), rewriter);
   auto bcast = [&](Value value) -> Value {
-    return broadcast(builder, value, *shape);
+    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(), *shape);
+  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 shape = vectorShape(op.operand().getType(), isF32);
+  if (!getElementTypeOrSelf(op.operand()).isF32())
+    return rewriter.notifyMatchFailure(op, "unsupported operand type");
+
+  ArrayRef<int64_t> shape = vectorShape(op.operand());
+
   // Only support already-vectorized rsqrt's.
-  if (!shape.hasValue() || shape->back() % 8 != 0)
+  if (shape.empty() || shape.back() % 8 != 0)
     return rewriter.notifyMatchFailure(op, "unsupported operand type");
 
   ImplicitLocOpBuilder builder(op->getLoc(), rewriter);
   auto bcast = [&](Value value) -> Value {
-    return broadcast(builder, value, *shape);
+    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);