[js/webgpu] Fix NAN caused by un-initialized buffer in instance-norm

[axinging]

The added case will be NAN because of the un-initialized buffer.

[axinging]

@qjia7 @xhcao @hujiajie @jzm-intel @gyagp PTAL

[gyagp]

@axinging Please hold this a bit and I will discuss further with you.

[fs-eire]

I think this fix is OK and we should go merge this fix. Let me explain it a little bit about my understanding.

The current implementation of InstanceNormalization is a multi-step implementation. It uses 3 shaders:

• InstanceNormComputeMean
  This program takes the original input (shape=[N, H, W, C]) with scale and bias, and generate a temporary intermediate output (shape=[N, C, 64, 2]
• InstanceNormComputeChannelScaleShift
  This program takes the output of the previous step as input with scale and bias, and generate a temporary intermediate output (shape=[N, C, 2])
• InstanceNormalizationNHWC
  This program takes the output of the previous step as input with scale and bias and generate the final output.

The problem is, the shader of step.1 does not always write into all the memory that it claims for the output, and step.2 uses the value that step.1 does not write to. This means any garbage data can be loaded into step.2 and generate wrong result.

This PR may fix this issue. But eventually we may need a better implementation that either we don't need 3 shaders (use workerBarriers maybe) or generate shader that make sure it writes to all memory that covers the output data. For now, it's better than a broken kernel.

@guschmue @satyajandhyala

[guschmue]

/azp run ONNX Runtime Web CI Pipeline

[guschmue]

/azp run Linux CPU CI Pipeline,Linux CPU Minimal Build E2E CI Pipeline,Linux GPU CI Pipeline,Linux GPU TensorRT CI Pipeline,Linux OpenVINO CI Pipeline,Linux QNN CI Pipeline,MacOS CI Pipeline,Windows ARM64 QNN CI Pipeline,Windows CPU CI Pipeline

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[guschmue]

/azp run Windows GPU CI Pipeline,Windows GPU TensorRT CI Pipeline,onnxruntime-binary-size-checks-ci-pipeline,orttraining-linux-ci-pipeline,orttraining-linux-gpu-ci-pipeline,orttraining-ortmodule-distributed,Windows x64 QNN CI Pipeline

[azure-pipelines]

Azure Pipelines could not run because the pipeline triggers exclude this branch/path.

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[fs-eire]

/azp run Big Models

[azure-pipelines]

Azure Pipelines could not run because the pipeline triggers exclude this branch/path.

[gyagp]

I think this fix is OK and we should go merge this fix. Let me explain it a little bit about my understanding.

The current implementation of InstanceNormalization is a multi-step implementation. It uses 3 shaders:

• InstanceNormComputeMean
  This program takes the original input (shape=[N, H, W, C]) with scale and bias, and generate a temporary intermediate output (shape=[N, C, 64, 2]
• InstanceNormComputeChannelScaleShift
  This program takes the output of the previous step as input with scale and bias, and generate a temporary intermediate output (shape=[N, C, 2])
• InstanceNormalizationNHWC
  This program takes the output of the previous step as input with scale and bias and generate the final output.

The problem is, the shader of step.1 does not always write into all the memory that it claims for the output, and step.2 uses the value that step.1 does not write to. This means any garbage data can be loaded into step.2 and generate wrong result.

This PR may fix this issue. But eventually we may need a better implementation that either we don't need 3 shaders (use workerBarriers maybe) or generate shader that make sure it writes to all memory that covers the output data. For now, it's better than a broken kernel.

@guschmue @satyajandhyala

My intuition is we should avoid reading out-of-scope data in step2, instead of filling them with all 0s in step1. Any concerns we can't do this?

[fs-eire]

I think this fix is OK and we should go merge this fix. Let me explain it a little bit about my understanding.
The current implementation of InstanceNormalization is a multi-step implementation. It uses 3 shaders:

• InstanceNormComputeMean
  This program takes the original input (shape=[N, H, W, C]) with scale and bias, and generate a temporary intermediate output (shape=[N, C, 64, 2]
• InstanceNormComputeChannelScaleShift
  This program takes the output of the previous step as input with scale and bias, and generate a temporary intermediate output (shape=[N, C, 2])
• InstanceNormalizationNHWC
  This program takes the output of the previous step as input with scale and bias and generate the final output.

The problem is, the shader of step.1 does not always write into all the memory that it claims for the output, and step.2 uses the value that step.1 does not write to. This means any garbage data can be loaded into step.2 and generate wrong result.
This PR may fix this issue. But eventually we may need a better implementation that either we don't need 3 shaders (use workerBarriers maybe) or generate shader that make sure it writes to all memory that covers the output data. For now, it's better than a broken kernel.
@guschmue @satyajandhyala

My intuition is we should avoid reading out-of-scope data in step2, instead of filling them with all 0s in step1. Any concerns we can't do this?

We absolutely should do this. Usually I agree with "avoid reading out-of-scope data". But in this specific implementation, the step.1 creates the output in shape [N, C, 64, 2], regardless of the W and H from the original input. The dim[2] is fixed to 64 for whatever reason (maybe to simplify the implementation).

Since step.1 write to [N, C, 64, 2] and step.2 read from the same sized buffer as input, there can be another point of view that the whole NxCx128 data is "in the scope". It's the internal steps of the operator's implementation anyway, so it's up to who implement the operator as long as the read/write range should match. So in this POV, fixing step.1 to make sure writting to all required memory is also a valid fix.

But you reminded me that I didn't check whether uniforms.H equals to N*C*128/sizeof(input_elem). This is important otherwise step.1 writes to out-of-scope data.

[axinging]

@fs-eire, @gyagp Boundary check added, PTAL.

But I think the instance-norm (NHWC) has room for improvement: It wastes a lot of memory and GPU compute unit.
For memory: [n, c, 64, 2] is a huge waste, indeed [n, c, 1, 2] should be enough for this case.
For compute unit, it scheduled n * c / components * 64 threads, I test that n * c / components * 1 also works.
I will do some perf test on this and if possible, I will fix this in another PR.

[fs-eire]

/azp run Windows ARM64 QNN CI Pipeline,Windows x64 QNN CI Pipeline,Windows CPU CI Pipeline,Windows GPU CI Pipeline,Windows GPU TensorRT CI Pipeline,ONNX Runtime Web CI Pipeline,Linux CPU CI Pipeline,Linux CPU Minimal Build E2E CI Pipeline,Linux GPU CI Pipeline,Linux GPU TensorRT CI Pipeline

[fs-eire]

/azp run Linux OpenVINO CI Pipeline,Linux QNN CI Pipeline,MacOS CI Pipeline,orttraining-amd-gpu-ci-pipeline,orttraining-linux-ci-pipeline,orttraining-linux-gpu-ci-pipeline,orttraining-ortmodule-distributed,onnxruntime-binary-size-checks-ci-pipeline,Big Models,Android CI Pipeline

[fs-eire]

/azp run iOS CI Pipeline,ONNX Runtime React Native CI Pipeline

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).