Layouts for targets whose data layout differs from Julia's

NEWS.md

@@ -71,3 +71,6 @@
   inference, with CodeInstances of its own that keep their sources, `collect_codeinfos`),
   cached by CodeInstance (so redefining a method recompiles what depends on it), with
   the objects the code refers to linked and kept alive, and thread-safe.
+- Layouts for targets whose data layout aligns a field differently than Julia does (e.g.
+  `Int128` on SPIR-V, Metal and BPF): a packed struct with the fields at Julia's offsets
+  and explicit padding.

src/constants.jl

@@ -88,14 +88,16 @@ function llvm_constant(fc::FnCtx, @nospecialize(x), @nospecialize(T))
         U = Dict(1 => UInt8, 2 => UInt16, 4 => UInt32, 8 => UInt64, 16 => UInt128)[sizeof(T)]
         return LLVM.ConstantInt(ty, reinterpret(U, x))
     end
-    elements = LLVM.Constant[]
+    # the elements, by their index: the fields, and zero padding (of a padded struct)
+    elements = ty isa LLVM.StructType ? LLVM.Constant[LLVM.null(t) for t in LLVM.elements(ty)] :
+               Vector{LLVM.Constant}(undef, count(>=(0), lay.fieldidx))
     for i in 1:fieldcount(T)
         lay.fieldidx[i] < 0 && continue
         s = lay.storage[i]
-        if s === FieldInline
-            push!(elements, llvm_constant(fc, getfield(x, i), fieldtype(T, i)))
+        elements[lay.fieldidx[i] + 1] = if s === FieldInline
+            llvm_constant(fc, getfield(x, i), fieldtype(T, i))
         elseif s === FieldPointer
-            push!(elements, LLVM.null(T_tracked()))
+            LLVM.null(T_tracked())
         else
             unsupported("a constant of ", T, " with a field stored as ", s)
         end

src/layout.jl

@@ -151,20 +151,52 @@ function compute_layout(ls::Layouts, @nospecialize T)
     ty = isempty(elements) ? LLVM.StructType(LLVM.LLVMType[]) :
          homogeneous ? LLVM.ArrayType(elements[1], length(elements)) :
          LLVM.StructType(elements)
-    if !isempty(elements)
-        check_size(ls, T, ty)
-        if ty isa LLVM.StructType
-            for i in 1:n
-                idx[i] < 0 && continue
-                LLVM.offsetof(ls.datalayout, ty, idx[i] + 1) == fieldoffset(T, i) ||
-                    unsupported("the layout of ", T, ": the offset of field ", i,
-                                " differs between Julia and LLVM")
-            end
-        end
+    if !isempty(elements) && !same_offsets(ls, T, ty, idx)
+        # the target's data layout aligns a field differently than Julia does (e.g. `i128`
+        # to 8 bytes on targets that don't say otherwise, where Julia aligns `Int128` to
+        # 16): a packed struct with the fields at Julia's offsets, and explicit padding
+        ty, idx = padded_struct(ls, T, idx, elements)
     end
+    isempty(elements) || check_size(ls, T, ty)
     return Layout(ty, idx, storage, Aggregate)
 end
 
+# whether the LLVM aggregate `ty` has the size of `T`, and the fields of `T` (the elements
+# `idx`) at Julia's offsets
+function same_offsets(ls::Layouts, @nospecialize(T), ty::LLVM.LLVMType, idx::Vector{Int})
+    dl = ls.datalayout
+    LLVM.abi_size(dl, ty) == sizeof(T) || return false
+    for i in 1:fieldcount(T)
+        idx[i] < 0 && continue
+        offset = ty isa LLVM.StructType ? LLVM.offsetof(dl, ty, idx[i] + 1) :
+                 idx[i] * LLVM.abi_size(dl, ty.element_type)
+        offset == fieldoffset(T, i) || return false
+    end
+    return true
+end
+
+# a packed struct of the `elements` of the fields of `T` (the elements `idx`) at Julia's
+# offsets, with arrays of bytes as the padding; and the fields' elements in it
+function padded_struct(ls::Layouts, @nospecialize(T), idx::Vector{Int},
+                       elements::Vector{LLVM.LLVMType})
+    dl = ls.datalayout
+    padded = LLVM.LLVMType[]
+    newidx = fill(-1, length(idx))
+    offset = 0
+    for i in sortperm(1:fieldcount(T); by=i -> fieldoffset(T, i))
+        idx[i] < 0 && continue
+        fo = Int(fieldoffset(T, i))
+        fo < offset && unsupported("the layout of ", T, ": overlapping fields on this target")
+        fo > offset && push!(padded, LLVM.ArrayType(LLVM.Int8Type(), fo - offset))
+        newidx[i] = length(padded)
+        el = elements[idx[i] + 1]
+        push!(padded, el)
+        offset = fo + Int(LLVM.abi_size(dl, el))
+    end
+    sizeof(T) > offset && push!(padded, LLVM.ArrayType(LLVM.Int8Type(), sizeof(T) - offset))
+    return LLVM.StructType(padded; packed=true), newidx
+end
+
 function check_size(ls::Layouts, @nospecialize(T), ty::LLVM.LLVMType)
     sz = LLVM.abi_size(ls.datalayout, ty)
     sz == sizeof(T) || unsupported("the layout of ", T, ": its LLVM type has ", sz,

test/jit.jl

@@ -110,13 +110,15 @@ test_fallbacks(@nospecialize(f), @nospecialize(types); config=test_config()) =
     test_compile(Tuple{Core.Typeof(f), types...}; config).fallbacks
 
 """
-    test_module_ir(f, types; triple) -> String
+    test_module_ir(f, types; triple, datalayout=nothing) -> String
 
 The (unoptimized) module that JuliaCodegen generates for `f(::types...)` for the target
-`triple`, which need not be the host's: the code is not run.
+`triple` (with the data layout `datalayout`, if given), which need not be the host's: the
+code is not run.
 """
-function test_module_ir(@nospecialize(f), @nospecialize(types); triple::AbstractString)
-    config = JuliaCodegen.CodegenConfig(; target=JuliaCodegen.Target(; triple),
+function test_module_ir(@nospecialize(f), @nospecialize(types); triple::AbstractString,
+                        datalayout=nothing)
+    config = JuliaCodegen.CodegenConfig(; target=JuliaCodegen.Target(; triple, datalayout),
                                         on_unsupported=:error, on_internal_error=:error)
     @dispose ctx=LLVM.Context() mod=LLVM.Module("jctarget") begin
         cgm = JuliaCodegen.CGModule(mod, config)

test/repr.jl

@@ -95,6 +95,37 @@ end
     end
 end
 
+struct R_PadPair; a::Float32; b::Int128; end
+struct R_Nested; x::R_PadPair; y::Int8; end
+r_padded(p::R_PadPair) = R_Nested(R_PadPair(p.a + 1f0, p.b), Int8(3))
+
+# targets whose data layout aligns a field differently than Julia: SPIR-V's aligns `i128`
+# to 8 bytes, Julia `Int128` to 16
+@testset "layouts on other targets" begin
+    triple = "spirv64-unknown-unknown"
+    datalayout = "e-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-G1-ni:10:11:12:13"
+    config = JC.CodegenConfig(; target=JC.Target(; triple, datalayout), on_unsupported=:error,
+                              on_internal_error=:error)
+    @dispose ctx=Context() begin
+        ls = JC.Layouts(config)
+        dl = ls.datalayout
+        lay = JC.layout(ls, R_PadPair)
+        @test string(lay.llvm) == "<{ float, [12 x i8], i128 }>"
+        @test lay.fieldidx == [0, 2]
+        lay = JC.layout(ls, R_Nested)
+        @test LLVM.abi_size(dl, lay.llvm) == sizeof(R_Nested)
+        @test LLVM.offsetof(dl, lay.llvm, lay.fieldidx[2] + 1) == fieldoffset(R_Nested, 2)
+        # the host's layout needs no padding
+        host = JC.Layouts(JC.CodegenConfig())
+        Sys.ARCH === :x86_64 && @test string(JC.layout(host, R_PadPair).llvm) == "{ float, i128 }"
+    end
+    # code with such values
+    @test @filecheck begin
+        @check "<{ float, [12 x i8], i128 }>"
+        test_module_ir(r_padded, (R_PadPair,); triple, datalayout)
+    end
+end
+
 @testset "layouts" begin
     config = JC.CodegenConfig()
     @dispose ctx=Context() begin