[HLSL] Initial codegen for SV_GroupIndex

Semantic parameters aren't passed as actual parameters, instead they are
populated from intrinsics which are generally lowered to reads from
dedicated hardware registers.

This change modifies clang CodeGen to emit the intrinsic calls and
populate the parameter's LValue with the result of the intrinsic call
for SV_GroupIndex.

The result of this is to make the actual passed argument ignored, which
will make it easy to clean up later in an IR pass.

Reviewed By: aaron.ballman

Differential Revision: https://reviews.llvm.org/D131203

clang/lib/AST/Mangle.cpp

@@ -133,10 +133,6 @@ bool MangleContext::shouldMangleDeclName(const NamedDecl *D) {
   if (isa<MSGuidDecl>(D))
     return true;
 
-  // HLSL shader entry function never need to be mangled.
-  if (getASTContext().getLangOpts().HLSL && D->hasAttr<HLSLShaderAttr>())
-    return false;
-
   return shouldMangleCXXName(D);
 }
 

clang/lib/CodeGen/CGHLSLRuntime.cpp

@@ -14,7 +14,9 @@
 
 #include "CGHLSLRuntime.h"
 #include "CodeGenModule.h"
+#include "clang/AST/Decl.h"
 #include "clang/Basic/TargetOptions.h"
+#include "llvm/IR/IntrinsicsDirectX.h"
 #include "llvm/IR/Metadata.h"
 #include "llvm/IR/Module.h"
 
@@ -87,11 +89,55 @@ void CGHLSLRuntime::annotateHLSLResource(const VarDecl *D, GlobalVariable *GV) {
             ConstantAsMetadata::get(B.getInt32(Counter))}));
 }
 
-void clang::CodeGen::CGHLSLRuntime::setHLSLFunctionAttributes(
-    llvm::Function *F, const FunctionDecl *FD) {
-  if (HLSLShaderAttr *ShaderAttr = FD->getAttr<HLSLShaderAttr>()) {
-    const StringRef ShaderAttrKindStr = "hlsl.shader";
-    F->addFnAttr(ShaderAttrKindStr,
-                 ShaderAttr->ConvertShaderTypeToStr(ShaderAttr->getType()));
-  }
+void clang::CodeGen::CGHLSLRuntime::setHLSLEntryAttributes(
+    const FunctionDecl *FD, llvm::Function *Fn) {
+  const auto *ShaderAttr = FD->getAttr<HLSLShaderAttr>();
+  assert(ShaderAttr && "All entry functions must have a HLSLShaderAttr");
+  const StringRef ShaderAttrKindStr = "hlsl.shader";
+  Fn->addFnAttr(ShaderAttrKindStr,
+                ShaderAttr->ConvertShaderTypeToStr(ShaderAttr->getType()));
+}
+
+llvm::Value *CGHLSLRuntime::emitInputSemantic(IRBuilder<> &B,
+                                              const ParmVarDecl &D) {
+  assert(D.hasAttrs() && "Entry parameter missing annotation attribute!");
+  if (D.hasAttr<HLSLSV_GroupIndexAttr>()) {
+    llvm::Function *DxGroupIndex =
+        CGM.getIntrinsic(Intrinsic::dx_flattened_thread_id_in_group);
+    return B.CreateCall(FunctionCallee(DxGroupIndex));
+  }
+  assert(false && "Unhandled parameter attribute");
+  return nullptr;
+}
+
+void CGHLSLRuntime::emitEntryFunction(const FunctionDecl *FD,
+                                      llvm::Function *Fn) {
+  llvm::Module &M = CGM.getModule();
+  llvm::LLVMContext &Ctx = M.getContext();
+  auto *EntryTy = llvm::FunctionType::get(llvm::Type::getVoidTy(Ctx), false);
+  Function *EntryFn =
+      Function::Create(EntryTy, Function::ExternalLinkage, FD->getName(), &M);
+
+  // Copy function attributes over, we have no argument or return attributes
+  // that can be valid on the real entry.
+  AttributeList NewAttrs = AttributeList::get(Ctx, AttributeList::FunctionIndex,
+                                              Fn->getAttributes().getFnAttrs());
+  EntryFn->setAttributes(NewAttrs);
+  setHLSLEntryAttributes(FD, EntryFn);
+
+  // Set the called function as internal linkage.
+  Fn->setLinkage(GlobalValue::InternalLinkage);
+
+  BasicBlock *BB = BasicBlock::Create(Ctx, "entry", EntryFn);
+  IRBuilder<> B(BB);
+  llvm::SmallVector<Value *> Args;
+  // FIXME: support struct parameters where semantics are on members.
+  for (const auto *Param : FD->parameters()) {
+    Args.push_back(emitInputSemantic(B, *Param));
+  }
+
+  CallInst *CI = B.CreateCall(FunctionCallee(Fn), Args);
+  (void)CI;
+  // FIXME: Handle codegen for return type semantics.
+  B.CreateRetVoid();
 }

clang/lib/CodeGen/CGHLSLRuntime.h

@@ -15,6 +15,8 @@
 #ifndef LLVM_CLANG_LIB_CODEGEN_CGHLSLRUNTIME_H
 #define LLVM_CLANG_LIB_CODEGEN_CGHLSLRUNTIME_H
 
+#include "llvm/IR/IRBuilder.h"
+
 #include "clang/Basic/HLSLRuntime.h"
 
 namespace llvm {
@@ -23,6 +25,7 @@ class Function;
 } // namespace llvm
 namespace clang {
 class VarDecl;
+class ParmVarDecl;
 
 class FunctionDecl;
 
@@ -36,6 +39,8 @@ protected:
   uint32_t ResourceCounters[static_cast<uint32_t>(
       hlsl::ResourceClass::NumClasses)] = {0};
 
+  llvm::Value *emitInputSemantic(llvm::IRBuilder<> &B, const ParmVarDecl &D);
+
 public:
   CGHLSLRuntime(CodeGenModule &CGM) : CGM(CGM) {}
   virtual ~CGHLSLRuntime() {}
@@ -44,7 +49,9 @@ public:
 
   void finishCodeGen();
 
-  void setHLSLFunctionAttributes(llvm::Function *, const FunctionDecl *);
+  void setHLSLEntryAttributes(const FunctionDecl *FD, llvm::Function *Fn);
+
+  void emitEntryFunction(const FunctionDecl *FD, llvm::Function *Fn);
 };
 
 } // namespace CodeGen

clang/lib/CodeGen/CodeGenFunction.cpp

@@ -16,6 +16,7 @@
 #include "CGCXXABI.h"
 #include "CGCleanup.h"
 #include "CGDebugInfo.h"
+#include "CGHLSLRuntime.h"
 #include "CGOpenMPRuntime.h"
 #include "CodeGenModule.h"
 #include "CodeGenPGO.h"
@@ -1137,6 +1138,10 @@ void CodeGenFunction::StartFunction(GlobalDecl GD, QualType RetTy,
   if (getLangOpts().OpenMP && CurCodeDecl)
     CGM.getOpenMPRuntime().emitFunctionProlog(*this, CurCodeDecl);
 
+  // Handle emitting HLSL entry functions.
+  if (D && D->hasAttr<HLSLShaderAttr>())
+    CGM.getHLSLRuntime().emitEntryFunction(FD, Fn);
+
   EmitFunctionProlog(*CurFnInfo, CurFn, Args);
 
   if (isa_and_nonnull<CXXMethodDecl>(D) &&

clang/lib/CodeGen/CodeGenModule.cpp

@@ -1703,10 +1703,6 @@ void CodeGenModule::SetLLVMFunctionAttributes(GlobalDecl GD,
                          /*AttrOnCallSite=*/false, IsThunk);
   F->setAttributes(PAL);
   F->setCallingConv(static_cast<llvm::CallingConv::ID>(CallingConv));
-  if (getLangOpts().HLSL) {
-    if (const FunctionDecl *FD = dyn_cast_or_null<FunctionDecl>(GD.getDecl()))
-      getHLSLRuntime().setHLSLFunctionAttributes(F, FD);
-  }
 }
 
 static void removeImageAccessQualifier(std::string& TyName) {

clang/test/CodeGenHLSL/semantics/GroupIndex-codegen.hlsl

@@ -0,0 +1,22 @@
+// RUN: %clang_cc1 -triple dxil-pc-shadermodel6.0-compute -x hlsl -emit-llvm -disable-llvm-passes -o - -hlsl-entry main %s
+
+[numthreads(1,1,1)]
+void main(unsigned GI : SV_GroupIndex) {
+  main(GI - 1);
+}
+
+// For HLSL entry functions, we are generating a C-export function that wraps
+// the C++-mangled entry function. The wrapper function can be used to populate
+// semantic parameters and provides the expected void(void) signature that
+// drivers expect for entry points.
+
+//CHECK: define void @main() #[[ENTRY_ATTR:#]]{
+//CHECK-NEXT: entry:
+//CHECK-NEXT:   %0 = call i32 @llvm.dx.flattened.thread.id.in.group()
+//CHECK-NEXT:   call void @"?main@@YAXI@Z"(i32 %0)
+//CHECK-NEXT:   ret void
+//CHECK-NEXT: }
+
+// Verify that the entry had the expected dx.shader attribute
+
+//CHECK: attributes #[[ENTRY_ATTR]] = { {{.*}}"dx.shader"="compute"{{.*}} }