[mlir] Support inlining into affine operations

[ftynse]

Summary

Introduce support for inlining into affine operations. This uses the generic
inline infrastructure and boils down to checking that, if applied, the inlining
doesn't violate the affine dimension/symbol value categorization. Given valid
IR, only the values that are valid dimensions/symbols thanks to being top-level
in their affine scope need special handling.

[rriddle]

Looks mostly fine to me, but will defer approval to those more closely following the affine dialect. (In case I forgot about some specific properties of the affine operations)

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

65

Is it still requiring ConstantOp directly and not using the more general m_Constant?

193

nit: You can merge all of these isa<> together: isa<AffineForOp, AffineParallelOp, ...>(parentOp)

[ftynse]

Address review

[dcaballe]

Thanks for working on this, Alex! LGTM! Just some minor comments.

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

40

remmains -> remains

58–59

Maybe we could do:

if (value.isa<BlockArgument>())
  return legalityCheck

so that a legal block argument doesn't continue? The condition below will always be false for a BlockArgument. We could even move value.isa<OpResult>() to an assert.

mlir/test/Dialect/Affine/inlining.mlir

75

could you also please add a test with an affine.vector_load or affine.vector_store?

84

should we capture %arg0 and match it here to make sure the output is correct (similar to IV)?

[bondhugula]

Looks great!

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

140

The doc comment is missing documentation of wouldBeClone.

[bondhugula]

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

144–146

I might have gotten it wrong, but shouldn't be a check for the op having an AffineScope trait here?

144–146

• shouldn't be -> shouldn't there be

[ftynse]

Address review

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

144–146

No, this is checking if we can inline into an affine control flow construct (for, if). These constructs are not affine scopes, but they are presumably in some scope already.

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

 //===- AffineOps.cpp - MLIR Affine 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Affine/IR/AffineValueMap.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Transforms/InliningUtils.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallBitVector.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/Debug.h"
 
 using namespace mlir;
 using llvm::dbgs;
 
 #define DEBUG_TYPE "affine-analysis"
 
+/// A utility function to check if a value is defined at the top level of
+/// `region` or is an argument of `region`. A value of index type defined at the
+/// top level of a `AffineScope` region is always a valid symbol for all
+/// uses in that region.
+static bool isTopLevelValue(Value value, Region *region) {
+  if (auto arg = value.dyn_cast<BlockArgument>())
+    return arg.getParentRegion() == region;
+  return value.getDefiningOp()->getParentRegion() == region;
+}
+
+/// Checks if `value` known to be a legal affine dimension or symbol in `src`
+/// region remains legal if the operation that uses it is inlined into `dest`
+/// with the given value mapping. `legalityCheck` is either `isValidDim` or
+/// `isValidSymbol`, depending on the value being required to remain a valid
+/// dimension or symbol.
+static bool
+remainsLegalAfterInline(Value value, Region *src, Region *dest,
+                        const BlockAndValueMapping &mapping,
+                        function_ref<bool(Value, Region *)> legalityCheck) {
+  // If the value is a valid dimension for any other reason than being
+  // a top-level value, it will remain valid: constants get inlined
+  // with the function, transitive affine applies also get inlined and
+  // will be checked themselves, etc.
+  if (!isTopLevelValue(value, src))
+    return true;
+
+  // If it's a top-level value because it's a block operand, i.e. a
+  // function argument, check whether the value replacing it after
+  // inlining is a valid dimension in the new region.
+  if (value.isa<BlockArgument>())
+    return legalityCheck(mapping.lookup(value), dest);
+
+  // If it's a top-level value beacuse it's defined in the region,
+  // it can only be inlined if the defining op is a constant or a
+  // `dim`, which can appear anywhere and be valid, since the defining
+  // op won't be top-level anymore after inlining.
+  Attribute operandCst;
+  return matchPattern(value.getDefiningOp(), m_Constant(&operandCst)) ||
+         value.getDefiningOp<DimOp>();
+}
+
+/// Checks if all values known to be legal affine dimensions or symbols in `src`
+/// remain so if their respective users are inlined into `dest`.
+static bool
+remainsLegalAfterInline(ValueRange values, Region *src, Region *dest,
+                        const BlockAndValueMapping &mapping,
+                        function_ref<bool(Value, Region *)> legalityCheck) {
+  return llvm::all_of(values, [&](Value v) {
+    return remainsLegalAfterInline(v, src, dest, mapping, legalityCheck);
+  });
+}
+
+/// Checks if an affine read or write operation remains legal after inlining
+/// from `src` to `dest`.
+template <typename OpTy>
+static bool remainsLegalAfterInline(OpTy op, Region *src, Region *dest,
+                                    const BlockAndValueMapping &mapping) {
+  static_assert(llvm::is_one_of<OpTy, AffineReadOpInterface,
+                                AffineWriteOpInterface>::value,
+                "only ops with affine read/write interface are supported");
+
+  AffineMap map = op.getAffineMap();
+  ValueRange dimOperands = op.getMapOperands().take_front(map.getNumDims());
+  ValueRange symbolOperands =
+      op.getMapOperands().take_back(map.getNumSymbols());
+  if (!remainsLegalAfterInline(
+          dimOperands, src, dest, mapping,
+          static_cast<bool (*)(Value, Region *)>(isValidDim)))
+    return false;
+  if (!remainsLegalAfterInline(
+          symbolOperands, src, dest, mapping,
+          static_cast<bool (*)(Value, Region *)>(isValidSymbol)))
+    return false;
+  return true;
+}
+
+/// Checks if an affine apply operation remains legal after inlining from `src`
+/// to `dest`.
+template <>
+bool remainsLegalAfterInline(AffineApplyOp op, Region *src, Region *dest,
+                             const BlockAndValueMapping &mapping) {
+  // If it's a valid dimension, we need to check that it remains so.
+  if (isValidDim(op.getResult(), src))
+    return remainsLegalAfterInline(
+        op.getMapOperands(), src, dest, mapping,
+        static_cast<bool (*)(Value, Region *)>(isValidDim));
+
+  // Otherwise it must be a valid symbol, check that it remains so.
+  return remainsLegalAfterInline(
+      op.getMapOperands(), src, dest, mapping,
+      static_cast<bool (*)(Value, Region *)>(isValidSymbol));
+}
+
 //===----------------------------------------------------------------------===//
 // AffineDialect Interfaces
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// This class defines the interface for handling inlining with affine
 /// operations.
 struct AffineInlinerInterface : public DialectInlinerInterface {
   using DialectInlinerInterface::DialectInlinerInterface;
 
   //===--------------------------------------------------------------------===//
   // Analysis Hooks
   //===--------------------------------------------------------------------===//
 
   /// Returns true if the given region 'src' can be inlined into the region
   /// 'dest' that is attached to an operation registered to the current dialect.
+  /// 'wouldBeCloned' is set if the region is cloned into its new location
+  /// rather than moved, indicating there may be other users.
   bool isLegalToInline(Region *dest, Region *src, bool wouldBeCloned,
                        BlockAndValueMapping &valueMapping) const final {
-    // Conservatively don't allow inlining into affine structures.
+    // We can inline into affine loops and conditionals if this doesn't break
+    // affine value categorization rules.
+    Operation *destOp = dest->getParentOp();
+    if (!isa<AffineParallelOp, AffineForOp, AffineIfOp>(destOp))
+      return false;
+
+    // Multi-block regions cannot be inlined into affine constructs, all of
+    // which require single-block regions.
+    if (!llvm::hasSingleElement(*src))
       return false;
+
+    // Side-effecting operations that the affine dialect cannot understand
+    // should not be inlined.
+    Block &srcBlock = src->front();
+    for (Operation &op : srcBlock) {
+      // Ops with no side effects are fine,
+      if (auto iface = dyn_cast<MemoryEffectOpInterface>(op)) {
+        if (iface.hasNoEffect())
+          continue;
+      }
+
+      // Assuming the inlined region is valid, we only need to check if the
+      // inlining would change it.
+      bool remainsValid =
+          llvm::TypeSwitch<Operation *, bool>(&op)
+              .Case<AffineApplyOp, AffineReadOpInterface,
+                    AffineWriteOpInterface>([&](auto op) {
+                return remainsLegalAfterInline(op, src, dest, valueMapping);
+              })
+              .Default([](Operation *) {
+                // Conservatively disallow inlining ops we cannot reason about.
+                return false;
+              });
+
+      if (!remainsValid)
+        return false;
+    }
+
+    return true;
   }
 
   /// Returns true if the given operation 'op', that is registered to this
   /// dialect, can be inlined into the given region, false otherwise.
   bool isLegalToInline(Operation *op, Region *region, bool wouldBeCloned,
                        BlockAndValueMapping &valueMapping) const final {
-    // Always allow inlining affine operations into the top-level region of a
-    // function. There are some edge cases when inlining *into* affine
-    // structures, but that is handled in the other 'isLegalToInline' hook
-    // above.
-    // TODO: We should be able to inline into other regions than functions.
-    return isa<FuncOp>(region->getParentOp());
+    // Always allow inlining affine operations into a region that is marked as
+    // affine scope, or into affine loops and conditionals. There are some edge
+    // cases when inlining *into* affine structures, but that is handled in the
+    // other 'isLegalToInline' hook above.
+    Operation *parentOp = region->getParentOp();
+    return parentOp->hasTrait<OpTrait::AffineScope>() ||
+           isa<AffineForOp, AffineParallelOp, AffineIfOp>(parentOp);
   }
 
   /// Affine regions should be analyzed recursively.
   bool shouldAnalyzeRecursively(Operation *op) const final { return true; }
 };
 } // end anonymous namespace
 
 //===----------------------------------------------------------------------===//
if (auto arg = value.dyn_cast<BlockArgument>()) {
     Operation *parentOp = arg.getOwner()->getParentOp();
     return parentOp && parentOp->hasTrait<OpTrait::AffineScope>();
   }
   // The defining Op may live in an unlinked block so its parent Op may be null.
   Operation *parentOp = value.getDefiningOp()->getParentOp();
   return parentOp && parentOp->hasTrait<OpTrait::AffineScope>();
 }
 
-/// A utility function to check if a value is defined at the top level of
-/// `region` or is an argument of `region`. A value of index type defined at the
-/// top level of a `AffineScope` region is always a valid symbol for all
-/// uses in that region.
-static bool isTopLevelValue(Value value, Region *region) {
-  if (auto arg = value.dyn_cast<BlockArgument>())
-    return arg.getParentRegion() == region;
-  return value.getDefiningOp()->getParentRegion() == region;
-}
-
 /// Returns the closest region enclosing `op` that is held by an operation with
 /// trait `AffineScope`; `nullptr` if there is no such region.
 //  TODO: getAffineScope should be publicly exposed for affine passes/utilities.
 static Region *getAffineScope(Operation *op) {
   auto *curOp = op;
   while (auto *parentOp = curOp->getParentOp()) {
     if (parentOp->hasTrait<OpTrait::AffineScope>())
       return curOp->getParentRegion();

mlir/test/Dialect/Affine/inlining.mlir

 func @not_inline_invalid_nest_op() {
   // CHECK: call @func_with_invalid_nested_op
   call @func_with_invalid_nested_op() : () -> ()
   return
 }
 
 // -----
 
-// Test that calls are not inlined into affine structures.
+// Test that calls are inlined into affine structures.
 func @func_noop() {
   return
 }
 
-// CHECK-LABEL: func @not_inline_into_affine_ops
-func @not_inline_into_affine_ops() {
-  // CHECK: call @func_noop
+// CHECK-LABEL: func @inline_into_affine_ops
+func @inline_into_affine_ops() {
+  // CHECK-NOT: call @func_noop
   affine.for %i = 1 to 10 {
     call @func_noop() : () -> ()
   }
   return
 }
+
+// -----
+
+// Test that calls with dimension arguments are properly inlined.
+func @func_dim(%arg0: index, %arg1: memref<?xf32>) {
+  affine.load %arg1[%arg0] : memref<?xf32>
+  return
+}
+
+// CHECK-LABEL: @inline_dimension
+// CHECK: (%[[ARG0:.*]]: memref<?xf32>)
+func @inline_dimension(%arg0: memref<?xf32>) {
+  // CHECK: affine.for %[[IV:.*]] =
+  affine.for %i = 1 to 42 {
+    // CHECK-NOT: call @func_dim
+    // CHECK: affine.load %[[ARG0]][%[[IV]]]
+    call @func_dim(%i, %arg0) : (index, memref<?xf32>) -> ()
+  }
+  return
+}
+
+// -----
+
+// Test that calls with vector operations are also inlined.
+func @func_vector_dim(%arg0: index, %arg1: memref<32xf32>) {
+  affine.vector_load %arg1[%arg0] : memref<32xf32>, vector<4xf32>
+  return
+}
+
+// CHECK-LABEL: @inline_dimension_vector
+// CHECK: (%[[ARG0:.*]]: memref<32xf32>)
+func @inline_dimension_vector(%arg0: memref<32xf32>) {
+  // CHECK: affine.for %[[IV:.*]] =
+  affine.for %i = 1 to 42 {
+    // CHECK-NOT: call @func_dim
+    // CHECK: affine.vector_load %[[ARG0]][%[[IV]]]
+    call @func_vector_dim(%i, %arg0) : (index, memref<32xf32>) -> ()
+  }
+  return
+}
+
+// -----
+
+// Test that calls that would result in violation of affine value
+// categorization (top-level value stop being top-level) are not inlined.
+func private @get_index() -> index
+
+func @func_top_level(%arg0: memref<?xf32>) {
+  %0 = call @get_index() : () -> index
+  affine.load %arg0[%0] : memref<?xf32>
+  return
+}
+
+// CHECK-LABEL: @no_inline_not_top_level
+func @no_inline_not_top_level(%arg0: memref<?xf32>) {
+  affine.for %i = 1 to 42 {
+    // CHECK: call @func_top_level
+    call @func_top_level(%arg0) : (memref<?xf32>) -> ()
+  }
+  return
+}