bug (or my mistake)

[FotisK]

Running it on a M1 iMac with 16GB memory, I get the following issue both with metal and when using SKIP_METAL=1 bash setup.sh. Does this happen to me only?

Sampling sparse structure: 100%|████████████████| 12/12 [06:23<00:00, 31.94s/it]
Sampling shape SLat: 0%| | 0/12 [00:04<?, ?it/s]
Traceback (most recent call last):
File "/Users/user/Desktop/Programming/trellis-mac/generate.py", line 234, in
main()
File "/Users/user/Desktop/Programming/trellis-mac/generate.py", line 88, in main
outputs = pipeline.run(
File "/Users/user/Desktop/Programming/trellis-mac/.venv/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 124, in decorate_context
return func(*args, **kwargs)
File "/Users/user/Desktop/Programming/trellis-mac/TRELLIS.2/trellis2/pipelines/trellis2_image_to_3d.py", line 547, in run
shape_slat = self.sample_shape_slat(
File "/Users/user/Desktop/Programming/trellis-mac/TRELLIS.2/trellis2/pipelines/trellis2_image_to_3d.py", line 260, in sample_shape_slat
slat = self.shape_slat_sampler.sample(
File "/Users/user/Desktop/Programming/trellis-mac/.venv/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 124, in decorate_context
return func(*args, **kwargs)
File "/Users/user/Desktop/Programming/trellis-mac/TRELLIS.2/trellis2/pipelines/samplers/flow_euler.py", line 208, in sample
return super().sample(model, noise, cond, steps, rescale_t, verbose, neg_cond=neg_cond, guidance_strength=guidance_strength, guidance_interval=guidance_interval, **kwargs)
File "/Users/user/Desktop/Programming/trellis-mac/.venv/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 124, in decorate_context
return func(*args, **kwargs)
File "/Users/user/Desktop/Programming/trellis-mac/TRELLIS.2/trellis2/pipelines/samplers/flow_euler.py", line 121, in sample
out = self.sample_once(model, sample, t, t_prev, cond, **kwargs)
File "/Users/user/Desktop/Programming/trellis-mac/.venv/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 124, in decorate_context
return func(*args, **kwargs)
File "/Users/user/Desktop/Programming/trellis-mac/TRELLIS.2/trellis2/pipelines/samplers/flow_euler.py", line 79, in sample_once
pred_x_0, pred_eps, pred_v = self._get_model_prediction(model, x_t, t, cond, **kwargs)
File "/Users/user/Desktop/Programming/trellis-mac/TRELLIS.2/trellis2/pipelines/samplers/flow_euler.py", line 49, in _get_model_prediction
pred_v = self._inference_model(model, x_t, t, cond, **kwargs)
File "/Users/user/Desktop/Programming/trellis-mac/TRELLIS.2/trellis2/pipelines/samplers/guidance_interval_mixin.py", line 11, in _inference_model
return super()._inference_model(model, x_t, t, cond, guidance_strength=guidance_strength, **kwargs)
File "/Users/user/Desktop/Programming/trellis-mac/TRELLIS.2/trellis2/pipelines/samplers/classifier_free_guidance_mixin.py", line 23, in _inference_model
std_pos = x_0_pos.std(dim=list(range(1, x_0_pos.ndim)), keepdim=True)
File "/Users/user/Desktop/Programming/trellis-mac/TRELLIS.2/trellis2/modules/sparse/basic.py", line 296, in std
mean = self.mean(dim=dim, keepdim=True)
File "/Users/user/Desktop/Programming/trellis-mac/TRELLIS.2/trellis2/modules/sparse/basic.py", line 287, in mean
return self.reduce(op='mean', dim=dim, keepdim=keepdim)
File "/Users/user/Desktop/Programming/trellis-mac/TRELLIS.2/trellis2/modules/sparse/basic.py", line 283, in reduce
red = torch.segment_reduce(red, reduce=op, lengths=self.seqlen)
NotImplementedError: The operator 'aten::segment_reduce' is not currently implemented for the MPS device. If you want this op to be considered for addition please comment on pytorch/pytorch#141287 and mention use-case, that resulted in missing op as well as commit hash 70d99e998b4955e0049d13a98d77ae1b14db1f45. As a temporary fix, you can set the environment variable PYTORCH_ENABLE_MPS_FALLBACK=1 to use the CPU as a fallback for this op. WARNING: this will be slower than running natively on MPS.

[FotisK]

running it with 'PYTORCH_ENABLE_MPS_FALLBACK=1' goes a bit further, but also hits a wall, this time at simplification of the mesh(?)

Generating 3D model (pipeline=512, seed=42)...
Sampling sparse structure: 100%|████████████████| 12/12 [06:18<00:00, 31.56s/it]
Sampling shape SLat: 0%| | 0/12 [00:00<?, ?it/s]/Users/user/Desktop/Programming/trellis-mac/TRELLIS.2/trellis2/modules/sparse/basic.py:283: UserWarning: The operator 'aten::segment_reduce' is not currently supported on the MPS backend and will fall back to run on the CPU. This may have performance implications. (Triggered internally at /Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/mps/MPSFallback.mm:34.)
red = torch.segment_reduce(red, reduce=op, lengths=self.seqlen)
Sampling shape SLat: 100%|██████████████████████| 12/12 [00:57<00:00, 4.77s/it]
Sampling texture SLat: 100%|████████████████████| 12/12 [00:31<00:00, 2.66s/it]

Mesh: 233,197 vertices, 477,252 triangles
Generation time: 553.4s

Baking PBR textures via Metal (512x512)...
Simplifying mesh: 477,252 -> ~200,000 faces
Original mesh: 94763 vertices, 200000 faces
After filling holes: 95159 vertices, 201421 faces
Building BVH for current mesh...Done
Cleaning mesh...
After inital simplification: 95159 vertices, 201421 faces
After initial cleanup: 112795 vertices, 203096 faces
Simplifying [thres=1.00e-08]: 0%| | 0/3096 [00:00<?, ?it/s]Traceback (most recent call last):
File "/Users/user/Desktop/Programming/trellis-mac/generate.py", line 234, in
main()
File "/Users/user/Desktop/Programming/trellis-mac/generate.py", line 172, in main
glb = o_voxel.postprocess.to_glb(
File "/Users/user/Desktop/Programming/trellis-mac/.venv/lib/python3.10/site-packages/o_voxel/postprocess.py", line 230, in to_glb
mesh.simplify(decimation_target, verbose=verbose)
File "/Users/user/Desktop/Programming/trellis-mac/.venv/lib/python3.10/site-packages/cumesh/metal_backend.py", line 211, in simplify
new_num_vert, new_num_face = self.simplify_step(
File "/Users/user/Desktop/Programming/trellis-mac/.venv/lib/python3.10/site-packages/cumesh/metal_backend.py", line 190, in simplify_step
return self._mesh.simplify_step(lambda_edge_length, lambda_skinny, threshold, timing)
RuntimeError: [MtlMesh] Failed to create pipeline for 'propagate_cost_kernel': Unsupported float atomic operation for given target.Unsupported float atomic operation for given target.
Simplifying [thres=1.00e-08]: 0%| | 0/3096 [00:00<?, ?it/s]
PYTORCH_ENABLE_MPS_FALLBACK=1 python generate.py ~/Downloads/totem1.png 512 144.33s user 44.20s system 25% cpu 12:13.13 total

[FotisK]

In the 'PYTORCH_ENABLE_MPS_FALLBACK=1' case, I disabled textures altogether ('--no-texture') and it produced the .obj and .glb successfully! good job

[shivampkumar]

Hey, sorry for the delay in reply. I'm glad it got figured out!

Thanks for the detailed report! Two separate issues here:

1. segment_reduce not on MPS: This was a bug - PYTORCH_ENABLE_MPS_FALLBACK was being set after
   flex_gemm already imported torch, so the fallback didn't activate. Just pushed a fix (ede0bdd) - pull latest
   and it should work without needing to set the env var manually.

2. Metal cumesh crash on M1 16GB: The simplify_step error ("Unsupported float atomic operation") is a
   Metal GPU family limitation - M1 (Apple7 GPU family) doesn't support float atomic operations that were added in M2 and later (Apple8+). The Metal cumesh port from pedronaugusto assumes newer GPU features.

Workaround for now: --no-texture as you found, or SKIP_METAL=1 bash setup.sh to use the pure-Python KDTree
texture baker instead. I'll add a try/catch so it falls back automatically instead of crashing.

Glad geometry generation works on M1 16GB tho, that's awesome!

[FotisK]

thanks for looking into it! I've run a few tests (without textures), and they are all very impressive! Next days I'll come back and pull the updated version.

A remark about the second/texture issue: my recent tests were run after having run 'SKIP_METAL=1 bash setup.sh' already (I never run setup.sh after) and apparently the KDTree(?) texture-baker failed. (the output above is actually from this run)

[shivampkumar]

Ahh, didn't realize that, sorry about that.

I think the issue is that SKIP_METAL=1 only skips installing the Metal packages - if they were already
installed from a previous setup run, generate.py still detects and tries to use them. Then cumesh crashes on
M1 because of the float atomic limitation.

Just pushed a quick one liner (8288a2c) that wraps the Metal bake in a try/catch so it falls back to the KDTree texture
baker if cumesh crashes, instead of dying.

Could you pull latest and try again with textures enabled? If my diagnosis is right it should fall back
cleanly now. If not, let me know what you see and I'll dig deeper.

Sorry about the earlier wrong diagnosis.

[FotisK]

No worries! Thank you for all the good work, I've long been hoping that I could run trellis2 on my mac!
So I pulled the update and it seems to be gracefully falling back to KDTree now, though I'm not sure if it works or not - It's been a couple hours that it's stuck in the below step of Unwrapping with xatlas, but I have no idea, maybe this is just normal for xatlas. I will leave it running and check back later when I'm home.
All the while in Activity Monitor, python is consistently maxing out 2 cores (199% CPU usage), ram utilization is 16GB + 21GB swap which I think is slightly higher compared to during my tests without the textures (I think swap never exceeded 16GB in those), and there is not any disk activity happening.

time PYTORCH_ENABLE_MPS_FALLBACK=1 python generate.py ~/Downloads/person.png --pipeline-type 512 --texture-size 512
============================================================
TRELLIS.2 on Apple Silicon
============================================================

Loading pipeline...
[SPARSE] Conv backend: flex_gemm; Attention backend: sdpa
[ATTENTION] Using backend: sdpa
Loading weights: 100%|██████████████████████| 415/415 [00:00<00:00, 4994.95it/s]
/Users/user/Desktop/Programming/trellis-mac/.venv/lib/python3.10/site-packages/timm/models/layers/__init__.py:49: FutureWarning: Importing from timm.models.layers is deprecated, please import via timm.layers
  warnings.warn(f"Importing from {__name__} is deprecated, please import via timm.layers", FutureWarning)
/Users/user/Desktop/Programming/trellis-mac/.venv/lib/python3.10/site-packages/timm/models/registry.py:4: FutureWarning: Importing from timm.models.registry is deprecated, please import via timm.models
  warnings.warn(f"Importing from {__name__} is deprecated, please import via timm.models", FutureWarning)
Loading weights: 100%|██████████████████████| 754/754 [00:00<00:00, 6516.64it/s]
Loaded in 162s
Device: MPS
Input: /Users/user/Downloads/person.png (1086x1448)

Generating 3D model (pipeline=512, seed=42)...
Sampling sparse structure: 100%|████████████████| 12/12 [06:37<00:00, 33.16s/it]
Sampling shape SLat:   0%|                               | 0/12 [00:00<?, ?it/s]/Users/user/Desktop/Programming/trellis-mac/TRELLIS.2/trellis2/modules/sparse/basic.py:283: UserWarning: The operator 'aten::segment_reduce' is not currently supported on the MPS backend and will fall back to run on the CPU. This may have performance implications. (Triggered internally at /Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/mps/MPSFallback.mm:34.)
  red = torch.segment_reduce(red, reduce=op, lengths=self.seqlen)
Sampling shape SLat: 100%|██████████████████████| 12/12 [12:46<00:00, 63.86s/it]
Sampling texture SLat: 100%|████████████████████| 12/12 [02:07<00:00, 10.64s/it]

Mesh: 809,968 vertices, 1,681,124 triangles
Generation time: 1412.4s

Baking PBR textures via Metal (512x512)...
  Simplifying mesh: 1,681,124 -> ~200,000 faces
Original mesh: 68922 vertices, 199998 faces
After filling holes: 69764 vertices, 202946 faces
Building BVH for current mesh...Done
Cleaning mesh...
After inital simplification: 69764 vertices, 202946 faces
After initial cleanup: 193457 vertices, 239613 faces
Simplifying [thres=1.00e-08]:   0%|                   | 0/39613 [00:00<?, ?it/s]
  Metal bake failed: [MtlMesh] Failed to create pipeline for 'propagate_cost_kernel': Unsupported float atomic operation for given target.Unsupported float atomic operation for given target.
  Falling back to KDTree texture baker...
Simplifying [thres=1.00e-08]:   0%|                   | 0/39613 [00:00<?, ?it/s]

Baking PBR textures via KDTree (512x512)...
  UV unwrapping with xatlas...

[shivampkumar]

Okay so the Metal fallback is working as intended. It caught the cumesh crash and fell to the KDTree baker.

The hang is xatlas trying to UV-unwrap your full 809K vertex / 1.68M triangle mesh. That's extremely slow, especially with 21GB swap on 16GB RAM. xatlas is CPU-only and single-threaded.

Just pushed a fix (d53f2f0): the KDTree path now simplifies the mesh to ~200K faces before UV unwrapping (same as the Metal path does). Should be much faster.

Pull latest and try again: would love to know if it completes on your M1.

[FotisK]

I left my computer to run another 24h but when i got back xatlas was still stuck at the same -- python CPU ~200%, real memory: 160MBs virtual memory: 412,4GB (swap used 21,5GB). I will try with the new update and get back to you!

[FotisK]

So, an update: After killing the xatlas process of the previous version (full-mesh res), I pulled the update, and run it again; this time xatlas completed in 2h (similar memory and CPU footprint) on a simplified mesh as seen below! The original photo is a portrait photo. However the UV map generated is garbled.

Baking PBR textures via KDTree (512x512)... Simplifying mesh: 1,681,124 -> ~200,000 faces UV unwrapping with xatlas... UV unwrap: 809,968 -> 350,667 vertices Voxels: 809,968, channels: 6 Building KDTree... Tree built in 0.2s Rasterizing 199,998 triangles into 512x512... Rasterizing: 100,000/199,998 Coverage: 37.9%, rasterized in 4.8s Querying 99,462 texels, k=8... Query done in 0.1s Weighting colors (max_dist = 0.0039)... Final coverage: 100.0%, total bake: 5.4s Saved: output_3d.glb Bake time: 6476s Saved: output_3d.obj

[shivampkumar]

Got it. Thanks for the screenshots, those tell the real story.

The garbled texture isn't a data corruption issue, it's a hole-fill artifact in the KDTree baker. Here's
what's happening I guess

1. xatlas on your 200K-face simplified mesh produces lots of small UV charts because the bust has many sharp
   features (hair, beard, hood folds = many seams).
2. Rasterizing those charts only fills ~38% of the 512×512 texture (the "Coverage: 37.9%" line).
3. The chart interiors get correct colors from KDTree.
4. The remaining 62% is filled by 8 iterations of dilation + 3×3 blur over the sparse channel — that's what
   produces the dim noise pattern in the rest of the texture. The blur of mostly-zero pixels gives dim averaged
   values that bleed outward from each tiny chart.

You can see this in your texture: the top-left coherent patches are the few bigger charts; everything else is
dilation bleed.

Two things we can try, in order of likely impact:

1. Replace the dilation+blur hole-fill with proper nearest-texel inpainting. For each unfilled texel, find
   the nearest filled texel and use its 3D position to query the KDTree directly. No bleed, no blur, hard chart
   boundaries.
2. Lower the default simplify target for the KDTree path (e.g. 50K faces instead of 200K). Fewer features
   → fewer xatlas charts → larger islands → less hole-fill needed. Tradeoff: slightly rougher silhouette.

It would be awesome if you can try these on your machine (Use Claude or smthg if not your domain), and see if anything helps. Then we can make a PR out of what works.

Also, even after my fixes, the KDTree path is going to be slow on 16GB simply because the full
attribute volume + KDTree don't fit comfortably. Worth trying --texture-size 256 to see if smaller textures
alleviate it.

If you can share the input photo, I can try it on my machine as well.

[FotisK]

thanks for delving into it shivampkumar! I run 10-20 tests, they all invariably suffer from the same issue. The link for a few of them (including the photo & model I posted above) can be found here .

I'll give it a shot with the help of codex (not really my domain but I feel lucky), try to see what makes it work and get back to you!

[FotisK]

small update, I managed (correction: codex did) to make metal texture baking work on M1 aswell (and hole-filling too, though let's see how this goes for big meshes), though as can be seen on the screenshot, it is still far from the results that trellis2.0 demo on HF is giving. So now I'm focusing on why this might be the case!
Didn't get into the xatlas/kdtree thing since it was too slow anyway - metal work completes in under 20seconds at 2048px UV maps, where xatlas/kdtree were completing in 2h for 512px and for 1024px leaving it run for 24h was not enough.

The updates are all pushed together in a local codex managed repository (my git skills are close to zero), so I'll do some untangling towards the end of the week and bring them here

[FotisK]

So this is the M1 patch so that it doesn't fallback to the KDTree baker: This patch adds Apple-GPU compatibility fallbacks to mtlmesh so the Metal bake path keeps working on older Apple Silicon targets (notably M1-class devices) instead of failing and forcing fallback to the much slower CPU/KDTree baker.

Parity notes:

• Relative to the existing Metal implementation on newer Apple GPU families:
  • Feature parity: yes
  • Algorithmic parity: mostly yes
  • Bitwise parity: no guarantee
• Relative to the original CUDA implementation:
  • Feature parity: yes
  • Algorithmic parity: mostly yes
  • Bitwise parity: no guarantee

More specifically:

• The atlas perimeter fallback replaces native float atomics with a CAS-loop float accumulation over atomic_uint.
  • This preserves the operation and intended behavior, but floating-point atomic accumulation is order-dependent, so exact bitwise parity is not expected versus either newer Metal targets or CUDA.
• The simplify fallback replaces the native propagation kernel with a direct per-face minimum over the union of edges adjacent to the face’s three vertices.
  • This is intended to preserve the same effective reduction logic as the original propagation step much more closely than the atlas case, though exact cross-backend bitwise parity is still not claimed.

diff --git a/src/metal/atlas.metal b/src/metal/atlas.metal
index ddc24d9..f83f368 100644
--- a/src/metal/atlas.metal
+++ b/src/metal/atlas.metal
@@ -8,6 +8,21 @@ inline ulong pack_kv(int key, float value) {
     return (ulong(v) << 32) | ulong(uint(key));
 }
 
+// Float atomics are not reliably available across older Apple GPU targets.
+// Emulate atomic add via a CAS loop on the underlying 32-bit representation.
+inline void atomicAddFloat(device atomic_uint* addr, float value) {
+    uint expected = atomic_load_explicit(addr, memory_order_relaxed);
+    while (true) {
+        float current = as_type<float>(expected);
+        uint desired = as_type<uint>(current + value);
+        if (atomic_compare_exchange_weak_explicit(
+                addr, &expected, desired,
+                memory_order_relaxed, memory_order_relaxed)) {
+            return;
+        }
+    }
+}
+
 // Initialize chart adjacency from face adjacency
 kernel void init_chart_adj_kernel(
     device const packed_float3* vertices [[buffer(0)]],
@@ -41,9 +56,29 @@ kernel void init_chart_adj_kernel(
     } else { length[tid] = 0; }
 }
 
-// Count chart-edge connectivity + accumulate perimeters
-// Uses atomic float add (Metal 3.0, Apple Silicon M1+)
+// Count chart-edge connectivity + accumulate perimeters.
+// Use CAS-based float accumulation instead of native atomic_float so the
+// kernel compiles on older Apple GPU targets such as M1.
 kernel void get_chart_edge_cnt_kernel(
+    device const ulong* chart_adj [[buffer(0)]],
+    device const float* chart_adj_length [[buffer(1)]],
+    constant int& E [[buffer(2)]],
+    device atomic_int* chart2edge_cnt [[buffer(3)]],
+    device atomic_uint* chart_perim [[buffer(4)]],
+    uint tid [[thread_position_in_grid]]
+) {
+    if (tid >= (uint)E) return;
+    ulong c = chart_adj[tid];
+    float l = chart_adj_length[tid];
+    int c0 = int(c >> 32), c1 = int(c & 0xFFFFFFFF);
+    atomic_fetch_add_explicit(&chart2edge_cnt[c0], 1, memory_order_relaxed);
+    atomic_fetch_add_explicit(&chart2edge_cnt[c1], 1, memory_order_relaxed);
+    atomicAddFloat(&chart_perim[c0], l);
+    atomicAddFloat(&chart_perim[c1], l);
+}
+
+// Native-float-atomic variant for newer Apple GPU families.
+kernel void get_chart_edge_cnt_kernel_native(
     device const ulong* chart_adj [[buffer(0)]],
     device const float* chart_adj_length [[buffer(1)]],
     constant int& E [[buffer(2)]],
diff --git a/src/metal/simplify.metal b/src/metal/simplify.metal
index d079eb9..7ae255e 100644
--- a/src/metal/simplify.metal
+++ b/src/metal/simplify.metal
@@ -152,6 +152,39 @@ kernel void propagate_cost_kernel(
     }
 }
 
+// Fallback for older Apple GPU targets: compute the propagated minimum cost
+// per face directly. The original kernel propagates each edge's cost to all
+// faces incident to either endpoint, so the face-local minimum is taken over
+// the union of edges adjacent to the face's three vertices.
+kernel void propagate_cost_per_face_kernel(
+    device const packed_int3* faces [[buffer(0)]],
+    device const int* vert2edge [[buffer(1)]],
+    device const int* vert2edge_offset [[buffer(2)]],
+    device const float* edge_collapse_costs [[buffer(3)]],
+    constant int& F [[buffer(4)]],
+    device ulong* propagated_costs [[buffer(5)]],
+    uint tid [[thread_position_in_grid]]
+) {
+    if (tid >= (uint)F) return;
+    int3 vids = to_int3(faces[tid]);
+    ulong best = ULONG_MAX;
+
+    for (int i = vert2edge_offset[vids.x]; i < vert2edge_offset[vids.x + 1]; i++) {
+        int eid = vert2edge[i];
+        best = min(best, pack_key_value_positive(eid, edge_collapse_costs[eid]));
+    }
+    for (int i = vert2edge_offset[vids.y]; i < vert2edge_offset[vids.y + 1]; i++) {
+        int eid = vert2edge[i];
+        best = min(best, pack_key_value_positive(eid, edge_collapse_costs[eid]));
+    }
+    for (int i = vert2edge_offset[vids.z]; i < vert2edge_offset[vids.z + 1]; i++) {
+        int eid = vert2edge[i];
+        best = min(best, pack_key_value_positive(eid, edge_collapse_costs[eid]));
+    }
+
+    propagated_costs[tid] = best;
+}
+
 // Collapse edges in parallel (conflict-free via propagated cost check)
 kernel void collapse_edges_kernel(
     device packed_float3* vertices [[buffer(0)]],
diff --git a/src/mtlmesh.mm b/src/mtlmesh.mm
index f663db1..b0c484d 100644
--- a/src/mtlmesh.mm
+++ b/src/mtlmesh.mm
@@ -13,6 +13,24 @@ namespace mtlmesh {
 
 template<typename T> static inline T* PTR(id<MTLBuffer> buf) { return (T*)[buf contents]; }
 
+static const char* chart_edge_count_kernel_name(id<MTLDevice> device) {
+    if (@available(macOS 14.0, *)) {
+        if ([device supportsFamily:MTLGPUFamilyApple9]) {
+            return "get_chart_edge_cnt_kernel_native";
+        }
+    }
+    return "get_chart_edge_cnt_kernel";
+}
+
+static const char* simplify_propagate_kernel_name(id<MTLDevice> device) {
+    if (@available(macOS 14.0, *)) {
+        if ([device supportsFamily:MTLGPUFamilyApple9]) {
+            return "propagate_cost_kernel";
+        }
+    }
+    return "propagate_cost_per_face_kernel";
+}
+
 // ========== Buffer helpers ==========
 
 id<MTLBuffer> MtlMesh::alloc(size_t bytes) {
@@ -1200,6 +1218,8 @@ void MtlMesh::unify_face_orientations() {
 std::tuple<int, int> MtlMesh::simplify_step(float lambda_edge_length, float lambda_skinny, float threshold, bool timing) {
     get_vertex_face_adjacency();
     get_edges();
+    get_edge_face_adjacency();
+    get_vertex_edge_adjacency();
     get_boundary_info();
 
     int V = num_verts, F = num_fcs, E = num_edges_;
@@ -1235,14 +1255,26 @@ std::tuple<int, int> MtlMesh::simplify_step(float lambda_edge_length, float lamb
     auto prop_costs = alloc(F * 8);
     memset([prop_costs contents], 0xFF, F * 8);  // UINT64_MAX
 
-    CTX.dispatch("propagate_cost_kernel", [&](id<MTLComputeCommandEncoder> enc) {
-        [enc setBuffer:edges offset:0 atIndex:0];
-        [enc setBuffer:vert2face offset:0 atIndex:1];
-        [enc setBuffer:vert2face_offset offset:0 atIndex:2];
-        [enc setBuffer:costs offset:0 atIndex:3];
-        [enc setBytes:&E length:sizeof(int) atIndex:4];
-        [enc setBuffer:prop_costs offset:0 atIndex:5];
-    }, E);
+    const char* prop_kernel = simplify_propagate_kernel_name(dev_);
+    if (strcmp(prop_kernel, "propagate_cost_kernel") == 0) {
+        CTX.dispatch(prop_kernel, [&](id<MTLComputeCommandEncoder> enc) {
+            [enc setBuffer:edges offset:0 atIndex:0];
+            [enc setBuffer:vert2face offset:0 atIndex:1];
+            [enc setBuffer:vert2face_offset offset:0 atIndex:2];
+            [enc setBuffer:costs offset:0 atIndex:3];
+            [enc setBytes:&E length:sizeof(int) atIndex:4];
+            [enc setBuffer:prop_costs offset:0 atIndex:5];
+        }, E);
+    } else {
+        CTX.dispatch(prop_kernel, [&](id<MTLComputeCommandEncoder> enc) {
+            [enc setBuffer:faces offset:0 atIndex:0];
+            [enc setBuffer:vert2edge offset:0 atIndex:1];
+            [enc setBuffer:vert2edge_offset offset:0 atIndex:2];
+            [enc setBuffer:costs offset:0 atIndex:3];
+            [enc setBytes:&F length:sizeof(int) atIndex:4];
+            [enc setBuffer:prop_costs offset:0 atIndex:5];
+        }, F);
+    }
 
     // Step 4: Collapse edges
     auto vert_kept = alloc(V * 4);
@@ -1483,7 +1515,7 @@ void MtlMesh::compute_charts(float threshold, int refine_iters, int global_iters
         // Build chart-edge CSR adjacency
         auto chart2edge_cnt = alloc_zero(num_charts * 4);
         auto chart_perims = alloc_zero(num_charts * 4);
-        CTX.dispatch("get_chart_edge_cnt_kernel", [&](id<MTLComputeCommandEncoder> enc) {
+        CTX.dispatch(chart_edge_count_kernel_name(dev_), [&](id<MTLComputeCommandEncoder> enc) {
             [enc setBuffer:unique_adj offset:0 atIndex:0];
             [enc setBuffer:agg_lengths offset:0 atIndex:1];
             [enc setBytes:&chart_E length:sizeof(int) atIndex:2];
@@ -1610,7 +1642,7 @@ void MtlMesh::compute_charts(float threshold, int refine_iters, int global_iters
 
             chart2edge_cnt = alloc_zero(num_charts * 4);
             chart_perims = alloc_zero(num_charts * 4);
-            CTX.dispatch("get_chart_edge_cnt_kernel", [&](id<MTLComputeCommandEncoder> enc) {
+            CTX.dispatch(chart_edge_count_kernel_name(dev_), [&](id<MTLComputeCommandEncoder> enc) {
                 [enc setBuffer:unique_adj offset:0 atIndex:0];
                 [enc setBuffer:agg_lengths offset:0 atIndex:1];
                 [enc setBytes:&chart_E length:sizeof(int) atIndex:2];