Add support for conservatively inlining Affine operations.

This commit defines an initial implementation of the DialectInlinerInterface for the AffineOps dialect. This change allows for affine operations to be inlined into any region that is not an affine region. Inlining into affine regions requires special handling for dimension/symbol identifiers that will be added in followups.

PiperOrigin-RevId: 267467078

mlir/lib/Dialect/AffineOps/AffineOps.cpp

@@ -24,6 +24,7 @@
 #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/Support/Debug.h"
@@ -32,6 +33,43 @@ using llvm::dbgs;
 
 #define DEBUG_TYPE "affine-analysis"
 
+//===----------------------------------------------------------------------===//
+// AffineOpsDialect 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.
+  bool isLegalToInline(Region *dest, Region *src,
+                       BlockAndValueMapping &valueMapping) const final {
+    // Conservatively don't allow inlining into affine structures.
+    return false;
+  }
+
+  /// 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,
+                       BlockAndValueMapping &valueMapping) const final {
+    // Always allow inlining affine operations. There are some edge cases when
+    // inlining *into* affine structures, but that is handled in the other
+    // 'isLegalToInline' hook above.
+    return true;
+  }
+
+  /// Affine regions should be analyzed recursively.
+  bool shouldAnalyzeRecursively(Operation *op) const final { return true; }
+};
+} // end anonymous namespace
+
 //===----------------------------------------------------------------------===//
 // AffineOpsDialect
 //===----------------------------------------------------------------------===//
@@ -43,6 +81,7 @@ AffineOpsDialect::AffineOpsDialect(MLIRContext *context)
 #define GET_OP_LIST
 #include "mlir/Dialect/AffineOps/AffineOps.cpp.inc"
                 >();
+  addInterfaces<AffineInlinerInterface>();
 }
 
 /// A utility function to check if a given region is attached to a function.

mlir/test/AffineOps/inlining.mlir

@@ -0,0 +1,69 @@
+// RUN: mlir-opt %s -inline | FileCheck %s
+
+// Basic test that functions within affine operations are inlined.
+func @func_with_affine_ops(%N: index) {
+  %c = constant 200 : index
+  affine.for %i = 1 to 10 {
+    affine.if (i)[N] : (i - 2 >= 0, 4 - i >= 0)(%i)[%c]  {
+      %w = affine.apply (d0,d1)[s0] -> (d0+d1+s0) (%i, %i) [%N]
+    }
+  }
+  return
+}
+
+// CHECK-LABEL: func @inline_with_affine_ops
+func @inline_with_affine_ops() {
+  %c = constant 1 : index
+
+  // CHECK: affine.for
+  // CHECK-NEXT: affine.if
+  // CHECK-NEXT: affine.apply
+  // CHECK-NOT: call
+  call @func_with_affine_ops(%c) : (index) -> ()
+  return
+}
+
+// CHECK-LABEL: func @not_inline_in_affine_op
+func @not_inline_in_affine_op() {
+  %c = constant 1 : index
+
+  // CHECK-NOT: affine.if
+  // CHECK: call
+  affine.for %i = 1 to 10 {
+    call @func_with_affine_ops(%c) : (index) -> ()
+  }
+  return
+}
+
+// -----
+
+// Test when an invalid operation is nested in an affine op.
+func @func_with_invalid_nested_op() {
+  affine.for %i = 1 to 10 {
+    "foo.opaque"() : () -> ()
+  }
+  return
+}
+
+// CHECK-LABEL: func @not_inline_invalid_nest_op
+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.
+func @func_noop() {
+  return
+}
+
+// CHECK-LABEL: func @not_inline_into_affine_ops
+func @not_inline_into_affine_ops() {
+  // CHECK: call @func_noop
+  affine.for %i = 1 to 10 {
+    call @func_noop() : () -> ()
+  }
+  return
+}