Skip to content

Commit 47fc4c3

Browse files
arsenmarsenm
authored
InferAddressSpaces: Handle masked load and store intrinsics (#102007)
1 parent f01a6f5 commit 47fc4c3

File tree

2 files changed

+77
-7
lines changed

2 files changed

+77
-7
lines changed

‎llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -369,13 +369,13 @@ bool InferAddressSpacesImpl::rewriteIntrinsicOperands(IntrinsicInst *II,
369369
Value *OldV,
370370
Value *NewV) const {
371371
Module *M = II->getParent()->getParent()->getParent();
372-
373-
switch (II->getIntrinsicID()) {
374-
case Intrinsic::objectsize: {
372+
Intrinsic::ID IID = II->getIntrinsicID();
373+
switch (IID) {
374+
case Intrinsic::objectsize:
375+
case Intrinsic::masked_load: {
375376
Type *DestTy = II->getType();
376377
Type *SrcTy = NewV->getType();
377-
Function *NewDecl =
378-
Intrinsic::getDeclaration(M, II->getIntrinsicID(), {DestTy, SrcTy});
378+
Function *NewDecl = Intrinsic::getDeclaration(M, IID, {DestTy, SrcTy});
379379
II->setArgOperand(0, NewV);
380380
II->setCalledFunction(NewDecl);
381381
return true;
@@ -386,12 +386,12 @@ bool InferAddressSpacesImpl::rewriteIntrinsicOperands(IntrinsicInst *II,
386386
case Intrinsic::masked_gather: {
387387
Type *RetTy = II->getType();
388388
Type *NewPtrTy = NewV->getType();
389-
Function *NewDecl =
390-
Intrinsic::getDeclaration(M, II->getIntrinsicID(), {RetTy, NewPtrTy});
389+
Function *NewDecl = Intrinsic::getDeclaration(M, IID, {RetTy, NewPtrTy});
391390
II->setArgOperand(0, NewV);
392391
II->setCalledFunction(NewDecl);
393392
return true;
394393
}
394+
case Intrinsic::masked_store:
395395
case Intrinsic::masked_scatter: {
396396
Type *ValueTy = II->getOperand(0)->getType();
397397
Type *NewPtrTy = NewV->getType();
@@ -429,11 +429,13 @@ void InferAddressSpacesImpl::collectRewritableIntrinsicOperands(
429429
appendsFlatAddressExpressionToPostorderStack(II->getArgOperand(0),
430430
PostorderStack, Visited);
431431
break;
432+
case Intrinsic::masked_load:
432433
case Intrinsic::masked_gather:
433434
case Intrinsic::prefetch:
434435
appendsFlatAddressExpressionToPostorderStack(II->getArgOperand(0),
435436
PostorderStack, Visited);
436437
break;
438+
case Intrinsic::masked_store:
437439
case Intrinsic::masked_scatter:
438440
appendsFlatAddressExpressionToPostorderStack(II->getArgOperand(1),
439441
PostorderStack, Visited);

‎llvm/test/Transforms/InferAddressSpaces/AMDGPU/masked-load-store.ll

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,68 @@
1+
; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
2+
; RUN: opt -S -mtriple=amdgcn-amd-amdhsa -passes=infer-address-spaces %s | FileCheck %s
3+
4+
define <32 x i32> @masked_load_v32i32_global_to_flat(ptr addrspace(1) %ptr, <32 x i1> %mask) {
5+
; CHECK-LABEL: define <32 x i32> @masked_load_v32i32_global_to_flat(
6+
; CHECK-SAME: ptr addrspace(1) [[PTR:%.*]], <32 x i1> [[MASK:%.*]]) {
7+
; CHECK-NEXT: [[LOAD:%.*]] = call <32 x i32> @llvm.masked.load.v32i32.p1(ptr addrspace(1) [[PTR]], i32 8, <32 x i1> [[MASK]], <32 x i32> zeroinitializer)
8+
; CHECK-NEXT: ret <32 x i32> [[LOAD]]
9+
;
10+
%cast = addrspacecast ptr addrspace(1) %ptr to ptr
11+
%load = call <32 x i32> @llvm.masked.load.v32i32.p0(ptr %cast, i32 8, <32 x i1> %mask, <32 x i32> zeroinitializer)
12+
ret <32 x i32> %load
13+
}
14+
define <32 x i32> @masked_load_v32i32_local_to_flat(ptr addrspace(3) %ptr, <32 x i1> %mask) {
15+
; CHECK-LABEL: define <32 x i32> @masked_load_v32i32_local_to_flat(
16+
; CHECK-SAME: ptr addrspace(3) [[PTR:%.*]], <32 x i1> [[MASK:%.*]]) {
17+
; CHECK-NEXT: [[LOAD:%.*]] = call <32 x i32> @llvm.masked.load.v32i32.p3(ptr addrspace(3) [[PTR]], i32 8, <32 x i1> [[MASK]], <32 x i32> zeroinitializer)
18+
; CHECK-NEXT: ret <32 x i32> [[LOAD]]
19+
;
20+
%cast = addrspacecast ptr addrspace(3) %ptr to ptr
21+
%load = call <32 x i32> @llvm.masked.load.v32i32.p0(ptr %cast, i32 8, <32 x i1> %mask, <32 x i32> zeroinitializer)
22+
ret <32 x i32> %load
23+
}
24+
25+
define <32 x i32> @masked_load_v32i32_private_to_flat(ptr addrspace(5) %ptr, <32 x i1> %mask) {
26+
; CHECK-LABEL: define <32 x i32> @masked_load_v32i32_private_to_flat(
27+
; CHECK-SAME: ptr addrspace(5) [[PTR:%.*]], <32 x i1> [[MASK:%.*]]) {
28+
; CHECK-NEXT: [[LOAD:%.*]] = call <32 x i32> @llvm.masked.load.v32i32.p5(ptr addrspace(5) [[PTR]], i32 8, <32 x i1> [[MASK]], <32 x i32> zeroinitializer)
29+
; CHECK-NEXT: ret <32 x i32> [[LOAD]]
30+
;
31+
%cast = addrspacecast ptr addrspace(5) %ptr to ptr
32+
%load = call <32 x i32> @llvm.masked.load.v32i32.p0(ptr %cast, i32 8, <32 x i1> %mask, <32 x i32> zeroinitializer)
33+
ret <32 x i32> %load
34+
}
35+
36+
define void @masked_store_v32i32_global_to_flat(ptr addrspace(1) %ptr, <32 x i1> %mask) {
37+
; CHECK-LABEL: define void @masked_store_v32i32_global_to_flat(
38+
; CHECK-SAME: ptr addrspace(1) [[PTR:%.*]], <32 x i1> [[MASK:%.*]]) {
39+
; CHECK-NEXT: tail call void @llvm.masked.store.v32i32.p1(<32 x i32> zeroinitializer, ptr addrspace(1) [[PTR]], i32 128, <32 x i1> [[MASK]])
40+
; CHECK-NEXT: ret void
41+
;
42+
%cast = addrspacecast ptr addrspace(1) %ptr to ptr
43+
tail call void @llvm.masked.store.v32i32.p0(<32 x i32> zeroinitializer, ptr %cast, i32 128, <32 x i1> %mask)
44+
ret void
45+
}
46+
47+
define void @masked_store_v32i32_local_to_flat(ptr addrspace(3) %ptr, <32 x i1> %mask) {
48+
; CHECK-LABEL: define void @masked_store_v32i32_local_to_flat(
49+
; CHECK-SAME: ptr addrspace(3) [[PTR:%.*]], <32 x i1> [[MASK:%.*]]) {
50+
; CHECK-NEXT: tail call void @llvm.masked.store.v32i32.p3(<32 x i32> zeroinitializer, ptr addrspace(3) [[PTR]], i32 128, <32 x i1> [[MASK]])
51+
; CHECK-NEXT: ret void
52+
;
53+
%cast = addrspacecast ptr addrspace(3) %ptr to ptr
54+
tail call void @llvm.masked.store.v32i32.p0(<32 x i32> zeroinitializer, ptr %cast, i32 128, <32 x i1> %mask)
55+
ret void
56+
}
57+
58+
define void @masked_store_v32i32_private_to_flat(ptr addrspace(5) %ptr, <32 x i1> %mask) {
59+
; CHECK-LABEL: define void @masked_store_v32i32_private_to_flat(
60+
; CHECK-SAME: ptr addrspace(5) [[PTR:%.*]], <32 x i1> [[MASK:%.*]]) {
61+
; CHECK-NEXT: tail call void @llvm.masked.store.v32i32.p5(<32 x i32> zeroinitializer, ptr addrspace(5) [[PTR]], i32 128, <32 x i1> [[MASK]])
62+
; CHECK-NEXT: ret void
63+
;
64+
%cast = addrspacecast ptr addrspace(5) %ptr to ptr
65+
tail call void @llvm.masked.store.v32i32.p0(<32 x i32> zeroinitializer, ptr %cast, i32 128, <32 x i1> %mask)
66+
ret void
67+
}
68+

0 commit comments

Comments
 (0)