microsoft / microsoft/onnxruntime

[Feature Request] Add CUDA kernel for the ScatterElements operator in opset 18

Open
#18,381 6 comments 2 reactions 0 assignees View on GitHub
ep:CUDA feature request
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.