Enable inlining for Linalg dialect

[ezhulenev]

Summary

Enable inlining for Linalg dialect.

NOTE: Please merge this when ready, I'm still don't have access.

[mehdi_amini]

Can you add various tests exercises the three methods?

[ezhulenev]

LinAlg->Linalg

[ezhulenev]

Add a test for inlining regions with Linalg ops

[ezhulenev]

I added a basic mlir test that inlines a function call with a Linalg operation in the body.

[nicolasvasilache]

Thanks @ezhulenev !

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

 
 #include "mlir/Dialect/Linalg/IR/LinalgTypes.h"
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/DialectImplementation.h"
 #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"
 
 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<
 #define GET_OP_LIST
 #include "mlir/Dialect/Linalg/IR/LinalgOps.cpp.inc"
       >();
   addOperations<
 #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;
   if (parser.parseKeyword(&keyword))
     return Type();
   MLIRContext *context = getContext();
 
   // Handle 'range' types.

mlir/test/Dialect/Linalg/inlining.mlir

+// 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
+}