Skip to content
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

Make MakeBatchPointers use the kernel registry #23853

Open
wants to merge 1 commit into
base: main
Choose a base branch
from

Conversation

copybara-service[bot]
Copy link

@copybara-service copybara-service bot commented Mar 18, 2025

Make MakeBatchPointers use the kernel registry

  • Moves MakeBatchPointers logic into backends/gpu/runtime since it's a runtime component.
  • Defines trait for the MakeBatchPointers kernel in stream_executor/gpu/
  • Moves the implementations of this kernel into stream_executor/{cuda|rocm} and registers them with the registry.
  • Makes MakeBatchPointers retrieve the kernel by using the kernel registry.
  • Add the kernel implementations as dependencies to the all_runtime targets for CUDA and ROCm.

@copybara-service copybara-service bot force-pushed the test_737892623 branch 4 times, most recently from c14ff31 to f51d6b8 Compare March 20, 2025 07:36
@copybara-service copybara-service bot force-pushed the test_737892623 branch 4 times, most recently from 7a54aca to d911d9a Compare April 2, 2025 07:43
- Moves `MakeBatchPointers` logic into backends/gpu/runtime since it's a runtime component.
- Defines trait for the MakeBatchPointers kernel in stream_executor/gpu/
- Moves the implementations of this kernel into stream_executor/{cuda|rocm} and registers them with the registry.
- Makes `MakeBatchPointers` retrieve the kernel by using the kernel registry.
- Add the kernel implementations as dependencies to the `all_runtime` targets for CUDA and ROCm.

PiperOrigin-RevId: 737892623
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.

1 participant