wehub-resource-sync
78 lines
2.5 KiB
Python
78 lines
2.5 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
"""Codegen registry for CUDA HW ops.
|
|
|
|
User-facing Python wrappers are hand-written in :mod:`tvm.tirx.op` so that
|
|
editors / static analyzers (Cursor, Pyright) can see their signatures. This
|
|
module only handles the backend codegen side.
|
|
"""
|
|
|
|
import functools
|
|
|
|
import tvm_ffi
|
|
|
|
CODEGEN_REGISTRY = {}
|
|
|
|
|
|
def _canonical_device_intrin_name(op_name: str) -> str:
|
|
if not op_name.startswith("tirx."):
|
|
return op_name
|
|
basename = op_name[len("tirx.") :]
|
|
if "." in basename:
|
|
return op_name
|
|
for prefix, namespace in (
|
|
("cuda_", "cuda"),
|
|
("ptx_", "ptx"),
|
|
("nvshmem_", "nvshmem"),
|
|
("nki_", "nki"),
|
|
):
|
|
if basename.startswith(prefix):
|
|
return f"tirx.{namespace}.{basename[len(prefix) :]}"
|
|
return op_name
|
|
|
|
|
|
@tvm_ffi.register_global_func("tirx.intrinsics.cuda.get_codegen")
|
|
def get_codegen(op):
|
|
"""get the codegen function for a given op"""
|
|
return CODEGEN_REGISTRY.get(op, None)
|
|
|
|
|
|
def register_codegen(op, backend="cuda"):
|
|
"""Register a codegen function for a given op.
|
|
|
|
The codegen function should return a ``cuda_func_call`` statement, and
|
|
optionally a list of tags that the codegen function needs.
|
|
"""
|
|
|
|
def decorator(func):
|
|
full_op_name = "tirx." + op
|
|
canonical_op_name = _canonical_device_intrin_name(full_op_name)
|
|
op_names = {full_op_name, canonical_op_name}
|
|
|
|
@functools.wraps(func)
|
|
def wrapper(arg_list):
|
|
res = func(*arg_list) # pylint: disable=not-callable
|
|
if isinstance(res, tuple):
|
|
return res[0], res[1]
|
|
return res, list()
|
|
|
|
for op_name in op_names:
|
|
CODEGEN_REGISTRY[op_name] = wrapper
|
|
return wrapper
|
|
|
|
return decorator
|