def __missing__(self, key):
translation_tables = [
_backend_specific_translations[p]
for p in xb.expand_platform_alias(key)
]
ret = self[key] = _TranslationRuleAdapter(translation_tables,
_wrap_old_translation)
return ret
def check_backend_matches(inner_backend, outer_backend):
# For nested calls, the outermost call sets the backend for all inner calls;
# it's an error if the inner call has a conflicting explicit backend spec.
if inner_backend is None:
return
if (inner_backend != outer_backend and
outer_backend not in xb.expand_platform_alias(inner_backend)):
raise ValueError(
f"Outer-jit backend specification {outer_backend} must match explicit "
f"inner-jit backend specification {inner_backend}.")
def register_translation(prim: core.Primitive, rule: TranslationRule, *,
platform: Optional[str] = None) -> None:
if platform is None:
_translations[prim] = rule
else:
# For backward compatibility reasons, we allow rules to be registered
# under "gpu" even though the platforms are now called "cuda" and "rocm".
# TODO(phawkins): fix up users to specify either "cuda" or "rocm" and remove
# this expansion.
for p in xb.expand_platform_alias(platform):
_backend_specific_translations[p][prim] = rule