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
60fn signed_remainder(a: i64, b: i64) -> i64 {
61    if b == -1 {
62        return 0;
63    }
64    let remainder = a % b;
65    if remainder != 0 && (remainder < 0) != (b < 0) {
66        remainder + b
67    } else {
68        remainder
69    }
70}
71
72pub fn int_remainder(lhs: HostTensor, rhs: HostTensor) -> HostTensor {
73    // U64 values > i64::MAX produce wrong results through i64 cast
74    if lhs.dtype() == DType::U64 {
75        let (lhs, rhs) = crate::expand::broadcast_binary(lhs, rhs);
76        return binary_op_typed(lhs, &rhs, |a: u64, b: u64| a % b);
77    }
78    // Python/PyTorch-style remainder: result has same sign as divisor
79    int_binary_op(lhs, rhs, signed_remainder)
80}
81
82pub fn int_remainder_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    // Python/PyTorch-style remainder: result has same sign as divisor
87    int_scalar_op(lhs, rhs.to_i64().unwrap(), signed_remainder)
88}
89
90pub fn bitwise_and(lhs: HostTensor, rhs: HostTensor) -> HostTensor {
91    int_binary_op(lhs, rhs, |a, b| a & b)
92}
93
94pub fn bitwise_and_scalar(lhs: HostTensor, rhs: Scalar) -> HostTensor {
95    if lhs.dtype() == DType::U64 {
96        return scalar_op_typed(lhs, rhs.to_u64().unwrap(), |a: u64, b: u64| a & b);
97    }
98    int_scalar_op(lhs, rhs.to_i64().unwrap(), |a, b| a & b)
99}
100
101pub fn bitwise_or(lhs: HostTensor, rhs: HostTensor) -> HostTensor {
102    int_binary_op(lhs, rhs, |a, b| a | b)
103}
104
105pub fn bitwise_or_scalar(lhs: HostTensor, rhs: Scalar) -> HostTensor {
106    if lhs.dtype() == DType::U64 {
107        return scalar_op_typed(lhs, rhs.to_u64().unwrap(), |a: u64, b: u64| a | b);
108    }
109    int_scalar_op(lhs, rhs.to_i64().unwrap(), |a, b| a | b)
110}
111
112pub fn bitwise_xor(lhs: HostTensor, rhs: HostTensor) -> HostTensor {
113    int_binary_op(lhs, rhs, |a, b| a ^ b)
114}
115
116pub fn bitwise_xor_scalar(lhs: HostTensor, rhs: Scalar) -> HostTensor {
117    if lhs.dtype() == DType::U64 {
118        return scalar_op_typed(lhs, rhs.to_u64().unwrap(), |a: u64, b: u64| a ^ b);
119    }
120    int_scalar_op(lhs, rhs.to_i64().unwrap(), |a, b| a ^ b)
121}
122
123pub fn bitwise_not(tensor: HostTensor) -> HostTensor {
124    // Use scalar op with dummy value, only applying NOT to lhs
125    int_scalar_op(tensor, 0, |a, _| !a)
126}
127
128pub fn bitwise_left_shift(lhs: HostTensor, rhs: HostTensor) -> HostTensor {
129    int_binary_op(lhs, rhs, |a, b| a.wrapping_shl(b as u32))
130}
131
132pub fn bitwise_left_shift_scalar(lhs: HostTensor, rhs: Scalar) -> HostTensor {
133    int_scalar_op(lhs, rhs.to_i64().unwrap(), |a, b| a.wrapping_shl(b as u32))
134}
135
136pub fn bitwise_right_shift(lhs: HostTensor, rhs: HostTensor) -> HostTensor {
137    int_binary_op(lhs, rhs, |a, b| a.wrapping_shr(b as u32))
138}
139
140pub fn bitwise_right_shift_scalar(lhs: HostTensor, rhs: Scalar) -> HostTensor {
141    int_scalar_op(lhs, rhs.to_i64().unwrap(), |a, b| a.wrapping_shr(b as u32))
142}
143
144pub fn int_powi(lhs: HostTensor, rhs: HostTensor) -> HostTensor {
145    int_binary_op(lhs, rhs, |a, b| a.wrapping_pow(b as u32))
146}
147
148pub fn int_powi_scalar(lhs: HostTensor, rhs: ruda_core::tensor::element::Scalar) -> HostTensor {
149    use num_traits::ToPrimitive;
150    match rhs.to_i64().unwrap() {
151        0 => crate::fill::int_ones(lhs.shape(), lhs.dtype().into()),
152        1 => lhs,
153        2 => crate::binary::dispatch_int::int_mul(lhs.clone(), lhs),
154        _ => crate::binary::dispatch_int::int_powi_scalar_impl(lhs, rhs),
155    }
156}
157
158pub fn int_powi_scalar_impl(lhs: HostTensor, rhs: ruda_core::tensor::element::Scalar) -> HostTensor {
159    use num_traits::ToPrimitive;
160    let exp = rhs.to_i64().unwrap() as u32;
161    if lhs.dtype() == DType::U64 {
162        return scalar_op_typed(lhs, exp as u64, move |x: u64, _| x.wrapping_pow(exp));
163    }
164    int_scalar_op(lhs, exp as i64, move |x, _| x.wrapping_pow(exp))
165}
ruprim_host/binary/dispatch_int.rs

ruprim_host/binary/
dispatch_int.rs
