fn register_bindings()

in crates/ratchet-core/src/ops/matmul/subgroup_gemv.rs [55:87]

• 31 lines of code
• 5 McCabe index (conditional complexity)

    fn register_bindings<P: WgslPrimitive>(
        &self,
        builder: &mut WgslKernelBuilder,
        _: bool,
    ) -> Result<(), OperationError> {
        let A = &self.lhs;
        let bias = &self.bias;
        let float_arr = Array::<P>::default();

        if A.dt().is_float() {
            builder.register_storage("mat", BindingMode::ReadOnly, float_arr);
            builder.register_storage("inVec", BindingMode::ReadOnly, float_arr);
            if bias.is_some() {
                builder.register_storage("bias", BindingMode::ReadOnly, float_arr);
            }
            builder.register_storage("outVec", BindingMode::ReadWrite, float_arr);
        } else if A.dt().is_quantized() {
            let scalar = Array::<Scalar<P::T>>::default();
            let u32_arr = Array::<Scalar<u32>>::default();
            builder.register_storage("mat", BindingMode::ReadOnly, u32_arr);
            builder.register_storage("scale", BindingMode::ReadOnly, scalar);
            builder.register_storage("inVec", BindingMode::ReadOnly, scalar);
            if bias.is_some() {
                builder.register_storage("bias", BindingMode::ReadOnly, scalar);
            }
            builder.register_storage("outVec", BindingMode::ReadWrite, scalar);
        } else {
            return Err(InvariantError::UnsupportedDType(A.dt()).into());
        }

        builder.register_uniform();
        Ok(())
    }