[clang] 0f77bdd - [HLSL] generate hlsl.wavesize attribute (#107176)

[via cfe-commits]

Author: Xiang Li
Date: 2024-09-17T23:00:13-04:00
New Revision: 0f77bdd2c9bed6db6c2b9e10e05b5e99fde13d66

URL: https://github.com/llvm/llvm-project/commit/0f77bdd2c9bed6db6c2b9e10e05b5e99fde13d66
DIFF: https://github.com/llvm/llvm-project/commit/0f77bdd2c9bed6db6c2b9e10e05b5e99fde13d66.diff

LOG: [HLSL] generate hlsl.wavesize attribute (#107176)

Generate function attribute hlsl.wavesize from [WaveSize].

For #70118

Added: 
    clang/test/CodeGenHLSL/wavesize.hlsl

Modified: 
    clang/lib/CodeGen/CGHLSLRuntime.cpp

Removed: 
    


################################################################################
diff  --git a/clang/lib/CodeGen/CGHLSLRuntime.cpp b/clang/lib/CodeGen/CGHLSLRuntime.cpp
index bec0a29e34fcb5..59d8fc830dcc8f 100644
--- a/clang/lib/CodeGen/CGHLSLRuntime.cpp
+++ b/clang/lib/CodeGen/CGHLSLRuntime.cpp
@@ -338,6 +338,13 @@ void clang::CodeGen::CGHLSLRuntime::setHLSLEntryAttributes(
                 NumThreadsAttr->getZ());
     Fn->addFnAttr(NumThreadsKindStr, NumThreadsStr);
   }
+  if (HLSLWaveSizeAttr *WaveSizeAttr = FD->getAttr<HLSLWaveSizeAttr>()) {
+    const StringRef WaveSizeKindStr = "hlsl.wavesize";
+    std::string WaveSizeStr =
+        formatv("{0},{1},{2}", WaveSizeAttr->getMin(), WaveSizeAttr->getMax(),
+                WaveSizeAttr->getPreferred());
+    Fn->addFnAttr(WaveSizeKindStr, WaveSizeStr);
+  }
   Fn->addFnAttr(llvm::Attribute::NoInline);
 }
 

diff  --git a/clang/test/CodeGenHLSL/wavesize.hlsl b/clang/test/CodeGenHLSL/wavesize.hlsl
new file mode 100644
index 00000000000000..fcb817ce0d0629
--- /dev/null
+++ b/clang/test/CodeGenHLSL/wavesize.hlsl
@@ -0,0 +1,34 @@
+// RUN: %clang_cc1 -std=hlsl2021 -finclude-default-header -x hlsl -triple \
+// RUN:   dxil-pc-shadermodel6.6-compute %s -DSM66 -hlsl-entry foo \
+// RUN:   -emit-llvm -disable-llvm-passes -o - | FileCheck %s
+
+// RUN: %clang_cc1 -std=hlsl2021 -finclude-default-header -x hlsl -triple \
+// RUN:   dxil-pc-shadermodel6.8-compute %s -DNO_PREFERR -hlsl-entry foo \
+// RUN:   -emit-llvm -disable-llvm-passes -o - | FileCheck %s --check-prefix=NO_PREFERR
+
+// RUN: %clang_cc1 -std=hlsl2021 -finclude-default-header -x hlsl -triple \
+// RUN:   dxil-pc-shadermodel6.8-compute %s -hlsl-entry foo \
+// RUN:   -emit-llvm -disable-llvm-passes -o - | FileCheck %s --check-prefix=CHECK-SM68
+
+
+// Make sure wavesize attribute get correct value for sm66 and sm68.
+// CHECK:define void @foo()
+// CHECK:"hlsl.wavesize"="8,0,0"
+
+// NO_PREFERR:define void @foo()
+// NO_PREFERR:"hlsl.wavesize"="8,128,0"
+
+// CHECK-SM68:define void @foo()
+// CHECK-SM68:"hlsl.wavesize"="8,128,64"
+
+[numthreads(16,8,1)]
+#ifdef SM66
+[WaveSize(8)]
+#elif NO_PREFERR
+[WaveSize(8, 128)]
+#else
+[WaveSize(8, 128, 64)]
+#endif
+void foo() {
+
+}