[HLSL] Add support to branch/flatten attributes to switch

[joaosaffran]

closes: #125754

[llvmbot]

@llvm/pr-subscribers-clang-codegen

@llvm/pr-subscribers-backend-directx

Author: None (joaosaffran)

Changes

closes: #125754

---

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

5 Files Affected:

• (modified) clang/include/clang/Basic/Attr.td (+2-2)
• (modified) clang/lib/CodeGen/CGStmt.cpp (+23)
• (modified) clang/test/AST/HLSL/HLSLControlFlowHint.hlsl (+59)
• (modified) llvm/test/CodeGen/DirectX/HLSLControlFlowHint.ll (+131)
• (modified) llvm/test/CodeGen/SPIRV/structurizer/HLSLControlFlowHint.ll (+122-3)

diff --git a/clang/include/clang/Basic/Attr.td b/clang/include/clang/Basic/Attr.td
index d9840dad6a5e2..948fc99c9b083 100644
--- a/clang/include/clang/Basic/Attr.td
+++ b/clang/include/clang/Basic/Attr.td
@@ -4394,8 +4394,8 @@ def HLSLControlFlowHint: StmtAttr {
   /// [branch]
   /// [flatten]
   let Spellings = [Microsoft<"branch">, Microsoft<"flatten">];
-  let Subjects = SubjectList<[IfStmt],
-                              ErrorDiag, "'if' statements">;
+  let Subjects = SubjectList<[IfStmt, SwitchStmt],
+                              ErrorDiag, "'if' and 'switch' statements">;
   let LangOpts = [HLSL];
   let Documentation = [InternalOnly];
 }
diff --git a/clang/lib/CodeGen/CGStmt.cpp b/clang/lib/CodeGen/CGStmt.cpp
index 9860f23dc9e28..ef963fc4984d4 100644
--- a/clang/lib/CodeGen/CGStmt.cpp
+++ b/clang/lib/CodeGen/CGStmt.cpp
@@ -2276,6 +2276,29 @@ void CodeGenFunction::EmitSwitchStmt(const SwitchStmt &S) {
   // failure.
   llvm::BasicBlock *DefaultBlock = createBasicBlock("sw.default");
   SwitchInsn = Builder.CreateSwitch(CondV, DefaultBlock);
+  switch (HLSLControlFlowAttr) {
+  case HLSLControlFlowHintAttr::Microsoft_branch:
+  case HLSLControlFlowHintAttr::Microsoft_flatten: {
+    llvm::MDBuilder MDHelper(CGM.getLLVMContext());
+
+    llvm::ConstantInt *BranchHintConstant =
+        HLSLControlFlowAttr ==
+                HLSLControlFlowHintAttr::Spelling::Microsoft_branch
+            ? llvm::ConstantInt::get(CGM.Int32Ty, 1)
+            : llvm::ConstantInt::get(CGM.Int32Ty, 2);
+
+    SmallVector<llvm::Metadata *, 2> Vals(
+        {MDHelper.createString("hlsl.controlflow.hint"),
+         MDHelper.createConstant(BranchHintConstant)});
+    SwitchInsn->setMetadata("hlsl.controlflow.hint",
+                            llvm::MDNode::get(CGM.getLLVMContext(), Vals));
+    break;
+  }
+  // This is required to avoid warnings during compilation
+  case HLSLControlFlowHintAttr::SpellingNotCalculated:
+    break;
+  }
+
   if (PGO.haveRegionCounts()) {
     // Walk the SwitchCase list to find how many there are.
     uint64_t DefaultCount = 0;
diff --git a/clang/test/AST/HLSL/HLSLControlFlowHint.hlsl b/clang/test/AST/HLSL/HLSLControlFlowHint.hlsl
index c1e6d969c8d31..18263bedbe3ec 100644
--- a/clang/test/AST/HLSL/HLSLControlFlowHint.hlsl
+++ b/clang/test/AST/HLSL/HLSLControlFlowHint.hlsl
@@ -41,3 +41,62 @@ export int no_attr(int X){
 
     return resp;
 }
+
+// CHECK: FunctionDecl {{.*}} used flatten_switch 'int (int)'
+// CHECK: AttributedStmt
+// CHECK-NEXT: HLSLControlFlowHintAttr {{.*}} flatten
+export int flatten_switch(int X){
+    int resp;
+    [flatten] 
+    switch (X) {
+        case 0:
+        resp = -X;
+        break;
+    case 1:
+         resp = X+X;
+        break;
+    case 2:
+        resp = X * X; break;
+    }
+
+    return resp;
+}
+
+// CHECK: FunctionDecl {{.*}} used branch_switch 'int (int)'
+// CHECK: AttributedStmt
+// CHECK-NEXT: HLSLControlFlowHintAttr {{.*}} branch
+export int branch_switch(int X){
+    int resp;
+    [branch] 
+    switch (X) {
+     case 0:
+        resp = -X;
+        break;
+    case 1:
+        resp = X+X;
+        break;
+    case 2:
+        resp = X * X; break;
+    }
+
+    return resp;
+}
+
+// CHECK: FunctionDecl {{.*}} used no_attr_switch 'int (int)'
+// CHECK-NOT: AttributedStmt
+// CHECK-NOT: HLSLControlFlowHintAttr
+export int no_attr_switch(int X){
+    int resp;
+    switch (X) {
+        case 0:
+        resp = -X;
+        break;
+    case 1:
+         resp = X+X;
+        break;
+    case 2:
+        resp = X * X; break;
+    }
+
+    return resp;
+}
diff --git a/llvm/test/CodeGen/DirectX/HLSLControlFlowHint.ll b/llvm/test/CodeGen/DirectX/HLSLControlFlowHint.ll
index 6a5274429930e..00dd374daf460 100644
--- a/llvm/test/CodeGen/DirectX/HLSLControlFlowHint.ll
+++ b/llvm/test/CodeGen/DirectX/HLSLControlFlowHint.ll
@@ -91,6 +91,137 @@ if.end:                                           ; preds = %if.else, %if.then
   %3 = load i32, ptr %resp, align 4
   ret i32 %3
 }
+
+; CHECK: define i32 @flatten_switch(i32 %X)
+; CHECK-NOT: hlsl.controlflow.hint
+; CHECK:      switch i32 %0, label %sw.epilog [
+; CHECK-NEXT:   i32 0, label %sw.bb
+; CHECK-NEXT:   i32 1, label %sw.bb1
+; CHECK-NEXT:   i32 2, label %sw.bb2
+; CHECK-NEXT: ], !dx.controlflow.hints [[HINT_FLATTEN]]
+define i32 @flatten_switch(i32 %X) #0 {
+entry:
+  %X.addr = alloca i32, align 4
+  %resp = alloca i32, align 4
+  store i32 %X, ptr %X.addr, align 4
+  %0 = load i32, ptr %X.addr, align 4
+  switch i32 %0, label %sw.epilog [
+    i32 0, label %sw.bb
+    i32 1, label %sw.bb1
+    i32 2, label %sw.bb2
+  ], !hlsl.controlflow.hint !1
+
+sw.bb:                                            ; preds = %entry
+  %1 = load i32, ptr %X.addr, align 4
+  %sub = sub nsw i32 0, %1
+  store i32 %sub, ptr %resp, align 4
+  br label %sw.epilog
+
+sw.bb1:                                           ; preds = %entry
+  %2 = load i32, ptr %X.addr, align 4
+  %3 = load i32, ptr %X.addr, align 4
+  %add = add nsw i32 %2, %3
+  store i32 %add, ptr %resp, align 4
+  br label %sw.epilog
+
+sw.bb2:                                           ; preds = %entry
+  %4 = load i32, ptr %X.addr, align 4
+  %5 = load i32, ptr %X.addr, align 4
+  %mul = mul nsw i32 %4, %5
+  store i32 %mul, ptr %resp, align 4
+  br label %sw.epilog
+
+sw.epilog:                                        ; preds = %entry, %sw.bb2, %sw.bb1, %sw.bb
+  %6 = load i32, ptr %resp, align 4
+  ret i32 %6
+}
+
+
+; CHECK: define i32 @branch_switch(i32 %X)
+; CHECK-NOT: hlsl.controlflow.hint
+; CHECK:      switch i32 %0, label %sw.epilog [
+; CHECK-NEXT:   i32 0, label %sw.bb
+; CHECK-NEXT:   i32 1, label %sw.bb1
+; CHECK-NEXT:   i32 2, label %sw.bb2
+; CHECK-NEXT: ], !dx.controlflow.hints [[HINT_BRANCH]]
+define i32 @branch_switch(i32 %X) #0 {
+entry:
+  %X.addr = alloca i32, align 4
+  %resp = alloca i32, align 4
+  store i32 %X, ptr %X.addr, align 4
+  %0 = load i32, ptr %X.addr, align 4
+  switch i32 %0, label %sw.epilog [
+    i32 0, label %sw.bb
+    i32 1, label %sw.bb1
+    i32 2, label %sw.bb2
+  ], !hlsl.controlflow.hint !0
+
+sw.bb:                                            ; preds = %entry
+  %1 = load i32, ptr %X.addr, align 4
+  %sub = sub nsw i32 0, %1
+  store i32 %sub, ptr %resp, align 4
+  br label %sw.epilog
+
+sw.bb1:                                           ; preds = %entry
+  %2 = load i32, ptr %X.addr, align 4
+  %3 = load i32, ptr %X.addr, align 4
+  %add = add nsw i32 %2, %3
+  store i32 %add, ptr %resp, align 4
+  br label %sw.epilog
+
+sw.bb2:                                           ; preds = %entry
+  %4 = load i32, ptr %X.addr, align 4
+  %5 = load i32, ptr %X.addr, align 4
+  %mul = mul nsw i32 %4, %5
+  store i32 %mul, ptr %resp, align 4
+  br label %sw.epilog
+
+sw.epilog:                                        ; preds = %entry, %sw.bb2, %sw.bb1, %sw.bb
+  %6 = load i32, ptr %resp, align 4
+  ret i32 %6
+}
+
+
+; CHECK: define i32 @no_attr_switch(i32 %X)
+; CHECK-NOT: hlsl.controlflow.hint
+; CHECK-NOT: !dx.controlflow.hints
+define i32 @no_attr_switch(i32 %X) #0 {
+entry:
+  %X.addr = alloca i32, align 4
+  %resp = alloca i32, align 4
+  store i32 %X, ptr %X.addr, align 4
+  %0 = load i32, ptr %X.addr, align 4
+  switch i32 %0, label %sw.epilog [
+    i32 0, label %sw.bb
+    i32 1, label %sw.bb1
+    i32 2, label %sw.bb2
+  ]
+
+sw.bb:                                            ; preds = %entry
+  %1 = load i32, ptr %X.addr, align 4
+  %sub = sub nsw i32 0, %1
+  store i32 %sub, ptr %resp, align 4
+  br label %sw.epilog
+
+sw.bb1:                                           ; preds = %entry
+  %2 = load i32, ptr %X.addr, align 4
+  %3 = load i32, ptr %X.addr, align 4
+  %add = add nsw i32 %2, %3
+  store i32 %add, ptr %resp, align 4
+  br label %sw.epilog
+
+sw.bb2:                                           ; preds = %entry
+  %4 = load i32, ptr %X.addr, align 4
+  %5 = load i32, ptr %X.addr, align 4
+  %mul = mul nsw i32 %4, %5
+  store i32 %mul, ptr %resp, align 4
+  br label %sw.epilog
+
+sw.epilog:                                        ; preds = %entry, %sw.bb2, %sw.bb1, %sw.bb
+  %6 = load i32, ptr %resp, align 4
+  ret i32 %6
+}
+
 ; CHECK-NOT: hlsl.controlflow.hint
 ; CHECK: [[HINT_BRANCH]] = !{!"dx.controlflow.hints", i32 1}
 ; CHECK: [[HINT_FLATTEN]] = !{!"dx.controlflow.hints", i32 2}
diff --git a/llvm/test/CodeGen/SPIRV/structurizer/HLSLControlFlowHint.ll b/llvm/test/CodeGen/SPIRV/structurizer/HLSLControlFlowHint.ll
index 848eaf70f5a19..9c6f977dc9b34 100644
--- a/llvm/test/CodeGen/SPIRV/structurizer/HLSLControlFlowHint.ll
+++ b/llvm/test/CodeGen/SPIRV/structurizer/HLSLControlFlowHint.ll
@@ -5,7 +5,7 @@
 define spir_func noundef i32 @test_branch(i32 noundef %X) {
 entry:
 ; CHECK-LABEL: ; -- Begin function test_branch
-; OpSelectionMerge %[[#]] DontFlatten
+; CHECK: OpSelectionMerge %[[#]] DontFlatten
   %X.addr = alloca i32, align 4
   %resp = alloca i32, align 4
   store i32 %X, ptr %X.addr, align 4
@@ -34,7 +34,7 @@ if.end:                                           ; preds = %if.else, %if.then
 define spir_func noundef i32 @test_flatten(i32 noundef %X) {
 entry:
 ; CHECK-LABEL: ; -- Begin function test_flatten
-; OpSelectionMerge %[[#]] Flatten
+; CHECK: OpSelectionMerge %[[#]] Flatten
   %X.addr = alloca i32, align 4
   %resp = alloca i32, align 4
   store i32 %X, ptr %X.addr, align 4
@@ -62,7 +62,7 @@ if.end:                                           ; preds = %if.else, %if.then
 define spir_func noundef i32 @test_no_attr(i32 noundef %X) {
 entry:
 ; CHECK-LABEL: ; -- Begin function test_no_attr
-; OpSelectionMerge %[[#]] None
+; CHECK: OpSelectionMerge %[[#]] None
   %X.addr = alloca i32, align 4
   %resp = alloca i32, align 4
   store i32 %X, ptr %X.addr, align 4
@@ -87,5 +87,124 @@ if.end:                                           ; preds = %if.else, %if.then
   ret i32 %3
 }
 
+define spir_func noundef i32 @flatten_switch(i32 noundef %X) {
+entry:
+; CHECK-LABEL: ; -- Begin function flatten_switch
+; CHECK: OpSelectionMerge %[[#]] Flatten
+  %X.addr = alloca i32, align 4
+  %resp = alloca i32, align 4
+  store i32 %X, ptr %X.addr, align 4
+  %0 = load i32, ptr %X.addr, align 4
+  switch i32 %0, label %sw.epilog [
+    i32 0, label %sw.bb
+    i32 1, label %sw.bb1
+    i32 2, label %sw.bb2
+  ], !hlsl.controlflow.hint !1
+
+sw.bb:                                            ; preds = %entry
+  %1 = load i32, ptr %X.addr, align 4
+  %sub = sub nsw i32 0, %1
+  store i32 %sub, ptr %resp, align 4
+  br label %sw.epilog
+
+sw.bb1:                                           ; preds = %entry
+  %2 = load i32, ptr %X.addr, align 4
+  %3 = load i32, ptr %X.addr, align 4
+  %add = add nsw i32 %2, %3
+  store i32 %add, ptr %resp, align 4
+  br label %sw.epilog
+
+sw.bb2:                                           ; preds = %entry
+  %4 = load i32, ptr %X.addr, align 4
+  %5 = load i32, ptr %X.addr, align 4
+  %mul = mul nsw i32 %4, %5
+  store i32 %mul, ptr %resp, align 4
+  br label %sw.epilog
+
+sw.epilog:                                        ; preds = %entry, %sw.bb2, %sw.bb1, %sw.bb
+  %6 = load i32, ptr %resp, align 4
+  ret i32 %6
+}
+
+
+define spir_func noundef i32 @branch_switch(i32 noundef %X) {
+  entry:
+  ; CHECK-LABEL: ; -- Begin function branch_switch
+  ; CHECK: OpSelectionMerge %[[#]] DontFlatten
+  %X.addr = alloca i32, align 4
+  %resp = alloca i32, align 4
+  store i32 %X, ptr %X.addr, align 4
+  %0 = load i32, ptr %X.addr, align 4
+  switch i32 %0, label %sw.epilog [
+    i32 0, label %sw.bb
+    i32 1, label %sw.bb1
+    i32 2, label %sw.bb2
+  ], !hlsl.controlflow.hint !0
+
+sw.bb:                                            ; preds = %entry
+  %1 = load i32, ptr %X.addr, align 4
+  %sub = sub nsw i32 0, %1
+  store i32 %sub, ptr %resp, align 4
+  br label %sw.epilog
+
+sw.bb1:                                           ; preds = %entry
+  %2 = load i32, ptr %X.addr, align 4
+  %3 = load i32, ptr %X.addr, align 4
+  %add = add nsw i32 %2, %3
+  store i32 %add, ptr %resp, align 4
+  br label %sw.epilog
+
+sw.bb2:                                           ; preds = %entry
+  %4 = load i32, ptr %X.addr, align 4
+  %5 = load i32, ptr %X.addr, align 4
+  %mul = mul nsw i32 %4, %5
+  store i32 %mul, ptr %resp, align 4
+  br label %sw.epilog
+
+sw.epilog:                                        ; preds = %entry, %sw.bb2, %sw.bb1, %sw.bb
+  %6 = load i32, ptr %resp, align 4
+  ret i32 %6
+}
+
+
+define spir_func noundef i32 @no_attr_switch(i32 noundef %X) {
+  ; CHECK-LABEL: ; -- Begin function no_attr_switch
+; CHECK: OpSelectionMerge %[[#]] None
+entry:
+  %X.addr = alloca i32, align 4
+  %resp = alloca i32, align 4
+  store i32 %X, ptr %X.addr, align 4
+  %0 = load i32, ptr %X.addr, align 4
+  switch i32 %0, label %sw.epilog [
+    i32 0, label %sw.bb
+    i32 1, label %sw.bb1
+    i32 2, label %sw.bb2
+  ]
+
+sw.bb:                                            ; preds = %entry
+  %1 = load i32, ptr %X.addr, align 4
+  %sub = sub nsw i32 0, %1
+  store i32 %sub, ptr %resp, align 4
+  br label %sw.epilog
+
+sw.bb1:                                           ; preds = %entry
+  %2 = load i32, ptr %X.addr, align 4
+  %3 = load i32, ptr %X.addr, align 4
+  %add = add nsw i32 %2, %3
+  store i32 %add, ptr %resp, align 4
+  br label %sw.epilog
+
+sw.bb2:                                           ; preds = %entry
+  %4 = load i32, ptr %X.addr, align 4
+  %5 = load i32, ptr %X.addr, align 4
+  %mul = mul nsw i32 %4, %5
+  store i32 %mul, ptr %resp, align 4
+  br label %sw.epilog
+
+sw.epilog:                                        ; preds = %entry, %sw.bb2, %sw.bb1, %sw.bb
+  %6 = load i32, ptr %resp, align 4
+  ret i32 %6
+}
+
 !0 = !{!"hlsl.controlflow.hint", i32 1}
 !1 = !{!"hlsl.controlflow.hint", i32 2}