[MLIR][NFC] Move expandAffineMap/Expr out to Affine utils

[bondhugula]

Summary

Move expandAffineMap and expandAffineApplyExpr out to AffineUtils. This
is a useful method. The child revision uses it. NFC.

mlir/include/mlir/Conversion/AffineToStandard/AffineToStandard.h

 //===- AffineToStandard.h - Convert Affine to Standard dialect --*- 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_CONVERSION_AFFINETOSTANDARD_AFFINETOSTANDARD_H
 #define MLIR_CONVERSION_AFFINETOSTANDARD_AFFINETOSTANDARD_H
 
 #include "mlir/Support/LLVM.h"
 
 namespace mlir {
-class AffineExpr;
 class AffineForOp;
-class AffineMap;
-class AffineParallelOp;
 class Location;
 struct LogicalResult;
 class OpBuilder;
 class Pass;
 class RewritePattern;
 class Value;
 class ValueRange;
 
 class RewritePatternSet;
 
-/// Emit code that computes the given affine expression using standard
-/// arithmetic operations applied to the provided dimension and symbol values.
-Value expandAffineExpr(OpBuilder &builder, Location loc, AffineExpr expr,
-                       ValueRange dimValues, ValueRange symbolValues);
-
-/// Create a sequence of operations that implement the `affineMap` applied to
-/// the given `operands` (as it it were an AffineApplyOp).
-Optional<SmallVector<Value, 8>> expandAffineMap(OpBuilder &builder,
-                                                Location loc,
-                                                AffineMap affineMap,
-                                                ValueRange operands);
-
 /// Collect a set of patterns to convert from the Affine dialect to the Standard
 /// dialect, in particular convert structured affine control flow into CFG
 /// branch-based control flow.
 void populateAffineToStdConversionPatterns(RewritePatternSet &patterns);
 
 /// Collect a set of patterns to convert vector-related Affine ops to the Vector
 /// dialect.
 void populateAffineToVectorConversionPatterns(RewritePatternSet &patterns);

mlir/include/mlir/Dialect/Affine/Utils.h

 /// In the following cases, `sliceOps` remains empty:
 ///   1. If none of opInst's operands were the result of an affine.apply
 ///      (i.e., there was no affine computation slice to create).
 ///   2. If all the affine.apply op's supplying operands to this opInst did not
 ///      have any uses other than those in this opInst.
 void createAffineComputationSlice(Operation *opInst,
                                   SmallVectorImpl<AffineApplyOp> *sliceOps);
 
+/// Emit code that computes the given affine expression using standard
+/// arithmetic operations applied to the provided dimension and symbol values.
+Value expandAffineExpr(OpBuilder &builder, Location loc, AffineExpr expr,
+                       ValueRange dimValues, ValueRange symbolValues);
+
+/// Create a sequence of operations that implement the `affineMap` applied to
+/// the given `operands` (as it it were an AffineApplyOp).
+Optional<SmallVector<Value, 8>> expandAffineMap(OpBuilder &builder,
+                                                Location loc,
+                                                AffineMap affineMap,
+                                                ValueRange operands);
+
 } // namespace mlir
 
 #endif // MLIR_DIALECT_AFFINE_UTILS_H

mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp

 // operations) within a function into their standard If and For equivalent ops.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Conversion/AffineToStandard/AffineToStandard.h"
 
 #include "../PassDetail.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
-#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Affine/Utils.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/SCF.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
-#include "mlir/IR/AffineExprVisitor.h"
 #include "mlir/IR/BlockAndValueMapping.h"
-#include "mlir/IR/Builders.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "mlir/Transforms/Passes.h"
 
 using namespace mlir;
 using namespace mlir::vector;
 
-namespace {
-/// Visit affine expressions recursively and build the sequence of operations
-/// that correspond to it.  Visitation functions return an Value of the
-/// expression subtree they visited or `nullptr` on error.
-class AffineApplyExpander
-    : public AffineExprVisitor<AffineApplyExpander, Value> {
-public:
-  /// This internal class expects arguments to be non-null, checks must be
-  /// performed at the call site.
-  AffineApplyExpander(OpBuilder &builder, ValueRange dimValues,
-                      ValueRange symbolValues, Location loc)
-      : builder(builder), dimValues(dimValues), symbolValues(symbolValues),
-        loc(loc) {}
-
-  template <typename OpTy>
-  Value buildBinaryExpr(AffineBinaryOpExpr expr) {
-    auto lhs = visit(expr.getLHS());
-    auto rhs = visit(expr.getRHS());
-    if (!lhs || !rhs)
-      return nullptr;
-    auto op = builder.create<OpTy>(loc, lhs, rhs);
-    return op.getResult();
-  }
-
-  Value visitAddExpr(AffineBinaryOpExpr expr) {
-    return buildBinaryExpr<arith::AddIOp>(expr);
-  }
-
-  Value visitMulExpr(AffineBinaryOpExpr expr) {
-    return buildBinaryExpr<arith::MulIOp>(expr);
-  }
-
-  /// Euclidean modulo operation: negative RHS is not allowed.
-  /// Remainder of the euclidean integer division is always non-negative.
-  ///
-  /// Implemented as
-  ///
-  ///     a mod b =
-  ///         let remainder = srem a, b;
-  ///             negative = a < 0 in
-  ///         select negative, remainder + b, remainder.
-  Value visitModExpr(AffineBinaryOpExpr expr) {
-    auto rhsConst = expr.getRHS().dyn_cast<AffineConstantExpr>();
-    if (!rhsConst) {
-      emitError(
-          loc,
-          "semi-affine expressions (modulo by non-const) are not supported");
-      return nullptr;
-    }
-    if (rhsConst.getValue() <= 0) {
-      emitError(loc, "modulo by non-positive value is not supported");
-      return nullptr;
-    }
-
-    auto lhs = visit(expr.getLHS());
-    auto rhs = visit(expr.getRHS());
-    assert(lhs && rhs && "unexpected affine expr lowering failure");
-
-    Value remainder = builder.create<arith::RemSIOp>(loc, lhs, rhs);
-    Value zeroCst = builder.create<arith::ConstantIndexOp>(loc, 0);
-    Value isRemainderNegative = builder.create<arith::CmpIOp>(
-        loc, arith::CmpIPredicate::slt, remainder, zeroCst);
-    Value correctedRemainder =
-        builder.create<arith::AddIOp>(loc, remainder, rhs);
-    Value result = builder.create<arith::SelectOp>(
-        loc, isRemainderNegative, correctedRemainder, remainder);
-    return result;
-  }
-
-  /// Floor division operation (rounds towards negative infinity).
-  ///
-  /// For positive divisors, it can be implemented without branching and with a
-  /// single division operation as
-  ///
-  ///        a floordiv b =
-  ///            let negative = a < 0 in
-  ///            let absolute = negative ? -a - 1 : a in
-  ///            let quotient = absolute / b in
-  ///                negative ? -quotient - 1 : quotient
-  Value visitFloorDivExpr(AffineBinaryOpExpr expr) {
-    auto rhsConst = expr.getRHS().dyn_cast<AffineConstantExpr>();
-    if (!rhsConst) {
-      emitError(
-          loc,
-          "semi-affine expressions (division by non-const) are not supported");
-      return nullptr;
-    }
-    if (rhsConst.getValue() <= 0) {
-      emitError(loc, "division by non-positive value is not supported");
-      return nullptr;
-    }
-
-    auto lhs = visit(expr.getLHS());
-    auto rhs = visit(expr.getRHS());
-    assert(lhs && rhs && "unexpected affine expr lowering failure");
-
-    Value zeroCst = builder.create<arith::ConstantIndexOp>(loc, 0);
-    Value noneCst = builder.create<arith::ConstantIndexOp>(loc, -1);
-    Value negative = builder.create<arith::CmpIOp>(
-        loc, arith::CmpIPredicate::slt, lhs, zeroCst);
-    Value negatedDecremented = builder.create<arith::SubIOp>(loc, noneCst, lhs);
-    Value dividend =
-        builder.create<arith::SelectOp>(loc, negative, negatedDecremented, lhs);
-    Value quotient = builder.create<arith::DivSIOp>(loc, dividend, rhs);
-    Value correctedQuotient =
-        builder.create<arith::SubIOp>(loc, noneCst, quotient);
-    Value result = builder.create<arith::SelectOp>(loc, negative,
-                                                   correctedQuotient, quotient);
-    return result;
-  }
-
-  /// Ceiling division operation (rounds towards positive infinity).
-  ///
-  /// For positive divisors, it can be implemented without branching and with a
-  /// single division operation as
-  ///
-  ///     a ceildiv b =
-  ///         let negative = a <= 0 in
-  ///         let absolute = negative ? -a : a - 1 in
-  ///         let quotient = absolute / b in
-  ///             negative ? -quotient : quotient + 1
-  Value visitCeilDivExpr(AffineBinaryOpExpr expr) {
-    auto rhsConst = expr.getRHS().dyn_cast<AffineConstantExpr>();
-    if (!rhsConst) {
-      emitError(loc) << "semi-affine expressions (division by non-const) are "
-                        "not supported";
-      return nullptr;
-    }
-    if (rhsConst.getValue() <= 0) {
-      emitError(loc, "division by non-positive value is not supported");
-      return nullptr;
-    }
-    auto lhs = visit(expr.getLHS());
-    auto rhs = visit(expr.getRHS());
-    assert(lhs && rhs && "unexpected affine expr lowering failure");
-
-    Value zeroCst = builder.create<arith::ConstantIndexOp>(loc, 0);
-    Value oneCst = builder.create<arith::ConstantIndexOp>(loc, 1);
-    Value nonPositive = builder.create<arith::CmpIOp>(
-        loc, arith::CmpIPredicate::sle, lhs, zeroCst);
-    Value negated = builder.create<arith::SubIOp>(loc, zeroCst, lhs);
-    Value decremented = builder.create<arith::SubIOp>(loc, lhs, oneCst);
-    Value dividend =
-        builder.create<arith::SelectOp>(loc, nonPositive, negated, decremented);
-    Value quotient = builder.create<arith::DivSIOp>(loc, dividend, rhs);
-    Value negatedQuotient =
-        builder.create<arith::SubIOp>(loc, zeroCst, quotient);
-    Value incrementedQuotient =
-        builder.create<arith::AddIOp>(loc, quotient, oneCst);
-    Value result = builder.create<arith::SelectOp>(
-        loc, nonPositive, negatedQuotient, incrementedQuotient);
-    return result;
-  }
-
-  Value visitConstantExpr(AffineConstantExpr expr) {
-    auto op = builder.create<arith::ConstantIndexOp>(loc, expr.getValue());
-    return op.getResult();
-  }
-
-  Value visitDimExpr(AffineDimExpr expr) {
-    assert(expr.getPosition() < dimValues.size() &&
-           "affine dim position out of range");
-    return dimValues[expr.getPosition()];
-  }
-
-  Value visitSymbolExpr(AffineSymbolExpr expr) {
-    assert(expr.getPosition() < symbolValues.size() &&
-           "symbol dim position out of range");
-    return symbolValues[expr.getPosition()];
-  }
-
-private:
-  OpBuilder &builder;
-  ValueRange dimValues;
-  ValueRange symbolValues;
-
-  Location loc;
-};
-} // namespace
-
-/// Create a sequence of operations that implement the `expr` applied to the
-/// given dimension and symbol values.
-mlir::Value mlir::expandAffineExpr(OpBuilder &builder, Location loc,
-                                   AffineExpr expr, ValueRange dimValues,
-                                   ValueRange symbolValues) {
-  return AffineApplyExpander(builder, dimValues, symbolValues, loc).visit(expr);
-}
-
-/// Create a sequence of operations that implement the `affineMap` applied to
-/// the given `operands` (as it it were an AffineApplyOp).
-Optional<SmallVector<Value, 8>> mlir::expandAffineMap(OpBuilder &builder,
-                                                      Location loc,
-                                                      AffineMap affineMap,
-                                                      ValueRange operands) {
-  auto numDims = affineMap.getNumDims();
-  auto expanded = llvm::to_vector<8>(
-      llvm::map_range(affineMap.getResults(),
-                      [numDims, &builder, loc, operands](AffineExpr expr) {
-                        return expandAffineExpr(builder, loc, expr,
-                                                operands.take_front(numDims),
-                                                operands.drop_front(numDims));
-                      }));
-  if (llvm::all_of(expanded, [](Value v) { return v; }))
-    return expanded;
-  return None;
-}
-
 /// Given a range of values, emit the code that reduces them with "min" or "max"
 /// depending on the provided comparison predicate.  The predicate defines which
 /// comparison to perform, "lt" for "min", "gt" for "max" and is used for the
 /// `cmpi` operation followed by the `select` operation:
 ///
 ///   %cond   = arith.cmpi "predicate" %v0, %v1
 ///   %result = select %cond, %v0, %v1
 ///

mlir/lib/Conversion/AffineToStandard/CMakeLists.txt

 add_mlir_conversion_library(MLIRAffineToStandard
   AffineToStandard.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Conversion/AffineToStandard
 
   DEPENDS
   MLIRConversionPassIncGen
 
   LINK_COMPONENTS
   Core
 
   LINK_LIBS PUBLIC
   MLIRAffine
+  MLIRAffineUtils
   MLIRArithmetic
+  MLIRIR
   MLIRMemRef
   MLIRSCF
   MLIRPass
   MLIRStandard
   MLIRTransforms
-  MLIRIR
   MLIRVector
   )

mlir/lib/Dialect/Affine/Utils/Utils.cpp

 
 #include "mlir/Dialect/Affine/Utils.h"
 
 #include "mlir/Dialect/Affine/Analysis/Utils.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Affine/IR/AffineValueMap.h"
 #include "mlir/Dialect/Affine/LoopUtils.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/IR/AffineExprVisitor.h"
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/Dominance.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 #define DEBUG_TYPE "affine-utils"
 
 using namespace mlir;
 
+namespace {
+/// Visit affine expressions recursively and build the sequence of operations
+/// that correspond to it.  Visitation functions return an Value of the
+/// expression subtree they visited or `nullptr` on error.
+class AffineApplyExpander
+    : public AffineExprVisitor<AffineApplyExpander, Value> {
+public:
+  /// This internal class expects arguments to be non-null, checks must be
+  /// performed at the call site.
+  AffineApplyExpander(OpBuilder &builder, ValueRange dimValues,
+                      ValueRange symbolValues, Location loc)
+      : builder(builder), dimValues(dimValues), symbolValues(symbolValues),
+        loc(loc) {}
+
+  template <typename OpTy>
+  Value buildBinaryExpr(AffineBinaryOpExpr expr) {
+    auto lhs = visit(expr.getLHS());
+    auto rhs = visit(expr.getRHS());
+    if (!lhs || !rhs)
+      return nullptr;
+    auto op = builder.create<OpTy>(loc, lhs, rhs);
+    return op.getResult();
+  }
+
+  Value visitAddExpr(AffineBinaryOpExpr expr) {
+    return buildBinaryExpr<arith::AddIOp>(expr);
+  }
+
+  Value visitMulExpr(AffineBinaryOpExpr expr) {
+    return buildBinaryExpr<arith::MulIOp>(expr);
+  }
+
+  /// Euclidean modulo operation: negative RHS is not allowed.
+  /// Remainder of the euclidean integer division is always non-negative.
+  ///
+  /// Implemented as
+  ///
+  ///     a mod b =
+  ///         let remainder = srem a, b;
+  ///             negative = a < 0 in
+  ///         select negative, remainder + b, remainder.
+  Value visitModExpr(AffineBinaryOpExpr expr) {
+    auto rhsConst = expr.getRHS().dyn_cast<AffineConstantExpr>();
+    if (!rhsConst) {
+      emitError(
+          loc,
+          "semi-affine expressions (modulo by non-const) are not supported");
+      return nullptr;
+    }
+    if (rhsConst.getValue() <= 0) {
+      emitError(loc, "modulo by non-positive value is not supported");
+      return nullptr;
+    }
+
+    auto lhs = visit(expr.getLHS());
+    auto rhs = visit(expr.getRHS());
+    assert(lhs && rhs && "unexpected affine expr lowering failure");
+
+    Value remainder = builder.create<arith::RemSIOp>(loc, lhs, rhs);
+    Value zeroCst = builder.create<arith::ConstantIndexOp>(loc, 0);
+    Value isRemainderNegative = builder.create<arith::CmpIOp>(
+        loc, arith::CmpIPredicate::slt, remainder, zeroCst);
+    Value correctedRemainder =
+        builder.create<arith::AddIOp>(loc, remainder, rhs);
+    Value result = builder.create<arith::SelectOp>(
+        loc, isRemainderNegative, correctedRemainder, remainder);
+    return result;
+  }
+
+  /// Floor division operation (rounds towards negative infinity).
+  ///
+  /// For positive divisors, it can be implemented without branching and with a
+  /// single division operation as
+  ///
+  ///        a floordiv b =
+  ///            let negative = a < 0 in
+  ///            let absolute = negative ? -a - 1 : a in
+  ///            let quotient = absolute / b in
+  ///                negative ? -quotient - 1 : quotient
+  Value visitFloorDivExpr(AffineBinaryOpExpr expr) {
+    auto rhsConst = expr.getRHS().dyn_cast<AffineConstantExpr>();
+    if (!rhsConst) {
+      emitError(
+          loc,
+          "semi-affine expressions (division by non-const) are not supported");
+      return nullptr;
+    }
+    if (rhsConst.getValue() <= 0) {
+      emitError(loc, "division by non-positive value is not supported");
+      return nullptr;
+    }
+
+    auto lhs = visit(expr.getLHS());
+    auto rhs = visit(expr.getRHS());
+    assert(lhs && rhs && "unexpected affine expr lowering failure");
+
+    Value zeroCst = builder.create<arith::ConstantIndexOp>(loc, 0);
+    Value noneCst = builder.create<arith::ConstantIndexOp>(loc, -1);
+    Value negative = builder.create<arith::CmpIOp>(
+        loc, arith::CmpIPredicate::slt, lhs, zeroCst);
+    Value negatedDecremented = builder.create<arith::SubIOp>(loc, noneCst, lhs);
+    Value dividend =
+        builder.create<arith::SelectOp>(loc, negative, negatedDecremented, lhs);
+    Value quotient = builder.create<arith::DivSIOp>(loc, dividend, rhs);
+    Value correctedQuotient =
+        builder.create<arith::SubIOp>(loc, noneCst, quotient);
+    Value result = builder.create<arith::SelectOp>(loc, negative,
+                                                   correctedQuotient, quotient);
+    return result;
+  }
+
+  /// Ceiling division operation (rounds towards positive infinity).
+  ///
+  /// For positive divisors, it can be implemented without branching and with a
+  /// single division operation as
+  ///
+  ///     a ceildiv b =
+  ///         let negative = a <= 0 in
+  ///         let absolute = negative ? -a : a - 1 in
+  ///         let quotient = absolute / b in
+  ///             negative ? -quotient : quotient + 1
+  Value visitCeilDivExpr(AffineBinaryOpExpr expr) {
+    auto rhsConst = expr.getRHS().dyn_cast<AffineConstantExpr>();
+    if (!rhsConst) {
+      emitError(loc) << "semi-affine expressions (division by non-const) are "
+                        "not supported";
+      return nullptr;
+    }
+    if (rhsConst.getValue() <= 0) {
+      emitError(loc, "division by non-positive value is not supported");
+      return nullptr;
+    }
+    auto lhs = visit(expr.getLHS());
+    auto rhs = visit(expr.getRHS());
+    assert(lhs && rhs && "unexpected affine expr lowering failure");
+
+    Value zeroCst = builder.create<arith::ConstantIndexOp>(loc, 0);
+    Value oneCst = builder.create<arith::ConstantIndexOp>(loc, 1);
+    Value nonPositive = builder.create<arith::CmpIOp>(
+        loc, arith::CmpIPredicate::sle, lhs, zeroCst);
+    Value negated = builder.create<arith::SubIOp>(loc, zeroCst, lhs);
+    Value decremented = builder.create<arith::SubIOp>(loc, lhs, oneCst);
+    Value dividend =
+        builder.create<arith::SelectOp>(loc, nonPositive, negated, decremented);
+    Value quotient = builder.create<arith::DivSIOp>(loc, dividend, rhs);
+    Value negatedQuotient =
+        builder.create<arith::SubIOp>(loc, zeroCst, quotient);
+    Value incrementedQuotient =
+        builder.create<arith::AddIOp>(loc, quotient, oneCst);
+    Value result = builder.create<arith::SelectOp>(
+        loc, nonPositive, negatedQuotient, incrementedQuotient);
+    return result;
+  }
+
+  Value visitConstantExpr(AffineConstantExpr expr) {
+    auto op = builder.create<arith::ConstantIndexOp>(loc, expr.getValue());
+    return op.getResult();
+  }
+
+  Value visitDimExpr(AffineDimExpr expr) {
+    assert(expr.getPosition() < dimValues.size() &&
+           "affine dim position out of range");
+    return dimValues[expr.getPosition()];
+  }
+
+  Value visitSymbolExpr(AffineSymbolExpr expr) {
+    assert(expr.getPosition() < symbolValues.size() &&
+           "symbol dim position out of range");
+    return symbolValues[expr.getPosition()];
+  }
+
+private:
+  OpBuilder &builder;
+  ValueRange dimValues;
+  ValueRange symbolValues;
+
+  Location loc;
+};
+} // namespace
+
+/// Create a sequence of operations that implement the `expr` applied to the
+/// given dimension and symbol values.
+mlir::Value mlir::expandAffineExpr(OpBuilder &builder, Location loc,
+                                   AffineExpr expr, ValueRange dimValues,
+                                   ValueRange symbolValues) {
+  return AffineApplyExpander(builder, dimValues, symbolValues, loc).visit(expr);
+}
+
+/// Create a sequence of operations that implement the `affineMap` applied to
+/// the given `operands` (as it it were an AffineApplyOp).
+Optional<SmallVector<Value, 8>> mlir::expandAffineMap(OpBuilder &builder,
+                                                      Location loc,
+                                                      AffineMap affineMap,
+                                                      ValueRange operands) {
+  auto numDims = affineMap.getNumDims();
+  auto expanded = llvm::to_vector<8>(
+      llvm::map_range(affineMap.getResults(),
+                      [numDims, &builder, loc, operands](AffineExpr expr) {
+                        return expandAffineExpr(builder, loc, expr,
+                                                operands.take_front(numDims),
+                                                operands.drop_front(numDims));
+                      }));
+  if (llvm::all_of(expanded, [](Value v) { return v; }))
+    return expanded;
+  return None;
+}
+
 /// Promotes the `then` or the `else` block of `ifOp` (depending on whether
 /// `elseBlock` is false or true) into `ifOp`'s containing block, and discards
 /// the rest of the op.
 static void promoteIfBlock(AffineIfOp ifOp, bool elseBlock) {
   if (elseBlock)
     assert(ifOp.hasElse() && "else block expected");
 
   Block *destBlock = ifOp->getBlock();