emx-onnx-cgen 0.3.7__py3-none-any.whl → 0.4.1.dev0__py3-none-any.whl

emx_onnx_cgen/templates/arg_reduce_op.c.j2

@@ -0,0 +1,18 @@
+static inline void {{ op_name }}({{ dim_args }}{{ params | join(', ') }}) {
+{% for dim in output_shape %}
+for (idx_t {{ output_loop_vars[loop.index0] }} = 0; {{ output_loop_vars[loop.index0] }} < {{ dim }}; ++{{ output_loop_vars[loop.index0] }}) {
+{% endfor %}
+{{ input_c_type }} best = {{ input0 }}{{ init_index_expr }};
+{{ output_c_type }} best_index = 0;
+for (idx_t {{ reduce_var }} = 1; {{ reduce_var }} < {{ reduce_dim }}; ++{{ reduce_var }}) {
+    {{ input_c_type }} candidate = {{ input0 }}{{ input_index_expr }};
+    if (candidate {{ compare_op }} best) {
+        best = candidate;
+        best_index = ({{ output_c_type }}){{ reduce_var }};
+    }
+}
+{{ output }}{{ output_index_expr }} = best_index;
+{% for _ in output_shape %}
+}
+{% endfor %}
+}

emx_onnx_cgen/templates/attention_op.c.j2

@@ -0,0 +1,189 @@
+static inline void {{ op_name }}({{ dim_args }}{{ params | join(', ') }}) {
+    const {{ c_type }} scale = {{ scale_literal }};
+    const {{ c_type }} softcap = {{ softcap_literal }};
+    {% if output_present_key %}
+    for (int b = 0; b < {{ batch }}; ++b) {
+        for (int h = 0; h < {{ kv_heads }}; ++h) {
+            for (int ki = 0; ki < {{ total_seq }}; ++ki) {
+                if (ki < {{ past_seq }}) {
+                    for (int d = 0; d < {{ qk_head_size }}; ++d) {
+                        {{ output_present_key }}[b][h][ki][d] = {{ input_past_key }}[b][h][ki][d];
+                    }
+                } else {
+                    int k_index = ki - {{ past_seq }};
+                    for (int d = 0; d < {{ qk_head_size }}; ++d) {
+                        {% if k_rank == 4 %}
+                        {{ output_present_key }}[b][h][ki][d] = {{ input_k }}[b][h][k_index][d];
+                        {% else %}
+                        {{ output_present_key }}[b][h][ki][d] = {{ input_k }}[b][k_index][h * {{ qk_head_size }} + d];
+                        {% endif %}
+                    }
+                }
+            }
+        }
+    }
+    {% endif %}
+    {% if output_present_value %}
+    for (int b = 0; b < {{ batch }}; ++b) {
+        for (int h = 0; h < {{ kv_heads }}; ++h) {
+            for (int ki = 0; ki < {{ total_seq }}; ++ki) {
+                if (ki < {{ past_seq }}) {
+                    for (int d = 0; d < {{ v_head_size }}; ++d) {
+                        {{ output_present_value }}[b][h][ki][d] = {{ input_past_value }}[b][h][ki][d];
+                    }
+                } else {
+                    int v_index = ki - {{ past_seq }};
+                    for (int d = 0; d < {{ v_head_size }}; ++d) {
+                        {% if v_rank == 4 %}
+                        {{ output_present_value }}[b][h][ki][d] = {{ input_v }}[b][h][v_index][d];
+                        {% else %}
+                        {{ output_present_value }}[b][h][ki][d] = {{ input_v }}[b][v_index][h * {{ v_head_size }} + d];
+                        {% endif %}
+                    }
+                }
+            }
+        }
+    }
+    {% endif %}
+    for (int b = 0; b < {{ batch }}; ++b) {
+        for (int h = 0; h < {{ q_heads }}; ++h) {
+            int kv_head = h / {{ head_group_size }};
+            for (int qi = 0; qi < {{ q_seq }}; ++qi) {
+                {{ c_type }} scores[{{ total_seq }}];
+                {{ c_type }} max_score = {{ min_literal }};
+                for (int ki = 0; ki < {{ total_seq }}; ++ki) {
+                    {{ c_type }} score = {{ zero_literal }};
+                    int k_index = ki - {{ past_seq }};
+                    for (int d = 0; d < {{ qk_head_size }}; ++d) {
+                        {{ c_type }} q_val;
+                        {{ c_type }} k_val;
+                        {% if q_rank == 4 %}
+                        q_val = {{ input_q }}[b][h][qi][d];
+                        {% else %}
+                        q_val = {{ input_q }}[b][qi][h * {{ qk_head_size }} + d];
+                        {% endif %}
+                        {% if input_past_key %}
+                        if (ki < {{ past_seq }}) {
+                            k_val = {{ input_past_key }}[b][kv_head][ki][d];
+                        } else {
+                            {% if k_rank == 4 %}
+                            k_val = {{ input_k }}[b][kv_head][k_index][d];
+                            {% else %}
+                            k_val = {{ input_k }}[b][k_index][kv_head * {{ qk_head_size }} + d];
+                            {% endif %}
+                        }
+                        {% else %}
+                        {% if k_rank == 4 %}
+                        k_val = {{ input_k }}[b][kv_head][k_index][d];
+                        {% else %}
+                        k_val = {{ input_k }}[b][k_index][kv_head * {{ qk_head_size }} + d];
+                        {% endif %}
+                        {% endif %}
+                        score += q_val * k_val;
+                    }
+                    score *= scale;
+                    {{ c_type }} bias = {{ zero_literal }};
+                    {% if input_attn_mask %}
+                    if (ki >= {{ mask_kv_seq }}) {
+                        bias = {{ min_literal }};
+                    } else {
+                        int mask_q = {{ '0' if mask_broadcast_q_seq else 'qi' }};
+                        {% if mask_rank == 2 %}
+                        {% if mask_is_bool %}
+                        bias = {{ input_attn_mask }}[mask_q][ki] ? {{ zero_literal }} : {{ min_literal }};
+                        {% else %}
+                        bias = {{ input_attn_mask }}[mask_q][ki];
+                        {% endif %}
+                        {% elif mask_rank == 3 %}
+                        int mask_b = {{ '0' if mask_broadcast_batch else 'b' }};
+                        {% if mask_is_bool %}
+                        bias = {{ input_attn_mask }}[mask_b][mask_q][ki] ? {{ zero_literal }} : {{ min_literal }};
+                        {% else %}
+                        bias = {{ input_attn_mask }}[mask_b][mask_q][ki];
+                        {% endif %}
+                        {% else %}
+                        int mask_b = {{ '0' if mask_broadcast_batch else 'b' }};
+                        int mask_h = {{ '0' if mask_broadcast_heads else 'h' }};
+                        {% if mask_is_bool %}
+                        bias = {{ input_attn_mask }}[mask_b][mask_h][mask_q][ki] ? {{ zero_literal }} : {{ min_literal }};
+                        {% else %}
+                        bias = {{ input_attn_mask }}[mask_b][mask_h][mask_q][ki];
+                        {% endif %}
+                        {% endif %}
+                    }
+                    {% endif %}
+                    {% if input_nonpad_kv_seqlen %}
+                    if (ki >= {{ input_nonpad_kv_seqlen }}[b]) {
+                        bias = {{ min_literal }};
+                    }
+                    {% endif %}
+                    if ({{ is_causal }} && ki > qi + {{ past_seq }}) {
+                        bias = {{ min_literal }};
+                    }
+                    {{ c_type }} score_bias = score + bias;
+                    {{ c_type }} score_softcap = score_bias;
+                    if (softcap != {{ zero_literal }}) {
+                        score_softcap = softcap * {{ tanh_fn }}(score_bias / softcap);
+                    }
+                    {% if output_qk_matmul %}
+                    if ({{ qk_matmul_output_mode }} == 0) {
+                        {{ output_qk_matmul }}[b][h][qi][ki] = score;
+                    } else if ({{ qk_matmul_output_mode }} == 1) {
+                        {{ output_qk_matmul }}[b][h][qi][ki] = score_bias;
+                    } else if ({{ qk_matmul_output_mode }} == 2) {
+                        {{ output_qk_matmul }}[b][h][qi][ki] = score_softcap;
+                    }
+                    {% endif %}
+                    scores[ki] = score_softcap;
+                    max_score = {{ max_fn }}(max_score, score_softcap);
+                }
+                {{ c_type }} weights[{{ total_seq }}];
+                {{ c_type }} sum = {{ zero_literal }};
+                for (int ki = 0; ki < {{ total_seq }}; ++ki) {
+                    {{ c_type }} weight = {{ zero_literal }};
+                    if (max_score != {{ min_literal }}) {
+                        weight = {{ exp_fn }}(scores[ki] - max_score);
+                    }
+                    weights[ki] = weight;
+                    sum += weight;
+                }
+                {{ c_type }} inv_sum = sum == {{ zero_literal }} ? {{ zero_literal }} : ({{ c_type }}){{ one_literal }} / sum;
+                {% if output_qk_matmul %}
+                if ({{ qk_matmul_output_mode }} == 3) {
+                    for (int ki = 0; ki < {{ total_seq }}; ++ki) {
+                        {{ output_qk_matmul }}[b][h][qi][ki] = weights[ki] * inv_sum;
+                    }
+                }
+                {% endif %}
+                for (int vd = 0; vd < {{ v_head_size }}; ++vd) {
+                    {{ c_type }} acc = {{ zero_literal }};
+                    for (int ki = 0; ki < {{ total_seq }}; ++ki) {
+                        {{ c_type }} weight = weights[ki] * inv_sum;
+                        {% if input_past_value %}
+                        if (ki < {{ past_seq }}) {
+                            acc += weight * {{ input_past_value }}[b][kv_head][ki][vd];
+                        } else {
+                            {% if v_rank == 4 %}
+                            acc += weight * {{ input_v }}[b][kv_head][ki - {{ past_seq }}][vd];
+                            {% else %}
+                            acc += weight * {{ input_v }}[b][ki - {{ past_seq }}][kv_head * {{ v_head_size }} + vd];
+                            {% endif %}
+                        }
+                        {% else %}
+                        {% if v_rank == 4 %}
+                        acc += weight * {{ input_v }}[b][kv_head][ki][vd];
+                        {% else %}
+                        acc += weight * {{ input_v }}[b][ki][kv_head * {{ v_head_size }} + vd];
+                        {% endif %}
+                        {% endif %}
+                    }
+                    {% if output_rank == 4 %}
+                    {{ output }}[b][h][qi][vd] = acc;
+                    {% else %}
+                    {{ output }}[b][qi][h * {{ v_head_size }} + vd] = acc;
+                    {% endif %}
+                }
+            }
+        }
+    }
+}

emx_onnx_cgen/templates/average_pool_op.c.j2

@@ -0,0 +1,126 @@
+static inline void {{ op_name }}({{ dim_args }}{{ params | join(', ') }}) {
+{% if spatial_rank == 3 %}
+    for (idx_t n = 0; n < {{ batch }}; ++n) {
+        for (idx_t c = 0; c < {{ channels }}; ++c) {
+            for (idx_t od = 0; od < {{ out_d }}; ++od) {
+                for (idx_t oh = 0; oh < {{ out_h }}; ++oh) {
+                    for (idx_t ow = 0; ow < {{ out_w }}; ++ow) {
+                        {{ c_type }} values[{{ kernel_d * kernel_h * kernel_w }}];
+                        idx_t count = 0;
+                        for (idx_t kd = 0; kd < {{ kernel_d }}; ++kd) {
+                            const idx_t id = od * {{ stride_d }} + kd * {{ dilation_d }} - {{ pad_front }};
+                            for (idx_t kh = 0; kh < {{ kernel_h }}; ++kh) {
+                                const idx_t ih = oh * {{ stride_h }} + kh * {{ dilation_h }} - {{ pad_top }};
+                                for (idx_t kw = 0; kw < {{ kernel_w }}; ++kw) {
+                                    const idx_t iw = ow * {{ stride_w }} + kw * {{ dilation_w }} - {{ pad_left }};
+                                    if (id >= 0 && id < {{ in_d }}
+                                        && ih >= 0 && ih < {{ in_h }}
+                                        && iw >= 0 && iw < {{ in_w }}) {
+                                        values[count++] = {{ input0 }}[n][c][id][ih][iw];
+                                    } else if ({{ count_include_pad }}) {
+                                        values[count++] = {{ zero_literal }};
+                                    }
+                                }
+                            }
+                        }
+                        const idx_t denom = {{ count_include_pad }}
+                            ? {{ kernel_d * kernel_h * kernel_w }}
+                            : count;
+                        {{ c_type }} acc = {{ zero_literal }};
+                        if (count > 0) {
+                            idx_t reduce_count = count;
+                            while (reduce_count > 1) {
+                                idx_t next = 0;
+                                for (idx_t i = 0; i + 1 < reduce_count; i += 2) {
+                                    values[next++] = values[i] + values[i + 1];
+                                }
+                                if (reduce_count % 2) {
+                                    values[next++] = values[reduce_count - 1];
+                                }
+                                reduce_count = next;
+                            }
+                            acc = values[0];
+                        }
+                        {{ output }}[n][c][od][oh][ow] = denom ? acc / ({{ c_type }})denom : {{ zero_literal }};
+                    }
+                }
+            }
+        }
+    }
+{% elif spatial_rank == 1 %}
+    for (idx_t n = 0; n < {{ batch }}; ++n) {
+        for (idx_t c = 0; c < {{ channels }}; ++c) {
+            for (idx_t ow = 0; ow < {{ out_w }}; ++ow) {
+                {{ c_type }} values[{{ kernel_w }}];
+                idx_t count = 0;
+                for (idx_t kw = 0; kw < {{ kernel_w }}; ++kw) {
+                    const idx_t iw = ow * {{ stride_w }} + kw * {{ dilation_w }} - {{ pad_left }};
+                    if (iw >= 0 && iw < {{ in_w }}) {
+                        values[count++] = {{ input0 }}[n][c][iw];
+                    } else if ({{ count_include_pad }}) {
+                        values[count++] = {{ zero_literal }};
+                    }
+                }
+                const idx_t denom = {{ count_include_pad }} ? {{ kernel_w }} : count;
+                {{ c_type }} acc = {{ zero_literal }};
+                if (count > 0) {
+                    idx_t reduce_count = count;
+                    while (reduce_count > 1) {
+                        idx_t next = 0;
+                        for (idx_t i = 0; i + 1 < reduce_count; i += 2) {
+                            values[next++] = values[i] + values[i + 1];
+                        }
+                        if (reduce_count % 2) {
+                            values[next++] = values[reduce_count - 1];
+                        }
+                        reduce_count = next;
+                    }
+                    acc = values[0];
+                }
+                {{ output }}[n][c][ow] = denom ? acc / ({{ c_type }})denom : {{ zero_literal }};
+            }
+        }
+    }
+{% else %}
+    for (idx_t n = 0; n < {{ batch }}; ++n) {
+        for (idx_t c = 0; c < {{ channels }}; ++c) {
+            for (idx_t oh = 0; oh < {{ out_h }}; ++oh) {
+                for (idx_t ow = 0; ow < {{ out_w }}; ++ow) {
+                    {{ c_type }} values[{{ kernel_h * kernel_w }}];
+                    idx_t count = 0;
+                    for (idx_t kh = 0; kh < {{ kernel_h }}; ++kh) {
+                        const idx_t ih = oh * {{ stride_h }} + kh * {{ dilation_h }} - {{ pad_top }};
+                        for (idx_t kw = 0; kw < {{ kernel_w }}; ++kw) {
+                            const idx_t iw = ow * {{ stride_w }} + kw * {{ dilation_w }} - {{ pad_left }};
+                            if (ih >= 0 && ih < {{ in_h }} && iw >= 0 && iw < {{ in_w }}) {
+                                values[count++] = {{ input0 }}[n][c][ih][iw];
+                            } else if ({{ count_include_pad }}) {
+                                values[count++] = {{ zero_literal }};
+                            }
+                        }
+                    }
+                    const idx_t denom = {{ count_include_pad }}
+                        ? {{ kernel_h * kernel_w }}
+                        : count;
+                    {{ c_type }} acc = {{ zero_literal }};
+                    if (count > 0) {
+                        idx_t reduce_count = count;
+                        while (reduce_count > 1) {
+                            idx_t next = 0;
+                            for (idx_t i = 0; i + 1 < reduce_count; i += 2) {
+                                values[next++] = values[i] + values[i + 1];
+                            }
+                            if (reduce_count % 2) {
+                                values[next++] = values[reduce_count - 1];
+                            }
+                            reduce_count = next;
+                        }
+                        acc = values[0];
+                    }
+                    {{ output }}[n][c][oh][ow] = denom ? acc / ({{ c_type }})denom : {{ zero_literal }};
+                }
+            }
+        }
+    }
+{% endif %}
+}

emx_onnx_cgen/templates/batch_norm_op.c.j2

@@ -0,0 +1,11 @@
+static inline void {{ op_name }}({{ dim_args }}{{ params | join(', ') }}) {
+{% for dim in shape %}
+for (idx_t {{ loop_vars[loop.index0] }} = 0; {{ loop_vars[loop.index0] }} < {{ dim }}; ++{{ loop_vars[loop.index0] }}) {
+{% endfor %}
+idx_t c = {{ loop_vars[1] }};
+{{ c_type }} denom = {{ sqrt_fn }}({{ variance }}[c] + {{ epsilon_literal }});
+{{ output }}{% for var in loop_vars %}[{{ var }}]{% endfor %} = ({{ input0 }}{% for var in loop_vars %}[{{ var }}]{% endfor %} - {{ mean }}[c]) / denom * {{ scale }}[c] + {{ bias }}[c];
+{% for _ in shape %}
+}
+{% endfor %}
+}

emx_onnx_cgen/templates/bernoulli_op.c.j2

@@ -0,0 +1,34 @@
+static inline uint64_t {{ op_name }}_rng_next_u64(uint64_t *state) {
+    uint64_t x = *state;
+    x ^= x >> 12;
+    x ^= x << 25;
+    x ^= x >> 27;
+    *state = x;
+    return x * 0x2545f4914f6cdd1dull;
+}
+
+static inline double {{ op_name }}_rng_next_double(uint64_t *state) {
+    return (double){{ op_name }}_rng_next_u64(state) * (1.0 / 18446744073709551616.0);
+}
+
+static inline void {{ op_name }}({{ dim_args }}{{ params | join(', ') }}) {
+    uint64_t rng_state = {{ seed }}ull;
+    if (!rng_state) {
+        rng_state = 0x243f6a8885a308d3ull;
+    }
+{% for dim in shape %}
+    for (idx_t {{ loop_vars[loop.index0] }} = 0; {{ loop_vars[loop.index0] }} < {{ dim }}; ++{{ loop_vars[loop.index0] }}) {
+{% endfor %}
+    const double prob = (double){{ input0 }}{{ input_index_expr }};
+    int is_one = 0;
+    if (prob >= 1.0) {
+        is_one = 1;
+    } else if (prob > 0.0) {
+        const double sample = {{ op_name }}_rng_next_double(&rng_state);
+        is_one = sample < prob;
+    }
+    {{ output }}{{ output_index_expr }} = is_one ? {{ one_literal }} : {{ zero_literal }};
+{% for _ in shape %}
+    }
+{% endfor %}
+}

emx_onnx_cgen/templates/binary_op.c.j2

@@ -0,0 +1,9 @@
+static inline void {{ op_name }}({{ dim_args }}{{ params | join(', ') }}) {
+{% for dim in shape %}
+for (idx_t {{ loop_vars[loop.index0] }} = 0; {{ loop_vars[loop.index0] }} < {{ dim }}; ++{{ loop_vars[loop.index0] }}) {
+{% endfor %}
+{{ output }}{% for var in loop_vars %}[{{ var }}]{% endfor %} = {% if operator_kind == "func" %}{{ operator }}({{ left_expr }}, {{ right_expr }}){% elif operator_kind == "expr" %}{{ operator_expr }}{% else %}{{ left_expr }} {{ operator }} {{ right_expr }}{% endif %};
+{% for _ in shape %}
+}
+{% endfor %}
+}

emx_onnx_cgen/templates/cast_op.c.j2

@@ -0,0 +1,9 @@
+static inline void {{ op_name }}({{ dim_args }}{{ params | join(', ') }}) {
+{% for dim in shape %}
+for (idx_t {{ loop_vars[loop.index0] }} = 0; {{ loop_vars[loop.index0] }} < {{ dim }}; ++{{ loop_vars[loop.index0] }}) {
+{% endfor %}
+{{ output }}{% for var in loop_vars %}[{{ var }}]{% endfor %} = ({{ output_c_type }}){{ input0 }}{% for var in loop_vars %}[{{ var }}]{% endfor %};
+{% for _ in shape %}
+}
+{% endfor %}
+}

emx_onnx_cgen/templates/clip_op.c.j2

@@ -0,0 +1,14 @@
+static inline void {{ op_name }}({{ dim_args }}{{ params | join(', ') }}) {
+{% for dim in shape %}
+for (idx_t {{ loop_vars[loop.index0] }} = 0; {{ loop_vars[loop.index0] }} < {{ dim }}; ++{{ loop_vars[loop.index0] }}) {
+{% endfor %}
+{{ output_c_type }} value = {{ input_expr }};
+const {{ output_c_type }} min_value = {{ min_expr }};
+const {{ output_c_type }} max_value = {{ max_expr }};
+value = {{ max_fn }}(value, min_value);
+value = {{ min_fn }}(value, max_value);
+{{ output }}{% for var in loop_vars %}[{{ var }}]{% endfor %} = value;
+{% for _ in shape %}
+}
+{% endfor %}
+}

emx_onnx_cgen/templates/concat_op.c.j2

@@ -0,0 +1,28 @@
+static inline void {{ op_name }}({{ dim_args }}{{ params | join(', ') }}) {
+    const void *inputs[] = { {% for input in inputs %}{{ input }}{% if not loop.last %}, {% endif %}{% endfor %} };
+    const idx_t axis_sizes[] = { {% for axis in axis_sizes %}{{ axis }}{% if not loop.last %}, {% endif %}{% endfor %} };
+    idx_t concat_axis = 0;
+    for (idx_t idx = 0; idx < {{ input_count }}; ++idx) {
+        concat_axis += axis_sizes[idx];
+    }
+    if (concat_axis == 0) {
+        return;
+    }
+    for (idx_t outer_idx = 0; outer_idx < {{ outer }}; ++outer_idx) {
+        idx_t output_offset = outer_idx * concat_axis * {{ inner }};
+        idx_t axis_offset = 0;
+        for (idx_t input_idx = 0; input_idx < {{ input_count }}; ++input_idx) {
+            idx_t axis = axis_sizes[input_idx];
+            idx_t copy_elems = axis * {{ inner }};
+            const unsigned char *input_bytes =
+                (const unsigned char *)inputs[input_idx];
+            idx_t input_offset = outer_idx * copy_elems;
+            memcpy(
+                ((unsigned char *){{ output }}) + (output_offset + axis_offset) * sizeof({{ c_type }}),
+                input_bytes + input_offset * sizeof({{ c_type }}),
+                copy_elems * sizeof({{ c_type }})
+            );
+            axis_offset += copy_elems;
+        }
+    }
+}

emx_onnx_cgen/templates/constant_of_shape_op.c.j2

@@ -0,0 +1,10 @@
+static inline void {{ op_name }}({{ dim_args }}{{ params | join(', ') }}) {
+    (void){{ input0 }};
+{% for dim in shape %}
+for (idx_t {{ loop_vars[loop.index0] }} = 0; {{ loop_vars[loop.index0] }} < {{ dim }}; ++{{ loop_vars[loop.index0] }}) {
+{% endfor %}
+{{ output }}{% for var in loop_vars %}[{{ var }}]{% endfor %} = {{ value_literal }};
+{% for _ in shape %}
+}
+{% endfor %}
+}

emx_onnx_cgen/templates/conv_integer_op.c.j2

@@ -0,0 +1,34 @@
+static inline void {{ op_name }}({{ dim_args }}{{ params | join(', ') }}) {
+    for (idx_t n = 0; n < {{ batch }}; ++n) {
+        for (idx_t g = 0; g < {{ group }}; ++g) {
+            for (idx_t oc = 0; oc < {{ group_out_channels }}; ++oc) {
+                const idx_t oc_global = g * {{ group_out_channels }} + oc;
+                {% for dim in range(spatial_rank) %}
+                for (idx_t {{ out_indices[dim] }} = 0; {{ out_indices[dim] }} < {{ out_spatial[dim] }}; ++{{ out_indices[dim] }}) {
+                {% endfor %}
+                    {{ acc_type }} acc = {{ acc_zero_literal }};
+                    for (idx_t ic = 0; ic < {{ group_in_channels }}; ++ic) {
+                        const idx_t ic_global = g * {{ group_in_channels }} + ic;
+                        {% for dim in range(spatial_rank) %}
+                        for (idx_t {{ kernel_indices[dim] }} = 0; {{ kernel_indices[dim] }} < {{ kernel_shape[dim] }}; ++{{ kernel_indices[dim] }}) {
+                        {% endfor %}
+                            {% for dim in range(spatial_rank) %}
+                            const idx_t {{ in_indices[dim] }} = {{ out_indices[dim] }} * {{ strides[dim] }} + {{ kernel_indices[dim] }} * {{ dilations[dim] }} - {{ pads_begin[dim] }};
+                            {% endfor %}
+                            if ({% for dim in range(spatial_rank) %}{{ in_indices[dim] }} >= 0 && {{ in_indices[dim] }} < {{ in_spatial[dim] }}{% if not loop.last %} && {% endif %}{% endfor %}) {
+                                {{ acc_type }} input_value = ({{ acc_type }}){{ input0 }}[n][ic_global]{% for dim in range(spatial_rank) %}[{{ in_indices[dim] }}]{% endfor %} - ({{ acc_type }})({{ x_zero_expr }});
+                                {{ acc_type }} weight_value = ({{ acc_type }}){{ weights }}[oc_global][ic]{% for dim in range(spatial_rank) %}[{{ kernel_indices[dim] }}]{% endfor %} - ({{ acc_type }})({{ w_zero_expr }});
+                                acc += input_value * weight_value;
+                            }
+                        {% for dim in range(spatial_rank) %}
+                        }
+                        {% endfor %}
+                    }
+                    {{ output }}[n][oc_global]{% for dim in range(spatial_rank) %}[{{ out_indices[dim] }}]{% endfor %} = ({{ c_type }})acc;
+                {% for dim in range(spatial_rank) %}
+                }
+                {% endfor %}
+            }
+        }
+    }
+}

emx_onnx_cgen/templates/conv_op.c.j2

@@ -0,0 +1,32 @@
+static inline void {{ op_name }}({{ dim_args }}{{ params | join(', ') }}) {
+    for (idx_t n = 0; n < {{ batch }}; ++n) {
+        for (idx_t g = 0; g < {{ group }}; ++g) {
+            for (idx_t oc = 0; oc < {{ group_out_channels }}; ++oc) {
+                const idx_t oc_global = g * {{ group_out_channels }} + oc;
+                {% for dim in range(spatial_rank) %}
+                for (idx_t {{ out_indices[dim] }} = 0; {{ out_indices[dim] }} < {{ out_spatial[dim] }}; ++{{ out_indices[dim] }}) {
+                {% endfor %}
+                    {{ acc_type }} acc = {% if bias %}({{ acc_type }}){{ bias }}[oc_global]{% else %}{{ acc_zero_literal }}{% endif %};
+                    for (idx_t ic = 0; ic < {{ group_in_channels }}; ++ic) {
+                        const idx_t ic_global = g * {{ group_in_channels }} + ic;
+                        {% for dim in range(spatial_rank) %}
+                        for (idx_t {{ kernel_indices[dim] }} = 0; {{ kernel_indices[dim] }} < {{ kernel_shape[dim] }}; ++{{ kernel_indices[dim] }}) {
+                        {% endfor %}
+                            {% for dim in range(spatial_rank) %}
+                            const idx_t {{ in_indices[dim] }} = {{ out_indices[dim] }} * {{ strides[dim] }} + {{ kernel_indices[dim] }} * {{ dilations[dim] }} - {{ pads_begin[dim] }};
+                            {% endfor %}
+                            if ({% for dim in range(spatial_rank) %}{{ in_indices[dim] }} >= 0 && {{ in_indices[dim] }} < {{ in_spatial[dim] }}{% if not loop.last %} && {% endif %}{% endfor %}) {
+                                acc += ({{ acc_type }}){{ input0 }}[n][ic_global]{% for dim in range(spatial_rank) %}[{{ in_indices[dim] }}]{% endfor %} * ({{ acc_type }}){{ weights }}[oc_global][ic]{% for dim in range(spatial_rank) %}[{{ kernel_indices[dim] }}]{% endfor %};
+                            }
+                        {% for dim in range(spatial_rank) %}
+                        }
+                        {% endfor %}
+                    }
+                    {{ output }}[n][oc_global]{% for dim in range(spatial_rank) %}[{{ out_indices[dim] }}]{% endfor %} = ({{ c_type }})acc;
+                {% for dim in range(spatial_rank) %}
+                }
+                {% endfor %}
+            }
+        }
+    }
+}

emx_onnx_cgen/templates/conv_transpose_op.c.j2

@@ -0,0 +1,43 @@
+static inline void {{ op_name }}({{ dim_args }}{{ params | join(', ') }}) {
+    for (idx_t n = 0; n < {{ batch }}; ++n) {
+        for (idx_t oc = 0; oc < {{ out_channels }}; ++oc) {
+            {% for dim in range(spatial_rank) %}
+            for (idx_t {{ out_indices[dim] }} = 0; {{ out_indices[dim] }} < {{ out_spatial[dim] }}; ++{{ out_indices[dim] }}) {
+            {% endfor %}
+                {{ output }}[n][oc]{% for dim in range(spatial_rank) %}[{{ out_indices[dim] }}]{% endfor %} = {% if bias %}{{ bias }}[oc]{% else %}{{ zero_literal }}{% endif %};
+            {% for dim in range(spatial_rank) %}
+            }
+            {% endfor %}
+        }
+    }
+    for (idx_t n = 0; n < {{ batch }}; ++n) {
+        for (idx_t g = 0; g < {{ group }}; ++g) {
+            for (idx_t oc = 0; oc < {{ group_out_channels }}; ++oc) {
+                const idx_t oc_global = g * {{ group_out_channels }} + oc;
+                {% for dim in range(spatial_rank) %}
+                for (idx_t {{ kernel_indices[dim] }} = 0; {{ kernel_indices[dim] }} < {{ kernel_shape[dim] }}; ++{{ kernel_indices[dim] }}) {
+                {% endfor %}
+                    {% for dim in range(spatial_rank) %}
+                    for (idx_t {{ in_indices[dim] }} = 0; {{ in_indices[dim] }} < {{ in_spatial[dim] }}; ++{{ in_indices[dim] }}) {
+                    {% endfor %}
+                        {{ c_type }} acc = {{ zero_literal }};
+                        for (idx_t ic = 0; ic < {{ group_in_channels }}; ++ic) {
+                            const idx_t ic_global = g * {{ group_in_channels }} + ic;
+                            acc += {{ input0 }}[n][ic_global]{% for dim in range(spatial_rank) %}[{{ in_indices[dim] }}]{% endfor %} * {{ weights }}[ic_global][oc]{% for dim in range(spatial_rank) %}[{{ kernel_indices[dim] }}]{% endfor %};
+                        }
+                        {% for dim in range(spatial_rank) %}
+                        const idx_t {{ out_indices[dim] }} = {{ in_indices[dim] }} * {{ strides[dim] }} + {{ kernel_indices[dim] }} * {{ dilations[dim] }} - {{ pads_begin[dim] }};
+                        {% endfor %}
+                        if ({% for dim in range(spatial_rank) %}{{ out_indices[dim] }} >= 0 && {{ out_indices[dim] }} < {{ out_spatial[dim] }}{% if not loop.last %} && {% endif %}{% endfor %}) {
+                            {{ output }}[n][oc_global]{% for dim in range(spatial_rank) %}[{{ out_indices[dim] }}]{% endfor %} += acc;
+                        }
+                    {% for dim in range(spatial_rank) %}
+                    }
+                    {% endfor %}
+                {% for dim in range(spatial_rank) %}
+                }
+                {% endfor %}
+            }
+        }
+    }
+}

emx_onnx_cgen/templates/cumsum_op.c.j2

@@ -0,0 +1,51 @@
+static inline void {{ op_name }}({{ dim_args }}const {{ c_type }} {{ input0 }}{{ input_suffix }}{% if axis_input %}, const {{ axis_c_type }} {{ axis_input }}{{ axis_suffix }}{% endif %}, {{ c_type }} {{ output }}{{ output_suffix }}) {
+    const {{ c_type }} *input_data = (const {{ c_type }} *){{ input0 }};
+    {{ c_type }} *output_data = ({{ c_type }} *){{ output }};
+    const idx_t dims[{{ rank }}] = { {% for dim in input_shape %}{{ dim }}{% if not loop.last %}, {% endif %}{% endfor %} };
+    int axis = {% if axis_literal is not none %}{{ axis_literal }}{% else %}(int){{ axis_input }}[0]{% endif %};
+    if (axis < 0) {
+        axis += {{ rank }};
+    }
+    if (axis < 0 || axis >= {{ rank }}) {
+        return;
+    }
+    idx_t outer = 1;
+    for (int i = 0; i < axis; ++i) {
+        outer *= dims[i];
+    }
+    idx_t inner = 1;
+    for (int i = axis + 1; i < {{ rank }}; ++i) {
+        inner *= dims[i];
+    }
+    idx_t axis_dim = dims[axis];
+    for (idx_t outer_index = 0; outer_index < outer; ++outer_index) {
+        for (idx_t inner_index = 0; inner_index < inner; ++inner_index) {
+            {{ c_type }} acc = ({{ c_type }})0;
+            idx_t base = (outer_index * axis_dim * inner) + inner_index;
+{% if reverse %}
+            for (idx_t axis_offset = 0; axis_offset < axis_dim; ++axis_offset) {
+                idx_t axis_index = axis_dim - 1 - axis_offset;
+                idx_t offset = base + axis_index * inner;
+{% if exclusive %}
+                output_data[offset] = acc;
+                acc += input_data[offset];
+{% else %}
+                acc += input_data[offset];
+                output_data[offset] = acc;
+{% endif %}
+            }
+{% else %}
+            for (idx_t axis_index = 0; axis_index < axis_dim; ++axis_index) {
+                idx_t offset = base + axis_index * inner;
+{% if exclusive %}
+                output_data[offset] = acc;
+                acc += input_data[offset];
+{% else %}
+                acc += input_data[offset];
+                output_data[offset] = acc;
+{% endif %}
+            }
+{% endif %}
+        }
+    }
+}

emx_onnx_cgen/templates/depth_to_space_op.c.j2

@@ -0,0 +1,26 @@
+static inline void {{ op_name }}({{ dim_args }}{{ params | join(', ') }}) {
+    const {{ c_type }} *input_data = (const {{ c_type }} *){{ input0 }};
+    {{ c_type }} *output_data = ({{ c_type }} *){{ output }};
+    idx_t output_index = 0;
+    for (idx_t n = 0; n < {{ batch }}; ++n) {
+        for (idx_t c_out = 0; c_out < {{ out_channels }}; ++c_out) {
+            for (idx_t h_out = 0; h_out < {{ out_h }}; ++h_out) {
+                idx_t h_in = h_out / {{ blocksize }};
+                idx_t offset_h = h_out % {{ blocksize }};
+                for (idx_t w_out = 0; w_out < {{ out_w }}; ++w_out) {
+                    idx_t w_in = w_out / {{ blocksize }};
+                    idx_t offset_w = w_out % {{ blocksize }};
+                    idx_t c_in;
+{% if mode == "DCR" %}
+                    c_in = (offset_h * {{ blocksize }} + offset_w) * {{ out_channels }} + c_out;
+{% else %}
+                    c_in = (c_out * {{ blocksize }} + offset_h) * {{ blocksize }} + offset_w;
+{% endif %}
+                    idx_t input_index = ((n * {{ in_channels }} + c_in) * {{ in_h }} + h_in) * {{ in_w }} + w_in;
+                    output_data[output_index] = input_data[input_index];
+                    output_index++;
+                }
+            }
+        }
+    }
+}