feat: add exp2 and log2 operations#1197
Conversation
Add Exp2 and Log2 as first-class operations in cubecl-ir, mirroring the existing Exp and Log operations. Wire them through all backends: - cubecl-core: trait impls for all float types, Float supertrait bound - cubecl-ir: Arithmetic enum variants with Display - cubecl-opt: visitor and constant propagation support - cubecl-spirv: GLSLstd450 Exp2/Log2 extension instructions - cubecl-wgpu: exp2()/log2() WGSL builtins
| impl<const POS: u8> Recip for ElemExpand<POS> {} | ||
| impl<const POS: u8> Erf for ElemExpand<POS> {} | ||
| impl<const POS: u8> Exp for ElemExpand<POS> {} | ||
| impl<const POS: u8> Exp2 for ElemExpand<POS> {} |
| + Exp | ||
| + Exp2 | ||
| + Log | ||
| + Log1p |
There was a problem hiding this comment.
Is there a reason Log2 was left out here
There was a problem hiding this comment.
Is there a reason Log2 was left out here
| let out = out.fmt_left(); | ||
| writeln!(f, "{out} = exp({input});") | ||
| } | ||
| Instruction::Exp2 { input, out } => { |
There was a problem hiding this comment.
It should also be implemented for the cpp compiler
|
Stale? |
|
Not abandoned, I've just been a bit busy with some personal matters the last couple weeks. Give me a moment. |
|
I believe everything has been resolved, but I'm not as familiar with the codebase as the owners. Is everything good here, or does it need more work? |
nathanielsimard
left a comment
nathanielsimard
left a comment
There was a problem hiding this comment.
I think the only missing piece is the unit tests. I believe they are in cubecl-core/src/runtime_tests
|
I wasn't able to find any matching unit tests for |
4 participants
Inside tight loops for inference engines that do stuff like sigmoid functions, it's often more efficient to precalculate a constant to multiply, and then use
exp2/log2instead ofexp. This is because the2variants generally only require some very simple math on the exponent. (Some Stable Diffusion implementations use this trick, and I want to mirror it.)Add Exp2 and Log2 as first-class operations in cubecl-ir, mirroring the existing Exp and Log operations. Wire them through all backends:
Validate your PR with burn.
This PR is purely additive, and doesn't introduce any issues in burn. (But I did fork it and check.)