From 07ceb0f035dd94a847969dcd81d73775efc50901 Mon Sep 17 00:00:00 2001
From: Sirui Mu <msrlancern@gmail.com>
Date: Tue, 12 Nov 2024 22:58:01 +0800
Subject: [PATCH] [mlir][LLVM] Add support for invariant group related
 intrinsics

This patch adds support for the following LLVM intrinsics:
  - `llvm.launder.invariant.group`
  - `llvm.strip.invariant.group`
---
 .../mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td   | 18 ++++++++++++++
 mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp | 24 +++++++++++++++++++
 mlir/test/Dialect/LLVMIR/mem2reg.mlir         | 12 ++++++++++
 mlir/test/Dialect/LLVMIR/roundtrip.mlir       | 10 ++++++++
 mlir/test/Target/LLVMIR/Import/intrinsic.ll   | 11 +++++++++
 .../test/Target/LLVMIR/llvmir-intrinsics.mlir |  9 +++++++
 6 files changed, 84 insertions(+)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td
index d07ebbacc6043..b0d6148726f61 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td
@@ -320,6 +320,24 @@ def LLVM_InvariantEndOp : LLVM_ZeroResultIntrOp<"invariant.end", [2],
       "qualified(type($ptr))";
 }
 
+def LLVM_LaunderInvariantGroupOp
+    : LLVM_OneResultIntrOp<"launder.invariant.group", [], [0],
+        [DeclareOpInterfaceMethods<PromotableOpInterface>,
+         SameOperandsAndResultType]> {
+  let arguments = (ins LLVM_AnyPointer:$ptr);
+  let results = (outs LLVM_AnyPointer:$res);
+  let assemblyFormat = "$ptr attr-dict `:` qualified(type($ptr))";
+}
+
+def LLVM_StripInvariantGroupOp
+    : LLVM_OneResultIntrOp<"strip.invariant.group", [], [0],
+        [DeclareOpInterfaceMethods<PromotableOpInterface>,
+         SameOperandsAndResultType]> {
+  let arguments = (ins LLVM_AnyPointer:$ptr);
+  let results = (outs LLVM_AnyPointer:$res);
+  let assemblyFormat = "$ptr attr-dict `:` qualified(type($ptr))";
+}
+
 // Constrained Floating-Point Intrinsics.
 
 class LLVM_ConstrainedIntr<string mnem, int numArgs,
diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp
index 5dc506c14ef96..dc68e51341851 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp
@@ -590,6 +590,30 @@ DeletionKind LLVM::InvariantEndOp::removeBlockingUses(
   return DeletionKind::Delete;
 }
 
+bool LLVM::LaunderInvariantGroupOp::canUsesBeRemoved(
+    const SmallPtrSetImpl<OpOperand *> &blockingUses,
+    SmallVectorImpl<OpOperand *> &newBlockingUses,
+    const DataLayout &dataLayout) {
+  return forwardToUsers(*this, newBlockingUses);
+}
+
+DeletionKind LLVM::LaunderInvariantGroupOp::removeBlockingUses(
+    const SmallPtrSetImpl<OpOperand *> &blockingUses, OpBuilder &builder) {
+  return DeletionKind::Delete;
+}
+
+bool LLVM::StripInvariantGroupOp::canUsesBeRemoved(
+    const SmallPtrSetImpl<OpOperand *> &blockingUses,
+    SmallVectorImpl<OpOperand *> &newBlockingUses,
+    const DataLayout &dataLayout) {
+  return forwardToUsers(*this, newBlockingUses);
+}
+
+DeletionKind LLVM::StripInvariantGroupOp::removeBlockingUses(
+    const SmallPtrSetImpl<OpOperand *> &blockingUses, OpBuilder &builder) {
+  return DeletionKind::Delete;
+}
+
 bool LLVM::DbgDeclareOp::canUsesBeRemoved(
     const SmallPtrSetImpl<OpOperand *> &blockingUses,
     SmallVectorImpl<OpOperand *> &newBlockingUses,
diff --git a/mlir/test/Dialect/LLVMIR/mem2reg.mlir b/mlir/test/Dialect/LLVMIR/mem2reg.mlir
index 38c836c139da6..e66b498ed4fa0 100644
--- a/mlir/test/Dialect/LLVMIR/mem2reg.mlir
+++ b/mlir/test/Dialect/LLVMIR/mem2reg.mlir
@@ -414,6 +414,18 @@ llvm.func @ignore_lifetime() {
 
 // -----
 
+// CHECK-LABEL: llvm.func @ignore_invariant_group
+// CHECK-NOT: llvm.alloca
+llvm.func @ignore_invariant_group() {
+  %0 = llvm.mlir.constant(1 : i32) : i32
+  %1 = llvm.alloca %0 x i32 {alignment = 4 : i64} : (i32) -> !llvm.ptr
+  %2 = llvm.intr.launder.invariant.group %1 : !llvm.ptr
+  %3 = llvm.intr.strip.invariant.group %2 : !llvm.ptr
+  llvm.return
+}
+
+// -----
+
 // CHECK-LABEL: llvm.func @ignore_discardable_tree
 // CHECK-NOT: = llvm.alloca
 llvm.func @ignore_discardable_tree() {
diff --git a/mlir/test/Dialect/LLVMIR/roundtrip.mlir b/mlir/test/Dialect/LLVMIR/roundtrip.mlir
index aa558bad2299c..f8cd23dc4e723 100644
--- a/mlir/test/Dialect/LLVMIR/roundtrip.mlir
+++ b/mlir/test/Dialect/LLVMIR/roundtrip.mlir
@@ -672,6 +672,16 @@ llvm.func @invariant(%p: !llvm.ptr) {
   llvm.return
 }
 
+// CHECK-LABEL: @invariant_group_intrinsics
+// CHECK-SAME: %[[P:.+]]: !llvm.ptr
+llvm.func @invariant_group_intrinsics(%p: !llvm.ptr) {
+  // CHECK: %{{.+}} = llvm.intr.launder.invariant.group %[[P]] : !llvm.ptr
+  %1 = llvm.intr.launder.invariant.group %p : !llvm.ptr
+  // CHECK: %{{.+}} = llvm.intr.strip.invariant.group %[[P]] : !llvm.ptr
+  %2 = llvm.intr.strip.invariant.group %p : !llvm.ptr
+  llvm.return
+}
+
 // CHECK-LABEL: @vararg_func
 llvm.func @vararg_func(%arg0: i32, ...) {
   // CHECK: %[[C:.*]] = llvm.mlir.constant(1 : i32)
diff --git a/mlir/test/Target/LLVMIR/Import/intrinsic.ll b/mlir/test/Target/LLVMIR/Import/intrinsic.ll
index 606b11175f572..64911b09ebc85 100644
--- a/mlir/test/Target/LLVMIR/Import/intrinsic.ll
+++ b/mlir/test/Target/LLVMIR/Import/intrinsic.ll
@@ -802,6 +802,15 @@ define void @invariant(ptr %0) {
   ret void
 }
 
+; CHECK-LABEL: llvm.func @invariant_group
+define void @invariant_group(ptr %0) {
+  ; CHECK: %{{.+}} = llvm.intr.launder.invariant.group %{{.*}} : !llvm.ptr
+  %2 = call ptr @llvm.launder.invariant.group.p0(ptr %0)
+  ; CHECK: %{{.+}} = llvm.intr.strip.invariant.group %{{.*}} : !llvm.ptr
+  %3 = call ptr @llvm.strip.invariant.group.p0(ptr %0)
+  ret void
+}
+
 ; CHECK-LABEL: llvm.func @vector_insert
 define void @vector_insert(<vscale x 4 x float> %0, <4 x float> %1) {
   ; CHECK: llvm.intr.vector.insert %{{.*}}, %{{.*}}[4] : vector<4xf32> into !llvm.vec<? x 4 x  f32>
@@ -1191,6 +1200,8 @@ declare void @llvm.lifetime.start.p0(i64 immarg, ptr nocapture)
 declare void @llvm.lifetime.end.p0(i64 immarg, ptr nocapture)
 declare ptr @llvm.invariant.start.p0(i64 immarg, ptr nocapture)
 declare void @llvm.invariant.end.p0(ptr, i64 immarg, ptr nocapture)
+declare ptr @llvm.launder.invariant.group.p0(ptr nocapture)
+declare ptr @llvm.strip.invariant.group.p0(ptr nocapture)
 
 declare void @llvm.assume(i1)
 declare float @llvm.ssa.copy.f32(float returned)
diff --git a/mlir/test/Target/LLVMIR/llvmir-intrinsics.mlir b/mlir/test/Target/LLVMIR/llvmir-intrinsics.mlir
index cb712eb4e1262..059abed34cd5d 100644
--- a/mlir/test/Target/LLVMIR/llvmir-intrinsics.mlir
+++ b/mlir/test/Target/LLVMIR/llvmir-intrinsics.mlir
@@ -1020,6 +1020,15 @@ llvm.func @invariant(%p: !llvm.ptr) {
   llvm.return
 }
 
+// CHECK-LABEL: @invariant_group
+llvm.func @invariant_group(%p: !llvm.ptr) {
+  // CHECK: call ptr @llvm.launder.invariant.group
+  %1 = llvm.intr.launder.invariant.group %p : !llvm.ptr
+  // CHECK: call ptr @llvm.strip.invariant.group
+  %2 = llvm.intr.strip.invariant.group %p : !llvm.ptr
+  llvm.return
+}
+
 // CHECK-LABEL: @ssa_copy
 llvm.func @ssa_copy(%arg: f32) -> f32 {
   // CHECK: call float @llvm.ssa.copy