mirror of
https://github.com/NVIDIA/cutlass.git
synced 2026-07-11 17:48:57 +00:00
119 lines
4.1 KiB
Python
119 lines
4.1 KiB
Python
# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
|
|
#
|
|
# Use of this software is governed by the terms and conditions of the
|
|
# NVIDIA End User License Agreement (EULA), available at:
|
|
# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
|
|
#
|
|
# Any use, reproduction, disclosure, or distribution of this software
|
|
# and related documentation outside the scope permitted by the EULA
|
|
# is strictly prohibited.
|
|
|
|
import enum
|
|
import types
|
|
from collections.abc import Callable
|
|
|
|
|
|
class DialectAutoConvertProxy:
|
|
"""
|
|
Proxy that wraps a raw MLIR dialect module, auto-converting DSL types
|
|
(anything with an ``.ir_value()`` method) to ``ir.Value`` when calling
|
|
dialect operations.
|
|
|
|
This enables users to write cleaner code without explicit
|
|
``.ir_value()`` calls::
|
|
|
|
# Before (raw dialect module):
|
|
nvvm.shfl_sync(T.i32(), Int32(mask).ir_value(), ...)
|
|
|
|
# After (proxied):
|
|
nvvm.shfl_sync(T.i32(), Int32(mask), ...)
|
|
|
|
Non-callable attributes and enum classes are passed through unchanged
|
|
so that attribute access like ``nvvm.ShflKind.idx`` still works.
|
|
|
|
Parameters
|
|
----------
|
|
dialect_module
|
|
The raw MLIR dialect module to wrap
|
|
(e.g. ``cutlass._mlir.dialects.nvvm``).
|
|
"""
|
|
|
|
def __init__(self, dialect_module: types.ModuleType) -> None:
|
|
self._module = dialect_module
|
|
self._wrapped_cache: dict[str, Callable[..., object]] = {}
|
|
|
|
@staticmethod
|
|
def _convert_arg(
|
|
arg: object,
|
|
loc: object | None,
|
|
ip: object | None,
|
|
) -> object:
|
|
"""Convert DSL value wrappers to ir.Value; pass primitives/containers through.
|
|
|
|
Numeric primitives (int/float/bool) are immediates/attributes, not
|
|
operands, so they pass through unchanged.
|
|
DSL ``Numeric`` types are not int/float subclasses, so genuine operands
|
|
still convert below.
|
|
"""
|
|
if isinstance(arg, (bool, int, float)):
|
|
return arg
|
|
if hasattr(arg, "ir_value") and callable(arg.ir_value):
|
|
try:
|
|
return arg.ir_value(loc=loc, ip=ip)
|
|
except TypeError:
|
|
# Some ir_value() methods (e.g. Array) don't accept loc/ip.
|
|
return arg.ir_value()
|
|
if isinstance(arg, (list, tuple)):
|
|
converted = [
|
|
DialectAutoConvertProxy._convert_arg(item, loc, ip) for item in arg
|
|
]
|
|
return type(arg)(converted)
|
|
return arg
|
|
|
|
def __getattr__(self, name: str) -> object:
|
|
attr = getattr(self._module, name)
|
|
|
|
# Non-callable attributes and enum classes pass through
|
|
# unchanged. Enum classes need attribute access (e.g.
|
|
# ShflKind.idx), but MLIR operation classes should be
|
|
# wrapped for argument conversion.
|
|
if not callable(attr) or isinstance(attr, enum.EnumMeta):
|
|
return attr
|
|
|
|
# Use cache for wrapped callables
|
|
if name not in self._wrapped_cache:
|
|
|
|
def _make_wrapper(
|
|
func: Callable[..., object],
|
|
) -> Callable[..., object]:
|
|
def wrapped(
|
|
*args: object,
|
|
loc: object | None = None,
|
|
ip: object | None = None,
|
|
**kwargs: object,
|
|
) -> object:
|
|
converted_args = tuple(
|
|
DialectAutoConvertProxy._convert_arg(arg, loc, ip)
|
|
for arg in args
|
|
)
|
|
converted_kwargs = {
|
|
k: DialectAutoConvertProxy._convert_arg(v, loc, ip)
|
|
for k, v in kwargs.items()
|
|
}
|
|
return func(
|
|
*converted_args,
|
|
loc=loc,
|
|
ip=ip,
|
|
**converted_kwargs,
|
|
)
|
|
|
|
return wrapped
|
|
|
|
self._wrapped_cache[name] = _make_wrapper(attr)
|
|
|
|
return self._wrapped_cache[name]
|
|
|
|
def __dir__(self) -> list[str]:
|
|
return dir(self._module)
|