[PATCH] D150070: [InferAddressSpaces] Support intrinsic masked gather & masked scatter

[CaprYang via Phabricator via llvm-commits]

CaprYang created this revision.
Herald added a subscriber: hiraditya.
Herald added a project: All.
CaprYang requested review of this revision.
Herald added a project: LLVM.
Herald added a subscriber: llvm-commits.

Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D150070

Files:
  llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp
  llvm/test/Transforms/InferAddressSpaces/masked-gather-scatter.ll


Index: llvm/test/Transforms/InferAddressSpaces/masked-gather-scatter.ll
===================================================================
--- /dev/null
+++ llvm/test/Transforms/InferAddressSpaces/masked-gather-scatter.ll
@@ -0,0 +1,25 @@
+; RUN: opt -S -passes=infer-address-spaces -assume-default-is-flat-addrspace %s | FileCheck %s
+
+; CHECK-LABEL: @masked_gather_inferas(
+; CHECK: tail call <4 x i32> @llvm.masked.gather.v4i32.v4p1
+define <4 x i32> @masked_gather_inferas(ptr addrspace(1) %out, <4 x i64> %index) {
+entry:
+  %out.1 = addrspacecast ptr addrspace(1) %out to ptr
+  %ptrs = getelementptr inbounds i32, ptr %out.1, <4 x i64> %index
+  %value = tail call <4 x i32> @llvm.masked.gather.v4i32.v4p0(<4 x ptr> %ptrs, i32 4, <4 x i1> <i1 true, i1 true, i1 true, i1 true>, <4 x i32> poison)
+  ret <4 x i32> %value
+}
+
+; CHECK-LABEL: @masked_scatter_inferas(
+; CHECK: tail call void @llvm.masked.scatter.v4i32.v4p1
+define void @masked_scatter_inferas(ptr addrspace(1) %out, <4 x i64> %index, <4 x i32> %value) {
+entry:
+  %out.1 = addrspacecast ptr addrspace(1) %out to ptr
+  %ptrs = getelementptr inbounds i32, ptr %out.1, <4 x i64> %index
+  tail call void @llvm.masked.scatter.v4i32.v4p0(<4 x i32> %value, <4 x ptr> %ptrs, i32 4, <4 x i1> <i1 true, i1 true, i1 true, i1 true>)
+  ret void
+}
+
+declare <4 x i32> @llvm.masked.gather.v4i32.v4p0(<4 x ptr>, i32 immarg, <4 x i1>, <4 x i32>)
+
+declare void @llvm.masked.scatter.v4i32.v4p0(<4 x i32>, <4 x ptr>, i32 immarg, <4 x i1>)
\ No newline at end of file
Index: llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp
===================================================================
--- llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp
+++ llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp
@@ -416,6 +416,24 @@
   case Intrinsic::ptrmask:
     // This is handled as an address expression, not as a use memory operation.
     return false;
+  case Intrinsic::masked_gather: {
+    Type *RetTy = II->getType();
+    Type *NewPtrTy = NewV->getType();
+    Function *NewDecl =
+        Intrinsic::getDeclaration(M, II->getIntrinsicID(), {RetTy, NewPtrTy});
+    II->setArgOperand(0, NewV);
+    II->setCalledFunction(NewDecl);
+    return true;
+  }
+  case Intrinsic::masked_scatter: {
+    Type *ValueTy = II->getOperand(0)->getType();
+    Type *NewPtrTy = NewV->getType();
+    Function *NewDecl =
+        Intrinsic::getDeclaration(M, II->getIntrinsicID(), {ValueTy, NewPtrTy});
+    II->setArgOperand(1, NewV);
+    II->setCalledFunction(NewDecl);
+    return true;
+  }
   default: {
     Value *Rewrite = TTI->rewriteIntrinsicWithAddressSpace(II, OldV, NewV);
     if (!Rewrite)
@@ -437,6 +455,14 @@
     appendsFlatAddressExpressionToPostorderStack(II->getArgOperand(0),
                                                  PostorderStack, Visited);
     break;
+  case Intrinsic::masked_gather:
+    appendsFlatAddressExpressionToPostorderStack(II->getArgOperand(0),
+                                                 PostorderStack, Visited);
+    break;
+  case Intrinsic::masked_scatter:
+    appendsFlatAddressExpressionToPostorderStack(II->getArgOperand(1),
+                                                 PostorderStack, Visited);
+    break;
   default:
     SmallVector<int, 2> OpIndexes;
     if (TTI->collectFlatAddressOperands(OpIndexes, IID)) {


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D150070.520191.patch
Type: text/x-patch
Size: 3336 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20230507/cf639a1b/attachment-0001.bin>