Ensure to emit 32-bit type for bitfield operations for SPIRV (#7046)

Close #7014

2 files changed, 76 insertions, 0 deletions

diff --git a/source/slang/slang-ir-spirv-legalize.cpp b/source/slang/slang-ir-spirv-legalize.cpp
index 82b424daa..0287ae81a 100644
--- a/source/slang/slang-ir-spirv-legalize.cpp
+++ b/source/slang/slang-ir-spirv-legalize.cpp
@@ -1471,6 +1471,39 @@ struct SPIRVLegalizationContext : public SourceEmitterBase
         }
     }
 
+    // TODO: Currently SPIRV doesn't support non-32-bit integer types for bitfield extract and
+    // insert. We will relax this restriction once this is done:
+    // https://github.com/shader-slang/slang/issues/7015.
+    void processBitFieldOp(IRInst* inst)
+    {
+        auto dataType = inst->getDataType();
+        IRVectorType* vectorType = as<IRVectorType>(dataType);
+        Slang::IRType* elementType = dataType;
+        if (vectorType)
+            elementType = vectorType->getElementType();
+
+        const IntInfo i = getIntTypeInfo(elementType);
+
+        // SPIRV doesn't support non-32bit integer types, so we need to convert
+        if (i.width < 32)
+        {
+            IRBuilder builder(inst);
+            builder.setInsertBefore(inst);
+            IRType* intType = i.isSigned ? builder.getIntType() : builder.getUIntType();
+            auto targetType = vectorType
+                                  ? builder.getVectorType(intType, vectorType->getElementCount())
+                                  : intType;
+            auto baseInst = builder.emitCast(targetType, inst->getOperand(0));
+            builder.replaceOperand(inst->getOperands(), baseInst);
+            if (inst->getOp() == kIROp_BitfieldInsert)
+            {
+                auto insertInst = builder.emitCast(targetType, inst->getOperand(1));
+                builder.replaceOperand(inst->getOperands() + 1, insertInst);
+            }
+            inst->setFullType(intType);
+        }
+    }
+
     void legalizeSPIRVEntryPoint(IRFunc* func, IREntryPointDecoration* entryPointDecor)
     {
         auto stage = entryPointDecor->getProfile().getStage();
@@ -1704,6 +1737,11 @@ struct SPIRVLegalizationContext : public SourceEmitterBase
             case kIROp_SPIRVAsm:
                 processSPIRVAsm(as<IRSPIRVAsm>(inst));
                 break;
+            case kIROp_BitfieldExtract:
+            case kIROp_BitfieldInsert:
+                processBitFieldOp(inst);
+                break;
+
             case kIROp_DebugValue:
                 if (!isSimpleDataType(as<IRDebugValue>(inst)->getDebugVar()->getDataType()))
                     inst->removeAndDeallocate();
diff --git a/tests/language-feature/bitfield/bitfield-insert-extract-non32bit.slang b/tests/language-feature/bitfield/bitfield-insert-extract-non32bit.slang
new file mode 100644
index 000000000..0edbb049a
--- /dev/null
+++ b/tests/language-feature/bitfield/bitfield-insert-extract-non32bit.slang
@@ -0,0 +1,38 @@
+//TEST:COMPARE_COMPUTE(filecheck-buffer=BUF): -vk -output-using-type
+//TEST_INPUT:ubuffer(data=[0 0 0 0 0 0 0 0], stride=4):out,name=outputBuffer
+RWStructuredBuffer<int32_t> outputBuffer;
+
+//TEST_INPUT:set input = ubuffer(data=[0xC0A0], stride=4)
+RWStructuredBuffer<uint32_t> input;
+
+[shader("compute")]
+[numthreads(1, 1, 1)]
+void computeMain(
+    uint3 tid : SV_DispatchThreadID)
+{
+    uint16_t a = uint16_t(input[0] & 0xFFFF);
+    int16_t b = bit_cast<int16_t>(a);
+
+    outputBuffer[0] = bitfieldExtract(a, 14, 2);
+    outputBuffer[1] = bitfieldExtract(b, 13, 3);    // 0b110, with sign extend, it will be 0xFFFFFFFE => -2
+    // BUF: 3
+    // BUF: -2
+
+    uint8_t c = uint8_t(input[0] & 0xFF);
+    int8_t d = bit_cast<int8_t>(c);
+
+    outputBuffer[2] = bitfieldExtract(c, 5, 3);     // b101 = 5
+    outputBuffer[3] = bitfieldExtract(d, 5, 3);     // 0b101, with sign extend, it will be 0xFFFFFFFD => -3
+    // BUF: 5
+    // BUF: -3
+
+    outputBuffer[4] = bitfieldInsert(a, 0x03U, 8, 2);   // 0xC3A0 = 50080 (unsigned)
+    outputBuffer[5] = bitfieldInsert(b, 0x03,  8, 2);    // 0xC3A0 = -15456 (signed)
+    // BUF: 50080
+    // BUF: -15456
+
+    outputBuffer[6] = bitfieldInsert(c, 0x03U, 3, 2);   // 0xB8 = 184  (unsigned)
+    outputBuffer[7] = bitfieldInsert(d, 0x03,  3, 2);    // 0xB8 = -72 (signed)
+    // BUF: 184
+    // BUF: -72
+}