[cute dsl] optimize cute dsl make_ptr perf (#1607)

Author: limin2021

flashinfer/cute_dsl/blockscaled_gemm.py

@@ -40,7 +40,8 @@
 import functools
 from cutlass._mlir import ir
 from cutlass.cute.nvgpu import cpasync, tcgen05
-from cutlass.cute.runtime import from_dlpack, make_ptr
+from cutlass.cute.runtime import from_dlpack
+
 from cutlass.cutlass_dsl import (
     Int32,
     Integer,
@@ -49,7 +50,7 @@
     new_from_mlir_values,
 )
 from cutlass.utils.static_persistent_tile_scheduler import WorkTileInfo
-from .utils import get_cutlass_dtype, cutlass_to_torch_dtype, get_num_sm
+from .utils import get_cutlass_dtype, cutlass_to_torch_dtype, get_num_sm, make_ptr
 from typing import Callable, List
 
 

flashinfer/cute_dsl/utils.py

@@ -14,10 +14,16 @@
 limitations under the License.
 """
 
+import ctypes
+import functools
+import importlib.util
+from typing import Union
+
 import cutlass
+import cutlass._mlir.dialects.cute as _cute_ir
 import torch
-import importlib.util
-import functools
+from cutlass._mlir import ir
+from cutlass.cute.typing import AddressSpace, Numeric, Pointer, Type
 
 
 def is_cute_dsl_available() -> bool:
@@ -67,3 +73,151 @@ def cutlass_to_torch_dtype(cutlass_dtype):
 def get_num_sm(device: torch.device) -> int:
     # get the compute capability of the device, which would be cached
     return torch.cuda.get_device_properties(device).multi_processor_count
+
+
+# WAR for CuTeDSL make_ptr implementation for flashinfer
+class _Pointer(Pointer):
+    """Runtime representation of a pointer that can inter-operate with
+    various data structures, including numpy arrays and device memory.
+
+    :param pointer: The pointer to the data
+    :type pointer: int or pointer-like object
+    :param dtype: Data type of the elements pointed to
+    :type dtype: Type
+    :param mem_space: Memory space where the pointer resides, defaults generic
+    :type mem_space: _cute_ir.AddressSpace, optional
+    :param assumed_align: Alignment of input pointer in bytes, defaults None
+    :type assumed_align: int, optional
+
+    :ivar _pointer: The underlying pointer
+    :ivar _dtype: Data type of the elements
+    :ivar _addr_space: Memory space of the pointer
+    :ivar _assumed_align: Alignment of the pointer in bytes
+    :ivar _desc: C-type descriptor for the pointer
+    :ivar _c_pointer: C-compatible pointer representation
+    """
+
+    def __init__(
+        self,
+        pointer,
+        dtype,
+        mem_space: _cute_ir.AddressSpace = _cute_ir.AddressSpace.generic,
+        assumed_align=None,
+    ):
+        self._pointer = pointer
+        self._dtype = dtype
+        self._addr_space = mem_space
+
+        if assumed_align is None:
+            self._assumed_align = dtype.width // 8
+        else:
+            self._assumed_align = assumed_align
+
+        self._desc = None
+        self._c_pointer = None
+        assert int(self._pointer) % self._assumed_align == 0, (
+            f"pointer must be {self._assumed_align} bytes aligned"
+        )
+
+    def size_in_bytes(self) -> int:
+        return ctypes.sizeof(ctypes.c_void_p(int(self._pointer)))
+
+    def __get_mlir_types__(self):
+        return [self.mlir_type]
+
+    def __c_pointers__(self):
+        if self._c_pointer is None:
+            self._desc = ctypes.c_void_p(int(self._pointer))
+            self._c_pointer = ctypes.addressof(self._desc)
+        return [self._c_pointer]
+
+    def __new_from_mlir_values__(self, values):
+        assert len(values) == 1
+        return values[0]
+
+    # Move mlir Type out of __init__ to decouple with mlir Context
+    @property
+    def mlir_type(self) -> ir.Type:
+        return _cute_ir.PtrType.get(
+            self._dtype.mlir_type, self._addr_space, self._assumed_align
+        )
+
+    @property
+    def dtype(self) -> Type[Numeric]:
+        return self._dtype
+
+    @property
+    def memspace(self):
+        return self._addr_space
+
+    def align(self, min_align: int, *, loc=None, ip=None) -> Pointer:
+        raise NotImplementedError("align is not supported in runtime")
+
+    def verify(self, expected_py_type):
+        # if expected_py_type is Pointer:
+        #     return True
+        # elif isinstance(expected_py_type, ir.Value) and expected_py_type.ty is Pointer:
+        #     return True
+        if expected_py_type is Pointer or (
+            isinstance(expected_py_type, ir.Value) and expected_py_type.ty is Pointer
+        ):
+            return True
+
+        return False
+
+    def __str__(self) -> str:
+        return f"Ptr<0x{int(self._pointer):016x}@{self._addr_space}>"
+
+    def __repr__(self):
+        return self.__str__()
+
+
+def make_ptr(
+    dtype: Type[Numeric],
+    value: Union[int, ctypes._Pointer],
+    mem_space: AddressSpace = AddressSpace.generic,
+    assumed_align=None,
+) -> Pointer:
+    """Create a pointer from a memory address
+
+    :param dtype: Data type of the pointer elements
+    :type dtype: Type[Numeric]
+    :param value: Memory address as integer or ctypes pointer
+    :type value: Union[int, ctypes._Pointer]
+    :param mem_space: Memory address space, defaults to AddressSpace.generic
+    :type mem_space: AddressSpace, optional
+    :param assumed_align: Alignment in bytes, defaults to None
+    :type assumed_align: int, optional
+    :return: A pointer object
+    :rtype: Pointer
+
+    .. code-block:: python
+
+        import numpy as np
+        import ctypes
+
+        from cutlass import Float32
+        from cutlass.cute.runtime import make_ptr
+
+        # Create a numpy array
+        a = np.random.randn(16, 32).astype(np.float32)
+
+        # Get pointer address as integer
+        ptr_address = a.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
+
+        # Create pointer from address
+        y = make_ptr(cutlass.Float32, ptr_address)
+    """
+    # check if value is int or ctypes.POINTER
+    if isinstance(value, int):
+        address_value = value
+    elif isinstance(value, ctypes._Pointer):
+        # get address value
+        address_value = ctypes.cast(value, ctypes.c_void_p).value
+        assert address_value is not None, "Pointer address is None"
+    else:
+        raise TypeError(
+            f"Expect int or ctypes.POINTER for value but got {type(value)=}"
+        )
+
+    return _Pointer(address_value, dtype, mem_space, assumed_align=assumed_align)