microsoft / microsoft/onnxruntime
[Feature Request] Add CUDA kernel for the ScatterElements operator in opset 18
- Dominant language
- C++
- Stars
- 21.9k
- Forks
- 4.2k
- Avg merge
- 4d 11h
- Merged PRs (30d)
- 184
Description
### Describe the feature request
It seems that the operator `ScatterElements` in not implemented in CUDA when using opset 16,17 or 18 and "add" reduction.
We get a message `CUDA kernel not found in registries for Op type: ScatterElements node name: /ScatterElements` in the log when loading the onnx model and "CUDAExecutionProvider" in the profiling file.
Note that the operator` ScatterElements` is available when using opset 15 but provides wrong results (see https://github.com/onnx/onnx/pull/3484)
Here is some minimal python code to reproduce the problem using `torch==2.0.0+cu118` or `torch==2.1.0+cu118` and `onnxruntime-gpu==1.16.2` using a NVIDIA GeForce GTX 1050
```
import io
import json
import numpy as np
import onnxruntime
import torch
import torch.nn as nn
class ScatterAdd(nn.Module):
"""Point cloud renderer"""
def __init__(self, length: int):
super(ScatterAdd, self).__init__()
self.length = length
def forward(self, indices: torch.Tensor, weights: torch.Tensor) -> torch.Tensor:
result = torch.zeros((self.length), device="cuda", dtype=torch.float)
result = result.scatter_add(0, indices, weights)
return result
def main():
opset_version = 16
onnx_provider = "CUDAExecutionProvider"
indices = torch.Tensor([0, 0, 1, 2]).long().cuda()
weights = torch.Tensor([1.0, 3.0, 5.0, 7.0]).cuda()
scatter_add = ScatterAdd(length=3)
result = scatter_add(indices=indices, weights=weights)
assert np.allclose(result.cpu().numpy(), [4.0, 5.0, 7.0])
bytes_io = io.BytesIO()
torch.onnx.export(
scatter_add, (indices, weights), bytes_io, opset_version=opset_version, input_names=["indices", "weights"]
)
onnxruntime.set_default_logger_severity(1)
sess_options = onnxruntime.SessionOptions()
sess_options.enable_profiling = True
ort_session = onnxruntime.InferenceSession(
bytes_io.getvalue(),
providers=[
onnx_provider,
],
sess_options=sess_options,
)
# when using CUDAExecutionProvider with opset_version in 16, 17,=ir 18 getting in the log:
# CUDA kernel not found in registries for Op type: ScatterElements node name: /ScatterElements
numpy_inputs = {
"indices": np.array([0, 0, 1, 2], dtype=np.int64),
"weights": np.array([1.0, 3.0, 5.0, 7.0], dtype=np.float32),
}
result = ort_session.run(None, numpy_inputs)
prof_file = ort_session.end_profiling()
with open(prof_file) as f:
sess_time = json.load(f)
# fails when using opset_version=15 with both CUDAExecutionProvider
assert np.allclose(result, [4.0, 5.0, 7.0])
# fails when using opset_version=16 or 17 or 18 with both CUDAExecutionProvider
assert sess_time[3]["args"]["provider"] == "CUDAExecutionProvider"
if __name__ == "__main__":
main()
```
### Describe scenario use case
This is used in an image processing pipeline.
Contributor guide
Research direction
Start by running the provided Python reproducer with onnxruntime-gpu and the CUDAExecutionProvider, confirming the missing ScatterElements kernel for opsets 16–18 with add reduction. Then locate the ScatterElements CUDA implementation and registration, and compare its behavior with the opset 15 path. Done means the model runs on CUDA, produces [4.0, 5.0, 7.0], and profiling shows CUDAExecutionProvider.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp, python, pytorch
- Domain
- backend, machine-learning
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100