ernie-renderer 0.2.0.dev145__tar.gz → 0.2.0.dev149__tar.gz

{ernie_renderer-0.2.0.dev145 → ernie_renderer-0.2.0.dev149}/CMakeLists.txt

@@ -15,7 +15,7 @@ project(
 )
 
 # Pre-release suffix such as .dev12, written by CI.
-set(ERNIE_VERSION_SUFFIX ".dev145" CACHE STRING "Pre-release suffix of Ernie's version")
+set(ERNIE_VERSION_SUFFIX ".dev149" CACHE STRING "Pre-release suffix of Ernie's version")
 set(ERNIE_VERSION "${PROJECT_VERSION}${ERNIE_VERSION_SUFFIX}")
 
 option(ERNIE_BUILD_TESTS "Build Ernie's tests" OFF)

{ernie_renderer-0.2.0.dev145 → ernie_renderer-0.2.0.dev149}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ernie-renderer
-Version: 0.2.0.dev145
+Version: 0.2.0.dev149
 Summary: Efficient Renderer for Natural and Interplanetary Environments: a physically based, spectral Vulkan ray tracer
 Author: Alex Maietta
 License-Expression: BSD-3-Clause

{ernie_renderer-0.2.0.dev145 → ernie_renderer-0.2.0.dev149}/ernie/app/viewer.cpp

@@ -221,10 +221,11 @@ graph::Image Viewer::AddViewportPass(graph::Graph& graph) {
         CameraOutputs outputs{.colour = viewport,
                               .colourIndex = viewportImage_,
                               .background = glm::vec3(shell_->GetTheme().viewport)};
+        const glm::uvec2 size(viewport_->Description().extent);
         camera_->Add(graph,
                      {.cameraToWorld = orbit.Transform(),
-                      .verticalFov = orbit.verticalFov,
-                      .size = glm::uvec2(viewport_->Description().extent)},
+                      .size = size,
+                      .lens = MakePinhole(orbit.verticalFov, size)},
                      outputs);
         return *outputs.colour;
     }

{ernie_renderer-0.2.0.dev145 → ernie_renderer-0.2.0.dev149}/ernie/python/device.cpp

@@ -82,11 +82,8 @@ struct Context {
         context->WaitIdle();
     }
 
-    CudaInterop& Cuda() const {
-        if (!cuda) {
-            throw nb::buffer_error(("No CUDA device pairs with " + context->Info().name).c_str());
-        }
-        return *cuda;
+    Interop Pair() {
+        return {context.get(), cuda.get()};
     }
     Arenas& Shared() const {
         if (!shared) {
@@ -96,43 +93,55 @@ struct Context {
     }
 };
 
-/// A column, and the ticket of CUDA's writes Ernie has yet to see.
-struct Column {
-    Context& owner;
-    std::unique_ptr<DeviceColumn> column;
-    /// The tensor a wrapped column shares, which `__dlpack__` hands back as it is.
-    std::optional<Tensor> wrapped;
-    Ticket written;
+}  // namespace
+
+CudaInterop& Interop::Cuda() const {
+    if (cuda == nullptr) {
+        throw nb::buffer_error(("No CUDA device pairs with " + context->Info().name).c_str());
+    }
+    return *cuda;
+}
 
-    Column(Context& context, std::unique_ptr<DeviceColumn> device,
-           std::optional<Tensor> tensor = std::nullopt)
-        : owner(context), column(std::move(device)), wrapped(std::move(tensor)) {}
+Column::Column(Interop interop, std::unique_ptr<DeviceColumn> device, std::optional<Tensor> tensor)
+    : owner(interop), owned(std::move(device)), column(owned.get()), wrapped(std::move(tensor)) {}
 
-    nb::object Dlpack(const nb::object& stream, const nb::object& maxVersion) {
-        CudaInterop& cuda = owner.Cuda();
-        // -1 asks for no synchronization.
-        if (stream.is_none() || nb::cast<std::intptr_t>(stream) != -1) {
-            for (const Queue* queue : owner.context->Queues()) {
-                cuda.Wait(ToStream(stream), queue->Last());
-            }
+Column::Column(Interop interop, const DeviceColumn& device, std::vector<std::int64_t> view,
+               Ticket ready)
+    : owner(interop), column(&device), shape(std::move(view)), written(ready) {}
+
+std::vector<std::int64_t> Column::Shape() const {
+    return shape.empty() ? column->Shape() : shape;
+}
+
+nb::object Column::Dlpack(const nb::object& stream, const nb::object& maxVersion) {
+    CudaInterop& cuda = owner.Cuda();
+    // -1 asks for no synchronization.
+    if (stream.is_none() || nb::cast<std::intptr_t>(stream) != -1) {
+        for (const Queue* queue : owner.context->Queues()) {
+            cuda.Wait(ToStream(stream), queue->Last());
         }
-        const bool versioned =
-            !maxVersion.is_none() && nb::cast<std::tuple<int, int>>(maxVersion) >=
-                                         std::tuple<int, int>(DLPACK_MAJOR_VERSION, 0);
-        return Capsule((wrapped ? *wrapped : cuda.View(*column)).ToDlpack(), versioned);
     }
+    const bool versioned =
+        !maxVersion.is_none() &&
+        nb::cast<std::tuple<int, int>>(maxVersion) >= std::tuple<int, int>(DLPACK_MAJOR_VERSION, 0);
+    const Tensor view = wrapped ? *wrapped : cuda.View(*column).Reshaped(Shape());
+    return Capsule(view.ToDlpack(), versioned);
+}
 
-    void Assign(const nb::object& source, const nb::object& stream) {
-        written = owner.Cuda().Assign(*column, TakeTensor(source, stream), ToStream(stream));
-    }
+void Column::Assign(const nb::object& source, const nb::object& stream) {
+    written = owner.Cuda().Assign(*column, TakeTensor(source, stream), ToStream(stream));
+}
 
-    nb::bytes Read() {
-        Readback readback(*owner.context, owner.context->Main(), column->Storage(),
-                          column->Range().offset, column->Bytes(), std::span(&written, 1));
-        const std::span<const std::byte> bytes = readback.Bytes();
-        return nb::bytes(reinterpret_cast<const char*>(bytes.data()), column->Bytes());
-    }
-};
+Tensor Column::Read() const {
+    Readback readback(*owner.context, owner.context->Main(), column->Storage(),
+                      column->Range().offset, column->Bytes(), std::span(&written, 1));
+    Tensor tensor = Tensor::Zeros(column->Type(), Shape());
+    std::ranges::copy(readback.Bytes().first(column->Bytes()),
+                      static_cast<std::byte*>(tensor.Data()));
+    return tensor;
+}
+
+namespace {
 
 /// @return The allocations `ernie_cuda_malloc` made, or none without a CUDA driver: PyTorch's
 /// allocator hook is process-wide, so they are too.
@@ -175,7 +184,7 @@ void BindDevice(nb::module_& module) {
                 if (shape.empty() || shape.front() < 0) {
                     throw nb::value_error("A column's shape starts with its row count");
                 }
-                return new Column(self,
+                return new Column(self.Pair(),
                                   std::make_unique<DeviceColumn>(
                                       self.Shared(), name, ParseType(dtype),
                                       std::vector<std::int64_t>(shape.begin() + 1, shape.end()),
@@ -187,7 +196,7 @@ void BindDevice(nb::module_& module) {
             "wrap",
             [](Context& self, const nb::object& tensor, const std::string& name,
                const nb::object& stream) {
-                CudaInterop& cuda = self.Cuda();
+                CudaInterop& cuda = self.Pair().Cuda();
                 Tensor source = TakeTensor(tensor, stream);
                 std::unique_ptr<DeviceColumn> column = cuda.Wrap(source, name);
                 if (!column) {
@@ -195,7 +204,7 @@ void BindDevice(nb::module_& module) {
                         "The tensor's memory cannot be shared: allocate it with "
                         "ernie.cuda_allocator(), or copy it into a column with assign()");
                 }
-                auto* wrapped = new Column(self, std::move(column), std::move(source));
+                auto* wrapped = new Column(self.Pair(), std::move(column), std::move(source));
                 wrapped->written = cuda.Signal(ToStream(stream));
                 return wrapped;
             },
@@ -204,9 +213,10 @@ void BindDevice(nb::module_& module) {
             "A column over a CUDA tensor's own memory, in place, when it is exportable.");
 
     nb::class_<Column>(module, "DeviceColumn",
-                       "Rows on the device that CUDA, PyTorch and Warp read and write in place.")
+                       "Rows on the device that CUDA, PyTorch and Warp read and write in place.",
+                       nb::is_weak_referenceable())
         .def_prop_ro("name", [](const Column& self) { return self.column->Name(); })
-        .def_prop_ro("shape", [](const Column& self) { return self.column->Shape(); })
+        .def_prop_ro("shape", &Column::Shape)
         .def_prop_ro("dtype", [](const Column& self) { return TypeName(self.column->Type()); })
         .def(
             "__dlpack__",
@@ -234,7 +244,9 @@ void BindDevice(nb::module_& module) {
         .def("assign", &Column::Assign, nb::arg("tensor"), nb::arg("stream") = nb::none(),
              "Takes a DLPack tensor's values: in place for this column's own view, else by one "
              "copy on the device.")
-        .def("read", &Column::Read, "Copies the rows to the host, after CUDA's writes.");
+        .def(
+            "read", [](const Column& self) { return ToNumpy(self.Read()); },
+            "Copies the rows to the host as a NumPy array of its shape, after pending writes.");
 }
 
 }  // namespace ernie::python

ernie_renderer-0.2.0.dev149/ernie/python/device.h

@@ -0,0 +1,60 @@
+// SPDX-License-Identifier: BSD-3-Clause
+/// @file device.h
+/// @brief Binding the device to Python: columns any DLPack library shares in place.
+#pragma once
+
+#include "ernie/device/column.h"
+#include "ernie/device/cuda.h"
+#include "ernie/device/queue.h"
+#include "ernie/util/tensor.h"
+
+#include <nanobind/nanobind.h>
+
+#include <cstdint>
+#include <memory>
+#include <optional>
+#include <vector>
+
+namespace ernie::python {
+
+/// @brief A device and the CUDA device paired with it, if any.
+struct Interop {
+    DeviceContext* context = nullptr;
+    CudaInterop* cuda = nullptr;
+
+    /// @throws nanobind::buffer_error without a paired CUDA device.
+    CudaInterop& Cuda() const;
+};
+
+/// @brief A column Python shares through DLPack, and the ticket of the writes, CUDA's or a
+/// render's, Ernie or a reader has yet to wait for.
+struct Column {
+    Interop owner;
+    /// The column, owned by this or by a renderer.
+    std::unique_ptr<DeviceColumn> owned;
+    const DeviceColumn* column = nullptr;
+    /// The tensor a wrapped column shares, which `__dlpack__` hands back as it is.
+    std::optional<Tensor> wrapped;
+    /// The shape it shows; empty for the column's own.
+    std::vector<std::int64_t> shape;
+    Ticket written;
+    /// What keeps a column another object owns alive.
+    nanobind::object keeper;
+
+    Column(Interop interop, std::unique_ptr<DeviceColumn> device,
+           std::optional<Tensor> tensor = std::nullopt);
+    /// @brief A view of a column another object owns, shown as `view`.
+    Column(Interop interop, const DeviceColumn& device, std::vector<std::int64_t> view,
+           Ticket ready);
+
+    std::vector<std::int64_t> Shape() const;
+    nanobind::object Dlpack(const nanobind::object& stream, const nanobind::object& maxVersion);
+    void Assign(const nanobind::object& source, const nanobind::object& stream);
+    /// @return The rows copied to the host after the pending writes.
+    Tensor Read() const;
+};
+
+/// @brief Adds `DeviceContext` and `DeviceColumn` to the module.
+void BindDevice(nanobind::module_& module);
+
+}  // namespace ernie::python

{ernie_renderer-0.2.0.dev145 → ernie_renderer-0.2.0.dev149}/ernie/python/dlpack.cpp

@@ -3,12 +3,14 @@
 /// @brief DLPack capsules, versioned and not, per the DLPack Python specification.
 #include "ernie/python/dlpack.h"
 
+#include <nanobind/ndarray.h>
 #include <nanobind/stl/tuple.h>
 
 #include <Python.h>
 
 #include <cstdint>
 #include <tuple>
+#include <vector>
 
 namespace nb = nanobind;
 
@@ -95,4 +97,16 @@ Tensor TakeTensor(const nb::object& source, const nb::object& given) {
     return FromCapsule(capsule);
 }
 
+nb::object ToNumpy(Tensor tensor) {
+    DLManagedTensorVersioned* managed = tensor.ToDlpack();
+    const DLDataType type = managed->dl_tensor.dtype;
+    managed->deleter(managed);
+    auto* owned = new Tensor(tensor.Contiguous());
+    const nb::capsule owner(owned, [](void* held) noexcept { delete static_cast<Tensor*>(held); });
+    std::vector<std::size_t> shape(owned->Shape().begin(), owned->Shape().end());
+    return nb::cast(nb::ndarray<nb::numpy>(owned->Data(), shape.size(), shape.data(), owner,
+                                           nullptr,
+                                           nb::dlpack::dtype{type.code, type.bits, type.lanes}));
+}
+
 }  // namespace ernie::python

{ernie_renderer-0.2.0.dev145 → ernie_renderer-0.2.0.dev149}/ernie/python/dlpack.h

@@ -22,4 +22,7 @@ nanobind::object Capsule(DLManagedTensorVersioned* managed, bool versioned);
 /// @return The tensor, sharing the producer's memory.
 Tensor TakeTensor(const nanobind::object& source, const nanobind::object& stream);
 
+/// @return A host tensor as a NumPy array that owns it.
+nanobind::object ToNumpy(Tensor tensor);
+
 }  // namespace ernie::python

ernie_renderer-0.2.0.dev149/ernie/python/render.cpp

@@ -0,0 +1,160 @@
+// SPDX-License-Identifier: BSD-3-Clause
+/// @file render.cpp
+/// @brief Python bindings of the renderer: a scene's products rendered to NumPy images, or stepped
+/// into device columns PyTorch and Warp read in place.
+#include "ernie/python/render.h"
+
+#include "ernie/device/context.h"
+#include "ernie/python/device.h"
+#include "ernie/python/dlpack.h"
+#include "ernie/render/renderer.h"
+
+#include <nanobind/stl/string.h>
+#include <nanobind/stl/tuple.h>
+#include <nanobind/stl/vector.h>
+
+#include <algorithm>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+namespace nb = nanobind;
+
+namespace ernie::python {
+namespace {
+
+/// A renderer, the CUDA device paired with its device, and the columns bound to it, whose writes
+/// each step waits for. Members go in reverse: CUDA first.
+struct PythonRenderer {
+    Renderer renderer;
+    std::unique_ptr<CudaInterop> cuda;
+    std::vector<nb::weakref> bound;
+
+    explicit PythonRenderer(const RendererOptions& options)
+        : renderer(options), cuda(CudaInterop::Create(renderer.Context())) {}
+
+    Interop Pair() {
+        return {&renderer.Context(), cuda.get()};
+    }
+
+    /// Makes the next step wait for the bound columns' writes, forgetting columns gone.
+    void WaitForBound() {
+        std::erase_if(bound, [&](const nb::weakref& reference) {
+            const nb::object column = reference();
+            if (column.is_none()) {
+                return true;
+            }
+            renderer.WaitFor(nb::cast<const Column&>(column).written);
+            return false;
+        });
+    }
+};
+
+/// A batch of a step: its products and each var's columns, valid until the next step.
+struct PythonBatch {
+    std::vector<std::string> products;
+    std::tuple<std::uint32_t, std::uint32_t> size;
+    std::vector<std::pair<std::string, nb::object>> vars;
+
+    nb::object Get(const std::string& name) const {
+        for (const auto& [var, column] : vars) {
+            if (var == name) {
+                return column;
+            }
+        }
+        throw nb::key_error(("The batch holds no '" + name + "'").c_str());
+    }
+};
+
+std::vector<PythonBatch> Step(PythonRenderer& self, const nb::object& owner, Scene& scene) {
+    self.WaitForBound();
+    const std::span<const RenderBatch> batches = self.renderer.Step(scene);
+    const Interop interop = self.Pair();
+    const Ticket ready = self.renderer.Context().Main().Last();
+    std::vector<PythonBatch> result;
+    for (const RenderBatch& batch : batches) {
+        PythonBatch& next = result.emplace_back();
+        next.size = {batch.size.x, batch.size.y};
+        for (const Path& product : batch.products) {
+            next.products.push_back(product.String());
+        }
+        for (const RenderVarOutput& var : batch.vars) {
+            auto* column = new Column(interop, *var.column, var.shape, ready);
+            column->keeper = owner;
+            next.vars.emplace_back(var.name, nb::cast(column, nb::rv_policy::take_ownership));
+        }
+    }
+    return result;
+}
+
+}  // namespace
+
+void BindRender(nb::module_& module) {
+    nb::class_<PythonBatch>(module, "RenderBatch",
+                            "Products of one size and vars rendered in one dispatch: each var a "
+                            "DeviceColumn of shape (views, height, width, channels...), a view a "
+                            "product, valid until the renderer's next step.")
+        .def_ro("products", &PythonBatch::products)
+        .def_ro("size", &PythonBatch::size, "(width, height) in pixels.")
+        .def("__getitem__", &PythonBatch::Get, nb::arg("name"))
+        .def("__contains__",
+             [](const PythonBatch& self, const std::string& name) {
+                 return std::ranges::any_of(self.vars,
+                                            [&](const auto& var) { return var.first == name; });
+             })
+        .def("keys", [](const PythonBatch& self) {
+            std::vector<std::string> names;
+            for (const auto& var : self.vars) {
+                names.push_back(var.first);
+            }
+            return names;
+        });
+
+    nb::class_<PythonRenderer>(module, "Renderer",
+                               "A device and what it keeps of the scene it last rendered, so "
+                               "rendering the same scene again costs only what changed.")
+        .def(
+            "__init__",
+            [](PythonRenderer* self, const std::string& device, bool validation) {
+                new (self)
+                    PythonRenderer(RendererOptions{.device = device, .validation = validation});
+            },
+            nb::kw_only(), nb::arg("device") = "", nb::arg("validation") = false,
+            "Opens a device that ray traces: by UUID or part of its name, or the best one.")
+        .def(
+            "render",
+            [](PythonRenderer& self, Scene& scene) { return ToNumpy(self.renderer.Render(scene)); },
+            nb::arg("scene"),
+            "Renders the scene's first product: LdrColor as a (height, width, 4) uint8 NumPy "
+            "array.")
+        .def(
+            "step",
+            [](nb::object self, Scene& scene) {
+                return Step(nb::cast<PythonRenderer&>(self), self, scene);
+            },
+            nb::arg("scene"),
+            "Renders every product of the scene's first RenderSettings, products alike in size "
+            "and vars in one dispatch, and returns the RenderBatches without waiting: their vars "
+            "are device columns that torch.from_dlpack views in place, or read() copies.")
+        .def(
+            "bind_transform_op",
+            [](nb::object self, Scene& scene, const std::string& prim,
+               const std::string& attribute) {
+                auto& renderer = nb::cast<PythonRenderer&>(self);
+                const DeviceColumn& bound =
+                    renderer.renderer.BindTransformOp(scene, Path(prim), attribute);
+                auto* column = new Column(renderer.Pair(), bound, {}, Ticket{});
+                column->keeper = self;
+                nb::object handle = nb::cast(column, nb::rv_policy::take_ownership);
+                renderer.bound.emplace_back(handle);
+                return handle;
+            },
+            nb::arg("scene"), nb::arg("prim"), nb::arg("attribute"),
+            "Hands a transform op (such as xformOp:translate) of every prim in the prim's bucket "
+            "to the device: returns its DeviceColumn, a row a prim in the bucket's order, which "
+            "PyTorch or Warp write in place; call ready(stream) after writing, and the next step "
+            "renders the new poses with no copy through the host.");
+}
+
+}  // namespace ernie::python

{ernie_renderer-0.2.0.dev145 → ernie_renderer-0.2.0.dev149}/ernie/python/scene.cpp

@@ -327,6 +327,20 @@ void BindScene(nb::module_& module) {
                 return Map(self.Prims(type), &Path::String);
             },
             nb::arg("type") = "", "The paths of every prim of a type or one deriving from it.")
+        .def(
+            "bucket_paths",
+            [](const Scene& self, const Target& prim) {
+                const PrimStore& store = self.Store();
+                const Bucket& bucket = store.GetBucket(store.PrimBucket(Resolve(self, prim)));
+                std::vector<std::string> paths;
+                for (const PrimId row : bucket.prims) {
+                    paths.push_back(self.PrimPath(row).String());
+                }
+                return paths;
+            },
+            nb::arg("prim"),
+            "The paths of the prims stored with this one (its type and applied schemas), by row: "
+            "the rows of a column bound for them. Adding or removing such prims reorders them.")
         .def(
             "apply_api",
             [](Scene& self, const Target& prim, const std::string& api) {

{ernie_renderer-0.2.0.dev145 → ernie_renderer-0.2.0.dev149}/ernie/render/CMakeLists.txt

@@ -6,6 +6,7 @@ add_library(
     geometry.cpp
     instances.cpp
     prim_rows.cpp
+    products.cpp
     renderer.cpp
     rt/bottom_levels.cpp
     rt/camera_pass.cpp

{ernie_renderer-0.2.0.dev145 → ernie_renderer-0.2.0.dev149}/ernie/render/geometry.cpp

@@ -120,15 +120,16 @@ glm::dmat4 ShapeTransform(const Scene& scene, PrimId prim, GeometryKind kind) {
 }
 
 /// What a changed attribute does to a gprim: reshapes it, moves a mesh's points, remakes its
-/// topology, or recolours it.
+/// topology, recolours it, or moves it to another world.
 struct Touched {
     bool shape = false;
     bool points = false;
     bool topology = false;
     bool colour = false;
+    bool world = false;
 
     bool Any() const {
-        return shape || points || topology || colour;
+        return shape || points || topology || colour || world;
     }
 };
 
@@ -136,6 +137,10 @@ Touched ClassifyAttribute(std::string_view attribute) {
     if (attribute == Gprim::kDisplayColor) {
         return {.colour = true};
     }
+    // Inherited: an ancestor's change reaches the gprim as its resolved value's.
+    if (attribute == ResolvedName(Imageable::kWorld)) {
+        return {.world = true};
+    }
     if (attribute == Mesh::kPoints) {
         return {.points = true};
     }
@@ -302,12 +307,13 @@ void GeometryPool::Place(PrimId prim, bool points, bool topology, Changes& chang
     }
     next.shape = ShapeTransform(scene_, prim, *kind);
     next.colour = DisplayColour(scene_, prim);
+    next.world = ReadWorld(scene_, prim);
     if (!draws(next.geometry)) {
         if (drew) {
             changes.removed.push_back(prim);
         }
     } else if (!drew || placed->geometry != next.geometry || placed->shape != next.shape ||
-               placed->colour != next.colour) {
+               placed->colour != next.colour || placed->world != next.world) {
         changes.placed.push_back(next);
     }
     placed = next;

{ernie_renderer-0.2.0.dev145 → ernie_renderer-0.2.0.dev149}/ernie/render/geometry.h

@@ -61,6 +61,8 @@ struct PrimGeometry {
     glm::dmat4 shape{1.0};
     /// Its albedo in linear RGB, from `primvars:displayColor`, until materials bind one.
     glm::vec3 colour{0.5f};
+    /// The batch world it is in, from its resolved `ernie:world`; -1 for every world.
+    std::int32_t world = -1;
 };
 
 /// @brief The scene's gprims as device geometry. Implicit shapes share one unit geometry a kind

{ernie_renderer-0.2.0.dev145 → ernie_renderer-0.2.0.dev149}/ernie/render/instances.cpp

@@ -24,6 +24,7 @@ InstanceTable::InstanceTable(DeviceContext& context, Arenas& arenas)
     geometries_ = column("instance geometries", DType::UInt32, {});
     shapes_ = column("instance shapes", DType::Float32, {3, 4});
     colours_ = column("instance colours", DType::Float32, {3});
+    sources_ = column("top-level sources", DType::UInt32, {});
 }
 
 std::optional<std::uint32_t> InstanceTable::Row(PrimId prim) const {
@@ -40,6 +41,8 @@ void InstanceTable::Set(std::uint32_t row, const PrimGeometry& placed) {
     hostGeometries_[row] = placed.geometry;
     hostShapes_[row] = ToAffineRows(placed.shape);
     hostColours_[row] = placed.colour;
+    relayout_ = relayout_ || hostWorlds_[row] != placed.world;
+    hostWorlds_[row] = placed.world;
     rows_[placed.prim.slot] = row;
     dirty_.push_back(row);
 }
@@ -58,6 +61,34 @@ void InstanceTable::Remove(PrimId prim) {
     }
     rows_[prim.slot] = kNoRow;
     ForEachHostColumn([](auto& column) { column.pop_back(); });
+    relayout_ = true;
+}
+
+void InstanceTable::LayOutWorlds() {
+    std::int32_t highest = 0;
+    std::vector<std::uint32_t> shared;
+    for (std::uint32_t row = 0; row < hostWorlds_.size(); ++row) {
+        highest = std::max(highest, hostWorlds_[row]);
+        if (hostWorlds_[row] < 0) {
+            shared.push_back(row);
+        }
+    }
+    const std::size_t count =
+        std::max<std::size_t>(static_cast<std::size_t>(highest) + 1, reserved_);
+    std::vector<std::vector<std::uint32_t>> own(count);
+    for (std::uint32_t row = 0; row < hostWorlds_.size(); ++row) {
+        if (hostWorlds_[row] >= 0) {
+            own[static_cast<std::size_t>(hostWorlds_[row])].push_back(row);
+        }
+    }
+    worlds_.assign(count, {});
+    hostSources_.clear();
+    for (std::size_t world = 0; world < count; ++world) {
+        worlds_[world] = {.first = static_cast<std::uint32_t>(hostSources_.size()),
+                          .count = static_cast<std::uint32_t>(shared.size() + own[world].size())};
+        hostSources_.insert(hostSources_.end(), shared.begin(), shared.end());
+        hostSources_.insert(hostSources_.end(), own[world].begin(), own[world].end());
+    }
 }
 
 void InstanceTable::Apply(const GeometryPool::Changes& changes) {
@@ -75,6 +106,7 @@ void InstanceTable::Apply(const GeometryPool::Changes& changes) {
             continue;
         }
         ForEachHostColumn([](auto& column) { column.emplace_back(); });
+        relayout_ = true;
         Set(static_cast<std::uint32_t>(prims_.size() - 1), placed);
     }
 }
@@ -89,8 +121,15 @@ void InstanceTable::AddUploads(graph::Graph& graph) {
     StageRows(uploader_, *shapes_, hostShapes_, dirty_);
     StageRows(uploader_, *colours_, hostColours_, dirty_);
     dirty_.clear();
+    if (relayout_) {
+        LayOutWorlds();
+        sources_->Resize(hostSources_.size(), context_.Main());
+        uploader_.Write(sources_->Storage(), sources_->Range().offset,
+                        std::span<const std::uint32_t>(hostSources_));
+        relayout_ = false;
+    }
     std::array columns{graph.Import(*nodes_), graph.Import(*geometries_), graph.Import(*shapes_),
-                       graph.Import(*colours_)};
+                       graph.Import(*colours_), graph.Import(*sources_)};
     graph::AddUploadPass(graph, "Upload instances", uploader_, columns);
 }