Skip to content

Fix CUDA global argument detection for aliased kernel parameters #1843

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

Open
MdSaifAliMolla wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Fix CUDA global argument detection for aliased kernel parameters #1843

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
5 changes: 1 addition & 4 deletions lib/Differentiator/DiffPlanner.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1302,11 +1302,8 @@ static QualType GetDerivedFunctionType(const CallExpr* CE) {
if (!utils::isArrayOrPointerType(paramDecl->getType()) ||
IsDifferentiableArg)
request.DVI.push_back(paramDecl);
// FIXME: If we cannot deduce whether the argument is
// differentiable, we should still add it to CUDAGlobalArgsIndexes.
// i.e. remove `&& PVD`.
// We know we should use atomic ops here
if (useCUDA && PVD && IsDifferentiableArg)
if (useCUDA && IsDifferentiableArg)
request.CUDAGlobalArgsIndexes.push_back(i);
}
}
Expand Down
39 changes: 31 additions & 8 deletions lib/Differentiator/ReverseModeVisitor.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2050,6 +2050,9 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) {
// take arg by reference, we can request a derivative w.r.t. to this arg
// using the forward mode.
bool asGrad = !utils::canUsePushforwardInRevMode(FD);
// Method operators have a base like methods do but it's included in the
// call arguments so we have to shift the indexing of call arguments.
bool isMethodOperatorCall = MD && isa<CXXOperatorCallExpr>(CE);

// Build the DiffRequest
DiffRequest pullbackRequest{};
Expand Down Expand Up @@ -2078,9 +2081,23 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) {
size_t i = 0;
for (const ParmVarDecl* PVD : FD->parameters()) {
pullbackRequest.DVI.push_back(PVD);
// To determine if this callee parameter corresponds to a
// CUDA global arg, check the actual call argument (not the
// callee's param name) to find the caller's ParmVarDecl.
std::size_t argIdx =
i + static_cast<std::size_t>(isMethodOperatorCall);
if (!m_DiffReq.CUDAGlobalArgsIndexes.empty() &&
m_DiffReq.HasIndependentParameter(PVD))
pullbackRequest.CUDAGlobalArgsIndexes.push_back(i);
argIdx < CE->getNumArgs()) {
const Expr* arg = CE->getArg(argIdx)->IgnoreParenImpCasts();
const ParmVarDecl* callerPVD = nullptr;
if (const auto* DRE = dyn_cast<DeclRefExpr>(arg))
callerPVD = dyn_cast<ParmVarDecl>(DRE->getDecl());
// If we can identify the caller's param and it is an
// independent parameter (i.e. a CUDA global arg), or if
// we cannot identify it, conservatively mark it.
if (!callerPVD || m_DiffReq.HasIndependentParameter(callerPVD))
pullbackRequest.CUDAGlobalArgsIndexes.push_back(i);
}
++i;
}
}
Expand All @@ -2094,9 +2111,6 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) {
// statements there later.
std::size_t insertionPoint = getCurrentBlock(direction::reverse).size();

// Method operators have a base like methods do but it's included in the
// call arguments so we have to shift the indexing of call arguments.
bool isMethodOperatorCall = MD && isa<CXXOperatorCallExpr>(CE);
// The set of arguments to be used in a ``reverse_forw`` after the original
// args.
llvm::SmallVector<Expr*, 16> revForwAdjointArgs{};
Expand Down Expand Up @@ -2222,9 +2236,18 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) {
size_t offset =
(bool)MD && MD->isInstance() && !isLambdaCallOperator(MD);
if (CallArgDx[i + offset]) {
if (!m_DiffReq.CUDAGlobalArgsIndexes.empty() &&
m_DiffReq.HasIndependentParameter(PVD))
pullbackRequest.CUDAGlobalArgsIndexes.push_back(i);
if (!m_DiffReq.CUDAGlobalArgsIndexes.empty()) {
std::size_t argIdx =
i + static_cast<std::size_t>(isMethodOperatorCall);
if (argIdx < CE->getNumArgs()) {
const Expr* arg = CE->getArg(argIdx)->IgnoreParenImpCasts();
const ParmVarDecl* callerPVD = nullptr;
if (const auto* DRE = dyn_cast<DeclRefExpr>(arg))
callerPVD = dyn_cast<ParmVarDecl>(DRE->getDecl());
if (!callerPVD || m_DiffReq.HasIndependentParameter(callerPVD))
pullbackRequest.CUDAGlobalArgsIndexes.push_back(i);
}
}
pullbackRequest.DVI.push_back(PVD);
} else
hasDynamicNonDiffParams = true;
Expand Down
57 changes: 54 additions & 3 deletions test/CUDA/GradientKernels.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -262,7 +262,7 @@ __global__ void kernel_with_device_call(double *out, const double *in, double va
out[index] = device_fn(in[index], val);
}

// CHECK: __attribute__((device)) void device_fn_pullback_1(const double in, double val, double _d_y, double *_d_in, double *_d_val) {
// CHECK: __attribute__((device)) void device_fn_pullback_0_1(const double in, double val, double _d_y, double *_d_in, double *_d_val) {
//CHECK-NEXT: {
//CHECK-NEXT: *_d_in += _d_y;
//CHECK-NEXT: *_d_val += _d_y;
Expand All @@ -280,7 +280,7 @@ __global__ void kernel_with_device_call(double *out, const double *in, double va
//CHECK-NEXT: _d_out[index0] = 0.;
//CHECK-NEXT: double _r0 = 0.;
//CHECK-NEXT: double _r1 = 0.;
//CHECK-NEXT: device_fn_pullback_1(in[index0], val, _r_d0, &_r0, &_r1);
//CHECK-NEXT: device_fn_pullback_0_1(in[index0], val, _r_d0, &_r0, &_r1);
//CHECK-NEXT: atomicAdd(_d_val, _r1);
//CHECK-NEXT: }
//CHECK-NEXT:}
Expand Down Expand Up @@ -444,7 +444,7 @@ __global__ void fn1(double *out, const double *in, double val) {
// CHECK-NEXT: _d_out[index0] = 0.;
// CHECK-NEXT: double _r0 = 0.;
// CHECK-NEXT: double _r1 = 0.;
// CHECK-NEXT: device_fn_pullback_1(in[index0], temp, _r_d0, &_r0, &_r1);
// CHECK-NEXT: device_fn_pullback_0_1(in[index0], temp, _r_d0, &_r0, &_r1);
// CHECK-NEXT: _d_temp += _r1;
// CHECK-NEXT: }
// CHECK-NEXT: atomicAdd(_d_val, _d_temp);
Expand Down Expand Up @@ -872,6 +872,53 @@ __global__ void injective_reassignment_loop(int *a) {
//CHECK-NEXT: }
//CHECK-NEXT: }

__device__ void device_int_add(int *a, int *b) { *b += *a; }
__device__ void device_mul_addr(int *a, int *b) { *b *= *a; }

__global__ void alias_kernel(int *a) { device_int_add(a, a); }
__global__ void offset_kernel(int *a) { device_int_add(a, a + 1); }
__global__ void addr_kernel(int *a) { int b = 5; device_mul_addr(a, &b); }


// CHECK: __attribute__((device)) void device_int_add_pullback_0_1(int *a, int *b, int *_d_a, int *_d_b) {
// CHECK-NEXT: int _t0 = *b;
// CHECK-NEXT: *b += *a;
// CHECK-NEXT: {
// CHECK-NEXT: *b = _t0;
// CHECK-NEXT: int _r_d0 = *_d_b;
// CHECK-NEXT: atomicAdd(_d_a, _r_d0);
// CHECK-NEXT: }
// CHECK-NEXT: }

// CHECK: void alias_kernel_grad(int *a, int *_d_a) {
// CHECK-NEXT: device_int_add(a, a);
// CHECK-NEXT: device_int_add_pullback_0_1(a, a, _d_a, _d_a);
// CHECK-NEXT: }

// CHECK: void offset_kernel_grad(int *a, int *_d_a) {
// CHECK-NEXT: device_int_add(a, a + 1);
// CHECK-NEXT: device_int_add_pullback_0_1(a, a + 1, _d_a, _d_a + 1);
// CHECK-NEXT: }

// CHECK: __attribute__((device)) void device_mul_addr_pullback_0_1(int *a, int *b, int *_d_a, int *_d_b) {
// CHECK-NEXT: int _t0 = *b;
// CHECK-NEXT: *b *= *a;
// CHECK-NEXT: {
// CHECK-NEXT: *b = _t0;
// CHECK-NEXT: int _r_d0 = *_d_b;
// CHECK-NEXT: *_d_b = 0;
// CHECK-NEXT: *_d_b += _r_d0 * *a;
// CHECK-NEXT: atomicAdd(_d_a, *b * _r_d0);
// CHECK-NEXT: }
// CHECK-NEXT: }

// CHECK: void addr_kernel_grad(int *a, int *_d_a) {
// CHECK-NEXT: int _d_b = 0;
// CHECK-NEXT: int b = 5;
// CHECK-NEXT: device_mul_addr(a, &b);
// CHECK-NEXT: device_mul_addr_pullback_0_1(a, &b, _d_a, &_d_b);
// CHECK-NEXT: }

#define TEST(F, grid, block, shared_mem, use_stream, x, dx, N) \
{ \
int *fives = (int*)malloc(N * sizeof(int)); \
Expand Down Expand Up @@ -1153,6 +1200,10 @@ int main(void) {
TEST(injective_reassignment, dim3(1), dim3(1), 0, false, n, d_n, 2); // CHECK-EXEC: 1
TEST(injective_reassignment_loop, dim3(1), dim3(1), 0, false, n, d_n,2); // CHECK-EXEC: 1

TEST(alias_kernel, dim3(1), dim3(1), 0, false, n, d_n, 1); // CHECK-EXEC: 2
TEST(offset_kernel, dim3(1), dim3(1), 0, false, n, d_n, 2); // CHECK-EXEC: 2, 1
TEST(addr_kernel, dim3(1), dim3(1), 0, false, n, d_n, 1); // CHECK-EXEC: 1

cudaFree(dummy_in);
cudaFree(dummy_out);
cudaFree(d_out);
Expand Down
Loading