cera 0.5.3

Rust-native LLM inference engine
Documentation
#include <metal_stdlib>
#include <metal_math>
#include <metal_texture>
using namespace metal;

#line 58 "cera/src/backend/shaders/slang/mel_norm.slang"
struct KernelContext_0
{
    packed_uint4 device* par_buf_0;
    float device* dst_buf_0;
    float device* mel_buf_0;
    array<float, int(256)> threadgroup* scratch_0;
};


#line 52
float block_sum_0(uint tid_0, float v_0, KernelContext_0 thread* kernelContext_0)
{



    float sg_0 = simd_sum(v_0);
    if((tid_0 & 31U) == 0U)
    {

#line 58
        (*kernelContext_0->scratch_0)[tid_0 >> 5U] = sg_0;

#line 58
    }
    threadgroup_barrier(mem_flags::mem_threadgroup);

#line 59
    float lane_0;
    if(tid_0 < 8U)
    {

#line 60
        lane_0 = (*kernelContext_0->scratch_0)[tid_0];

#line 60
    }
    else
    {

#line 60
        lane_0 = 0.0f;

#line 60
    }
    float total_0 = simd_sum(lane_0);
    if(tid_0 == 0U)
    {

#line 62
        (*kernelContext_0->scratch_0)[int(0)] = total_0;

#line 62
    }
    threadgroup_barrier(mem_flags::mem_threadgroup);
    float _S1 = (*kernelContext_0->scratch_0)[int(0)];

#line 79
    return _S1;
}



[[kernel]] void mel_norm(uint3 lid_0 [[thread_position_in_threadgroup]], uint3 wid_0 [[threadgroup_position_in_grid]], packed_uint4 device* par_buf_1 [[buffer(2)]], float device* dst_buf_1 [[buffer(1)]], float device* mel_buf_1 [[buffer(0)]])
{

#line 84
    thread KernelContext_0 kernelContext_1;

#line 84
    (&kernelContext_1)->par_buf_0 = par_buf_1;

#line 84
    (&kernelContext_1)->dst_buf_0 = dst_buf_1;

#line 84
    (&kernelContext_1)->mel_buf_0 = mel_buf_1;

#line 84
    threadgroup array<float, int(256)> scratch_1;

#line 84
    (&kernelContext_1)->scratch_0 = &scratch_1;
    uint tid_1 = lid_0.x;
    uint mi_0 = wid_0.x;
    uint _S2 = (uint4(*(par_buf_1+int(0))) ).x;
    uint n_frames_0 = (uint4(*(par_buf_1+int(0))) ).y;
    uint eff_0 = (uint4(*(par_buf_1+int(0))) ).z;
    float eps_0 = (as_type<float>(((uint4(*(par_buf_1+int(0))) ).w)));
    uint _S3 = mi_0 * n_frames_0;

#line 91
    uint t_0;



    if(eff_0 <= 1U)
    {

#line 95
        t_0 = tid_1;
        for(;;)
        {

#line 96
            if(t_0 < n_frames_0)
            {
            }
            else
            {

#line 96
                break;
            }

#line 97
            *((&kernelContext_1)->dst_buf_0+(t_0 * _S2 + mi_0)) = 0.0f;

#line 96
            t_0 = t_0 + 256U;

#line 96
        }


        return;
    }

#line 99
    t_0 = tid_1;

#line 99
    float partial_0 = 0.0f;

#line 104
    for(;;)
    {

#line 104
        if(t_0 < eff_0)
        {
        }
        else
        {

#line 104
            break;
        }

#line 105
        float partial_1 = partial_0 + (&kernelContext_1)->mel_buf_0[_S3 + t_0];

#line 104
        t_0 = t_0 + 256U;

#line 104
        partial_0 = partial_1;

#line 104
    }

#line 104
    float _S4 = block_sum_0(tid_1, partial_0, &kernelContext_1);


    float _S5 = _S4 / float(eff_0);


    threadgroup_barrier(mem_flags::mem_threadgroup);

#line 110
    t_0 = tid_1;

#line 110
    partial_0 = 0.0f;



    for(;;)
    {

#line 114
        if(t_0 < eff_0)
        {
        }
        else
        {

#line 114
            break;
        }

#line 115
        float d_0 = (&kernelContext_1)->mel_buf_0[_S3 + t_0] - _S5;
        float partial_2 = partial_0 + d_0 * d_0;

#line 114
        t_0 = t_0 + 256U;

#line 114
        partial_0 = partial_2;

#line 114
    }

#line 114
    float _S6 = block_sum_0(tid_1, partial_0, &kernelContext_1);



    float _S7 = 1.0f / sqrt(_S6 / float(eff_0 - 1U) + eps_0);

#line 118
    t_0 = tid_1;


    for(;;)
    {

#line 121
        if(t_0 < n_frames_0)
        {
        }
        else
        {

#line 121
            break;
        }

#line 121
        float v_1;
        if(t_0 < eff_0)
        {

#line 122
            v_1 = ((&kernelContext_1)->mel_buf_0[_S3 + t_0] - _S5) * _S7;

#line 122
        }
        else
        {

#line 122
            v_1 = 0.0f;

#line 122
        }
        *((&kernelContext_1)->dst_buf_0+(t_0 * _S2 + mi_0)) = v_1;

#line 121
        t_0 = t_0 + 256U;

#line 121
    }



    return;
}