Helper methods for WebGPU runtime

[visheratin]

Hi, @fs-eire. In the TensorView interface definition, there is the getFloat32Array() function. The comment says that the tensor data must be on CPU. For implementing the Reshape operator, I wanted to create an additional function that looks like this:

getInt64Array(): BigInt64Array {
    return new BigInt64Array(this.module.HEAP8.buffer, this.data, ShapeUtil.size(this.dims));
}

The motivation is that in the spec, the shape is a tensor input. The idea was to get the new shape from the input tensor and implement Reshape along these lines:

export const reshape = (context: ComputeContext, attributes: ReshapeAttributes): void => {
  const inputs = context.inputs;
  if (!inputs || inputs.length !== 2) {
    throw new Error('Reshape requires 2 inputs  - input tensor and shape tensor.');
  }
  const newDims = inputs[1].getInt64Array();
  validateInputs(inputs[0].dims, newDims);
  context.compute({
    ...reshapeProgramMetadata,
    cacheHint: attributes.cacheKey,
    get: () => createReshapeProgramInfo(context.inputs[0], newDims)
  });
};

The problem is that during execution, the tensors are located in GPU memory, and the tensor views we are working with are like pointers to them (if my understanding of JSEP is correct). In this case, considering the comment in the TensorView definition, the getInt64Array() method won't work because the tensor data is not in the buffer of the WASM module. Could you please help to better understand how this aspect of runtime works?

[visheratin]

An alternative for that would be to run the checks inside the WGSL function and raise an error in case of mismatching shapes, but in this case, it is not clear how to get dims for the TensorInfo output.

[fs-eire]

@visheratin I appreciated your help on trying to implement webgpu operators. However, onnxruntime-web WebGPU is implemented using JSEP (JavaScript Execution Provider, some design overview in PR summary), so operators like Reshape, Squeeze, Unsqueeze are actually only updating tensor metadata instead of modifying any tensor data. If you try to debug onnxruntime-web with a model with Reshape operator, you will find that it does not run into the JavaScript code.

You can see that in current code, Reshape is registered here:

onnxruntime/onnxruntime/core/providers/js/js_execution_provider.cc

Lines 249 to 251 in 896a963

	BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 5, 12, Reshape)>,
	BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 13, 13, Reshape)>,
	BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 14, Reshape)>,

its implementation is here: https://github.com/microsoft/onnxruntime/blob/896a963492fc406866a796eeb9c4a422f2b2f6e2/onnxruntime/core/providers/js/operators/reshape.cc
https://github.com/microsoft/onnxruntime/blob/896a963492fc406866a796eeb9c4a422f2b2f6e2/onnxruntime/core/providers/js/operators/reshape.h

In the kernel declaration, we have this: .Alias(0, 0) so this means that we assume Reshape operator's output[0] is an alias of input[0], which allows the allocator to create OrtValue for output[0] as the same data that used as input[0]. This part could be a little bit difficult to understand ( ORT's allocator and graph partitioner have quite complicated design ); so in short - Reshape is considered as a no-op for the tensor data and is handled all inside C++ code.

[fs-eire]

according to ORT's design, tensor's data may be on CPU or on GPU; but tensor's metadata (dims/type) will always be on CPU. As Reshape only modify the dims, usually it should not need a GPU kernel.

[visheratin]

Thank you for the detailed answer! I have a number of follow-up questions:

1. What does InputMemoryType(OrtMemTypeCPU, 1) mean in the reshape declaration? Is it a condition that the shape tensor must be on CPU or instruction to copy the shape tensor to CPU prior to the execution?
2. When saying that Reshape is a no-op, do you mean that it doesn't touch the data that is on GPU? Do I assume right that because of .Alias(0, 0), target and source here are always the same?
3. Do I understand right that Squeeze and Unsqueeze operators should be implemented in a similar manner?

[fs-eire]

A1 - it's enforcing the data of input[1] ( which is the shape ) to be on CPU. so the framework ensures that input[1] will be on CPU before executing operator Reshape.
A2 - I don't fully understand this part. My understanding is that .Alias(0, 0) will let ORT to try to make input[0] and output[0] to use the same data. Not sure if that is always same or not. So in code I also write a GPU to GPU copy if they are not same.
A3 - yes.

[visheratin]

Got it. The last question is about testing new operators. What is the best way to test/debug WebGPU-related code? I tried to look up some instructions in the contribution guidelines but couldn't find them.

[fs-eire]

Usually I launch npm test in folder /js/web with '--debug' flag to debug the typescript code.

[fs-eire]

I don't add a getInt64Array function to TensorView type because it is for representing a GPU data, and we don't have the requirement to download (u)int64 from GPU.

[visheratin]

Is the getFloat32Array() method used for extracting the data from GPU?

[fs-eire]

No. getFloat32Array() is used in JS code when we know the corresponding tensor represented by the TensorView has its data on CPU. It does not trigger any data transfering between CPU and GPU.