[CUDA][HIP] Do not treat host var address as constant in device compilation

[yaxunl]

Summary

Currently clang treats host var address as constant in device compilation,
which causes const vars initialized with host var address promoted to
device variables incorrectly and results in undefined symbols.

This patch fixes that.

[yaxunl]

fix test

[tra]

LGTM.

Do we need to do anything special about __managed__ vars?

clang/lib/AST/ExprConstant.cpp

2224

Nit: Negations are my pet peeve. I think if (!(A||B||C||D...)) would be easier to read -- not (this or that or that) has only N+1 operations to keep in one's head vs 2*N in not this, and not that, and not that, and..... In this case it's not too bad, so I'll leave it up to you which form you prefer.

[yaxunl]

In D118153#3270122, @tra wrote:

LGTM.

Do we need to do anything special about __managed__ vars?

Right __managed__ var is special. Its address is set by runtime, therefore it is not a constant. nvcc does not treat it as constant either. https://godbolt.org/z/jK534MxfG I will fix it.

BTW I just found a regression related to templates in CodeGenCUDA/host-used-device-var.cu. I will fix that and update this patch.

[yaxunl]

Fix the regression in lit tests.

Basically in device compilation we still evaluate constant expression for host functions or host template instantiation. If we just disallow host variable in any constant expressions we will get errors in template class instantiation which use host variables as non-type template arguments.

Therefore we should only disallow host variables in constant expressions in situations when we are sure that allowing them will lead to issues, e.g. when promoting const variables.

A CUDAConstantEvaluationContext is introduced in ASTContext to control this.

[tra]

@rsmith -- is there anything else we need to worry about when it comes to treating pointers as constant values (or not)?

clang/lib/AST/ExprConstant.cpp

2227

Does it mean that we currently treat address of a __shared__ variable as a constant? Looks like we do: https://godbolt.org/z/eG4vG1rbf

__shared__ int s;
__device__ const int *const p2 = &s;

__device__ bool f() { return p2 == &s; }

f() currently always returns true, even though there will be multiple instances of s and we probably should not have allowed to init p2 to start with.

I think we should always return false for __shared__ vars.

clang/test/SemaCUDA/const-var.cu

2–3

Would it make sense to add a host-side test to illustrate a case where some pointer expressions will be const on the host side, but not on the device side?

18

I'd add a comment that global const variables are treated as __constant__. Otherwise no error on taking address of a on device side looks wrong.

Makes me wonder what would be result of B::p == B::p2? On device side, it should be possible to const evaluate it to true.
However, on the host, B::p would have a real value, but B::p2 will be represented by a proxy var. Off the top of my head I'm not sure whether it will be initialized or left undefined.

I'd add a codegen test to make sure B::p == B::p2 does get evaluated consistently on host and device.

[yaxunl]

clang/lib/AST/ExprConstant.cpp

2227

Will do. This is similar to a TLS var in C++, which is treated as non-constant by clang.

clang/test/SemaCUDA/const-var.cu

2–3

will do.

18

will do.

B::p == B::p2 evaluates to 1 on both device and host side.

[yaxunl]

Revised by Artem's comments.

clang/include/clang/AST/ASTContext.h

public:
   const TranslationUnitKind TUKind;
   mutable DeclarationNameTable DeclarationNames;
   IntrusiveRefCntPtr<ExternalASTSource> ExternalSource;
   ASTMutationListener *Listener = nullptr;
 
   /// Returns the clang bytecode interpreter context.
   interp::Context &getInterpContext();
 
+  struct CUDAConstantEvalContext {
+    /// Do not allow wrong-sided variables in constant expressions.
+    bool NoWrongSidedVars = false;
+  } CUDAConstantEvalCtx;
+  struct CUDAConstantEvalContextRAII {
+    ASTContext &Ctx;
+    CUDAConstantEvalContext SavedCtx;
+    CUDAConstantEvalContextRAII(ASTContext &Ctx_, bool NoWrongSidedVars)
+        : Ctx(Ctx_), SavedCtx(Ctx_.CUDAConstantEvalCtx) {
+      Ctx_.CUDAConstantEvalCtx.NoWrongSidedVars = NoWrongSidedVars;
+    }
+    ~CUDAConstantEvalContextRAII() { Ctx.CUDAConstantEvalCtx = SavedCtx; }
+  };
+
   /// Returns the dynamic AST node parent map context.
   ParentMapContext &getParentMapContext();
 
   // A traversal scope limits the parts of the AST visible to certain analyses.
   // RecursiveASTVisitor only visits specified children of TranslationUnitDecl.
   // getParents() will only observe reachable parent edges.
   //
   // The scope is defined by a set of "top-level" declarations which will be

clang/lib/AST/ExprConstant.cpp

EvalInfo(const ASTContext &C, Expr::EvalStatus &S, EvaluationMode Mode)
           EvaluatingDeclValue(nullptr), HasActiveDiagnostic(false),
           HasFoldFailureDiagnostic(false), InConstantContext(false),
           EvalMode(Mode) {}
 
     ~EvalInfo() {
       discardCleanups();
     }
 
+    ASTContext &getCtx() const override { return Ctx; }
+
     void setEvaluatingDecl(APValue::LValueBase Base, APValue &Value,
                            EvaluatingDeclKind EDK = EvaluatingDeclKind::Ctor) {
       EvaluatingDecl = Base;
       IsEvaluatingDecl = EDK;
       EvaluatingDeclValue = &Value;
     }
 
     bool CheckCallLimit(SourceLocation Loc) {
private:
     void setActiveDiagnostic(bool Flag) override { HasActiveDiagnostic = Flag; }
 
     void setFoldFailureDiagnostic(bool Flag) override {
       HasFoldFailureDiagnostic = Flag;
     }
 
     Expr::EvalStatus &getEvalStatus() const override { return EvalStatus; }
 
-    ASTContext &getCtx() const override { return Ctx; }
-
     // If we have a prior diagnostic, it will be noting that the expression
     // isn't a constant expression. This diagnostic is more important,
     // unless we require this evaluation to produce a constant expression.
     //
     // FIXME: We might want to show both diagnostics to the user in
     // EM_ConstantFold mode.
     bool hasPriorDiagnostic() override {
       if (!EvalStatus.Diag->empty()) {
if (const VarDecl *Var = dyn_cast<const VarDecl>(BaseVD)) {
         // FIXME: Diagnostic!
         return false;
 
       // A dllimport variable never acts like a constant, unless we're
       // evaluating a value for use only in name mangling.
       if (!isForManglingOnly(Kind) && Var->hasAttr<DLLImportAttr>())
         // FIXME: Diagnostic!
         return false;
+
+      // In CUDA/HIP device compilation, only device side variables have
+      // constant addresses.
+      if (Info.getCtx().getLangOpts().CUDA &&
+          Info.getCtx().getLangOpts().CUDAIsDevice &&
+          Info.getCtx().CUDAConstantEvalCtx.NoWrongSidedVars) {
+        if ((!Var->hasAttr<CUDADeviceAttr>() &&
+             !Var->hasAttr<CUDAConstantAttr>() &&
+             !Var->getType()->isCUDADeviceBuiltinSurfaceType() &&
+             !Var->getType()->isCUDADeviceBuiltinTextureType()) ||
+            Var->hasAttr<HIPManagedAttr>())
+          return false;
+      }
     }
     if (const auto *FD = dyn_cast<const FunctionDecl>(BaseVD)) {
       // __declspec(dllimport) must be handled very carefully:
       // We must never initialize an expression with the thunk in C++.
       // Doing otherwise would allow the same id-expression to yield
       // different addresses for the same function in different translation
       // units.  However, this means that we must dynamically initialize the
       // expression with the contents of the import address table at runtime.

clang/lib/Sema/SemaCUDA.cpp

if (!Init)
       return true;
     if (const auto *CE = dyn_cast<CXXConstructExpr>(Init)) {
       return S.isEmptyCudaConstructor(VD->getLocation(), CE->getConstructor());
     }
     return false;
   };
   auto IsConstantInit = [&](const Expr *Init) {
     assert(Init);
+    ASTContext::CUDAConstantEvalContextRAII EvalCtx(S.Context,
+                                                    /*NoWronSidedVars=*/true);
     return Init->isConstantInitializer(S.Context,
                                        VD->getType()->isReferenceType());
   };
   auto HasEmptyDtor = [&](VarDecl *VD) {
     if (const auto *RD = VD->getType()->getAsCXXRecordDecl())
       return S.isEmptyCudaDestructor(VD->getLocation(), RD->getDestructor());
     return true;
   };

clang/test/CodeGenCUDA/const-var.cu

+// REQUIRES: amdgpu-registered-target
+// RUN: %clang_cc1 -triple amdgcn-amd-amdhsa -fcuda-is-device -x hip %s \
+// RUN:   -emit-llvm -o - | FileCheck -check-prefix=DEV %s
+// RUN: %clang_cc1 -triple x86_64-unknown-linux-gnu -x hip %s \
+// RUN:   -emit-llvm -o - | FileCheck -check-prefix=HOST %s
+
+// Negative tests.
+
+// RUN: %clang_cc1 -triple amdgcn-amd-amdhsa -fcuda-is-device -x hip %s \
+// RUN:   -emit-llvm -o - | FileCheck -check-prefix=DEV-NEG %s
+
+#include "Inputs/cuda.h"
+
+// Test const var initialized with address of a const var.
+// Both are promoted to device side.
+
+// DEV-DAG: @_ZN5Test1L1aE = internal addrspace(4) constant i32 1
+// DEV-DAG: @_ZN5Test11B2p1E = addrspace(4) externally_initialized constant i32* addrspacecast (i32 addrspace(4)* @_ZN5Test1L1aE to i32*)
+// DEV-DAG: @_ZN5Test11B2p2E = addrspace(4) externally_initialized constant i32* addrspacecast (i32 addrspace(4)* @_ZN5Test1L1aE to i32*)
+// DEV-DAG: @_ZN5Test12b2E = addrspace(1) externally_initialized global i32 1
+// HOST-DAG: @_ZN5Test1L1aE = internal constant i32 1
+// HOST-DAG: @_ZN5Test11B2p1E = constant i32* @_ZN5Test1L1aE
+// HOST-DAG: @_ZN5Test11B2p2E = internal constant i32* undef
+// HOST-DAG: @_ZN5Test12b1E = global i32 1
+// HOST-DAG: @_ZN5Test12b2E = internal global i32 undef
+namespace Test1 {
+const int a = 1;
+
+struct B {
+    static const int *const p1;
+    static __device__ const int *const p2;
+};
+const int *const B::p1 = &a;
+__device__ const int *const B::p2 = &a;
+int b1 = B::p1 == B::p2;
+__device__ int b2 = B::p1 == B::p2;
+}
+
+// Test const var initialized with address of a non-cost var.
+// Neither is promoted to device side.
+
+// DEV-NEG-NOT: @_ZN5Test2L1aE
+// DEV-NEG-NOT: @_ZN5Test21B1pE
+// HOST-DAG: @_ZN5Test21aE = global i32 1
+// HOST-DAG: @_ZN5Test21B1pE = constant i32* @_ZN5Test21aE
+
+namespace Test2 {
+int a = 1;
+
+struct B {
+    static int *const p;
+};
+int *const B::p = &a;
+}

clang/test/SemaCUDA/const-var.cu

+// RUN: %clang_cc1 -triple amdgcn-amd-amdhsa -fcuda-is-device -x hip %s \
+// RUN:   -fsyntax-only -verify
+// RUN: %clang_cc1 -triple x86_64 -x hip %s \
+// RUN:   -fsyntax-only -verify=host
+
+// host-no-diagnostics
+
+#include "Inputs/cuda.h"
+
+// Test const var initialized with address of a const var.
+// Both are promoted to device side.
+
+namespace Test1 {
+const int a = 1;
+
+struct B {
+    static const int *const p;
+    __device__ static const int *const p2;
+};
+const int *const B::p = &a;
+// Const variable 'a' is treated as __constant__ on device side,
+// therefore its address can be used as initializer for another
+// device variable.
+__device__ const int *const B::p2 = &a;
+
+__device__ void f() {
+  int y = a;
+  const int *x = B::p;
+  const int *z = B::p2;
+}
+}
+
+// Test const var initialized with address of a non-cost var.
+// Neither is promoted to device side.
+
+namespace Test2 {
+int a = 1;
+// expected-note@-1{{host variable declared here}}
+
+struct B {
+    static int *const p;
+};
+int *const B::p = &a;
+// expected-note@-1{{const variable cannot be emitted on device side due to dynamic initialization}}
+
+__device__ void f() {
+  int y = a;
+  // expected-error@-1{{reference to __host__ variable 'a' in __device__ function}}
+  const int *x = B::p;
+  // expected-error@-1{{reference to __host__ variable 'p' in __device__ function}}
+}
+}
+
+// Test device var initialized with address of a non-const host var, __shared var,
+// __managed__ var, __device__ var, __constant__ var, texture var, surface var.
+
+namespace Test3 {
+struct textureReference {
+  int desc;
+};
+
+enum ReadMode {
+  ElementType = 0,
+  NormalizedFloat = 1
+};
+
+template <typename T, int dim = 1, enum ReadMode mode = ElementType>
+struct __attribute__((device_builtin_texture_type)) texture : public textureReference {
+};
+
+struct surfaceReference {
+  int desc;
+};
+
+template <typename T, int dim = 1>
+struct __attribute__((device_builtin_surface_type)) surface : public surfaceReference {
+};
+
+// Partial specialization over `void`.
+template<int dim>
+struct __attribute__((device_builtin_surface_type)) surface<void, dim> : public surfaceReference {
+};
+
+texture<float, 2, ElementType> tex;
+surface<void, 2> surf;
+
+int a = 1;
+__shared__ int b;
+__managed__ int c = 1;
+__device__ int d = 1;
+__constant__ int e = 1;
+struct B {
+    __device__ static int *const p1;
+    __device__ static int *const p2;
+    __device__ static int *const p3;
+    __device__ static int *const p4;
+    __device__ static int *const p5;
+    __device__ static texture<float, 2, ElementType> *const p6;
+    __device__ static surface<void, 2> *const p7;
+};
+__device__ int *const B::p1 = &a;
+// expected-error@-1{{dynamic initialization is not supported for __device__, __constant__, __shared__, and __managed__ variables}}
+__device__ int *const B::p2 = &b;
+// expected-error@-1{{dynamic initialization is not supported for __device__, __constant__, __shared__, and __managed__ variables}}
+__device__ int *const B::p3 = &c;
+// expected-error@-1{{dynamic initialization is not supported for __device__, __constant__, __shared__, and __managed__ variables}}
+__device__ int *const B::p4 = &d;
+__device__ int *const B::p5 = &e;
+__device__ texture<float, 2, ElementType> *const B::p6 = &tex;
+__device__ surface<void, 2> *const B::p7 = &surf;
+}