     }
 
     // If the cell has not been expanded, we will not process it next time
-    cell_flags[id] = (k << axis) | (cell_flags[id] & ~(1 << axis));
+    if (!last_iter) cell_flags[id] = (k << axis) | (cell_flags[id] & ~(1 << axis));
 
     store_cell(new_cells + id, cell);
 }
 
-template <bool subset_only, typename Primitive>
+template <bool last_iter, typename Primitive>
 void expansion_iter(Grid& grid, const Primitive* prims, Cell*& new_cells, int* cell_flags) {
-    overlap_step<0, subset_only><<<round_div(grid.num_cells, 64), 64>>>(grid.entries, grid.ref_ids, prims, grid.cells, new_cells, cell_flags, grid.num_cells);
+    overlap_step<0, last_iter><<<round_div(grid.num_cells, 64), 64>>>(grid.entries, grid.ref_ids, prims, grid.cells, new_cells, cell_flags, grid.num_cells);
     std::swap(new_cells, grid.cells);
-    overlap_step<1, subset_only><<<round_div(grid.num_cells, 64), 64>>>(grid.entries, grid.ref_ids, prims, grid.cells, new_cells, cell_flags, grid.num_cells);
+    overlap_step<1, last_iter><<<round_div(grid.num_cells, 64), 64>>>(grid.entries, grid.ref_ids, prims, grid.cells, new_cells, cell_flags, grid.num_cells);
     std::swap(new_cells, grid.cells);
-    overlap_step<2, subset_only><<<round_div(grid.num_cells, 64), 64>>>(grid.entries, grid.ref_ids, prims, grid.cells, new_cells, cell_flags, grid.num_cells);
+    overlap_step<2, last_iter><<<round_div(grid.num_cells, 64), 64>>>(grid.entries, grid.ref_ids, prims, grid.cells, new_cells, cell_flags, grid.num_cells);
     std::swap(new_cells, grid.cells);
 }
 

static __constant__ ivec3 grid_dims;
static __constant__ vec3  grid_min;
static __constant__ vec3  cell_size;
static __constant__ int   grid_shift;

/// Gets the component of the given vector on the specified axis

    return j == c1;
}

/// Computes the amount of overlap possible for a cell and a given primitive
template <int axis, bool dir, typename Primitive>
__device__ int compute_overlap(const Primitive& prim, const Cell& cell, const BBox& cell_bbox, int d) {
    constexpr int axis1 = (axis + 1) % 3;
    constexpr int axis2 = (axis + 2) % 3;
    auto prim_bbox = prim.bbox();

    if (get<axis1>(prim_bbox.min) <= get<axis1>(cell_bbox.max) &&
        get<axis1>(prim_bbox.max) >= get<axis1>(cell_bbox.min) &&
        get<axis2>(prim_bbox.min) <= get<axis2>(cell_bbox.max) &&
        get<axis2>(prim_bbox.max) >= get<axis2>(cell_bbox.min)) {
        int prim_d = get<axis>(((dir ? prim_bbox.min : prim_bbox.max) - grid_min) * grid_inv);
        d = dir
            ? min(d, prim_d - get<axis>(cell.max))
            : max(d, prim_d - get<axis>(cell.min) + 1);
        d = dir ? max(d, 0) : min(d, 0);
    }
    return d;
}

/// Finds the maximum overlap possible for one cell
template <int axis, bool dir, bool subset_only, typename Primitive>
__device__ int find_overlap(const Entry* __restrict__ entries,
