[Mlir-commits] [flang] [mlir] [mlir][LLVMIR] Set memory effects for llvm.intrin.stacksave/restore. (PR #191918)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-llvm

Author: Slava Zakharin (vzakhari)

<details>
<summary>Changes</summary>

These operations are too conservative currently.
Adding effects on `AutomaticAllocationScopeResource` improves
Flang LICM.


---
Full diff: https://github.com/llvm/llvm-project/pull/191918.diff


3 Files Affected:

- (added) flang/test/Transforms/licm-stacksaverestore.mlir (+35) 
- (modified) mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td (+27-2) 
- (added) mlir/test/Dialect/LLVMIR/side-effects.mlir (+13) 


``````````diff
diff --git a/flang/test/Transforms/licm-stacksaverestore.mlir b/flang/test/Transforms/licm-stacksaverestore.mlir
new file mode 100644
index 0000000000000..e408ee9b09085
--- /dev/null
+++ b/flang/test/Transforms/licm-stacksaverestore.mlir
@@ -0,0 +1,35 @@
+// RUN: fir-opt -flang-licm --split-input-file %s | FileCheck %s
+
+// Verify that an invariant load is hoisted out of the loop
+// when stacksave/stackrestore are present.
+//
+// Original code:
+// subroutine test(x,y)
+//   real :: x(10), y
+//   do i=1,10
+//      block
+//        x(i) = y
+//      end block
+//   end do
+// end subroutine test
+// CHECK-LABEL:   func.func @_QPtest(
+// CHECK:           %[[DECLARE_1:.*]] = fir.declare %{{.*}} dummy_scope %{{.*}} arg 2 {uniq_name = "_QFtestEy"} : (!fir.ref<f32>, !fir.dscope) -> !fir.ref<f32>
+// CHECK:           %[[LOAD_0:.*]] = fir.load %[[DECLARE_1]] : !fir.ref<f32>
+// CHECK:           fir.do_loop
+// CHECK-NOT: fir.load
+func.func @_QPtest(%arg0: !fir.ref<!fir.array<10xf32>> {fir.bindc_name = "x"}, %arg1: !fir.ref<f32> {fir.bindc_name = "y"}) {
+  %c1 = arith.constant 1 : index
+  %c10 = arith.constant 10 : index
+  %0 = fir.dummy_scope : !fir.dscope
+  %3 = fir.shape %c10 : (index) -> !fir.shape<1>
+  %4 = fir.declare %arg0(%3) dummy_scope %0 arg 1 {uniq_name = "_QFtestEx"} : (!fir.ref<!fir.array<10xf32>>, !fir.shape<1>, !fir.dscope) -> !fir.ref<!fir.array<10xf32>>
+  %5 = fir.declare %arg1 dummy_scope %0 arg 2 {uniq_name = "_QFtestEy"} : (!fir.ref<f32>, !fir.dscope) -> !fir.ref<f32>
+  fir.do_loop %arg2 = %c1 to %c10 step %c1 {
+    %8 = llvm.intr.stacksave : !llvm.ptr
+    %9 = fir.load %5 : !fir.ref<f32>
+    %12 = fir.array_coor %4(%3) %arg2 : (!fir.ref<!fir.array<10xf32>>, !fir.shape<1>, index) -> !fir.ref<f32>
+    fir.store %9 to %12 : !fir.ref<f32>
+    llvm.intr.stackrestore %8 : !llvm.ptr
+  }
+  return
+}
diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td
index 688bc19cbf18a..7945cf9fae71f 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td
@@ -893,11 +893,36 @@ def LLVM_EhTypeidForOp : LLVM_OneResultIntrOp<"eh.typeid.for", [], [0]> {
 // Stack save/restore intrinsics.
 //
 
-def LLVM_StackSaveOp : LLVM_OneResultIntrOp<"stacksave", [0]> {
+// Technically, llvm.intrin.stacksave does not free the stack,
+// but MemFree<AutomaticAllocationScopeResource> seems to be
+// an appropriate effect to guarantee that the allocations
+// between stacksave/stackrestore cannot be moved out of
+// the stacksave/stackrestore "scope".
+// For example,
+//   %ptr = llvm.intrin.stacksave : !llvm.ptr
+//   %alloca = llvm.alloca %zero x i32 : (i32) -> !llvm.ptr
+//   ... = llvm.load %alloca : !llvm.ptr -> i32
+//   llvm.store %x, %alloca : i32, !llvm.ptr
+//   llvm.intrin.stackrestore %ptr : !llvm.ptr
+//
+// The 'llvm.alloca' should not be movable neither before
+// llvm.intrin.stacksave nor after llvm.intrin.stackrestore.
+// Moreover, llvm.intrin.stackrestore cannot be moved
+// above the load the the store.
+//
+// The 'Free' affect on 'AutomaticAllocationScopeResource'
+// for both stacksave and stackrestore should guarantee that.
+def LLVM_StackSaveOp
+    : LLVM_OneResultIntrOp<
+          "stacksave", [0], [],
+          [MemoryEffects<[MemFree<AutomaticAllocationScopeResource>]>]> {
   let assemblyFormat = "attr-dict `:` qualified(type($res))";
 }
 
-def LLVM_StackRestoreOp : LLVM_ZeroResultIntrOp<"stackrestore", [0]> {
+def LLVM_StackRestoreOp
+    : LLVM_ZeroResultIntrOp<
+          "stackrestore", [0],
+          [MemoryEffects<[MemFree<AutomaticAllocationScopeResource>]>]> {
   let arguments = (ins LLVM_AnyPointer:$ptr);
   let assemblyFormat = "$ptr attr-dict `:` qualified(type($ptr))";
 }
diff --git a/mlir/test/Dialect/LLVMIR/side-effects.mlir b/mlir/test/Dialect/LLVMIR/side-effects.mlir
new file mode 100644
index 0000000000000..b3e10a95d2868
--- /dev/null
+++ b/mlir/test/Dialect/LLVMIR/side-effects.mlir
@@ -0,0 +1,13 @@
+// RUN: mlir-opt %s --test-side-effects --verify-diagnostics
+
+func.func @test_stacksave_stackrestore() {
+  // expected-remark @below {{found an instance of 'free' on resource 'AutomaticAllocationScope'}}
+  %ptr = llvm.intr.stacksave : !llvm.ptr
+  // expected-remark @below {{operation has no memory effects}}
+  %c1 = llvm.mlir.constant(1 : i32) : i32
+  // expected-remark @below {{found an instance of 'allocate' on op result 0, on resource 'AutomaticAllocationScope'}}
+  %alloca = llvm.alloca %c1 x i32 : (i32) -> !llvm.ptr
+  // expected-remark @below {{found an instance of 'free' on resource 'AutomaticAllocationScope'}}
+  llvm.intr.stackrestore %ptr : !llvm.ptr
+  llvm.return
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/191918