feat: Move the ttype conditional to templating to avoid conditional in kernel

Branch: ggml-cumsum-tri

Signed-off-by: Gabe Goodhart <ghart@us.ibm.com>

Author: gabe-l-hart

ggml/src/ggml-metal/ggml-metal-device.cpp

@@ -363,8 +363,9 @@ ggml_metal_pipeline_t ggml_metal_library_get_pipeline_tri(ggml_metal_library_t l
     char name[256];
 
     const char * op_str = "tri";
+    const int ttype = op->op_params[0];
 
-    snprintf(base, 256, "kernel_%s_%s", op_str, ggml_type_name(op->src[0]->type));
+    snprintf(base, 256, "kernel_%s_%s_%d", op_str, ggml_type_name(op->src[0]->type), ttype);
 
     snprintf(name, 256, "%s", base);
 

ggml/src/ggml-metal/ggml-metal-impl.h

@@ -848,7 +848,6 @@ typedef struct {
     uint64_t nb1;
     uint64_t nb2;
     uint64_t nb3;
-    uint32_t ttype;
 } ggml_metal_kargs_tri;
 
 typedef struct {

ggml/src/ggml-metal/ggml-metal-ops.cpp

@@ -3914,8 +3914,6 @@ int ggml_metal_op_tri(ggml_metal_op_t ctx, int idx) {
     GGML_TENSOR_LOCALS( int32_t, ne,  op,         ne);
     GGML_TENSOR_LOCALS(uint32_t, nb,  op,         nb);
 
-    const ggml_tri_type ttype = (ggml_tri_type) op->op_params[0];
-
     ggml_metal_kargs_tri args = {
         /*.ne00  =*/ ne00,
         /*.ne01  =*/ ne01,
@@ -3933,7 +3931,6 @@ int ggml_metal_op_tri(ggml_metal_op_t ctx, int idx) {
         /*.nb1   =*/ nb1,
         /*.nb2   =*/ nb2,
         /*.nb3   =*/ nb3,
-        /*.ttype =*/ static_cast<uint32_t>(ttype)
     };
 
     ggml_metal_pipeline_t pipeline = ggml_metal_library_get_pipeline_tri(lib, op);

ggml/src/ggml-metal/ggml-metal.metal

@@ -1943,17 +1943,31 @@ typedef decltype(kernel_cumsum_add<float>) kernel_cumsum_add_t;
 
 template [[host_name("kernel_cumsum_add_f32")]] kernel kernel_cumsum_add_t kernel_cumsum_add<float>;
 
-inline static bool _ggml_vec_tri_cmp(const int i, const int r, const uint32_t type) {
-    switch (type) {
-        // ggml.h:620
-        case /* GGML_TRI_TYPE_LOWER      */ 3: return i < r; break;
-        case /* GGML_TRI_TYPE_LOWER_DIAG */ 2: return i <= r; break;
-        case /* GGML_TRI_TYPE_UPPER      */ 1: return i > r; break;
-        case /* GGML_TRI_TYPE_UPPER_DIAG */ 0: return i >= r; break;
-    }
+
+template<uint32_t ttype>
+bool _ggml_vec_tri_cmp(const int i, const int r);
+
+template<>
+bool _ggml_vec_tri_cmp</* GGML_TRI_TYPE_LOWER */ 3>(const int i, const int r) {
+    return i < r;
 }
 
-template<typename T>
+template<>
+bool _ggml_vec_tri_cmp</* GGML_TRI_TYPE_LOWER_DIAG */ 2>(const int i, const int r) {
+    return i <= r;
+}
+
+template<>
+bool _ggml_vec_tri_cmp</* GGML_TRI_TYPE_UPPER */ 1>(const int i, const int r) {
+    return i > r;
+}
+
+template<>
+bool _ggml_vec_tri_cmp</* GGML_TRI_TYPE_UPPER_DIAG */ 0>(const int i, const int r) {
+    return i >= r;
+}
+
+template<typename T, int ttype>
 kernel void kernel_tri(
         constant ggml_metal_kargs_tri & args,
         device const char * src0,
@@ -1978,16 +1992,25 @@ 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(i0, i1, args.ttype) ? src_row[i0] : static_cast<T>(0.f);
+        dst_row[i0] = _ggml_vec_tri_cmp<ttype>(i0, i1) ? src_row[i0] : static_cast<T>(0.f);
     }
 }
 
-typedef decltype(kernel_tri<float>) kernel_tri_t;
+typedef decltype(kernel_tri<float, 0>) kernel_tri_t;
 
-template [[host_name("kernel_tri_f32")]] kernel kernel_tri_t kernel_tri<float>;
-template [[host_name("kernel_tri_f16")]] kernel kernel_tri_t kernel_tri<half>;
+template [[host_name("kernel_tri_f32_0")]] kernel kernel_tri_t kernel_tri<float, 0>;
+template [[host_name("kernel_tri_f32_1")]] kernel kernel_tri_t kernel_tri<float, 1>;
+template [[host_name("kernel_tri_f32_2")]] kernel kernel_tri_t kernel_tri<float, 2>;
+template [[host_name("kernel_tri_f32_3")]] kernel kernel_tri_t kernel_tri<float, 3>;
+template [[host_name("kernel_tri_f16_0")]] kernel kernel_tri_t kernel_tri<half, 0>;
+template [[host_name("kernel_tri_f16_1")]] kernel kernel_tri_t kernel_tri<half, 1>;
+template [[host_name("kernel_tri_f16_2")]] kernel kernel_tri_t kernel_tri<half, 2>;
+template [[host_name("kernel_tri_f16_3")]] kernel kernel_tri_t kernel_tri<half, 3>;
 #if defined(GGML_METAL_HAS_BF16)
-template [[host_name("kernel_tri_bf16")]] kernel kernel_tri_t kernel_tri<bfloat>;
+template [[host_name("kernel_tri_bf16_0")]] kernel kernel_tri_t kernel_tri<bfloat, 0>;
+template [[host_name("kernel_tri_bf16_1")]] kernel kernel_tri_t kernel_tri<bfloat, 1>;
+template [[host_name("kernel_tri_bf16_2")]] kernel kernel_tri_t kernel_tri<bfloat, 2>;
+template [[host_name("kernel_tri_bf16_3")]] kernel kernel_tri_t kernel_tri<bfloat, 3>;
 #endif
 
 template<typename T>