Merge branch 'main' into feat/bitwise-ops

crates/burn-jit/src/kernel/binary.rs

@@ -24,12 +24,45 @@ pub(crate) struct SubOp;
 pub(crate) struct MulOp;
 pub(crate) struct DivOp;
 pub(crate) struct RemainderOp;
-pub(crate) struct PowOp;
 pub(crate) struct BitwiseAndOp;
 pub(crate) struct BitwiseOrOp;
 pub(crate) struct BitwiseXorOp;
 pub(crate) struct BitwiseNotOp;
 
+/// Since Powf only works on float, but we still want to implement the numeric binary op family, we
+/// set another precision in the family type to cast, when necessary, the input value to a valid
+/// float.
+///
+/// Because of this we won't benefit from the cubecl rust compilation speed improvement from using
+/// the family pattern for [PowOp], but at least we don't duplicate code.
+pub(crate) struct PowOp<F: Float> {
+    _f: PhantomData<F>,
+}
+
+impl BinaryOpFamily for AddOp {
+    type BinaryOp<C: Numeric> = Self;
+}
+
+impl BinaryOpFamily for SubOp {
+    type BinaryOp<C: Numeric> = Self;
+}
+
+impl BinaryOpFamily for MulOp {
+    type BinaryOp<C: Numeric> = Self;
+}
+
+impl BinaryOpFamily for DivOp {
+    type BinaryOp<C: Numeric> = Self;
+}
+
+impl BinaryOpFamily for RemainderOp {
+    type BinaryOp<C: Numeric> = Self;
+}
+
+impl<F: Float> BinaryOpFamily for PowOp<F> {
+    type BinaryOp<C: Numeric> = Self;
+}
+
 #[cube]
 impl<N: Numeric> BinaryOp<N> for AddOp {
     fn execute(lhs: Line<N>, rhs: Line<N>) -> Line<N> {
@@ -76,6 +109,7 @@ impl<N: Numeric, F: Float> BinaryOp<N> for PowOp<F> {
     }
 }
 
+
 #[cube]
 impl<N: Numeric> BinaryOp<N> for BitwiseAndOp {
     fn execute(lhs: Line<N>, rhs: Line<N>) -> Line<N> {
@@ -105,8 +139,9 @@ impl<N: Numeric> BinaryOp<N> for BitwiseNotOp {
     }
 }
 
-#[cube(launch)]
-pub(crate) fn kernel_scalar_binop<C: Numeric, O: BinaryOp<C>>(
+
+#[cube(launch_unchecked)]
+pub(crate) fn kernel_scalar_binop<C: Numeric, O: BinaryOpFamily>(
     input: &Tensor<Line<C>>,
     scalar: C,
     output: &mut Tensor<Line<C>>,