Add struct types to Halide

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 \

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_<OutputFileType>(m, "OutputFileType")
         .value("assembly", OutputFileType::assembly)

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<FuncRef, Expr>();
     py::implicitly_convertible<FuncTupleElementRef, Expr>();
+    py::implicitly_convertible<FieldRef, Expr>();
     py::implicitly_convertible<Param<>, Expr>();
     py::implicitly_convertible<RDom, Expr>();
     py::implicitly_convertible<RVar, Expr>();

python_bindings/halide/src/halide_/PyIROperator.cpp

@@ -1,7 +1,9 @@
 #include "PyIROperator.h"
 
+#include <pybind11/functional.h>
 #include <utility>
 
+#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_<FieldRef>(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<FieldRef (*)(const Expr &, const std::string &)>(&field),
+          py::arg("struct_value"), py::arg("name"));
+    m.def("field", static_cast<FieldRef (*)(const Expr &, int)>(&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<StructFieldInit> 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<FieldRef>(item)) {
+                    inits.emplace_back(item.cast<FieldRef>());
+                } else {
+                    inits.emplace_back(item.cast<Expr>());
+                }
+            }
+            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<Expr (*)(const std::vector<Expr> &)>(&scatter));
     m.def("gather", static_cast<Expr (*)(const std::vector<Expr> &)>(&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<Expr (*)(int, const std::function<Expr(Expr)> &)>(&gather),
+          py::arg("extent"), py::arg("gen"));
     m.def("extract_bits", static_cast<Expr (*)(Type, const Expr &, const Expr &)>(&extract_bits));
     m.def("concat_bits", &concat_bits);
     m.def("widen_right_add", &widen_right_add);

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_<StructField>(m, "StructField")
+        .def(py::init([](const std::string &name, const Type &type, const py::object &array_extent) -> StructField {
+                 std::optional<int> extent;
+                 if (!array_extent.is_none()) {
+                     extent = array_extent.cast<int>();
+                 }
+                 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<int> extent;
+            if (t.size() == 3 && !t[2].is_none()) {
+                extent = t[2].cast<int>();
+            }
+            return StructField{t[0].cast<std::string>(), t[1].cast<Type>(), 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<int>());
+            })
+        .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::tuple, StructField>();
+
+    py::class_<StructTypeInfo>(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_<Type>(m, "Type")
         .def(py::init<>())
         .def(py::init<halide_type_code_t, int, int>(), 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; })
 

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

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()

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();
         }
     }

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