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.

Sign up for GitHub

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

Jump to bottom

jagged_dense_elementwise_mul add Meta and Autograd backend #1469

Closed
wants to merge 4 commits into from
Closed

jagged_dense_elementwise_mul add Meta and Autograd backend #1469

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
214aa0f
add bf16 support for jagged tensor op (#1472)
jianyuh Nov 23, 2022
e6fb259
Jagged to padded dense
brad-mengchi Nov 23, 2022
878fddf
Jagged-dense-dense elementwise jagged autograd and meta backend
brad-mengchi Nov 23, 2022
44c686c
jagged_dense_elementwise_mul add Meta and Autograd backend (#1469)
brad-mengchi Nov 23, 2022
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
2 changes: 2 additions & 0 deletions fbgemm_gpu/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -290,6 +290,8 @@ set(fbgemm_gpu_sources_cpu
codegen/embedding_bounds_check_host_cpu.cpp
src/permute_pooled_embedding_ops_split_cpu.cpp
src/cpu_utils.cpp
src/jagged_tensor_ops_autograd.cpp
src/jagged_tensor_ops_meta.cpp
src/jagged_tensor_ops_cpu.cpp
src/input_combine_cpu.cpp
src/layout_transform_ops_cpu.cpp
Expand Down
71 changes: 27 additions & 44 deletions fbgemm_gpu/include/fbgemm_gpu/sparse_ops.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -365,23 +365,6 @@ at::Tensor batched_unary_embeddings_backward_cuda(
const at::Tensor& offsets,
const at::Tensor& indices);

///@ingroup sparse-data-cuda
std::tuple<std::vector<at::Tensor>, std::vector<at::Tensor>>
stacked_jagged_2d_to_dense_forward_cuda(
at::Tensor values,
at::Tensor lengths,
const std::vector<int64_t>& offset_per_key,
const std::vector<int64_t>& max_lengths_per_key);

///@ingroup sparse-data-cuda
at::Tensor stacked_jagged_2d_to_dense_backward_cuda(
int64_t B,
int64_t D,
int64_t total_L,
const std::vector<at::Tensor>& grad_padded_values_per_key,
const std::vector<at::Tensor>& offsets_tensor_per_key,
const std::vector<int64_t>& offset_per_key);

///@ingroup sparse-data-cpu
std::vector<at::Tensor> stacked_jagged_2d_to_dense_cpu(
at::Tensor values,
Expand All @@ -390,32 +373,47 @@ std::vector<at::Tensor> stacked_jagged_2d_to_dense_cpu(
const std::vector<int64_t>& max_lengths_per_key,
int64_t padding_value);

///@ingroup sparse-data-cuda
at::Tensor jagged_1d_to_dense_gpu(
at::Tensor jagged_to_padded_dense_autograd(
const at::Tensor& values,
const std::vector<at::Tensor>& offsets,
const std::vector<std::int64_t>& max_lengths,
const double padding_value);

at::Tensor jagged_dense_elementwise_add_autograd(
const at::Tensor& x_values,
const std::vector<at::Tensor>& x_offsets,
const at::Tensor& y);

at::Tensor jagged_1d_to_dense_autograd(
at::Tensor values,
at::Tensor offsets,
int64_t max_L,
int64_t padding_value);

///@ingroup sparse-data-cpu
at::Tensor jagged_1d_to_dense_cpu(
at::Tensor jagged_2d_to_dense_autograd(
at::Tensor values,
at::Tensor offsets,
int64_t max_L,
int64_t padding_value);
int64_t max_sequence_length);

std::tuple<at::Tensor, std::vector<at::Tensor>>
jagged_dense_dense_elementwise_add_jagged_output_autograd(
const at::Tensor& x_values,
const std::vector<at::Tensor>& x_offsets,
const at::Tensor& y_0,
const at::Tensor& y_1);

std::tuple<at::Tensor, std::vector<at::Tensor>>
jagged_dense_elementwise_mul_autograd(
const at::Tensor& x_values,
const std::vector<at::Tensor>& x_offsets,
const at::Tensor& y);

///@ingroup sparse-data-cpu
at::Tensor jagged_2d_to_dense_forward_cpu(
at::Tensor values,
at::Tensor offsets,
int64_t max_L);

///@ingroup sparse-data-cuda
at::Tensor jagged_2d_to_dense_gpu(
at::Tensor values,
at::Tensor offsets,
int64_t max_sequence_length);

///@ingroup sparse-data-cuda
at::Tensor jagged_2d_to_dense_gpu_forward(
at::Tensor values,
Expand All @@ -428,21 +426,6 @@ at::Tensor jagged_2d_to_dense_gpu_backward(
at::Tensor offsets,
int64_t max_lengths);

///@ingroup sparse-data-gpu
std::vector<at::Tensor> stacked_jagged_1d_to_dense_gpu(
at::Tensor values,
at::Tensor lengths,
const std::vector<int64_t>& offset_per_key,
const std::vector<int64_t>& max_lengths_per_key,
int64_t padding_value);

///@ingroup sparse-data-cpu
std::vector<at::Tensor> stacked_jagged_1d_to_dense_cpu(
at::Tensor values,
at::Tensor lengths,
const std::vector<int64_t>& offset_per_key,
const std::vector<int64_t>& max_lengths_per_key,
int64_t padding_value);
#endif

///@ingroup sparse-data-cpu
Expand Down
Loading