simple pass at giving union fields an optimized layout

unlike codegen, only bitstypes (!isptr) fields are permitted in the union
and the offset count starts from 0 instead of 1
but otherwise the tindex counter is compatible

Author: vtjnash

src/cgutils.cpp

@@ -1131,6 +1131,33 @@ static Value *emit_bounds_check(jl_codectx_t &ctx, const jl_cgval_t &ainfo, jl_v
 
 // --- loading and storing ---
 
+static Value *compute_box_tindex(Value *datatype, jl_value_t *supertype, jl_value_t *ut, jl_codectx_t *ctx)
+{
+    Value *tindex = ConstantInt::get(T_int8, 0);
+    unsigned counter = 0;
+    for_each_uniontype_small(
+            [&](unsigned idx, jl_datatype_t *jt) {
+                if (jl_subtype((jl_value_t*)jt, supertype)) {
+                    Value *cmp = builder.CreateICmpEQ(literal_pointer_val((jl_value_t*)jt), datatype);
+                    tindex = builder.CreateSelect(cmp, ConstantInt::get(T_int8, idx), tindex);
+                }
+            },
+            ut,
+            counter);
+    return tindex;
+}
+
+// get the runtime tindex value
+static Value *compute_tindex_unboxed(const jl_cgval_t &val, jl_value_t *typ, jl_codectx_t *ctx)
+{
+    if (val.constant)
+        return ConstantInt::get(T_int8, get_box_tindex((jl_datatype_t*)jl_typeof(val.constant), typ));
+    if (val.isboxed)
+        return compute_box_tindex(emit_typeof_boxed(val, ctx), val.typ, typ, ctx);
+    assert(val.TIndex);
+    return builder.CreateAnd(val.TIndex, ConstantInt::get(T_int8, 0x7f));
+}
+
 // If given alignment is 0 and LLVM's assumed alignment for a load/store via ptr
 // might be stricter than the Julia alignment for jltype, return the alignment of jltype.
 // Otherwise return the given alignment.
@@ -1436,6 +1463,9 @@ static jl_cgval_t emit_getfield_knownidx(jl_codectx_t &ctx, const jl_cgval_t &st
                 addr = ctx.builder.CreateStructGEP(lt, ptr, idx);
             }
         }
+        int align = jl_field_offset(jt, idx);
+        align |= 16;
+        align &= -align;
         if (jl_field_isptr(jt, idx)) {
             bool maybe_null = idx >= (unsigned)jt->ninitialized;
             Instruction *Load = maybe_mark_load_dereferenceable(
@@ -1447,6 +1477,29 @@ static jl_cgval_t emit_getfield_knownidx(jl_codectx_t &ctx, const jl_cgval_t &st
                 null_pointer_check(ctx, fldv);
             return mark_julia_type(ctx, fldv, true, jfty, strct.gcroot || !strct.isimmutable);
         }
+        else if (jl_is_uniontype(jfty)) {
+            int fsz = jl_field_size(jt, idx);
+            Value *ptindex = builder.CreateGEP(LLVM37_param(T_int8) emit_bitcast(addr, T_pint8), ConstantInt::get(T_size, fsz - 1));
+            Value *tindex = builder.CreateNUWAdd(ConstantInt::get(T_int8, 1), builder.CreateLoad(ptindex));
+            bool isimmutable = strct.isimmutable;
+            Value *gcroot = strct.gcroot;
+            if (jt->mutabl) {
+                // move value to an immutable stack slot
+                Type *AT = ArrayType::get(IntegerType::get(jl_LLVMContext, 8 * align), (fsz + align - 2) / align);
+                AllocaInst *lv = emit_static_alloca(AT, ctx);
+                if (align > 1)
+                    lv->setAlignment(align);
+                Value *nbytes = ConstantInt::get(T_size, fsz - 1);
+                builder.CreateMemCpy(lv, addr, nbytes, align);
+                addr = lv;
+                isimmutable = true;
+                gcroot = NULL;
+            }
+            jl_cgval_t fieldval = mark_julia_slot(addr, jfty, tindex, strct.tbaa);
+            fieldval.isimmutable = isimmutable;
+            fieldval.gcroot = gcroot;
+            return fieldval;
+        }
         else if (!jt->mutabl) {
             // just compute the pointer and let user load it when necessary
             jl_cgval_t fieldval = mark_julia_slot(addr, jfty, NULL, strct.tbaa);
@@ -2065,7 +2118,7 @@ static void emit_unionmove(jl_codectx_t &ctx, Value *dest, const jl_cgval_t &src
         jl_value_t *typ = src.constant ? jl_typeof(src.constant) : src.typ;
         Type *store_ty = julia_type_to_llvm(typ);
         assert(skip || jl_isbits(typ));
-        if (jl_isbits(typ)) {
+        if (jl_isbits(typ) && jl_datatype_size(typ) > 0) {
             if (!src.ispointer() || src.constant) {
                 emit_unbox(ctx, store_ty, src, typ, dest, isVolatile);
             }
@@ -2236,11 +2289,24 @@ static void emit_setfield(jl_codectx_t &ctx,
                 emit_checked_write_barrier(ctx, boxed(ctx, strct), r);
         }
         else {
-            int align = jl_field_offset(sty, idx0);
-            align |= 16;
-            align &= -align;
-            typed_store(ctx, addr, ConstantInt::get(T_size, 0), rhs, jfty,
-                strct.tbaa, data_pointer(ctx, strct, T_pjlvalue), align);
+            if (jl_is_uniontype(jfty)) {
+                int fsz = jl_field_size(sty, idx0);
+                // compute tindex from rhs
+                jl_cgval_t rhs_union = convert_julia_type(rhs, jfty, ctx);
+                Value *ptindex = builder.CreateGEP(LLVM37_param(T_int8) emit_bitcast(addr, T_pint8), ConstantInt::get(T_size, fsz - 1));
+                Value *tindex = compute_tindex_unboxed(rhs_union, jfty, ctx);
+                tindex = builder.CreateNUWSub(tindex, ConstantInt::get(T_int8, 1));
+                builder.CreateStore(tindex, ptindex);
+                // copy data
+                emit_unionmove(addr, rhs, NULL, false, NULL, ctx);
+            }
+            else {
+                int align = jl_field_offset(sty, idx0);
+                align |= 16;
+                align &= -align;
+                typed_store(addr, ConstantInt::get(T_size, 0), rhs, jfty, ctx,
+                    strct.tbaa, data_pointer(ctx, strct, T_pjlvalue), align);
+            }
         }
     }
     else {

src/codegen.cpp

@@ -743,6 +743,8 @@ static inline jl_cgval_t update_julia_type(jl_codectx_t &ctx, const jl_cgval_t &
     return jl_cgval_t(v, typ, NULL);
 }
 
+static jl_cgval_t convert_julia_type(const jl_cgval_t &v, jl_value_t *typ, jl_codectx_t *ctx, bool needsroot = true);
+
 // --- allocating local variables ---
 
 static jl_sym_t *slot_symbol(jl_codectx_t &ctx, int s)
@@ -823,7 +825,7 @@ static void jl_rethrow_with_add(const char *fmt, ...)
 }
 
 // given a value marked with type `v.typ`, compute the mapping and/or boxing to return a value of type `typ`
-static jl_cgval_t convert_julia_type(jl_codectx_t &ctx, const jl_cgval_t &v, jl_value_t *typ, bool needsroot = true)
+static jl_cgval_t convert_julia_type(jl_codectx_t &ctx, const jl_cgval_t &v, jl_value_t *typ, bool needsroot)
 {
     if (typ == (jl_value_t*)jl_typeofbottom_type)
         return ghostValue(typ); // normalize TypeofBottom to Type{Union{}}
@@ -3544,33 +3546,6 @@ static Value *try_emit_union_alloca(jl_codectx_t &ctx, jl_uniontype_t *ut, bool
     return NULL;
 }
 
-static Value *compute_box_tindex(jl_codectx_t &ctx, Value *datatype, jl_value_t *supertype, jl_value_t *ut)
-{
-    Value *tindex = ConstantInt::get(T_int8, 0);
-    unsigned counter = 0;
-    for_each_uniontype_small(
-            [&](unsigned idx, jl_datatype_t *jt) {
-                if (jl_subtype((jl_value_t*)jt, supertype)) {
-                    Value *cmp = ctx.builder.CreateICmpEQ(maybe_decay_untracked(literal_pointer_val(ctx, (jl_value_t*)jt)), datatype);
-                    tindex = ctx.builder.CreateSelect(cmp, ConstantInt::get(T_int8, idx), tindex);
-                }
-            },
-            ut,
-            counter);
-    return tindex;
-}
-
-// get the runtime tindex value
-static Value *compute_tindex_unboxed(jl_codectx_t &ctx, const jl_cgval_t &val, jl_value_t *typ)
-{
-    if (val.constant)
-        return ConstantInt::get(T_int8, get_box_tindex((jl_datatype_t*)jl_typeof(val.constant), typ));
-    if (val.isboxed)
-        return compute_box_tindex(ctx, emit_typeof_boxed(ctx, val), val.typ, typ);
-    assert(val.TIndex);
-    return ctx.builder.CreateAnd(val.TIndex, ConstantInt::get(T_int8, 0x7f));
-}
-
 static void emit_assignment(jl_codectx_t &ctx, jl_value_t *l, jl_value_t *r)
 {
     if (jl_is_ssavalue(l)) {

src/datatype.c

@@ -216,6 +216,38 @@ unsigned jl_special_vector_alignment(size_t nfields, jl_value_t *t)
     return alignment;
 }
 
+static int jl_layout_isbits(jl_value_t *ty)
+{
+    if (jl_isbits(ty) && jl_is_leaf_type(ty)) {
+        if (((jl_datatype_t*)ty)->layout) // layout check handles possible layout recursion
+            return 1;
+    }
+    return 0;
+}
+
+static unsigned jl_union_isbits(jl_value_t *ty, size_t *nbytes, size_t *align)
+{
+    if (jl_is_uniontype(ty)) {
+        unsigned na = jl_union_isbits(((jl_uniontype_t*)ty)->a, nbytes, align);
+        if (na == 0)
+            return 0;
+        unsigned nb = jl_union_isbits(((jl_uniontype_t*)ty)->b, nbytes, align);
+        if (nb == 0)
+            return 0;
+        return na + nb;
+    }
+    if (jl_layout_isbits(ty)) {
+        size_t sz = jl_datatype_size(ty);
+        size_t al = ((jl_datatype_t*)ty)->layout->alignment;
+        if (*nbytes < sz)
+            *nbytes = sz;
+        if (*align < al)
+            *align = al;
+        return 1;
+    }
+    return 0;
+}
+
 void jl_compute_field_offsets(jl_datatype_t *st)
 {
     size_t sz = 0, alignm = 1;
@@ -277,16 +309,22 @@ void jl_compute_field_offsets(jl_datatype_t *st)
 
     for (size_t i = 0; i < nfields; i++) {
         jl_value_t *ty = jl_field_type(st, i);
-        size_t fsz, al;
-        if (jl_isbits(ty) && jl_is_leaf_type(ty) && ((jl_datatype_t*)ty)->layout) {
-            fsz = jl_datatype_size(ty);
+        size_t fsz = 0, al = 0;
+        unsigned countbits = jl_union_isbits(ty, &fsz, &al);
+        if (countbits > 0 && countbits < 127) {
             // Should never happen
             if (__unlikely(fsz > max_size))
                 goto throw_ovf;
             al = jl_datatype_align(ty);
             desc[i].isptr = 0;
-            if (((jl_datatype_t*)ty)->layout->haspadding)
+            if (jl_is_uniontype(ty)) {
                 haspadding = 1;
+                fsz += 1; // selector byte
+            }
+            else { // isbits struct
+                if (((jl_datatype_t*)ty)->layout->haspadding)
+                    haspadding = 1;
+            }
         }
         else {
             fsz = sizeof(void*);
@@ -312,7 +350,7 @@ void jl_compute_field_offsets(jl_datatype_t *st)
             goto throw_ovf;
         sz += fsz;
     }
-    if (homogeneous && lastty!=NULL && jl_is_tuple_type(st)) {
+    if (homogeneous && lastty != NULL && jl_is_tuple_type(st)) {
         // Some tuples become LLVM vectors with stronger alignment than what was calculated above.
         unsigned al = jl_special_vector_alignment(nfields, lastty);
         assert(al % alignm == 0);
@@ -326,10 +364,12 @@ void jl_compute_field_offsets(jl_datatype_t *st)
     if (st->size > sz)
         haspadding = 1;
     st->layout = jl_get_layout(nfields, alignm, haspadding, desc);
-    if (descsz >= jl_page_size) free(desc);
+    if (descsz >= jl_page_size)
+        free(desc);
     return;
  throw_ovf:
-    if (descsz >= jl_page_size) free(desc);
+    if (descsz >= jl_page_size)
+        free(desc);
     jl_throw(jl_overflow_exception);
 }
 
@@ -704,46 +744,69 @@ JL_DLLEXPORT jl_value_t *jl_get_nth_field(jl_value_t *v, size_t i)
 {
     jl_datatype_t *st = (jl_datatype_t*)jl_typeof(v);
     assert(i < jl_datatype_nfields(st));
-    size_t offs = jl_field_offset(st,i);
-    if (jl_field_isptr(st,i)) {
+    size_t offs = jl_field_offset(st, i);
+    if (jl_field_isptr(st, i)) {
         return *(jl_value_t**)((char*)v + offs);
     }
-    return jl_new_bits(jl_field_type(st,i), (char*)v + offs);
+    jl_value_t *ty = jl_field_type(st, i);
+    if (jl_is_uniontype(ty)) {
+        uint8_t sel = ((uint8_t*)v)[offs + jl_field_size(st, i) - 1];
+        ty = jl_nth_union_component(ty, sel);
+    }
+    return jl_new_bits(ty, (char*)v + offs);
 }
 
 JL_DLLEXPORT jl_value_t *jl_get_nth_field_checked(jl_value_t *v, size_t i)
 {
     jl_datatype_t *st = (jl_datatype_t*)jl_typeof(v);
     if (i >= jl_datatype_nfields(st))
-        jl_bounds_error_int(v, i+1);
-    size_t offs = jl_field_offset(st,i);
-    if (jl_field_isptr(st,i)) {
+        jl_bounds_error_int(v, i + 1);
+    size_t offs = jl_field_offset(st, i);
+    if (jl_field_isptr(st, i)) {
         jl_value_t *fval = *(jl_value_t**)((char*)v + offs);
         if (fval == NULL)
             jl_throw(jl_undefref_exception);
         return fval;
     }
-    return jl_new_bits(jl_field_type(st,i), (char*)v + offs);
+    jl_value_t *ty = jl_field_type(st, i);
+    if (jl_is_uniontype(ty)) {
+        size_t fsz = jl_field_size(st, i);
+        uint8_t sel = ((uint8_t*)v)[offs + fsz - 1];
+        ty = jl_nth_union_component(ty, sel);
+        if (jl_is_datatype_singleton((jl_datatype_t*)ty))
+            return ((jl_datatype_t*)ty)->instance;
+    }
+    return jl_new_bits(ty, (char*)v + offs);
 }
 
 JL_DLLEXPORT void jl_set_nth_field(jl_value_t *v, size_t i, jl_value_t *rhs)
 {
     jl_datatype_t *st = (jl_datatype_t*)jl_typeof(v);
-    size_t offs = jl_field_offset(st,i);
-    if (jl_field_isptr(st,i)) {
+    size_t offs = jl_field_offset(st, i);
+    if (jl_field_isptr(st, i)) {
         *(jl_value_t**)((char*)v + offs) = rhs;
         if (rhs != NULL) jl_gc_wb(v, rhs);
     }
     else {
+        jl_value_t *ty = jl_field_type(st, i);
+        if (jl_is_uniontype(ty)) {
+            uint8_t *psel = &((uint8_t*)v)[offs + jl_field_size(st, i) - 1];
+            unsigned nth = 0;
+            if (!jl_find_union_component(ty, jl_typeof(rhs), &nth))
+                assert(0 && "invalid field assignment to isbits union");
+            *psel = nth;
+            if (jl_is_datatype_singleton((jl_datatype_t*)ty))
+                return;
+        }
         jl_assign_bits((char*)v + offs, rhs);
     }
 }
 
 JL_DLLEXPORT int jl_field_isdefined(jl_value_t *v, size_t i)
 {
     jl_datatype_t *st = (jl_datatype_t*)jl_typeof(v);
-    size_t offs = jl_field_offset(st,i);
-    if (jl_field_isptr(st,i)) {
+    size_t offs = jl_field_offset(st, i);
+    if (jl_field_isptr(st, i)) {
         return *(jl_value_t**)((char*)v + offs) != NULL;
     }
     return 1;