// slang-ir-byte-address-legalize.cpp
#include "slang-ir-byte-address-legalize.h"

// This file implements an IR pass that translates load/store operations
// on byte-address buffers to be legal for a chosen target.
//
// Currently this pass only applies to the operations generated for
// the generic `*ByteAddressBuffer.Load<T>` and `.Store<T>` operations,
// and not the non-generic versions that traffic in `uint` (e.g.,
// `Load2` or `Store3`).

#include "slang-ir-insts.h"
#include "slang-ir-layout.h"
#include "slang-rich-diagnostics.h"

namespace Slang
{

bool isCPUTarget(TargetRequest* targetReq);

// As is typical for IR passes in Slang, we will encapsulate the state
// while we process the code in a context type.
//
struct ByteAddressBufferLegalizationContext
{
    // We need access to the original session, as well as the options
    // that control what constructs we legalize, and how.
    //
    Session* m_session = nullptr;
    TargetProgram* m_targetProgram = nullptr;
    TargetRequest* m_target = nullptr;
    ByteAddressBufferLegalizationOptions m_options;

    DiagnosticSink* m_sink = nullptr;
    // We will also use a central IR builder when generating new
    // code as part of legalization (rather than create/destroy
    // IR builders on the fly).
    //
    IRModule* m_module;
    IRBuilder m_builder;

    // Track ByteAddressBuffer parameters that have been replaced by StructuredBuffer.
    // These need to be removed at the end to avoid leaving invalid types in the IR
    // for targets like SPIR-V that don't support ByteAddressBuffer directly.
    HashSet<IRGlobalParam*> m_replacedByteAddressBufferParams;

    // Everything starts with a request to process a module,
    // which delegates to the central recursive walk of the IR.
    //
    void processModule(IRModule* module)
    {
        m_module = module;
        m_builder = IRBuilder(m_module);

        processInstRec(module->getModuleInst());

        // After processing, remove ByteAddressBuffer parameters that have been
        // replaced by StructuredBuffer parameters. These are no longer needed
        // and would cause errors in backends like SPIR-V that don't support
        // ByteAddressBuffer types directly.
        for (auto param : m_replacedByteAddressBufferParams)
        {
            // Only remove if it has no uses (its uses should have been replaced)
            if (!param->hasUses())
            {
                param->removeAndDeallocate();
            }
        }
    }

    // We recursively walk the entire IR structure (except
    // for decorations), and process any byte-address buffer
    // load or store operations.
    //
    void processInstRec(IRInst* inst)
    {
        switch (inst->getOp())
        {
        case kIROp_ByteAddressBufferLoad:
            processLoad(inst);
            break;

        case kIROp_ByteAddressBufferStore:
            processStore(inst);
            break;

        case kIROp_GetEquivalentStructuredBuffer:
            processGetEquivalentStructuredBuffer(inst);
            break;
        }

        for (IRInst* child : inst->getModifiableChildren())
        {
            processInstRec(child);
        }
    }

    void processGetEquivalentStructuredBuffer(IRInst* inst)
    {
        m_builder.setInsertBefore(inst);

        // We need to see what type it is to be interpreted as.
        auto type = inst->getDataType();

        // We want to determine the element type
        auto structuredBufferType = as<IRHLSLStructuredBufferTypeBase>(type);
        auto elementType = structuredBufferType->getElementType();

        // The buffer is operand 0
        auto buffer = inst->getOperand(0);
        // Get the equivalent structured buffer for the buffer.
        if (auto structuredBuffer = getEquivalentStructuredBuffer(elementType, buffer))
        {
            // We want to replace the the inst, with the equivalent structured buffer reference
            inst->replaceUsesWith(structuredBuffer);
            // Once replaced we don't need anymore
            inst->removeAndDeallocate();
        }
    }

    // The logic for both the load and store cases is similar,
    // so we will present the entire load case first and then
    // move on to stores.
    //
    void processLoad(IRInst* load)
    {
        // What we want to do with a load depends on the type
        // being loaded.
        //
        auto type = load->getDataType();

        // Validate the load's `location`/`alignment` contract up front so the diagnostic
        // fires regardless of whether the loaded type happens to need legalization.
        validateExplicitAlignment(load->getOperand(1), load->getOperand(2), type);

        // We start by looking at the type being loaded so
        // that we can opt out if it is legal.
        //
        if (isTypeLegalForByteAddressLoadStore(type))
            return;

        // If the type is one that requires legalization,
        // then we will set up to insert new IR instructions
        // to replace it.
        //
        m_builder.setInsertBefore(load);

        // We then emit a "legal load" with the same buffer
        // and byte offset from the original.
        //
        auto buffer = load->getOperand(0);
        auto offset = load->getOperand(1);
        auto alignment = load->getOperand(2);
        auto legalLoad = emitLegalLoad(type, buffer, offset, 0, alignment);

        // If it currently possible for the legalization
        // to fail (perhaps because of something else that
        // is invalid in the IR), so we will defensively
        // leave the code along in that case.
        //
        if (!legalLoad)
            return;

        // If we were able to generate a legal load operation,
        // then the value it yields can be used to fully
        // replace the previous illegal load.
        //
        load->replaceUsesWith(legalLoad);
        load->removeAndDeallocate();
    }

    bool isTypeLegalForByteAddressLoadStore(IRType* type)
    {
        // Whether or not a type is legal to use for
        // byte-address buffer load/store depends on
        // properties of the target, which will have
        // been passed into this pass via its options.
        //
        // If we are expected to translate all byte-address
        // operations to equivalent structured-buffer
        // operations, then that means *no* type is
        // legal for byte-address load/store.
        //
        if (m_options.translateToStructuredBufferOps)
            return false;

        // Basic types are usually legal to load/store
        // on all targets.
        //
        if (auto basicType = as<IRBasicType>(type))
        {
            // On targets that require translation to
            // make all load/store use `uint` values,
            // any scalar type that isn't `uint` is
            // illegal.
            //
            if (m_options.useBitCastFromUInt && basicType->getBaseType() != BaseType::UInt)
            {
                return false;
            }

            // For Metal targets, 64-bit integer types need special handling
            // because Metal doesn't support as_type casts from 64-bit to 32-bit.
            // These types should be lowered to two 32-bit operations.
            //
            if (m_options.lowerBasicTypeOps)
            {
                // getSameSizeUIntBaseType should convert any 64-bit types to UInt64
                auto unsignedBaseType = getSameSizeUIntBaseType(type->getOp());
                if (unsignedBaseType == BaseType::UInt64)
                {
                    return false; // Force legalization for 64-bit integer types
                }
            }

            // Otherwise, scalar types are assumed
            // legal for load/store.
            //
            return true;
        }

        // Vector types also depend on the options.
        //
        if (as<IRVectorType>(type))
        {
            // If we've been asked to scalarize all
            // vector load/store, then we need to
            // treat them as illegal.
            //
            if (m_options.scalarizeVectorLoadStore)
                return false;
        }

        // All other types are treated as always illegal,
        // so that we will legalize the load/store ops
        // in all cases.
        //
        // Note: recent builds of dxc (perhaps coupled with
        // recent shader models) support byte-address load/store
        // of more complex types, but it is simpler for Slang
        // to just legalize all the composite cases rather
        // than rely on a downstream compiler.
        //
        return false;
    }

    // Validate the explicit `alignment`/`location` contract of a `LoadAligned`/`StoreAligned` call
    // (issue #11545). Checks, short-circuiting so one call reports at most one diagnostic:
    //   - `alignment == 0`: the plain `Load`/`Store` "no promise" sentinel; nothing to validate.
    //   - `alignment` is a power of two, else E41301.
    //   - `alignment` is at least the access type's scalar-component alignment, else E41300.
    //   - a compile-time-constant `location` is a multiple of `alignment`, else E41303.
    // `alignment` is a `constexpr` parameter, so it is always an `IRIntLit` here.
    //
    // The diagnostic is the enforcement; a diagnosed access still lowers (it does not abort
    // legalization), so later errors can also surface.
    void validateExplicitAlignment(IRInst* baseOffset, IRInst* alignment, IRType* accessType)
    {
        // `alignment` is a `constexpr` parameter, so a non-constant promise is rejected at the
        // front end and the operand is always a literal by the time legalization runs.
        auto alignLit = as<IRIntLit>(alignment);
        SLANG_RELEASE_ASSERT(alignLit);

        auto alignmentVal = alignLit->getValue();
        if (alignmentVal == 0)
            return;

        if ((alignmentVal & (alignmentVal - 1)) != 0)
        {
            m_sink->diagnose(Diagnostics::ByteAddressBufferAlignmentNotPowerOfTwo{
                .alignment = alignmentVal,
                .location = baseOffset->sourceLoc,
            });
            return;
        }

        // The promise must be at least the access type's scalar-component alignment. Example:
        // `LoadAligned<float4>(0, 2)` promises 2, but a `float` component needs 4, so it is
        // rejected (E41300). A composite's natural alignment is the max of its members', so
        // checking the whole type covers every scalar leaf. Query *natural* (C) layout, not the
        // target's buffer layout, since the contract is defined against natural layout.
        IRSizeAndAlignment sizeAlignment;
        if (SLANG_SUCCEEDED(getNaturalSizeAndAlignment(m_target, accessType, &sizeAlignment)) &&
            sizeAlignment.alignment > 0 && (alignmentVal % sizeAlignment.alignment) != 0)
        {
            m_sink->diagnose(Diagnostics::ByteAddressBufferUnaligned{
                .alignment = alignmentVal,
                .requiredAlignment = sizeAlignment.alignment,
                .location = baseOffset->sourceLoc,
            });
            return;
        }

        // The alignment promise is itself valid; a compile-time-constant location must honor it.
        if (auto baseOffsetLit = as<IRIntLit>(baseOffset))
        {
            if ((baseOffsetLit->getValue() % alignmentVal) != 0)
            {
                m_sink->diagnose(Diagnostics::ByteAddressBufferLocationNotAligned{
                    .offset = baseOffsetLit->getValue(),
                    .alignment = alignmentVal,
                    .location = baseOffset->sourceLoc,
                });
            }
        }
    }

    // Returns true if a vectorized `accessSize`-byte access at `baseOffset + immediateOffset` is
    // known to be sufficiently aligned, deciding from the caller's alignment *promise* alone.
    //
    // `promisedAlignment` is the `alignment` argument of `LoadAligned`/`Store`, which is a
    // `constexpr` and therefore always an `IRIntLit`. A non-zero promise is the sole input to the
    // wide-vs-scalarized decision, so a constant `location` and a runtime `location` carrying the
    // same promise legalize identically (issue #11545's consistency goal). `promise == 0` is the
    // "no promise" sentinel of the plain `Load`/`Store` forms; those fall back to a
    // compile-time-constant base offset (a runtime no-promise access cannot be proven aligned and
    // stays scalarized). This predicate is pure: it never diagnoses (the alignment contract is
    // validated separately).
    bool isWideAccessAligned(
        IRInst* baseOffset,
        IRIntegerValue immediateOffset,
        IRInst* promisedAlignment,
        IRIntegerValue accessSize)
    {
        if (accessSize <= 0)
            return false;

        auto alignLit = as<IRIntLit>(promisedAlignment);
        SLANG_RELEASE_ASSERT(alignLit);
        auto promisedValue = alignLit->getValue();
        if (promisedValue != 0)
        {
            return (promisedValue % accessSize) == 0 && (immediateOffset % accessSize) == 0;
        }

        // No promise (plain `Load`/`Store`): trust a compile-time-constant base offset only.
        if (auto baseOffsetVal = as<IRIntLit>(baseOffset))
        {
            return ((baseOffsetVal->getValue() + immediateOffset) % accessSize) == 0;
        }
        return false;
    }

    IRType* getDescriptorHandleStorageType()
    {
        if (!m_options.translateToStructuredBufferOps)
            return nullptr;

        if (m_target->getTargetCaps().implies(CapabilityAtom::spvBindlessTextureNV))
            return m_builder.getUInt64Type();

        return m_builder.getVectorType(m_builder.getUIntType(), 2);
    }

    IROp getCastDescriptorHandleToStorageOp()
    {
        return m_target->getTargetCaps().implies(CapabilityAtom::spvBindlessTextureNV)
                   ? kIROp_CastDescriptorHandleToUInt64
                   : kIROp_CastDescriptorHandleToUInt2;
    }

    IROp getCastStorageToDescriptorHandleOp()
    {
        return m_target->getTargetCaps().implies(CapabilityAtom::spvBindlessTextureNV)
                   ? kIROp_CastUInt64ToDescriptorHandle
                   : kIROp_CastUInt2ToDescriptorHandle;
    }

    SlangResult getOffset(TargetProgram* target, IRStructField* field, IRIntegerValue* outOffset)
    {
        if (target->getHLSLToVulkanLayoutOptions() &&
            target->getHLSLToVulkanLayoutOptions()->shouldUseGLLayout())
        {
            return getStd430Offset(target->getTargetReq(), field, outOffset);
        }
        return getNaturalOffset(target->getTargetReq(), field, outOffset);
    }

    SlangResult getSizeAndAlignment(
        TargetProgram* target,
        IRType* type,
        IRSizeAndAlignment* outSizeAlignment)
    {
        if (target->getHLSLToVulkanLayoutOptions() &&
            target->getHLSLToVulkanLayoutOptions()->shouldUseGLLayout())
        {
            return getStd430SizeAndAlignment(target->getTargetReq(), type, outSizeAlignment);
        }
        return getNaturalSizeAndAlignment(target->getTargetReq(), type, outSizeAlignment);
    }

    // The core workhorse routine for the load case is `emitLegalLoad`,
    // which tries to emit load operations that read a value of the
    // given `type` from the given `buffer` at the required `baseOffset`
    // plus the `immediateOffset` if any.
    //
    IRInst* emitLegalLoad(
        IRType* type,
        IRInst* buffer,
        IRInst* baseOffset,
        IRIntegerValue immediateOffset,
        IRInst* alignment)
    {
        // The right way to load a value depends primarily
        // on the type, and secondarily on the options
        // that have been specified for this pass.
        //
        if (auto structType = as<IRStructType>(type))
        {
            // When loading a value of `struct` type, we will
            // load each field with its own operation.
            //
            // Note: A more "clever" implementation might try
            // to emit a minimal number of loads of whatever
            // is the largest supported type matching the
            // alignment of `structType`, and then break those
            // loaded values into fields with bit-level ops
            // once they are in registers.
            //
            // Such an approach could conceivably allow more
            // types to be loadable even on targets that
            // don't directly support them (e.g., a structure
            // with an `int` and two `int16_t` could be loadable
            // even when targetting DXBC).
            //
            // The flip side to such an approach would be that
            // it would complicate the generated code, and also
            // make the rules about when a type is supported
            // for byte-address load/store much more complicated.

            // We collect the loaded per-field values into an
            // array, which we will then use to construct the
            // full value of the `struct` type.
            //
            List<IRInst*> fieldVals;
            for (auto field : structType->getFields())
            {
                auto fieldType = field->getFieldType();

                // The relative offset of each field is calculated using
                // the IR-based layout subsystem, which works with the
                // "natural" in-memory layout of types.
                //
                // It is possible for layout computation to fail (e.g.,
                // if the field type somehow wasn't one that can be
                // laid out "naturally"). If the layout process fails,
                // then we fail to legalize this load.
                //
                IRIntegerValue fieldOffset = 0;
                SLANG_RETURN_NULL_ON_FAIL(getOffset(m_targetProgram, field, &fieldOffset));

                // Otherwise, we load the field by recursively calling this function
                // on the field type, with an adjusted immediate offset.
                //
                // If legalizing the field load fails, then we fail the load
                // of the structure as well. Any loads that were generated
                // for earlier fields will be left behind but can be eliminated
                // as dead code.
                //
                auto fieldVal = emitLegalLoad(
                    fieldType,
                    buffer,
                    baseOffset,
                    immediateOffset + fieldOffset,
                    alignment);
                if (!fieldVal)
                    return nullptr;

                fieldVals.add(fieldVal);
            }

            // Once all the field values have been loaded, we can bind
            // then together to make a single value of the `struct` type,
            // representing the result of the legalized load.
            //
            return m_builder.emitMakeStruct(type, fieldVals);
        }
        else if (auto arrayType = as<IRArrayTypeBase>(type))
        {
            // Loading a value of array type amounts to loading each
            // of its elements. There is shared logic between the
            // array, matrix, and vector cases, so we factor it into
            // a subroutine that we will explain later.
            //
            // We need a known constant number of elements in an array
            // to be able to emit per-element loads, so we skip
            // legalization if the array type isn't in the right form
            // for us to proceed.
            //

            if (auto elementCountInst = as<IRIntLit>(arrayType->getElementCount()))
            {
                // Emit an aligned load operation on an array when using a LoadAligned inst.
                // Else, fallback to scalarizing the loads.
                IRSizeAndAlignment elementLayout;
                SLANG_RELEASE_ASSERT(!getNaturalSizeAndAlignment(
                    m_target,
                    arrayType->getElementType(),
                    &elementLayout));
                IRIntegerValue elementStride = elementLayout.getStride();
                auto alignmentVal = elementStride * elementCountInst->getValue();
                if (!isWideAccessAligned(baseOffset, immediateOffset, alignment, alignmentVal))
                {
                    return emitLegalSequenceLoad(
                        type,
                        buffer,
                        baseOffset,
                        immediateOffset,
                        kIROp_MakeArray,
                        arrayType->getElementType(),
                        elementCountInst->getValue(),
                        alignment);
                }
                else
                {
                    return emitSimpleLoad(type, buffer, baseOffset, immediateOffset);
                }
            }
        }
        else if (auto descriptorHandleType = as<IRDescriptorHandleType>(type))
        {
            if (auto storageType = getDescriptorHandleStorageType())
            {
                auto storageVal =
                    emitLegalLoad(storageType, buffer, baseOffset, immediateOffset, alignment);
                if (!storageVal)
                    return nullptr;
                IRInst* args[] = {storageVal};
                return m_builder.emitIntrinsicInst(
                    descriptorHandleType,
                    getCastStorageToDescriptorHandleOp(),
                    1,
                    args);
            }
        }
        else if (auto matType = as<IRMatrixType>(type))
        {
            // Handling a matrix is largely like an array, with the
            // small detail that we need to construct the row type
            // that we expect to load for each "element."
            //
            if (getIntVal(matType->getLayout()) != SLANG_MATRIX_LAYOUT_COLUMN_MAJOR)
            {
                auto rowCountInst = as<IRIntLit>(matType->getRowCount());
                if (rowCountInst)
                {
                    auto rowType = m_builder.getVectorType(
                        matType->getElementType(),
                        matType->getColumnCount());
                    return emitLegalSequenceLoad(
                        type,
                        buffer,
                        baseOffset,
                        immediateOffset,
                        kIROp_MakeMatrix,
                        rowType,
                        rowCountInst->getValue(),
                        alignment);
                }
            }
            else
            {
                List<IRInst*> elements;
                auto colCount = (Index)getIntVal(matType->getColumnCount());
                auto rowCount = (Index)getIntVal(matType->getRowCount());
                auto colVectorType = m_builder.getVectorType(matType->getElementType(), rowCount);
                IRSizeAndAlignment colVectorSizeAlignment;
                getSizeAndAlignment(m_targetProgram, colVectorType, &colVectorSizeAlignment);
                for (Index c = 0; c < colCount; c++)
                {
                    auto colVector = emitLegalLoad(
                        colVectorType,
                        buffer,
                        baseOffset,
                        immediateOffset,
                        alignment);
                    for (Index r = 0; r < rowCount; r++)
                    {
                        elements.add(m_builder.emitElementExtract(colVector, (IRIntegerValue)r));
                    }
                    immediateOffset += colVectorSizeAlignment.getStride();
                }
                List<IRInst*> args;
                for (Index r = 0; r < rowCount; r++)
                {
                    for (Index c = 0; c < colCount; c++)
                    {
                        auto index = c * rowCount + r;
                        args.add(elements[index]);
                    }
                }
                return m_builder.emitMakeMatrix(matType, (UInt)args.getCount(), args.getBuffer());
            }
        }
        else if (auto vecType = as<IRVectorType>(type))
        {
            // One of the options that can vary per-target is whether to
            // scalarize vetor load/store operations. When that option
            // is turned on, we can treat a vector load just like an
            // array load.
            //
            if (auto elementCountInst = as<IRIntLit>(vecType->getElementCount()))
            {
                // Emit an aligned vector load operation when using a LoadAligned inst.
                // Else, fallback to scalarizing the loads.
                IRSizeAndAlignment elementLayout;
                SLANG_RELEASE_ASSERT(!getNaturalSizeAndAlignment(
                    m_target,
                    vecType->getElementType(),
                    &elementLayout));
                IRIntegerValue elementStride = elementLayout.getStride();
                auto alignmentVal = elementStride * elementCountInst->getValue();
                if (m_options.scalarizeVectorLoadStore ||
                    !isWideAccessAligned(baseOffset, immediateOffset, alignment, alignmentVal))
                {
                    return emitLegalSequenceLoad(
                        type,
                        buffer,
                        baseOffset,
                        immediateOffset,
                        kIROp_MakeVector,
                        vecType->getElementType(),
                        elementCountInst->getValue(),
                        alignment);
                }
                else
                {
                    return emitSimpleLoad(type, buffer, baseOffset, immediateOffset);
                }
            }

            // If we aren't scalarizing a vetor load then we next need
            // to consider the case where the target might only support
            // byte-address load/store of unsigned integer data (e.g.,
            // this is the case for D3D11/DXBC).
            //
            // We can still support loads of vectors with other element
            // types by first loading the data as unsigned integers, and
            // then bit-casting it to the correct type (e.g., load a
            // `uint4` with `Load4()` and then bit-cast to `float4` using
            // `asfloat()`).
            //
            if (m_options.useBitCastFromUInt)
            {
                // We will look at the element type of the vector (which must
                // be a basic type for this to work).
                //
                if (auto elementType = as<IRBasicType>(vecType->getElementType()))
                {
                    // If there is a distinct unsigned integer type of the
                    // same size as the element type, then we can use that
                    // for our load.
                    //
                    if (auto unsignedElementType = getSameSizeUIntType(elementType))
                    {
                        // We form the appropriate unsigned-integer vector type,
                        // and then emit a load for it.
                        //
                        auto unsignedVecType = m_builder.getVectorType(
                            unsignedElementType,
                            vecType->getElementCount());
                        auto unsignedVecVal =
                            emitSimpleLoad(unsignedVecType, buffer, baseOffset, immediateOffset);

                        // Once we have loaded the bits into a temporary,
                        // we can bit-cast it to the correct tyep and
                        // we have our result.
                        //
                        return m_builder.emitBitCast(vecType, unsignedVecVal);
                    }
                }
            }

            // Any cases of vectors not handled above are allowed to fall through
            // and be handled in the catch-all logic below.
        }
        else if (auto basicType = as<IRBasicType>(type))
        {
            // Only split misaligned 8-byte scalar loads for targets that translate
            // byte-address buffers into typed structured buffers. Targets that keep
            // byte-address-buffer operations use the simple/lowerBasicTypeOps paths.
            if (m_options.translateToStructuredBufferOps)
            {
                IRSizeAndAlignment sizeAlignment;
                SLANG_RETURN_NULL_ON_FAIL(
                    getNaturalSizeAndAlignment(m_target, type, &sizeAlignment));
                if (sizeAlignment.size == 8 && !isWideAccessAligned(
                                                   baseOffset,
                                                   immediateOffset,
                                                   alignment,
                                                   sizeAlignment.getStride()))
                {
                    return emitLegalUnaligned64BitLoadFromTwoUInts(
                        type,
                        buffer,
                        baseOffset,
                        immediateOffset,
                        alignment);
                }
            }

            // Most basic scalar types can be handled directly on targets,
            // but as described above for vectors, the D3D11/DXBC target
            // only support loading `uint` values, so we need to emulate
            // loads of other types (like `float`) by first loading a
            // `uint` and then bit-casting.
            //
            if (m_options.useBitCastFromUInt)
            {
                if (auto unsignedType = getSameSizeUIntType(basicType))
                {
                    auto unsignedVal =
                        emitSimpleLoad(unsignedType, buffer, baseOffset, immediateOffset);
                    return m_builder.emitBitCast(basicType, unsignedVal);
                }
            }
        }

        // If none of the many special cases above was triggered, then we
        // are in the base case and assume we want to emit a single load
        // for the type we were given.
        //
        return emitSimpleLoad(type, buffer, baseOffset, immediateOffset);
    }

    // Loading of sequences for arrays, matrices, and vectors is
    // bottlenecked through a single function.
    //
    IRInst* emitLegalSequenceLoad(
        IRType* type,
        IRInst* buffer,
        IRInst* baseOffset,
        IRIntegerValue immediateOffset,
        IROp op,
        IRType* elementType,
        IRIntegerValue elementCount,
        IRInst* alignment)
    {
        // Or goal here is to produce a value of the given `type`, loaded from `buffer`
        // at `baseOffset` plus `immediateOffset`.
        //
        // We will do this by emitting `elementCount` loads for the elements of
        // the given `elementType`, and then grouping them into the final sequence
        // using the given `op` (which will be something like `kIROp_MakeArray`).

        // To know how many bytes to step between loads, we must compute
        // the "stride" of the element type.
        //
        IRSizeAndAlignment elementLayout;
        SLANG_RETURN_NULL_ON_FAIL(
            getNaturalSizeAndAlignment(m_target, elementType, &elementLayout));
        IRIntegerValue elementStride = elementLayout.getStride();

        // We will collect all the element values into an array so
        // that we can construct the sequence when we are done.
        //
        List<IRInst*> elementVals;
        for (IRIntegerValue ii = 0; ii < elementCount; ++ii)
        {
            auto elementVal = emitLegalLoad(
                elementType,
                buffer,
                baseOffset,
                immediateOffset + ii * elementStride,
                alignment);
            if (!elementVal)
                return nullptr;

            elementVals.add(elementVal);
        }

        // Once we are done loading the elements we construct the sequence value.
        //
        return m_builder
            .emitIntrinsicInst(type, op, elementVals.getCount(), elementVals.getBuffer());
    }

    IRInst* emitLegalUnaligned64BitLoadFromTwoUInts(
        IRType* type,
        IRInst* buffer,
        IRInst* baseOffset,
        IRIntegerValue immediateOffset,
        IRInst* alignment)
    {
        auto uintType = m_builder.getUIntType();
        auto uint64Type = m_builder.getUInt64Type();

        auto loLoad = emitLegalLoad(uintType, buffer, baseOffset, immediateOffset, alignment);
        if (!loLoad)
            return nullptr;

        auto hiLoad = emitLegalLoad(uintType, buffer, baseOffset, immediateOffset + 4, alignment);
        if (!hiLoad)
            return nullptr;

        auto lo64 = m_builder.emitCast(uint64Type, loLoad);
        auto hi64 = m_builder.emitCast(uint64Type, hiLoad);
        auto shift = m_builder.emitShl(uint64Type, hi64, m_builder.getIntValue(uint64Type, 32));
        auto fullValue = m_builder.emitBitOr(uint64Type, lo64, shift);
        return m_builder.emitBitCast(type, fullValue);
    }

    // All of the loading operations above eventually bottom out at `emitSimpleLoad`,
    // which is meant to handle the base case where we do *not* want to
    // recurse on the structure of `type`.
    //
    IRInst* emitSimpleLoad(
        IRType* type,
        IRInst* buffer,
        IRInst* baseOffset,
        IRIntegerValue immediateOffset)
    {
        // For all of the operations above this in the call chain we have been
        // tracking a pair of a `baseOffset` as an IR instruction, and an
        // `immediateOffset` value. Keeping things split avoided introducing
        // a bunch of `add` instructions that could be constant-folded away.
        //
        // Instead, now that we are about to emit a load "for real"
        // we want to turn those two offset values into one.
        //
        IRInst* offset = emitOffsetAddIfNeeded(baseOffset, immediateOffset);

        // At this point there is one last (major) detail we need to
        // get into, which is that some targets (currently just GLSL)
        // do not actually have anything like byte-address buffers
        // as a built-in feature.
        //
        // Instead, GLSL has "shader storage buffers" which are
        // tied to a particular element type when declared. E.g.,:
        //
        //      buffer MyBuffer { uint _data[]; } myBuffer;
        //
        // The `myBuffer` declaration above can be used to load
        // `uint` values, but isn't much use if you want to load/store
        // a `half` or a `double` efficiently (and atomically,
        // where possible/guaranteed).
        //
        // Shader storage buffers like this are closer in spirit to
        // HLSL/Slang "structured buffers," so we think of this code
        // path as converting byte-address buffer operations into
        // structured-buffer operations.
        //
        // To make things work for GLSL output, we need to generate
        // multiple `buffer` declarations that all alias one another
        // (accomplished by giving them the same `binding`), but that
        // declare buffers with different element types.
        //
        if (m_options.translateToStructuredBufferOps)
        {
            // In order to emit a suitable structured-buffer load,
            // we need to find or create a global declaration for
            // a structured buffer that is "equivalent" to `buffer`,
            // but has `type` as its element type.
            //
            // That operation could conceivably fail, and when it
            // does we will fall back to the default handling of
            // emitting a byte-address buffer load (which will
            // then fail to generate valid GLSL code).
            //
            if (auto structuredBuffer = getEquivalentStructuredBuffer(type, buffer))
            {
                // The `offset` instruction represents the byte offset of
                // the thing we are trying to load, and we need to translate
                // that into an *index* for use with a structured buffer.
                //
                // We convert the offset to an index by dividing by the
                // stride of `type` as computed with our "natural layout" rules.
                //
                // This logic will be invalid if `offset` isn't a multiple of
                // the stride of `type`, but that case would have been
                // undefined behavior anyway.
                //
                auto offsetType = offset->getDataType();

                IRSizeAndAlignment typeLayout;
                SLANG_RETURN_NULL_ON_FAIL(getNaturalSizeAndAlignment(m_target, type, &typeLayout));
                auto typeStrideVal = typeLayout.getStride();

                auto typeStrideInst = m_builder.getIntValue(offsetType, typeStrideVal);
                IRInst* divArgs[] = {offset, typeStrideInst};
                auto index = m_builder.emitIntrinsicInst(offsetType, kIROp_Div, 2, divArgs);

                IRInst* args[] = {structuredBuffer, index};
                return m_builder.emitIntrinsicInst(type, kIROp_StructuredBufferLoad, 2, args);
            }
        }

        if (m_options.lowerBasicTypeOps)
        {
            // Some platforms e.g. Metal does not allow loading basic types that are not 4-byte
            // sized. We need to lower such loads.
            IRSizeAndAlignment sizeAlignment;
            SLANG_RETURN_NULL_ON_FAIL(getNaturalSizeAndAlignment(m_target, type, &sizeAlignment));
            if (sizeAlignment.size == 8)
            {
                // We need to load the value as two 4-byte values and then combine them.
                auto loOffset = offset;
                auto hiOffset = emitOffsetAddIfNeeded(offset, 4);
                IRInst* loadLoArgs[] = {buffer, loOffset};
                IRInst* loadHiArgs[] = {buffer, hiOffset};
                auto loLoad = m_builder.emitIntrinsicInst(
                    m_builder.getUIntType(),
                    kIROp_ByteAddressBufferLoad,
                    2,
                    loadLoArgs);
                auto hiLoad = m_builder.emitIntrinsicInst(
                    m_builder.getUIntType(),
                    kIROp_ByteAddressBufferLoad,
                    2,
                    loadHiArgs);
                auto lo64 = m_builder.emitCast(m_builder.getUInt64Type(), loLoad);
                auto hi64 = m_builder.emitCast(m_builder.getUInt64Type(), hiLoad);
                auto shift = m_builder.emitShl(
                    m_builder.getUInt64Type(),
                    hi64,
                    m_builder.getIntValue(m_builder.getUInt64Type(), 32));
                auto fullValue = m_builder.emitBitOr(m_builder.getUInt64Type(), lo64, shift);
                // For pointer types, Metal doesn't allow as_type casts from integers to pointers,
                // so we use proper cast operations instead of bit casts.
                if (type->getOp() == kIROp_PtrType || type->getOp() == kIROp_RawPointerType)
                {
                    // Use proper cast operation instead of bit cast for pointers
                    return m_builder.emitCastIntToPtr(type, fullValue);
                }
                else
                {
                    // For non-pointer 64-bit types (including IntPtr/UIntPtr)
                    return m_builder.emitBitCast(type, fullValue);
                }
            }
            else if (sizeAlignment.size < 4)
            {
                auto alignedOffset = m_builder.emitDiv(
                    offset->getDataType(),
                    offset,
                    m_builder.getIntValue(offset->getDataType(), 4));
                alignedOffset = m_builder.emitMul(
                    offset->getDataType(),
                    alignedOffset,
                    m_builder.getIntValue(offset->getDataType(), 4));
                IRInst* loadArgs[] = {buffer, alignedOffset};
                auto val = m_builder.emitIntrinsicInst(
                    m_builder.getUIntType(),
                    kIROp_ByteAddressBufferLoad,
                    2,
                    loadArgs);
                auto shiftAmount = m_builder.emitSub(offset->getDataType(), offset, alignedOffset);
                shiftAmount = m_builder.emitMul(
                    offset->getDataType(),
                    shiftAmount,
                    m_builder.getIntValue(offset->getDataType(), 8));
                IRInst* mask = nullptr;
                switch (sizeAlignment.size)
                {
                case 1:
                    mask = m_builder.getIntValue(m_builder.getUIntType(), 0xFF);
                    break;
                case 2:
                    mask = m_builder.getIntValue(m_builder.getUIntType(), 0xFFFF);
                    break;
                default:
                    SLANG_ASSERT(!"Unexpected size");
                    break;
                }
                auto shift = m_builder.emitShr(m_builder.getUIntType(), val, shiftAmount);
                auto masked = m_builder.emitBitAnd(m_builder.getUIntType(), shift, mask);
                IRInst* casted = nullptr;
                switch (sizeAlignment.size)
                {
                case 1:
                    casted = m_builder.emitCast(m_builder.getUInt8Type(), masked);
                    break;
                case 2:
                    casted = m_builder.emitCast(m_builder.getUInt16Type(), masked);
                    break;
                default:
                    SLANG_ASSERT(!"Unexpected size");
                    break;
                }
                return m_builder.emitBitCast(type, casted);
            }
        }
        // When we finally run out of special cases to handle, we just emit
        // a byte-address buffer load operation directly, assuming it will
        // work for the chosen target.
        //
        {
            IRInst* loadArgs[] = {buffer, offset};
            return m_builder.emitIntrinsicInst(type, kIROp_ByteAddressBufferLoad, 2, loadArgs);
        }
    }

    IRInst* emitOffsetAddIfNeeded(IRInst* baseOffset, IRIntegerValue immediateOffset)
    {
        // We need to create an instruction to represent
        // `baseOffset` plus `immediateOffset`.
        //
        // An important special case is when `immediateOffset` is zero:
        //
        if (immediateOffset == 0)
            return baseOffset;

        // Otherwise, we emit an `add` instruction of the appropriate type
        //
        auto type = baseOffset->getDataType();
        IRInst* args[] = {baseOffset, m_builder.getIntValue(type, immediateOffset)};
        return m_builder.emitIntrinsicInst(type, kIROp_Add, 2, args);
    }

    // At this point we have gone through the main logic of the load path,
    // and before we turn our attention to the store path we can go
    // ahead and define some of the utility functions that the code above
    // requires.

    // In order to handle interesting types on D3D11/DXBC, we need to
    // be able to map a base type to another type of the same size.
    //
    BaseType getSameSizeUIntBaseType(IROp op)
    {
        // For now we are only handling the 32-bit types here, because
        // the D3D11/DXBC target will not be able to handle 16- or
        // 64-bit types anyway. This could be improved over time
        // if needed.
        //
        switch (op)
        {
        case kIROp_IntType:
        case kIROp_FloatType:
        case kIROp_BoolType:
            // The basic 32-bit types (and `bool`) can be handled by
            // loading `uint` values and then bit-casting.
            //
            // Note: We aren't listing `kIROp_UIntType` here because
            // we don't want to introduce a bit-cast in the case where
            // the load was already for a `uint`.
            //
            return BaseType::UInt;
        case kIROp_Int8Type:
        case kIROp_UInt8Type:
            return BaseType::UInt8;
        case kIROp_Int16Type:
        case kIROp_UInt16Type:
        case kIROp_HalfType:
            return BaseType::UInt16;
        case kIROp_Int64Type:
        case kIROp_UInt64Type:
        case kIROp_DoubleType:
            return BaseType::UInt64;
        case kIROp_IntPtrType:
        case kIROp_UIntPtrType:
        case kIROp_RawPointerType:
        case kIROp_PtrType:
            if (isCPUTarget(m_target) && sizeof(void*) == 4)
                return BaseType::UInt;
            else
                return BaseType::UInt64;
        default:
            // All other types map to a sentinel value of `Void` to
            // indicate that a bit-cast solution shouldn't be attempted:
            // either load the original type, or fail.
            //
            return BaseType::Void;
        }
    }
    IRBasicType* getSameSizeUIntType(IRType* type)
    {
        auto unsignedBaseType = getSameSizeUIntBaseType(type->getOp());
        if (unsignedBaseType == BaseType::Void)
            return nullptr;

        return m_builder.getBasicType(unsignedBaseType);
    }

    // When replacing byte-address buffer load/store operations with
    // structured buffer ones, we need to be able to map an IR instruction
    // that represents a byte-address buffer to one that represents an
    // "equivalent" structured buffer.
    //
    // An important/tricky detail here is that the byte-address buffer
    // might have been passed in as a function parameter, or be indexed
    // from an array, etc.
    //
    // The logic here assumes this pass has run after a full legalization
    // pass on resource parameter usage, so that any references to
    // buffers in an instruction are "grounded" in a known global shader
    // parameter.

    IRInst* getEquivalentStructuredBuffer(IRType* elementType, IRInst* byteAddressBuffer)
    {
        if (this->m_options.treatGetEquivalentStructuredBufferAsGetThis)
            return byteAddressBuffer;

        if (!elementType)
        {
            return nullptr;
        }
        if (as<IRHLSLStructuredBufferTypeBase>(byteAddressBuffer->getDataType()))
            return byteAddressBuffer;
        // The simple case for replacement is when the byte-address buffer to
        // be replaced is a global shader parameter. That path will get its
        // own routine.
        if (auto byteAddressBufferParam = as<IRGlobalParam>(byteAddressBuffer))
        {
            return getEquivalentStructuredBufferParam(elementType, byteAddressBufferParam);
        }
        else if (auto castDynamicResource = as<IRCastDynamicResource>(byteAddressBuffer))
        {
            // If the underlying structured buffer is a CastDynamicResource,
            // we can simply cast the dynamic resource into the byte address buffer type instead.
            auto arg = castDynamicResource->getOperand(0);
            return m_builder.emitIntrinsicInst(
                getEquivalentStructuredBufferParamType(
                    elementType,
                    byteAddressBuffer->getDataType()),
                kIROp_CastDynamicResource,
                1,
                &arg);
        }
        else if (auto loadDescriptor = as<IRSPIRVLoadDescriptorFromHeap>(byteAddressBuffer))
        {
            auto structuredBufferType =
                getEquivalentStructuredBufferParamType(elementType, loadDescriptor->getDataType());
            if (!structuredBufferType)
                return nullptr;

            return m_builder.emitLoadDescriptorFromHeap(
                structuredBufferType,
                loadDescriptor->getHeap(),
                loadDescriptor->getIndex());
        }
        else if (auto resourceDescriptor = as<IRLoadResourceDescriptorFromHeap>(byteAddressBuffer))
        {
            auto structuredBufferType = getEquivalentStructuredBufferParamType(
                elementType,
                resourceDescriptor->getDataType());
            if (!structuredBufferType)
                return nullptr;

            auto index = resourceDescriptor->getIndex();
            return m_builder.emitIntrinsicInst(
                structuredBufferType,
                kIROp_LoadResourceDescriptorFromHeap,
                1,
                &index);
        }

        if (byteAddressBuffer->getOp() == kIROp_GetElement)
        {
            // If the code is fetching the byte-address buffer from an
            // array, then we need to create an "equivalent" structured
            // buffer array, and then index into that.
            //
            auto byteAddressBufferArray = byteAddressBuffer->getOperand(0);
            auto index = byteAddressBuffer->getOperand(1);

            auto structuredBufferArray =
                getEquivalentStructuredBuffer(elementType, byteAddressBufferArray);
            if (!structuredBufferArray)
                return nullptr;

            auto structuredBufferArrayType =
                as<IRArrayTypeBase>(structuredBufferArray->getDataType());
            if (!structuredBufferArrayType)
                return nullptr;

            // If we succeeded in creating a declaration for an array of
            // structured buffers to index into, we can now emit a new
            // operation to index into that array instead, and the result
            // will work as our "equivalent" structured buffer.
            //
            return m_builder.emitElementExtract(
                structuredBufferArrayType->getElementType(),
                structuredBufferArray,
                index);
        }

        if (auto babParam = as<IRParam>(byteAddressBuffer))
        {
            if (!m_options.translateToStructuredBufferOps)
                return nullptr;

            IRType* babType = babParam->getDataType();
            IRType* structuredBufferType =
                getEquivalentStructuredBufferParamType(elementType, babType);
            babType->replaceUsesWith(structuredBufferType);
            return babParam;
        }

        // If we failed to pattern-match the byte-address buffer operand
        // against something we can handle, then we need to bail out
        // of our attempt to legalize things here.
        //
        // TODO: Should we make this case an error?
        //
        return nullptr;
    }

    // Our seach for an "equivalent" structured buffer should bottom out when
    // we find a global shader parameter of byte-address buffer type, or an
    // array (of array of array of ...) byte-address buffer type.
    //
    // We then want to create an equivalent shader parameter of a matching
    // structured buffer (or array...) type.
    //
    // To avoid creating too many buffers (e.g., one per load), we will cache and
    // re-use the buffers we declare in this way. Note that we do *not* need
    // to worry if the deduplication is perfect, because we are already assuming
    // that the target will handle multiple buffers with the same `binding`
    // correctly.
    //
    Dictionary<KeyValuePair<IRInst*, IRInst*>, IRGlobalParam*> m_cachedStructuredBuffers;
    IRGlobalParam* getEquivalentStructuredBufferParam(
        IRType* elementType,
        IRGlobalParam* byteAddressBufferParam)
    {
        KeyValuePair<IRInst*, IRInst*> key(elementType, byteAddressBufferParam);

        IRGlobalParam* structuredBufferParam;
        if (!m_cachedStructuredBuffers.tryGetValue(key, structuredBufferParam))
        {
            structuredBufferParam =
                createEquivalentStructuredBufferParam(elementType, byteAddressBufferParam);
            m_cachedStructuredBuffers.add(key, structuredBufferParam);
        }
        return structuredBufferParam;
    }

    void cloneBufferDecorations(IRBuilder& builder, IRInst* dest, IRInst* src)
    {
        List<IRDecoration*> decorationsToRemove;
        for (auto decoration : src->getDecorations())
        {
            switch (decoration->getOp())
            {
            case kIROp_MemoryQualifierSetDecoration:
                builder.addMemoryQualifierSetDecoration(
                    dest,
                    as<IRMemoryQualifierSetDecoration>(decoration)->getMemoryQualifierBit());
                break;
            case kIROp_KeepAliveDecoration:
                // Transfer KeepAlive to the new buffer and mark for removal from original
                // This is needed when -preserve-params is used: the original ByteAddressBuffer
                // should not be kept alive since it's replaced by the StructuredBuffer.
                builder.addKeepAliveDecoration(dest);
                decorationsToRemove.add(decoration);
                break;
            case kIROp_ExportDecoration:
                // Transfer Export decoration to the new buffer and mark for removal from original
                {
                    auto exportDecor = as<IRExportDecoration>(decoration);
                    builder.addExportDecoration(dest, exportDecor->getMangledName());
                    decorationsToRemove.add(decoration);
                }
                break;
            default:
                break;
            }
        }
        // Remove decorations from the original that were transferred to the new buffer
        for (auto decoration : decorationsToRemove)
        {
            decoration->removeAndDeallocate();
        }
    }

    IRGlobalParam* createEquivalentStructuredBufferParam(
        IRType* elementType,
        IRGlobalParam* byteAddressBufferParam)
    {
        // When we need to create a new structured buffer to stand in for
        // some byte-address buffer (with a new `elementType` being used
        // for load/store), we need to figure out the "equivalent" type
        // to use for the new buffer.
        //
        auto byteAddressBufferParamType = byteAddressBufferParam->getDataType();
        auto structuredBufferParamType =
            getEquivalentStructuredBufferParamType(elementType, byteAddressBufferParamType);
        if (!structuredBufferParamType)
            return nullptr;

        // Next we will create a global shader parameter using the new
        // type.
        //
        // Note: we are creating a new `IRBuilder` here rather than using
        // `m_builder` because this logic could get called during the middle
        // of legalizing a load or store, and we don't want to mess with
        // the insertion location of `m_builder`.
        //
        IRBuilder paramBuilder(m_module);
        paramBuilder.setInsertBefore(byteAddressBufferParam);

        auto structuredBufferParam = paramBuilder.createGlobalParam(structuredBufferParamType);
        if (auto nameHint = byteAddressBufferParam->findDecoration<IRNameHintDecoration>())
            paramBuilder.addNameHintDecoration(structuredBufferParam, nameHint->getName());
        if (auto typeHint = byteAddressBufferParam->findDecoration<IRUserTypeNameDecoration>())
            paramBuilder.addUserTypeNameDecoration(
                structuredBufferParam,
                typeHint->getUserTypeName());

        // The new parameter needs to be given a layout to match the existing
        // parameter, so that it is given the same `binding` in the generated code.
        //
        if (auto layoutDecoration = byteAddressBufferParam->findDecoration<IRLayoutDecoration>())
        {
            paramBuilder.addLayoutDecoration(structuredBufferParam, layoutDecoration->getLayout());
        }
        cloneBufferDecorations(paramBuilder, structuredBufferParam, byteAddressBufferParam);

        // Track that this ByteAddressBuffer parameter has been replaced.
        // It will be removed at the end of the pass if it has no remaining uses.
        m_replacedByteAddressBufferParams.add(byteAddressBufferParam);

        return structuredBufferParam;
    }

    IRType* getEquivalentStructuredBufferParamType(
        IRType* elementType,
        IRType* byteAddressBufferType)
    {
        if (as<IRHLSLStructuredBufferTypeBase>(byteAddressBufferType))
            return byteAddressBufferType;

        // Our task in this function is to compute the type for
        // a structure buffer that is equivalent to `byteAddressBufferType`,
        // but with the given `elementType`.
        IRInst* typeOperands[2] = {elementType, m_builder.getType(kIROp_DefaultBufferLayoutType)};
        switch (byteAddressBufferType->getOp())
        {
            // The basic `*ByteAddressBuffer` types map directly to the
            // `*StructuredBuffer<elementType>` cases.
        case kIROp_HLSLByteAddressBufferType:
            return m_builder.getType(kIROp_HLSLStructuredBufferType, 2, typeOperands);
        case kIROp_HLSLRWByteAddressBufferType:
            return m_builder.getType(kIROp_HLSLRWStructuredBufferType, 2, typeOperands);
        case kIROp_HLSLRasterizerOrderedByteAddressBufferType:
            return m_builder.getType(
                kIROp_HLSLRasterizerOrderedStructuredBufferType,
                2,
                typeOperands);

        case kIROp_ArrayType:
        case kIROp_UnsizedArrayType:
            {
                // Array types (both sized and unsized) need to translate
                // their element type to an equivalent structured buffer
                // and build a new array type with the same element count.
                //
                auto arrayType = cast<IRArrayTypeBase>(byteAddressBufferType);
                return m_builder.getArrayTypeBase(
                    byteAddressBufferType->getOp(),
                    getEquivalentStructuredBufferParamType(
                        elementType,
                        arrayType->getElementType()),
                    arrayType->getElementCount());
            }

        default:
            return nullptr;
        }
    }

    // At this point we've covered all the logic for the load case down
    // to the last detail.
    //
    // All that remains is to go over the equivalent logic for the case
    // of byte-address buffer stores, which mostly parallels code we've
    // already discussed.

    void processStore(IRInst* store)
    {
        // Just as for loads, the logic for stores is base on the type
        // being used, but unlike in the load case we don't care about
        // the type of the store operation, but instead the operand
        // that represents the value to be stored.
        //
        auto value = store->getOperand(3);
        auto type = value->getDataType();

        // Validate the store's `location`/`alignment` contract up front so the diagnostic
        // fires regardless of whether the stored type happens to need legalization.
        validateExplicitAlignment(store->getOperand(1), store->getOperand(2), type);

        // Types that are already legal to use don't require any processing.
        //
        if (isTypeLegalForByteAddressLoadStore(type))
            return;

        // Otherwise we set up to try and emit a replacement.
        //
        m_builder.setInsertBefore(store);

        // It is possible that our attempt to emit a replacement will fail
        // (this should only happen if we run into types that shouldn't
        // actually be allowed on a target), and in those cases we will
        // leave the original store around as well (this is at worst a
        // performance issue, but we should still consider trying to
        // tighten this up and make all uhandled cases be hard errors).
        //
        auto result = emitLegalStore(
            type,
            store->getOperand(0),
            store->getOperand(1),
            0,
            store->getOperand(2),
            value);
        if (SLANG_FAILED(result))
            return;

        store->removeAndDeallocate();
    }

    Result emitLegalStore(
        IRType* type,
        IRInst* buffer,
        IRInst* baseOffset,
        IRIntegerValue immediateOffset,
        IRInst* alignment,
        IRInst* value)
    {
        // The flow for emitting a legal store is very similar to that for
        // legal loads; we will recurse on the structure of `type` and
        // emit stores for fields/elements as needed.

        if (auto structType = as<IRStructType>(type))
        {
            // To store a structure, we store each of its fields at
            // the appropriate relative offset.
            //
            for (auto field : structType->getFields())
            {
                auto fieldType = field->getFieldType();

                IRIntegerValue fieldOffset;
                SLANG_RETURN_ON_FAIL(getOffset(m_targetProgram, field, &fieldOffset));

                auto fieldVal = m_builder.emitFieldExtract(fieldType, value, field->getKey());
                SLANG_RETURN_ON_FAIL(emitLegalStore(
                    fieldType,
                    buffer,
                    baseOffset,
                    immediateOffset + fieldOffset,
                    alignment,
                    fieldVal));
            }
            return SLANG_OK;
        }
        else if (auto arrayType = as<IRArrayTypeBase>(type))
        {
            // Arrays and other sequences bottleneck through a helper
            // function, which we will cover later.
            //
            if (auto elementCountInst = as<IRIntLit>(arrayType->getElementCount()))
            {
                // Emit an aligned store operation on an array when using a StoreAligned inst.
                // Else, fallback to scalarizing the stores.
                IRSizeAndAlignment elementLayout;
                SLANG_RELEASE_ASSERT(!getNaturalSizeAndAlignment(
                    m_target,
                    arrayType->getElementType(),
                    &elementLayout));
                IRIntegerValue elementStride = elementLayout.getStride();
                auto alignmentVal = elementStride * elementCountInst->getValue();
                if (!isWideAccessAligned(baseOffset, immediateOffset, alignment, alignmentVal))
                {
                    return emitLegalSequenceStore(
                        buffer,
                        baseOffset,
                        immediateOffset,
                        value,
                        arrayType->getElementType(),
                        elementCountInst->getValue(),
                        alignment);
                }
                else
                {
                    return emitSimpleStore(
                        value->getDataType(),
                        buffer,
                        baseOffset,
                        immediateOffset,
                        value);
                }
            }
        }
        else if (as<IRDescriptorHandleType>(type))
        {
            if (auto storageType = getDescriptorHandleStorageType())
            {
                IRInst* args[] = {value};
                auto storageVal = m_builder.emitIntrinsicInst(
                    storageType,
                    getCastDescriptorHandleToStorageOp(),
                    1,
                    args);
                return emitLegalStore(
                    storageType,
                    buffer,
                    baseOffset,
                    immediateOffset,
                    alignment,
                    storageVal);
            }
        }
        else if (auto matType = as<IRMatrixType>(type))
        {
            auto layout = getIntVal(matType->getLayout());
            if (layout != SLANG_MATRIX_LAYOUT_COLUMN_MAJOR)
            {
                auto rowCountInst = as<IRIntLit>(matType->getRowCount());
                if (rowCountInst)
                {
                    auto rowType = m_builder.getVectorType(
                        matType->getElementType(),
                        matType->getColumnCount());
                    return emitLegalSequenceStore(
                        buffer,
                        baseOffset,
                        immediateOffset,
                        value,
                        rowType,
                        rowCountInst->getValue(),
                        alignment);
                }
            }
            else
            {
                auto colCount = (Index)getIntVal(matType->getColumnCount());
                auto rowCount = (Index)getIntVal(matType->getRowCount());
                List<IRInst*> srcRows;
                for (Index r = 0; r < rowCount; r++)
                    srcRows.add(m_builder.emitElementExtract(value, (IRIntegerValue)r));
                for (Index c = 0; c < colCount; c++)
                {
                    List<IRInst*> colVectorArgs;
                    for (Index r = 0; r < rowCount; r++)
                    {
                        auto rowVector = srcRows[r];
                        auto element = m_builder.emitElementExtract(rowVector, (IRIntegerValue)c);
                        colVectorArgs.add(element);
                    }
                    auto colVectorType =
                        m_builder.getVectorType(matType->getElementType(), rowCount);
                    auto colVector = m_builder.emitMakeVector(colVectorType, colVectorArgs);
                    IRSizeAndAlignment colVectorSizeAlignment;
                    getSizeAndAlignment(m_targetProgram, colVectorType, &colVectorSizeAlignment);
                    emitLegalStore(
                        colVectorType,
                        buffer,
                        baseOffset,
                        immediateOffset,
                        alignment,
                        colVector);
                    immediateOffset += colVectorSizeAlignment.getStride();
                }
                return SLANG_OK;
            }
        }
        else if (auto vecType = as<IRVectorType>(type))
        {
            if (auto elementCountInst = as<IRIntLit>(vecType->getElementCount()))
            {
                // Emit an aligned vector store operation when using a StoreAligned inst.
                // Else, fallback to scalarizing the stores.

                IRSizeAndAlignment elementLayout;
                SLANG_RELEASE_ASSERT(!getNaturalSizeAndAlignment(
                    m_target,
                    vecType->getElementType(),
                    &elementLayout));
                IRIntegerValue elementStride = elementLayout.getStride();
                auto alignmentVal = elementStride * elementCountInst->getValue();
                if (m_options.scalarizeVectorLoadStore ||
                    !isWideAccessAligned(baseOffset, immediateOffset, alignment, alignmentVal))
                {
                    return emitLegalSequenceStore(
                        buffer,
                        baseOffset,
                        immediateOffset,
                        value,
                        vecType->getElementType(),
                        elementCountInst->getValue(),
                        alignment);
                }
                else
                {
                    return emitSimpleStore(
                        value->getDataType(),
                        buffer,
                        baseOffset,
                        immediateOffset,
                        value);
                }
            }

            if (m_options.useBitCastFromUInt)
            {
                auto elementType = as<IRBasicType>(vecType->getElementType());
                if (auto unsignedElementType = getSameSizeUIntType(elementType))
                {
                    // The bit-cast case for stores is similar to the case
                    // for loads, except that we cast the value before
                    // storing it (instead of casting a value after loading).
                    //
                    auto unsignedVecType =
                        m_builder.getVectorType(unsignedElementType, vecType->getElementCount());
                    auto unsignedVecVal = m_builder.emitBitCast(unsignedVecType, value);
                    return emitSimpleStore(
                        unsignedVecType,
                        buffer,
                        baseOffset,
                        immediateOffset,
                        unsignedVecVal);
                }
            }
        }
        else if (auto basicType = as<IRBasicType>(type))
        {
            // Match the load path: this split is only needed when byte-address buffers
            // are translated into typed structured buffers.
            if (m_options.translateToStructuredBufferOps)
            {
                IRSizeAndAlignment sizeAlignment;
                SLANG_RETURN_ON_FAIL(getNaturalSizeAndAlignment(m_target, type, &sizeAlignment));
                if (sizeAlignment.size == 8 && !isWideAccessAligned(
                                                   baseOffset,
                                                   immediateOffset,
                                                   alignment,
                                                   sizeAlignment.getStride()))
                {
                    return emitLegalUnaligned64BitStoreAsTwoUInts(
                        buffer,
                        baseOffset,
                        immediateOffset,
                        alignment,
                        value);
                }
            }

            if (m_options.useBitCastFromUInt)
            {
                if (auto unsignedType = getSameSizeUIntType(basicType))
                {
                    auto unsignedVal = m_builder.emitBitCast(unsignedType, value);
                    return emitSimpleStore(
                        unsignedType,
                        buffer,
                        baseOffset,
                        immediateOffset,
                        unsignedVal);
                }
            }
        }

        return emitSimpleStore(type, buffer, baseOffset, immediateOffset, value);
    }

    Result emitLegalUnaligned64BitStoreAsTwoUInts(
        IRInst* buffer,
        IRInst* baseOffset,
        IRIntegerValue immediateOffset,
        IRInst* alignment,
        IRInst* value)
    {
        auto uintType = m_builder.getUIntType();
        auto uint64Type = m_builder.getUInt64Type();

        auto uint64Val = m_builder.emitBitCast(uint64Type, value);
        auto loVal = m_builder.emitCast(uintType, uint64Val);
        auto hiVal = m_builder.emitCast(
            uintType,
            m_builder.emitShr(uint64Type, uint64Val, m_builder.getIntValue(uint64Type, 32)));

        SLANG_RETURN_ON_FAIL(
            emitLegalStore(uintType, buffer, baseOffset, immediateOffset, alignment, loVal));
        SLANG_RETURN_ON_FAIL(
            emitLegalStore(uintType, buffer, baseOffset, immediateOffset + 4, alignment, hiVal));
        return SLANG_OK;
    }

    Result emitSimpleStore(
        IRType* type,
        IRInst* buffer,
        IRInst* baseOffset,
        IRIntegerValue immediateOffset,
        IRInst* value)
    {
        IRInst* offset = emitOffsetAddIfNeeded(baseOffset, immediateOffset);
        if (m_options.translateToStructuredBufferOps)
        {
            if (auto structuredBuffer = getEquivalentStructuredBuffer(type, buffer))
            {
                // Similar to the load case, if we are replacing byte-address
                // buffers with structured buffers, then once we find the
                // "equivalent" buffer to use, we emit a structured-buffer store,
                // with an index computed by dividing the offset by the stride.
                //
                auto indexType = offset->getDataType();

                IRSizeAndAlignment typeLayout;
                SLANG_RETURN_ON_FAIL(getNaturalSizeAndAlignment(m_target, type, &typeLayout));

                auto typeStride = m_builder.getIntValue(indexType, typeLayout.getStride());

                IRInst* divArgs[] = {offset, typeStride};
                auto index = m_builder.emitIntrinsicInst(indexType, kIROp_Div, 2, divArgs);

                IRInst* args[] = {structuredBuffer, index, value};
                m_builder.emitIntrinsicInst(type, kIROp_RWStructuredBufferStore, 3, args);
                return SLANG_OK;
            }
        }
        if (m_options.lowerBasicTypeOps)
        {
            // Some platforms e.g. Metal does not allow storing basic types that are not 4-byte
            // sized. We need to lower such stores.
            IRSizeAndAlignment sizeAlignment;
            SLANG_RETURN_ON_FAIL(getNaturalSizeAndAlignment(m_target, type, &sizeAlignment));
            if (sizeAlignment.size == 8)
            {
                // We need to store the value as two 4-byte values.
                // For pointer types, Metal doesn't allow as_type casts from pointers to integers,
                // so we use proper cast operations instead of bit casts.
                IRInst* loVal;
                IRInst* hiVal;
                IRInst* uint64Val;
                if (type->getOp() == kIROp_PtrType || type->getOp() == kIROp_RawPointerType)
                {
                    // Use proper cast operation instead of bit cast for pointers
                    uint64Val = m_builder.emitCastPtrToInt(value);
                }
                else
                {
                    // For non-pointer 64-bit types (including IntPtr/UIntPtr)
                    uint64Val = m_builder.emitBitCast(m_builder.getUInt64Type(), value);
                }
                loVal = m_builder.emitCast(m_builder.getUIntType(), uint64Val);
                hiVal = m_builder.emitCast(
                    m_builder.getUIntType(),
                    m_builder.emitShr(
                        m_builder.getUInt64Type(),
                        uint64Val,
                        m_builder.getIntValue(m_builder.getUInt64Type(), 32)));
                auto loOffset = offset;
                auto hiOffset = emitOffsetAddIfNeeded(offset, 4);
                IRInst* alignment = m_builder.getIntValue(m_builder.getUIntType(), 0);
                m_builder.emitByteAddressBufferStore(buffer, loOffset, alignment, loVal);
                m_builder.emitByteAddressBufferStore(buffer, hiOffset, alignment, hiVal);
                return SLANG_OK;
            }
            else if (sizeAlignment.size < 4)
            {
                IRInst* loadArgs[] = {buffer, offset};
                auto existingVal = m_builder.emitIntrinsicInst(
                    m_builder.getUIntType(),
                    kIROp_ByteAddressBufferLoad,
                    2,
                    loadArgs);
                auto alignedOffset = m_builder.emitDiv(
                    offset->getDataType(),
                    offset,
                    m_builder.getIntValue(offset->getDataType(), 4));
                alignedOffset = m_builder.emitMul(
                    offset->getDataType(),
                    alignedOffset,
                    m_builder.getIntValue(offset->getDataType(), 4));
                auto shiftAmount = m_builder.emitSub(offset->getDataType(), offset, alignedOffset);
                shiftAmount = m_builder.emitMul(
                    offset->getDataType(),
                    shiftAmount,
                    m_builder.getIntValue(offset->getDataType(), 8));
                auto uintVal = m_builder.emitCast(
                    m_builder.getUIntType(),
                    m_builder.emitBitCast(getSameSizeUIntType(value->getDataType()), value));
                auto shiftedData = m_builder.emitShl(m_builder.getUIntType(), uintVal, shiftAmount);
                IRInst* mask = nullptr;
                switch (sizeAlignment.size)
                {
                case 1:
                    mask = m_builder.getIntValue(m_builder.getUIntType(), 0xFF);
                    break;
                case 2:
                    mask = m_builder.getIntValue(m_builder.getUIntType(), 0xFFFF);
                    break;
                default:
                    SLANG_ASSERT(!"Unexpected size");
                    return SLANG_FAIL;
                }
                mask = m_builder.emitShl(m_builder.getUIntType(), mask, shiftAmount);
                mask = m_builder.emitBitNot(m_builder.getUIntType(), mask);
                auto maskedData = m_builder.emitBitAnd(m_builder.getUIntType(), existingVal, mask);
                auto newData =
                    m_builder.emitBitOr(m_builder.getUIntType(), maskedData, shiftedData);
                m_builder.emitByteAddressBufferStore(buffer, alignedOffset, newData);
                return SLANG_OK;
            }
        }
        {
            m_builder.emitByteAddressBufferStore(buffer, offset, value);
            return SLANG_OK;
        }
    }

    Result emitLegalSequenceStore(
        IRInst* buffer,
        IRInst* baseOffset,
        IRIntegerValue immediateOffset,
        IRInst* value,
        IRType* elementType,
        IRIntegerValue elementCount,
        IRInst* alignment)
    {
        // The store case for sequences is similar to the load case.
        //
        // We iterate over the elements and fetch then store each one.
        //
        IRSizeAndAlignment elementLayout;
        SLANG_RETURN_ON_FAIL(getNaturalSizeAndAlignment(m_target, elementType, &elementLayout));
        IRIntegerValue elementStride = elementLayout.getStride();

        auto indexType = m_builder.getIntType();
        for (IRIntegerValue ii = 0; ii < elementCount; ++ii)
        {
            auto elementIndex = m_builder.getIntValue(indexType, ii);
            auto elementVal = m_builder.emitElementExtract(elementType, value, elementIndex);
            SLANG_RETURN_ON_FAIL(emitLegalStore(
                elementType,
                buffer,
                baseOffset,
                immediateOffset + ii * elementStride,
                alignment,
                elementVal));
        }

        return SLANG_OK;
    }
};


void legalizeByteAddressBufferOps(
    IRModule* module,
    Session* session,
    TargetProgram* program,
    DiagnosticSink* sink,
    ByteAddressBufferLegalizationOptions const& options)
{
    ByteAddressBufferLegalizationContext context;
    context.m_session = session;
    context.m_target = program->getTargetReq();
    context.m_options = options;
    context.m_targetProgram = program;
    context.m_sink = sink;
    context.processModule(module);
}


} // namespace Slang
