feat: bitwise ops implementation for line (#284)

* feat: bitwise ops implementation for line * feat: bitwise-ops-fix * fix: bitwise-not
tracel-ai · Nov 21, 2024 · 6d37f80 · 6d37f80
1 parent 8174ab7
commit 6d37f80
Showing 1 changed file with 122 additions and 0 deletions.
diff --git a/crates/cubecl-core/src/frontend/container/line/ops.rs b/crates/cubecl-core/src/frontend/container/line/ops.rs
@@ -119,6 +119,128 @@ where
     }
 }
 
+impl<P> core::ops::BitAnd<Self> for Line<P>
+where
+    P: CubePrimitive,
+    P: core::ops::BitAnd<P, Output = P>,
+{
+    type Output = Self;
+
+    fn bitand(self, rhs: Self) -> Self::Output {
+        Self::new(self.val & rhs.val)
+    }
+}
+
+impl<P> core::ops::BitOr<Self> for Line<P>
+where
+    P: CubePrimitive,
+    P: core::ops::BitOr<P, Output = P>,
+{
+    type Output = Self;
+
+    fn bitor(self, rhs: Self) -> Self::Output {
+        Self::new(self.val | rhs.val)
+    }
+}
+
+impl<P> core::ops::BitXor<Self> for Line<P>
+where
+    P: CubePrimitive,
+    P: core::ops::BitXor<P, Output = P>,
+{
+    type Output = Self;
+
+    fn bitxor(self, rhs: Self) -> Self::Output {
+        Self::new(self.val ^ rhs.val)
+    }
+}
+
+impl<P> core::ops::Not for Line<P>
+where
+    P: CubePrimitive,
+    P: core::ops::Not<Output = P>,
+{
+    type Output = Self;
+
+    fn not(self) -> Self::Output {
+        Self::new(!self.val)
+    }
+}
+
+impl<P> core::ops::Shl<Self> for Line<P>
+where
+    P: CubePrimitive,
+    P: core::ops::Shl<P, Output = P>,
+{
+    type Output = Self;
+
+    fn shl(self, rhs: Self) -> Self::Output {
+        Self::new(self.val << rhs.val)
+    }
+}
+
+impl<P> core::ops::Shr<Self> for Line<P>
+where
+    P: CubePrimitive,
+    P: core::ops::Shr<P, Output = P>,
+{
+    type Output = Self;
+
+    fn shr(self, rhs: Self) -> Self::Output {
+        Self::new(self.val >> rhs.val)
+    }
+}
+
+impl<P> core::ops::BitAndAssign<Self> for Line<P>
+where
+    P: CubePrimitive,
+    P: core::ops::BitAndAssign,
+{
+    fn bitand_assign(&mut self, rhs: Self) {
+        self.val &= rhs.val;
+    }
+}
+
+impl<P> core::ops::BitOrAssign<Self> for Line<P>
+where
+    P: CubePrimitive,
+    P: core::ops::BitOrAssign,
+{
+    fn bitor_assign(&mut self, rhs: Self) {
+        self.val |= rhs.val;
+    }
+}
+
+impl<P> core::ops::BitXorAssign<Self> for Line<P>
+where
+    P: CubePrimitive,
+    P: core::ops::BitXorAssign,
+{
+    fn bitxor_assign(&mut self, rhs: Self) {
+        self.val ^= rhs.val;
+    }
+}
+
+impl<P> core::ops::ShlAssign<Self> for Line<P>
+where
+    P: CubePrimitive,
+    P: core::ops::ShlAssign,
+{
+    fn shl_assign(&mut self, rhs: Self) {
+        self.val <<= rhs.val;
+    }
+}
+
+impl<P> core::ops::ShrAssign<Self> for Line<P>
+where
+    P: CubePrimitive,
+    P: core::ops::ShrAssign,
+{
+    fn shr_assign(&mut self, rhs: Self) {
+        self.val >>= rhs.val;
+    }
+}
+
 impl<P: CubePrimitive + Abs> Abs for Line<P> {}
 impl<P: CubePrimitive + Max> Max for Line<P> {}
 impl<P: CubePrimitive + Min> Min for Line<P> {}