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

[feat] Add GPU operator: sparse_fill_empty_rows and sparse_reshape #87

Merged
merged 1 commit into from
Jun 2, 2021
Merged

[feat] Add GPU operator: sparse_fill_empty_rows and sparse_reshape #87

merged 1 commit into from
Jun 2, 2021

Conversation

funsimple
Copy link
Contributor

@funsimple funsimple commented May 27, 2021 •
edited
Loading

Description

Add GPU operator: sparse_fill_empty_rows and sparse_reshape

Fixes # (issue)

Type of change

  • Bug fix
  • New Tutorial
  • Updated or additional documentation
  • Additional Testing
  • New Feature

Checklist:

  • I've properly formatted my code according to the guidelines
    • By running Black + Flake8
    • By running pre-commit hooks
  • This PR addresses an already submitted issue for TensorFlow Recommenders Addons
  • I have made corresponding changes to the documentation
  • I have added tests that prove my fix is effective or that my feature works
  • This PR contains modifications to C++ custom-ops

How Has This Been Tested?

If you're adding a bugfix or new feature please describe the tests that you ran to verify your changes:
*

@funsimple funsimple requested a review from rhdong as a code owner May 27, 2021 07:51
@funsimple funsimple force-pushed the gpu-dev branch 5 times, most recently from 385de86 to f8d71f1 Compare May 28, 2021 03:57
@rhdong rhdong requested a review from Lifann May 28, 2021 04:34
@funsimple funsimple force-pushed the gpu-dev branch 11 times, most recently from 225a2aa to a9d93b1 Compare May 31, 2021 07:06
rhdong
rhdong reviewed Jun 1, 2021
View reviewed changes
tensorflow_recommenders_addons/dynamic_embedding/core/ops/math_ops.cc Outdated
@@ -115,4 +115,65 @@ REGISTER_OP("TFRA>SparseSegmentSumWithNumSegments")
.SetShapeFn(SparseSegmentReductionWithNumSegmentsShapeFn);
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM

#if GOOGLE_CUDA
REGISTER_OP("TFRA>SparseFillEmptyRows")
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pls replace ‘TFRA>’ by 'Tfra', for bad compatiblity with TF1.x

rhdong
rhdong reviewed Jun 1, 2021
View reviewed changes
tensorflow_recommenders_addons/dynamic_embedding/core/kernels/sparse_fill_empty_rows_op.cc
.Device(DEVICE_GPU) \
.TypeConstraint<type>("T"), \
SparseFillEmptyRowsOp<GPUDevice, type>)
TF_CALL_int32(REGISTER_KERNELS);
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We need float, half, int32 and int8 to work with GPU hash table

tensorflow_recommenders_addons/dynamic_embedding/core/kernels/sparse_fill_empty_rows_op.cu.cc
#define DEFINE_GPU_KERNELS(type) \
template struct SparseFillEmptyRowsFunctor<GPUDevice, type>;

TF_CALL_int32(DEFINE_GPU_KERNELS);
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same problem

tensorflow_recommenders_addons/dynamic_embedding/core/kernels/sparse_reshape_op.cc Outdated
};

#if GOOGLE_CUDA
REGISTER_KERNEL_BUILDER(Name("TFRA>SparseReshape").Device(DEVICE_GPU),
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Tfra

tensorflow_recommenders_addons/dynamic_embedding/core/ops/math_ops.cc Outdated
return Status::OK();
});

REGISTER_OP("TFRA>SparseReshape")
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Tfra

tensorflow_recommenders_addons/dynamic_embedding/python/ops/math_ops.py Outdated
def sparse_fill_empty_rows(sp_input, default_value, name=None):
"""Fills empty rows in the input 2-D `SparseTensor` with a default value.

It do same things as `tf.sparse.fill_empty_rows`. Here we provide GPU impl.
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It does

tensorflow_recommenders_addons/dynamic_embedding/python/ops/math_ops.py Outdated

It do same things as `tf.sparse.fill_empty_rows`. Here we provide GPU impl.

Go [tf api](https://www.tensorflow.org/api_docs/python/tf/sparse/fill_empty_rows)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

TF API

tensorflow_recommenders_addons/dynamic_embedding/python/ops/math_ops.py Outdated
Args:
sp_input: A `SparseTensor` with shape `[N, M]`.
default_value: The value to fill for empty rows, with the same type as
`sp_input.`
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

sp_input. ->sp_input

tensorflow_recommenders_addons/dynamic_embedding/python/ops/math_ops.py Outdated

It do same things as `tf.sparse.reshape`. Here we provide GPU impl.

Go [tf api](https://www.tensorflow.org/api_docs/python/tf/sparse/reshape)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

TF API

tensorflow_recommenders_addons/dynamic_embedding/python/ops/math_ops.py Outdated
def sparse_reshape(sp_input, shape, name=None):
"""Reshapes a `SparseTensor` to represent values in a new dense shape.

It do same things as `tf.sparse.reshape`. Here we provide GPU impl.
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

implement

tensorflow_recommenders_addons/dynamic_embedding/python/ops/math_ops.py
try:
return _sparse_reshape_gpu(sp_input, shape, name=name)
except errors.NotFoundError:
tf_logging.warn('`tfra.dynamic_embedding.sparse_reshape` is not'
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe this log will be printed too frequently on eager mode ? @Lifann

Copy link
Member

@Lifann Lifann Jun 2, 2021 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It only happens when trying to run the TFRA CPU version on GPU enabled environment. It reminds someone to install the GPU version when it is enabled. But you're right, it could be reconsidered to get better warning info.

Mr-Nineteen
Mr-Nineteen reviewed Jun 2, 2021
View reviewed changes
tensorflow_recommenders_addons/dynamic_embedding/core/BUILD
@@ -23,15 +23,23 @@ custom_op_library(
)

custom_op_library(
name = "_segment_reduction_ops.so",
name = "_math_ops.so",
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Change the so name, and also change the readme.

@funsimple funsimple force-pushed the gpu-dev branch 3 times, most recently from af8c11e to 7177565 Compare June 2, 2021 03:57
@rhdong
Copy link
Member

rhdong commented Jun 2, 2021

LGTM

@rhdong rhdong merged commit 2d53422 into tensorflow:master Jun 2, 2021
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@Lifann Lifann Lifann left review comments

@Mr-Nineteen Mr-Nineteen Mr-Nineteen left review comments

@rhdong rhdong rhdong left review comments

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
4 participants
@funsimple @rhdong @Mr-Nineteen @Lifann