Repository: shader-slang/slang
File: source/slang/slang-ir-transform-params-to-constref.cpp (C++)

Code before:

            // Skip entry-point and pseudo-entry-point functions
            // since we cannot legalize the input parameters.
            if (as<IREntryPointDecoration>(decoration) || as<IRCudaKernelDecoration>(decoration) ||
                as<IRAutoPyBindCudaDecoration>(decoration))
                return false;
        }

        // Skip functions with `kIROp_GenericAsm` since
        // these instructions inject target specific code
        // using parameters in an unpredictable way, relying
        // on assumptions that parameters do not change type.
        for (auto block : func->getBlocks())
        {
            for (auto inst : block->getChildren())
            {
                if (!as<IRGenericAsm>(inst))
                    continue;
                return false;
            }
        }

        return true;
    }

    // Eliminate unnecessary load+store pairs that can be optimized away
    // Only optimize patterns related to the updated parameters
    void eliminateLoadStorePairs(IRFunc* func, HashSet<IRParam*>& updatedParams)
    {
        List<IRInst*> toRemove;

        // Look for patterns: load(ptr) followed by store(var, load_result)
        // This can be optimized to direct pointer usage
        for (auto block : func->getBlocks())
        {
            for (auto inst : block->getChildren())
            {
                if (auto storeInst = as<IRStore>(inst))
                {
                    auto storedValue = storeInst->getVal();
                    auto destPtr = storeInst->getPtr();

                    // Check if we're storing a load result
                    if (auto loadInst = as<IRLoad>(storedValue))
                    {
                        auto loadedPtr = loadInst->getPtr();

                        // Only optimize if the loaded pointer is related to our updated parameters
                        bool isRelatedToUpdatedParam = false;
                        for (auto param : updatedParams)
                        {

Review comment (suggestion):
I don’t think the conditions being checked here is comprehensive.

It shouldn’t matter whether or not the load is used once. However it matters if the source address is immutable, and the callee isn’t modifying var.

Code after:

            // Skip entry-point and pseudo-entry-point functions
            // since we cannot legalize the input parameters.
            if (as<IREntryPointDecoration>(decoration) || as<IRCudaKernelDecoration>(decoration) ||
                as<IRAutoPyBindCudaDecoration>(decoration))
                return false;
        }

        // Skip functions with `kIROp_GenericAsm` since
        // these instructions inject target specific code
        // using parameters in an unpredictable way, relying
        // on assumptions that parameters do not change type.
        for (auto block : func->getBlocks())
        {
            for (auto inst : block->getChildren())
            {
                if (!as<IRGenericAsm>(inst))
                    continue;
                return false;
            }
        }

        return true;
    }

    // Process a single function
    void processFunc(IRFunc* func)
    {
        HashSet<IRParam*> updatedParams;

        // First pass: Transform parameter types
        for (auto param = func->getFirstParam(); param; param = param->getNextParam())
        {
            if (shouldTransformParam(param))
            {
                // Our goal here is to transform `in T` parameters to const-ref.
                // We are selective about what we will transform for a few reasons:
                // 1. no reason to transform simple primitives like `int`.
                // 2. not every type makes sense as constref. For example, `ParameterBlock`.
                // 3. constref is not 100% stable, so we need to be selective on what we let
                //    transform into constref.
                //
                // This allows us to pass the address of variables directly into a function,
                // giving us the choice to remove copies into a parameter.
                auto paramType = param->getDataType();
                auto constRefType = builder.getConstRefType(paramType, AddressSpace::ThreadLocal);
                param->setFullType(constRefType);

                changed = true;
                updatedParams.add(param);
            }