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

diff --git a/mlir/lib/Dialect/AffineOps/AffineOps.cpp b/mlir/lib/Dialect/AffineOps/AffineOps.cpp
index 1c46b77fcfee5561202ff5091d80bc890eb33ba0..c6abc05b966284d413debf1d50d1aa7079e32211 100644 (file)
--- a/mlir/lib/Dialect/AffineOps/AffineOps.cpp
+++ b/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.

diff --git a/mlir/test/AffineOps/inlining.mlir b/mlir/test/AffineOps/inlining.mlir
new file mode 100644 (file)
index 0000000..53de793
--- /dev/null
+++ b/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
+}