diff --git a/src/metal.jl b/src/metal.jl index c3b47ed0..0f3b6f36 100644 --- a/src/metal.jl +++ b/src/metal.jl @@ -773,6 +773,11 @@ function lower_air!(@nospecialize(job::CompilerJob{MetalCompilerTarget}), mod::L end end + # the AGX back-end cannot handle integers whose width is not a power of two (see + # `legalize_odd_width_integers!`). run after the cleanups above, as InstCombine would + # narrow the promoted operations back, but before merging the byte GEPs this introduces. + legalize_odd_width_integers!(mod) + # flatten chained byte GEPs that the AGX back-end miscompiles for 1-byte accesses on a # 2-threadgroup grid (see `merge_byte_gep_chains!`). run last, after the intrinsic-lowering # cleanup above, so the merged form is what reaches the AIR downgrader / back-end. @@ -802,6 +807,285 @@ function load_thread_pointers_as_words!(mod::LLVM.Module) return end +# Legalization of odd-width integers +# +# The AGX back-end cannot handle integers wider than a byte whose width is not a power of two +# (Metal.jl#1125). Arithmetic on them crashes the compiler, except for a `zext` that it folds +# into a naturally-aligned load; underaligned loads of them crash too, and naturally-aligned +# stores are widened to the next power of two, overwriting the bytes after the value. LLVM's +# SROA introduces such integers when it splits or widens an aggregate, e.g. copying the 3 +# padding bytes of `{i32, i8}` as one `i24`, or assembling a 7-byte alloca in an `i56`. +# +# Legalize them like SelectionDAG's type legalizer does for other back-ends: promote every +# odd-width value to the next power of two, leaving its upper bits undefined, and split loads +# and stores into power-of-two pieces. Operations whose result depends on the upper bits +# (right shifts, division, comparisons and extensions) first clear or sign-fill them. Uses +# that are not legalized (e.g., calls) get the value truncated back, and odd-width values +# that are not legalized (e.g., call results) are extended where they are defined. +# +# This has to run after the last InstCombine, which narrows promoted operations back. + +is_odd_width(T::LLVMType) = + T isa LLVM.IntegerType && 8 < LLVM.width(T) && !ispow2(LLVM.width(T)) +promoted_type(T::LLVM.IntegerType) = LLVM.IntType(nextpow(2, LLVM.width(T))) + +function legalize_odd_width_integers!(mod::LLVM.Module) + changed = false + for f in mod.functions + isdeclaration(f) && continue + changed |= legalize_odd_width_integers!(f) + end + return changed +end + +function legalize_odd_width_integers!(f::LLVM.Function) + odd(v) = is_odd_width(v.value_type) + any(LLVM.blocks(f)) do bb + any(inst -> odd(inst) || any(odd, LLVM.operands(inst)), LLVM.instructions(bb)) + end || return false + + # visit the blocks in reverse post-order, so that values are legalized before their uses + # (except in phis, which are completed last), and then the unreachable ones + order = LLVM.BasicBlock[] + visited = Set{LLVM.BasicBlock}([LLVM.entry(f)]) + stack = [(LLVM.entry(f), 1)] + while !isempty(stack) + bb, i = stack[end] + succs = collect(LLVM.successors(LLVM.terminator(bb))) + if i <= length(succs) + stack[end] = (bb, i + 1) + succ = succs[i] + if !(succ in visited) + push!(visited, succ) + push!(stack, (succ, 1)) + end + else + pop!(stack) + push!(order, bb) + end + end + reverse!(order) + append!(order, Iterators.filter(bb -> !(bb in visited), LLVM.blocks(f))) + + typed_ptrs = supports_typed_pointers(context()) + T_i8 = LLVM.Int8Type() + aa_kinds = (MD_tbaa, MD_alias_scope, MD_noalias) + + wide = Dict{LLVM.Value,LLVM.Value}() # odd-width value => its promoted value + producers = LLVM.Instruction[] # legalized odd-width instructions + consumers = LLVM.Instruction[] # legalized instructions using odd-width values + phis = LLVM.PHIInst[] + + function position_before_non_phi!(builder, inst) + while inst isa LLVM.PHIInst + inst = LLVM.next(inst) + end + position!(builder, LLVM.before(inst)) + end + position_after!(builder, inst) = position_before_non_phi!(builder, LLVM.next(inst)) + + # the promoted value of an odd-width `v` + function promote(v) + haskey(wide, v) && return wide[v] + @dispose builder=IRBuilder() begin + if v isa LLVM.Instruction + position_after!(builder, v) + else + position_before_non_phi!(builder, first(LLVM.instructions(LLVM.entry(f)))) + end + return wide[v] = zext!(builder, v, promoted_type(v.value_type)) + end + end + + @dispose builder=IRBuilder() begin + # the promoted value of `v` with the bits above `T` cleared, or filled with its sign + function zext_in_reg(v, T) + W = promoted_type(T) + return and!(builder, promote(v), ConstantInt(W, (big(1) << LLVM.width(T)) - 1)) + end + function sext_in_reg(v, T) + W = promoted_type(T) + amount = ConstantInt(W, LLVM.width(W) - LLVM.width(T)) + return ashr!(builder, shl!(builder, promote(v), amount), amount) + end + resize(v, T, extend) = + LLVM.width(T) > LLVM.width(v.value_type) ? extend(builder, v, T) : + LLVM.width(T) < LLVM.width(v.value_type) ? trunc!(builder, v, T) : v + + # little-endian (offset, size) pieces of an access of `T`, in bytes + function pieces(T) + nbytes = cld(LLVM.width(T), 8) + ps = Tuple{Int,Int}[] + offset = 0 + while offset < nbytes + size = prevpow(2, nbytes - offset) + push!(ps, (offset, size)) + offset += size + end + return ps + end + function piece_pointer(inst, offset, T_piece) + ptr = inst.pointer_operand + as = ptr.value_type.addrspace + typed_ptrs && (ptr = bitcast!(builder, ptr, LLVM.PointerType(T_i8, as))) + offset > 0 && + (ptr = inbounds_gep!(builder, T_i8, ptr, [ConstantInt(Int64(offset))])) + typed_ptrs && (ptr = bitcast!(builder, ptr, LLVM.PointerType(T_piece, as))) + align = offset == 0 ? inst.alignment : + min(inst.alignment, 1 << trailing_zeros(offset)) + return ptr, align + end + function copy_metadata!(dst, src) + md = src.metadata + for kind in aa_kinds + haskey(md, kind) && (dst.metadata[kind] = md[kind]) + end + end + + # `op` of `src` to `T`, where either is an odd-width integer + function legalize_conversion(op, src, T) + extend = op == LLVM.API.LLVMSExt ? sext! : zext! + if odd(src) + src = op == LLVM.API.LLVMZExt ? zext_in_reg(src, src.value_type) : + op == LLVM.API.LLVMSExt ? sext_in_reg(src, src.value_type) : + promote(src) + end + return resize(src, T, extend) + end + + # the promoted version of an odd-width instruction, or `nothing` + function legalize_producer(inst) + T = inst.value_type + W = promoted_type(T) + op = LLVM.opcode(inst) + ops = LLVM.operands(inst) + if inst isa LLVM.LoadInst + # AIR atomics are lowered separately, and Julia only emits power-of-two ones + isatomic(inst) && return nothing + val = nothing + for (offset, size) in pieces(T) + T_piece = LLVM.IntType(8 * size) + ptr, align = piece_pointer(inst, offset, T_piece) + part = load!(builder, T_piece, ptr; align, volatile=inst.volatile) + copy_metadata!(part, inst) + part = zext!(builder, part, W) + offset > 0 && (part = shl!(builder, part, ConstantInt(W, 8 * offset))) + val = val === nothing ? part : or!(builder, val, part) + end + return val + elseif op in (LLVM.API.LLVMAdd, LLVM.API.LLVMSub, LLVM.API.LLVMMul, + LLVM.API.LLVMAnd, LLVM.API.LLVMOr, LLVM.API.LLVMXor) + build = op == LLVM.API.LLVMAdd ? add! : op == LLVM.API.LLVMSub ? sub! : + op == LLVM.API.LLVMMul ? mul! : op == LLVM.API.LLVMAnd ? and! : + op == LLVM.API.LLVMOr ? or! : xor! + return build(builder, promote(ops[1]), promote(ops[2])) + elseif op == LLVM.API.LLVMShl + return shl!(builder, promote(ops[1]), zext_in_reg(ops[2], T)) + elseif op == LLVM.API.LLVMLShr + return lshr!(builder, zext_in_reg(ops[1], T), zext_in_reg(ops[2], T)) + elseif op == LLVM.API.LLVMAShr + return ashr!(builder, sext_in_reg(ops[1], T), zext_in_reg(ops[2], T)) + elseif op in (LLVM.API.LLVMUDiv, LLVM.API.LLVMURem) + build = op == LLVM.API.LLVMUDiv ? udiv! : urem! + return build(builder, zext_in_reg(ops[1], T), zext_in_reg(ops[2], T)) + elseif op in (LLVM.API.LLVMSDiv, LLVM.API.LLVMSRem) + build = op == LLVM.API.LLVMSDiv ? sdiv! : srem! + return build(builder, sext_in_reg(ops[1], T), sext_in_reg(ops[2], T)) + elseif op == LLVM.API.LLVMSelect && !(ops[1].value_type isa LLVM.VectorType) + return select!(builder, ops[1], promote(ops[2]), promote(ops[3])) + elseif op == LLVM.API.LLVMFreeze + return LLVM.freeze!(builder, promote(ops[1])) + elseif op in (LLVM.API.LLVMZExt, LLVM.API.LLVMSExt, LLVM.API.LLVMTrunc) + return legalize_conversion(op, ops[1], W) + end + return nothing + end + + # the legalized version of an instruction using odd-width values, or `nothing` + function legalize_consumer(inst) + op = LLVM.opcode(inst) + ops = LLVM.operands(inst) + if inst isa LLVM.StoreInst && odd(inst.value_operand) && !isatomic(inst) + T = inst.value_operand.value_type + W = promoted_type(T) + val = promote(inst.value_operand) + local part + for (offset, size) in pieces(T) + T_piece = LLVM.IntType(8 * size) + ptr, align = piece_pointer(inst, offset, T_piece) + part = offset == 0 ? val : + lshr!(builder, val, ConstantInt(W, 8 * offset)) + part = store!(builder, trunc!(builder, part, T_piece), ptr; align, + volatile=inst.volatile) + copy_metadata!(part, inst) + end + return part + elseif inst isa LLVM.ICmpInst && odd(ops[1]) + pred = LLVM.predicate(inst) + T = ops[1].value_type + ext = LLVM.is_signed_predicate(pred) ? sext_in_reg : zext_in_reg + return icmp!(builder, pred, ext(ops[1], T), ext(ops[2], T)) + elseif op in (LLVM.API.LLVMZExt, LLVM.API.LLVMSExt, LLVM.API.LLVMTrunc) && + odd(ops[1]) + return legalize_conversion(op, ops[1], inst.value_type) + end + return nothing + end + + for bb in order, inst in collect(LLVM.instructions(bb)) + if inst isa LLVM.PHIInst + odd(inst) || continue + position!(builder, LLVM.before(inst)) + wide[inst] = phi!(builder, promoted_type(inst.value_type)) + push!(phis, inst) + push!(producers, inst) + continue + end + # (positioning the builder gives new instructions the original's debug location) + position!(builder, LLVM.before(inst)) + if odd(inst) + new = legalize_producer(inst) + new === nothing && continue + wide[inst] = new + push!(producers, inst) + elseif any(odd, LLVM.operands(inst)) + new = legalize_consumer(inst) + new === nothing && continue + inst isa LLVM.StoreInst || replace_uses!(inst, new) + push!(consumers, inst) + end + end + end + + for phi in phis + append!(LLVM.incoming(wide[phi]), + [(promote(val), bb) for (val, bb) in LLVM.incoming(phi)]) + end + + # remaining (not legalized) uses of legalized values get them truncated back + truncs = LLVM.Instruction[] + @dispose builder=IRBuilder() begin + for inst in producers + isempty(LLVM.uses(inst)) && continue + w = wide[inst] + if w isa LLVM.Instruction + position_after!(builder, w) + else + position!(builder, LLVM.before(inst)) + end + narrow = trunc!(builder, w, inst.value_type) + narrow isa LLVM.Instruction && push!(truncs, narrow) + replace_uses!(inst, narrow) + end + end + foreach(erase!, consumers) + foreach(erase!, producers) + for narrow in truncs + isempty(LLVM.uses(narrow)) && erase!(narrow) + end + return true +end + # lowering of LLVM atomics # # Metal has no LLVM back-end, so the work that `AtomicExpand` and instruction selection do for diff --git a/test/metal.jl b/test/metal.jl index 658a79e6..496ae34d 100644 --- a/test/metal.jl +++ b/test/metal.jl @@ -1053,6 +1053,131 @@ end end end +@testset "odd-width integers" begin + # The AGX back-end cannot handle integers whose width is not a power of two: arithmetic + # and underaligned loads crash it, and naturally-aligned stores overwrite the bytes after + # the value (Metal.jl#1125). Such integers are promoted, and their accesses split into + # power-of-two pieces. + p(T) = opaque_ptrs ? "ptr addrspace(1)" : "$T addrspace(1)*" + accesses = [("i24", 1), ("i24", 4), ("i12", 1), ("i56", 1), ("i72", 8)] + ir = join(map(enumerate(accesses)) do (i, (T, align)) + """ + define void @f$i($(p(T)) %p, $(p(T)) %q) { + %v = load $T, $(p(T)) %p, align $align + store $T %v, $(p(T)) %q, align $align + ret void + } + """ + end, "\n") + + # the expected (offset, width, alignment) of the pieces + pieces = Dict( + 1 => [(0, 16, 1), (2, 8, 1)], + 2 => [(0, 16, 4), (2, 8, 2)], + 3 => [(0, 16, 1)], + 4 => [(0, 32, 1), (4, 16, 1), (6, 8, 1)], + 5 => [(0, 64, 8), (8, 8, 8)], + ) + + @dispose ctx=Context() begin + mod = parse(LLVM.Module, ir) + @test GPUCompiler.legalize_odd_width_integers!(mod) + @test !GPUCompiler.legalize_odd_width_integers!(mod) + @test LLVM.verify(mod) === nothing + + piece_offset(ptr) = + ptr isa LLVM.BitCastInst ? piece_offset(ptr.operands[1]) : + ptr isa LLVM.GetElementPtrInst ? convert(Int, ptr.operands[2]) : 0 + for (i, expected) in pieces + insts = collect(only(mod.functions["f$i"].blocks).instructions) + loads = filter(inst -> inst isa LLVM.LoadInst, insts) + stores = filter(inst -> inst isa LLVM.StoreInst, insts) + @test map(loads) do ld + (piece_offset(ld.pointer_operand), LLVM.width(ld.value_type), ld.alignment) + end == expected + @test map(stores) do st + (piece_offset(st.pointer_operand), LLVM.width(st.value_operand.value_type), + st.alignment) + end == expected + end + end + + # operations on odd-width values, including in a loop, with no odd-width value left + ir = """ + define i64 @f(i64 %x, i64 %y, i64 %n) { + entry: + %a = trunc i64 %x to i56 + %b = trunc i64 %y to i56 + br label %loop + loop: + %i = phi i64 [ 0, %entry ], [ %i.next, %loop ] + %acc = phi i56 [ %a, %entry ], [ %acc.next, %loop ] + %sum = add i56 %acc, %b + %shr = ashr i56 %sum, 3 + %div = udiv i56 %shr, 7 + %big = icmp slt i56 %div, %b + %acc.next = select i1 %big, i56 %div, i56 %sum + %i.next = add i64 %i, 1 + %done = icmp eq i64 %i.next, %n + br i1 %done, label %exit, label %loop + exit: + %r = sext i56 %acc.next to i64 + ret i64 %r + } + + define i32 @g($(p("i24")) %p) { + entry: + br label %exit + dead: + %v = load i24, $(p("i24")) %p, align 1 + br label %exit + exit: + %r = phi i24 [ 0, %entry ], [ %v, %dead ] + %z = zext i24 %r to i32 + ret i32 %z + } + """ + @dispose ctx=Context() begin + mod = parse(LLVM.Module, ir) + @test GPUCompiler.legalize_odd_width_integers!(mod) + @test LLVM.verify(mod) === nothing + odd(v) = GPUCompiler.is_odd_width(v.value_type) + @testset "$(LLVM.name(f))" for f in mod.functions + @test !any(bb -> any(inst -> odd(inst) || any(odd, LLVM.operands(inst)), + LLVM.instructions(bb)), + LLVM.blocks(f)) + end + end + + # ... as part of the AIR lowering + mod = @eval module $(gensym()) + const ir = $(opaque_ptrs ? """ + define void @entry(ptr addrspace(1) %p) #0 { + %q = getelementptr i8, ptr addrspace(1) %p, i64 5 + %v = load i24, ptr addrspace(1) %q, align 1 + store i24 %v, ptr addrspace(1) %p, align 4 + ret void + } + attributes #0 = { alwaysinline }""" : """ + define void @entry(i8 addrspace(1)* %p) #0 { + %q = getelementptr i8, i8 addrspace(1)* %p, i64 5 + %qq = bitcast i8 addrspace(1)* %q to i24 addrspace(1)* + %v = load i24, i24 addrspace(1)* %qq, align 1 + %pp = bitcast i8 addrspace(1)* %p to i24 addrspace(1)* + store i24 %v, i24 addrspace(1)* %pp, align 4 + ret void + } + attributes #0 = { alwaysinline }""") + function kernel(p::Core.LLVMPtr{UInt8,1}) + Base.llvmcall((ir, "entry"), Nothing, Tuple{Core.LLVMPtr{UInt8,1}}, p) + return + end + end + air = sprint(io -> Metal.code_native(io, mod.kernel, Tuple{Core.LLVMPtr{UInt8,1}}; + kernel=true)) + @test !occursin("i24", air) +end + @testset "boxed constant classification" begin JuliaContext() do ctx mod = LLVM.Module("boxed-constant-classification")