Skip to main content

is_fusable_elementwise

Function is_fusable_elementwise 

Source
pub fn is_fusable_elementwise(op: &Op) -> bool
Expand description

Whether op is eligible to appear inside an Op::Fused chain.

Deliberately conservative: a fixed allowlist of simple, single-input/single-output activation ops that are (a) genuinely shape-preserving and (b) already dispatched through a uniform, fixed-block-size elementwise kernel (BLOCK_SIZE = 1024) in teeny-kernels, so every candidate is known to share the same CTA/thread structure without needing to consult a RuntimeOp (which doesn’t exist yet at this point in the pipeline — fuse_elementwise_chains() runs on the teeny-core::Op graph, before teeny-kernels’ lowering constructs any RuntimeOp). Excludes normalization ops (BatchNorm*, LayerNorm, …) even though they’re also shape-preserving: several of their RuntimeOp impls are unimplemented stubs today, and some aren’t the plain fixed-block-size elementwise shape this allowlist is scoped to.