[Mlir-commits] [mlir] 758be97 - [mlir:LLVM] Rudimentary inlining support for LLVM load store.

[Tobias Gysi]

Author: Johannes de Fine Licht
Date: 2023-01-09T10:28:21+01:00
New Revision: 758be971dc11379ba2ee7c47099b59871467d56e

URL: https://github.com/llvm/llvm-project/commit/758be971dc11379ba2ee7c47099b59871467d56e
DIFF: https://github.com/llvm/llvm-project/commit/758be971dc11379ba2ee7c47099b59871467d56e.diff

LOG: [mlir:LLVM] Rudimentary inlining support for LLVM load store.

Conservatively only allow inlining for loads and stores that don't carry
any attributes that require handling while inlining. This can later be
relaxed when proper handling is introduced.

Reviewed By: Dinistro, gysit

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

Added: 
    

Modified: 
    mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
    mlir/test/Dialect/LLVMIR/inlining.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
index c32ca2b8f5a0b..cc6213a9817ec 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
@@ -2830,10 +2830,25 @@ namespace {
 struct LLVMInlinerInterface : public DialectInlinerInterface {
   using DialectInlinerInterface::DialectInlinerInterface;
 
-  /// Conservatively only allow inlining of pure ops.
+  /// Conservative allowlist-based inlining of operations supported so far.
   bool isLegalToInline(Operation *op, Region *, bool,
                        BlockAndValueMapping &) const final {
-    return isPure(op);
+    if (isPure(op))
+      return true;
+    return llvm::TypeSwitch<Operation *, bool>(op)
+        .Case<LLVM::LoadOp, LLVM::StoreOp>([&](auto memOp) {
+          // Some attributes on load and store operations require handling
+          // during inlining. Since this is not yet implemented, refuse to
+          // inline memory operations that have any of these attributes.
+          if (memOp.getAccessGroups())
+            return false;
+          if (memOp.getAliasScopes())
+            return false;
+          if (memOp.getNoaliasScopes())
+            return false;
+          return true;
+        })
+        .Default([](auto) { return false; });
   }
 };
 } // end anonymous namespace

diff  --git a/mlir/test/Dialect/LLVMIR/inlining.mlir b/mlir/test/Dialect/LLVMIR/inlining.mlir
index ffe03ee0e9738..0ba983f408a0b 100644
--- a/mlir/test/Dialect/LLVMIR/inlining.mlir
+++ b/mlir/test/Dialect/LLVMIR/inlining.mlir
@@ -1,16 +1,29 @@
-// RUN: mlir-opt %s -inline | FileCheck %s
+// RUN: mlir-opt %s -inline -split-input-file | FileCheck %s
 
-// CHECK-LABEL: func.func @test_inline() -> i32 {
-// CHECK-NEXT: %[[RES:.*]] = llvm.mlir.constant(42 : i32) : i32
+func.func @inner_func_inlinable(%ptr : !llvm.ptr) -> i32 {
+  %0 = llvm.mlir.constant(42 : i32) : i32
+  llvm.store %0, %ptr { alignment = 8 } : i32, !llvm.ptr
+  %1 = llvm.load %ptr { alignment = 8 } : !llvm.ptr -> i32
+  return %1 : i32
+}
+
+// CHECK-LABEL: func.func @test_inline(
+// CHECK-SAME: %[[PTR:[a-zA-Z0-9_]+]]
+// CHECK-NEXT: %[[CST:.*]] = llvm.mlir.constant(42 : i32) : i32
+// CHECK-NEXT: llvm.store %[[CST]], %[[PTR]]
+// CHECK-NEXT: %[[RES:.+]] = llvm.load %[[PTR]]
 // CHECK-NEXT: return %[[RES]] : i32
-func.func @test_inline() -> i32 {
-  %0 = call @inner_func_inlinable() : () -> i32
+func.func @test_inline(%ptr : !llvm.ptr) -> i32 {
+  %0 = call @inner_func_inlinable(%ptr) : (!llvm.ptr) -> i32
   return %0 : i32
 }
 
-func.func @inner_func_inlinable() -> i32 {
-  %0 = llvm.mlir.constant(42 : i32) : i32
-  return %0 : i32
+// -----
+
+func.func @inner_func_not_inlinable() -> !llvm.ptr<f64> {
+  %0 = llvm.mlir.constant(0 : i32) : i32
+  %1 = llvm.alloca %0 x f64 : (i32) -> !llvm.ptr<f64>
+  return %1 : !llvm.ptr<f64>
 }
 
 // CHECK-LABEL: func.func @test_not_inline() -> !llvm.ptr<f64> {
@@ -21,8 +34,25 @@ func.func @test_not_inline() -> !llvm.ptr<f64> {
   return %0 : !llvm.ptr<f64>
 }
 
-func.func @inner_func_not_inlinable() -> !llvm.ptr<f64> {
-  %0 = llvm.mlir.constant(0 : i32) : i32
-  %1 = llvm.alloca %0 x f64 : (i32) -> !llvm.ptr<f64>
-  return %1 : !llvm.ptr<f64>
+// -----
+
+llvm.metadata @metadata {
+  llvm.access_group @group
+  llvm.return
+}
+
+func.func private @with_mem_attr(%ptr : !llvm.ptr) -> () {
+  %0 = llvm.mlir.constant(42 : i32) : i32
+  // Do not inline load/store operations that carry attributes requiring
+  // handling while inlining, until this is supported by the inliner.
+  llvm.store %0, %ptr { access_groups = [@metadata::@group] }: i32, !llvm.ptr
+  return
+}
+
+// CHECK-LABEL: func.func @test_not_inline
+// CHECK-NEXT: call @with_mem_attr
+// CHECK-NEXT: return
+func.func @test_not_inline(%ptr : !llvm.ptr) -> () {
+  call @with_mem_attr(%ptr) : (!llvm.ptr) -> ()
+  return
 }