diff --git a/Makefile b/Makefile index c66ac327781f..7d4543a44300 100644 --- a/Makefile +++ b/Makefile @@ -545,6 +545,7 @@ SOURCE_FILES = \ Lower.cpp \ LowerParallelTasks.cpp \ LowerSMEStreamingTasks.cpp \ + LowerStructTypes.cpp \ LowerWarpShuffles.cpp \ Memoization.cpp \ Module.cpp \ @@ -753,6 +754,7 @@ HEADER_FILES = \ Lower.h \ LowerParallelTasks.h \ LowerSMEStreamingTasks.h \ + LowerStructTypes.h \ LowerWarpShuffles.h \ Memoization.h \ Module.h \ diff --git a/python_bindings/halide/src/halide_/PyEnums.cpp b/python_bindings/halide/src/halide_/PyEnums.cpp index d6327abab3ba..36353e601f57 100644 --- a/python_bindings/halide/src/halide_/PyEnums.cpp +++ b/python_bindings/halide/src/halide_/PyEnums.cpp @@ -250,7 +250,8 @@ void define_enums(py::module &m) { .value("Int", Type::Int) .value("UInt", Type::UInt) .value("Float", Type::Float) - .value("Handle", Type::Handle); + .value("Handle", Type::Handle) + .value("Struct", Type::StructKind); py::enum_(m, "OutputFileType") .value("assembly", OutputFileType::assembly) diff --git a/python_bindings/halide/src/halide_/PyExpr.cpp b/python_bindings/halide/src/halide_/PyExpr.cpp index 473af42a2d8c..36cc646d5e80 100644 --- a/python_bindings/halide/src/halide_/PyExpr.cpp +++ b/python_bindings/halide/src/halide_/PyExpr.cpp @@ -39,6 +39,7 @@ void define_expr(py::module &m) { // for implicitly_convertible .def(py::init([](const FuncRef &f) -> Expr { return f; })) .def(py::init([](const FuncTupleElementRef &f) -> Expr { return f; })) + .def(py::init([](const FieldRef &f) -> Expr { return f; })) .def(py::init([](const Param<> &p) -> Expr { return p; })) .def(py::init([](const RDom &r) -> Expr { return r; })) .def(py::init([](const RVar &r) -> Expr { return r; })) @@ -77,6 +78,7 @@ void define_expr(py::module &m) { // There must be an Expr() ctor available for each of these py::implicitly_convertible(); py::implicitly_convertible(); + py::implicitly_convertible(); py::implicitly_convertible, Expr>(); py::implicitly_convertible(); py::implicitly_convertible(); diff --git a/python_bindings/halide/src/halide_/PyIROperator.cpp b/python_bindings/halide/src/halide_/PyIROperator.cpp index 55e19786b7cf..d7861432de99 100644 --- a/python_bindings/halide/src/halide_/PyIROperator.cpp +++ b/python_bindings/halide/src/halide_/PyIROperator.cpp @@ -1,7 +1,9 @@ #include "PyIROperator.h" +#include #include +#include "PyBinaryOperators.h" #include "PyTuple.h" namespace Halide { @@ -72,6 +74,39 @@ py::object py_select(const py::args &args) { } // namespace void define_operators(py::module &m) { + auto field_ref_class = + py::class_(m, "FieldRef") + .def("__getitem__", &FieldRef::operator[]) + .def("size", &FieldRef::size) + .def("__len__", &FieldRef::size); + add_binary_operators(field_ref_class); + + m.def("field", static_cast(&field), + py::arg("struct_value"), py::arg("name")); + m.def("field", static_cast(&field), + py::arg("struct_value"), py::arg("index")); + // Per-field pack_struct: one initializer per struct field, in declaration + // order. Each item is an Expr (a scalar field's value; a gather() packet or + // a single expression with one `_` placeholder swept over an array field's + // extent) or a FieldRef (to copy a whole same-typed field). Mirrors the C++ + // per-field pack_struct overload. + m.def( + "pack_struct", + [](const Type &type, const py::iterable &fields) -> Expr { + std::vector inits; + for (py::handle item : fields) { + // Check FieldRef first: a scalar FieldRef also converts to Expr, + // but here it means "copy this field", not "use its value". + if (py::isinstance(item)) { + inits.emplace_back(item.cast()); + } else { + inits.emplace_back(item.cast()); + } + } + return pack_struct(type, inits); + }, + py::arg("type"), py::arg("fields")); + m.def("max", [](const py::args &args) -> Expr { if (args.size() < 2) { throw py::value_error("max() must have at least 2 arguments"); @@ -189,6 +224,11 @@ void define_operators(py::module &m) { m.def("strict_float", &strict_float); m.def("scatter", static_cast &)>(&scatter)); m.def("gather", static_cast &)>(&gather)); + // Generator form: gather(extent, gen) builds a packet of `extent` elements + // by calling gen(k) for k in [0, extent), e.g. to fill a struct array field. + m.def("gather", + static_cast &)>(&gather), + py::arg("extent"), py::arg("gen")); m.def("extract_bits", static_cast(&extract_bits)); m.def("concat_bits", &concat_bits); m.def("widen_right_add", &widen_right_add); diff --git a/python_bindings/halide/src/halide_/PyType.cpp b/python_bindings/halide/src/halide_/PyType.cpp index 31404826be05..de962e74110b 100644 --- a/python_bindings/halide/src/halide_/PyType.cpp +++ b/python_bindings/halide/src/halide_/PyType.cpp @@ -38,6 +38,14 @@ std::string halide_type_to_string(const Type &type) { stream << "handle"; } break; + case halide_type_struct: + // A struct's size is a byte count, not a bit width; report it and + // return early so the trailing bits() isn't appended. + stream << "struct" << std::to_string(type.bytes()); + if (type.lanes() > 1) { + stream << "x" + std::to_string(type.lanes()); + } + return stream.str(); default: stream << "#unknown"; break; @@ -51,6 +59,48 @@ std::string halide_type_to_string(const Type &type) { } void define_type(py::module &m) { + py::class_(m, "StructField") + .def(py::init([](const std::string &name, const Type &type, const py::object &array_extent) -> StructField { + std::optional extent; + if (!array_extent.is_none()) { + extent = array_extent.cast(); + } + return StructField{name, type, extent}; + }), + py::arg("name"), py::arg("type"), py::arg("array_extent") = py::none()) + .def(py::init([](const py::tuple &t) -> StructField { + if (t.size() < 2 || t.size() > 3) { + throw py::value_error("StructField requires (name, type) or (name, type, array_extent)"); + } + std::optional extent; + if (t.size() == 3 && !t[2].is_none()) { + extent = t[2].cast(); + } + return StructField{t[0].cast(), t[1].cast(), extent}; + })) + .def_readwrite("name", &StructField::name) + .def_readwrite("type", &StructField::type) + .def_property( + "array_extent", + [](const StructField &f) -> py::object { + return f.array_extent ? py::cast(*f.array_extent) : py::none(); + }, + [](StructField &f, const py::object &v) { + f.array_extent = v.is_none() ? std::nullopt : std::make_optional(v.cast()); + }) + .def("__eq__", [](const StructField &a, const StructField &b) -> bool { return a == b; }) + .def("__ne__", [](const StructField &a, const StructField &b) -> bool { return !(a == b); }); + + // Allow e.g. Type.Struct([("d", Float(16)), ("qs", UInt(8), 16)]), + // mirroring the tuple -> Range convenience conversion above. + py::implicitly_convertible(); + + py::class_(m, "StructTypeInfo") + .def_readonly("fields", &StructTypeInfo::fields) + .def_readonly("offsets", &StructTypeInfo::offsets) + .def_readonly("total_bytes", &StructTypeInfo::total_bytes) + .def("find_field", &StructTypeInfo::find_field, py::arg("name")); + py::class_(m, "Type") .def(py::init<>()) .def(py::init(), py::arg("code"), py::arg("bits"), py::arg("lanes")) @@ -73,6 +123,13 @@ void define_type(py::module &m) { .def("is_handle", &Type::is_handle) .def("same_handle_type", &Type::same_handle_type, py::arg("other")) + .def_static("Struct", &Type::Struct, py::arg("fields")) + .def("is_struct", &Type::is_struct) + .def("same_struct_type", &Type::same_struct_type, py::arg("other")) + .def_property_readonly("struct_type", [](const Type &t) -> py::object { + return t.struct_type() ? py::cast(StructTypeInfo(*t.struct_type())) : py::none(); + }) + .def("__eq__", [](const Type &value, Type *value2) -> bool { return value2 && value == *value2; }) .def("__ne__", [](const Type &value, Type *value2) -> bool { return !value2 || value != *value2; }) diff --git a/python_bindings/halide/test/correctness/CMakeLists.txt b/python_bindings/halide/test/correctness/CMakeLists.txt index b9231af7e5c3..284a2952feb5 100644 --- a/python_bindings/halide/test/correctness/CMakeLists.txt +++ b/python_bindings/halide/test/correctness/CMakeLists.txt @@ -29,6 +29,7 @@ set(tests realize_warnings.py runtime_prefixes.py serialization.py + struct_type.py target.py tuple_select.py type.py diff --git a/python_bindings/halide/test/correctness/struct_type.py b/python_bindings/halide/test/correctness/struct_type.py new file mode 100644 index 000000000000..90b1a23f236c --- /dev/null +++ b/python_bindings/halide/test/correctness/struct_type.py @@ -0,0 +1,163 @@ +import halide as hl + + +# A block_q4_0-shaped struct, per doc/StructTypeDesign.md: one float32 scale +# followed by 8 packed byte-wide "codes". This mirrors the C++ +# test/correctness/struct_type.cpp fixture (simplified to float32, since the +# Python bindings have no float16 literal support). +def make_block_type(): + return hl.Type.Struct([("d", hl.Float(32)), ("qs", hl.UInt(8), 8)]) + + +def test_type_struct_basics(): + block_a = make_block_type() + block_b = make_block_type() # independently constructed, same layout + reordered = hl.Type.Struct([("qs", hl.UInt(8), 8), ("d", hl.Float(32))]) + different_field = hl.Type.Struct([("d", hl.Float(32)), ("qs", hl.UInt(8), 4)]) + + assert block_a.is_struct() + assert block_b.is_struct() + assert block_a.bytes() == 4 + 8 + assert block_a.code() == hl.TypeCode.Struct + assert not block_a.is_uint() + assert block_a.lanes() == 1 + + assert block_a == block_b + assert block_a.same_struct_type(block_b) + assert block_a != reordered + assert block_a != different_field + assert block_a != hl.UInt(8) + + # A struct's type code carries its struct-ness, so with_bits/with_lanes + # (which don't change the code) always preserve it; only with_code to a + # different code drops it. + assert block_a.with_bits(8).is_struct() + assert block_a.with_bits(16).is_struct() + assert block_a.with_lanes(1).is_struct() + assert not block_a.with_code(hl.TypeCode.UInt).is_struct() + + info = block_a.struct_type + assert info is not None + assert info.total_bytes == 12 + assert [f.name for f in info.fields] == ["d", "qs"] + assert info.fields[0].type == hl.Float(32) + assert info.fields[1].type == hl.UInt(8) + assert info.fields[1].array_extent == 8 + assert info.fields[0].array_extent is None + assert list(info.offsets) == [0, 4] + assert info.find_field("qs") == 1 + assert info.find_field("nope") == -1 + + assert hl.Int(32).struct_type is None + assert not hl.Int(32).is_struct() + + +# Struct types are treated as an ordinary opaque element type throughout the +# compiler: a Func whose value is pack_struct(...), consumed via field() by +# another Func, must produce identical results whether the producer is left +# at its default (inlined) schedule or explicitly compute_root()'d (which +# forces a real struct-typed buffer to be materialized). +def _build_pipeline(schedule): + block_t = make_block_type() + blk = hl.Var("blk") + k = hl.Var("k") + + producer = hl.Func("producer") + # Per-field: a scalar `d`, and the array field `qs` filled with a gather() + # generator (extent inferred from the struct field). + producer[blk] = hl.pack_struct( + block_t, [hl.f32(blk) + 0.5, hl.gather(8, lambda kk: hl.u8(blk * 3 + kk))] + ) + + schedule(producer) + + delta = hl.Func("delta") + delta[blk] = hl.f32(hl.field(producer[blk], "d")) + + codes = hl.Func("codes") + codes[k, blk] = hl.i32(hl.field(producer[blk], "qs")[k]) + + return delta, codes + + +def _check_results(delta, codes, num_blocks): + delta_buf = delta.realize([num_blocks]) + codes_buf = codes.realize([8, num_blocks]) + for b in range(num_blocks): + assert abs(delta_buf[b] - (b + 0.5)) < 1e-6 + for k in range(8): + assert codes_buf[k, b] == (b * 3 + k) % 256 + + +def test_struct_pack_and_field_inlined(): + delta, codes = _build_pipeline(lambda f: None) + _check_results(delta, codes, num_blocks=6) + + +def test_struct_pack_and_field_compute_root(): + delta, codes = _build_pipeline(lambda f: f.compute_root()) + _check_results(delta, codes, num_blocks=6) + + +# FieldRef itself: scalar fields implicitly convert to Expr, array fields +# support __getitem__/size()/__len__. +def test_field_ref_array_indexing(): + block_t = make_block_type() + blk = hl.Var("blk") + + producer = hl.Func("producer") + # Array field filled by sweeping the `_` placeholder over its extent, with + # index arithmetic (blk + _). + producer[blk] = hl.pack_struct(block_t, [hl.f32(blk), hl.u8(blk + hl._)]) + + qs_field = hl.field(producer[blk], "qs") + assert qs_field.size() == 8 + assert len(qs_field) == 8 + + out = hl.Func("out") + k = hl.Var("k") + out[k, blk] = hl.i32(qs_field[k]) + + out_buf = out.realize([8, 3]) + for b in range(3): + for k_ in range(8): + assert out_buf[k_, b] == b + k_ + + +# pack_struct can copy a whole array field from another struct via field(), +# and fill a scalar field from an explicit gather() over an existing list. +def test_pack_struct_field_copy(): + block_t = make_block_type() + blk = hl.Var("blk") + k = hl.Var("k") + + src = hl.Func("src") + qs = [hl.u8(blk * 2 + kk) for kk in range(8)] + src[blk] = hl.pack_struct(block_t, [hl.f32(blk), hl.gather(qs)]) + src.compute_root() + + # Re-pack: replace `d`, keep `qs` by copying the whole array field. + s = src[blk] + repacked = hl.Func("repacked") + repacked[blk] = hl.pack_struct(block_t, [hl.f32(blk) + 100.0, hl.field(s, "qs")]) + + out_d = hl.Func("out_d") + out_q = hl.Func("out_q") + out_d[blk] = hl.f32(hl.field(repacked[blk], "d")) + out_q[k, blk] = hl.i32(hl.field(repacked[blk], "qs")[k]) + + num_blocks = 5 + d_buf = out_d.realize([num_blocks]) + q_buf = out_q.realize([8, num_blocks]) + for b in range(num_blocks): + assert abs(d_buf[b] - (b + 100.0)) < 1e-6 + for k_ in range(8): + assert q_buf[k_, b] == (b * 2 + k_) % 256 + + +if __name__ == "__main__": + test_type_struct_basics() + test_struct_pack_and_field_inlined() + test_struct_pack_and_field_compute_root() + test_field_ref_array_indexing() + test_pack_struct_field_copy() diff --git a/src/Bounds.cpp b/src/Bounds.cpp index 7aeff2f87bad..d2f82b8a8ab5 100644 --- a/src/Bounds.cpp +++ b/src/Bounds.cpp @@ -256,6 +256,8 @@ class Bounds : public IRVisitor { if ((t.is_uint() || t.is_int()) && t.bits() <= 16) { interval = Interval(t.min(), t.max()); } else { + // Handle and struct types (and wider integers/floats) have no + // meaningful bounded numeric range. interval = Interval::everything(); } } diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 02215600ef8e..511917c9f6b7 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -159,6 +159,7 @@ target_sources( Lower.h LowerParallelTasks.h LowerSMEStreamingTasks.h + LowerStructTypes.h LowerWarpShuffles.h Memoization.h Module.h @@ -340,6 +341,7 @@ target_sources( Lower.cpp LowerParallelTasks.cpp LowerSMEStreamingTasks.cpp + LowerStructTypes.cpp LowerWarpShuffles.cpp Memoization.cpp Module.cpp diff --git a/src/CodeGen_C.cpp b/src/CodeGen_C.cpp index cd537062185f..3e201d82d63a 100644 --- a/src/CodeGen_C.cpp +++ b/src/CodeGen_C.cpp @@ -2383,6 +2383,11 @@ void CodeGen_C::visit(const Allocate *op) { string op_name = print_name(op->name); string op_type = print_type(op->type, AppendSpace); + string op_bytes = std::to_string(op->type.bytes()); + + // A struct's C storage element is a raw uint8_t (see type_to_c_type), so + // its per-element size must come from the true byte count, not sizeof. + string elem_size_expr = op->type.is_struct() ? op_bytes : ("sizeof(" + op_type + ")"); // For sizes less than 8k, do a stack allocation bool on_stack = false; @@ -2437,8 +2442,8 @@ void CodeGen_C::visit(const Allocate *op) { } stream << get_indent() << "if ((" << size_id << " > ((int64_t(1) << 31) - 1)) || ((" - << size_id << " * sizeof(" - << op_type << ")) > ((int64_t(1) << 31) - 1)))\n"; + << size_id << " * " + << elem_size_expr << ") > ((int64_t(1) << 31) - 1)))\n"; open_scope(); stream << get_indent(); // TODO: call halide_error_buffer_allocation_too_large() here instead @@ -2468,8 +2473,10 @@ void CodeGen_C::visit(const Allocate *op) { stream << get_indent() << op_type; if (on_stack) { + // For a struct, op_type is "uint8_t", so size_id is really an element count. + string declared_size = op->type.is_struct() ? (size_id + " * " + op_bytes) : size_id; stream << op_name - << "[" << size_id << "];\n"; + << "[" << declared_size << "];\n"; } else { // Shouldn't ever currently be possible to have !on_stack && size_id.empty(), // but reality-check in case things change in the future. @@ -2478,9 +2485,9 @@ void CodeGen_C::visit(const Allocate *op) { << op_name << " = (" << op_type - << " *)halide_malloc(_ucon, sizeof(" - << op_type - << ")*" << size_id << ");\n"; + << " *)halide_malloc(_ucon, " + << elem_size_expr + << "*" << size_id << ");\n"; heap_allocations.push(op->name); } } diff --git a/src/CodeGen_D3D12Compute_Dev.cpp b/src/CodeGen_D3D12Compute_Dev.cpp index f58b43452271..9c4c0f488a04 100644 --- a/src/CodeGen_D3D12Compute_Dev.cpp +++ b/src/CodeGen_D3D12Compute_Dev.cpp @@ -1431,6 +1431,12 @@ void CodeGen_D3D12Compute_Dev::CodeGen_D3D12Compute_C::visit(const Allocate *op) << "Only fixed-size allocations are supported on the gpu. " << "Try storing into shared memory instead."; + // A struct is stored as a raw byte array (print_storage_type emits a byte + // element type), so size is an element count and must be scaled to bytes. + if (op->type.is_struct()) { + size *= op->type.bytes(); + } + stream << get_indent() << print_storage_type(op->type) << " " << print_name(op->name) << "[" << size << "];\n"; stream << get_indent(); diff --git a/src/CodeGen_LLVM.cpp b/src/CodeGen_LLVM.cpp index 67b91e9b1775..c2a3da1ab7fd 100644 --- a/src/CodeGen_LLVM.cpp +++ b/src/CodeGen_LLVM.cpp @@ -1078,7 +1078,7 @@ void CodeGen_LLVM::compile_buffer(const Buffer<> &buf) { Constant *type_fields[] = { ConstantInt::get(i8_t, buf.type().code()), ConstantInt::get(i8_t, buf.type().bits()), - ConstantInt::get(i16_t, 0)}; + ConstantInt::get(i16_t, buf.type().to_abi().reserved)}; // struct byte size, else 0 Constant *shape = nullptr; if (buf.dimensions()) { @@ -1273,7 +1273,7 @@ llvm::Function *CodeGen_LLVM::embed_metadata_getter(const std::string &metadata_ Constant *type_fields[] = { ConstantInt::get(i8_t, args[arg].type.code()), ConstantInt::get(i8_t, args[arg].type.bits()), - ConstantInt::get(i16_t, 0)}; // reserved (formerly lanes); must be 0 + ConstantInt::get(i16_t, args[arg].type.to_abi().reserved)}; // struct byte size, else 0 Constant *type = ConstantStruct::get(type_t_type, type_fields); auto argument_estimates = args[arg].argument_estimates; @@ -6031,6 +6031,9 @@ int CodeGen_LLVM::get_vector_num_elements(const llvm::Value *v) { llvm::Type *CodeGen_LLVM::llvm_type_of(LLVMContext *c, Halide::Type t, int effective_vscale) const { if (t.lanes() == 1) { + if (t.is_struct()) { + return llvm::ArrayType::get(llvm::Type::getInt8Ty(*c), t.bytes()); + } if (t.is_float() && !t.is_bfloat()) { switch (t.bits()) { case 16: diff --git a/src/CodeGen_Metal_Dev.cpp b/src/CodeGen_Metal_Dev.cpp index 08145925aa10..b25b342b2280 100644 --- a/src/CodeGen_Metal_Dev.cpp +++ b/src/CodeGen_Metal_Dev.cpp @@ -557,8 +557,15 @@ void CodeGen_Metal_Dev::CodeGen_Metal_C::visit(const Allocate *op) { << "Only fixed-size allocations are supported on the gpu. " << "Try storing into shared memory instead."; + // A struct is stored as a raw byte array (print_storage_type emits a byte + // element type), so size is an element count and must be scaled to bytes. + if (op->type.is_struct()) { + size *= op->type.bytes(); + } + stream << get_indent() << print_storage_type(op->type) << " " << print_name(op->name) << "[" << size << "];\n"; + stream << get_indent() << "#define " << get_memory_space(op->name) << " thread\n"; Allocation alloc; diff --git a/src/CodeGen_OpenCL_Dev.cpp b/src/CodeGen_OpenCL_Dev.cpp index 88e858355f1e..41b40bdfa238 100644 --- a/src/CodeGen_OpenCL_Dev.cpp +++ b/src/CodeGen_OpenCL_Dev.cpp @@ -798,8 +798,15 @@ void CodeGen_OpenCL_Dev::CodeGen_OpenCL_C::visit(const Allocate *op) { << "Only fixed-size allocations are supported on the gpu. " << "Try storing into shared memory instead."; + // A struct is stored as a raw byte array (print_storage_type emits a byte + // element type), so size is an element count and must be scaled to bytes. + if (op->type.is_struct()) { + size *= op->type.bytes(); + } + stream << get_indent() << print_type(op->type) << " " << print_name(op->name) << "[" << size << "];\n"; + stream << get_indent() << "#define " << get_memory_space(op->name) << " __private\n"; Allocation alloc; diff --git a/src/CodeGen_Vulkan_Dev.cpp b/src/CodeGen_Vulkan_Dev.cpp index b7d8dd844667..c23a8b062196 100644 --- a/src/CodeGen_Vulkan_Dev.cpp +++ b/src/CodeGen_Vulkan_Dev.cpp @@ -1888,6 +1888,11 @@ void CodeGen_Vulkan_Dev::SPIRV_Emitter::visit(const Allocate *op) { // static fixed size allocation if (op->extents.size() == 1 && is_const(op->extents[0])) { array_size = op->constant_allocation_size(); + // A struct is stored as a raw byte array (print_storage_type emits a byte + // element type), so size is an element count and must be scaled to bytes. + if (op->type.is_struct()) { + array_size *= op->type.bytes(); + } array_type_id = builder.declare_type(op->type, array_size); builder.add_symbol(variable_name + "_array_type", array_type_id, builder.current_module().id()); debug(2) << "Vulkan: Allocate (fixed-size) " << op->name << " type=" << op->type << " array_size=" << array_size << " in shared memory on device in global scope\n"; @@ -1895,6 +1900,8 @@ void CodeGen_Vulkan_Dev::SPIRV_Emitter::visit(const Allocate *op) { } else { // dynamic allocation with unknown size at compile time ... + // TODO: what to do for struct types here? It doesn't seem to depend on op->type. + // declare the array size as a specialization constant (which will get overridden at runtime) Type array_size_type = UInt(32); array_size = std::max(workgroup_size[0], uint32_t(1)); // use one item per workgroup as an initial guess @@ -1934,6 +1941,12 @@ void CodeGen_Vulkan_Dev::SPIRV_Emitter::visit(const Allocate *op) { << "Allocation " << op->name << " has a dynamic size. " << "Only fixed-size local allocations are supported with Vulkan."; + // A struct is stored as a raw byte array (print_storage_type emits a byte + // element type), so size is an element count and must be scaled to bytes. + if (op->type.is_struct()) { + array_size *= op->type.bytes(); + } + debug(2) << "Vulkan: Allocate " << op->name << " type=" << op->type << " size=" << array_size << " on device in function scope\n"; array_type_id = builder.declare_type(op->type, array_size); diff --git a/src/CodeGen_WebGPU_Dev.cpp b/src/CodeGen_WebGPU_Dev.cpp index 23db34fa02ad..089cf36d7e8a 100644 --- a/src/CodeGen_WebGPU_Dev.cpp +++ b/src/CodeGen_WebGPU_Dev.cpp @@ -428,6 +428,12 @@ void CodeGen_WebGPU_Dev::CodeGen_WGSL::visit(const Allocate *op) { << "Only fixed-size allocations are supported on the gpu. " << "Try storing into shared memory instead."; + // A struct is stored as a raw byte array (print_storage_type emits a byte + // element type), so size is an element count and must be scaled to bytes. + if (op->type.is_struct()) { + size *= op->type.bytes(); + } + stream << get_indent() << "var " << print_name(op->name) << " : array<" << print_type(op->type) << ", " << size << ">;\n"; diff --git a/src/FuseGPUThreadLoops.cpp b/src/FuseGPUThreadLoops.cpp index a444273c9a46..fa3a48b5ea1a 100644 --- a/src/FuseGPUThreadLoops.cpp +++ b/src/FuseGPUThreadLoops.cpp @@ -997,13 +997,15 @@ class ExtractSharedAndHeapAllocations : public IRMutator { internal_assert(ratio != 0) << "alloc_type should have been at most as wide as the widest type in group\n"; // Sizes here are counted in units of one type or another, and - // are converted between them by dividing, so the types have to - // be whole multiples of each other. - internal_assert(is_power_of_two(alloc_type.bytes()) && - is_power_of_two(alloc.widest_type.bytes())) - << "Allocation types must be a power of two bytes wide, but these " - << "are " << alloc_type.bytes() << " and " - << alloc.widest_type.bytes() << "\n"; + // are converted between them by dividing, so the widest type + // has to be a whole multiple of alloc_type -- not necessarily + // a power of two, e.g. when alloc_type is UInt(8) (always the + // case when may_merge_allocs_of_different_type is set) any + // byte count, including a struct's, divides evenly. + internal_assert(alloc.widest_type.bytes() % alloc_type.bytes() == 0) + << "Allocation type " << alloc.widest_type << " (" << alloc.widest_type.bytes() + << " bytes) is not a whole multiple of the cluster's allocation type " + << alloc_type << " (" << alloc_type.bytes() << " bytes)\n"; total_size += align_up(alloc.max_size * ratio, async_copy_alignment / alloc_type.bytes()); } diff --git a/src/IR.cpp b/src/IR.cpp index a5ff626ed7e6..d1d3f8ccd567 100644 --- a/src/IR.cpp +++ b/src/IR.cpp @@ -912,6 +912,8 @@ constexpr const char *intrinsic_op_names[] = { "strict_mul", "strict_sub", "stringify", + "struct_field_read", + "struct_pack", "target_arch_is", "target_bits", "target_has_feature", diff --git a/src/IR.h b/src/IR.h index 113b1876a830..2204515f381c 100644 --- a/src/IR.h +++ b/src/IR.h @@ -890,6 +890,15 @@ struct Call : public ExprNode { strict_sub, // Convert a list of Exprs to a string stringify, + // Read one element of one field out of a struct-typed value. + // args = {struct value, field index as an IntImm, element index}. + // The field's name is resolved to its index eagerly; the element index + // is 0 for a scalar field. + struct_field_read, + // Construct a value of a struct type from its field elements, in + // declaration order, flattened element-by-element (an array field + // contributes array_extent args). + struct_pack, // Query properties of the compiled-for target (resolved at compile-time) target_arch_is, target_bits, diff --git a/src/IROperator.cpp b/src/IROperator.cpp index 710883026201..abf2398bf560 100644 --- a/src/IROperator.cpp +++ b/src/IROperator.cpp @@ -18,6 +18,7 @@ #include "IRVisitor.h" #include "Interval.h" #include "StrictifyFloat.h" +#include "Substitute.h" #include "Util.h" #include "Var.h" @@ -1041,7 +1042,7 @@ void split_into_ands(const Expr &cond, std::vector &result) { } Expr BufferBuilder::build() const { - std::vector args(10); + std::vector args(11); if (buffer_memory.defined()) { args[0] = buffer_memory; } else { @@ -1079,7 +1080,11 @@ Expr BufferBuilder::build() const { args[5] = (int)type.code(); args[6] = type.bits(); - args[7] = dimensions; + // A struct element type carries its packed byte size in the ABI's reserved + // field; thread it through so the runtime buffer's type is faithful (and + // its element stride/allocation size correct). Zero for ordinary types. + args[7] = (int)type.to_abi().reserved; + args[8] = dimensions; std::vector shape; for (size_t i = 0; i < (size_t)dimensions; i++) { @@ -1107,11 +1112,11 @@ Expr BufferBuilder::build() const { } Expr shape_arg = Call::make(type_of(), Call::make_struct, shape, Call::Intrinsic); if (shape_memory.defined()) { - args[8] = shape_arg; + args[9] = shape_arg; } else if (dimensions == 0) { - args[8] = make_zero(type_of()); + args[9] = make_zero(type_of()); } else { - args[8] = shape_var; + args[9] = shape_var; } Expr flags = make_zero(UInt(64)); @@ -1125,7 +1130,7 @@ Expr BufferBuilder::build() const { make_const(UInt(64), halide_buffer_flag_device_dirty), make_zero(UInt(64))); } - args[9] = flags; + args[10] = flags; Expr e = Call::make(type_of(), Call::buffer_init, args, Call::Extern); @@ -2956,6 +2961,16 @@ Expr gather(const std::vector &args) { return make_scatter_gather(args); } +Expr gather(int extent, const std::function &gen) { + user_assert(extent >= 1) << "gather() extent must be at least 1, got " << extent << ".\n"; + std::vector elems; + elems.reserve(extent); + for (int k = 0; k < extent; k++) { + elems.push_back(gen(make_const(Int(32), k))); + } + return make_scatter_gather(elems); +} + Expr extract_bits(Type t, const Expr &e, const Expr &lsb) { return Call::make(t, Call::extract_bits, {e, lsb}, Call::Intrinsic); } @@ -2970,6 +2985,189 @@ Expr concat_bits(const std::vector &e) { return Call::make(t.with_bits(t.bits() * (int)e.size()), Call::concat_bits, e, Call::Intrinsic); } +Expr FieldRef::read(const Expr &elem_index) const { + return Call::make(elem_type, Call::struct_field_read, + {struct_value, make_const(Int(32), field_index), elem_index}, Call::PureIntrinsic); +} + +FieldRef::FieldRef(Expr struct_value, int field_index, Type elem_type, std::optional array_extent) + : struct_value(std::move(struct_value)), field_index(field_index), elem_type(elem_type), array_extent(array_extent) { +} + +FieldRef::operator Expr() const { + user_assert(!array_extent.has_value()) + << "This struct field is an array field; use operator[] to access an element of it, " + << "not an implicit conversion to Expr.\n"; + return read(make_const(Int(32), 0)); +} + +Expr FieldRef::operator[](const Expr &i) const { + user_assert(array_extent.has_value()) + << "This struct field is a scalar field; it can't be indexed with operator[].\n"; + Expr idx = cast(i); + if (auto ci = as_const_int(idx)) { + user_assert(*ci >= 0 && *ci < *array_extent) + << "Struct array field index " << *ci << " is out of range [0, " << *array_extent << ").\n"; + } + return read(idx); +} + +namespace { +FieldRef make_field_ref(const Expr &struct_value, int index) { + user_assert(struct_value.defined() && struct_value.type().is_struct()) + << "field() requires an Expr of a struct type.\n"; + const StructTypeInfo *info = struct_value.type().struct_type(); + user_assert(index >= 0 && index < (int)info->fields.size()) + << "Struct field index " << index << " is out of range; this struct has " + << info->fields.size() << " fields.\n"; + const StructField &f = info->fields[index]; + return FieldRef(struct_value, index, f.type, f.array_extent); +} +} // namespace + +FieldRef field(const Expr &struct_value, const std::string &name) { + user_assert(struct_value.defined() && struct_value.type().is_struct()) + << "field() requires an Expr of a struct type (see Type::Struct).\n"; + int index = struct_value.type().struct_type()->find_field(name); + user_assert(index >= 0) << "Struct type has no field named \"" << name << "\".\n"; + return make_field_ref(struct_value, index); +} + +FieldRef field(const Expr &struct_value, int index) { + return make_field_ref(struct_value, index); +} + +Expr pack_struct(const Type &t, const std::vector &field_values) { + user_assert(t.is_struct()) << "pack_struct() requires a struct type.\n"; + const StructTypeInfo &info = *t.struct_type(); + + size_t expected = 0; + for (const auto &f : info.fields) { + expected += (size_t)f.array_extent.value_or(1); + } + user_assert(field_values.size() == expected) + << "pack_struct() for a " << expected << "-element struct (counting array fields " + << "element-by-element) was given " << field_values.size() << " values.\n"; + + size_t idx = 0; + for (const auto &f : info.fields) { + int n = f.array_extent.value_or(1); + for (int i = 0; i < n; i++) { + user_assert(field_values[idx].defined() && field_values[idx].type() == f.type) + << "pack_struct(): value " << idx << " has type " << field_values[idx].type() + << " but struct field \"" << f.name << "\" has type " << f.type << ".\n"; + idx++; + } + } + + return Call::make(t, Call::struct_pack, field_values, Call::PureIntrinsic); +} + +namespace { +// Collect the distinct implicit-var names (`_`, `_0`, ...) appearing in an +// Expr, in first-encountered order. These are the placeholder(s) a pack_struct +// array-field initializer sweeps over its extent. +class FindImplicitVars : public IRGraphVisitor { + using IRGraphVisitor::visit; + std::set seen; + + void visit(const Variable *op) override { + if (Var::is_implicit(op->name) && seen.insert(op->name).second) { + names.push_back(op->name); + } + } + +public: + std::vector names; +}; + +std::vector find_implicit_vars(const Expr &e) { + FindImplicitVars finder; + e.accept(&finder); + return finder.names; +} +} // namespace + +Expr pack_struct(const Type &t, const std::vector &fields) { + user_assert(t.is_struct()) << "pack_struct() requires a struct type.\n"; + const StructTypeInfo &info = *t.struct_type(); + + user_assert(fields.size() == info.fields.size()) + << "pack_struct() for a " << info.fields.size() << "-field struct type was given " + << fields.size() << " field initializers. Supply exactly one initializer per field " + << "(an array field's elements go in a single gather() or a swept `_` expression).\n"; + + std::vector flat; + for (size_t fi = 0; fi < info.fields.size(); fi++) { + const StructField &f = info.fields[fi]; + const int extent = f.array_extent.value_or(1); + const bool is_array = f.array_extent.has_value(); + const StructFieldInit &init = fields[fi]; + + if (init.source) { + // Copy a whole same-typed field out of another struct. + const FieldRef &r = *init.source; + user_assert(r.element_type() == f.type) + << "pack_struct(): field \"" << f.name << "\" of type " << f.type + << " is initialized from a field of type " << r.element_type() << ".\n"; + user_assert(r.size() == extent) + << "pack_struct(): field \"" << f.name << "\" has " << extent + << " element(s), but is initialized from a field with " << r.size() << ".\n"; + if (is_array) { + for (int k = 0; k < extent; k++) { + flat.push_back(r[k]); + } + } else { + flat.push_back(Expr(r)); + } + continue; + } + + const Expr &e = *init.value; + if (const Call *g = e.as(); g != nullptr && g->is_intrinsic(Call::scatter_gather)) { + // An explicit gather() packet fills an array field's elements. + user_assert(is_array) + << "pack_struct(): scalar field \"" << f.name << "\" was given a gather() packet; " + << "a gather fills an array field.\n"; + user_assert((int)g->args.size() == extent) + << "pack_struct(): array field \"" << f.name << "\" has extent " << extent + << ", but its gather() supplies " << g->args.size() << " element(s).\n"; + flat.insert(flat.end(), g->args.begin(), g->args.end()); + continue; + } + + std::vector implicit = find_implicit_vars(e); + user_assert(implicit.size() <= 1) + << "pack_struct(): the initializer for field \"" << f.name << "\" contains " + << implicit.size() << " placeholders (" << (implicit.empty() ? "" : implicit[0]) + << " ...); an array field may sweep at most one.\n"; + + if (!implicit.empty()) { + // A single `_` placeholder swept over the field's extent. + user_assert(is_array) + << "pack_struct(): scalar field \"" << f.name << "\" was given a swept `_` " + << "expression; a placeholder sweep fills an array field.\n"; + for (int k = 0; k < extent; k++) { + flat.push_back(substitute(implicit[0], make_const(Int(32), k), e)); + } + } else { + // A plain value: only a scalar field. + user_assert(!is_array) + << "pack_struct(): array field \"" << f.name << "\" needs its " << extent + << " elements supplied via gather(), a swept `_` expression, or a field() copy, " + << "not a single value.\n"; + flat.push_back(e); + } + } + + // The flattened form does the final per-element type check and builds the node. + return pack_struct(t, flat); +} + +Expr pack_struct(const Type &t, std::initializer_list fields) { + return pack_struct(t, std::vector(fields)); +} + Expr target_arch_is(Target::Arch arch) { return Call::make(Bool(), Call::target_arch_is, {Expr((int)arch)}, Call::PureIntrinsic); } diff --git a/src/IROperator.h b/src/IROperator.h index 6489a580e106..4f65d39b45ba 100644 --- a/src/IROperator.h +++ b/src/IROperator.h @@ -8,8 +8,10 @@ */ #include +#include #include #include +#include #include #include "Bounds.h" @@ -1610,22 +1612,34 @@ f(select(p, scatter(3, 5, 5), scatter(1, 2, 3))) = f(select(p, gather(5, 3, 3), * * Note that in the p == true case, we redundantly load from 3 and write * to 5 twice. +* +* A gather is also the way to supply the elements of a struct's array field to +* pack_struct(): the packet's values become that field's elements, in order. +* See pack_struct and \ref gather(int, const std::function &). */ //@{ Expr scatter(const std::vector &args); Expr gather(const std::vector &args); -template +template && ...)>> Expr scatter(const Expr &e, Args &&...args) { return scatter({e, std::forward(args)...}); } -template +template && ...)>> Expr gather(const Expr &e, Args &&...args) { return gather({e, std::forward(args)...}); } // @} +/** Build a gather packet of `extent` elements by evaluating `gen(k)` for each + * `k` in `[0, extent)` (passed as an int32 constant). This is the general way + * to fill a struct array field with a computed value per element when the fill + * can't be written as a single placeholder sweep; see pack_struct. */ +Expr gather(int extent, const std::function &gen); + /** Extract a contiguous subsequence of the bits of 'e', starting at the bit * index given by 'lsb', where zero is the least-significant bit, returning a * value of type 't'. Any out-of-range bits requested are filled with zeros. @@ -1674,6 +1688,93 @@ f32.vectorize(x, 8); */ Expr concat_bits(const std::vector &e); +/** A reference to one field of a struct-typed Expr, returned by field(). + * Implicitly converts to Expr for a scalar field; supports operator[] for + * an array field. Using the wrong one of these for how the field was actually + * declared is a user_error. + */ +class FieldRef { + Expr struct_value; + int field_index; + Type elem_type; + std::optional array_extent; + + Expr read(const Expr &elem_index) const; + +public: + FieldRef(Expr struct_value, int field_index, Type elem_type, std::optional array_extent); + + /** Valid only for a scalar field. */ + operator Expr() const; + + /** Valid only for an array field. i may be a runtime Expr. */ + Expr operator[](const Expr &i) const; + + /** The number of elements, for an array field. 1 for a scalar field. */ + int size() const { + return array_extent.value_or(1); + } + + /** Whether this refers to an array field (vs. a scalar field). */ + bool is_array() const { + return array_extent.has_value(); + } + + /** The element type of the field. */ + Type element_type() const { + return elem_type; + } +}; + +/** Extract field `name` (or `index`) from a struct-typed Expr. The field's + * offset and type are resolved eagerly. + */ +// @{ +FieldRef field(const Expr &struct_value, const std::string &name); +FieldRef field(const Expr &struct_value, int index); +// @} + +/** Construct a value of struct type `t` (see Type::Struct) from one Expr + * per field, in declaration order. All fields must be supplied; array + * fields take `array_extent` Exprs, flattened into the same list. This is the + * low-level form; prefer the per-field overload below. */ +Expr pack_struct(const Type &t, const std::vector &field_values); + +/** One field's worth of initializer for the per-field pack_struct() overload. + * Implicitly constructible from anything convertible to Expr -- a scalar + * field's value; a gather() packet whose elements fill an array field; or a + * single expression containing one implicit-var placeholder `_`, swept over the + * array field's extent (index arithmetic on `_` is allowed) -- and from a + * FieldRef, to copy a whole same-typed field out of another struct. */ +class StructFieldInit { +public: + template && + !std::is_same_v, FieldRef>>> + StructFieldInit(T &&e) + : value(Expr(std::forward(e))) { + } + StructFieldInit(const FieldRef &f) + : source(f) { + } + +private: + friend Expr pack_struct(const Type &t, const std::vector &fields); + // Exactly one of these is engaged. + std::optional value; + std::optional source; +}; + +/** Construct a value of struct type `t` with one initializer per field, in + * declaration order (contrast the flattened form above). A scalar field takes a + * scalar Expr; an array field takes a gather() packet, a single expression with + * one `_` placeholder swept over the field extent, or a FieldRef to a + * same-typed array field to copy element-by-element. */ +// @{ +Expr pack_struct(const Type &t, const std::vector &fields); +Expr pack_struct(const Type &t, std::initializer_list fields); +// @} + /** Below is a collection of intrinsics for fixed-point programming. Most of * them can be expressed via other means, but this is more natural for some, as * it avoids ghost widened intermediates that don't (or shouldn't) actually show diff --git a/src/IRPrinter.cpp b/src/IRPrinter.cpp index b4ca72d54e3a..5e9de62b5a8a 100644 --- a/src/IRPrinter.cpp +++ b/src/IRPrinter.cpp @@ -46,8 +46,22 @@ ostream &operator<<(ostream &out, const Type &type) { case Type::BFloat: out << "bfloat"; break; + case Type::StructKind: + out << "struct{"; + if (const StructTypeInfo *info = type.struct_type()) { + const char *sep = ""; + for (const auto &f : info->fields) { + out << sep << f.name << ": " << f.type; + if (f.array_extent) { + out << "[" << *f.array_extent << "]"; + } + sep = ", "; + } + } + out << "}"; + break; } - if (!type.is_handle()) { + if (!type.is_handle() && !type.is_struct()) { out << type.bits(); } if (type.lanes() > 1) { diff --git a/src/ImageParam.cpp b/src/ImageParam.cpp index 615ae1e6013b..a61fff5bfeb3 100644 --- a/src/ImageParam.cpp +++ b/src/ImageParam.cpp @@ -41,7 +41,7 @@ Func ImageParam::create_func() const { void ImageParam::set(const Buffer<> &b) { if (b.defined()) { - user_assert(b.type() == type()) + user_assert(type().is_compatible_for_buffer_bind(b.type())) << "Can't bind ImageParam " << name() << " of type " << type() << " to Buffer " << b.name() diff --git a/src/Lower.cpp b/src/Lower.cpp index 21acfff8df2d..d6d7da79e465 100644 --- a/src/Lower.cpp +++ b/src/Lower.cpp @@ -46,6 +46,7 @@ #include "LoopCarry.h" #include "LowerParallelTasks.h" #include "LowerSMEStreamingTasks.h" +#include "LowerStructTypes.h" #include "LowerWarpShuffles.h" #include "Memoization.h" #include "OffloadGPULoops.h" @@ -305,6 +306,10 @@ void lower_impl(const vector &output_funcs, s = storage_flattening(s, outputs, env, t); log("Lowering after storage flattening:", s); + debug(1) << "Lowering struct type field access...\n"; + s = lower_struct_types(s); + log("Lowering after lowering struct type field access:", s); + debug(1) << "Adding atomic mutex allocation...\n"; s = add_atomic_mutex(s, outputs); log("Lowering after adding atomic mutex allocation:", s); diff --git a/src/LowerStructTypes.cpp b/src/LowerStructTypes.cpp new file mode 100644 index 000000000000..f3d6c1de2cfa --- /dev/null +++ b/src/LowerStructTypes.cpp @@ -0,0 +1,330 @@ +#include "LowerStructTypes.h" +#include "IRMutator.h" +#include "IROperator.h" +#include "IRVisitor.h" +#include "Scope.h" + +namespace Halide { +namespace Internal { + +using std::vector; + +namespace { + +// Number of pack_struct() args occupied by all fields before field_index. +int flat_start_of_field(const StructTypeInfo &info, int field_index) { + int start = 0; + for (int i = 0; i < field_index; i++) { + start += info.fields[i].array_extent.value_or(1); + } + return start; +} + +class LowerStructTypesMutator : public IRMutator { + using IRMutator::visit; + + // Struct-typed Let bindings are tracked here because project_field() needs + // direct syntactic access to the underlying struct_pack()/Select/Load. Because + // struct-typed values are never materialized (only individual fields are), + // inlining it back in at each field() use site costs nothing at runtime. + Scope struct_lets; + + Expr resolve_struct_value(Expr e) const { + while (const Variable *v = e.as()) { + internal_assert(v->type.is_struct()); + user_assert(struct_lets.contains(v->name)) + << "Struct-typed variable \"" << v->name << "\" is unbound; a struct-typed Expr " + << "may only be the value of a struct-typed Func/Store, or a direct field() " + << "argument, in v1.\n"; + e = struct_lets.get(v->name); + } + return e; + } + + static Expr fold_field_of_pack(const Call *pack, const StructTypeInfo &info, int field_index, const Expr &elem_index) { + const StructField &f = info.fields[field_index]; + int flat_start = flat_start_of_field(info, field_index); + int extent = f.array_extent.value_or(1); + + if (auto ci = as_const_int(elem_index)) { + int flat = flat_start + (int)*ci; + internal_assert(flat >= 0 && flat < (int)pack->args.size()); + return pack->args[flat]; + } + + internal_assert(extent >= 1); + Expr result = pack->args[flat_start + extent - 1]; + for (int i = extent - 2; i >= 0; i--) { + result = select(elem_index == i, pack->args[flat_start + i], result); + } + return result; + } + + // Rewrite `field(e, field_index)[elem_index]`, where e is a real (already + // flattened) struct-typed Load into byte-addressed Loads plus a Reinterpret. + Expr lower_field_read_from_load(const Load *load, const Type &field_type, int field_base_offset, const Expr &elem_index_in) { + // A field whose own type is itself Type::Struct (a sub-struct field) + // can't be read this way: the byte-combine-and-Reinterpret trick + // below fundamentally assumes field_type's *bits* describe its + // *bytes* (true for every ordinary scalar type), but a struct's bits() + // is not its (packed, possibly not power-of-two) byte size, so + // Reinterpret would be asked to reinterpret between mismatched bit + // widths. Nested struct fields read from a literal pack_struct() go + // through fold_field_of_pack instead and don't hit this at all; only + // reading a nested field from a (non-inlined) struct-typed buffer is + // unsupported. + user_assert(!field_type.is_struct()) + << "Reading a nested struct field (a field whose own type is Type::Struct) directly " + << "from a struct-typed buffer/Func that isn't fully inlined is not supported. " + << "Field \"" << field_type << "\" at byte offset " << field_base_offset + << " within " << load->name << " was requested this way.\n"; + + Expr elem_index = mutate(elem_index_in); + Expr flat_index = mutate(load->index); + Type index_type = flat_index.type(); + int elem_bytes = field_type.bytes(); + + Expr byte_offset = flat_index * make_const(index_type, load->type.bytes()) + + make_const(index_type, field_base_offset) + + cast(index_type, elem_index) * make_const(index_type, elem_bytes); + + auto byte_load = [&](int b) { + Expr idx = b == 0 ? byte_offset : byte_offset + make_const(index_type, b); + return Load::make(UInt(8), load->name, idx, load->image, load->param, + const_true(), ModulusRemainder(), load->is_streaming); + }; + + if (elem_bytes == 1) { + Expr bl = byte_load(0); + return field_type == UInt(8) ? bl : Reinterpret::make(field_type, bl); + } + + vector bytes; + bytes.reserve(elem_bytes); + for (int b = 0; b < elem_bytes; b++) { + bytes.push_back(byte_load(b)); + } + + if (field_type.is_float()) { + // A float field (e.g. a packed fp16 delta) is built as a genuine + // shift/or chain of the bytes rather than through concat_bits's + // vector-shuffle-then-reinterpret lowering. Some AArch64 backends + // (LLVM < 23) cannot legalize a bitcast straight from a vector to + // a scalar half, and unlike the integer case below, there's no way + // to route around that once the value is vector-shaped: any chain + // of pure bitcasts collapses back to the direct (illegal) one + // during optimization. Building the packed bits with real ALU ops + // keeps the final Reinterpret's input a plain scalar integer, so + // the last step is an always-legal scalar-to-scalar bitcast; LLVM's + // load-combiner still typically recovers a single wide load. + Type uint_t = UInt(8 * elem_bytes); + Expr combined = cast(uint_t, bytes[0]); + for (int b = 1; b < elem_bytes; b++) { + combined = combined | (cast(uint_t, bytes[b]) << (8 * b)); + } + return Reinterpret::make(field_type, combined); + } + + // Keep a packed scalar field as a bit-concatenation of adjacent bytes + // rather than immediately expanding it to a shift/or tree. The + // concat_bits lowering turns this into a dense byte shuffle followed + // by a vector reinterpret, which gives LLVM enough structure to issue + // one (possibly unaligned) wide load. Expanding here obscures that the + // byte loads are adjacent once individual extracts of the field have + // simplified, and commonly leaves one scalar load per byte. + Expr combined = concat_bits(bytes); + return field_type == combined.type() ? combined : Reinterpret::make(field_type, combined); + } + + // Project field field_index's element elem_index (0 for a scalar + // field) out of a struct-typed Expr. Recurses through the only shapes a + // struct-typed Expr can legally have at this point in lowering: a + // literal struct_pack(), a Select between two struct-typed branches, a + // struct-typed Let (via struct_lets), or a genuine flattened Load. + Expr project_field(const Expr &struct_expr, int field_index, const Expr &elem_index) { + Expr e = resolve_struct_value(struct_expr); + const StructTypeInfo *info = e.type().struct_type(); + internal_assert(info != nullptr) << "project_field applied to a non-struct-typed Expr.\n"; + internal_assert(field_index >= 0 && field_index < (int)info->fields.size()); + + if (const Call *pack = e.as(); pack != nullptr && pack->is_intrinsic(Call::struct_pack)) { + return mutate(fold_field_of_pack(pack, *info, field_index, elem_index)); + } + + // e is itself an unresolved nested field() access (e.g. field(field(outer, "inner"), "a"), + // written as two separate front-end field() calls). We resolve the inner one first, then + // continue projecting the requested field out of *that* result. + // This is what makes nested struct fields work. + if (const Call *nested = e.as(); nested != nullptr && nested->is_intrinsic(Call::struct_field_read)) { + auto nested_field_index = as_const_int(nested->args[1]); + internal_assert(nested_field_index) << "struct_field_read's field index must be a constant.\n"; + Expr resolved_inner = project_field(nested->args[0], (int)*nested_field_index, nested->args[2]); + return project_field(resolved_inner, field_index, elem_index); + } + + if (const Select *sel = e.as