Skip to content

Pre-commit AMDGPU tests for masked load/store/scatter/gather #104645

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Aug 19, 2024
Merged

Pre-commit AMDGPU tests for masked load/store/scatter/gather #104645

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
084c02d
Pre-commit AMDGPU tests for masked load/store/scatter/gather
krzysz00 Aug 16, 2024
d775b76
Add a few more tests
krzysz00 Aug 19, 2024
807dbe3
Uniform value tests
krzysz00 Aug 19, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
170 changes: 170 additions & 0 deletions llvm/test/Transforms/ScalarizeMaskedMemIntrin/AMDGPU/expamd-masked-load.ll
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,170 @@
; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
; RUN: opt -S %s -passes=scalarize-masked-mem-intrin -mtriple=amdgcn-amd-amdhsa | FileCheck %s

; COM: Test that, unlike on CPU targets, the mask doesn't get bitcast to a scalar,
; COM: since, on GPUs, each i1 takes up at least one register and so they should
; COM: be treated separately.

define <2 x i32> @scalarize_v2i32(ptr %p, <2 x i1> %mask, <2 x i32> %passthru) {
; CHECK-LABEL: define <2 x i32> @scalarize_v2i32(
; CHECK-SAME: ptr [[P:%.*]], <2 x i1> [[MASK:%.*]], <2 x i32> [[PASSTHRU:%.*]]) {
; CHECK-NEXT: [[SCALAR_MASK:%.*]] = bitcast <2 x i1> [[MASK]] to i2
; CHECK-NEXT: [[TMP1:%.*]] = and i2 [[SCALAR_MASK]], 1
; CHECK-NEXT: [[TMP2:%.*]] = icmp ne i2 [[TMP1]], 0
; CHECK-NEXT: br i1 [[TMP2]], label %[[COND_LOAD:.*]], label %[[ELSE:.*]]
; CHECK: [[COND_LOAD]]:
; CHECK-NEXT: [[TMP3:%.*]] = getelementptr inbounds i32, ptr [[P]], i32 0
; CHECK-NEXT: [[TMP4:%.*]] = load i32, ptr [[TMP3]], align 4
; CHECK-NEXT: [[TMP5:%.*]] = insertelement <2 x i32> [[PASSTHRU]], i32 [[TMP4]], i64 0
; CHECK-NEXT: br label %[[ELSE]]
; CHECK: [[ELSE]]:
; CHECK-NEXT: [[RES_PHI_ELSE:%.*]] = phi <2 x i32> [ [[TMP5]], %[[COND_LOAD]] ], [ [[PASSTHRU]], [[TMP0:%.*]] ]
; CHECK-NEXT: [[TMP6:%.*]] = and i2 [[SCALAR_MASK]], -2
; CHECK-NEXT: [[TMP7:%.*]] = icmp ne i2 [[TMP6]], 0
; CHECK-NEXT: br i1 [[TMP7]], label %[[COND_LOAD1:.*]], label %[[ELSE2:.*]]
; CHECK: [[COND_LOAD1]]:
; CHECK-NEXT: [[TMP8:%.*]] = getelementptr inbounds i32, ptr [[P]], i32 1
; CHECK-NEXT: [[TMP9:%.*]] = load i32, ptr [[TMP8]], align 4
; CHECK-NEXT: [[TMP10:%.*]] = insertelement <2 x i32> [[RES_PHI_ELSE]], i32 [[TMP9]], i64 1
; CHECK-NEXT: br label %[[ELSE2]]
; CHECK: [[ELSE2]]:
; CHECK-NEXT: [[RES_PHI_ELSE3:%.*]] = phi <2 x i32> [ [[TMP10]], %[[COND_LOAD1]] ], [ [[RES_PHI_ELSE]], %[[ELSE]] ]
; CHECK-NEXT: ret <2 x i32> [[RES_PHI_ELSE3]]
;
%ret = call <2 x i32> @llvm.masked.load.v2i32.p0(ptr %p, i32 128, <2 x i1> %mask, <2 x i32> %passthru)
ret <2 x i32> %ret
}

define <2 x i32> @scalarize_v2i32_splat_mask(ptr %p, i1 %mask, <2 x i32> %passthrough) {
; CHECK-LABEL: define <2 x i32> @scalarize_v2i32_splat_mask(
; CHECK-SAME: ptr [[P:%.*]], i1 [[MASK:%.*]], <2 x i32> [[PASSTHROUGH:%.*]]) {
; CHECK-NEXT: [[MASK_VEC:%.*]] = insertelement <2 x i1> poison, i1 [[MASK]], i32 0
; CHECK-NEXT: [[MASK_SPLAT:%.*]] = shufflevector <2 x i1> [[MASK_VEC]], <2 x i1> poison, <2 x i32> zeroinitializer
; CHECK-NEXT: [[MASK_SPLAT_FIRST:%.*]] = extractelement <2 x i1> [[MASK_SPLAT]], i64 0
; CHECK-NEXT: br i1 [[MASK_SPLAT_FIRST]], label %[[COND_LOAD:.*]], label %[[BB1:.*]]
; CHECK: [[COND_LOAD]]:
; CHECK-NEXT: [[RET_COND_LOAD:%.*]] = load <2 x i32>, ptr [[P]], align 8
; CHECK-NEXT: br label %[[BB1]]
; CHECK: [[BB1]]:
; CHECK-NEXT: [[RET:%.*]] = phi <2 x i32> [ [[RET_COND_LOAD]], %[[COND_LOAD]] ], [ [[PASSTHROUGH]], [[TMP0:%.*]] ]
; CHECK-NEXT: ret <2 x i32> [[RET]]
;
%mask.vec = insertelement <2 x i1> poison, i1 %mask, i32 0
%mask.splat = shufflevector <2 x i1> %mask.vec, <2 x i1> poison, <2 x i32> zeroinitializer
%ret = call <2 x i32> @llvm.masked.load.v2i32.p0(ptr %p, i32 8, <2 x i1> %mask.splat, <2 x i32> %passthrough)
ret <2 x i32> %ret
}

define <2 x half> @scalarize_v2f16(ptr %p, <2 x i1> %mask, <2 x half> %passthru) {
; CHECK-LABEL: define <2 x half> @scalarize_v2f16(
; CHECK-SAME: ptr [[P:%.*]], <2 x i1> [[MASK:%.*]], <2 x half> [[PASSTHRU:%.*]]) {
; CHECK-NEXT: [[SCALAR_MASK:%.*]] = bitcast <2 x i1> [[MASK]] to i2
; CHECK-NEXT: [[TMP1:%.*]] = and i2 [[SCALAR_MASK]], 1
; CHECK-NEXT: [[TMP2:%.*]] = icmp ne i2 [[TMP1]], 0
; CHECK-NEXT: br i1 [[TMP2]], label %[[COND_LOAD:.*]], label %[[ELSE:.*]]
; CHECK: [[COND_LOAD]]:
; CHECK-NEXT: [[TMP3:%.*]] = getelementptr inbounds half, ptr [[P]], i32 0
; CHECK-NEXT: [[TMP4:%.*]] = load half, ptr [[TMP3]], align 2
; CHECK-NEXT: [[TMP5:%.*]] = insertelement <2 x half> [[PASSTHRU]], half [[TMP4]], i64 0
; CHECK-NEXT: br label %[[ELSE]]
; CHECK: [[ELSE]]:
; CHECK-NEXT: [[RES_PHI_ELSE:%.*]] = phi <2 x half> [ [[TMP5]], %[[COND_LOAD]] ], [ [[PASSTHRU]], [[TMP0:%.*]] ]
; CHECK-NEXT: [[TMP6:%.*]] = and i2 [[SCALAR_MASK]], -2
; CHECK-NEXT: [[TMP7:%.*]] = icmp ne i2 [[TMP6]], 0
; CHECK-NEXT: br i1 [[TMP7]], label %[[COND_LOAD1:.*]], label %[[ELSE2:.*]]
; CHECK: [[COND_LOAD1]]:
; CHECK-NEXT: [[TMP8:%.*]] = getelementptr inbounds half, ptr [[P]], i32 1
; CHECK-NEXT: [[TMP9:%.*]] = load half, ptr [[TMP8]], align 2
; CHECK-NEXT: [[TMP10:%.*]] = insertelement <2 x half> [[RES_PHI_ELSE]], half [[TMP9]], i64 1
; CHECK-NEXT: br label %[[ELSE2]]
; CHECK: [[ELSE2]]:
; CHECK-NEXT: [[RES_PHI_ELSE3:%.*]] = phi <2 x half> [ [[TMP10]], %[[COND_LOAD1]] ], [ [[RES_PHI_ELSE]], %[[ELSE]] ]
; CHECK-NEXT: ret <2 x half> [[RES_PHI_ELSE3]]
;
%ret = call <2 x half> @llvm.masked.load.v2f16.p0(ptr %p, i32 128, <2 x i1> %mask, <2 x half> %passthru)
ret <2 x half> %ret
}

define <2 x i32> @scalarize_v2i32_p3(ptr addrspace(3) %p, <2 x i1> %mask, <2 x i32> %passthru) {
; CHECK-LABEL: define <2 x i32> @scalarize_v2i32_p3(
; CHECK-SAME: ptr addrspace(3) [[P:%.*]], <2 x i1> [[MASK:%.*]], <2 x i32> [[PASSTHRU:%.*]]) {
; CHECK-NEXT: [[SCALAR_MASK:%.*]] = bitcast <2 x i1> [[MASK]] to i2
; CHECK-NEXT: [[TMP1:%.*]] = and i2 [[SCALAR_MASK]], 1
; CHECK-NEXT: [[TMP2:%.*]] = icmp ne i2 [[TMP1]], 0
; CHECK-NEXT: br i1 [[TMP2]], label %[[COND_LOAD:.*]], label %[[ELSE:.*]]
; CHECK: [[COND_LOAD]]:
; CHECK-NEXT: [[TMP3:%.*]] = getelementptr inbounds i32, ptr addrspace(3) [[P]], i32 0
; CHECK-NEXT: [[TMP4:%.*]] = load i32, ptr addrspace(3) [[TMP3]], align 4
; CHECK-NEXT: [[TMP5:%.*]] = insertelement <2 x i32> [[PASSTHRU]], i32 [[TMP4]], i64 0
; CHECK-NEXT: br label %[[ELSE]]
; CHECK: [[ELSE]]:
; CHECK-NEXT: [[RES_PHI_ELSE:%.*]] = phi <2 x i32> [ [[TMP5]], %[[COND_LOAD]] ], [ [[PASSTHRU]], [[TMP0:%.*]] ]
; CHECK-NEXT: [[TMP6:%.*]] = and i2 [[SCALAR_MASK]], -2
; CHECK-NEXT: [[TMP7:%.*]] = icmp ne i2 [[TMP6]], 0
; CHECK-NEXT: br i1 [[TMP7]], label %[[COND_LOAD1:.*]], label %[[ELSE2:.*]]
; CHECK: [[COND_LOAD1]]:
; CHECK-NEXT: [[TMP8:%.*]] = getelementptr inbounds i32, ptr addrspace(3) [[P]], i32 1
; CHECK-NEXT: [[TMP9:%.*]] = load i32, ptr addrspace(3) [[TMP8]], align 4
; CHECK-NEXT: [[TMP10:%.*]] = insertelement <2 x i32> [[RES_PHI_ELSE]], i32 [[TMP9]], i64 1
; CHECK-NEXT: br label %[[ELSE2]]
; CHECK: [[ELSE2]]:
; CHECK-NEXT: [[RES_PHI_ELSE3:%.*]] = phi <2 x i32> [ [[TMP10]], %[[COND_LOAD1]] ], [ [[RES_PHI_ELSE]], %[[ELSE]] ]
; CHECK-NEXT: ret <2 x i32> [[RES_PHI_ELSE3]]
;
%ret = call <2 x i32> @llvm.masked.load.v2i32.p3(ptr addrspace(3) %p, i32 128, <2 x i1> %mask, <2 x i32> %passthru)
ret <2 x i32> %ret
}

define <2 x i32> @scalarize_v2i32_lane_mask(ptr %p, <2 x i32> %passthrough) {
; CHECK-LABEL: define <2 x i32> @scalarize_v2i32_lane_mask(
; CHECK-SAME: ptr [[P:%.*]], <2 x i32> [[PASSTHROUGH:%.*]]) {
; CHECK-NEXT: [[ITEM_ID:%.*]] = call i32 @llvm.amdgcn.workitem.id.x()
; CHECK-NEXT: [[MASK:%.*]] = icmp ult i32 [[ITEM_ID]], 16
; CHECK-NEXT: [[MASK_VEC:%.*]] = insertelement <2 x i1> poison, i1 [[MASK]], i32 0
; CHECK-NEXT: [[MASK_SPLAT:%.*]] = shufflevector <2 x i1> [[MASK_VEC]], <2 x i1> poison, <2 x i32> zeroinitializer
; CHECK-NEXT: [[MASK_SPLAT_FIRST:%.*]] = extractelement <2 x i1> [[MASK_SPLAT]], i64 0
; CHECK-NEXT: br i1 [[MASK_SPLAT_FIRST]], label %[[COND_LOAD:.*]], label %[[BB1:.*]]
; CHECK: [[COND_LOAD]]:
; CHECK-NEXT: [[RET_COND_LOAD:%.*]] = load <2 x i32>, ptr [[P]], align 8
; CHECK-NEXT: br label %[[BB1]]
; CHECK: [[BB1]]:
; CHECK-NEXT: [[RET:%.*]] = phi <2 x i32> [ [[RET_COND_LOAD]], %[[COND_LOAD]] ], [ [[PASSTHROUGH]], [[TMP0:%.*]] ]
; CHECK-NEXT: ret <2 x i32> [[RET]]
;
%item.id = call i32 @llvm.amdgcn.workitem.id.x()
%mask = icmp ult i32 %item.id, 16
%mask.vec = insertelement <2 x i1> poison, i1 %mask, i32 0
%mask.splat = shufflevector <2 x i1> %mask.vec, <2 x i1> poison, <2 x i32> zeroinitializer
%ret = call <2 x i32> @llvm.masked.load.v2i32.p0(ptr %p, i32 8, <2 x i1> %mask.splat, <2 x i32> %passthrough)
ret <2 x i32> %ret
}

define <2 x i32> @scalarize_v2i32_group_mask(ptr %p, <2 x i32> %passthrough) {
; CHECK-LABEL: define <2 x i32> @scalarize_v2i32_group_mask(
; CHECK-SAME: ptr [[P:%.*]], <2 x i32> [[PASSTHROUGH:%.*]]) {
; CHECK-NEXT: [[GROUP_ID:%.*]] = call i32 @llvm.amdgcn.workgroup.id.x()
; CHECK-NEXT: [[MASK:%.*]] = icmp ult i32 [[GROUP_ID]], 4
; CHECK-NEXT: [[MASK_VEC:%.*]] = insertelement <2 x i1> poison, i1 [[MASK]], i32 0
; CHECK-NEXT: [[MASK_SPLAT:%.*]] = shufflevector <2 x i1> [[MASK_VEC]], <2 x i1> poison, <2 x i32> zeroinitializer
; CHECK-NEXT: [[MASK_SPLAT_FIRST:%.*]] = extractelement <2 x i1> [[MASK_SPLAT]], i64 0
; CHECK-NEXT: br i1 [[MASK_SPLAT_FIRST]], label %[[COND_LOAD:.*]], label %[[BB1:.*]]
; CHECK: [[COND_LOAD]]:
; CHECK-NEXT: [[RET_COND_LOAD:%.*]] = load <2 x i32>, ptr [[P]], align 8
; CHECK-NEXT: br label %[[BB1]]
; CHECK: [[BB1]]:
; CHECK-NEXT: [[RET:%.*]] = phi <2 x i32> [ [[RET_COND_LOAD]], %[[COND_LOAD]] ], [ [[PASSTHROUGH]], [[TMP0:%.*]] ]
; CHECK-NEXT: ret <2 x i32> [[RET]]
;
%group.id = call i32 @llvm.amdgcn.workgroup.id.x()
%mask = icmp ult i32 %group.id, 4
%mask.vec = insertelement <2 x i1> poison, i1 %mask, i32 0
%mask.splat = shufflevector <2 x i1> %mask.vec, <2 x i1> poison, <2 x i32> zeroinitializer
%ret = call <2 x i32> @llvm.masked.load.v2i32.p0(ptr %p, i32 8, <2 x i1> %mask.splat, <2 x i32> %passthrough)
ret <2 x i32> %ret
}

declare <2 x i32> @llvm.masked.load.v2i32.p0(ptr, i32, <2 x i1>, <2 x i32>)
declare <2 x half> @llvm.masked.load.v2f16.p0(ptr, i32, <2 x i1>, <2 x half>)
declare <2 x i32> @llvm.masked.load.v2i32.p3(ptr addrspace(3), i32, <2 x i1>, <2 x i32>)
declare noundef i32 @llvm.amdgcn.workitem.id.x()
declare noundef i32 @llvm.amdgcn.workgroup.id.x()
38 changes: 38 additions & 0 deletions llvm/test/Transforms/ScalarizeMaskedMemIntrin/AMDGPU/expand-masked-gather.ll
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,38 @@
; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
; RUN: opt -S %s -passes=scalarize-masked-mem-intrin -mtriple=amdgcn-amd-amdhsa | FileCheck %s

; COM: Test that, unlike on CPU targets, the mask doesn't get bitcast to a scalar,
; COM: since, on GPUs, each i1 takes up at least one register and so they should
; COM: be treated separately.

define <2 x i32> @scalarize_v2i32(<2 x ptr> %p, <2 x i1> %mask, <2 x i32> %passthru) {
; CHECK-LABEL: define <2 x i32> @scalarize_v2i32(
; CHECK-SAME: <2 x ptr> [[P:%.*]], <2 x i1> [[MASK:%.*]], <2 x i32> [[PASSTHRU:%.*]]) {
; CHECK-NEXT: [[SCALAR_MASK:%.*]] = bitcast <2 x i1> [[MASK]] to i2
; CHECK-NEXT: [[TMP1:%.*]] = and i2 [[SCALAR_MASK]], 1
; CHECK-NEXT: [[TMP2:%.*]] = icmp ne i2 [[TMP1]], 0
; CHECK-NEXT: br i1 [[TMP2]], label %[[COND_LOAD:.*]], label %[[ELSE:.*]]
; CHECK: [[COND_LOAD]]:
; CHECK-NEXT: [[PTR0:%.*]] = extractelement <2 x ptr> [[P]], i64 0
; CHECK-NEXT: [[LOAD0:%.*]] = load i32, ptr [[PTR0]], align 8
; CHECK-NEXT: [[RES0:%.*]] = insertelement <2 x i32> [[PASSTHRU]], i32 [[LOAD0]], i64 0
; CHECK-NEXT: br label %[[ELSE]]
; CHECK: [[ELSE]]:
; CHECK-NEXT: [[RES_PHI_ELSE:%.*]] = phi <2 x i32> [ [[RES0]], %[[COND_LOAD]] ], [ [[PASSTHRU]], [[TMP0:%.*]] ]
; CHECK-NEXT: [[TMP3:%.*]] = and i2 [[SCALAR_MASK]], -2
; CHECK-NEXT: [[TMP4:%.*]] = icmp ne i2 [[TMP3]], 0
; CHECK-NEXT: br i1 [[TMP4]], label %[[COND_LOAD1:.*]], label %[[ELSE2:.*]]
; CHECK: [[COND_LOAD1]]:
; CHECK-NEXT: [[PTR1:%.*]] = extractelement <2 x ptr> [[P]], i64 1
; CHECK-NEXT: [[LOAD1:%.*]] = load i32, ptr [[PTR1]], align 8
; CHECK-NEXT: [[RES1:%.*]] = insertelement <2 x i32> [[RES_PHI_ELSE]], i32 [[LOAD1]], i64 1
; CHECK-NEXT: br label %[[ELSE2]]
; CHECK: [[ELSE2]]:
; CHECK-NEXT: [[RES_PHI_ELSE3:%.*]] = phi <2 x i32> [ [[RES1]], %[[COND_LOAD1]] ], [ [[RES_PHI_ELSE]], %[[ELSE]] ]
; CHECK-NEXT: ret <2 x i32> [[RES_PHI_ELSE3]]
;
%ret = call <2 x i32> @llvm.masked.gather.v2i32.v2p0(<2 x ptr> %p, i32 8, <2 x i1> %mask, <2 x i32> %passthru)
ret <2 x i32> %ret
}

declare <2 x i32> @llvm.masked.gather.v2i32.v2p0(<2 x ptr>, i32, <2 x i1>, <2 x i32>)
36 changes: 36 additions & 0 deletions llvm/test/Transforms/ScalarizeMaskedMemIntrin/AMDGPU/expand-masked-scatter.ll
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,36 @@
; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
; RUN: opt -S %s -passes=scalarize-masked-mem-intrin -mtriple=amdgcn-amd-amdhsa | FileCheck %s

; COM: Test that, unlike on CPU targets, the mask doesn't get bitcast to a scalar,
; COM: since, on GPUs, each i1 takes up at least one register and so they should
; COM: be treated separately.

define void @scalarize_v2i32(<2 x ptr> %p, <2 x i1> %mask, <2 x i32> %value) {
; CHECK-LABEL: define void @scalarize_v2i32(
; CHECK-SAME: <2 x ptr> [[P:%.*]], <2 x i1> [[MASK:%.*]], <2 x i32> [[VALUE:%.*]]) {
; CHECK-NEXT: [[SCALAR_MASK:%.*]] = bitcast <2 x i1> [[MASK]] to i2
; CHECK-NEXT: [[TMP1:%.*]] = and i2 [[SCALAR_MASK]], 1
; CHECK-NEXT: [[TMP2:%.*]] = icmp ne i2 [[TMP1]], 0
; CHECK-NEXT: br i1 [[TMP2]], label %[[COND_STORE:.*]], label %[[ELSE:.*]]
; CHECK: [[COND_STORE]]:
; CHECK-NEXT: [[ELT0:%.*]] = extractelement <2 x i32> [[VALUE]], i64 0
; CHECK-NEXT: [[PTR0:%.*]] = extractelement <2 x ptr> [[P]], i64 0
; CHECK-NEXT: store i32 [[ELT0]], ptr [[PTR0]], align 8
; CHECK-NEXT: br label %[[ELSE]]
; CHECK: [[ELSE]]:
; CHECK-NEXT: [[TMP3:%.*]] = and i2 [[SCALAR_MASK]], -2
; CHECK-NEXT: [[TMP4:%.*]] = icmp ne i2 [[TMP3]], 0
; CHECK-NEXT: br i1 [[TMP4]], label %[[COND_STORE1:.*]], label %[[ELSE2:.*]]
; CHECK: [[COND_STORE1]]:
; CHECK-NEXT: [[ELT1:%.*]] = extractelement <2 x i32> [[VALUE]], i64 1
; CHECK-NEXT: [[PTR1:%.*]] = extractelement <2 x ptr> [[P]], i64 1
; CHECK-NEXT: store i32 [[ELT1]], ptr [[PTR1]], align 8
; CHECK-NEXT: br label %[[ELSE2]]
; CHECK: [[ELSE2]]:
; CHECK-NEXT: ret void
;
call void @llvm.masked.scatter.v2i32.v2p0(<2 x i32> %value, <2 x ptr> %p, i32 8, <2 x i1> %mask)
ret void
}

declare void @llvm.masked.scatter.v2i32.v2p0(<2 x i32>, <2 x ptr>, i32, <2 x i1>)
Loading
Loading