mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-19 05:16:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

[PJRT:C] Rename PJRT_Structure_* to PJRT_Extension_*

Browse Source 
This is to make it more clear what the APIs pertain to (extensions to
the PJRT C API). This is a source-incompatible but ABI compatible
change. It should only affect plugins that use extensions.

This is part of the PJRT C API extension improvements described in
https://docs.google.com/document/d/1PsSIkCbHujPfiZphyQ3U5wB9VGBLKBjxGrTtmziJwEw/edit?usp=sharing

PiperOrigin-RevId: 606802550
This commit is contained in:
Skye Wanderman-Milne 2024-02-13 17:42:17 -08:00 committed by jax authors
parent b9c7296754
commit cbe23b43f7
1 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
jaxlib/cuda_plugin_extension.cc
View File

@ -44,11 +44,11 @@ Status RegisterCustomCallTarget(const PJRT_Api* c_api, nb::str fn_name,
if (c_api->extension_start == nullptr) {
return Unimplemented("The plugin does not have extension.");
}
const PJRT_Structure_Base* next =
reinterpret_cast<const PJRT_Structure_Base*>(c_api->extension_start);
const PJRT_Extension_Base* next =
reinterpret_cast<const PJRT_Extension_Base*>(c_api->extension_start);
while (next != nullptr &&
next->type !=
PJRT_Structure_Type::PJRT_Structure_Type_Gpu_Custom_Call) {
PJRT_Extension_Type::PJRT_Extension_Type_Gpu_Custom_Call) {
next = next->next;
}
if (next == nullptr) {