[Mlir-commits] [mlir] 8c0dc1e - Enable inlining for Linalg dialect

[Nicolas Vasilache]

Author: Eugene Zhulenev
Date: 2020-09-16T10:19:13-04:00
New Revision: 8c0dc1e38b6c1a2d35c66ac4b0c1ccd616dd1685

URL: https://github.com/llvm/llvm-project/commit/8c0dc1e38b6c1a2d35c66ac4b0c1ccd616dd1685
DIFF: https://github.com/llvm/llvm-project/commit/8c0dc1e38b6c1a2d35c66ac4b0c1ccd616dd1685.diff

LOG: Enable inlining for Linalg dialect

Enable inlining for Linalg dialect.

Differential Revision: https://reviews.llvm.org/D87567

Added: 
    mlir/test/Dialect/Linalg/inlining.mlir

Modified: 
    mlir/lib/Dialect/Linalg/IR/LinalgTypes.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Linalg/IR/LinalgTypes.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgTypes.cpp
index b8bffd35f5a1..abc82f300f63 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgTypes.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgTypes.cpp
@@ -17,6 +17,7 @@
 #include "mlir/IR/StandardTypes.h"
 #include "mlir/Parser.h"
 #include "mlir/Support/LLVM.h"
+#include "mlir/Transforms/InliningUtils.h"
 
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/Support/raw_ostream.h"
@@ -24,6 +25,38 @@
 using namespace mlir;
 using namespace mlir::linalg;
 
+//===----------------------------------------------------------------------===//
+// LinalgDialect Dialect Interfaces
+//===----------------------------------------------------------------------===//
+
+namespace {
+
+struct LinalgInlinerInterface : public DialectInlinerInterface {
+  using DialectInlinerInterface::DialectInlinerInterface;
+
+  // We don't have any special restrictions on what can be inlined into
+  // destination regions (e.g. while/conditional bodies). Always allow it.
+  bool isLegalToInline(Region *dest, Region *src,
+                       BlockAndValueMapping &valueMapping) const final {
+    return true;
+  }
+  // Operations in Linalg dialect are always legal to inline.
+  bool isLegalToInline(Operation *, Region *,
+                       BlockAndValueMapping &) const final {
+    return true;
+  }
+  // Handle the given inlined terminator by replacing it with a new operation
+  // as necessary. Required when the region has only one block.
+  void handleTerminator(Operation *op,
+                        ArrayRef<Value> valuesToRepl) const final {}
+};
+
+} // end anonymous namespace
+
+//===----------------------------------------------------------------------===//
+// LinalgDialect
+//===----------------------------------------------------------------------===//
+
 void mlir::linalg::LinalgDialect::initialize() {
   addTypes<RangeType>();
   addOperations<
@@ -34,7 +67,9 @@ void mlir::linalg::LinalgDialect::initialize() {
 #define GET_OP_LIST
 #include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
       >();
+  addInterfaces<LinalgInlinerInterface>();
 }
+
 Type mlir::linalg::LinalgDialect::parseType(DialectAsmParser &parser) const {
   // Parse the main keyword for the type.
   StringRef keyword;

diff  --git a/mlir/test/Dialect/Linalg/inlining.mlir b/mlir/test/Dialect/Linalg/inlining.mlir
new file mode 100644
index 000000000000..1e5af263eb83
--- /dev/null
+++ b/mlir/test/Dialect/Linalg/inlining.mlir
@@ -0,0 +1,31 @@
+// RUN: mlir-opt %s -inline | FileCheck %s
+
+// These tests verify that regions with operations from Lingalg dialect
+// can be inlined.
+
+#accesses = [
+  affine_map<(i) -> (i)>,
+  affine_map<(i) -> (i)>
+]
+
+#trait = {
+  args_in = 1,
+  args_out = 1,
+  indexing_maps = #accesses,
+  iterator_types = ["parallel"]
+}
+
+func @inline_into(%arg0: memref<?xf32>) {
+  // CHECK: linalg.generic
+  call @inlined_fn(%arg0) : (memref<?xf32>) -> ()
+  return
+}
+
+func @inlined_fn(%arg0: memref<?xf32>) {
+  // CHECK: linalg.generic
+  linalg.generic #trait %arg0, %arg0 {
+    ^bb(%0 : f32, %1 : f32) :
+      linalg.yield %0 : f32
+  } : memref<?xf32>, memref<?xf32>
+  return
+}