[CUDA] Add #pragma clang force_cuda_host_device_{begin,end} pragmas.

[jlebar]

Summary

These cause us to consider all functions in-between to be host
device.

You can nest these pragmas; you just can't have more 'end's than
'begin's.

[tra]

LGTM. Should we add new pragma description to docs/LanguageExtensions.rst ?

[rsmith]

clang/lib/Parse/ParsePragma.cpp

172

I think this would be more consistent with the behavior of our other pragmas as two tokens rather than one:

#pragma clang force_cuda_host_device begin
#pragma clang force_cuda_host_device end

It's also more extensible this way, for instance if we later want some kind of push/pop mechanism to temporarily leave force-host-device mode. (Maybe we could even use #pragma clang cuda force_host_device begin, but I'm fine with either choice there.)

clang/lib/Sema/SemaCUDA.cpp

466–472

Does this apply within template instantiations? I'm concerned about two cases:

1. Eagerly-instantiated declarations whose instantiation is triggered within one of these scopes:

template<typename T> auto f() { return T(); }
template<typename T> struct X { void f(); };
#pragma clang force_cuda_host_device_begin
int n = f<int>(); // is f<int> implicitly __host__ __device__?
X<int> x; // is X<int>::f implicitly __host__ __device__?
#pragma clang force_cuda_host_device_end

2. Template instantiation at the end of the TU if we have a begin but no end:

template<typename T> int f() { return T(); }
#pragma clang force_cuda_host_device_begin
int n = f<int>(); // is f<int> implicitly __host__ __device__?

I would expect the answer in each case to be that the function is host. Even if these do work correctly already, please add some testcases.

[jlebar]

Address Richard Smith's review comments:

• Change macro format.
• Add tests (these Just Worked).

[rsmith]

Please add serialisation code for the push count for PCH. Should it be an error if the count is nonzero at the end of the TU?

What happens if there are trailing tokens after the pragma?

clang/include/clang/Basic/DiagnosticParseKinds.td

1027

Diagnostics should start with a lowercase letter.

clang/test/Parser/cuda-force-host-device-templates.cu

8 ↗

(On Diff #72717)

You need this to return auto to trigger the eager instantiation codepath I was concerned about.

[jlebar]

Address Richard's comments.

I'm fairly neutral on whether we want to make it an error not to match all of
your "begin" pragmas with "end"s. I checked pragma push_macro, and it looks
like it's not an error to pop those, so with that prior art, and since it was
simpler not to check for matching begin/ends, I did the same. But like I say,
I don't feel strongly either way (or even if we wanted to make these new
pragmas not-nestable).

[jlebar]

What happens if there are trailing tokens after the pragma?

Added code to make this an error.

[rsmith]

Please add a test to test/PCH for the serialization code. Otherwise, LGTM.

clang/include/clang/Basic/DiagnosticParseKinds.td

1032

Diagnostics should not end with a period. =)

clang/include/clang/Basic/DiagnosticSemaKinds.td

6701–6702 ↗

(On Diff #72734)

This appears to be unused.

[jlebar]

In D24975#565054, @rsmith wrote:

Please add a test to test/PCH for the serialization code. Otherwise, LGTM.

Test added. It caught a bug, too. :)

Thank you for the review.

clang/include/clang/Basic/DiagnosticParseKinds.td

def err_pragma_loop_invalid_option : Error<
   "vectorize_width, interleave, interleave_count, unroll, unroll_count, or distribute">;
 def err_pragma_invalid_keyword : Error<
   "invalid argument; expected 'enable'%select{|, 'full'}0%select{|, 'assume_safety'}1 or 'disable'">;
 
 // Pragma unroll support.
 def warn_pragma_unroll_cuda_value_in_parens : Warning<
   "argument to '#pragma unroll' should not be in parentheses in CUDA C/C++">,
   InGroup<CudaCompat>;
+
+def err_pragma_cannot_end_force_cuda_host_device : Error<
+  "force_cuda_host_device_end pragma without matching "
+  "force_cuda_host_device_begin.">;
 } // end of Parse Issue category.
 
 let CategoryName = "Modules Issue" in {
 def err_expected_module_interface_decl : Error<
   "expected module declaration at start of module interface">;
 def err_unexpected_module_decl : Error<
   "module declaration must be the first declaration in the translation unit">;
 def err_module_expected_ident : Error<
   "expected a module name after module%select{| import}0">;
 def err_unexpected_module_kind : Error<
   "unexpected module kind %0; expected 'implementation' or 'partition'">;
 def err_attribute_not_module_attr : Error<

clang/include/clang/Parse/Parser.h

class Parser : public CodeCompletionHandler {
   std::unique_ptr<PragmaHandler> MSInitSeg;
   std::unique_ptr<PragmaHandler> MSDataSeg;
   std::unique_ptr<PragmaHandler> MSBSSSeg;
   std::unique_ptr<PragmaHandler> MSConstSeg;
   std::unique_ptr<PragmaHandler> MSCodeSeg;
   std::unique_ptr<PragmaHandler> MSSection;
   std::unique_ptr<PragmaHandler> MSRuntimeChecks;
   std::unique_ptr<PragmaHandler> MSIntrinsic;
+  std::unique_ptr<PragmaHandler> CUDAForceHostDeviceStartHandler;
+  std::unique_ptr<PragmaHandler> CUDAForceHostDeviceEndHandler;
   std::unique_ptr<PragmaHandler> OptimizeHandler;
   std::unique_ptr<PragmaHandler> LoopHintHandler;
   std::unique_ptr<PragmaHandler> UnrollHintHandler;
   std::unique_ptr<PragmaHandler> NoUnrollHintHandler;
 
   std::unique_ptr<CommentHandler> CommentSemaHandler;
 
   /// Whether the '>' token acts as an operator or not. This will be

clang/include/clang/Sema/Sema.h

   /// VerifyBitField - verifies that a bit field expression is an ICE and has
   /// the correct width, and that the field type is valid.
   /// Returns false on success.
   /// Can optionally return whether the bit-field is of width 0
   ExprResult VerifyBitField(SourceLocation FieldLoc, IdentifierInfo *FieldName,
                             QualType FieldTy, bool IsMsStruct,
                             Expr *BitWidth, bool *ZeroWidth = nullptr);
 
+private:
+  unsigned ForceCUDAHostDeviceDepth = 0;
+
+public:
+  /// Increments our count of the number of times we've seen a pragma forcing
+  /// functions to be __host__ __device__.  So long as this count is greater
+  /// than zero, all functions encountered will be __host__ __device__.
+  void PushForceCUDAHostDevice();
+
+  /// Decrements our count of the number of times we've seen a pragma forcing
+  /// functions to be __host__ __device__.  Returns false if the count is 0
+  /// before incrementing, so you can emit an error.
+  bool PopForceCUDAHostDevice();
+
   enum CUDAFunctionTarget {
     CFT_Device,
     CFT_Global,
     CFT_Host,
     CFT_HostDevice,
     CFT_InvalidTarget
   };
 

clang/lib/Parse/ParsePragma.cpp

 };
 
 struct PragmaMSIntrinsicHandler : public PragmaHandler {
   PragmaMSIntrinsicHandler() : PragmaHandler("intrinsic") {}
   void HandlePragma(Preprocessor &PP, PragmaIntroducerKind Introducer,
                     Token &FirstToken) override;
 };
 
+struct PragmaForceCUDAHostDeviceStartHandler : public PragmaHandler {
+  PragmaForceCUDAHostDeviceStartHandler(Sema &Actions)
+      : PragmaHandler("force_cuda_host_device_begin"), Actions(Actions) {}
+  void HandlePragma(Preprocessor &PP, PragmaIntroducerKind Introducer,
+                    Token &NameTok) override;
+
+private:
+  Sema &Actions;
+};
+
+struct PragmaForceCUDAHostDeviceEndHandler : public PragmaHandler {
+  PragmaForceCUDAHostDeviceEndHandler(Sema &Actions)
+      : PragmaHandler("force_cuda_host_device_end"), Actions(Actions) {}
+  void HandlePragma(Preprocessor &PP, PragmaIntroducerKind Introducer,
+                    Token &NameTok) override;
+
+private:
+  Sema &Actions;
+};
+
 }  // end namespace
 
 void Parser::initializePragmaHandlers() {
   AlignHandler.reset(new PragmaAlignHandler());
   PP.AddPragmaHandler(AlignHandler.get());
 
   GCCVisibilityHandler.reset(new PragmaGCCVisibilityHandler());
   PP.AddPragmaHandler("GCC", GCCVisibilityHandler.get());
if (getLangOpts().MicrosoftExt) {
     MSSection.reset(new PragmaMSPragma("section"));
     PP.AddPragmaHandler(MSSection.get());
     MSRuntimeChecks.reset(new PragmaMSRuntimeChecksHandler());
     PP.AddPragmaHandler(MSRuntimeChecks.get());
     MSIntrinsic.reset(new PragmaMSIntrinsicHandler());
     PP.AddPragmaHandler(MSIntrinsic.get());
   }
 
+  if (getLangOpts().CUDA) {
+    CUDAForceHostDeviceStartHandler.reset(
+        new PragmaForceCUDAHostDeviceStartHandler(Actions));
+    PP.AddPragmaHandler("clang", CUDAForceHostDeviceStartHandler.get());
+    CUDAForceHostDeviceEndHandler.reset(
+        new PragmaForceCUDAHostDeviceEndHandler(Actions));
+    PP.AddPragmaHandler("clang", CUDAForceHostDeviceEndHandler.get());
+  }
+
   OptimizeHandler.reset(new PragmaOptimizeHandler(Actions));
   PP.AddPragmaHandler("clang", OptimizeHandler.get());
 
   LoopHintHandler.reset(new PragmaLoopHintHandler());
   PP.AddPragmaHandler("clang", LoopHintHandler.get());
 
   UnrollHintHandler.reset(new PragmaUnrollHintHandler("unroll"));
   PP.AddPragmaHandler(UnrollHintHandler.get());
if (getLangOpts().MicrosoftExt) {
     PP.RemovePragmaHandler(MSSection.get());
     MSSection.reset();
     PP.RemovePragmaHandler(MSRuntimeChecks.get());
     MSRuntimeChecks.reset();
     PP.RemovePragmaHandler(MSIntrinsic.get());
     MSIntrinsic.reset();
   }
 
+  if (getLangOpts().CUDA) {
+    PP.RemovePragmaHandler("clang", CUDAForceHostDeviceStartHandler.get());
+    CUDAForceHostDeviceStartHandler.reset();
+    PP.RemovePragmaHandler("clang", CUDAForceHostDeviceEndHandler.get());
+    CUDAForceHostDeviceEndHandler.reset();
+  }
+
   PP.RemovePragmaHandler("STDC", FPContractHandler.get());
   FPContractHandler.reset();
 
   PP.RemovePragmaHandler("clang", OptimizeHandler.get());
   OptimizeHandler.reset();
 
   PP.RemovePragmaHandler("clang", LoopHintHandler.get());
   LoopHintHandler.reset();
if (Tok.isNot(tok::r_paren)) {
     return;
   }
   PP.Lex(Tok);
 
   if (Tok.isNot(tok::eod))
     PP.Diag(Tok.getLocation(), diag::warn_pragma_extra_tokens_at_eol)
         << "intrinsic";
 }
+void PragmaForceCUDAHostDeviceStartHandler::HandlePragma(
+    Preprocessor &PP, PragmaIntroducerKind Introducer, Token &NameTok) {
+  Actions.PushForceCUDAHostDevice();
+}
+void PragmaForceCUDAHostDeviceEndHandler::HandlePragma(
+    Preprocessor &PP, PragmaIntroducerKind Introducer, Token &NameTok) {
+  if (!Actions.PopForceCUDAHostDevice())
+    PP.Diag(NameTok.getLocation(),
+            diag::err_pragma_cannot_end_force_cuda_host_device);
+}

clang/lib/Sema/SemaCUDA.cpp

 #include "clang/Sema/Lookup.h"
 #include "clang/Sema/Sema.h"
 #include "clang/Sema/SemaDiagnostic.h"
 #include "clang/Sema/Template.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/SmallVector.h"
 using namespace clang;
 
+void Sema::PushForceCUDAHostDevice() {
+  assert(getLangOpts().CUDA && "May be called only for CUDA compilations.");
+  ForceCUDAHostDeviceDepth++;
+}
+
+bool Sema::PopForceCUDAHostDevice() {
+  assert(getLangOpts().CUDA && "May be called only for CUDA compilations.");
+  if (ForceCUDAHostDeviceDepth == 0)
+    return false;
+  ForceCUDAHostDeviceDepth--;
+  return true;
+}
+
 ExprResult Sema::ActOnCUDAExecConfigExpr(Scope *S, SourceLocation LLLLoc,
                                          MultiExprArg ExecConfig,
                                          SourceLocation GGGLoc) {
   FunctionDecl *ConfigDecl = Context.getcudaConfigureCallDecl();
   if (!ConfigDecl)
     return ExprError(Diag(LLLLoc, diag::err_undeclared_var_use)
                      << "cudaConfigureCall");
   QualType ConfigQTy = ConfigDecl->getType();

 
 // With -fcuda-host-device-constexpr, an unattributed constexpr function is
 // treated as implicitly __host__ __device__, unless:
 //  * it is a variadic function (device-side variadic functions are not
 //    allowed), or
 //  * a __device__ function with this signature was already declared, in which
 //    case in which case we output an error, unless the __device__ decl is in a
 //    system header, in which case we leave the constexpr function unattributed.
+//
+// In addition, all function decls are treated as __host__ __device__ when
+// ForceCUDAHostDeviceDepth > 0 (corresponding to code within a
+//   #pragma clang force_cuda_host_device_begin/end
+// pair).
 void Sema::maybeAddCUDAHostDeviceAttrs(Scope *S, FunctionDecl *NewD,
                                        const LookupResult &Previous) {
   assert(getLangOpts().CUDA && "May be called only for CUDA compilations.");
+
+  if (ForceCUDAHostDeviceDepth > 0) {
+    if (!NewD->hasAttr<CUDAHostAttr>())
+      NewD->addAttr(CUDAHostAttr::CreateImplicit(Context));
+    if (!NewD->hasAttr<CUDADeviceAttr>())
+      NewD->addAttr(CUDADeviceAttr::CreateImplicit(Context));
+    return;
+  }
+
   if (!getLangOpts().CUDAHostDeviceConstexpr || !NewD->isConstexpr() ||
       NewD->isVariadic() || NewD->hasAttr<CUDAHostAttr>() ||
       NewD->hasAttr<CUDADeviceAttr>() || NewD->hasAttr<CUDAGlobalAttr>())
     return;
 
   // Is D a __device__ function with the same signature as NewD, ignoring CUDA
   // attributes?
   auto IsMatchingDeviceFn = [&](NamedDecl *D) {

clang/test/Parser/cuda-force-host-device.cu

+// RUN: %clang_cc1 -fsyntax-only -verify %s
+
+// Check the force_cuda_host_device_{begin,end} pragmas.
+
+#pragma clang force_cuda_host_device_begin
+void f();
+#pragma clang force_cuda_host_device_begin
+void g();
+#pragma clang force_cuda_host_device_end
+void h();
+#pragma clang force_cuda_host_device_end
+
+void i(); // expected-note {{not viable}}
+
+void host() {
+  f();
+  g();
+  h();
+  i();
+}
+
+__attribute__((device)) void device() {
+  f();
+  g();
+  h();
+  i(); // expected-error {{no matching function}}
+}