[SPIR-V] Add support for 64-bit switch sel (microsoft#6049)

Allow 64-bit integer variables as switch statement selectors and clarify
the error message for literal selectors (related to microsoft#4338).

Author: sudonatalie

tools/clang/include/clang/SPIRV/SpirvBuilder.h

@@ -377,7 +377,7 @@ class SpirvBuilder {
   void
   createSwitch(SpirvBasicBlock *mergeLabel, SpirvInstruction *selector,
                SpirvBasicBlock *defaultLabel,
-               llvm::ArrayRef<std::pair<uint32_t, SpirvBasicBlock *>> target,
+               llvm::ArrayRef<std::pair<llvm::APInt, SpirvBasicBlock *>> target,
                SourceLocation, SourceRange);
 
   /// \brief Creates a fragment-shader discard via by emitting OpKill.

tools/clang/include/clang/SPIRV/SpirvInstruction.h

@@ -824,7 +824,7 @@ class SpirvSwitch : public SpirvBranching {
   SpirvSwitch(
       SourceLocation loc, SpirvInstruction *selector,
       SpirvBasicBlock *defaultLabel,
-      llvm::ArrayRef<std::pair<uint32_t, SpirvBasicBlock *>> &targetsVec);
+      llvm::ArrayRef<std::pair<llvm::APInt, SpirvBasicBlock *>> &targetsVec);
 
   DEFINE_RELEASE_MEMORY_FOR_CLASS(SpirvSwitch)
 
@@ -837,7 +837,7 @@ class SpirvSwitch : public SpirvBranching {
 
   SpirvInstruction *getSelector() const { return selector; }
   SpirvBasicBlock *getDefaultLabel() const { return defaultLabel; }
-  llvm::ArrayRef<std::pair<uint32_t, SpirvBasicBlock *>> getTargets() const {
+  llvm::ArrayRef<std::pair<llvm::APInt, SpirvBasicBlock *>> getTargets() const {
     return targets;
   }
   // Returns the branch label that will be taken for the given literal.
@@ -853,7 +853,7 @@ class SpirvSwitch : public SpirvBranching {
 private:
   SpirvInstruction *selector;
   SpirvBasicBlock *defaultLabel;
-  llvm::SmallVector<std::pair<uint32_t, SpirvBasicBlock *>, 4> targets;
+  llvm::SmallVector<std::pair<llvm::APInt, SpirvBasicBlock *>, 4> targets;
 };
 
 /// \brief OpUnreachable instruction

tools/clang/lib/SPIRV/EmitVisitor.cpp

@@ -846,7 +846,7 @@ bool EmitVisitor::visit(SpirvSwitch *inst) {
   curInst.push_back(
       getOrAssignResultId<SpirvBasicBlock>(inst->getDefaultLabel()));
   for (const auto &target : inst->getTargets()) {
-    curInst.push_back(target.first);
+    typeHandler.emitIntLiteral(target.first, curInst);
     curInst.push_back(getOrAssignResultId<SpirvBasicBlock>(target.second));
   }
   finalizeInstruction(&mainBinary);
@@ -2611,6 +2611,12 @@ template <typename vecType>
 void EmitTypeHandler::emitIntLiteral(const SpirvConstantInteger *intLiteral,
                                      vecType &outInst) {
   const auto &literalVal = intLiteral->getValue();
+  emitIntLiteral(literalVal, outInst);
+}
+
+template <typename vecType>
+void EmitTypeHandler::emitIntLiteral(const llvm::APInt &literalVal,
+                                     vecType &outInst) {
   bool positive = !literalVal.isNegative();
   if (literalVal.getBitWidth() <= 32) {
     outInst.push_back(positive ? literalVal.getZExtValue()

tools/clang/lib/SPIRV/EmitVisitor.h

@@ -114,6 +114,8 @@ class EmitTypeHandler {
   void emitFloatLiteral(const SpirvConstantFloat *, vecType &outInst);
   template <typename vecType>
   void emitIntLiteral(const SpirvConstantInteger *, vecType &outInst);
+  template <typename vecType>
+  void emitIntLiteral(const llvm::APInt &literalVal, vecType &outInst);
 
 private:
   void initTypeInstruction(spv::Op op);

tools/clang/lib/SPIRV/SpirvBuilder.cpp

@@ -761,7 +761,7 @@ SpirvSelect *SpirvBuilder::createSelect(QualType resultType,
 void SpirvBuilder::createSwitch(
     SpirvBasicBlock *mergeLabel, SpirvInstruction *selector,
     SpirvBasicBlock *defaultLabel,
-    llvm::ArrayRef<std::pair<uint32_t, SpirvBasicBlock *>> target,
+    llvm::ArrayRef<std::pair<llvm::APInt, SpirvBasicBlock *>> target,
     SourceLocation loc, SourceRange range) {
   assert(insertPoint && "null insert point");
   // Create the OpSelectioMerege.

tools/clang/lib/SPIRV/SpirvEmitter.cpp

@@ -25,8 +25,10 @@
 #include "clang/SPIRV/AstTypeProbe.h"
 #include "clang/SPIRV/String.h"
 #include "clang/Sema/Sema.h"
+#include "llvm/ADT/APInt.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/StringExtras.h"
+#include "llvm/Support/Casting.h"
 
 #ifdef SUPPORT_QUERY_GIT_COMMIT_INFO
 #include "clang/Basic/Version.h"
@@ -13422,7 +13424,7 @@ bool SpirvEmitter::allSwitchCasesAreIntegerLiterals(const Stmt *root) {
 
 void SpirvEmitter::discoverAllCaseStmtInSwitchStmt(
     const Stmt *root, SpirvBasicBlock **defaultBB,
-    std::vector<std::pair<uint32_t, SpirvBasicBlock *>> *targets) {
+    std::vector<std::pair<llvm::APInt, SpirvBasicBlock *>> *targets) {
   if (!root)
     return;
 
@@ -13442,7 +13444,7 @@ void SpirvEmitter::discoverAllCaseStmtInSwitchStmt(
   }
 
   std::string caseLabel;
-  uint32_t caseValue = 0;
+  llvm::APInt caseValue;
   if (defaultStmt) {
     // This is the default branch.
     caseLabel = "switch.default";
@@ -13452,15 +13454,10 @@ void SpirvEmitter::discoverAllCaseStmtInSwitchStmt(
     // case <literal_integer>: {...; break;}
     const Expr *caseExpr = caseStmt->getLHS();
     assert(caseExpr && caseExpr->isEvaluatable(astContext));
-    auto bitWidth = astContext.getIntWidth(caseExpr->getType());
-    if (bitWidth != 32)
-      emitError(
-          "non-32bit integer case value in switch statement unimplemented",
-          caseExpr->getExprLoc());
     Expr::EvalResult evalResult;
     caseExpr->EvaluateAsRValue(evalResult, astContext);
-    const int64_t value = evalResult.Val.getInt().getSExtValue();
-    caseValue = static_cast<uint32_t>(value);
+    caseValue = evalResult.Val.getInt();
+    const int64_t value = caseValue.getSExtValue();
     caseLabel = "switch." + std::string(value < 0 ? "n" : "") +
                 llvm::itostr(std::abs(value));
   }
@@ -13535,6 +13532,13 @@ void SpirvEmitter::processSwitchStmtUsingSpirvOpSwitch(
     doDeclStmt(condVarDeclStmt);
 
   auto *cond = switchStmt->getCond();
+  if (llvm::dyn_cast<IntegerLiteral>(cond)) {
+    emitError(
+        "integer literal selectors in switch statements not yet implemented",
+        cond->getLocStart());
+    return;
+  }
+
   auto *selector = doExpr(cond);
 
   // We need a merge block regardless of the number of switch cases.
@@ -13547,7 +13551,7 @@ void SpirvEmitter::processSwitchStmtUsingSpirvOpSwitch(
   auto *defaultBB = mergeBB;
 
   // (literal, labelId) pairs to pass to the OpSwitch instruction.
-  std::vector<std::pair<uint32_t, SpirvBasicBlock *>> targets;
+  std::vector<std::pair<llvm::APInt, SpirvBasicBlock *>> targets;
   discoverAllCaseStmtInSwitchStmt(switchStmt->getBody(), &defaultBB, &targets);
 
   // Create the OpSelectionMerge and OpSwitch.

tools/clang/lib/SPIRV/SpirvEmitter.h

@@ -919,7 +919,7 @@ class SpirvEmitter : public ASTConsumer {
   /// method panics if it finds a case value that is not an integer literal.
   void discoverAllCaseStmtInSwitchStmt(
       const Stmt *root, SpirvBasicBlock **defaultBB,
-      std::vector<std::pair<uint32_t, SpirvBasicBlock *>> *targets);
+      std::vector<std::pair<llvm::APInt, SpirvBasicBlock *>> *targets);
 
   /// Flattens structured AST of the given switch statement into a vector of AST
   /// nodes and stores into flatSwitch.

tools/clang/lib/SPIRV/SpirvInstruction.cpp

@@ -384,7 +384,7 @@ SpirvReturn::SpirvReturn(SourceLocation loc, SpirvInstruction *retVal,
 SpirvSwitch::SpirvSwitch(
     SourceLocation loc, SpirvInstruction *selectorInst,
     SpirvBasicBlock *defaultLbl,
-    llvm::ArrayRef<std::pair<uint32_t, SpirvBasicBlock *>> &targetsVec)
+    llvm::ArrayRef<std::pair<llvm::APInt, SpirvBasicBlock *>> &targetsVec)
     : SpirvBranching(IK_Switch, spv::Op::OpSwitch, loc), selector(selectorInst),
       defaultLabel(defaultLbl), targets(targetsVec.begin(), targetsVec.end()) {}
 

tools/clang/test/CodeGenSPIRV_Lit/cf.switch.opswitch.hlsl

@@ -357,4 +357,31 @@ void main() {
     break;
   }
 
+
+  /////////////////////////////////////////
+  // 64-bit integer variable as selector //
+  /////////////////////////////////////////
+  int64_t longsel;
+// CHECK:       [[longSelector:%[0-9]+]] = OpLoad %long %longsel
+// CHECK-NEXT:                          OpSelectionMerge %switch_merge_10 None
+// CHECK-NEXT:                          OpSwitch [[longSelector]] %switch_merge_10 -1 %switch_n1
+  switch (longsel) {
+  case -1:
+    result = 0;
+    break;
+  }
+
+
+  //////////////////////////////////////////////////
+  // 64-bit unsigned integer variable as selector //
+  //////////////////////////////////////////////////
+  uint64_t ulongsel;
+// CHECK:      [[ulongSelector:%[0-9]+]] = OpLoad %ulong %ulongsel
+// CHECK-NEXT:                          OpSelectionMerge %switch_merge_11 None
+// CHECK-NEXT:                          OpSwitch [[ulongSelector]] %switch_merge_11 12345678910 %switch_12345678910
+  switch (ulongsel) {
+  case 12345678910:
+    result = 0;
+    break;
+  }
 }

tools/clang/test/CodeGenSPIRV_Lit/cf.switch.opswitch.literal.hlsl

@@ -0,0 +1,11 @@
+// RUN: not %dxc -T ps_6_0 -fcgl -spirv %s 2>&1 | FileCheck %s
+
+// CHECK: error: integer literal selectors in switch statements not yet implemented
+
+float main() : SV_TARGET {
+  switch (0) {
+  case 0:
+    return 1;
+  }
+  return 0;
+}