Skip to content

Commit

Permalink
Add device-agnostic NestedTensor Ops (#1258)
Browse files Browse the repository at this point in the history
Device-agnostic ops. CUDA, CPU dispatched to the same ops as well. 

- `_nested_compute_contiguous_strides_offsets`
- `_nested_tensor_from_mask`
- `_nested_tensor_from_mask_left_aligned`
- `_nested_view_from_buffer`
  • Loading branch information
min-jean-cho authored Jan 9, 2025
1 parent 5480a8a commit 14ead62
Showing 1 changed file with 29 additions and 0 deletions.
29 changes: 29 additions & 0 deletions yaml/native/native_functions.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -8612,3 +8612,32 @@
XPU: _weight_int4pack_mm_xpu
# autogen: _weight_int4pack_mm.out
# tags: core

- func: _nested_compute_contiguous_strides_offsets(Tensor nested_size) -> (Tensor, Tensor)
variants: function
device_check: NoCheck
dispatch:
XPU: _nested_compute_contiguous_strides_offsets

- func: _nested_tensor_from_mask(Tensor t, Tensor mask, bool mask_check=True) -> Tensor
dispatch:
XPU: NestedTensor_nested_tensor_from_mask
autogen: _nested_tensor_from_mask.out

- func: _nested_tensor_from_mask_left_aligned(Tensor t, Tensor mask) -> bool
dispatch:
XPU: NestedTensor_nested_tensor_from_mask_left_aligned

- func: _nested_view_from_buffer(Tensor(a) self, Tensor nested_size, Tensor nested_strides, Tensor offsets) -> Tensor(a)
variants: function
device_check: NoCheck
dispatch:
XPU: _nested_view_from_buffer

- func: _nested_view_from_buffer_copy(Tensor self, Tensor nested_size, Tensor nested_strides, Tensor offsets) -> Tensor
variants: function
device_check: NoCheck
tags: view_copy
dispatch:
CompositeExplicitAutogradNonFunctional: _nested_view_from_buffer_copy
autogen: _nested_view_from_buffer_copy.out

0 comments on commit 14ead62

Please sign in to comment.