[HLSL] Initial codegen for SV_GroupIndex

[beanz]

Summary

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.

[python3kgae]

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

5

Does this change plan to support cases where SV_GroupIndex is inside a struct?
If yes, we need a test for that.

[beanz]

This is a fairly significant reworking of this patch to handle entry functions being called (which I hadn't anticipated). There may be some additional tweaks required, but the general flow here is:

(1) Don't prevent mangling of entry functions, this allows us to generate calls based on the mangled function name which makes things happy.
(2) Do codegen the function with parameters following the C++ codegen paths with no alterations
(3) Also generate a c-exported, unmanged, void(void) function for each entry which populates semantic arguments and calls into the mangled entry function

This solution results in code generation that covers the required cases, and the optimizer can (and will) eliminate the wrapper function in all cases that it can, and in any case that it can't the compiler will explode because recurison is not allowed in HLSL unless it can be translated to a loop by the compiler.

Fun right?

This now depends on D131625, which is a small change to enforce that all entry parameters have semantic annotations.

[python3kgae]

clang/lib/CodeGen/CGHLSLRuntime.cpp

19

Does this mean CGHLSLRuntime is only for DirectX backend?

147

FIXME: support struct parameter which may not have Annotation attribute.

163

Need to move all attributes from Fn to EntryFn.
And make Fn a normal function.

[beanz]

clang/lib/CodeGen/CGHLSLRuntime.cpp

19

No, it is fairly normal to include target intrinsics throughout codegen. That said, I'm not handling non-DirectX intrinsics nor am I verifying the appropriate target, so I should make some adjustment here.

163

Yes on the function attributes.

What do you mean by "normal function" as written it is an external linkage function, which matches what we expect in DXIL.

[python3kgae]

clang/lib/CodeGen/CGHLSLRuntime.cpp

163

By "normal function" I mean remove the special attribute like HLSLEntry, numofthreads.

[beanz]

Handling function attributes appropriately.

[aaron.ballman]

clang/lib/CodeGen/CGHLSLRuntime.cpp

92–136

105–106

120–139

133

clang/lib/CodeGen/CGHLSLRuntime.h

52–54

These parameters should have names so people don't think the parameters aren't used.

[beanz]

Updates based on review feedback.

[aaron.ballman]

LGTM!

clang/lib/AST/Mangle.cpp

bool MangleContext::shouldMangleDeclName(const NamedDecl *D) {
   // over all other naming in the .o file.
   if (D->hasAttr<AsmLabelAttr>())
     return true;
 
   // Declarations that don't have identifier names always need to be mangled.
   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);
 }
 
 void MangleContext::mangleName(GlobalDecl GD, raw_ostream &Out) {
   const ASTContext &ASTContext = getASTContext();
   const NamedDecl *D = cast<NamedDecl>(GD.getDecl());
 
   // Any decl can be declared with __asm("foo") on it, and this takes precedence

clang/lib/CodeGen/CGHLSLRuntime.h

 // subclasses of this implement code generation for specific HLSL
 // runtime libraries.
 //
 //===----------------------------------------------------------------------===//
 
 #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 {
 class GlobalVariable;
 class Function;
 } // namespace llvm
 namespace clang {
 class VarDecl;
+class ParmVarDecl;
 
 class FunctionDecl;
 
 namespace CodeGen {
 
 class CodeGenModule;
 
 class CGHLSLRuntime {
 protected:
   CodeGenModule &CGM;
   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() {}
 
   void annotateHLSLResource(const VarDecl *D, llvm::GlobalVariable *GV);
 
   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
 } // namespace clang
 
 #endif

clang/lib/CodeGen/CGHLSLRuntime.cpp

 //===----- CGHLSLRuntime.cpp - Interface to HLSL Runtimes -----------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This provides an abstract class for HLSL code generation.  Concrete
 // subclasses of this implement code generation for specific HLSL
 // runtime libraries.
 //
 //===----------------------------------------------------------------------===//
 
 #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"
 
 using namespace clang;
 using namespace CodeGen;
 using namespace hlsl;
 using namespace llvm;
 
void CGHLSLRuntime::annotateHLSLResource(const VarDecl *D, GlobalVariable *GV) {
   auto &Ctx = CGM.getModule().getContext();
   IRBuilder<> B(Ctx);
   QualType QT(Ty, 0);
   ResourceMD->addOperand(MDNode::get(
       Ctx, {ValueAsMetadata::get(GV), MDString::get(Ctx, QT.getAsString()),
             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");
+  Fn->addFnAttr(ShaderAttrKindStr,
+
+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/CodeGenFunction.cpp

 //===----------------------------------------------------------------------===//
 
 #include "CodeGenFunction.h"
 #include "CGBlocks.h"
 #include "CGCUDARuntime.h"
 #include "CGCXXABI.h"
 #include "CGCleanup.h"
 #include "CGDebugInfo.h"
+#include "CGHLSLRuntime.h"
 #include "CGOpenMPRuntime.h"
 #include "CodeGenModule.h"
 #include "CodeGenPGO.h"
 #include "TargetInfo.h"
 #include "clang/AST/ASTContext.h"
 #include "clang/AST/ASTLambda.h"
 #include "clang/AST/Attr.h"
 #include "clang/AST/Decl.h"
#undef SANITIZER
   EmitStartEHSpec(CurCodeDecl);
 
   PrologueCleanupDepth = EHStack.stable_begin();
 
   // Emit OpenMP specific initialization of the device functions.
   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) &&
       cast<CXXMethodDecl>(D)->isInstance()) {
     CGM.getCXXABI().EmitInstanceFunctionProlog(*this);
     const CXXMethodDecl *MD = cast<CXXMethodDecl>(D);
     if (MD->getParent()->isLambda() &&
         MD->getOverloadedOperator() == OO_Call) {

clang/lib/CodeGen/CodeGenModule.cpp

void CodeGenModule::SetLLVMFunctionAttributes(GlobalDecl GD,
                                               const CGFunctionInfo &Info,
                                               llvm::Function *F, bool IsThunk) {
   unsigned CallingConv;
   llvm::AttributeList PAL;
   ConstructAttributeList(F->getName(), Info, GD, PAL, CallingConv,
                          /*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) {
   std::string ReadOnlyQual("__read_only");
   std::string::size_type ReadOnlyPos = TyName.find(ReadOnlyQual);
   if (ReadOnlyPos != std::string::npos)
     // "+ 1" for the space after access qualifier.
     TyName.erase(ReadOnlyPos, ReadOnlyQual.size() + 1);

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

+// 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"{{.*}} }