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

[webgpu] Use i32 as coords type in WGSL #5628

Merged
merged 6 commits into from Sep 17, 2021
Merged

Conversation

axinging
Copy link
Contributor

@axinging axinging commented Sep 15, 2021

To see the logs from the Cloud Build CI, please join either our discussion or announcement mailing list.


This change is Reviewable

@google-cla google-cla bot added the cla: yes label Sep 15, 2021
@axinging axinging force-pushed the use_i32 branch 3 times, most recently from 2fbd478 to 4b8d5be Compare September 15, 2021 08:42
@axinging axinging marked this pull request as ready for review September 15, 2021 08:46
@axinging
Copy link
Contributor Author

axinging commented Sep 15, 2021

@qjia7 @haoyunfeix @xhcao @gyagp PTAL

A comment, we can use

  1. getXXXAtOutCoordsByGlobalId(vec3(globalId), index) to get data.
  2. Or we can define an extra globalIdInt:
let globalId = vec3<i32>globalIdU32;

Then use getXXXAtOutCoordsByGlobalId(globalId , index) to get data.

As to the option 1, it will save a vec3 in each program. As to the option 2, it will save some code changes. Current 1 is applied.

@axinging axinging force-pushed the use_i32 branch 2 times, most recently from cd0bbab to d2c56ff Compare September 16, 2021 02:17
Copy link
Collaborator

@qjia7 qjia7 left a comment

Choose a reason for hiding this comment

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

I vote keeping globalId, localId's type as function parameter types. And do the corresponding cast as needed. And for others, we all unify to use i32 instead of u32 expect for dispatchSize. DispatchSize or numWorkGroups is similarly like globalId, localId which is u32 instead i32.

tfjs-backend-webgpu/src/kernels/conv2d_mm_webgpu.ts Outdated Show resolved Hide resolved
tfjs-backend-webgpu/src/shader_preprocessor_wgsl.ts Outdated Show resolved Hide resolved
tfjs-backend-webgpu/src/shader_preprocessor_wgsl.ts Outdated Show resolved Hide resolved
Copy link
Collaborator

@qjia7 qjia7 left a comment

Choose a reason for hiding this comment

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

Looks better, thanks. LGTM with two nits.

Copy link
Collaborator

@lina128 lina128 left a comment

Choose a reason for hiding this comment

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

Thank you for the refactor, looks great!

Reviewable status: :shipit: complete! 2 of 1 approvals obtained (waiting on @axinging and @pyu10055)

@qjia7 qjia7 merged commit a9da3ff into tensorflow:master Sep 17, 2021
@axinging axinging deleted the use_i32 branch September 22, 2021 06:50
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

3 participants