Codegen: Optional Expression Rewriter

[eugenebokhan]

What changed

• OPTIONAL(...) conditions like OPTIONAL(quant_method == QuantizationMethod::AWQ) now are available to use.
• Replaced use_zero_points / use_mlx_quant bool pair in Qmv, QmvFast, QmmTransposed kernels with single quant_method: QuantizationMethod enum.
• Split bindgen unreadable spaghetti into a typed sum-of-products module.

Bindgen decoupling

build/metal/bindgen/
├── mod.rs                       # orchestrator
├── arguments.rs                 # ArgumentEmission { Buffer | Constant | IndirectDispatch }
├── specialize.rs                # SpecializeEmission
├── variants.rs                  # VariantBind
├── variant_path_rewriter.rs     # variant ident → self.field rewriter
├── dispatch.rs                  # DispatchEmission (axis / direct / indirect)
└── trait_wiring.rs              # TraitWiring

Each kernel argument is parsed once into a typed variant; each emission stream is one iter().filter_map() pass:

pub enum ArgumentEmission {
    Buffer(BufferArgument), 
    Constant(ConstantArgument), 
    IndirectDispatch(IndirectDispatchArgument),
}

impl ArgumentEmission {
    pub fn struct_field(&self)               -> Option<TokenStream>;
    pub fn struct_initializer(&self)         -> Option<TokenStream>;
    pub fn encode_argument_definition(&self) -> TokenStream;
    pub fn encode_lifetime(&self)            -> Option<TokenStream>;
    pub fn encode_deconstruct(&self)         -> Option<TokenStream>;
    pub fn encode_access(&self)              -> Option<TokenStream>;
    pub fn encode_set(&self)                 -> TokenStream;
}

Replaced a .multiunzip() over a 7-tuple of Vec<Option<TokenStream>> and three mut accumulators (arg_count, indirect_flag, referenced_variants) with typed fields and a VariantPathRewriter consumed via .finish().

mod.rs orchestrator — parse → emit → assemble:

let variant_binds         = variants::parse(kernel)?;
let argument_emissions    = arguments::parse(kernel, enum_path_rewriter)?;
let specialize_emission   = specialize::parse(kernel, base_index, kernel_name)?;
let trait_wiring          = trait_wiring::build(kernel, &trait_name, &struct_name);
let mut variant_rewriter  = VariantPathRewriter::new(&variant_binds, kernel_name);
let dispatch_emission     = dispatch::parse(kernel, &mut variant_rewriter)?;
let referenced            = variant_rewriter.finish();

// One pass per output stream
let conditional_buffer_fields: Vec<_> =
    argument_emissions.iter().filter_map(|a| a.struct_field()).collect();
// …etc

quote! { /* assemble struct + impl from the named pieces */ }

[CC-Yeh]

Should we let it fail earlier instead of getting cryptic metal compile error?

[CC-Yeh]

How about ZeroPoint and Bias?

[CC-Yeh]

Curious why we use #define/#undef here instead of having rewrite_for_metal substitute bare specialize names with their prefixed counterparts, wouldn't that keep all name resolution in one place?