Metal kernel reads zeros from constant globals with dynamic indexing

[haakon-e]

Metal kernels silently read 0.0 from LLVM constant globals that contain aggregate data (tuples, SMatrix, arrays) when accessed with dynamic indices. Scalar constants and statically-indexed aggregates work correctly.

Reproducer

using Metal

function kernel_broken(out)
    i = Metal.thread_position_in_grid_1d()
    t = (1.0f0, 2.0f0, 3.0f0, 4.0f0)
    @inbounds out[i] = t[i]
    return nothing
end

out = MtlArray(ones(Float32, 4))
@metal threads=4 kernel_broken(out)
Array(out)  # [0.0, 0.0, 0.0, 0.0]  <-- should be [1.0, 2.0, 3.0, 4.0]

Replacing the dynamic index t[i] with a static index t[2] produces correct results. Passing the tuple as a kernel argument also works.

Environment

• Metal.jl: v1.9.3
• GPUCompiler: v1.9.1
• Julia 1.12.5 (LLVM 18.1.7)
• macOS (Apple M2)

I have a detailed IR-level analysis of the root cause and a working local fix in GPUCompiler. Happy to share if useful. However, if a fix is more appropriate elsewhere, please let me know.

[maleadt]

I have a detailed IR-level analysis of the root cause and a working local fix in GPUCompiler.

If it's a bug in GPUCompiler.jl, please share.

[maleadt]

I took a quick look as well, and the problem is with the add_global_address_spaces! pass only switching outermost global accesses to address space 2. With a dynamic index, we get a GEP, cast, and load back from the old address space 0.

We could augment add_global_address_spaces! to rewrite recursively, but the real solution is relying on the existing InferAddressSpaces pass in LLVM. It's currently a no-op since we lack a TTI for Metal. I'm exploring ways to feed the necessary info to LLVM (either with a pipeline parsing callback, or a custom TTI interface).

[haakon-e]

As far as I can tell, add_global_address_spaces! in src/metal.jl correctly moves constant globals into address space 2 (constant memory). However, rather than updating all uses to work with AS 2 pointers directly, it wraps each global in a
const_addrspacecast back to AS 0 (generic). As I understand it, this is done to avoid having to recursively update pointer types through the clone function. The result is that GEP and load instructions still operate on AS 0 pointers.

[haakon-e]

My local workaround follows. It has some clear weaknesses, including:

• It's a second pass cleaning up after the first (add_global_address_spaces!)
• It reimplements a subset of what InferAddressSpacesPass does generically, by pattern-matching specific instruction sequences
• It only handles GEP+load. bitcast, ptrtoint, call with pointer arg are not addressed

code

written with assistance of AI

# src/metal.jl L164
# add kernel metadata
    if job.config.kernel
        entry = add_parameter_address_spaces!(job, mod, entry)
        entry = add_global_address_spaces!(job, mod, entry)
        rewrite_constant_gep_loads_to_as2!(mod)  # <-- new method
        # [...]

# function definition
function rewrite_constant_gep_loads_to_as2!(mod::LLVM.Module)
    # collect AS 2 constant globals
    as2_globals = Set{LLVM.Value}()
    for gv in globals(mod)
        isconstant(gv) || continue
        addrspace(value_type(gv)) == 2 || continue
        push!(as2_globals, gv)
    end
    isempty(as2_globals) && return

    # helper: strip addrspacecast ConstantExpr to find the AS 2 source
    function strip_addrspacecast(val)
        if val isa LLVM.ConstantExpr && opcode(val) == LLVM.API.LLVMAddrSpaceCast
            return operands(val)[1]
        end
        return nothing
    end

    # helper: get GEP source element type via C API
    gep_source_element_type(inst) =
        LLVM.LLVMType(LLVM.API.LLVMGetGEPSourceElementType(inst))

    # helper: check if GEP is inbounds via C API
    gep_is_inbounds(inst) = LLVM.API.LLVMIsInBounds(inst) != 0

    for fun in functions(mod)
        isempty(blocks(fun)) && continue

        # pass 1: find GEP instructions whose base is addrspacecast(AS2 → AS0)
        worklist = Tuple{LLVM.Instruction, LLVM.Value}[]
        for bb in blocks(fun), inst in instructions(bb)
            opcode(inst) == LLVM.API.LLVMGetElementPtr || continue
            base = operands(inst)[1]
            src = strip_addrspacecast(base)
            if src !== nothing && src in as2_globals
                push!(worklist, (inst, src))
            end
        end
        isempty(worklist) && continue

        # pass 2: rebuild root GEPs in AS 2
        replacements = Dict{LLVM.Value, LLVM.Value}()
        @dispose builder=IRBuilder() begin
            for (gep_inst, as2_base) in worklist
                position!(builder, gep_inst)
                ops = operands(gep_inst)
                indices = LLVM.Value[ops[i] for i in 2:length(ops)]
                elem_ty = gep_source_element_type(gep_inst)
                new_gep = if gep_is_inbounds(gep_inst)
                    inbounds_gep!(builder, elem_ty, as2_base, indices)
                else
                    gep!(builder, elem_ty, as2_base, indices)
                end
                replacements[gep_inst] = new_gep
            end
        end

        # pass 3: propagate through secondary GEPs and rewrite loads
        # visited_loads tracks rewritten loads separately from replacements (which is
        # strictly for pointer-value substitutions) to avoid type confusion in the dict.
        visited_loads = Set{LLVM.Value}()
        changed = true
        @dispose builder=IRBuilder() begin
            while changed
                changed = false
                for bb in blocks(fun), inst in collect(instructions(bb))
                    haskey(replacements, inst) && continue
                    inst in visited_loads && continue
                    oc = opcode(inst)
                    if oc == LLVM.API.LLVMGetElementPtr
                        base = operands(inst)[1]
                        haskey(replacements, base) || continue
                        position!(builder, inst)
                        ops = operands(inst)
                        indices = LLVM.Value[ops[i] for i in 2:length(ops)]
                        elem_ty = gep_source_element_type(inst)
                        new_gep = if gep_is_inbounds(inst)
                            inbounds_gep!(builder, elem_ty, replacements[base], indices)
                        else
                            gep!(builder, elem_ty, replacements[base], indices)
                        end
                        replacements[inst] = new_gep
                        changed = true
                    elseif oc == LLVM.API.LLVMLoad
                        ptr_op = operands(inst)[1]
                        haskey(replacements, ptr_op) || continue
                        position!(builder, inst)
                        new_load = load!(builder, value_type(inst), replacements[ptr_op])
                        alignment!(new_load, alignment(inst))
                        replace_uses!(inst, new_load)
                        push!(visited_loads, inst)
                        changed = true
                    end
                end
            end
        end

        # pass 4: erase dead old instructions (loads first, then GEPs)
        for old in visited_loads
            isempty(uses(old)) && erase!(old)
        end
        for old in keys(replacements)
            isempty(uses(old)) && erase!(old)
        end
    end
end

[maleadt]

TTI-based fix in #782