[HLSL] Add SV_DispatchThreadID

[python3kgae]

Summary

Support SV_DispatchThreadID attribute.
Translate it into dx.thread.id in clang codeGen.

[python3kgae]

Fix warning.

[python3kgae]

Rebase

[aaron.ballman]

There are no tests for applying this to a global variable, so those should be added.

clang/include/clang/Basic/AttrDocs.td

6596–6598

It doesn't specify what happens when applied to a global variable.

clang/lib/CodeGen/CGHLSLRuntime.cpp

113

120–121

clang/lib/Sema/SemaDeclAttr.cpp

6925

Are there other restrictions we need to care about, like what the type of the parameter/global is, etc?

[python3kgae]

Add type check for dispatch thread id.

[python3kgae]

In D133983#3805761, @aaron.ballman wrote:

There are no tests for applying this to a global variable, so those should be added.

The global variable in the Subjects is wrong.
It should be Field.
Support for semantic on field is a bigger change.
Created https://github.com/llvm/llvm-project/issues/57889 to track it.

clang/include/clang/Basic/AttrDocs.td

6596–6598

The global variable is wrong. Change it to Field and added document.
But adding support for field is a big change.
Created https://github.com/llvm/llvm-project/issues/57889 to track it.

clang/lib/Sema/SemaDeclAttr.cpp

6925

Yes. Added the check for type.

[aaron.ballman]

In D133983#3807284, @python3kgae wrote:

In D133983#3805761, @aaron.ballman wrote:

There are no tests for applying this to a global variable, so those should be added.

The global variable in the Subjects is wrong.
It should be Field.

I wondered if that was the case, thanks for the fix. :-)

Support for semantic on field is a bigger change.
Created https://github.com/llvm/llvm-project/issues/57889 to track it.

Thanks!

You should add additional test coverage: applying the attribute to a static data member, to a non-static data member, to a variable, to an unnamed parameter, and to a lambda parameter. The tests should demonstrate that we give good diagnostics where appropriate, instead of crashing, asserting, or silently accepting.

clang/include/clang/Basic/AttrDocs.td

6596–6598

You should also wrap the docs to the usual 80-col limit.

Can you apply the attribute to a static data member in a struct, or only fields?

clang/include/clang/Basic/DiagnosticSemaKinds.td

11656 ↗

(On Diff #462039)

clang/lib/Sema/SemaDeclAttr.cpp

6913

6937

This will cause an assert when the attribute appears on a field -- you need to check for a field decl above and give a diagnostic about it not being supported yet.

6940

clang/test/SemaHLSL/Semantics/entry_parameter.hlsl

4–5

clang/test/SemaHLSL/Semantics/invalid_entry_parameter.hlsl

3–13 ↗

(On Diff #462039)

clang/test/SemaHLSL/Semantics/valid_entry_parameter.hlsl

3–16 ↗

(On Diff #462039)

[python3kgae]

Add more tests.

[python3kgae]

Add newline at end of file

[python3kgae]

In D133983#3808302, @aaron.ballman wrote:

In D133983#3807284, @python3kgae wrote:

In D133983#3805761, @aaron.ballman wrote:

There are no tests for applying this to a global variable, so those should be added.

The global variable in the Subjects is wrong.
It should be Field.

I wondered if that was the case, thanks for the fix. :-)

Support for semantic on field is a bigger change.
Created https://github.com/llvm/llvm-project/issues/57889 to track it.

Thanks!

You should add additional test coverage: applying the attribute to a static data member, to a non-static data member, to a variable, to an unnamed parameter, and to a lambda parameter. The tests should demonstrate that we give good diagnostics where appropriate, instead of crashing, asserting, or silently accepting.

Added test which only gets error now.
Create issue https://github.com/llvm/llvm-project/issues/57916 to change the error to a warning.
lambda is not supported in HLSL, so not add test for the lambda parameter.

clang/include/clang/Basic/AttrDocs.td

6596–6598

Not sure if static data member in a struct should be allowed for HLSL or not yet.

If it is allowed, it should be ignored with a warning. The issue is tracked by issue https://github.com/llvm/llvm-project/issues/57916

[aaron.ballman]

LGTM aside from a think-o that I introduced with an earlier suggested change. I'm not super qualified for the CodeGen changes. They look correct to me, but @beanz would probably be able to better validate the correctness there.

clang/include/clang/Basic/AttrDocs.td

6597

Sorry, I had a typo in my earlier suggestion!

[python3kgae]

Fix typo.

[beanz]

clang/lib/CodeGen/CGHLSLRuntime.cpp

166

I might be missing something, but I'm not seeing a test that exercises the sret case.

[python3kgae]

clang/lib/CodeGen/CGHLSLRuntime.cpp

166

Issue https://github.com/llvm/llvm-project/issues/57874 is to track this.
Cannot create a legal test on compute shader which only has input.

[beanz]

clang/lib/CodeGen/CGHLSLRuntime.cpp

166

My point is that if this code can't be reached and tested, we shouldn't add it.

[python3kgae]

Insert position value for sret output to make a test.

[python3kgae]

clang/lib/CodeGen/CGHLSLRuntime.cpp

166

Added a test.
It is still illegal HLSL, marked FIXME in the test.

[beanz]

clang/lib/CodeGen/CGHLSLRuntime.cpp

135

nit: this comment doesn't add value

165

Was dropping the const here intentional? Getting the parameter's type shouldn't modify it.

clang/test/CodeGenHLSL/sret_output.hlsl

5 ↗

(On Diff #464392)

18 ↗

(On Diff #464392)

add newline to the end of the file

[python3kgae]

Fix test fail after rebase and code cleanup to match comments.

[python3kgae]

Check sret parameter exist.

[python3kgae]

More check on the ir.

[beanz]

LGTM

clang/include/clang/Basic/Attr.td

 
 def HLSLSV_GroupIndex: HLSLAnnotationAttr {
   let Spellings = [HLSLSemantic<"SV_GroupIndex">];
   let Subjects = SubjectList<[ParmVar, GlobalVar]>;
   let LangOpts = [HLSL];
   let Documentation = [HLSLSV_GroupIndexDocs];
 }
 
+def HLSLSV_DispatchThreadID: HLSLAnnotationAttr {
+  let Spellings = [HLSLSemantic<"SV_DispatchThreadID">];
+  let Subjects = SubjectList<[ParmVar, GlobalVar]>;
+  let LangOpts = [HLSL];
+  let Documentation = [HLSLSV_DispatchThreadIDDocs];
+}
+
 def HLSLShader : InheritableAttr {
   let Spellings = [Microsoft<"shader">];
   let Subjects = SubjectList<[HLSLEntry]>;
   let LangOpts = [HLSL];
   // NOTE:
   // order for the enum should match order in llvm::Triple::EnvironmentType.
   // ShaderType will be converted to llvm::Triple::EnvironmentType like
   //   (llvm::Triple::EnvironmentType)((uint32_t)ShaderType +

clang/include/clang/Basic/AttrDocs.td

 The ``SV_GroupIndex`` semantic, when applied to an input parameter, specifies a
 data binding to map the group index to the specified parameter. This attribute
 is only supported in compute shaders.
 
 The full documentation is available here: https://docs.microsoft.com/en-us/windows/win32/direct3dhlsl/sv-groupindex
   }];
 }
 
+def HLSLSV_DispatchThreadIDDocs : Documentation {
+  let Category = DocCatFunction;
+  let Content = [{
+The ``SV_DispatchThreadID`` semantic, when applied to an input parameter, specifies a
+data binding to map global thread offset within the Dispatch call(per dimension of the group) to the specified parameter.
+This attribute is only supported in compute shaders.
+
+The full documentation is available here: https://docs.microsoft.com/en-us/windows/win32/direct3dhlsl/sv-dispatchthreadid
+  }];
+}
+
 def AnnotateTypeDocs : Documentation {
   let Category = DocCatType;
   let Heading = "annotate_type";
   let Content = [{
 This attribute is used to add annotations to types, typically for use by static
 analysis tools that are not integrated into the core Clang compiler (e.g.,
 Clang-Tidy checks or out-of-tree Clang-based tools). It is a counterpart to the
 `annotate` attribute, which serves the same purpose, but for declarations.

clang/lib/CodeGen/CGHLSLRuntime.h

 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);
+  llvm::Value *emitInputSemantic(llvm::IRBuilder<> &B, const ParmVarDecl &D,
+                                 llvm::Type *Ty);
 
 public:
   CGHLSLRuntime(CodeGenModule &CGM) : CGM(CGM) {}
   virtual ~CGHLSLRuntime() {}
 
   void annotateHLSLResource(const VarDecl *D, llvm::GlobalVariable *GV);
   void generateGlobalCtorDtorCalls();
 

clang/lib/CodeGen/CGHLSLRuntime.cpp

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()));
 }
 
+static Value *buildVectorInput(IRBuilder<> &B, Function *F, llvm::Type *Ty) {
+  if (FixedVectorType *VT = dyn_cast<FixedVectorType>(Ty)) {
+    Value *Result = PoisonValue::get(Ty);
+    for (unsigned I = 0; I < VT->getNumElements(); ++I) {
+      Value *Elt = B.CreateCall(F, {B.getInt32(I)});
+      Result = B.CreateInsertElement(Result, Elt, I);
+    }
+    return Result;
+  } else
+    return B.CreateCall(F, {B.getInt32(0)});
+}
+
 llvm::Value *CGHLSLRuntime::emitInputSemantic(IRBuilder<> &B,
-                                              const ParmVarDecl &D) {
+                                              const ParmVarDecl &D,
+                                              llvm::Type *Ty) {
   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));
   }
+  if (D.hasAttr<HLSLSV_DispatchThreadIDAttr>()) {
+    llvm::Function *DxThreadID = CGM.getIntrinsic(Intrinsic::dx_thread_id);
+    // dx_thread_id
+    return buildVectorInput(B, DxThreadID, Ty);
+  }
   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();
void CGHLSLRuntime::emitEntryFunction(const FunctionDecl *FD,
 
   // 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));
+  unsigned SRetOffset = 0;
+  for (const auto &Param : Fn->args()) {
+    if (Param.hasStructRetAttr()) {
+      // FIXME: support output.
+      SRetOffset = 1;
+      continue;
+    }
+    const ParmVarDecl *PD = FD->getParamDecl(Param.getArgNo() - SRetOffset);
+    Args.push_back(emitInputSemantic(B, *PD, Param.getType()));
   }
 
   CallInst *CI = B.CreateCall(FunctionCallee(Fn), Args);
   (void)CI;
   // FIXME: Handle codegen for return type semantics.
   B.CreateRetVoid();
 }
 

clang/lib/Sema/SemaDeclAttr.cpp

if (Target.getEnvironment() != Triple::Compute) {
     S.Diag(AL.getLoc(), diag::err_hlsl_attr_unsupported_in_stage)
         << AL << Pipeline << "Compute";
     return;
   }
 
   D->addAttr(::new (S.Context) HLSLSV_GroupIndexAttr(S.Context, AL));
 }
 
+static void handleHLSLSV_DispatchThreadIDAttr(Sema &S, Decl *D,
+                                              const ParsedAttr &AL) {
+  using llvm::Triple;
+  Triple Target = S.Context.getTargetInfo().getTriple();
+  // FIXME: it is OK for a compute shader entry and pixel shader entry live in
+  // same HLSL file.
+  if (Target.getEnvironment() != Triple::Compute &&
+      Target.getEnvironment() != Triple::Library) {
+    uint32_t Pipeline =
+        (uint32_t)S.Context.getTargetInfo().getTriple().getEnvironment() -
+        (uint32_t)llvm::Triple::Pixel;
+    S.Diag(AL.getLoc(), diag::err_hlsl_attr_unsupported_in_stage)
+        << AL << Pipeline << "Compute";
+    return;
+  }
+
+  D->addAttr(::new (S.Context) HLSLSV_DispatchThreadIDAttr(S.Context, AL));
+}
+
 static void handleHLSLShaderAttr(Sema &S, Decl *D, const ParsedAttr &AL) {
   StringRef Str;
   SourceLocation ArgLoc;
   if (!S.checkStringLiteralArgumentAttr(AL, 0, Str, &ArgLoc))
     return;
 
   HLSLShaderAttr::ShaderType ShaderType;
   if (!HLSLShaderAttr::ConvertStrToShaderType(Str, ShaderType) ||
       // Library is added to help convert HLSLShaderAttr::ShaderType to
       // llvm::Triple::EnviromentType. It is not a legal
       // HLSLShaderAttr::ShaderType.
       ShaderType == HLSLShaderAttr::Library) {
     S.Diag(AL.getLoc(), diag::warn_attribute_type_not_supported)
         << AL << Str << ArgLoc;
     return;
   }
 
   // FIXME: check function match the shader stage.
 
   HLSLShaderAttr *NewAttr = S.mergeHLSLShaderAttr(D, AL, ShaderType);
ProcessDeclAttribute(Sema &S, Scope *scope, Decl *D, const ParsedAttr &AL,
 
   // HLSL attributes:
   case ParsedAttr::AT_HLSLNumThreads:
     handleHLSLNumThreadsAttr(S, D, AL);
     break;
   case ParsedAttr::AT_HLSLSV_GroupIndex:
     handleHLSLSVGroupIndexAttr(S, D, AL);
     break;
+  case ParsedAttr::AT_HLSLSV_DispatchThreadID:
+    handleHLSLSV_DispatchThreadIDAttr(S, D, AL);
+    break;
   case ParsedAttr::AT_HLSLShader:
     handleHLSLShaderAttr(S, D, AL);
     break;
 
   case ParsedAttr::AT_AbiTag:
     handleAbiTagAttr(S, D, AL);
     break;
   case ParsedAttr::AT_CFGuard:

clang/test/CodeGenHLSL/semantics/DispatchThreadID.hlsl

+// RUN: %clang_cc1 -triple dxil-pc-shadermodel6.3-library -x hlsl -emit-llvm -finclude-default-header -disable-llvm-passes -o - %s
+
+// Make sure SV_DispatchThreadID translated into dx.thread.id.
+
+const RWBuffer<float> In;
+RWBuffer<float> Out;
+
+// CHECK: define void @foo()
+// CHECK: call i32 @llvm.dx.thread.id(i32 0)
+// CHECK: call void @"?foo@@YAXH@Z"(i32 %{{.*}})
+[shader("compute")]
+[numthreads(8,8,1)]
+void foo(int Idx : SV_DispatchThreadID) {
+  Out[Idx] = In[Idx];
+}
+
+// CHECK: define void @bar()
+// CHECK: call i32 @llvm.dx.thread.id(i32 0)
+// CHECK: call i32 @llvm.dx.thread.id(i32 1)
+// CHECK: call void @"?bar@@YAXT?$__vector@H$01@__clang@@@Z"(<2 x i32> %{{.*}})
+[shader("compute")]
+[numthreads(8,8,1)]
+void bar(int2 Idx : SV_DispatchThreadID) {
+  Out[Idx.y] = In[Idx.x];
+}
+

clang/test/SemaHLSL/Semantics/entry_parameter.hlsl

 // RUN: %clang_cc1 -triple dxil-pc-shadermodel6.0-compute -x hlsl -ast-dump -o - %s | FileCheck %s 
 // RUN: %clang_cc1 -triple dxil-pc-shadermodel6.0-mesh -x hlsl -ast-dump -verify -o - %s 
 
 [numthreads(8,8, 1)]
-// expected-error@+1 {{attribute 'SV_GroupIndex' is unsupported in Mesh shaders, requires Compute}}
-void CSMain(int GI : SV_GroupIndex) {
-// CHECK: FunctionDecl 0x{{[0-9a-fA-F]+}} <{{.*}}> line:[[@LINE-1]]:6 CSMain 'void (int)'
+// expected-error@+2 {{attribute 'SV_GroupIndex' is unsupported in Mesh shaders, requires Compute}}
+// expected-error@+1 {{attribute 'SV_DispatchThreadID' is unsupported in Mesh shaders, requires Compute}}
+void CSMain(int GI : SV_GroupIndex, int ID : SV_DispatchThreadID) {
+// CHECK: FunctionDecl 0x{{[0-9a-fA-F]+}} <{{.*}}> line:[[@LINE-1]]:6 CSMain 'void (int, int)'
 // CHECK-NEXT: ParmVarDecl 0x{{[0-9a-fA-F]+}} <{{.*}}> col:17 GI 'int'
 // CHECK-NEXT: HLSLSV_GroupIndexAttr
+// CHECK-NEXT: ParmVarDecl 0x{{[0-9a-fA-F]+}} <{{.*}}> col:41 ID 'int'
+// CHECK-NEXT: HLSLSV_DispatchThreadIDAttr
 }