Skip to content

[webgpu] Further tweak vectorized NaN handling in binary ops #7185

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

Merged
merged 3 commits into from
Dec 21, 2022

Conversation

hujiajie
Copy link
Contributor

@hujiajie hujiajie commented Dec 17, 2022

The big picture will be:

  fn binaryOperation(a : vec4<f32>, b : vec4<f32>) -> vec4<f32> {
    let isNaN = false;
    {
      // op-specific code begins
      let result = ...;
      let isNaN = ...;  // optional, declare this only if the op defines
                        // its own extra rules about NaN, say mod(0, 0)
      // op-specific code ends

      return select(result, vec4<f32>(valueForNaN),
                    vec4<bool>(isNaN) | isnanVec4(a) | isnanVec4(b));
    }
  }

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


This change is Reviewable

@hujiajie hujiajie marked this pull request as draft December 17, 2022 14:39
The big picture will be:

```
  fn binaryOperation(a : vec4<f32>, b : vec4<f32>) -> vec4<f32> {
    let isNaN = false;
    {
      // op-specific code begins
      let result = ...;
      let isNaN = ...;  // optional, declare this only if the op defines
                        // its own extra rules about NaN, say mod(0, 0)
      // op-specific code ends

      return select(result, vec4<f32>(valueForNaN),
                    vec4<bool>(isNaN) | isnanVec4(a) | isnanVec4(b));
    }
  }
```
@hujiajie hujiajie marked this pull request as ready for review December 18, 2022 04:49
return select(resultTemp, vec4<f32>(uniforms.NAN), isNaN);
let valueForNaN = uniforms.NAN;
${CHECK_NAN_SNIPPET_VEC4}
return resultTemp;
Copy link
Contributor

Choose a reason for hiding this comment

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

The original code only checks isNan. Here you also add isnanVec4(a) | isnanVec4(b). Is it expected?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

It's almost certain binOp(NaN, b) or binOp(a, NaN) will return valueForNaN. I guess previously this case was just quietly covered by the pow() builtin.

@@ -87,7 +87,10 @@ export class BinaryOpProgram implements WebGPUProgram {
const dType = this.isVec4 ? 'vec4<f32>' : 'f32';
const opFnStr = `
fn binaryOperation(a : ${dType}, b : ${dType}) -> ${dType} {
${getBinaryOpString(this.op, this.isVec4)}
let isNaN = false;
Copy link
Contributor

Choose a reason for hiding this comment

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

Can you add a comment to isNaN in the code? Like you did in the commit message. And in MOD_VEC4 and POW_VEC4. You can directly use isNaN = ....

     // Define op own extra rules about NaN, say mod(0, 0).
      var isNaN = ${this.isVec4 ?  'vec4<bool>(false)' : 'false'};

Copy link
Contributor Author

Choose a reason for hiding this comment

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

I feel it's better to have a combined comment after valueForNaN is restructured similarly, or even later after everything calms down.

The let before isNaN in op-specific code is an intentional design. IMO, the snippet looks more self-contained that way, so I won't wonder where is the declaration of this isNaN identifier while reading MOD_VEC4.

Copy link
Contributor

Choose a reason for hiding this comment

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

Sounds good.

@qjia7 qjia7 requested a review from gyagp December 19, 2022 08:24
Copy link
Contributor

@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.

Still LGTM, thanks.

@@ -87,7 +87,10 @@ export class BinaryOpProgram implements WebGPUProgram {
const dType = this.isVec4 ? 'vec4<f32>' : 'f32';
const opFnStr = `
fn binaryOperation(a : ${dType}, b : ${dType}) -> ${dType} {
${getBinaryOpString(this.op, this.isVec4)}
let isNaN = false;
Copy link
Contributor

Choose a reason for hiding this comment

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

Sounds good.

Copy link
Contributor

@xhcao xhcao left a comment

Choose a reason for hiding this comment

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

LGTM, thanks.

Copy link

@gyagp gyagp left a comment

Choose a reason for hiding this comment

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

LGTM

@gyagp gyagp merged commit 2c95a67 into tensorflow:master Dec 21, 2022
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

4 participants