ruprim_host/binary/
dispatch_int.rs

1use ruda_core::tensor::{DType, host::HostTensor, element::Scalar, TensorMetadata};
2use num_traits::ToPrimitive;
3use super::{binary_op_typed, int_binary_op, int_scalar_op, scalar_op_typed};
4
5pub fn int_add(lhs: HostTensor, rhs: HostTensor) -> HostTensor {
6    int_binary_op(lhs, rhs, |a, b| a + b)
7}
8
9pub fn int_add_scalar(lhs: HostTensor, rhs: Scalar) -> HostTensor {
10    if lhs.dtype() == DType::U64 {
11        return scalar_op_typed(lhs, rhs.to_u64().unwrap(), |a: u64, b: u64| {
12            a.wrapping_add(b)
13        });
14    }
15    int_scalar_op(lhs, rhs.to_i64().unwrap(), |a, b| a + b)
16}
17
18pub fn int_sub(lhs: HostTensor, rhs: HostTensor) -> HostTensor {
19    int_binary_op(lhs, rhs, |a, b| a - b)
20}
21
22pub fn int_sub_scalar(lhs: HostTensor, rhs: Scalar) -> HostTensor {
23    if lhs.dtype() == DType::U64 {
24        return scalar_op_typed(lhs, rhs.to_u64().unwrap(), |a: u64, b: u64| {
25            a.wrapping_sub(b)
26        });
27    }
28    int_scalar_op(lhs, rhs.to_i64().unwrap(), |a, b| a - b)
29}
30
31pub fn int_mul(lhs: HostTensor, rhs: HostTensor) -> HostTensor {
32    int_binary_op(lhs, rhs, |a, b| a * b)
33}
34
35pub fn int_mul_scalar(lhs: HostTensor, rhs: Scalar) -> HostTensor {
36    if lhs.dtype() == DType::U64 {
37        return scalar_op_typed(lhs, rhs.to_u64().unwrap(), |a: u64, b: u64| {
38            a.wrapping_mul(b)
39        });
40    }
41    int_scalar_op(lhs, rhs.to_i64().unwrap(), |a, b| a * b)
42}
43
44pub fn int_div(lhs: HostTensor, rhs: HostTensor) -> HostTensor {
45    // U64 values > i64::MAX produce wrong results through i64 cast
46    if lhs.dtype() == DType::U64 {
47        let (lhs, rhs) = crate::expand::broadcast_binary(lhs, rhs);
48        return binary_op_typed(lhs, &rhs, |a: u64, b: u64| a / b);
49    }
50    int_binary_op(lhs, rhs, |a, b| a / b)
51}
52
53pub fn int_div_scalar(lhs: HostTensor, rhs: Scalar) -> HostTensor {
54    if lhs.dtype() == DType::U64 {
55        return scalar_op_typed(lhs, rhs.to_u64().unwrap(), |a: u64, b: u64| a / b);
56    }
57    int_scalar_op(lhs, rhs.to_i64().unwrap(), |a, b| a / b)
58}
59
60pub fn int_remainder(lhs: HostTensor, rhs: HostTensor) -> HostTensor {
61    // U64 values > i64::MAX produce wrong results through i64 cast
62    if lhs.dtype() == DType::U64 {
63        let (lhs, rhs) = crate::expand::broadcast_binary(lhs, rhs);
64        return binary_op_typed(lhs, &rhs, |a: u64, b: u64| a % b);
65    }
66    // Python/PyTorch-style remainder: result has same sign as divisor
67    int_binary_op(lhs, rhs, |a, b| ((a % b) + b) % b)
68}
69
70pub fn int_remainder_scalar(lhs: HostTensor, rhs: Scalar) -> HostTensor {
71    if lhs.dtype() == DType::U64 {
72        return scalar_op_typed(lhs, rhs.to_u64().unwrap(), |a: u64, b: u64| a % b);
73    }
74    // Python/PyTorch-style remainder: result has same sign as divisor
75    int_scalar_op(lhs, rhs.to_i64().unwrap(), |a, b| ((a % b) + b) % b)
76}
77
78pub fn bitwise_and(lhs: HostTensor, rhs: HostTensor) -> HostTensor {
79    int_binary_op(lhs, rhs, |a, b| a & b)
80}
81
82pub fn bitwise_and_scalar(lhs: HostTensor, rhs: Scalar) -> HostTensor {
83    if lhs.dtype() == DType::U64 {
84        return scalar_op_typed(lhs, rhs.to_u64().unwrap(), |a: u64, b: u64| a & b);
85    }
86    int_scalar_op(lhs, rhs.to_i64().unwrap(), |a, b| a & b)
87}
88
89pub fn bitwise_or(lhs: HostTensor, rhs: HostTensor) -> HostTensor {
90    int_binary_op(lhs, rhs, |a, b| a | b)
91}
92
93pub fn bitwise_or_scalar(lhs: HostTensor, rhs: Scalar) -> HostTensor {
94    if lhs.dtype() == DType::U64 {
95        return scalar_op_typed(lhs, rhs.to_u64().unwrap(), |a: u64, b: u64| a | b);
96    }
97    int_scalar_op(lhs, rhs.to_i64().unwrap(), |a, b| a | b)
98}
99
100pub fn bitwise_xor(lhs: HostTensor, rhs: HostTensor) -> HostTensor {
101    int_binary_op(lhs, rhs, |a, b| a ^ b)
102}
103
104pub fn bitwise_xor_scalar(lhs: HostTensor, rhs: Scalar) -> HostTensor {
105    if lhs.dtype() == DType::U64 {
106        return scalar_op_typed(lhs, rhs.to_u64().unwrap(), |a: u64, b: u64| a ^ b);
107    }
108    int_scalar_op(lhs, rhs.to_i64().unwrap(), |a, b| a ^ b)
109}
110
111pub fn bitwise_not(tensor: HostTensor) -> HostTensor {
112    // Use scalar op with dummy value, only applying NOT to lhs
113    int_scalar_op(tensor, 0, |a, _| !a)
114}
115
116pub fn bitwise_left_shift(lhs: HostTensor, rhs: HostTensor) -> HostTensor {
117    int_binary_op(lhs, rhs, |a, b| a.wrapping_shl(b as u32))
118}
119
120pub fn bitwise_left_shift_scalar(lhs: HostTensor, rhs: Scalar) -> HostTensor {
121    int_scalar_op(lhs, rhs.to_i64().unwrap(), |a, b| a.wrapping_shl(b as u32))
122}
123
124pub fn bitwise_right_shift(lhs: HostTensor, rhs: HostTensor) -> HostTensor {
125    int_binary_op(lhs, rhs, |a, b| a.wrapping_shr(b as u32))
126}
127
128pub fn bitwise_right_shift_scalar(lhs: HostTensor, rhs: Scalar) -> HostTensor {
129    int_scalar_op(lhs, rhs.to_i64().unwrap(), |a, b| a.wrapping_shr(b as u32))
130}
131
132pub fn int_powi(lhs: HostTensor, rhs: HostTensor) -> HostTensor {
133    int_binary_op(lhs, rhs, |a, b| a.wrapping_pow(b as u32))
134}
135
136pub fn int_powi_scalar(lhs: HostTensor, rhs: ruda_core::tensor::element::Scalar) -> HostTensor {
137    use num_traits::ToPrimitive;
138    match rhs.to_i64().unwrap() {
139        0 => crate::fill::int_ones(lhs.shape(), lhs.dtype().into()),
140        1 => lhs,
141        2 => crate::binary::dispatch_int::int_mul(lhs.clone(), lhs),
142        _ => crate::binary::dispatch_int::int_powi_scalar_impl(lhs, rhs),
143    }
144}
145
146pub fn int_powi_scalar_impl(lhs: HostTensor, rhs: ruda_core::tensor::element::Scalar) -> HostTensor {
147    use num_traits::ToPrimitive;
148    let exp = rhs.to_i64().unwrap() as u32;
149    if lhs.dtype() == DType::U64 {
150        return scalar_op_typed(lhs, exp as u64, move |x: u64, _| x.wrapping_pow(exp));
151    }
152    int_scalar_op(lhs, exp as i64, move |x, _| x.wrapping_pow(exp))
153}
154
ruprim_host/binary/dispatch_int.rs

ruprim_host/binary/
dispatch_int.rs
