[CUDA][HIP] Fix capturing reference to host variable

[yaxunl]

Summary

In C++ when a reference variable is captured by copy, the lambda
is supposed to make a copy of the referenced variable in the captures
and refer to the copy in the lambda. Therefore, it is valid to capture
a reference to a host global variable in a device lambda since the
device lambda will refer to the copy of the host global variable instead
of access the host global variable directly.

However, clang tries to avoid capturing of reference to a host global variable
if it determines the use of the reference variable in the lambda function is
not odr-use. Clang also tries to emit load of the reference to a global variable
as load of the global variable if it determines that the reference variable is
a compile-time constant.

For a device lambda to capture a reference variable to host global variable
and use the captured value, clang needs to be taught that in such cases the use of the reference
variable is odr-use and the reference variable is not compile-time constant.

This patch fixes that.

[yaxunl]

remove debug code

[tra]

The use case in the tests makes sense, but I'd like to have a second opinion on whether it may have unintended consequences. Lambdas tend to bring interesting corner cases.

@rsmith Richard, do you see any issues with this approach?

clang/test/CodeGenCUDA/lambda-reference-var.cu

62

Do we have current Sema tests that verify that we we would not allow accessing host vars here?

[yaxunl]

clang/test/CodeGenCUDA/lambda-reference-var.cu

62

currently clang does not diagnose access to host global variable in device functions. do we want to diagnose that?

[tra]

Ouch. This is definitely a bug. https://godbolt.org/z/faPWcv

Granted, the compilation will eventually fail in most cases due to ptxas failing to resolve the reference to the variable we do not generate on the device, but it's something we should catch in Sema.

Filed https://bugs.llvm.org/show_bug.cgi?id=48121

[yaxunl]

added diagnosing referencing host variable in device functions

[tra]

In D91088#2388496, @yaxunl wrote:

added diagnosing referencing host variable in device functions

Thank you for fixing this. The new changes look good, but I think they should be in a separate patch.

clang/lib/Sema/SemaExpr.cpp

357

This could use a comment about why we only check D->H references.
H->D references will still work (sort of) because on the host side we have matching shadow variables.

[yaxunl]

Separate diagnose of host variable to another patch.

[yaxunl]

ping

clang/lib/Sema/SemaExpr.cpp

357

the comment has been added in the separated patch for diagnostics.

[tra]

LGTM in general.

clang/lib/Sema/SemaExpr.cpp

1973

Nit: I'd make it a free function.
The large-ish lambda handling a niche case makes the original simple function less readable. I think keeping capture check separate would be a bit cleaner.

1984

This could use some Sema tests.

[yaxunl]

extract lambda as a function

[yaxunl]

clang/lib/Sema/SemaExpr.cpp

1973

done

1984

the negative cases of the condition that incur diagnostics are referencing host variable in device or host device functions. The sema tests have been added in https://reviews.llvm.org/D91281.

The other negative cases of the condition (capturing host reference in host lambda, referencing host var in host function, capturing device reference in host lambda, referencing device var in host function, etc) do not cause diagnostics, which are covered by codegen tests.

[tra]

clang/lib/Sema/SemaExpr.cpp

1939

const Sema& ?

[yaxunl]

clang/lib/Sema/SemaExpr.cpp

1939

will do when committing

clang/lib/CodeGen/CGExpr.cpp

CodeGenFunction::tryEmitAsConstant(DeclRefExpr *refExpr) {
   } else {
     return ConstantEmission();
   }
 
   // In any case, if the initializer has side-effects, abandon ship.
   if (result.HasSideEffects)
     return ConstantEmission();
 
+  // In CUDA/HIP device compilation, a lambda may capture a reference variable
+  // referencing a global host variable by copy. In this case the lambda should
+  // make a copy of the value of the global host variable. The DRE of the
+  // captured reference variable cannot be emitted as load from the host
+  // global variable as compile time constant, since the host variable is not
+  // accessible on device. The DRE of the captured reference variable has to be
+  // loaded from captures.
+  if (CGM.getLangOpts().CUDAIsDevice &&
+      refExpr->refersToEnclosingVariableOrCapture()) {
+    auto *MD = dyn_cast_or_null<CXXMethodDecl>(CurCodeDecl);
+    if (MD && MD->getParent()->isLambda() &&
+        MD->getOverloadedOperator() == OO_Call) {
+      const APValue::LValueBase &base = result.Val.getLValueBase();
+      if (const ValueDecl *D = base.dyn_cast<const ValueDecl *>()) {
+        if (const VarDecl *VD = dyn_cast<const VarDecl>(D)) {
+          if (!VD->hasAttr<CUDADeviceAttr>()) {
+            return ConstantEmission();
+          }
+        }
+      }
+    }
+  }
+
   // Emit as a constant.
   auto C = ConstantEmitter(*this).emitAbstract(refExpr->getLocation(),
                                                result.Val, resultType);
 
   // Make sure we emit a debug reference to the global variable.
   // This should probably fire even for
   if (isa<VarDecl>(value)) {
     if (!getContext().DeclMustBeEmitted(cast<VarDecl>(value)))

clang/lib/Sema/SemaExpr.cpp

bool Sema::DiagnoseUseOfDecl(NamedDecl *D, ArrayRef<SourceLocation> Locs,
 
   DiagnoseAvailabilityOfDecl(D, Locs, UnknownObjCClass, ObjCPropertyAccess,
                              AvoidPartialAvailabilityChecks, ClassReceiver);
 
   DiagnoseUnusedOfDecl(*this, D, Loc);
 
   diagnoseUseOfInternalDeclInInlineFunction(*this, D, Loc);
 
   if (LangOpts.SYCLIsDevice || (LangOpts.OpenMP && LangOpts.OpenMPIsDevice)) {
     if (const auto *VD = dyn_cast<ValueDecl>(D))
       checkDeviceDecl(VD, Loc);
 
     if (!Context.getTargetInfo().isTLSSupported())
       if (const auto *VD = dyn_cast<VarDecl>(D))
         if (VD->getTLSKind() != VarDecl::TLS_None)
           targetDiag(*Locs.begin(), diag::err_thread_unsupported);
   }
Sema::BuildDeclRefExpr(ValueDecl *D, QualType Ty, ExprValueKind VK,
                        SourceLocation TemplateKWLoc,
                        const TemplateArgumentListInfo *TemplateArgs) {
   NestedNameSpecifierLoc NNS =
       SS ? SS->getWithLocInContext(Context) : NestedNameSpecifierLoc();
   return BuildDeclRefExpr(D, Ty, VK, NameInfo, NNS, FoundD, TemplateKWLoc,
                           TemplateArgs);
 }
 
+// CUDA/HIP: Check whether a captured reference variable is referencing a
+// host variable in a device or host device lambda.
+static bool isCapturingReferenceToHostVarInCUDADeviceLambda(const Sema &S,
+                                                            VarDecl *VD) {
+  if (!S.getLangOpts().CUDA || !VD->hasInit())
+    return false;
+  assert(VD->getType()->isReferenceType());
+
+  // Check whether the reference variable is referencing a host variable.
+  auto *DRE = dyn_cast<DeclRefExpr>(VD->getInit());
+  if (!DRE)
+    return false;
+  auto *Referee = dyn_cast<VarDecl>(DRE->getDecl());
+  if (!Referee || !Referee->hasGlobalStorage() ||
+      Referee->hasAttr<CUDADeviceAttr>())
+    return false;
+
+  // Check whether the current function is a device or host device lambda.
+  // Check whether the reference variable is a capture by getDeclContext()
+  // since refersToEnclosingVariableOrCapture() is not ready at this point.
+  auto *MD = dyn_cast_or_null<CXXMethodDecl>(S.CurContext);
+  if (MD && MD->getParent()->isLambda() &&
+      MD->getOverloadedOperator() == OO_Call && MD->hasAttr<CUDADeviceAttr>() &&
+      VD->getDeclContext() != MD)
+    return true;
+
+  return false;
+}
+
 NonOdrUseReason Sema::getNonOdrUseReasonInCurrentContext(ValueDecl *D) {
   // A declaration named in an unevaluated operand never constitutes an odr-use.
   if (isUnevaluatedContext())
     return NOUR_Unevaluated;
 
   // C++2a [basic.def.odr]p4:
   //   A variable x whose name appears as a potentially-evaluated expression e
   //   is odr-used by e unless [...] x is a reference that is usable in
   //   constant expressions.
+  // CUDA/HIP:
+  //   If a reference variable referencing a host variable is captured in a
+  //   device or host device lambda, the value of the referee must be copied
+  //   to the capture and the reference variable must be treated as odr-use
+  //   since the value of the referee is not known at compile time and must
+  //   be loaded from the captured.
   if (VarDecl *VD = dyn_cast<VarDecl>(D)) {
     if (VD->getType()->isReferenceType() &&
         !(getLangOpts().OpenMP && isOpenMPCapturedDecl(D)) &&
+        !isCapturingReferenceToHostVarInCUDADeviceLambda(*this, VD) &&
         VD->isUsableInConstantExpressions(Context))
       return NOUR_Constant;
   }
 
   // All remaining non-variable cases constitute an odr-use. For variables, we
   // need to wait and see how the expression is used.
   return NOUR_None;
 }

clang/test/CodeGenCUDA/lambda-reference-var.cu

+// RUN: %clang_cc1 -x hip -emit-llvm -std=c++11 %s -o - \
+// RUN:   -triple x86_64-linux-gnu \
+// RUN:   | FileCheck -check-prefix=HOST %s
+// RUN: %clang_cc1 -x hip -emit-llvm -std=c++11 %s -o - \
+// RUN:   -triple amdgcn-amd-amdhsa -fcuda-is-device \
+// RUN:   | FileCheck -check-prefix=DEV %s
+
+#include "Inputs/cuda.h"
+
+// HOST: %[[T1:.*]] = type <{ i32*, i32, [4 x i8] }>
+// HOST: %[[T2:.*]] = type { i32*, i32** }
+// HOST: %[[T3:.*]] = type <{ i32*, i32, [4 x i8] }>
+// DEV: %[[T1:.*]] = type { i32* }
+// DEV: %[[T2:.*]] = type { i32** }
+// DEV: %[[T3:.*]] = type <{ i32*, i32, [4 x i8] }>
+int global_host_var;
+__device__ int global_device_var;
+
+template<class F>
+__global__ void kern(F f) { f(); }
+
+// DEV-LABEL: @_ZZ27dev_capture_dev_ref_by_copyPiENKUlvE_clEv(
+// DEV: %[[VAL:.*]] = load i32, i32* addrspacecast (i32 addrspace(1)* @global_device_var to i32*)
+// DEV: store i32 %[[VAL]]
+__device__ void dev_capture_dev_ref_by_copy(int *out) {
+  int &ref = global_device_var;
+  [=](){ *out = ref;}();
+}
+
+// DEV-LABEL: @_ZZ26dev_capture_dev_ref_by_refPiENKUlvE_clEv(
+// DEV: %[[VAL:.*]] = load i32, i32* addrspacecast (i32 addrspace(1)* @global_device_var to i32*)
+// DEV: %[[VAL2:.*]] = add nsw i32 %[[VAL]], 1
+// DEV: store i32 %[[VAL2]], i32* addrspacecast (i32 addrspace(1)* @global_device_var to i32*)
+// DEV: %[[VAL:.*]] = load i32, i32* addrspacecast (i32 addrspace(1)* @global_device_var to i32*)
+// DEV: store i32 %[[VAL]]
+__device__ void dev_capture_dev_ref_by_ref(int *out) {
+  int &ref = global_device_var;
+  [&](){ ref++; *out = ref;}();
+}
+
+// DEV-LABEL: define void @_Z7dev_refPi(
+// DEV: %[[VAL:.*]] = load i32, i32* addrspacecast (i32 addrspace(1)* @global_device_var to i32*)
+// DEV: %[[VAL2:.*]] = add nsw i32 %[[VAL]], 1
+// DEV: store i32 %[[VAL2]], i32* addrspacecast (i32 addrspace(1)* @global_device_var to i32*)
+// DEV: %[[VAL:.*]] = load i32, i32* addrspacecast (i32 addrspace(1)* @global_device_var to i32*)
+// DEV: store i32 %[[VAL]]
+__device__ void dev_ref(int *out) {
+  int &ref = global_device_var;
+  ref++;
+  *out = ref;
+}
+
+// DEV-LABEL: @_ZZ14dev_lambda_refPiENKUlvE_clEv(
+// DEV: %[[VAL:.*]] = load i32, i32* addrspacecast (i32 addrspace(1)* @global_device_var to i32*)
+// DEV: %[[VAL2:.*]] = add nsw i32 %[[VAL]], 1
+// DEV: store i32 %[[VAL2]], i32* addrspacecast (i32 addrspace(1)* @global_device_var to i32*)
+// DEV: %[[VAL:.*]] = load i32, i32* addrspacecast (i32 addrspace(1)* @global_device_var to i32*)
+// DEV: store i32 %[[VAL]]
+__device__ void dev_lambda_ref(int *out) {
+  [=](){
+    int &ref = global_device_var;
+    ref++;
+    *out = ref;
+  }();
+}
+
+// HOST-LABEL: @_ZZ29host_capture_host_ref_by_copyPiENKUlvE_clEv(
+// HOST: %[[VAL:.*]] = load i32, i32* @global_host_var
+// HOST: store i32 %[[VAL]]
+void host_capture_host_ref_by_copy(int *out) {
+  int &ref = global_host_var;
+  [=](){ *out = ref;}();
+}
+
+// HOST-LABEL: @_ZZ28host_capture_host_ref_by_refPiENKUlvE_clEv(
+// HOST: %[[CAP:.*]] = getelementptr inbounds %[[T2]], %[[T2]]* %this1, i32 0, i32 0
+// HOST: %[[REF:.*]] = load i32*, i32** %[[CAP]]
+// HOST: %[[VAL:.*]] = load i32, i32* %[[REF]]
+// HOST: %[[VAL2:.*]] = add nsw i32 %[[VAL]], 1
+// HOST: store i32 %[[VAL2]], i32* %[[REF]]
+// HOST: %[[VAL:.*]] = load i32, i32* @global_host_var
+// HOST: store i32 %[[VAL]]
+void host_capture_host_ref_by_ref(int *out) {
+  int &ref = global_host_var;
+  [&](){ ref++; *out = ref;}();
+}
+
+// HOST-LABEL: define void @_Z8host_refPi(
+// HOST: %[[VAL:.*]] = load i32, i32* @global_host_var
+// HOST: %[[VAL2:.*]] = add nsw i32 %[[VAL]], 1
+// HOST: store i32 %[[VAL2]], i32* @global_host_var
+// HOST: %[[VAL:.*]] = load i32, i32* @global_host_var
+// HOST: store i32 %[[VAL]]
+void host_ref(int *out) {
+  int &ref = global_host_var;
+  ref++;
+  *out = ref;
+}
+
+// HOST-LABEL: @_ZZ15host_lambda_refPiENKUlvE_clEv(
+// HOST: %[[VAL:.*]] = load i32, i32* @global_host_var
+// HOST: %[[VAL2:.*]] = add nsw i32 %[[VAL]], 1
+// HOST: store i32 %[[VAL2]], i32* @global_host_var
+// HOST: %[[VAL:.*]] = load i32, i32* @global_host_var
+// HOST: store i32 %[[VAL]]
+void host_lambda_ref(int *out) {
+  [=](){
+    int &ref = global_host_var;
+    ref++;
+    *out = ref;
+  }();
+}
+
+// HOST-LABEL: define void @_Z28dev_capture_host_ref_by_copyPi(
+// HOST: %[[CAP:.*]] = getelementptr inbounds %[[T3]], %[[T3]]* %{{.*}}, i32 0, i32 1
+// HOST: %[[VAL:.*]] = load i32, i32* @global_host_var
+// HOST: store i32 %[[VAL]], i32* %[[CAP]]
+// DEV-LABEL: define internal void @_ZZ28dev_capture_host_ref_by_copyPiENKUlvE_clEv(
+// DEV: %[[CAP:.*]] = getelementptr inbounds %[[T3]], %[[T3]]* %this1, i32 0, i32 1
+// DEV: %[[VAL:.*]] = load i32, i32* %[[CAP]]
+// DEV: store i32 %[[VAL]]
+void dev_capture_host_ref_by_copy(int *out) {
+  int &ref = global_host_var;
+  kern<<<1, 1>>>([=]__device__() { *out = ref;});
+}
+