Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: bit shift #599

Merged
merged 2 commits into from
Apr 22, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions docs/framework/operators/tensor/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -133,6 +133,7 @@ use orion::operators::tensor::TensorTrait;
| [`tensor.dynamic_quantize_linear`](tensor.dynamic\_quantize\_linear.md) | Computes the Scale, Zero Point and FP32->8Bit conversion of FP32 Input data. |
| [`tensor.scatter_nd`](tensor.scatter\_nd.md) | The output of the operation is produced by creating a copy of the input data, and then updating its value to values specified by updates at specific index positions specified by indices. Its output shape is the same as the shape of data |
| [`tensor.label_encoder`](tensor.label\_encoder.md) | Maps each element in the input tensor to another value. |
| [`tensor.bit_shift`](tensor.bit\_shift.md) | Bitwise shift operator performs element-wise operation. For each input element, if the attribute "direction" is "RIGHT", this operator moves its binary representation toward the right side so that the input value is effectively decreased. If the attribute "direction" is "LEFT", bits of binary representation moves toward the left side, which results the increase of its actual value. |
| [`tensor.eye_like`](tensor.eye\_like.md) | Generate a 2D tensor (matrix) with ones on the diagonal and zeros everywhere else. Only 2D tensors are supported, i.e. input T1 must be of rank 2. The shape of the output tensor is the same as the input tensor. |

## Arithmetic Operations
Expand Down
68 changes: 68 additions & 0 deletions docs/framework/operators/tensor/tensor.bit_shift.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,68 @@
# TensorTrait::bit_shift

```rust
fn bit_shift(tensor1: @Tensor<T>, tensor2: @Tensor<T>, direction: felt252) -> Tensor<T>;
```

Bitwise shift operator performs element-wise operation. For each input element, if the attribute "direction" is "RIGHT", this operator moves its binary representation toward the right side so that the input value is effectively decreased. If the attribute "direction" is "LEFT", bits of binary representation moves toward the left side, which results the increase of its actual value.
The input tensor1 is the tensor to be shifted and another input tensor2 specifies the amounts of shifting.

## Args

* `tensor1`(`@Tensor<T>`) - First operand, input to be shifted.
* `tensor2`(`@Tensor<T>`) - Second operand, amounts of shift.
* `direction`(@Tensor<T>) - Direction of moving bits. It can be either "RIGHT" (for right shift) or "LEFT" (for left shift).

## Returns

* Output tensor

## Examples

```rust

use orion::operators::tensor::{U32Tensor, U32TensorAdd};
use core::array::{ArrayTrait, SpanTrait};
use orion::operators::tensor::{TensorTrait, Tensor};
use orion::utils::{assert_eq, assert_seq_eq};
use orion::operators::tensor::U32TensorPartialEq;


fn example() -> Tensor<u32> {
let mut shape = ArrayTrait::<usize>::new();
shape.append(3);
shape.append(3);

let mut data = ArrayTrait::new();
data.append(21);
data.append(7);
data.append(18);
data.append(43);
data.append(49);
data.append(49);
data.append(4);
data.append(28);
data.append(24);
let input_0 = TensorTrait::new(shape.span(), data.span());

let mut shape1 = ArrayTrait::<usize>::new();
shape1.append(3);
shape1.append(3);

let mut data1 = ArrayTrait::new();
data1.append(4);
data1.append(0);
data1.append(0);
data1.append(3);
data1.append(1);
data1.append(1);
data1.append(1);
data1.append(1);
data1.append(0);
let input_1 = TensorTrait::new(shape1.span(), data1.span());


return TensorTrait::bit_shift(@input_0, @input_1, 'LEFT');
}
>>> [336 7 18 344 98 98 8 56 24]
```
32 changes: 32 additions & 0 deletions nodegen/node/bit_shift.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,32 @@
import numpy as np
from nodegen.node import RunAll
from ..helpers import make_test, to_fp, Tensor, Dtype, FixedImpl, Trait, get_data_statement

class Bit_shift(RunAll):

@staticmethod
def left_u32():
x = np.random.randint(0, 60, (3, 3)).astype(np.uint32)
y = np.random.randint(0, 6, (3, 3)).astype(np.uint32)
z = x << y

x = Tensor(Dtype.U32, x.shape, x.flatten())
y = Tensor(Dtype.U32, y.shape, y.flatten())
z = Tensor(Dtype.U32, z.shape, z.flatten())

name = "bit_shift_left_u32"
make_test([x, y], z, "TensorTrait::bit_shift(@input_0, @input_1, 'LEFT')", name)

@staticmethod
def right_u32():
x = np.random.randint(0, 60, (3, 3)).astype(np.uint32)
y = np.random.randint(0, 6, (3, 3)).astype(np.uint32)
z = x >> y

x = Tensor(Dtype.U32, x.shape, x.flatten())
y = Tensor(Dtype.U32, y.shape, y.flatten())
z = Tensor(Dtype.U32, z.shape, z.flatten())

name = "bit_shift_right_u32"
make_test([x, y], z, "TensorTrait::bit_shift(@input_0, @input_1, 'RIGHT')", name)

71 changes: 71 additions & 0 deletions src/operators/tensor/core.cairo
Original file line number Diff line number Diff line change
Expand Up @@ -137,6 +137,7 @@ impl TensorSerde<T, impl TSerde: Serde<T>, impl TDrop: Drop<T>> of Serde<Tensor<
/// scatter_nd - The output of the operation is produced by creating a copy of the input data, and then updating its value to values specified by updates at specific index positions specified by indices. Its output shape is the same as the shape of data
/// center_crop_pad - Center crop or pad an input to given dimensions.
/// label_encoder - Maps each element in the input tensor to another value.
/// bit_shift - Bitwise shift operator performs element-wise operation. For each input element, if the attribute "direction" is "RIGHT", this operator moves its binary representation toward the right side so that the input value is effectively decreased. If the attribute "direction" is "LEFT", bits of binary representation moves toward the left side, which results the increase of its actual value.
/// eye_like - Generate a 2D tensor (matrix) with ones on the diagonal and zeros everywhere else. Only 2D tensors are supported, i.e. input T1 must be of rank 2. The shape of the output tensor is the same as the input tensor.
trait TensorTrait<T> {
/// # tensor.new
Expand Down Expand Up @@ -6089,6 +6090,76 @@ trait TensorTrait<T> {
values: Option<Span<T>>,
values_tensor: Option<Tensor<T>>
) -> Tensor<T>;
/// # TensorTrait::bit_shift
///
/// ```rust
/// fn bit_shift(tensor1: @Tensor<T>, tensor2: @Tensor<T>, direction: felt252) -> Tensor<T>;
/// ```
///
/// Bitwise shift operator performs element-wise operation. For each input element, if the attribute "direction" is "RIGHT", this operator moves its binary representation toward the right side so that the input value is effectively decreased. If the attribute "direction" is "LEFT", bits of binary representation moves toward the left side, which results the increase of its actual value.
/// The input tensor1 is the tensor to be shifted and another input tensor2 specifies the amounts of shifting.
///
/// ## Args
///
/// * `tensor1`(`@Tensor<T>`) - First operand, input to be shifted.
/// * `tensor2`(`@Tensor<T>`) - Second operand, amounts of shift.
/// * `direction`(@Tensor<T>) - Direction of moving bits. It can be either "RIGHT" (for right shift) or "LEFT" (for left shift).
///
/// ## Returns
///
/// * Output tensor
///
/// ## Examples
///
/// ```rust
///
/// use orion::operators::tensor::{U32Tensor, U32TensorAdd};
/// use core::array::{ArrayTrait, SpanTrait};
/// use orion::operators::tensor::{TensorTrait, Tensor};
/// use orion::utils::{assert_eq, assert_seq_eq};
/// use orion::operators::tensor::U32TensorPartialEq;
///
///
/// fn example() -> Tensor<u32> {
/// let mut shape = ArrayTrait::<usize>::new();
/// shape.append(3);
/// shape.append(3);
///
/// let mut data = ArrayTrait::new();
/// data.append(21);
/// data.append(7);
/// data.append(18);
/// data.append(43);
/// data.append(49);
/// data.append(49);
/// data.append(4);
/// data.append(28);
/// data.append(24);
/// let input_0 = TensorTrait::new(shape.span(), data.span());
///
/// let mut shape1 = ArrayTrait::<usize>::new();
/// shape1.append(3);
/// shape1.append(3);
///
/// let mut data1 = ArrayTrait::new();
/// data1.append(4);
/// data1.append(0);
/// data1.append(0);
/// data1.append(3);
/// data1.append(1);
/// data1.append(1);
/// data1.append(1);
/// data1.append(1);
/// data1.append(0);
/// let input_1 = TensorTrait::new(shape1.span(), data1.span());
///
///
/// return TensorTrait::bit_shift(@input_0, @input_1, 'LEFT');
/// }
/// >>> [336 7 18 344 98 98 8 56 24]
/// ```
///
fn bit_shift(tensor1: @Tensor<T>, tensor2: @Tensor<T>, direction: felt252) -> Tensor<T>;
/// # TensorTrait::eye_like
///
/// ```rust
Expand Down
6 changes: 6 additions & 0 deletions src/operators/tensor/implementations/tensor_bool.cairo
Original file line number Diff line number Diff line change
Expand Up @@ -581,6 +581,12 @@ impl BoolTensor of TensorTrait<bool> {
panic(array!['not supported!'])
}

fn bit_shift(
tensor1: @Tensor<bool>, tensor2: @Tensor<bool>, direction: felt252
) -> Tensor<bool> {
panic(array!['not supported!'])
}

fn eye_like(self: @Tensor<bool>, k: Option<i32>) -> Tensor<bool> {
panic(array!['not supported!'])
}
Expand Down
6 changes: 6 additions & 0 deletions src/operators/tensor/implementations/tensor_complex64.cairo
Original file line number Diff line number Diff line change
Expand Up @@ -624,6 +624,12 @@ impl Complex64Tensor of TensorTrait<complex64> {
panic(array!['not supported!'])
}

fn bit_shift(
tensor1: @Tensor<complex64>, tensor2: @Tensor<complex64>, direction: felt252
) -> Tensor<complex64> {
panic(array!['not supported!'])
}

fn eye_like(self: @Tensor<complex64>, k: Option<i32>) -> Tensor<complex64> {
panic(array!['not supported!'])
}
Expand Down
6 changes: 6 additions & 0 deletions src/operators/tensor/implementations/tensor_fp16x16.cairo
Original file line number Diff line number Diff line change
Expand Up @@ -692,6 +692,12 @@ impl FP16x16Tensor of TensorTrait<FP16x16> {
)
}

fn bit_shift(
tensor1: @Tensor<FP16x16>, tensor2: @Tensor<FP16x16>, direction: felt252
) -> Tensor<FP16x16> {
panic(array!['not supported!'])
}

fn eye_like(self: @Tensor<FP16x16>, k: Option<i32>) -> Tensor<FP16x16> {
math::eye_like::eye_like(self, k)
}
Expand Down
6 changes: 6 additions & 0 deletions src/operators/tensor/implementations/tensor_fp16x16wide.cairo
Original file line number Diff line number Diff line change
Expand Up @@ -633,6 +633,12 @@ impl FP16x16WTensor of TensorTrait<FP16x16W> {
)
}

fn bit_shift(
tensor1: @Tensor<FP16x16W>, tensor2: @Tensor<FP16x16W>, direction: felt252
) -> Tensor<FP16x16W> {
panic(array!['not supported!'])
}

fn eye_like(self: @Tensor<FP16x16W>, k: Option<i32>) -> Tensor<FP16x16W> {
math::eye_like::eye_like(self, k)
}
Expand Down
6 changes: 6 additions & 0 deletions src/operators/tensor/implementations/tensor_fp32x32.cairo
Original file line number Diff line number Diff line change
Expand Up @@ -688,6 +688,12 @@ impl FP32x32Tensor of TensorTrait<FP32x32> {
)
}

fn bit_shift(
tensor1: @Tensor<FP32x32>, tensor2: @Tensor<FP32x32>, direction: felt252
) -> Tensor<FP32x32> {
panic(array!['not supported!'])
}

fn eye_like(self: @Tensor<FP32x32>, k: Option<i32>) -> Tensor<FP32x32> {
math::eye_like::eye_like(self, k)
}
Expand Down
6 changes: 6 additions & 0 deletions src/operators/tensor/implementations/tensor_fp64x64.cairo
Original file line number Diff line number Diff line change
Expand Up @@ -688,6 +688,12 @@ impl FP64x64Tensor of TensorTrait<FP64x64> {
)
}

fn bit_shift(
tensor1: @Tensor<FP64x64>, tensor2: @Tensor<FP64x64>, direction: felt252
) -> Tensor<FP64x64> {
panic(array!['not supported!'])
}

fn eye_like(self: @Tensor<FP64x64>, k: Option<i32>) -> Tensor<FP64x64> {
math::eye_like::eye_like(self, k)
}
Expand Down
6 changes: 6 additions & 0 deletions src/operators/tensor/implementations/tensor_fp8x23.cairo
Original file line number Diff line number Diff line change
Expand Up @@ -684,6 +684,12 @@ impl FP8x23Tensor of TensorTrait<FP8x23> {
)
}

fn bit_shift(
tensor1: @Tensor<FP8x23>, tensor2: @Tensor<FP8x23>, direction: felt252
) -> Tensor<FP8x23> {
panic(array!['not supported!'])
}

fn eye_like(self: @Tensor<FP8x23>, k: Option<i32>) -> Tensor<FP8x23> {
math::eye_like::eye_like(self, k)
}
Expand Down
6 changes: 6 additions & 0 deletions src/operators/tensor/implementations/tensor_fp8x23wide.cairo
Original file line number Diff line number Diff line change
Expand Up @@ -610,6 +610,12 @@ impl FP8x23WTensor of TensorTrait<FP8x23W> {
)
}

fn bit_shift(
tensor1: @Tensor<FP8x23W>, tensor2: @Tensor<FP8x23W>, direction: felt252
) -> Tensor<FP8x23W> {
panic(array!['not supported!'])
}

fn eye_like(self: @Tensor<FP8x23W>, k: Option<i32>) -> Tensor<FP8x23W> {
math::eye_like::eye_like(self, k)
}
Expand Down
6 changes: 6 additions & 0 deletions src/operators/tensor/implementations/tensor_i32.cairo
Original file line number Diff line number Diff line change
Expand Up @@ -652,6 +652,12 @@ impl I32Tensor of TensorTrait<i32> {
)
}

fn bit_shift(
tensor1: @Tensor<i32>, tensor2: @Tensor<i32>, direction: felt252
) -> Tensor<i32> {
panic(array!['not supported!'])
}

fn eye_like(self: @Tensor<i32>, k: Option<i32>) -> Tensor<i32> {
math::eye_like::eye_like(self, k)
}
Expand Down
6 changes: 6 additions & 0 deletions src/operators/tensor/implementations/tensor_i8.cairo
Original file line number Diff line number Diff line change
Expand Up @@ -655,6 +655,12 @@ impl I8Tensor of TensorTrait<i8> {
)
}

fn bit_shift(
tensor1: @Tensor<i8>, tensor2: @Tensor<i8>, direction: felt252
) -> Tensor<i8> {
panic(array!['not supported!'])
}

fn eye_like(self: @Tensor<i8>, k: Option<i32>) -> Tensor<i8> {
math::eye_like::eye_like(self, k)
}
Expand Down
6 changes: 6 additions & 0 deletions src/operators/tensor/implementations/tensor_u32.cairo
Original file line number Diff line number Diff line change
Expand Up @@ -580,6 +580,12 @@ impl U32Tensor of TensorTrait<u32> {
)
}

fn bit_shift(
tensor1: @Tensor<u32>, tensor2: @Tensor<u32>, direction: felt252
) -> Tensor<u32> {
math::bit_shift::bit_shift(tensor1, tensor2, direction)
}

fn eye_like(self: @Tensor<u32>, k: Option<i32>) -> Tensor<u32> {
math::eye_like::eye_like(self, k)
}
Expand Down
1 change: 1 addition & 0 deletions src/operators/tensor/math.cairo
Original file line number Diff line number Diff line change
Expand Up @@ -68,4 +68,5 @@ mod hann_window;
mod hamming_window;
mod blackman_window;
mod scatter_nd;
mod bit_shift;
mod eye_like;
Loading
Loading