Skip to main content

select_primitive

Function select_primitive 

Source
pub fn select_primitive(
    m: usize,
    k: usize,
    n: usize,
    has_tensor_cores: bool,
) -> ComputePrimitive
Expand description

SELECT the compute primitive FROM THE PROBLEM: use the tensor cores when the GPU has them AND the shape is WMMA-legal (16-aligned); otherwise fall back to DP4A on the CUDA cores. Correct by construction — we never emit a WMMA kernel on a shape or GPU that cannot run it. (Which of the two is actually faster when both are legal is settled by measurement on the frontier, but WMMA is the strictly-more-powerful default.)