[202x][unroll] Only run the default LLVM loop unroll for 202x onward

docs/ReleaseNotes.md

@@ -33,6 +33,9 @@ line upon naming the release. Refer to previous for appropriate section names.
 
 #### HLSL Language
 
+- Starting with HLSL 202x, the count in `[unroll(N)]` is a partial-unroll hint
+  and no longer limits the number of loop iterations
+  [#8789](https://github.com/microsoft/DirectXShaderCompiler/issues/8789).
 - Casting a scalar to a struct or array containing a resource is now an error
   instead of crashing
   [#6661](https://github.com/microsoft/DirectXShaderCompiler/issues/6661).

include/llvm/Transforms/IPO/PassManagerBuilder.h

@@ -129,6 +129,7 @@ class PassManagerBuilder {
   bool HLSLHighLevel = false; // HLSL Change
   bool HLSLAllowPreserveValues = false; // HLSL Change
   bool HLSLOnlyWarnOnUnrollFail = false; // HLSL Change
+  bool HLSLUnrollCountIsHint = false;    // HLSL Change
   hlsl::HLSLExtensionsCodegenHelper *HLSLExtensionsCodeGen = nullptr; // HLSL Change
   bool HLSLResMayAlias = false; // HLSL Change
   unsigned ScanLimit = 0; // HLSL Change

include/llvm/Transforms/Scalar.h

@@ -129,7 +129,9 @@ void initializeDxilFixConstArrayInitializerPass(PassRegistry&);
 Pass *createDxilConditionalMem2RegPass(bool NoOpt);
 void initializeDxilConditionalMem2RegPass(PassRegistry&);
 
-Pass *createDxilLoopUnrollPass(unsigned MaxIterationAttempt, bool OnlyWarnOnFail, bool StructurizeLoopExits);
+Pass *createDxilLoopUnrollPass(unsigned MaxIterationAttempt,
+                               bool OnlyWarnOnFail, bool StructurizeLoopExits,
+                               bool UnrollCountIsHint);
 void initializeDxilLoopUnrollPass(PassRegistry&);
 
 Pass *createDxilEraseDeadRegionPass();

lib/Transforms/IPO/PassManagerBuilder.cpp

@@ -308,10 +308,12 @@ void PassManagerBuilder::addHLSLPasses(legacy::PassManagerBase &MPM) {
   // struct members.
   // Needs to happen before resources are lowered and before HL
   // module is gone.
-  MPM.add(createDxilLoopUnrollPass(1024, HLSLOnlyWarnOnUnrollFail, StructurizeLoopExitsForUnroll));
+  MPM.add(createDxilLoopUnrollPass(1024, HLSLOnlyWarnOnUnrollFail,
+                                   StructurizeLoopExitsForUnroll,
+                                   HLSLUnrollCountIsHint));
 
-  // Default unroll pass. This is purely for optimizing loops without
-  // attributes.
+  // Default unroll pass. In HLSL 202x, this consumes [unroll(N)] as a
+  // partial-unroll hint; it otherwise optimizes loops without attributes.
   if (OptLevel > 2) {
     MPM.add(createLoopUnrollPass(-1, -1, -1, -1, StructurizeLoopExitsForUnroll));
   }

lib/Transforms/Scalar/DxilLoopUnroll.cpp

@@ -109,12 +109,15 @@ class DxilLoopUnroll : public LoopPass {
   unsigned MaxIterationAttempt = 0;
   bool OnlyWarnOnFail = false;
   bool StructurizeLoopExits = false;
+  bool UnrollCountIsHint = false;
 
   DxilLoopUnroll(unsigned MaxIterationAttempt = 1024,
-                 bool OnlyWarnOnFail = false, bool StructurizeLoopExits = false)
+                 bool OnlyWarnOnFail = false, bool StructurizeLoopExits = false,
+                 bool UnrollCountIsHint = false)
       : LoopPass(ID), MaxIterationAttempt(MaxIterationAttempt),
         OnlyWarnOnFail(OnlyWarnOnFail),
-        StructurizeLoopExits(StructurizeLoopExits) {
+        StructurizeLoopExits(StructurizeLoopExits),
+        UnrollCountIsHint(UnrollCountIsHint) {
     initializeDxilLoopUnrollPass(*PassRegistry::getPassRegistry());
   }
   StringRef getPassName() const override { return "Dxil Loop Unroll"; }
@@ -138,12 +141,14 @@ class DxilLoopUnroll : public LoopPass {
                           false);
     GetPassOptionBool(O, "OnlyWarnOnFail", &OnlyWarnOnFail, false);
     GetPassOptionBool(O, "StructurizeLoopExits", &StructurizeLoopExits, false);
+    GetPassOptionBool(O, "UnrollCountIsHint", &UnrollCountIsHint, false);
   }
   void dumpConfig(raw_ostream &OS) override {
     LoopPass::dumpConfig(OS);
     OS << ",MaxIterationAttempt=" << MaxIterationAttempt;
     OS << ",OnlyWarnOnFail=" << OnlyWarnOnFail;
     OS << ",StructurizeLoopExits=" << StructurizeLoopExits;
+    OS << ",UnrollCountIsHint=" << UnrollCountIsHint;
   }
   void RecursivelyRemoveLoopOnSuccess(LPPassManager &LPM, Loop *L);
   void RecursivelyRecreateSubLoopForIteration(LPPassManager &LPM, LoopInfo *LI,
@@ -790,6 +795,8 @@ bool DxilLoopUnroll::runOnLoop(Loop *L, LPPassManager &LPM) {
       return false;
     }
     ExplicitUnrollCount = (unsigned)ExplicitUnrollCountSigned;
+    if (UnrollCountIsHint)
+      return false;
   }
 
   if (!IsLoopSafeToClone(L))
@@ -1256,9 +1263,10 @@ bool DxilLoopUnroll::doFinalization() {
 
 Pass *llvm::createDxilLoopUnrollPass(unsigned MaxIterationAttempt,
                                      bool OnlyWarnOnFail,
-                                     bool StructurizeLoopExits) {
+                                     bool StructurizeLoopExits,
+                                     bool UnrollCountIsHint) {
   return new DxilLoopUnroll(MaxIterationAttempt, OnlyWarnOnFail,
-                            StructurizeLoopExits);
+                            StructurizeLoopExits, UnrollCountIsHint);
 }
 
 INITIALIZE_PASS_BEGIN(DxilLoopUnroll, "dxil-loop-unroll", "Dxil Unroll loops",

tools/clang/lib/CodeGen/BackendUtil.cpp

@@ -342,6 +342,8 @@ void EmitAssemblyHelper::CreatePasses() {
   PMBuilder.HLSLHighLevel = CodeGenOpts.HLSLHighLevel;
   PMBuilder.HLSLAllowPreserveValues = CodeGenOpts.HLSLAllowPreserveValues;
   PMBuilder.HLSLOnlyWarnOnUnrollFail = CodeGenOpts.HLSLOnlyWarnOnUnrollFail;
+  PMBuilder.HLSLUnrollCountIsHint =
+      LangOpts.HLSLVersion >= hlsl::LangStd::v202x;
   PMBuilder.HLSLExtensionsCodeGen = CodeGenOpts.HLSLExtensionsCodegen.get();
   PMBuilder.HLSLResMayAlias = CodeGenOpts.HLSLResMayAlias;
   PMBuilder.ScanLimit = CodeGenOpts.ScanLimit;

tools/clang/test/HLSLFileCheck/hlsl/control_flow/attributes/unroll/count_cbuff.hlsl

@@ -1,4 +1,4 @@
-// RUN: %dxc -E main -T ps_6_0 %s | FileCheck %s
+// RUN: %dxc -E main -T ps_6_0 -HV 2021 %s | FileCheck %s
 
 // CHECK: call float @dx.op.dot3
 // CHECK: call float @dx.op.dot3

tools/clang/test/HLSLFileCheck/hlsl/control_flow/attributes/unroll/count_greater_than_i.hlsl

@@ -1,4 +1,4 @@
-// RUN: %dxc -E main -T ps_6_0 %s | FileCheck %s
+// RUN: %dxc -E main -T ps_6_0 -HV 2021 %s | FileCheck %s
 
 // CHECK: call float @dx.op.dot3
 // CHECK: call float @dx.op.dot3

tools/clang/test/HLSLFileCheck/hlsl/control_flow/attributes/unroll/count_hint_202x.hlsl

@@ -0,0 +1,24 @@
+// RUN: %dxc -E main -T ps_6_0 -HV 202x -O2 %s | FileCheck %s -check-prefix=O2
+// RUN: %dxc -E main -T ps_6_0 -HV 202x -O3 %s | FileCheck %s -check-prefix=O3
+
+// O2: call float @dx.op.dot3
+// O2-NOT: call float @dx.op.dot3
+// O2: br i1 {{.*}}, label {{.*}}, label {{.*}}, !llvm.loop
+// O2: !{!"llvm.loop.unroll.count", i32 3}
+
+// O3: call float @dx.op.dot3
+// O3: call float @dx.op.dot3
+// O3: call float @dx.op.dot3
+// O3: call float @dx.op.dot3
+// O3-NOT: call float @dx.op.dot3
+// O3: br i1 {{.*}}, label {{.*}}, label {{.*}}, !llvm.loop
+// O3: !{!"llvm.loop.unroll.disable"}
+
+float main(float3 a : A, float3 b : B) : SV_Target {
+  float result = 0;
+  [unroll(3)]
+  for (int i = 0; i < 10; i++) {
+    result += dot(a * i, b);
+  }
+  return result;
+}

tools/clang/test/HLSLFileCheck/hlsl/control_flow/attributes/unroll/count_less_than_i.hlsl

@@ -1,4 +1,4 @@
-// RUN: %dxc -E main -T ps_6_0 %s | FileCheck %s
+// RUN: %dxc -E main -T ps_6_0 -HV 2021 %s | FileCheck %s
 
 // CHECK: call float @dx.op.dot3
 // CHECK: call float @dx.op.dot3

tools/clang/test/HLSLFileCheck/hlsl/control_flow/attributes/unroll/count_runtime_hint_202x.hlsl

@@ -0,0 +1,46 @@
+// RUN: %dxc -E main -T ps_6_0 -HV 202x -O2 %s | FileCheck %s -check-prefixes=COMMON,O2
+// RUN: %dxc -E main -T ps_6_0 -HV 202x -O3 %s | FileCheck %s -check-prefixes=COMMON,O3
+// RUN: %dxc -E mainPowerOfTwo -T ps_6_0 -HV 202x -O2 %s | FileCheck %s -check-prefixes=COMMON,POWER-O2
+// RUN: %dxc -E mainPowerOfTwo -T ps_6_0 -HV 202x -O3 %s | FileCheck %s -check-prefix=POWER-O3
+
+// O2 Doesn't unroll and O3 Runtime unrolling only supports 
+// power-of-two factors, so the count hint is
+// consumed without cloning the loop body.
+// COMMON: call float @dx.op.dot3
+// COMMON-NOT: call float @dx.op.dot3
+// COMMON: br i1 {{.*}}, label {{.*}}, label {{.*}}, !llvm.loop
+// O2: !{!"llvm.loop.unroll.count", i32 3}
+// O3: !{!"llvm.loop.unroll.disable"}
+// POWER-O2: !{!"llvm.loop.unroll.count", i32 4}
+
+// The O3 output has one remainder-loop body and four unrolled main-loop bodies.
+// POWER-O3: and i32 {{.*}}, 3
+// POWER-O3: call float @dx.op.dot3
+// POWER-O3: call float @dx.op.dot3
+// POWER-O3: call float @dx.op.dot3
+// POWER-O3: call float @dx.op.dot3
+// POWER-O3: call float @dx.op.dot3
+// POWER-O3-NOT: call float @dx.op.dot3
+// POWER-O3: add i32 {{.*}}, 4
+// POWER-O3: br i1 {{.*}}, label {{.*}}, label {{.*}}, !llvm.loop
+// POWER-O3: !{!"llvm.loop.unroll.disable"}
+
+float main(float3 a : A, float3 b : B,
+           uint iterationCount : COUNT) : SV_Target {
+  float result = 0;
+  [unroll(3)]
+  for (uint i = 0; i < iterationCount; i++) {
+    result += dot(a * i, b);
+  }
+  return result;
+}
+
+float mainPowerOfTwo(float3 a : A, float3 b : B,
+                     uint iterationCount : COUNT) : SV_Target {
+  float result = 0;
+  [unroll(4)]
+  for (uint i = 0; i < iterationCount; i++) {
+    result += dot(a * i, b);
+  }
+  return result;
+}

tools/clang/test/HLSLFileCheck/hlsl/control_flow/attributes/unroll/full_hint_202x.hlsl

@@ -0,0 +1,25 @@
+// RUN: %dxc -E main -T ps_6_0 -HV 202x -O2 %s | FileCheck %s -check-prefix=CHECK
+// RUN: %dxc -E main -T ps_6_0 -HV 202x -O3 %s | FileCheck %s -check-prefix=CHECK
+
+// CHECK: call float @dx.op.dot3
+// CHECK: call float @dx.op.dot3
+// CHECK: call float @dx.op.dot3
+// CHECK: call float @dx.op.dot3
+// CHECK: call float @dx.op.dot3
+// CHECK: call float @dx.op.dot3
+// CHECK: call float @dx.op.dot3
+// CHECK: call float @dx.op.dot3
+// CHECK: call float @dx.op.dot3
+// CHECK: call float @dx.op.dot3
+// CHECK-NOT: call float @dx.op.dot3
+// CHECK-NOT: br i1
+// CHECK-NOT: !llvm.loop
+
+float main(float3 a : A, float3 b : B) : SV_Target {
+  float result = 0;
+  [unroll]
+  for (int i = 0; i < 10; i++) {
+    result += dot(a * i, b);
+  }
+  return result;
+}

utils/hct/hctdb.py

@@ -7580,6 +7580,12 @@ def add_pass(name, type_name, doc, opts):
                     "c": 1,
                     "d": "Whether the unroller should try to structurize loop exits first.",
                 },
+                {
+                    "n": "UnrollCountIsHint",
+                    "t": "bool",
+                    "c": 1,
+                    "d": "Whether an explicit unroll count should be treated as a hint.",
+                },
             ],
         )
         add_pass(