refactor: Branchless version of tri using _ggml_vec_tri_cmp as a mask

gabe-l-hart · gabe-l-hart · commit 7690808999fa · 2025-12-03T21:21:12.000-07:00
Branch: ggml-cumsum-tri

Signed-off-by: Gabe Goodhart &lt;ghart@us.ibm.com&gt;
diff --git a/ggml/src/ggml-metal/ggml-metal.metal b/ggml/src/ggml-metal/ggml-metal.metal
@@ -2038,7 +2038,8 @@ kernel void kernel_tri(
     // threadgroup, so this will loop once for each index that this thread is
     // responsible for
     for (int64_t i0 = tpitg.x; i0 < args.ne00; i0 += ntg.x) {
-        dst_row[i0] = _ggml_vec_tri_cmp<ttype>(i0, i1) ? src_row[i0] : static_cast<T>(0.f);
+        // Use the comparison as a mask for branchless
+        dst_row[i0] = static_cast<T>(_ggml_vec_tri_cmp<ttype>(i0, i1)) * src_row[i0];
     }
 }
 

Original file line number	Diff line number	Diff line change
`@@ -2038,7 +2038,8 @@ kernel void kernel_tri(`
`2038`	`2038`	`// threadgroup, so this will loop once for each index that this thread is`
`2039`	`2039`	`// responsible for`
`2040`	`2040`	`for (int64_t i0 = tpitg.x; i0 < args.ne00; i0 += ntg.x) {`
`2041`		`- dst_row[i0] = _ggml_vec_tri_cmp<ttype>(i0, i1) ? src_row[i0] : static_cast<T>(0.f);`
	`2041`	`+ // Use the comparison as a mask for branchless`
	`2042`	`+ dst_row[i0] = static_cast<T>(_ggml_vec_tri_cmp<ttype>(i0, i1)) * src_row[i0];`
`2042`	`2043`	`}`
`2043`	`2044`	`}`
`2044`	`2045`