[HLSL] add loop unroll

[farzonl]

spec: microsoft/hlsl-specs#263

• Attr.td - Define the HLSL loop attribute hints (unroll and loop)
• AttrDocs.td - Add documentation for unroll and loop
• CGLoopInfo.cpp - Add codegen for HLSL unroll that maps to clang unroll expectations
• ParseStmt.cpp - For statements if HLSL define DeclSpecAttrs via MaybeParseMicrosoftAttributes
• SemaStmtAttr.cpp - Add the HLSL loop unroll handeling

resolves #70114

dxc examples:

• for loop: https://hlsl.godbolt.org/z/8EK6Pa139
• while loop: https://hlsl.godbolt.org/z/ebr5MvEcK
• do while: https://hlsl.godbolt.org/z/be8cedoTs

Documentation:

[llvmbot]

@llvm/pr-subscribers-hlsl
@llvm/pr-subscribers-clang

@llvm/pr-subscribers-clang-codegen

Author: Farzon Lotfi (farzonl)

Changes

• Attr.td - Define the HLSL loop attribute hints (unroll and loop)
• AttrDocs.td - Add documentation for unroll and loop
• CGLoopInfo.cpp - Add codegen for HLSL unroll that maps to clang unroll expectations
• ParseStmt.cpp - For statements if HLSL define DeclSpecAttrs via MaybeParseMicrosoftAttributes
• SemaStmtAttr.cpp - Add the HLSL loop unroll handeling

resolves ##70114

---

Full diff: https://github.com/llvm/llvm-project/pull/93879.diff

6 Files Affected:

• (modified) clang/include/clang/Basic/Attr.td (+13)
• (modified) clang/include/clang/Basic/AttrDocs.td (+47)
• (modified) clang/lib/CodeGen/CGLoopInfo.cpp (+13-2)
• (modified) clang/lib/Parse/ParseStmt.cpp (+7-4)
• (modified) clang/lib/Sema/SemaStmtAttr.cpp (+29)
• (added) clang/test/CodeGenHLSL/loops/unroll.hlsl (+99)

diff --git a/clang/include/clang/Basic/Attr.td b/clang/include/clang/Basic/Attr.td
index 2665b7353ca4a..e319db08ba168 100644
--- a/clang/include/clang/Basic/Attr.td
+++ b/clang/include/clang/Basic/Attr.td
@@ -4114,6 +4114,19 @@ def LoopHint : Attr {
   let HasCustomParsing = 1;
 }
 
+/// The HLSL loop attributes
+def HLSLLoopHint: StmtAttr {
+  /// [unroll(directive)]
+  /// [loop]
+  let Spellings = [Microsoft<"unroll">, Microsoft<"loop">];
+  let Args = [UnsignedArgument<"directive">];
+  let Subjects = SubjectList<[ForStmt, WhileStmt, DoStmt],
+                              ErrorDiag, "'for', 'while', and 'do' statements">;
+  let LangOpts = [HLSL];
+  let Documentation = [HLSLLoopHintDocs, HLSLUnrollHintDocs];
+  let HasCustomParsing = 1;
+}
+
 def CapturedRecord : InheritableAttr {
   // This attribute has no spellings as it is only ever created implicitly.
   let Spellings = [];
diff --git a/clang/include/clang/Basic/AttrDocs.td b/clang/include/clang/Basic/AttrDocs.td
index a313e811c9d21..d84af9402d6db 100644
--- a/clang/include/clang/Basic/AttrDocs.td
+++ b/clang/include/clang/Basic/AttrDocs.td
@@ -7342,6 +7342,53 @@ where shaders must be compiled into a library and linked at runtime.
   }];
 }
 
+def HLSLLoopHintDocs : Documentation {
+  let Category = DocCatStmt;
+  let Heading = "#[loop]";
+  let Content = [{
+The ``[loop]`` directive allows loop optimization hints to be
+specified for the subsequent loop. The directive allows unrolling to
+be disabled and is not compatible with [unroll(x)]. 
+
+Specifying the parameter, ``[loop]``, directs the
+unroller to not unroll the loop. 
+
+.. code-block:: hlsl
+
+  [loop]
+  for (...) {
+    ...
+  }
+
+See `hlsl loop extensions
+<https://learn.microsoft.com/en-us/windows/win32/direct3dhlsl/dx-graphics-hlsl-for>`_
+for details.
+  }];
+}
+
+def HLSLUnrollHintDocs : Documentation {
+  let Category = DocCatStmt;
+  let Heading = "[unroll(x)]";
+  let Content = [{
+Loop unrolling optimization hints can be specified with ``[unroll(x)]``
+. The attribute is placed immediately before a for, while,
+or do-while.
+Specifying the parameter, ``[unroll(_value_)]``, directs the
+unroller to unroll the loop ``_value_`` times. Note: [unroll(x)] is not compatible with [loop].
+
+.. code-block:: hlsl
+
+  [unroll(4)]
+  for (...) {
+    ...
+  }
+See
+`hlsl loop extensions
+<https://learn.microsoft.com/en-us/windows/win32/direct3dhlsl/dx-graphics-hlsl-for>`_
+for details.
+  }];
+}
+
 def ClangRandomizeLayoutDocs : Documentation {
   let Category = DocCatDecl;
   let Heading = "randomize_layout, no_randomize_layout";
diff --git a/clang/lib/CodeGen/CGLoopInfo.cpp b/clang/lib/CodeGen/CGLoopInfo.cpp
index 0d4800b90a2f2..6b886bd6b6d2c 100644
--- a/clang/lib/CodeGen/CGLoopInfo.cpp
+++ b/clang/lib/CodeGen/CGLoopInfo.cpp
@@ -612,9 +612,9 @@ void LoopInfoStack::push(BasicBlock *Header, clang::ASTContext &Ctx,
     const LoopHintAttr *LH = dyn_cast<LoopHintAttr>(Attr);
     const OpenCLUnrollHintAttr *OpenCLHint =
         dyn_cast<OpenCLUnrollHintAttr>(Attr);
-
+    const HLSLLoopHintAttr *HLSLLoopHint = dyn_cast<HLSLLoopHintAttr>(Attr);
     // Skip non loop hint attributes
-    if (!LH && !OpenCLHint) {
+    if (!LH && !OpenCLHint && !HLSLLoopHint) {
       continue;
     }
 
@@ -635,6 +635,17 @@ void LoopInfoStack::push(BasicBlock *Header, clang::ASTContext &Ctx,
         Option = LoopHintAttr::UnrollCount;
         State = LoopHintAttr::Numeric;
       }
+    } else if (HLSLLoopHint) {
+      ValueInt = HLSLLoopHint->getDirective();
+      if (HLSLLoopHint->getSemanticSpelling() ==
+          HLSLLoopHintAttr::Spelling::Microsoft_unroll) {
+        if (ValueInt == 0)
+          State = LoopHintAttr::Enable;
+        if (ValueInt > 0) {
+          Option = LoopHintAttr::UnrollCount;
+          State = LoopHintAttr::Numeric;
+        }
+      }
     } else if (LH) {
       auto *ValueExpr = LH->getValue();
       if (ValueExpr) {
diff --git a/clang/lib/Parse/ParseStmt.cpp b/clang/lib/Parse/ParseStmt.cpp
index c25203243ee49..b8a717820c418 100644
--- a/clang/lib/Parse/ParseStmt.cpp
+++ b/clang/lib/Parse/ParseStmt.cpp
@@ -114,18 +114,21 @@ Parser::ParseStatementOrDeclaration(StmtVector &Stmts,
   // here because we don't want to allow arbitrary orderings.
   ParsedAttributes CXX11Attrs(AttrFactory);
   MaybeParseCXX11Attributes(CXX11Attrs, /*MightBeObjCMessageSend*/ true);
-  ParsedAttributes GNUAttrs(AttrFactory);
+  ParsedAttributes DeclSpecAttrs(AttrFactory);
   if (getLangOpts().OpenCL)
-    MaybeParseGNUAttributes(GNUAttrs);
+    MaybeParseGNUAttributes(DeclSpecAttrs);
+
+  if (getLangOpts().HLSL)
+    MaybeParseMicrosoftAttributes(DeclSpecAttrs);
 
   StmtResult Res = ParseStatementOrDeclarationAfterAttributes(
-      Stmts, StmtCtx, TrailingElseLoc, CXX11Attrs, GNUAttrs);
+      Stmts, StmtCtx, TrailingElseLoc, CXX11Attrs, DeclSpecAttrs);
   MaybeDestroyTemplateIds();
 
   // Attributes that are left should all go on the statement, so concatenate the
   // two lists.
   ParsedAttributes Attrs(AttrFactory);
-  takeAndConcatenateAttrs(CXX11Attrs, GNUAttrs, Attrs);
+  takeAndConcatenateAttrs(CXX11Attrs, DeclSpecAttrs, Attrs);
 
   assert((Attrs.empty() || Res.isInvalid() || Res.isUsable()) &&
          "attributes on empty statement");
diff --git a/clang/lib/Sema/SemaStmtAttr.cpp b/clang/lib/Sema/SemaStmtAttr.cpp
index 6f538ed55cb72..35a20c113f628 100644
--- a/clang/lib/Sema/SemaStmtAttr.cpp
+++ b/clang/lib/Sema/SemaStmtAttr.cpp
@@ -16,6 +16,7 @@
 #include "clang/Basic/TargetInfo.h"
 #include "clang/Sema/DelayedDiagnostic.h"
 #include "clang/Sema/Lookup.h"
+#include "clang/Sema/ParsedAttr.h"
 #include "clang/Sema/ScopeInfo.h"
 #include "clang/Sema/SemaInternal.h"
 #include "llvm/ADT/StringExtras.h"
@@ -584,6 +585,32 @@ static Attr *handleOpenCLUnrollHint(Sema &S, Stmt *St, const ParsedAttr &A,
   return ::new (S.Context) OpenCLUnrollHintAttr(S.Context, A, UnrollFactor);
 }
 
+static Attr *handleHLSLLoopHintAttr(Sema &S, Stmt *St, const ParsedAttr &A,
+                                    SourceRange Range) {
+  unsigned UnrollFactor = 0;
+  if (A.getNumArgs() == 1) {
+    Expr *E = A.getArgAsExpr(0);
+    std::optional<llvm::APSInt> ArgVal;
+
+    if (!(ArgVal = E->getIntegerConstantExpr(S.Context))) {
+      S.Diag(A.getLoc(), diag::err_attribute_argument_type)
+          << A << AANT_ArgumentIntegerConstant << E->getSourceRange();
+      return nullptr;
+    }
+
+    int Val = ArgVal->getSExtValue();
+    if (Val <= 0) {
+      S.Diag(A.getRange().getBegin(),
+             diag::err_attribute_requires_positive_integer)
+          << A << /* positive */ 0;
+      return nullptr;
+    }
+    UnrollFactor = static_cast<unsigned>(Val);
+  }
+
+  return ::new (S.Context) HLSLLoopHintAttr(S.Context, A, UnrollFactor);
+}
+
 static Attr *ProcessStmtAttribute(Sema &S, Stmt *St, const ParsedAttr &A,
                                   SourceRange Range) {
   if (A.isInvalid() || A.getKind() == ParsedAttr::IgnoredAttribute)
@@ -618,6 +645,8 @@ static Attr *ProcessStmtAttribute(Sema &S, Stmt *St, const ParsedAttr &A,
     return handleFallThroughAttr(S, St, A, Range);
   case ParsedAttr::AT_LoopHint:
     return handleLoopHintAttr(S, St, A, Range);
+  case ParsedAttr::AT_HLSLLoopHint:
+    return handleHLSLLoopHintAttr(S, St, A, Range);
   case ParsedAttr::AT_OpenCLUnrollHint:
     return handleOpenCLUnrollHint(S, St, A, Range);
   case ParsedAttr::AT_Suppress:
diff --git a/clang/test/CodeGenHLSL/loops/unroll.hlsl b/clang/test/CodeGenHLSL/loops/unroll.hlsl
new file mode 100644
index 0000000000000..0ebe5b0e847e4
--- /dev/null
+++ b/clang/test/CodeGenHLSL/loops/unroll.hlsl
@@ -0,0 +1,99 @@
+// RUN: %clang_cc1 -std=hlsl2021 -finclude-default-header -x hlsl -triple \
+// RUN: dxil-pc-shadermodel6.3-library %s -emit-llvm -o - | FileCheck %s
+
+/*** for ***/
+void for_count()
+{
+// CHECK-LABEL: for_count
+    [unroll(8)]
+    for( int i = 0; i < 1000; ++i);
+// CHECK: br label %{{.*}}, !llvm.loop ![[FOR_DISTINCT:.*]]
+}
+
+void for_disable()
+{
+// CHECK-LABEL: for_disable
+    [loop]
+    for( int i = 0; i < 1000; ++i);
+// CHECK: br label %{{.*}}, !llvm.loop ![[FOR_DISABLE:.*]]
+}
+
+void for_enable()
+{
+// CHECK-LABEL: for_enable
+    [unroll]
+    for( int i = 0; i < 1000; ++i);
+// CHECK: br label %{{.*}}, !llvm.loop ![[FOR_ENABLE:.*]]
+}
+
+/*** while ***/
+void while_count()
+{
+// CHECK-LABEL: while_count
+    int i = 1000;
+    [unroll(4)]
+    while(i-->0);
+// CHECK: br label %{{.*}}, !llvm.loop ![[WHILE_DISTINCT:.*]]
+}
+
+void while_disable()
+{
+// CHECK-LABEL: while_disable
+    int i = 1000;
+    [loop]
+    while(i-->0);
+// CHECK: br label %{{.*}}, !llvm.loop ![[WHILE_DISABLE:.*]]
+}
+
+void while_enable()
+{
+// CHECK-LABEL: while_enable
+    int i = 1000;
+    [unroll]
+    while(i-->0);
+// CHECK: br label %{{.*}}, !llvm.loop ![[WHILE_ENABLE:.*]]
+}
+
+/*** do ***/
+void do_count()
+{
+// CHECK-LABEL: do_count
+    int i = 1000;
+    [unroll(16)]
+    do {} while(i--> 0);
+// CHECK: br i1 %{{.*}}, label %{{.*}}, label %{{.*}}, !llvm.loop ![[DO_DISTINCT:.*]]
+}
+
+void do_disable()
+{
+// CHECK-LABEL: do_disable
+    int i = 1000;
+    [loop]
+    do {} while(i--> 0);
+// CHECK: br i1 %{{.*}}, label %{{.*}}, label %{{.*}}, !llvm.loop ![[DO_DISABLE:.*]]
+}
+
+void do_enable()
+{
+// CHECK-LABEL: do_enable
+    int i = 1000;
+    [unroll]
+    do {} while(i--> 0);
+// CHECK: br i1 %{{.*}}, label %{{.*}}, label %{{.*}}, !llvm.loop ![[DO_ENABLE:.*]]
+}
+
+
+// CHECK: ![[FOR_DISTINCT]]     =  distinct !{![[FOR_DISTINCT]],  ![[FOR_COUNT:.*]]}
+// CHECK: ![[FOR_COUNT]]         =  !{!"llvm.loop.unroll.count", i32 8}
+// CHECK: ![[FOR_DISABLE]]   =  distinct !{![[FOR_DISABLE]],  ![[DISABLE:.*]]}
+// CHECK: ![[DISABLE]]       =  !{!"llvm.loop.unroll.disable"}
+// CHECK: ![[FOR_ENABLE]]      =  distinct !{![[FOR_ENABLE]],  ![[ENABLE:.*]]}
+// CHECK: ![[ENABLE]]          =  !{!"llvm.loop.unroll.enable"}
+// CHECK: ![[WHILE_DISTINCT]]   =  distinct !{![[WHILE_DISTINCT]],    ![[WHILE_COUNT:.*]]}
+// CHECK: ![[WHILE_COUNT]]         =  !{!"llvm.loop.unroll.count", i32 4}
+// CHECK: ![[WHILE_DISABLE]] =  distinct !{![[WHILE_DISABLE]],  ![[DISABLE]]}
+// CHECK: ![[WHILE_ENABLE]]    =  distinct !{![[WHILE_ENABLE]],     ![[ENABLE]]}
+// CHECK: ![[DO_DISTINCT]]      =  distinct !{![[DO_DISTINCT]],       ![[DO_COUNT:.*]]}
+// CHECK: ![[DO_COUNT]]         =  !{!"llvm.loop.unroll.count", i32 16}
+// CHECK: ![[DO_DISABLE]]    =  distinct !{![[DO_DISABLE]],     ![[DISABLE]]}
+// CHECK: ![[DO_ENABLE]]       =  distinct !{![[DO_ENABLE]],        ![[ENABLE]]}

[bob80905]

As stated in your comment, CheckLoopHintExpr will check each of these asserts, and return true in each case, so nullptr will be returned by this function on line 608. That said, is there really a purpose for these asserts? I don't see how they can trigger when CheckLoopHintExpr confirms these conditions.

[farzonl]

Good question. You are correct it is not needed, I added it as a form of future proofing. Let me explain my thinking.

CheckLoopHintExpr is the sema checks initially created for clang loop unrolling. We are just leveraging it. There are typically two ways of enforcing a contract on expected behavior, unit tests and asserts. I typically like to include both asserts and unit tests. of the two asserts are my preference because they encode your expectations in the code base and are not tangential to it like unit tests.

As a thought experiment lets imagine a change to CheckLoopHintExpr. This is unlikely to happen, but for the sake of argument lets say the behavior of CheckLoopHintExpr changes to one day support non constant ints. Then the nullptr would not return and we would get a fall through. We want alerts to this regression in behavior on the HLSL portion and asserts are my prefered way of knowing what happend for the reasons mentioned above but also because you get a callstack. I use to write a bunch of c# code, we had code contracts to define expectations. Thats kind of what i'm trying to do here.