Bit extract (#5847)

* promoting bitfield extraction and insertion to become intrinsics for internal compiler use

* removing duplicate intrinsics from glsl.meta.slang

* refactor: update function signatures of bitfield extraction and insertion to use uint as the parameter type for offset and bits.

---------

Co-authored-by: Nate Morrical <natemorrical@gmail.com>
Co-authored-by: Yong He <yonghe@outlook.com>

1 files changed, 376 insertions, 0 deletions

diff --git a/source/slang/slang-emit-c-like.cpp b/source/slang/slang-emit-c-like.cpp
index 50fb6655e..f38adc67e 100644
--- a/source/slang/slang-emit-c-like.cpp
+++ b/source/slang/slang-emit-c-like.cpp
@@ -2862,6 +2862,16 @@ void CLikeSourceEmitter::defaultEmitInstExpr(IRInst* inst, const EmitOpInfo& inO
             m_writer->emit(")");
             break;
         }
+    case kIROp_BitfieldExtract:
+        {
+            emitBitfieldExtractImpl(inst);
+            break;
+        }
+    case kIROp_BitfieldInsert:
+        {
+            emitBitfieldInsertImpl(inst);
+            break;
+        }
     case kIROp_PackAnyValue:
         {
             m_writer->emit("packAnyValue<");
@@ -3839,6 +3849,372 @@ void CLikeSourceEmitter::emitFuncDecorationsImpl(IRFunc* func)
     }
 }
 
+bool CLikeSourceEmitter::tryGetIntInfo(IRType* elementType, bool& isSigned, int& bitWidth)
+{
+    Slang::IROp type = elementType->getOp();
+    if (!(type >= kIROp_Int8Type && type <= kIROp_UInt64Type))
+        return false;
+    isSigned = (type >= kIROp_Int8Type && type <= kIROp_Int64Type);
+
+    Slang::IROp stype = (isSigned) ? type : Slang::IROp(type - 4);
+    bitWidth = 8 << (stype - kIROp_Int8Type);
+    return true;
+}
+
+void CLikeSourceEmitter::emitVecNOrScalar(
+    IRVectorType* vectorType,
+    std::function<void()> emitComponentLogic)
+{
+    if (vectorType)
+    {
+        int N = int(getIntVal(vectorType->getElementCount()));
+        Slang::IRType* elementType = vectorType->getElementType();
+
+        // Special handling required for CUDA target
+        if (isCUDATarget(getTargetReq()))
+        {
+            m_writer->emit("make_");
+
+            switch (elementType->getOp())
+            {
+            case kIROp_Int8Type:
+                m_writer->emit("char");
+                break;
+            case kIROp_Int16Type:
+                m_writer->emit("short");
+                break;
+            case kIROp_IntType:
+                m_writer->emit("int");
+                break;
+            case kIROp_Int64Type:
+                m_writer->emit("longlong");
+                break;
+            case kIROp_UInt8Type:
+                m_writer->emit("uchar");
+                break;
+            case kIROp_UInt16Type:
+                m_writer->emit("ushort");
+                break;
+            case kIROp_UIntType:
+                m_writer->emit("uint");
+                break;
+            case kIROp_UInt64Type:
+                m_writer->emit("ulonglong");
+                break;
+            default:
+                SLANG_ABORT_COMPILATION("Unhandled type emitting CUDA vector");
+            }
+
+            m_writer->emitRawText(std::to_string(N).c_str());
+        }
+
+        // In other languages, we can output the Slang vector type directly
+        else
+        {
+            emitType(vectorType);
+        }
+
+        m_writer->emit("(");
+        for (int i = 0; i < N; ++i)
+        {
+            emitType(elementType);
+            m_writer->emit("(");
+            emitComponentLogic();
+            m_writer->emit(")");
+            if (i != N - 1)
+                m_writer->emit(", ");
+        }
+        m_writer->emit(")");
+    }
+    else
+    {
+        m_writer->emit("(");
+        emitComponentLogic();
+        m_writer->emit(")");
+    }
+}
+
+String CLikeSourceEmitter::_emitLiteralOneWithType(int bitWidth)
+{
+    if (getTarget() == CodeGenTarget::WGSL)
+    {
+        if (bitWidth != 32)
+        {
+            SLANG_DIAGNOSE_UNEXPECTED(getSink(), SourceLoc(), "unexpected bit width");
+            return String();
+        }
+        else
+        {
+            String one;
+            one = "u32(1)";
+            return one;
+        }
+    }
+
+    String one;
+    switch (bitWidth)
+    {
+    case 8:
+        one = "uint8_t(1)";
+        break;
+    case 16:
+        one = "uint16_t(1)";
+        break;
+    case 32:
+        one = "uint32_t(1)";
+        break;
+    case 64:
+        one = "uint64_t(1)";
+        break;
+    default:
+        SLANG_DIAGNOSE_UNEXPECTED(getSink(), SourceLoc(), "unexpected bit width");
+    }
+    return one;
+}
+
+void CLikeSourceEmitter::emitBitfieldExtractImpl(IRInst* inst)
+{
+    // If unsigned, bfue := ((val>>off)&((1u<<bts)-1))
+    // Else signed, bfse := (((val>>off)&((1u<<bts)-1))<<(nbts-bts)>>(nbts-bts));
+    //
+    // Note: In WGSL, the data type for bit operators are more restricted than in other languages.
+    // The number of bits to shift must be a u32 or vecN<u32>, therefore we have to cast this
+    // operand to u32 always. Another constraint is that for "&" and "|" operators, the operands
+    // must have the same type.
+    // TODO: We can consider to bring the logic to WGSLSourceEmitter::emitBitfieldExtractImpl so
+    // that we don't have to have those special handling here.
+    Slang::IRType* dataType = inst->getDataType();
+    Slang::IRInst* val = inst->getOperand(0);
+    Slang::IRInst* off = inst->getOperand(1);
+    Slang::IRInst* bts = inst->getOperand(2);
+
+    Slang::IRType* elementType = dataType;
+    IRVectorType* vectorType = as<IRVectorType>(elementType);
+    IRVectorType* vectorTypeForShiftNumber = nullptr;
+
+    if (vectorType)
+    {
+        elementType = vectorType->getElementType();
+
+        if (getTarget() == CodeGenTarget::WGSL)
+        {
+            IRBuilder builder(elementType);
+            vectorTypeForShiftNumber =
+                builder.getVectorType(builder.getUIntType(), vectorType->getElementCount());
+        }
+        else
+        {
+            vectorTypeForShiftNumber = vectorType;
+        }
+    }
+
+    bool isSigned;
+    int bitWidth;
+    if (!tryGetIntInfo(elementType, isSigned, bitWidth))
+    {
+        SLANG_DIAGNOSE_UNEXPECTED(
+            getSink(),
+            SourceLoc(),
+            "non-integer element type given to bitfieldExtract");
+        return;
+    }
+
+    String one = _emitLiteralOneWithType(bitWidth);
+
+    // Emit open paren and type cast for later sign extension
+    if (isSigned)
+    {
+        m_writer->emit("(");
+        emitType(inst->getDataType());
+        m_writer->emit("(");
+    }
+
+    // Emit bitfield extraction ( (val >> off) & ((1u << bts) - 1) )
+    m_writer->emit("(");
+
+    // In WGSL, "&" operator requires the operands to have the same type, since the
+    // right operand '((1u << bts) - 1)' is known to be u32, we need to cast the left operand to
+    // u32.
+    if (getTarget() == CodeGenTarget::WGSL)
+    {
+        (vectorTypeForShiftNumber != nullptr) ? emitType(vectorTypeForShiftNumber)
+                                              : m_writer->emit("u32");
+    }
+
+    m_writer->emit("(");
+
+    emitOperand(val, getInfo(EmitOp::General));
+    m_writer->emit(">>");
+    emitVecNOrScalar(
+        vectorTypeForShiftNumber,
+        [&]() { emitOperand(off, getInfo(EmitOp::General)); });
+
+    m_writer->emit(")&(");
+    emitVecNOrScalar(
+        vectorTypeForShiftNumber,
+        [&]()
+        {
+            m_writer->emit("((" + one + "<<");
+            emitOperand(bts, getInfo(EmitOp::General));
+            m_writer->emit(")-" + one + ")");
+        });
+    m_writer->emit("))");
+
+    // Emit sign extension logic
+    // ( type(bitfield << (numBits - bts) ) >> (numBits - bts) )
+    //           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+    if (isSigned)
+    {
+        m_writer->emit("<<");
+        emitVecNOrScalar(
+            vectorTypeForShiftNumber,
+            [&]()
+            {
+                m_writer->emit("(");
+                m_writer->emit(bitWidth);
+                m_writer->emit("-");
+                emitOperand(bts, getInfo(EmitOp::General));
+                m_writer->emit(")");
+            });
+        m_writer->emit(")>>");
+        emitVecNOrScalar(
+            vectorTypeForShiftNumber,
+            [&]()
+            {
+                m_writer->emit("(");
+                m_writer->emit(bitWidth);
+                m_writer->emit("-");
+                emitOperand(bts, getInfo(EmitOp::General));
+                m_writer->emit(")");
+            });
+        m_writer->emit(")");
+    }
+}
+
+void CLikeSourceEmitter::emitBitfieldInsertImpl(IRInst* inst)
+{
+    // uint clearMask = ~(((1u << bits) - 1u) << offset);
+    // uint clearedBase = base & clearMask;
+    // uint maskedInsert = (insert & ((1u << bits) - 1u)) << offset;
+    // BitfieldInsert := T(uint(clearedBase) | uint(maskedInsert));
+    Slang::IRType* dataType = inst->getDataType();
+    Slang::IRInst* base = inst->getOperand(0);
+    Slang::IRInst* insert = inst->getOperand(1);
+    Slang::IRInst* off = inst->getOperand(2);
+    Slang::IRInst* bts = inst->getOperand(3);
+
+
+    Slang::IRType* elementType = dataType;
+    IRVectorType* vectorType = as<IRVectorType>(elementType);
+    IRVectorType* vectorTypeForShiftNumber = nullptr;
+
+    if (vectorType)
+    {
+        elementType = vectorType->getElementType();
+
+        if (getTarget() == CodeGenTarget::WGSL)
+        {
+            IRBuilder builder(elementType);
+            vectorTypeForShiftNumber =
+                builder.getVectorType(builder.getUIntType(), vectorType->getElementCount());
+        }
+        else
+        {
+            vectorTypeForShiftNumber = vectorType;
+        }
+    }
+
+    bool isSigned;
+    int bitWidth;
+    if (!tryGetIntInfo(elementType, isSigned, bitWidth))
+    {
+        SLANG_DIAGNOSE_UNEXPECTED(
+            getSink(),
+            SourceLoc(),
+            "non-integer element type given to bitfieldInsert");
+        return;
+    }
+
+    String one = _emitLiteralOneWithType(bitWidth);
+
+    if (isSigned)
+    {
+        emitType(inst->getDataType());
+        m_writer->emit("(");
+    }
+    m_writer->emit("(");
+
+    // emit clearedBase := uint( base & ~( ((1u << bts) - 1u) << off ) )
+
+    // In WGSL, "&" operator requires the operands to have the same type, since the
+    // right operand '~( ((1u << bts) - 1u) << off )' is known to be u32, we need to
+    // cast the left operand to u32.
+    if (getTarget() == CodeGenTarget::WGSL)
+    {
+        (vectorTypeForShiftNumber != nullptr) ? emitType(vectorTypeForShiftNumber)
+                                              : m_writer->emit("u32");
+    }
+
+    m_writer->emit("(");
+    emitOperand(base, getInfo(EmitOp::General));
+    m_writer->emit(")");
+
+    m_writer->emit("&");
+    emitVecNOrScalar(
+        vectorTypeForShiftNumber,
+        [&]()
+        {
+            m_writer->emit("~(((" + one + "<<");
+            emitOperand(bts, getInfo(EmitOp::General));
+
+            m_writer->emit(")-" + one + ")<<");
+
+            emitOperand(off, getInfo(EmitOp::General));
+            m_writer->emit(")");
+        });
+
+    // bitwise or clearedBase with maskedInsert
+    m_writer->emit(")|(");
+
+    // Emit maskedInsert := ((insert & ((1u << bits) - 1u)) << offset);
+
+    // - first emit mask := (insert & ((1u << bits) - 1u))
+    m_writer->emit("(");
+
+    // For the same reason as above, we need to cast the left operand to u32 for WGSL target.
+    if (getTarget() == CodeGenTarget::WGSL)
+    {
+        (vectorTypeForShiftNumber != nullptr) ? emitType(vectorTypeForShiftNumber)
+                                              : m_writer->emit("u32");
+    }
+    m_writer->emit("(");
+    emitOperand(insert, getInfo(EmitOp::General));
+    m_writer->emit(")");
+
+    m_writer->emit("&");
+    emitVecNOrScalar(
+        vectorTypeForShiftNumber,
+        [&]()
+        {
+            m_writer->emit("(" + one + "<<");
+            emitOperand(bts, getInfo(EmitOp::General));
+            m_writer->emit(")-" + one);
+        });
+    m_writer->emit(")");
+
+    // then emit shift := << offset
+    m_writer->emit("<<");
+    emitVecNOrScalar(
+        vectorTypeForShiftNumber,
+        [&]() { emitOperand(off, getInfo(EmitOp::General)); });
+    m_writer->emit(")");
+
+    if (isSigned)
+    {
+        m_writer->emit(")");
+    }
+}
+
 void CLikeSourceEmitter::emitStruct(IRStructType* structType)
 {
     ensureTypePrelude(structType);