llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT--> @llvm/pr-subscribers-clang-codegen Author: Joseph Huber (jhuber6) <details> <summary>Changes</summary> Summary: First attempt to make some AMD specific behavior for UBSan checks. The language manual states that using non-uniform values to readlane or non-present indices to shuffling is undefined. There may be a few more of these, but I thought this was low hanging fruit, hopefully uncontroversial. --- Full diff: https://github.com/llvm/llvm-project/pull/231485.diff 7 Files Affected: - (modified) clang/lib/CodeGen/CodeGenFunction.h (+2) - (modified) clang/lib/CodeGen/TargetBuiltins/AMDGPU.cpp (+46-6) - (added) clang/test/CodeGen/AMDGPU/ubsan-lane-builtins.c (+56) - (modified) compiler-rt/lib/ubsan/ubsan_handlers.cpp (+8) - (modified) compiler-rt/lib/ubsan/ubsan_handlers.h (+2) - (modified) compiler-rt/test/Inputs/hip.h (+37) - (added) compiler-rt/test/ubsan/AMDGPU/lanes.hip (+55) ``````````diff diff --git a/clang/lib/CodeGen/CodeGenFunction.h b/clang/lib/CodeGen/CodeGenFunction.h index 00294ad547d5ec..9c3b9ddf99761d 100644 --- a/clang/lib/CodeGen/CodeGenFunction.h +++ b/clang/lib/CodeGen/CodeGenFunction.h @@ -5448,6 +5448,8 @@ class CodeGenFunction : public CodeGenTypeCache { BCK_CTZPassedZero, BCK_CLZPassedZero, BCK_AssumePassedFalse, + BCK_NonUniformLane, + BCK_InactiveLane, }; /// Emits an argument for a call to a builtin. If the builtin sanitizer is diff --git a/clang/lib/CodeGen/TargetBuiltins/AMDGPU.cpp b/clang/lib/CodeGen/TargetBuiltins/AMDGPU.cpp index 2f29e9557ba56f..be819da4bb77d9 100644 --- a/clang/lib/CodeGen/TargetBuiltins/AMDGPU.cpp +++ b/clang/lib/CodeGen/TargetBuiltins/AMDGPU.cpp @@ -11,6 +11,7 @@ //===----------------------------------------------------------------------===// #include "CGBuiltin.h" +#include "CGDebugInfo.h" #include "CodeGenFunction.h" #include "TargetInfo.h" #include "clang/Basic/DiagnosticFrontend.h" @@ -555,6 +556,21 @@ static Intrinsic::ID getIntrinsicIDforWaveReduction(unsigned BuiltinID) { } } +static void emitBuiltinCheck(CodeGenFunction &CGF, const CallExpr *E, + CodeGenFunction::BuiltinCheckKind Kind, + llvm::function_ref<Value *()> EmitCond) { + if (!CGF.SanOpts.has(SanitizerKind::Builtin)) + return; + + auto CheckOrdinal = SanitizerKind::SO_Builtin; + auto CheckHandler = SanitizerHandler::InvalidBuiltin; + SanitizerDebugLocation SanScope(&CGF, {CheckOrdinal}, CheckHandler); + CGF.EmitCheck( + std::make_pair(EmitCond(), CheckOrdinal), CheckHandler, + {CGF.EmitCheckSourceLocation(E->getExprLoc()), CGF.Builder.getInt8(Kind)}, + {}); +} + Value *CodeGenFunction::EmitAMDGPUBuiltinExpr(unsigned BuiltinID, const CallExpr *E) { llvm::AtomicOrdering AO = llvm::AtomicOrdering::SequentiallyConsistent; @@ -683,12 +699,36 @@ Value *CodeGenFunction::EmitAMDGPUBuiltinExpr(unsigned BuiltinID, case AMDGPU::BI__builtin_amdgcn_permlane64: return emitBuiltinWithOneOverloadedType<1>(*this, E, Intrinsic::amdgcn_permlane64); - case AMDGPU::BI__builtin_amdgcn_readlane: - return emitBuiltinWithOneOverloadedType<2>(*this, E, - Intrinsic::amdgcn_readlane); - case AMDGPU::BI__builtin_amdgcn_wave_shuffle: - return emitBuiltinWithOneOverloadedType<2>(*this, E, - Intrinsic::amdgcn_wave_shuffle); + case AMDGPU::BI__builtin_amdgcn_readlane: { + Value *Src = EmitScalarExpr(E->getArg(0)); + Value *Lane = EmitScalarExpr(E->getArg(1)); + // The lane index must be uniform across the wavefront. + emitBuiltinCheck(*this, E, BCK_NonUniformLane, [&] { + Value *First = Builder.CreateIntrinsic(Intrinsic::amdgcn_readfirstlane, + {Lane->getType()}, {Lane}); + return Builder.CreateICmpEQ(Lane, First); + }); + return Builder.CreateIntrinsic(Intrinsic::amdgcn_readlane, {Src->getType()}, + {Src, Lane}); + } + case AMDGPU::BI__builtin_amdgcn_wave_shuffle: { + Value *Src = EmitScalarExpr(E->getArg(0)); + Value *Idx = EmitScalarExpr(E->getArg(1)); + // The source lane, wrapped to the wavefront size, must be active. + emitBuiltinCheck(*this, E, BCK_InactiveLane, [&] { + Value *Active = Builder.CreateIntrinsic(Intrinsic::amdgcn_ballot, + {Int64Ty}, {Builder.getTrue()}); + Value *Size = + Builder.CreateIntrinsic(Int32Ty, Intrinsic::amdgcn_wavefrontsize, {}); + Value *Lane = + Builder.CreateAnd(Idx, Builder.CreateSub(Size, Builder.getInt32(1))); + Value *Bit = + Builder.CreateLShr(Active, Builder.CreateZExt(Lane, Int64Ty)); + return Builder.CreateTrunc(Bit, Builder.getInt1Ty()); + }); + return Builder.CreateIntrinsic(Intrinsic::amdgcn_wave_shuffle, + {Src->getType()}, {Src, Idx}); + } case AMDGPU::BI__builtin_amdgcn_readfirstlane: return emitBuiltinWithOneOverloadedType<1>(*this, E, Intrinsic::amdgcn_readfirstlane); diff --git a/clang/test/CodeGen/AMDGPU/ubsan-lane-builtins.c b/clang/test/CodeGen/AMDGPU/ubsan-lane-builtins.c new file mode 100644 index 00000000000000..b7f0d7105803fa --- /dev/null +++ b/clang/test/CodeGen/AMDGPU/ubsan-lane-builtins.c @@ -0,0 +1,56 @@ +// NOTE: Assertions have been autogenerated by utils/update_cc_test_checks.py UTC_ARGS: --version 6 +// RUN: %clang_cc1 -triple amdgpu11.00-amd-amdhsa -emit-llvm -fsanitize=builtin \ +// RUN: -o - %s | FileCheck %s + +// CHECK-LABEL: define dso_local i32 @check_readlane( +// CHECK-SAME: i32 noundef [[X:%.*]], i32 noundef [[LANE:%.*]]) #[[ATTR0:[0-9]+]] { +// CHECK-NEXT: [[ENTRY:.*:]] +// CHECK-NEXT: [[X_ADDR:%.*]] = alloca i32, align 4, addrspace(5) +// CHECK-NEXT: [[LANE_ADDR:%.*]] = alloca i32, align 4, addrspace(5) +// CHECK-NEXT: [[X_ADDR_ASCAST:%.*]] = addrspacecast ptr addrspace(5) [[X_ADDR]] to ptr +// CHECK-NEXT: [[LANE_ADDR_ASCAST:%.*]] = addrspacecast ptr addrspace(5) [[LANE_ADDR]] to ptr +// CHECK-NEXT: store i32 [[X]], ptr [[X_ADDR_ASCAST]], align 4 +// CHECK-NEXT: store i32 [[LANE]], ptr [[LANE_ADDR_ASCAST]], align 4 +// CHECK-NEXT: [[TMP0:%.*]] = load i32, ptr [[X_ADDR_ASCAST]], align 4 +// CHECK-NEXT: [[TMP1:%.*]] = load i32, ptr [[LANE_ADDR_ASCAST]], align 4 +// CHECK-NEXT: [[TMP2:%.*]] = call i32 @llvm.amdgcn.readfirstlane.i32(i32 [[TMP1]]), !nosanitize [[META2:![0-9]+]] +// CHECK-NEXT: [[TMP3:%.*]] = icmp eq i32 [[TMP1]], [[TMP2]], !nosanitize [[META2]] +// CHECK-NEXT: br i1 [[TMP3]], label %[[CONT:.*]], label %[[HANDLER_INVALID_BUILTIN:.*]], !prof [[PROF3:![0-9]+]], !nosanitize [[META2]] +// CHECK: [[HANDLER_INVALID_BUILTIN]]: +// CHECK-NEXT: call void @__ubsan_handle_invalid_builtin_abort(ptr addrspacecast (ptr addrspace(1) @[[GLOB0:[0-9]+]] to ptr)) #[[ATTR6:[0-9]+]], !nosanitize [[META2]] +// CHECK-NEXT: unreachable, !nosanitize [[META2]] +// CHECK: [[CONT]]: +// CHECK-NEXT: [[TMP4:%.*]] = call i32 @llvm.amdgcn.readlane.i32(i32 [[TMP0]], i32 [[TMP1]]) +// CHECK-NEXT: ret i32 [[TMP4]] +int check_readlane(int x, int lane) { + return __builtin_amdgcn_readlane(x, lane); +} + +// CHECK-LABEL: define dso_local i32 @check_wave_shuffle( +// CHECK-SAME: i32 noundef [[X:%.*]], i32 noundef [[IDX:%.*]]) #[[ATTR0]] { +// CHECK-NEXT: [[ENTRY:.*:]] +// CHECK-NEXT: [[X_ADDR:%.*]] = alloca i32, align 4, addrspace(5) +// CHECK-NEXT: [[IDX_ADDR:%.*]] = alloca i32, align 4, addrspace(5) +// CHECK-NEXT: [[X_ADDR_ASCAST:%.*]] = addrspacecast ptr addrspace(5) [[X_ADDR]] to ptr +// CHECK-NEXT: [[IDX_ADDR_ASCAST:%.*]] = addrspacecast ptr addrspace(5) [[IDX_ADDR]] to ptr +// CHECK-NEXT: store i32 [[X]], ptr [[X_ADDR_ASCAST]], align 4 +// CHECK-NEXT: store i32 [[IDX]], ptr [[IDX_ADDR_ASCAST]], align 4 +// CHECK-NEXT: [[TMP0:%.*]] = load i32, ptr [[X_ADDR_ASCAST]], align 4 +// CHECK-NEXT: [[TMP1:%.*]] = load i32, ptr [[IDX_ADDR_ASCAST]], align 4 +// CHECK-NEXT: [[TMP2:%.*]] = call i64 @llvm.amdgcn.ballot.i64(i1 true), !nosanitize [[META2]] +// CHECK-NEXT: [[TMP3:%.*]] = call i32 @llvm.amdgcn.wavefrontsize(), !nosanitize [[META2]] +// CHECK-NEXT: [[TMP4:%.*]] = sub i32 [[TMP3]], 1, !nosanitize [[META2]] +// CHECK-NEXT: [[TMP5:%.*]] = and i32 [[TMP1]], [[TMP4]], !nosanitize [[META2]] +// CHECK-NEXT: [[TMP6:%.*]] = zext i32 [[TMP5]] to i64, !nosanitize [[META2]] +// CHECK-NEXT: [[TMP7:%.*]] = lshr i64 [[TMP2]], [[TMP6]], !nosanitize [[META2]] +// CHECK-NEXT: [[TMP8:%.*]] = trunc i64 [[TMP7]] to i1, !nosanitize [[META2]] +// CHECK-NEXT: br i1 [[TMP8]], label %[[CONT:.*]], label %[[HANDLER_INVALID_BUILTIN:.*]], !prof [[PROF3]], !nosanitize [[META2]] +// CHECK: [[HANDLER_INVALID_BUILTIN]]: +// CHECK-NEXT: call void @__ubsan_handle_invalid_builtin_abort(ptr addrspacecast (ptr addrspace(1) @[[GLOB1:[0-9]+]] to ptr)) #[[ATTR6]], !nosanitize [[META2]] +// CHECK-NEXT: unreachable, !nosanitize [[META2]] +// CHECK: [[CONT]]: +// CHECK-NEXT: [[TMP9:%.*]] = call i32 @llvm.amdgcn.wave.shuffle.i32(i32 [[TMP0]], i32 [[TMP1]]) +// CHECK-NEXT: ret i32 [[TMP9]] +int check_wave_shuffle(int x, int idx) { + return __builtin_amdgcn_wave_shuffle(x, idx); +} diff --git a/compiler-rt/lib/ubsan/ubsan_handlers.cpp b/compiler-rt/lib/ubsan/ubsan_handlers.cpp index cbb02892984e42..d60606ace7718d 100644 --- a/compiler-rt/lib/ubsan/ubsan_handlers.cpp +++ b/compiler-rt/lib/ubsan/ubsan_handlers.cpp @@ -659,6 +659,14 @@ void __ubsan::handleInvalidBuiltin(InvalidBuiltinData *Data, if (Data->Kind == BCK_AssumePassedFalse) Diag(Loc, DL_Error, ET, "assumption is violated during execution"); + else if (Data->Kind == BCK_NonUniformLane) + Diag(Loc, DL_Error, ET, + "lane index passed to __builtin_amdgcn_readlane() is not uniform " + "across active lanes"); + else if (Data->Kind == BCK_InactiveLane) + Diag(Loc, DL_Error, ET, + "lane index passed to __builtin_amdgcn_wave_shuffle() refers to an " + "inactive lane"); else Diag(Loc, DL_Error, ET, "passing zero to __builtin_%0(), which is not a valid argument") diff --git a/compiler-rt/lib/ubsan/ubsan_handlers.h b/compiler-rt/lib/ubsan/ubsan_handlers.h index 521caa96bc771b..9420a1b69187ae 100644 --- a/compiler-rt/lib/ubsan/ubsan_handlers.h +++ b/compiler-rt/lib/ubsan/ubsan_handlers.h @@ -163,6 +163,8 @@ enum BuiltinCheckKind : unsigned char { BCK_CTZPassedZero, BCK_CLZPassedZero, BCK_AssumePassedFalse, + BCK_NonUniformLane, + BCK_InactiveLane, }; struct InvalidBuiltinData { diff --git a/compiler-rt/test/Inputs/hip.h b/compiler-rt/test/Inputs/hip.h index 6177ef94205bfb..1837b6581fa0cc 100644 --- a/compiler-rt/test/Inputs/hip.h +++ b/compiler-rt/test/Inputs/hip.h @@ -37,6 +37,43 @@ int hipLaunchKernel(const void *Func, dim3 GridDim, dim3 BlockDim, void **Args, int printf(const char *, ...); } +__device__ inline int __hip_warp_size() { + return __builtin_amdgcn_wavefrontsize(); +} +#define warpSize __hip_warp_size() + +__device__ inline int __lane_id() { + return __builtin_amdgcn_mbcnt_hi(~0u, __builtin_amdgcn_mbcnt_lo(~0u, 0)); +} + +__device__ inline int __shfl(int Var, int SrcLane, int Width = warpSize) { + int Self = __lane_id(); + int Index = (SrcLane & (Width - 1)) + (Self & ~(Width - 1)); + return __builtin_amdgcn_wave_shuffle(Var, Index); +} + +__device__ inline int __shfl_up(int Var, unsigned Delta, int Width = warpSize) { + int Self = __lane_id(); + int Index = Self - Delta; + Index = Index < (Self & ~(Width - 1)) ? Self : Index; + return __builtin_amdgcn_wave_shuffle(Var, Index); +} + +__device__ inline int __shfl_down(int Var, unsigned Delta, + int Width = warpSize) { + int Self = __lane_id(); + int Index = + (int)((Self & (Width - 1)) + Delta) >= Width ? Self : Self + Delta; + return __builtin_amdgcn_wave_shuffle(Var, Index); +} + +__device__ inline int __shfl_xor(int Var, int LaneMask, int Width = warpSize) { + int Self = __lane_id(); + int Index = Self ^ LaneMask; + Index = Index >= ((Self + Width) & ~(Width - 1)) ? Self : Index; + return __builtin_amdgcn_wave_shuffle(Var, Index); +} + #define CHECK_HIP(Expr) \ do { \ if ((Expr) != 0) { \ diff --git a/compiler-rt/test/ubsan/AMDGPU/lanes.hip b/compiler-rt/test/ubsan/AMDGPU/lanes.hip new file mode 100644 index 00000000000000..6f3937fe574742 --- /dev/null +++ b/compiler-rt/test/ubsan/AMDGPU/lanes.hip @@ -0,0 +1,55 @@ +// REQUIRES: ubsan-hip +// RUN: %clang_hip -fsanitize=builtin %s -o %t %hip_libs +// RUN: %run %t 2>&1 | FileCheck %s + +#define TID __builtin_amdgcn_workitem_id_x() + +__device__ int reduce(int V) { + for (int D = warpSize / 2; D > 0; D /= 2) + V += __shfl_down(V, D); + return V; +} + +// CHECK-NOT: runtime error +__global__ void clean(int *Out) { + int Lane = __lane_id(); + int X = Out[TID]; + int V = __builtin_amdgcn_readlane(X, 5); + V += __shfl(X, Lane + 1) + __shfl_up(X, 1) + __shfl_xor(X, 1) + reduce(X); + if (Lane < 16) + V += __shfl(X, Lane ^ 1); + Out[TID] = V; +} + +__global__ void readlane(int *Out) { + int X = Out[TID]; + // CHECK: lanes.hip:[[#@LINE+1]]:14: runtime error: lane index passed to __builtin_amdgcn_readlane() is not uniform across active lanes + Out[TID] = __builtin_amdgcn_readlane(X, TID % 4); +} + +// The reported location is the builtin inside the hip.h wrapper. +__global__ void shuffle(int *Out) { + int Lane = __lane_id(); + int X = Out[TID]; + // CHECK: hip.h:[[#]]:[[#]]: runtime error: lane index passed to __builtin_amdgcn_wave_shuffle() refers to an inactive lane + if (Lane < 16) + Out[TID] = __shfl(X, Lane + 16); +} + +// CHECK: hip.h:[[#]]:[[#]]: runtime error: lane index passed to __builtin_amdgcn_wave_shuffle() refers to an inactive lane +__global__ void partial_reduce(int *Out) { Out[TID] = reduce(Out[TID]); } + +int main() { + int *P = nullptr; + CHECK_HIP(hipMalloc((void **)&P, 64 * sizeof(int))); + clean<<<1, 64>>>(P); + CHECK_HIP(hipDeviceSynchronize()); + readlane<<<1, 64>>>(P); + CHECK_HIP(hipDeviceSynchronize()); + shuffle<<<1, 64>>>(P); + CHECK_HIP(hipDeviceSynchronize()); + partial_reduce<<<1, 20>>>(P); + CHECK_HIP(hipDeviceSynchronize()); + hipFree(P); + return 0; +} `````````` </details> https://github.com/llvm/llvm-project/pull/231485 _______________________________________________ cfe-commits mailing list [email protected] https://lists.llvm.org/cgi-bin/mailman/listinfo/cfe-commits
