v4.3 tag release update. (#2789)

python/CuTeDSL/cutlass/base_dsl/_mlir_helpers/arith.py

@@ -689,7 +689,21 @@ def _min(lhs, rhs, *, loc=None, ip=None):
         if not is_dynamic_expression(rhs):
             rhs = arith.constant(lhs.type, rhs, loc=loc, ip=ip)
 
-    if arith._is_integer_like_type(lhs.type):
+    # Handle vector types
+    if isinstance(lhs.type, T.VectorType):
+        elem_type = lhs.type.element_type
+        if arith._is_integer_like_type(elem_type):
+            assert hasattr(
+                lhs, "signed"
+            ), "Should have attribute `signed`, must be a bug"
+            if lhs.signed != False:
+                return arith.minsi(lhs, rhs, loc=loc, ip=ip)
+            else:
+                return arith.minui(lhs, rhs, loc=loc, ip=ip)
+        else:
+            return arith.minimumf(lhs, rhs, loc=loc, ip=ip)
+    elif arith._is_integer_like_type(lhs.type):
+        assert hasattr(lhs, "signed"), "Should have attribute `signed`, must be a bug"
         if lhs.signed != False:
             return arith.minsi(lhs, rhs, loc=loc, ip=ip)
         else:
@@ -705,7 +719,6 @@ def _max(lhs, rhs, *, loc=None, ip=None):
     Assuming the operands have the same type
     """
     from ..dsl import is_dynamic_expression
-
     if not is_dynamic_expression(lhs):
         if not is_dynamic_expression(rhs):
             return max(lhs, rhs)
@@ -714,8 +727,21 @@ def _max(lhs, rhs, *, loc=None, ip=None):
     else:
         if not is_dynamic_expression(rhs):
             rhs = arith.constant(lhs.type, rhs, loc=loc, ip=ip)
-
-    if arith._is_integer_like_type(lhs.type):
+    # Handle vector types
+    if isinstance(lhs.type, T.VectorType):
+        elem_type = lhs.type.element_type
+        if isinstance(elem_type, ir.IntegerType):
+            assert hasattr(
+                lhs, "signed"
+            ), "Should have attribute `signed`, must be a bug"
+            if lhs.signed != False:
+                return arith.maxsi(lhs, rhs, loc=loc, ip=ip)
+            else:
+                return arith.maxui(lhs, rhs, loc=loc, ip=ip)
+        else:
+            return arith.maximumf(lhs, rhs, loc=loc, ip=ip)
+    elif arith._is_integer_like_type(lhs.type):
+        assert hasattr(lhs, "signed"), "Should have attribute `signed`, must be a bug"
         if lhs.signed != False:
             return arith.maxsi(lhs, rhs, loc=loc, ip=ip)
         else:

python/CuTeDSL/cutlass/base_dsl/arch.py

@@ -11,7 +11,7 @@
 
 from enum import Enum
 import re
-from typing import Callable, List
+from typing import Callable, List, Tuple
 
 
 class Arch(Enum):
@@ -49,16 +49,47 @@ class Arch(Enum):
 
     @classmethod
     def _missing_(cls, value):
-        """Support creating Arch enum from (major, minor, suffix) tuple"""
-        if not isinstance(value, tuple):
-            raise ValueError(f"invalid arguments for Arch: {value}")
-        major, minor, suffix = None, None, None
-        if len(value) == 2:
+        if isinstance(value, tuple) and len(value) == 2:
+            # Support creating Arch enum from (major, minor) tuple
+            # Arch(major, minor) is equivalent to Arch(major, minor, "")
             major, minor, suffix = *value, ""
+            return cls((major, minor, suffix))
         else:
             raise ValueError(f"invalid arguments for Arch: {value}")
+        
 
-        return cls(major, minor, suffix)
+    # attributes to get arch list of specific families
+    @classmethod
+    def AmpereArchs(cls) -> Tuple["Arch"]:
+        return (Arch.sm_80, Arch.sm_86, Arch.sm_87)
+
+    @classmethod
+    def AdaArchs(cls) -> Tuple["Arch"]:
+        return (Arch.sm_89,)
+
+    @classmethod
+    def HopperArchs(cls) -> Tuple["Arch"]:
+        return (Arch.sm_90, Arch.sm_90a)
+
+    @classmethod
+    def BlackwellArchs(cls) -> Tuple["Arch"]:
+        return (
+            Arch.sm_100,
+            Arch.sm_100a,
+            Arch.sm_100f,
+            Arch.sm_101,
+            Arch.sm_101a,
+            Arch.sm_101f,
+            Arch.sm_110,
+            Arch.sm_110a,
+            Arch.sm_110f,
+            Arch.sm_120,
+            Arch.sm_120a,
+            Arch.sm_120f,
+            Arch.sm_121,
+            Arch.sm_121a,
+            Arch.sm_121f,
+        )
 
     def __repr__(self):
         return self.__str__()

python/CuTeDSL/cutlass/base_dsl/ast_preprocessor.py

@@ -526,6 +526,10 @@ class DSLPreprocessor(ast.NodeTransformer):
                 self.generic_visit(node)
                 self.loop_nest_level -= 1
 
+            def visit_FunctionDef(self, node):
+                # Stop at nested function def
+                return
+
         checker = EarlyExitChecker(kind)
         checker.generic_visit(tree)
         if not checker.has_early_exit:
@@ -1325,6 +1329,7 @@ class DSLPreprocessor(ast.NodeTransformer):
                     node,
                 )
             elif func.id in ["min", "max"]:
+                self.import_top_module = True
                 return ast.copy_location(
                     ast.Call(
                         func=self._create_module_attribute(
@@ -1334,7 +1339,7 @@ class DSLPreprocessor(ast.NodeTransformer):
                             lineno=node.lineno,
                             col_offset=node.col_offset,
                         ),
-                        args=[node.args[0], node.args[1]],
+                        args=node.args,
                         keywords=[],
                     ),
                     node,

python/CuTeDSL/cutlass/base_dsl/common.py

@@ -134,7 +134,7 @@ def _get_friendly_cuda_error_message(error_code, error_name):
     # Add target architecture info
     target_arch = os.getenv("CUTE_DSL_ARCH", "unknown")
 
-    error_messages = {
+    additional_info = {
         "CUDA_ERROR_INVALID_SOURCE": (
             f"{Colors.RED}❌ Failed to load CUDA kernel - likely architecture mismatch.{Colors.RESET}\n\n"
         ),
@@ -157,6 +157,9 @@ def _get_friendly_cuda_error_message(error_code, error_name):
             f"{Colors.RED}⚠️ Invalid parameter passed to CUDA operation.{Colors.RESET}\n\n"
             f"{Colors.YELLOW}This is likely a bug - please report it with:{Colors.RESET}"
         ),
+        "CUDA_ERROR_INVALID_CLUSTER_SIZE": (
+            f"{Colors.RED}❌ Invalid cluster size.{Colors.RESET}\n\n"
+        ),
     }
 
     error_suggestions = {
@@ -194,12 +197,12 @@ def _get_friendly_cuda_error_message(error_code, error_name):
         ),
     }
 
-    message = error_messages.get(
-        error_name, f"{Colors.RED}Unknown CUDA error{Colors.RESET}"
-    )
+    message = f"{error_name} (error code: {error_code}) \n" \
+              f"{additional_info.get(error_name, '')} \n\n{Colors.RESET}"
 
     # Add debug information
     debug_info = f"\n- {Colors.BOLD}Error name: {error_name}\n"
+    debug_info += f"- Error code: {error_code}\n"
     debug_info += f"- CUDA_TOOLKIT_PATH: {os.getenv('CUDA_TOOLKIT_PATH', 'not set')}\n"
     debug_info += (
         f"- Target SM ARCH: {os.getenv('CUTE_DSL_ARCH', 'not set')}{Colors.RESET}\n"

python/CuTeDSL/cutlass/base_dsl/compiler.py

@@ -15,7 +15,7 @@ and executes it using MLIR's ExecutionEngine.
 
 """
 
-from typing import Sequence, Optional, Tuple
+from typing import Sequence, Optional, Tuple, Callable
 import os
 import sys
 import inspect
@@ -99,6 +99,8 @@ class Compiler:
         self.execution_engine = execution_engine
         # Flag to track if CUDA dependencies have been checked once in this process
         self._cuda_dependencies_checked = False
+        # Post-compile hook to run on Module
+        self._post_compile_hook: Optional[Callable[[ir.Module], None]] = None
 
     def _process_error(self, error_msg: str) -> Tuple[Optional[str], Optional[str]]:
         """Process error message to extract NVVM error and IR context"""
@@ -158,11 +160,13 @@ class Compiler:
                 ) from e
             raise e
 
+        if self._post_compile_hook:
+            self._post_compile_hook(module)
+
     def jit(self, module, opt_level: int = 2, shared_libs: Sequence[str] = ()):
         """Wraps the module in a JIT execution engine."""
         # Check CUDA driver and GPU dependencies before JIT execution (once per process)
         self._check_cuda_dependencies_once(shared_libs)
-
         # If pre-checks passed, attempt to create ExecutionEngine
         # Any failures at this point are likely non-CUDA related
         return self.execution_engine.ExecutionEngine(
@@ -216,6 +220,23 @@ class Compiler:
             )
 
 
+class PostCompileHookContext:
+    """Context manager for post-compile hook for a compiler."""
+
+    def __init__(self, compiler: Compiler, hook: Callable[[ir.Module], None]):
+        self.compiler = compiler
+        self.hook = hook
+        self.prev_post_compile_hook = None
+
+    def __enter__(self):
+        self.prev_post_compile_hook = self.compiler._post_compile_hook
+        self.compiler._post_compile_hook = self.hook
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.compiler._post_compile_hook = self.prev_post_compile_hook
+
+
 class CompileOption:
     """
     Base class for compile options.
@@ -276,6 +297,11 @@ class BooleanBasedFileDumpOption(CompileOption):
         return ""
 
 
+class EmptyCompileOption(CompileOption):
+    def serialize(self):
+        return ""
+
+
 class OptLevel(CompileOption):
     option_name = "opt-level"
 
@@ -329,6 +355,10 @@ class GPUArch(StringCompileOption):
         self._value = value
 
 
+class EnableTVMFFI(EmptyCompileOption):
+    option_name = "enable-tvm-ffi"
+
+
 class CompileOptions:
     """
     This class encapsulates compilation options to configure the JIT compilation.
@@ -349,6 +379,7 @@ class CompileOptions:
             KeepPTX: KeepPTX(False),
             GPUArch: GPUArch(""),
             LinkLibraries: LinkLibraries(""),
+            EnableTVMFFI: EnableTVMFFI(False),
         }
 
         if options is not None:
@@ -376,15 +407,28 @@ class CompileOptions:
             self.options[EnableAssertions].value = True
         if envar.lineinfo:
             self.options[GenerateLineInfo].value = True
+        if envar.enable_tvm_ffi:
+            self.options[EnableTVMFFI].value = True
 
         # Update the dump path if the option is set
+        arch = (
+            envar.arch
+            if self.options[GPUArch].value == ""
+            else self.options[GPUArch].value
+        )
         if self.options[KeepPTX].value:
             self.options[KeepPTX].dump_path = os.path.join(
-                envar.dump_dir, f"{function_name}.ptx"
+                envar.dump_dir, f"{function_name}"
+            )
+            self.options[KeepPTX].full_ptx_path = os.path.join(
+                envar.dump_dir, f"{function_name}.{arch}.ptx"
             )
         if self.options[KeepCUBIN].value:
             self.options[KeepCUBIN].dump_path = os.path.join(
-                envar.dump_dir, f"{function_name}.cubin"
+                envar.dump_dir, f"{function_name}"
+            )
+            self.options[KeepCUBIN].full_cubin_path = os.path.join(
+                envar.dump_dir, f"{function_name}.{arch}.cubin"
             )
 
     @property
@@ -399,12 +443,38 @@ class CompileOptions:
     def dump_ptx_path(self) -> str | None:
         return self.options[KeepPTX].dump_path if self.options[KeepPTX].value else None
 
+    @property
+    def full_ptx_path(self) -> str | None:
+        return (
+            self.options[KeepPTX].full_ptx_path if self.options[KeepPTX].value else None
+        )
+
     @property
     def dump_cubin_path(self) -> str | None:
         return (
             self.options[KeepCUBIN].dump_path if self.options[KeepCUBIN].value else None
         )
 
+    @property
+    def full_cubin_path(self) -> str | None:
+        return (
+            self.options[KeepCUBIN].full_cubin_path
+            if self.options[KeepCUBIN].value
+            else None
+        )
+
+    @property
+    def enable_tvm_ffi(self) -> bool:
+        ret = self.options[EnableTVMFFI].value
+        if ret:
+            try:
+                import tvm_ffi
+            except ModuleNotFoundError:
+                raise DSLRuntimeError(
+                    "TVM FFI is not installed, please install it via `pip install apache-tvm-ffi`"
+                )
+        return ret
+
     def to_str(self) -> str:
         """
         Generate a string representation of all compilation options
@@ -433,6 +503,7 @@ def _parse_compile_options_from_str(options: str) -> CompileOptions:
             "keep_cubin": KeepCUBIN,
             "keep_ptx": KeepPTX,
             "gpu_arch": GPUArch,
+            "enable_tvm_ffi": EnableTVMFFI,
         }
         return mapping[option_str]
 
@@ -448,6 +519,7 @@ def _parse_compile_options_from_str(options: str) -> CompileOptions:
     parser.add_argument("--keep-ptx", action="store_true", default=False)
     parser.add_argument("--ptxas-options", type=str, default="")
     parser.add_argument("--gpu-arch", type=str, default="")
+    parser.add_argument("--enable-tvm-ffi", action="store_true", default=False)
     compile_options = CompileOptions()
     try:
         # Use shlex to properly handle options with spaces
@@ -553,6 +625,9 @@ class CompileCallable:
 
         # process compile options, extract the options and remove them from the kwargs
         options = kwargs.pop("options", None)
+        if isinstance(options, str) and len(options) == 0:
+            options = None
+
         if options is not None and isinstance(options, str):
             compile_options = _parse_compile_options_from_str(options)
         else:

python/CuTeDSL/cutlass/base_dsl/dsl.py

@@ -27,6 +27,7 @@ import re
 import inspect
 import argparse
 import hashlib
+import weakref
 from functools import lru_cache, wraps
 from collections import namedtuple, OrderedDict
 from abc import ABC, abstractmethod
@@ -52,7 +53,7 @@ from .runtime.jit_arg_adapters import is_argument_constexpr, JitArgAdapterRegist
 
 from .ast_preprocessor import DSLPreprocessor
 from .common import *
-from .typing import get_c_pointers, get_mlir_types, Integer
+from .typing import get_c_pointers, get_mlir_types, Integer, arg_compatible_with_tvm_ffi
 from .arch import Arch
 
 # =============================================================================
@@ -345,6 +346,7 @@ class BaseDSL:
 
         if preprocess:
             self.preprocessor = DSLPreprocessor(dsl_package_name)
+
         log().info(f"Initializing {name} DSL")
         log().debug(f"Logger initialized for {self.name}")
 
@@ -756,6 +758,7 @@ class BaseDSL:
         args_spec: inspect.FullArgSpec,
         *,
         is_host=True,
+        compile_only=False,
     ):
         """Generate JIT function arguments."""
 
@@ -806,7 +809,14 @@ class BaseDSL:
                     jit_adapted_args.append(arg)
 
                 if is_host:
-                    jit_exec_arg.extend(get_c_pointers(arg))
+                    if self.envar.enable_tvm_ffi:
+                        if not arg_compatible_with_tvm_ffi(arg):
+                            raise DSLRuntimeError(
+                                f"Argument #{i + 1} ({arg_name}) is not a TVM FFI argument."
+                            )
+                        jit_exec_arg.extend([arg])
+                    else:
+                        jit_exec_arg.extend(get_c_pointers(arg))
                     jit_arg_type.extend(get_mlir_types(arg))
                 else:
                     dyn_vals = extract_mlir_values(arg)
@@ -814,7 +824,10 @@ class BaseDSL:
                     jit_arg_type.extend([v.type for v in dyn_vals])
 
                 if not jit_arg_type or not jit_exec_arg:
-                    if (is_host and hasattr(arg, "__c_pointers__")) or (
+                    # when it is compile only, we don't have to prepare the executable arguments.
+                    if (
+                        is_host and (compile_only or hasattr(arg, "__c_pointers__"))
+                    ) or (
                         not is_host
                         and hasattr(arg, "__extract_mlir_values__")
                         and hasattr(arg, "__new_from_mlir_values__")
@@ -845,20 +858,32 @@ class BaseDSL:
         return jit_exec_args, jit_arg_types, jit_arg_attrs, jit_adapted_args
 
     def generate_mlir_function_types(
-        self, func, function_name, input_args, kwargs, args_spec: inspect.FullArgSpec
+        self,
+        func,
+        function_name,
+        input_args,
+        kwargs,
+        args_spec: inspect.FullArgSpec,
+        compile_only=False,
     ):
         """Convert input arguments to MLIR function signature also convert numpy arrays to memref."""
 
         exe_args, types, attrs, adapted_args = self._generate_jit_func_args(
-            func, function_name, input_args, kwargs, args_spec, is_host=True
+            func,
+            function_name,
+            input_args,
+            kwargs,
+            args_spec,
+            is_host=True,
+            compile_only=compile_only,
         )
 
         log().debug("Execution Arguments: %s", ", ".join(map(str, exe_args)))
         log().debug("Types: %s", ", ".join(map(str, types)))
 
-        assert len(exe_args) == len(types), (
-            "expects the same number of arguments and function parameters"
-        )
+        assert (
+            compile_only or self.envar.enable_tvm_ffi or len(exe_args) == len(types)
+        ), "expects the same number of arguments and function parameters"
 
         return exe_args, types, adapted_args
 
@@ -867,6 +892,7 @@ class BaseDSL:
         cluster: list = None
         grid: list = field(default_factory=lambda: [1, 1, 1])
         block: list = field(default_factory=lambda: [1, 1, 1])
+        max_number_threads: list = field(default_factory=lambda: [0, 0, 0])
         smem: int = None
         async_deps: list = field(default_factory=list)
         has_cluster: bool = False
@@ -906,6 +932,13 @@ class BaseDSL:
             elif len(self.cluster) != 3:
                 raise DSLRuntimeError(f"Expect 3d cluster!")
 
+        def has_max_number_threads(self):
+            """Check if max_number_threads is given by user"""
+            return all(
+                value == 0 if not is_dynamic_expression(value) else False
+                for value in self.max_number_threads
+            )
+
     def diagnostic(self):
         """Check command line parameters and enables diagnostic"""
         # Check command line arguments "-diagnostic"
@@ -1029,7 +1062,6 @@ class BaseDSL:
                 pipeline = re.sub(r"{.+}", opt_str, pipeline)
             else:
                 pipeline = pipeline.rstrip(")") + f"{{{opt_str}}})"
-        log().debug(f"Using pipeline = {pipeline}")
         return pipeline
 
     def get_shared_libs(self) -> list:
@@ -1101,6 +1133,18 @@ class BaseDSL:
 
         return module
 
+    def get_return_types(self) -> List[ir.Type]:
+        """
+        Get the return types of the host function.
+        """
+        return []
+
+    def generate_default_return_values(self, ip=None) -> List[ir.Value]:
+        """
+        Generate the default return values of the host function.
+        """
+        return []
+
     def generate_original_ir(
         self,
         ir,
@@ -1115,27 +1159,35 @@ class BaseDSL:
         frame=None,
     ):
         def build_ir_module():
-            module = ir.Module.create(loc=self.get_location(frame))
+            loc = self.get_location(frame)
+            module = ir.Module.create(loc=loc)
             unit_attr = ir.UnitAttr.get()
             module.operation.attributes["gpu.container_module"] = unit_attr
 
             with ir.InsertionPoint(module.body):
                 # Always generate gpu module. It's canonicalized by the compiler when it's not used.
-                self._build_gpu_module(gpu_module_attrs, loc=self.get_location(frame))
+                self._build_gpu_module(gpu_module_attrs, loc=loc)
 
+                ret_types = self.get_return_types()
                 fop = func.FuncOp(
-                    function_name, (func_types, []), loc=self.get_location(frame)
+                    function_name, (func_types, ret_types), loc=loc
                 )
                 fop.attributes["llvm.emit_c_interface"] = ir.UnitAttr.get()
                 log().debug("Generated Function OP [%s]", fop)
-                with ir.InsertionPoint(fop.add_entry_block()):
+                # Attach per-argument source locations if supported by the FuncOp binding.
+                arg_locs = [loc for _ in range(len(func_types))]
+                entry_block = fop.add_entry_block(arg_locs=arg_locs)
+                with ir.InsertionPoint(entry_block):
                     ir_args, ir_kwargs = self.generate_execution_arguments(
                         args, kwargs, fop, args_spec
                     )
                     # Call user function body
                     try:
                         result = funcBody(*ir_args, **ir_kwargs)
-                        func.ReturnOp([], loc=self.get_location(frame))
+                        default_ret_values = self.generate_default_return_values(
+                            ir.InsertionPoint.current
+                        )
+                        func.ReturnOp(default_ret_values, loc=loc)
                     except NameError as name_error:
                         raise DSLRuntimeError(
                             f"💥💥💥 Error during runtime code generation for function `{funcBody.__name__}` 💥💥💥",
@@ -1165,6 +1217,10 @@ class BaseDSL:
         args_spec,
         no_cache,
         func_type=JitCompiledFunction,
+        *,
+        dynamic_args=None,
+        dynamic_kwargs=None,
+        original_function_name=None,
     ):
         # If `gpu-arch` is set by compile_options, use it. Otherwise, use the arch from the environment variable.
         compile_gpu_arch = (
@@ -1178,6 +1234,7 @@ class BaseDSL:
         pipeline = self.preprocess_pipeline(
             self._get_pipeline(pipeline), compile_gpu_arch
         )
+        log().debug(f"Using pipeline = {pipeline}")
         shared_libs = self.get_shared_libs()
         profiler = timer(enable=self.envar.jit_time_profiling)
         if (
@@ -1212,7 +1269,7 @@ class BaseDSL:
             )
         capi_func = profiler(engine.lookup)(function_name) if engine else None
 
-        fn = JitCompiledFunction(
+        fn = func_type(
             module,
             engine,
             capi_func,
@@ -1221,12 +1278,15 @@ class BaseDSL:
             self.kernel_info,
             jit_time_profiling=self.envar.jit_time_profiling,
             jit_function_artifacts=JitFunctionArtifacts(
-                PTX=self.compile_options.dump_ptx_path,
-                CUBIN=self.compile_options.dump_cubin_path,
+                PTX=self.compile_options.full_ptx_path,
+                CUBIN=self.compile_options.full_cubin_path,
                 MLIR=(self.dump_mlir_path if self.envar.keep_ir else None),
             ),
         )
 
+        # set dynamic arguments if the jit_function is a JitCompiledFunction for AOT generation.
+        fn.set_dynamic_args(dynamic_args, dynamic_kwargs)
+
         if not no_cache:
             # module stored in cache is compiled.
             self.jit_cache[module_hash] = fn
@@ -1243,6 +1303,34 @@ class BaseDSL:
         # reset the compile options after the compilation is done.
         self.compile_options = CompileOptions()
 
+    def extract_dynamic_args(self, funcBody, args, kwargs, args_spec):
+        """This function is used to extract the original dynamic arguments for AOT C header generation.
+        The dynamic argument is the argument which is not marked as `Constexpr` in the function signature.
+        """
+        dynamic_args = []
+        dynamic_kwargs = OrderedDict()
+        for i, arg in enumerate(args):
+            if not is_argument_constexpr(
+                arg,
+                args_spec.annotations.get(args_spec.args[i], None),
+                args_spec.args[i],
+                i,
+                funcBody,
+            ):
+                try:
+                    dynamic_args.append(weakref.proxy(arg))
+                except TypeError:
+                    # If arg cannot be weakly referenced (e.g., int, float)
+                    dynamic_args.append(arg)
+        for i, (k, v) in enumerate(kwargs.items()):
+            if not is_argument_constexpr(v, args_spec.kwonlyargs[i], k, i, funcBody):
+                try:
+                    dynamic_kwargs[k] = weakref.proxy(v)
+                except TypeError:
+                    # If v cannot be weakly referenced (e.g., int, float)
+                    dynamic_kwargs[k] = v
+        return dynamic_args, dynamic_kwargs
+
     def generate_mlir(
         self,
         funcBody,
@@ -1258,49 +1346,63 @@ class BaseDSL:
         frame=None,
     ):
         """Generate MLIR module and compile iself.T_provider."""
-        with ir.Context(), ir.Location.unknown():
-            # Convert input arguments to MLIR arguments
-            exe_args, func_types, adapted_args = self.generate_mlir_function_types(
-                funcBody, function_name, args, kwargs, args_spec
-            )
-
-            # Generate original ir module and its hash value.
-            module, module_hash, result = self.generate_original_ir(
-                ir,
-                func,
-                funcBody,
-                kwargs,
-                function_name,
-                func_types,
-                gpu_module_attrs,
-                args,
-                args_spec,
-                frame=frame,
-            )
-
-            # dryrun is used to only generate IR
-            if self.envar.dryrun:
-                return result
-
-            if (
-                no_cache
-                or module_hash not in self.jit_cache
-                or self.jit_cache[module_hash].capi_func is None
-            ):
-                # no cache or cache miss, do ir generation/compilation/jit engine
-                jit_function = self.compile_and_cache(
-                    module, module_hash, function_name, pipeline, args_spec, no_cache
+        with ir.Context(), self.get_location(frame):
+            try:
+                # Convert input arguments to MLIR arguments
+                exe_args, func_types, adapted_args = self.generate_mlir_function_types(
+                    funcBody, function_name, args, kwargs, args_spec, compile_only
                 )
-            else:
-                # cache hit
-                log().info(
-                    "JIT cache hit IN-MEMORY function=[%s] module_hash=[%s]",
+                dynamic_args, dynamic_kwargs = self.extract_dynamic_args(
+                    funcBody, args, kwargs, args_spec
+                )
+                original_function_name = funcBody.__name__
+
+                # Generate original ir module and its hash value.
+                module, module_hash, result = self.generate_original_ir(
+                    ir,
+                    func,
+                    funcBody,
+                    kwargs,
                     function_name,
-                    module_hash,
+                    func_types,
+                    gpu_module_attrs,
+                    args,
+                    args_spec,
+                    frame=frame,
                 )
-                jit_function = self.jit_cache[module_hash]
 
-            self.post_compilation_cleanup()
+                # dryrun is used to only generate IR
+                if self.envar.dryrun:
+                    return result
+
+                if (
+                    no_cache
+                    or module_hash not in self.jit_cache
+                    or self.jit_cache[module_hash].capi_func is None
+                ):
+                    # no cache or cache miss, do ir generation/compilation/jit engine
+                    jit_function = self.compile_and_cache(
+                        module,
+                        module_hash,
+                        function_name,
+                        pipeline,
+                        args_spec,
+                        no_cache,
+                        dynamic_args=dynamic_args,
+                        dynamic_kwargs=dynamic_kwargs,
+                        original_function_name=original_function_name,
+                    )
+                else:
+                    # cache hit
+                    log().info(
+                        "JIT cache hit IN-MEMORY function=[%s] module_hash=[%s]",
+                        function_name,
+                        module_hash,
+                    )
+                    jit_function = self.jit_cache[module_hash]
+            
+            finally:
+                self.post_compilation_cleanup()
 
         # If compile_only is set, bypass execution return the jit_executor directly
         if compile_only:
@@ -1569,13 +1671,13 @@ class BaseDSL:
         """
         ret = None
 
-        with ir.Context(), ir.Location.unknown():
+        with ir.Context(), self.get_location():
             loc = self.get_location()
             module = ir.Module.create(loc=loc)
             unit_attr = ir.UnitAttr.get()
             module.operation.attributes["gpu.container_module"] = unit_attr
             with ir.InsertionPoint(module.body):
-                self._build_gpu_module({})
+                self._build_gpu_module({}, loc=loc)
                 ret, kernel_name = kernel_generator()
                 log().debug(
                     f"Kernel generator returned: ret={ret}, kernel_name={kernel_name}"

python/CuTeDSL/cutlass/base_dsl/env_manager.py

@@ -245,8 +245,15 @@ def get_prefix_dsl_libs(prefix: str):
             ]
 
             libs_cand = find_libs_in_ancestors(start, target_libs, lib_folder_guesses)
+
+            optional_libs_cand = find_libs_in_ancestors(
+                start, {"cuda_dialect_runtime"}, lib_folder_guesses
+            )
+
             if libs_cand:
                 dsl_libs = ":".join(libs_cand)
+                if optional_libs_cand:
+                    dsl_libs += ":" + ":".join(optional_libs_cand)
                 return dsl_libs
 
             return None
@@ -312,6 +319,7 @@ class EnvironmentVarManager(LogEnvironmentManager):
     - [DSL_NAME]_DISABLE_FILE_CACHING: Disable file caching (default: False)
     - [DSL_NAME]_FILE_CACHING_CAPACITY: Limits the number of the cache save/load files (default: 1000)
     - [DSL_NAME]_LIBS: Path to dependent shared libraries (default: None)
+    - [DSL_NAME]_ENABLE_TVM_FFI: Enable TVM-FFI or not (default: False)
     """
 
     def __init__(self, prefix="DSL"):
@@ -358,3 +366,5 @@ class EnvironmentVarManager(LogEnvironmentManager):
 
         # whether to enable assert in host and device code
         self.enable_assertions = get_bool_env_var(f"{prefix}_ENABLE_ASSERTIONS", False)
+
+        self.enable_tvm_ffi = get_bool_env_var(f"{prefix}_ENABLE_TVM_FFI", False)

python/CuTeDSL/cutlass/base_dsl/export/__init__.py

@@ -0,0 +1,20 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 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/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.
+
+from .c_header_generator import CHeaderGenerator
+from .export import get_export_module, dump_to_object, export_to_c
+
+__all__ = [
+    "CHeaderGenerator",
+    "get_export_module",
+    "dump_to_object",
+    "export_to_c",
+]

python/CuTeDSL/cutlass/base_dsl/export/c_header_generator.py

@@ -0,0 +1,383 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 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/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.
+
+from ..typing import (
+    Numeric,
+    NumericMeta,
+    Int8,
+    Int16,
+    Int32,
+    Int64,
+    Uint8,
+    Uint16,
+    Uint32,
+    Uint64,
+    Boolean,
+    Float16,
+    Float32,
+    Float64,
+    TFloat32,
+)
+from ..dsl import is_dynamic_expression
+from ..common import DSLRuntimeError
+from ..jit_executor import ExecutionArgs
+from ..._mlir import ir
+
+from typing import Type, List, Any, Dict
+from inspect import isclass
+import cuda.bindings.driver as cuda
+
+# =============================================================================
+# C Header Generator for c/cpp AOT support
+# =============================================================================
+cubin_suffix = "cubin"
+
+
+class CHeaderGenerator:
+    """This class provides a Export C Header Generator for c/cpp AOT support."""
+
+    includes = """
+#pragma once
+
+#include <cuda.h>
+#include <cuda_fp16.h>
+#include <stdio.h>
+
+"""
+    cuda_error_check = r"""_CUDA_ERROR_CHECK(err) { \
+    if ((err) != CUDA_SUCCESS) { \
+        const char *pErrStr = NULL; \
+        if (cuGetErrorString(err, &pErrStr) != CUDA_SUCCESS) { \
+            printf("Error: cuGetErrorString failed\n"); \
+            exit(err); \
+        } \
+        const char *pErrName = NULL; \
+        if (cuGetErrorName(err, &pErrName) != CUDA_SUCCESS) { \
+            printf("Error: cuGetErrorName failed\n"); \
+            exit(err); \
+        } \
+        printf("Got Cuda Error: %s %s\n", pErrName, pErrStr); \
+        exit(err); \
+    } \
+}
+"""
+
+    numeric_to_c_type = {
+        Int8: "int8_t ",
+        Int16: "int16_t ",
+        Int32: "int32_t ",
+        Int64: "int64_t ",
+        Uint8: "uint8_t ",
+        Uint16: "uint16_t ",
+        Uint32: "uint32_t ",
+        Uint64: "uint64_t ",
+        Boolean: "uint8_t ",
+        Float32: "float ",
+        TFloat32: "float ",
+        Float64: "double ",
+        Float16: "__half_raw ",
+    }
+
+    def _count_dynamic_expression(self, arg):
+        """
+        Count the number of dynamic values in the argument.
+        """
+        if isinstance(arg, (list, tuple)):
+            return sum(self._count_dynamic_expression(x) for x in arg)
+        elif is_dynamic_expression(arg):
+            return 1
+        return 0
+
+    def _generate_numeric_argument(self, arg_name: str, arg_type: Type[Numeric]):
+        """
+        Generate the argument of the wrapper function.
+        """
+        if arg_type in self.numeric_to_c_type:
+            return self.numeric_to_c_type[arg_type] + arg_name
+        raise DSLRuntimeError(
+            f"Unsupported argument type for c function argument generation: {arg_type}"
+        )
+
+    def _generate_check_cuda(self, dsl_name: str):
+        check_cuda = (
+            f"""
+// Macro to check for cuda errors.
+#ifndef {dsl_name}_CUDA_ERROR_CHECK
+#define {dsl_name}"""
+            + self.cuda_error_check
+            + f"""
+#endif
+"""
+        )
+
+        return check_cuda
+
+    def _generate_kernel_metadata(
+        self, symbol_prefix: str, kernel_info: Dict[str, List], dsl_name: str
+    ):
+        """
+        Generate the kernel metadata for the compiled function.
+        """
+        function_declarations = []
+        function_loads = []
+        function_set_attributes = []
+        function_non_portable_cluster_size_allowed = []
+        for sym, attrs in kernel_info.items():
+            function_declaration = f"CUfunction {sym};"
+            function_load = f'{dsl_name}_CUDA_ERROR_CHECK(cuModuleGetFunction(&metadata->{sym}, metadata->module, "{sym}"));'
+            function_declarations.append(function_declaration)
+            function_loads.append(function_load)
+            for attr in attrs:
+                function_set_attributes.append(
+                    f"{dsl_name}_CUDA_ERROR_CHECK(cuFuncSetAttribute(metadata->{sym}, {attr.name}, {attr.value}));"
+                )
+            function_non_portable_cluster_size_allowed.append(
+                f"{dsl_name}_CUDA_ERROR_CHECK(cuFuncSetAttribute(metadata->{sym}, CU_FUNC_ATTRIBUTE_NON_PORTABLE_CLUSTER_SIZE_ALLOWED, 1));"
+            )
+        function_declarations_str = "\n    ".join(function_declarations)
+        kernel_metadata_struct = f"""
+typedef struct {{
+    CUmodule module;
+    {function_declarations_str}
+}} {symbol_prefix}_Kernel_Metadata_t;
+"""
+        function_loads_str = "\n    ".join(function_loads)
+        function_set_attributes_str = "\n    ".join(function_set_attributes)
+        function_non_portable_cluster_size_allowed_str = "\n        ".join(
+            function_non_portable_cluster_size_allowed
+        )
+        kernel_metadata_load = f"""
+static inline void {symbol_prefix}_Kernel_Metadata_Load({symbol_prefix}_Kernel_Metadata_t *metadata) {{
+    {dsl_name}_CUDA_ERROR_CHECK(cuModuleLoadData(&metadata->module, {symbol_prefix}_{cubin_suffix}));
+    {function_loads_str}
+    {function_set_attributes_str}
+    int driver_version;
+    {dsl_name}_CUDA_ERROR_CHECK(cuDriverGetVersion(&driver_version));
+    if (driver_version >= 11080) {{
+        {function_non_portable_cluster_size_allowed_str}
+    }}
+}}
+"""
+        kernel_metadata_unload = f"""
+static inline void {symbol_prefix}_Kernel_Metadata_Unload({symbol_prefix}_Kernel_Metadata_t *metadata) {{
+    {dsl_name}_CUDA_ERROR_CHECK(cuModuleUnload(metadata->module));
+}}
+"""
+
+        return kernel_metadata_struct + kernel_metadata_load + kernel_metadata_unload
+
+    def _generate_arguments(
+        self,
+        symbol_prefix: str,
+        args_spec: ExecutionArgs,
+        args: List[Any],
+        kwargs: Dict[str, Any],
+    ):
+        """
+        Generate the arguments of the wrapper function.
+        """
+        arguments = []
+        packed_args = []
+        declarations = []
+        # traverse the runtime args_spec and generate the arguments
+        rectified_args = args_spec.get_rectified_args(args, kwargs)
+        input_arg_names = args_spec.args_spec.args + args_spec.args_spec.kwonlyargs
+        for arg_name, arg in zip(input_arg_names, rectified_args):
+            arg_type = args_spec.args_spec.annotations.get(arg_name, None)
+
+            # process optional argument
+            if arg is None:
+                continue
+
+            # Generate basic numeric types
+            if isinstance(arg_type, NumericMeta):
+                arguments.append(self._generate_numeric_argument(arg_name, arg_type))
+                packed_args.append("&" + arg_name)
+            elif isclass(arg_type) and issubclass(arg_type, cuda.CUstream):
+                arguments.append("CUstream " + arg_name)
+                packed_args.append("&" + arg_name)
+            else:
+                raise DSLRuntimeError(
+                    f"Unsupported argument for c function argument generation: {arg} with type {arg_type}"
+                )
+
+        return arguments, packed_args, declarations
+
+    def _generate_wrapper_function(
+        self,
+        symbol_prefix: str,
+        args_spec: ExecutionArgs,
+        function_name: str,
+        kernel_info: Dict[str, List],
+        dynamic_args: list,
+        dynamic_kwargs: dict,
+    ):
+        """
+        Generate the wrapper function for the compiled function which is provided to users as the entry point.
+        It uses the `symbol_prefix` as the function name for identification. The host/device symbols are hidden under the bytecode.
+        """
+        # 1. Get the name of the function wrapper
+        wrapper_function_name = f"{symbol_prefix}_wrapper"
+        capi_function_name = f"_mlir_{symbol_prefix}__mlir_ciface_{function_name}"
+        # 2. Generate the signature of the wrapper function
+        arguments, packed_args, declarations = self._generate_arguments(
+            symbol_prefix, args_spec, dynamic_args, dynamic_kwargs
+        )
+        declarations = "\n".join(declarations)
+
+        # 3. Generate the wrapper function
+        kernel_symbols = tuple(kernel_info.keys())
+        kernel_symbols_str = ", ".join([f"&metadata->{sym}" for sym in kernel_symbols])
+        function = (
+            declarations
+            + f"""
+#ifdef __cplusplus
+extern "C"
+#endif
+void {capi_function_name}(void **args, int32_t num_args);
+
+static inline void {wrapper_function_name}({symbol_prefix}_Kernel_Metadata_t *metadata, {", ".join(arguments)}) {{
+    void *args[{len(packed_args) + len(kernel_symbols)}] = {{
+        {", ".join(packed_args)}, {kernel_symbols_str}
+    }};
+    {capi_function_name}(args, {len(packed_args) + len(kernel_symbols)});
+}}
+"""
+        )
+        return function
+
+    def _generate_binary_declaration(self, symbol_prefix: str):
+        """
+        Generate the binary of the compiled function.
+        """
+
+        varname = symbol_prefix + "_" + cubin_suffix
+        binary = f"""
+extern const unsigned char {varname}[];
+"""
+        return binary
+
+    def __call__(
+        self,
+        symbol_prefix: str,
+        export_module: ir.Module,
+        args_spec: ExecutionArgs,
+        function_name: str,
+        kernel_info: Dict[str, List],
+        dynamic_args: list,
+        dynamic_kwargs: dict,
+        dsl_name: str,
+    ) -> str:
+        if len(kernel_info) > 0:
+            check_cuda = self._generate_check_cuda(dsl_name)
+            kernel_metadata = self._generate_kernel_metadata(
+                symbol_prefix, kernel_info, dsl_name
+            )
+            binary = self._generate_binary_declaration(symbol_prefix)
+        else:
+            check_cuda = ""
+            kernel_metadata = ""
+            binary = ""
+        function = self._generate_wrapper_function(
+            symbol_prefix,
+            args_spec,
+            function_name,
+            kernel_info,
+            dynamic_args,
+            dynamic_kwargs,
+        )
+        header = self.includes + check_cuda + binary + kernel_metadata + function
+        return header
+
+
+# =============================================================================
+# Below is the example output of a complete header file.
+# =============================================================================
+# #pragma once
+
+# #include <cuda.h>
+# #include <cuda_fp16.h>
+# #include <stdio.h>
+
+
+# // Macro to check for cuda errors.
+# #ifndef CUTE_DSL_CUDA_ERROR_CHECK
+# #define CUTE_DSL_CUDA_ERROR_CHECK(err) { \
+#     if ((err) != CUDA_SUCCESS) { \
+#         const char *pErrStr = NULL; \
+#         if (cuGetErrorString(err, &pErrStr) != CUDA_SUCCESS) { \
+#             printf("Error: cuGetErrorString failed\n"); \
+#             exit(err); \
+#         } \
+#         const char *pErrName = NULL; \
+#         if (cuGetErrorName(err, &pErrName) != CUDA_SUCCESS) { \
+#             printf("Error: cuGetErrorName failed\n"); \
+#             exit(err); \
+#         } \
+#         printf("Got Cuda Error: %s %s\n", pErrName, pErrStr); \
+#         exit(err); \
+#     } \
+# }
+# #endif
+
+# extern const unsigned char gemm_cubin[];
+
+# typedef struct {
+#     CUmodule module;
+#     CUfunction kernel_cutlass_dummy_gemm_tensorptrf32generico1_tensorptrf32generico1_tensorptrf32generico1_0;
+# } gemm_Kernel_Metadata_t;
+
+# static inline void gemm_Kernel_Metadata_Load(gemm_Kernel_Metadata_t *metadata) {
+#     CUTE_DSL_CUDA_ERROR_CHECK(cuModuleLoadData(&metadata->module, gemm_cubin));
+#     CUTE_DSL_CUDA_ERROR_CHECK(cuModuleGetFunction(&metadata->kernel_cutlass_dummy_gemm_tensorptrf32generico1_tensorptrf32generico1_tensorptrf32generico1_0, metadata->module, "kernel_cutlass_dummy_gemm_tensorptrf32generico1_tensorptrf32generico1_tensorptrf32generico1_0"));
+
+#     int driver_version;
+#     CUTE_DSL_CUDA_ERROR_CHECK(cuDriverGetVersion(&driver_version));
+#     if (driver_version >= 11080) {
+#         CUTE_DSL_CUDA_ERROR_CHECK(cuFuncSetAttribute(metadata->kernel_cutlass_dummy_gemm_tensorptrf32generico1_tensorptrf32generico1_tensorptrf32generico1_0, CU_FUNC_ATTRIBUTE_NON_PORTABLE_CLUSTER_SIZE_ALLOWED, 1));
+#     }
+# }
+
+# static inline void gemm_Kernel_Metadata_Unload(gemm_Kernel_Metadata_t *metadata) {
+#     CUTE_DSL_CUDA_ERROR_CHECK(cuModuleUnload(metadata->module));
+# }
+
+# typedef struct {
+#     void *data;
+#     int32_t dynamic_shapes[2];
+#     int64_t dynamic_strides[1];
+# } gemm_Tensor_a_t;
+
+
+# typedef struct {
+#     void *data;
+#     int32_t dynamic_shapes[2];
+#     int64_t dynamic_strides[1];
+# } gemm_Tensor_b_t;
+
+
+# typedef struct {
+#     void *data;
+#     int32_t dynamic_shapes[2];
+#     int64_t dynamic_strides[1];
+# } gemm_Tensor_c_t;
+
+# #ifdef __cplusplus
+# extern "C"
+# #endif
+# void _mlir_gemm__mlir_ciface_cutlass_host_func_Tensorgenericoi641_Tensorgenericoi641_Tensorgenericoi641(void **args, int32_t num_args);
+
+# static inline void gemm_wrapper(gemm_Kernel_Metadata_t *metadata, gemm_Tensor_a_t *a, gemm_Tensor_b_t *b, gemm_Tensor_c_t *c) {
+#     void *args[4] = {
+#         a, b, c, &metadata->kernel_cutlass_dummy_gemm_tensorptrf32generico1_tensorptrf32generico1_tensorptrf32generico1_0
+#     };
+#     _mlir_gemm__mlir_ciface_cutlass_host_func_Tensorgenericoi641_Tensorgenericoi641_Tensorgenericoi641(args, 4);
+# }

python/CuTeDSL/cutlass/base_dsl/export/export.py

@@ -0,0 +1,199 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 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/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 io
+import os
+import array
+import tempfile
+
+from ..common import DSLRuntimeError
+from ..jit_executor import JitCompiledFunction, get_escaped_cubin_bytes
+from ...base_dsl.dsl import BaseDSL
+from ...base_dsl.typing import Int32, Int64, Float32, Float64
+from ..._mlir import ir
+from ..._mlir.dialects import llvm
+
+from .c_header_generator import CHeaderGenerator
+
+from typing import Union
+
+cubin_suffix = "cubin"
+
+def get_export_module(ir_module: ir.Module, symbol_prefix: str):
+    """Get the export module which is cloned from the original compiled ir module, and add the prefix
+    to avoid the symbol conflict.
+
+    @param ir_module: The original compiled ir module. Comes from the JitCompiledFunction.ir_module.
+    @param symbol_prefix: The prefix name of the function. This is the unique identifier name of the function to avoid symbol conflict in the generated object file.
+    @return: The export module of the function.
+    """
+    # Add prefix for symbol names to avoid conflict with other functions
+    defined_symbols = set()
+
+    def walk_llvm_func_op(op):
+        # not a declaration
+        if (
+            op.name == "llvm.func"
+            and len(op.opview.operation.regions) > 0
+            and len(op.opview.operation.regions[0].blocks) > 0
+        ):
+            defined_symbols.add(op.attributes["sym_name"].value)
+            op.attributes["sym_name"] = ir.StringAttr.get(
+                symbol_prefix + "_" + op.attributes["sym_name"].value
+            )
+        return ir.WalkResult.ADVANCE
+
+    def walk_llvm_call_op(op):
+        if op.name == "llvm.call" and op.attributes["callee"].value in defined_symbols:
+            op.attributes["callee"] = ir.FlatSymbolRefAttr.get(
+                symbol_prefix + "_" + op.attributes["callee"].value
+            )
+        return ir.WalkResult.ADVANCE
+
+    with ir.Context():
+        export_module = ir.Module.parse(str(ir_module))
+        export_module.operation.walk(walk_llvm_func_op)
+        export_module.operation.walk(walk_llvm_call_op)
+    return export_module
+
+
+
+def dump_to_object(
+    prefix_name: str,
+    export_module: ir.Module,
+    jit_function: Union[JitCompiledFunction, "CudaDialectJitCompiledFunction"],
+    dsl: BaseDSL,
+    use_gpu_dialect: bool,
+) -> bytes:
+    """Dump the compiled ir function to a bytes object with ELF format. The bytes object contains the host
+    launch entry function and cubin inside.
+
+    @param prefix_name: The prefix name of the function. This is the unique identifier name of the function to avoid symbol conflict in the generated object file.
+    @param export_module: The export module of the function. This is the module that contains the function to be exported.
+    @param jit_function: The jit-compiled function. To provided other metadata for the object file.
+    @param dsl: The dsl object. This is the dsl object to get the compiler provider and shared libs.
+    @return: The bytes object of the function.
+    """
+    if use_gpu_dialect:
+        cubin_data = None
+
+        def strip_gpu_binary_op(op):
+            if op.name == "gpu.binary":
+                s = io.BytesIO()
+                op.operation.write_bytecode(s)
+                nonlocal cubin_data
+                cubin_data = s.getvalue()
+                cubin_data = cubin_data.split(b'bin = "')[1].split(b'">')[0]
+                cubin_data = get_escaped_cubin_bytes(cubin_data)
+                op.erase()
+                return ir.WalkResult.ADVANCE
+            return ir.WalkResult.ADVANCE
+
+        # Strip gpu related to avoid the object file generating builtin module load/unload functions
+        export_module.operation.walk(strip_gpu_binary_op)
+
+        cubin_array = array.array("b", cubin_data)
+        with (
+            export_module.context,
+            ir.Location.unknown(),
+            ir.InsertionPoint(export_module.body),
+        ):
+            new_binary_global_op = llvm.GlobalOp(
+                sym_name="_".join([prefix_name, cubin_suffix]),
+                global_type=ir.Type.parse(f"!llvm.array<{len(cubin_array)} x i8>"),
+                linkage=ir.Attribute.parse("#llvm.linkage<external>"),
+                value=ir.DenseIntElementsAttr.get(cubin_array),
+                constant=True,
+            )
+    if "gpu.container_module" in export_module.operation.attributes:
+        del export_module.operation.attributes["gpu.container_module"]
+    # Generate the object file
+    export_engine = dsl.compiler_provider.jit(
+        export_module, shared_libs=dsl.get_shared_libs()
+    )
+    # This lookup is necessary to make sure the compilation is done.
+    entry_func = export_engine.raw_lookup(
+        "_".join([prefix_name, jit_function.function_name])
+    )
+    if not entry_func:
+        raise DSLRuntimeError(
+            f"Execution engine cannot find the entry function {prefix_name}_{jit_function.function_name}"
+        )
+    try:
+        with tempfile.NamedTemporaryFile() as tmp_object_file:
+            export_engine.dump_to_object_file(tmp_object_file.name)
+            with open(tmp_object_file.name, "rb") as f:
+                ret = f.read()
+            return ret
+    except Exception as e:
+        raise DSLRuntimeError(f"Error dumping object file: {e}") from e
+
+
+def export_to_c(
+    jit_function: Union[JitCompiledFunction, "CudaDialectJitCompiledFunction"],
+    file_path: str,
+    file_name: str,
+    dsl: BaseDSL,
+    c_header_generator: CHeaderGenerator,
+    use_gpu_dialect: bool,
+):
+    """Exports the jit-compiled function to a C compatible files(header/library).
+    This is used for c/cpp AOT support.
+    The `file_name` will be used as the symbol prefix of the generated functions, it is guaranteed by
+    the caller that the generated functions are unique. And the function will always overwrite the existing file.
+
+
+    The c header file is generated with following components:
+    1. The host launch entry function. And the structure definitions of the arguments.
+    2. The device metadata load/unload functions.
+    3. The cubin data array and len.
+
+    The library contains the binary of the underlying host launch entry function.
+
+    @param jit_function: The jit-compiled function from `cute.compile`.
+    @param file_path: The path to the directory where the header and object files will be saved.
+    @param file_name: The name of the function. This is the unique identifier name of the function to avoid symbol conflict in the generated object file.
+    @param dsl: The dsl object. This is the dsl object to get the compiler provider and shared libs.
+    @param c_header_generator: The c header generator. This is the c header generator to generate the c header file.
+    """
+    export_module = get_export_module(jit_function.ir_module, file_name)
+    # Generate the c header file
+    header_file_content = c_header_generator(
+        file_name,
+        export_module,
+        jit_function.args_spec,
+        jit_function.function_name,
+        jit_function.kernel_info,
+        jit_function.dynamic_args,
+        jit_function.dynamic_kwargs,
+        dsl.name,
+    )
+    try:
+        with open(os.path.join(file_path, file_name + ".h"), "w") as f:
+            f.write(header_file_content)
+    except Exception as e:
+        raise DSLRuntimeError(f"Error writing header file: {e}") from e
+
+    # Generate the object file
+    object_file_content = dump_to_object(
+        file_name,
+        export_module,
+        jit_function,
+        dsl,
+        use_gpu_dialect,
+    )
+    try:
+        with open(os.path.join(file_path, file_name + ".o"), "wb") as f:
+            f.write(object_file_content)
+    except Exception as e:
+        raise DSLRuntimeError(f"Error writing object file: {e}") from e
+
+

python/CuTeDSL/cutlass/base_dsl/jit_executor.py

@@ -28,14 +28,13 @@ from .._mlir import ir
 
 # Local modules imports
 from . import typing as t
-from .common import DSLRuntimeError
+from .common import DSLRuntimeError, DSLCudaRuntimeError
 from .runtime import cuda as cuda_helpers
 from .runtime.jit_arg_adapters import JitArgAdapterRegistry, is_arg_spec_constexpr
 from .typing import get_c_pointers
 from .utils.logger import log
 from .utils.timer import timer
 
-
 class CudaModuleAndKernel:
     """A loaded CUDA kernel and its metadata."""
 
@@ -143,13 +142,11 @@ class ExecutionArgs:
             self.args_spec = self.filter_runtime_arg_spec(spec)
         self.original_args_spec = spec
 
-    def generate_execution_args(self, args, kwargs):
+    def get_rectified_args(self, args, kwargs):
         """
-        This function is the prune version of `generate_mlir_function_types` which only generates execution args
-        to get rid of mlir context.
+        This function is used to rectify the args and kwargs to a final runtime argument list according to the args_spec.
         """
         args_spec = self.args_spec
-
         # Process positional arguments with defaults
         rectified_args = list(args)
         if args_spec.defaults and len(args) < len(args_spec.args):
@@ -163,7 +160,9 @@ class ExecutionArgs:
                     rectified_args.append(v)
 
         # Process keyword arguments
-        rectified_kwargs = {k: v for k, v in kwargs.items() if k not in args_spec.args}
+        rectified_kwargs = collections.OrderedDict(
+            (k, v) for k, v in kwargs.items() if k not in args_spec.args
+        )
         if args_spec.kwonlydefaults and len(rectified_kwargs) < len(
             args_spec.kwonlyargs
         ):
@@ -182,10 +181,18 @@ class ExecutionArgs:
                     "function signature kwonlyargs length": len(args_spec.kwonlyargs),
                 },
             )
+        return rectified_args + list(rectified_kwargs.values())
+
+    def generate_execution_args(self, args, kwargs):
+        """
+        This function is the prune version of `generate_mlir_function_types` which only generates execution args
+        to get rid of mlir context.
+        """
+        args_spec = self.args_spec
 
         exe_args = []
         adapted_args = []
-        input_args = rectified_args + list(rectified_kwargs.values())
+        input_args = self.get_rectified_args(args, kwargs)
         input_arg_names = args_spec.args + args_spec.kwonlyargs
         for arg, arg_name in zip(input_args, input_arg_names):
             # short-cut for args already converted
@@ -243,7 +250,7 @@ class ExecutionArgs:
         runtime_kwonlydefaults = {}
 
         if arg_spec.kwonlyargs:
-            for kwarg in arg_spec.kwonlyargs:
+            for i, kwarg in enumerate(arg_spec.kwonlyargs):
                 arg_type = arg_spec.annotations.get(kwarg, None)
 
                 # Apply same filtering logic
@@ -366,6 +373,8 @@ class JitModule:
             for m in set([m.cuda_module for m in self.cuda_modules]):
                 cuda_helpers.unload_library(m)
             self.cuda_modules.clear()
+        except Exception as e:
+            pass
         finally:
             self._unloaded = True
 
@@ -382,8 +391,8 @@ class JitExecutor:
 
     def __init__(
         self,
-        jit_module: JitModule,
-        exec_context: JitExecuteContext,
+        jit_module: Union[JitModule, "CudaDialectJitModule"],
+        exec_context: Optional[JitExecuteContext],
         jit_time_profiling: bool,
     ):
         # JitExecutor will keep JitCompiledFunction alive so that the underlying
@@ -393,12 +402,25 @@ class JitExecutor:
         self.exec_context = exec_context
         self.profiler = timer(enable=jit_time_profiling)
 
+        # Get the cuda result type from the capi function.
+        # This is only set to i32 if CudaDialectJitModule is used.
+        cuda_result_type = self.jit_module.capi_func.restype
+        self.cuda_result = cuda_result_type() if cuda_result_type is not None else None
+
     # Assume each execution args has type `c_void_p` to reduce the overhead of `ctypes.cast`.
     def _get_invoke_packed_args(self, exe_args):
-        exe_args += self.exec_context.kernel_functions_ptrs
+        # If expecting a cuda result, add a pointer to exe_args
+        if self.cuda_result is not None:
+            exe_args.append(ctypes.addressof(self.cuda_result))
+        if self.exec_context is not None:
+            exe_args += self.exec_context.kernel_functions_ptrs
         packed_args = (ctypes.c_void_p * len(exe_args))()
         for argNum in range(len(exe_args)):
-            packed_args[argNum] = exe_args[argNum]
+            arg = exe_args[argNum]
+            if isinstance(arg, ctypes.c_void_p):
+                packed_args[argNum] = arg
+            else:
+                packed_args[argNum] = ctypes.c_void_p(arg).value
         return packed_args
 
     def generate_execution_args(self, *args, **kwargs):
@@ -408,6 +430,17 @@ class JitExecutor:
         try:
             packed_args = self.profiler(self._get_invoke_packed_args)(exe_args)
             self.profiler(self.jit_module.capi_func)(packed_args)
+            if self.cuda_result is not None:
+                if self.cuda_result.value != 0:
+                    error_code = self.cuda_result.value
+                    error_name = cuda_helpers._cudaGetErrorEnum(
+                        cuda_helpers.cuda.CUresult(error_code)
+                    )
+                    raise DSLCudaRuntimeError(error_code, error_name)
+                return self.cuda_result.value
+            return None
+        except DSLCudaRuntimeError as e:
+            raise e
         except Exception as e:
             raise DSLRuntimeError(f"💥💥💥 Runtime Crash 💥💥💥", cause=e)
 
@@ -458,6 +491,8 @@ class JitCompiledFunction:
         kernel_info,
         jit_time_profiling,
         jit_function_artifacts,
+        prefix=None,
+        load_from_binary=False,
     ):
         self.ir_module = ir_module
         self.engine = engine
@@ -473,6 +508,8 @@ class JitCompiledFunction:
             or jit_function_artifacts is None
         )
         self.artifacts = jit_function_artifacts
+        self.prefix = prefix
+        self.load_from_binary = load_from_binary
 
         # This runtime state is stored here so that we can preserve the module
         # in the compiler cache. Callers can extend the lifetime of the module
@@ -496,6 +533,35 @@ class JitCompiledFunction:
         """Returns the MLIR code of the JIT-compiled function."""
         return self.artifacts.MLIR if self.artifacts is not None else None
 
+    def _deserializer(self):
+        """Load the cuda module from the binary execution engine. This function will be injected as the
+        JitCompiledFunction method which will be called by the jit executor to load the cuda module by AOT flow.
+        @param self: The JitCompiledFunction object. This is the JitCompiledFunction object to load the cuda module.
+        @param name: The name of the function. This is the unique identifier name of the function to avoid symbol conflict in the generated object file.
+        @param execution_engine: The binary execution engine. This is the execution engine to load the cuda module.
+        @param kernel_info: The kernel info. This is the kernel info to load the cuda module.
+        @return: The list of cuda modules.
+        """
+        cubin_suffix = "cubin"
+        if self.prefix is None:
+            raise DSLRuntimeError("prefix is required to be set for binary loading")
+        cubin_data = self.engine.lookup("_".join([self.prefix, cubin_suffix]))
+        if not cubin_data:
+            raise RuntimeError(
+                "Unknown function " + "_".join([self.prefix, cubin_suffix])
+            )
+        cubin_module = cuda_helpers.load_library_data(cubin_data)
+        # load cuda module/get function pointer from module and cache
+        kernel_modules = collections.OrderedDict()
+        for sym, attrs in self.kernel_info.items():
+            kernel = cuda_helpers.get_library_kernel(cubin_module, sym)
+            if cuda_helpers.get_driver_version() >= 11080:
+                attrs[
+                    cuda_helpers.cuda.CUfunction_attribute.CU_FUNC_ATTRIBUTE_NON_PORTABLE_CLUSTER_SIZE_ALLOWED
+                ] = 1
+            kernel_modules[sym] = CudaModuleAndKernel(sym, cubin_module, kernel, attrs)
+        return list(kernel_modules.values())
+
     def to(self, device=None) -> JitExecutor:
         """Returns an executable function bound to the given device.
 
@@ -510,9 +576,10 @@ class JitCompiledFunction:
         with self._executor_lock:
             # We need to ensure that the modules are loaded if not already
             if self.jit_module is None:
-                cuda_modules = load_kernels_from_ir_module(
-                    self.ir_module, self.kernel_info
-                )
+                if self.ir_module is not None:
+                    cuda_modules = load_kernels_from_ir_module(
+                        self.ir_module, self.kernel_info
+                    )
                 self.jit_module = JitModule(
                     self.engine, self.capi_func, self.args_spec, cuda_modules
                 )
@@ -522,6 +589,11 @@ class JitCompiledFunction:
             context = self.jit_module.get_device_execute_context(device)
             return JitExecutor(self.jit_module, context, self.jit_time_profiling)
 
+    def set_dynamic_args(self, dynamic_args, dynamic_kwargs):
+        """Sets the dynamic argument information required for export to c code generation."""
+        self.dynamic_args = dynamic_args
+        self.dynamic_kwargs = dynamic_kwargs
+
     def generate_execution_args(self, *args, **kwargs):
         return self.args_spec.generate_execution_args(args, kwargs)
 

python/CuTeDSL/cutlass/base_dsl/runtime/cuda.py

@@ -23,15 +23,9 @@ import ctypes
 import cuda.bindings.driver as cuda
 import cuda.bindings.nvrtc as nvrtc
 
-# MLIR imports
-from ..._mlir import ir
-from ..._mlir.dialects import gpu
-
 # Local module imports
 from ..utils.logger import log as _log
 from ..common import *
-from .jit_arg_adapters import JitArgAdapterRegistry
-
 
 # =============================================================================
 # Utils
@@ -504,18 +498,18 @@ def load_library_data(cubin_data):
     """
     Loads a CUBIN from data and returns the library.
     :param cubin_data: The binary data of the CUBIN.
-    :type cubin_data: bytes
+    :type cubin_data: bytes or ctypes.c_void_p
     :return: The library.
     :rtype: cuda.CUlibrary
     :raise DSLRuntimeError: If the CUDA operation fails.
     """
     # Load module data
-    _log().info(f"cuLibraryLoadData {np.char.array(cubin_data).ctypes.data}")
+    if isinstance(cubin_data, bytes):
+        cubin_data = np.char.array(cubin_data).ctypes.data
+    _log().info(f"cuLibraryLoadData {cubin_data}")
 
     library = checkCudaErrors(
-        cuda.cuLibraryLoadData(
-            np.char.array(cubin_data).ctypes.data, None, None, 0, None, None, 0
-        )
+        cuda.cuLibraryLoadData(cubin_data, None, None, 0, None, None, 0)
     )
     return library
 
@@ -803,24 +797,3 @@ def get_device_attribute(attribute, device_id: int = 0):
     """
     device = checkCudaErrors(cuda.cuDeviceGet(device_id))
     return checkCudaErrors(cuda.cuDeviceGetAttribute(attribute, device))
-
-
-@JitArgAdapterRegistry.register_jit_arg_adapter(cuda.CUstream)
-class StreamAdapter:
-    """
-    Convert a CUDA stream to a stream representation for JIT arg generation.
-    """
-
-    def __init__(self, arg):
-        self._arg = arg
-        self._c_pointer = self._arg.getPtr()
-
-    def __new_from_mlir_values__(self, values):
-        assert len(values) == 1
-        return values[0]
-
-    def __c_pointers__(self):
-        return [self._c_pointer]
-
-    def __get_mlir_types__(self):
-        return [gpu.AsyncTokenType.get()]

python/CuTeDSL/cutlass/base_dsl/runtime/stream_adapter.py

@@ -0,0 +1,43 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 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/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.
+
+"""
+This module provides CUDA Python helper functions
+"""
+
+import cuda.bindings.driver as cuda
+
+# MLIR imports
+from ..._mlir import ir
+from ..._mlir.dialects import gpu
+
+from .jit_arg_adapters import JitArgAdapterRegistry
+
+
+@JitArgAdapterRegistry.register_jit_arg_adapter(cuda.CUstream)
+class StreamAdapter:
+    """
+    Convert a CUDA stream to a stream representation for JIT arg generation.
+    """
+
+    def __init__(self, arg):
+        self._arg = arg
+        self._c_pointer = self._arg.getPtr()
+
+    def __new_from_mlir_values__(self, values):
+        assert len(values) == 1
+        return values[0]
+
+    def __c_pointers__(self):
+        return [self._c_pointer]
+
+    def __get_mlir_types__(self):
+        return [gpu.AsyncTokenType.get()]

python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/README.md

@@ -0,0 +1,84 @@
+# TVM FFI DSL Bridge
+
+This folder implements a binding generator that
+can be used by DSL compilers to generate functions in tvm-ffi ABI.
+
+- Interface specification: we provide a spec namespace that allows
+  developers to define the interface of the function, with possibly shape and type constraints.
+- Constraint checking: we will generate MLIR functions that exposes the tvm-ffi ABI convention
+  checks the overall constraints and reads all related variable fields.
+- User defined calling convention: user can provide a `CallProvider` instance
+  that takes `CallContext` and emit a call to internal DSL function.
+
+## Basic Example: Mixed Parameters
+
+```python
+
+from cutlass.base_dsl.tvm_ffi_builder import (
+  spec, attach_ffi_func, ExecutionEngine, NopProvider
+)
+from cutlass._mlir import ir
+import tvm_ffi
+import numpy as np
+
+# Define parameters: int + tensor with symbolic shape
+n = spec.Var("n", "int32")
+params = [
+    spec.Var("batch_size", "int32"),           # Integer parameter
+    spec.Tensor("data", [n, 128], "float32"),  # Tensor with symbolic first dimension
+]
+
+# Generate function
+# Function signature: process_data(batch_size: int32, data: Tensor([n, 128], float32))
+with ir.Context(), ir.Location.unknown():
+    module = ir.Module.create()
+    attach_ffi_func(module, "process_data", params, NopProvider())
+
+    # Execute
+    engine = ExecutionEngine(module, opt_level=2, shared_libs=[])
+    func = tvm_ffi.Function.__from_mlir_packed_safe_call__(
+      engine.raw_lookup("__tvm_ffi_process_data")
+    )
+
+    # Call function
+    data = tvm_ffi.from_dlpack(np.zeros((10, 128), dtype=np.float32))
+    func(42, data)  # batch_size=42, data shape=[10, 128]
+```
+
+## Matrix Multiplication Example
+
+```python
+# Define matrix multiplication with automatic constraint validation
+n = spec.Var("n", "int32")
+m = spec.Var("m", "int32")
+k = spec.Var("k", "int32")
+
+with spec.DefaultConfig(device="cpu"):
+  params = [
+      spec.Tensor("A", [n, k], "float32"),  # A: n×k
+      spec.Tensor("B", [k, m], "float32"),  # B: k×m
+      spec.Tensor("C", [n, m], "float32"),  # C: n×m
+  ]
+
+# Generate function
+# Function signature: matmul(A: Tensor([n, k], float32), B: Tensor([k, m], float32), C: Tensor([n, m], float32))
+with ir.Context(), ir.Location.unknown():
+    module = ir.Module.create()
+    attach_ffi_func(module, "matmul", params, NopProvider())
+
+    # Execute
+    engine = ExecutionEngine(module, opt_level=2, shared_libs=[])
+    func = tvm_ffi.Function.__from_mlir_packed_safe_call__(
+      engine.raw_lookup("__tvm_ffi_matmul")
+    )
+
+    # Valid call: 2×3 × 3×4 = 2×4
+    A = tvm_ffi.from_dlpack(np.zeros((2, 3), dtype=np.float32))
+    B = tvm_ffi.from_dlpack(np.zeros((3, 4), dtype=np.float32))
+    C = tvm_ffi.from_dlpack(np.zeros((2, 4), dtype=np.float32))
+    func(A, B, C)  # ✅ Valid dimensions
+
+    # Invalid call: dimension mismatch
+    A_wrong = tvm_ffi.from_dlpack(np.zeros((2, 4), dtype=np.float32))  # Wrong: 2×4
+    func(A_wrong, B, C)  # ❌ Error: A[1]=4 != B[0]=3
+```

python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/__init__.py

@@ -0,0 +1,32 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 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/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.
+
+"""Helper tool to build TVM-FFI functions using MLIR."""
+
+from .call_provider import DynamicParamPackCallProvider, NopCallProvider
+from .spec import Param, Var
+from .tvm_ffi_builder import (
+    CallContext,
+    CallProvider,
+    attach_ffi_func,
+    rename_tvm_ffi_function,
+)
+
+__all__ = [
+    "CallContext",
+    "CallProvider",
+    "DynamicParamPackCallProvider",
+    "NopCallProvider",
+    "Param",
+    "Var",
+    "attach_ffi_func",
+    "rename_tvm_ffi_function",
+]

python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/call_provider.py

@@ -0,0 +1,245 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 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/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.
+
+"""Call provider that implements a specific calling convention."""
+
+from dataclasses import dataclass
+from typing import Optional, Union
+
+from . import spec
+from ..._mlir import ir
+from ..._mlir.dialects import llvm
+from .tvm_ffi_builder import CallContext, CallProvider, TVMFFIBuilder
+
+
+class NopCallProvider(CallProvider):
+    """No-op call provider for testing purposes."""
+
+    def __call__(self, current_block: ir.Block, context: CallContext) -> ir.Block:
+        """No-op call provider that just returns the current block."""
+        return current_block
+
+
+class DynamicParamPackCallProvider(CallProvider, TVMFFIBuilder):
+    """Packs dynamic arguments to a struct then calls the function.
+
+    .. code-block:: c
+
+        void call(Tensor0 t0, Tensor1 t1) {
+            // packed arguments
+            void** packed_args[] = {&t0, &t1};
+            // call target
+            target_func(packed_args);
+        }
+
+    Parameters
+    ----------
+    target_func: str
+        The name of the target function.
+
+    include_num_args: bool
+        Whether to include the number of arguments in the packed arguments.
+
+    struct_call: bool
+        Whether to use the struct call convention.
+    """
+
+    def __init__(
+        self, target_func: str, include_num_args: bool = False, struct_call: bool = False,
+    ) -> None:
+        import tvm_ffi
+
+        TVMFFIBuilder.__init__(self)
+        self.target_func = target_func
+        self.include_num_args = include_num_args
+        self.struct_call = struct_call
+        self.float4x2_dtype = tvm_ffi.dtype("float4_e2m1fnx2")
+
+    def get_callee_struct_for_param_tensor(
+        self,
+        param: spec.Tensor,
+        current_block: ir.Block,
+        data: ir.Value,
+        shape: list[ir.Value],
+        strides: list[ir.Value],
+        flatten_struct: ir.Type,
+    ) -> ir.Type:
+        """Routine used to override tensor passsing struct conention."""
+        return flatten_struct
+
+    def pack_param_tensor(
+        self, current_block: ir.Block, context: CallContext, param: spec.Tensor
+    ) -> tuple[ir.Type, ir.Value]:
+        """Pack a tensor parameter to a struct."""
+        map_shape_value = lambda _, value: value
+        map_stride_value = lambda _, value: value
+
+        if param.map_tensor_dtype_f4x2_to_f4 and param.dtype == self.float4x2_dtype:
+            # specially handle f4x2 to f4 tensor conversion
+            # we multiply the stride by 2 for all dimensions except the one with stride=1
+            # we also multiply the shape by 2 for the specific dimension with stride=1
+            def find_stride_one_index() -> int:
+                if param.strides is None:
+                    return len(param.shape) - 1
+                for i, stride in enumerate(param.strides):
+                    if isinstance(stride, int) and stride == 1:
+                        return i
+                raise ValueError("stride=1 index not found, needed for f4 tensor")
+            stride_one_index = find_stride_one_index()
+
+            def map_shape_for_tensor_dtype_f4x2_to_f4(index: int, value: ir.Value) -> ir.Value:
+                if index == stride_one_index:
+                    with ir.InsertionPoint(current_block):
+                        return self.mul(value, self.integer_constant(value.type, 2))
+                return value
+
+            def map_stride_for_tensor_dtype_f4x2_to_f4(index, value: ir.Value) -> ir.Value:
+                if index != stride_one_index:
+                    with ir.InsertionPoint(current_block):
+                        return self.mul(value, self.integer_constant(value.type, 2))
+                return value
+
+            map_shape_value = map_shape_for_tensor_dtype_f4x2_to_f4
+            map_stride_value = map_stride_for_tensor_dtype_f4x2_to_f4
+
+        data = context.matched_var_binding[param.data]
+        shape = []
+        strides = []
+        # append all vars in shape
+        for index, dim in enumerate(param.shape):
+            if isinstance(dim, spec.Var):
+                shape.append(map_shape_value(index, context.matched_var_binding[dim]))
+        # append all vars in strides
+        if param.strides is not None:
+            for index, dim in enumerate(param.strides):
+                if isinstance(dim, spec.Var):
+                    strides.append(map_stride_value(index, context.matched_var_binding[dim]))
+        flatten_struct, alloca = self.pack_values_to_alloca(
+            current_block, context.entry_block, [data, *shape, *strides]
+        )
+        callee_struct = self.get_callee_struct_for_param_tensor(
+            param, current_block, data, shape, strides, flatten_struct
+        )
+
+        return callee_struct, alloca
+
+    def pack_param_var(
+        self, current_block: ir.Block, context: CallContext, param: spec.Var
+    ) -> tuple[ir.Type, ir.Value]:
+        """Pack a var parameter to a struct."""
+        value: ir.Value = context.matched_var_binding[param]
+        _, alloca = self.pack_values_to_alloca(
+            current_block, context.entry_block, [value]
+        )
+        return (value.type, alloca)
+
+    def pack_param_shape(
+        self, current_block: ir.Block, context: CallContext, param: spec.Shape
+    ) -> tuple[ir.Type, ir.Value]:
+        """Pack a shape parameter to a struct."""
+        dynamic_args: list[ir.Value] = []
+        for dim in param.shape:
+            if isinstance(dim, spec.Var):
+                dynamic_args.append(context.matched_var_binding[dim])
+        return self.pack_values_to_alloca(
+            current_block, context.entry_block, dynamic_args
+        )
+
+    def pack_params(
+        self, current_block: ir.Block, context: CallContext
+    ) -> list[tuple[ir.Type, ir.Value]]:
+        """Pack a parameter to a struct."""
+        packed_params = []
+        for param in context.params:
+            if isinstance(param, spec.Tensor):
+                packed_params.append(
+                    self.pack_param_tensor(current_block, context, param)
+                )
+            elif isinstance(param, spec.Var):
+                packed_params.append(self.pack_param_var(current_block, context, param))
+            elif isinstance(param, spec.Shape):
+                packed_params.append(
+                    self.pack_param_shape(current_block, context, param)
+                )
+            elif isinstance(param, (spec.Stream, spec.EnvStream)):
+                packed_params.append(
+                    self.pack_param_var(current_block, context, param.var)
+                )
+            elif isinstance(param, spec.DataPointer):
+                packed_params.append(
+                    self.pack_param_var(current_block, context, param.var)
+                )
+            else:
+                raise NotImplementedError(f"Unsupported parameter type: {type(param)}")
+        return packed_params
+
+    def generate_llvm_call(
+        self,
+        current_block: ir.Block,
+        call_operands: list[ir.Value],
+        context: CallContext,
+    ) -> ir.Block:
+        """Generate the LLVM call operation."""
+        with ir.InsertionPoint(current_block):
+            llvm.call(
+                result=None,
+                callee=self.target_func,
+                callee_operands=call_operands,
+                op_bundle_sizes=[],
+                op_bundle_operands=[],
+            )
+        return current_block
+
+    def load_to_call_operands(
+        self,
+        struct_type: Union[ir.Type, tuple[ir.Type]],
+        alloca: Union[ir.Value, tuple[ir.Value]],
+    ) -> list[ir.Value]:
+        """Load the packed parameters to the call operands."""
+        assert (isinstance(struct_type, ir.Type) and isinstance(alloca, ir.Value)) or (
+            isinstance(struct_type, tuple) and isinstance(alloca, tuple)
+        )
+        if isinstance(struct_type, tuple):
+            return [
+                llvm.load(struct_type[i], alloca[i]) for i in range(len(struct_type))
+            ]
+        return [llvm.load(struct_type, alloca)]
+
+    def __call__(self, current_block: ir.Block, context: CallContext) -> ir.Block:
+        """Alloca call provider that uses dynamic param pack call convention."""
+        packed_params = self.pack_params(current_block, context)
+
+        if self.struct_call:
+            # load back arguments as structs from alloca
+            call_operands = []
+            with ir.InsertionPoint(current_block):
+                for struct_type, alloca in packed_params:
+                    call_operands += self.load_to_call_operands(struct_type, alloca)
+        else:
+            # pack the values to an alloca that we can pass as void**
+            all_values = []
+            for _, value in packed_params:
+                if isinstance(value, tuple):
+                    all_values.extend(value)
+                else:
+                    all_values.append(value)
+            _, packed_args_value = self.pack_values_to_alloca(
+                current_block, context.entry_block, all_values
+            )
+
+            call_operands = [packed_args_value]
+            if self.include_num_args:
+                with ir.InsertionPoint(current_block):
+                    num_args = self.i32(len(all_values))
+                    call_operands.append(num_args)
+
+        current_block = self.generate_llvm_call(current_block, call_operands, context)
+        return current_block

python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/mlir_builder.py

@@ -0,0 +1,468 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 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/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.
+
+"""MLIR type builder and basic operations for LLVM dialect."""
+
+from collections.abc import Sequence
+from typing import Any, Optional
+
+from ..._mlir import ir
+from ..._mlir.dialects import llvm
+
+
+class MLIRTypeBuilder:
+    """Builder for MLIR types and basic operations."""
+
+    BRANCH_WEIGHTS_LIKELY = (2000, 1)
+    BRANCH_WEIGHTS_UNLIKELY = (1, 2000)
+
+    def __init__(self) -> None:
+        """Initialize the MLIR type builder."""
+        self.i32_type = ir.IntegerType.get_signless(32)
+        self.ui32_type = ir.IntegerType.get_unsigned(32)
+        self.i64_type = ir.IntegerType.get_signless(64)
+        self.i16_type = ir.IntegerType.get_signless(16)
+        self.i8_type = ir.IntegerType.get_signless(8)
+        self.i1_type = ir.IntegerType.get_signless(1)
+        self.f32_type = ir.Type.parse("f32")
+        self.f64_type = ir.Type.parse("f64")
+        self.ptr_type = llvm.PointerType.get()
+        self.gpu_ptr_type = llvm.PointerType.get(address_space=1)
+        # did not find a programmatic way to get the void type
+        self.void_type = ir.Type.parse("!llvm.void")
+
+    def ptr_type_with_address_space(
+        self, address_space: Optional[int] = None
+    ) -> ir.Type:
+        """Get the pointer type with the given address space."""
+        if address_space is None or address_space == 0:
+            return self.ptr_type
+        return llvm.PointerType.get(address_space=address_space)
+
+    def as_attr(self, tp: ir.Type) -> ir.TypeAttr:
+        """Convert the type to a type attribute."""
+        return ir.TypeAttr.get(tp)
+
+    def int_type(self, bits: int) -> ir.Type:
+        """Get the `i<bits>` type."""
+        return ir.IntegerType.get_signless(bits)
+
+    def uint_type(self, bits: int) -> ir.Type:
+        """Get the `ui<bits>` type."""
+        return ir.IntegerType.get_unsigned(bits)
+
+    def struct_type(
+        self,
+        *,
+        name: Optional[str] = None,
+        fields: Sequence[ir.Type] = (),
+        packed: bool = False,
+    ) -> ir.Type:
+        """Get or create a struct type.
+
+        Parameters
+        ----------
+        name : Optional[str]
+            The name of the struct type. If not provided, a `literal` struct type is created,
+            which is identified by its fields only.
+        fields : Sequence[ir.Type]
+            The fields of the struct type.
+        packed : bool
+            Whether to create a packed struct type. If `True`, there is no padding between fields.
+            Otherwise, there might be padding
+            between fields to ensure alignment.
+
+        See Also
+        --------
+        https://mlir.llvm.org/docs/Dialects/LLVM/#structure-types
+
+        """
+        if name is None:
+            return llvm.StructType.get_literal(fields, packed=packed)
+        else:
+            return llvm.StructType.new_identified(name, fields, packed=packed)
+
+    def identified_struct_type(self, name: str) -> ir.Type:
+        """Get a previously created identified struct type by its name."""
+        return llvm.StructType.get_identified(name)
+
+    def func_type(self, *, params: Sequence[ir.Type] = (), ret: ir.Type) -> ir.Type:
+        """Get a function type.
+
+        Parameters
+        ----------
+        params : Sequence[ir.Type]
+            The parameters of the function type.
+        ret : ir.Type
+            The return type of the function type.
+
+        See Also
+        --------
+        https://mlir.llvm.org/docs/Dialects/LLVM/#function-types
+
+        """
+        return ir.Type.parse(
+            "!llvm.func<{} ({})>".format(str(ret), ", ".join(map(str, params)))
+        )  # did not find a programmatic way to get the function type
+
+    def global_dtor_entry_type(self) -> ir.Type:
+        """Get the type of the global destructor entry.
+
+        See Also:
+        - https://llvm.org/docs/LangRef.html#the-llvm-global-dtors-global-variable
+        - https://mlir.llvm.org/docs/Dialects/LLVM/#llvmmlirglobal_dtors-llvmglobaldtorsop
+        """
+        return self.struct_type(fields=[self.i32_type, self.ptr_type, self.ptr_type])
+
+
+class MLIRBuilder(MLIRTypeBuilder):
+    """A builder for MLIR related types and operations.
+
+    Convention:
+    all statement-generation methods expect we are inside a insertion point of the current_block.
+    If the statement terminates the current block, it will assign the new block to the current_block
+    but not set the insersion point.
+    """
+
+    def __init__(self) -> None:
+        """Initialize the MLIR builder."""
+        super().__init__()
+        self.module: Optional[ir.Module] = None
+        self.const_str_table: dict[str, ir.Value] = {}
+        self.get_element_extra_kwargs: dict[str, Any] = {}
+
+    # create constants
+    def integer_constant(self, tp: ir.Type, value: int) -> ir.Value:
+        """Create an integer constant with the given type and value."""
+        return llvm.ConstantOp(tp, ir.IntegerAttr.get(tp, value)).res
+
+    def i32(self, value: int) -> ir.Value:
+        """Create an i32 constant with the given value."""
+        return self.integer_constant(self.i32_type, value)
+
+    def ui32(self, value: int) -> ir.Value:
+        """Create a ui32 constant with the given value."""
+        return self.integer_constant(self.ui32_type, value)
+
+    def i1(self, value: int) -> ir.Value:
+        """Create an i1 constant with the given value."""
+        return self.integer_constant(self.i1_type, value)
+
+    def i8(self, value: int) -> ir.Value:
+        """Create an i8 constant with the given value."""
+        return self.integer_constant(self.i8_type, value)
+
+    def i16(self, value: int) -> ir.Value:
+        """Create an i16 constant with the given value."""
+        return self.integer_constant(self.i16_type, value)
+
+    def i64(self, value: int) -> ir.Value:
+        """Create an i64 constant with the given value."""
+        return self.integer_constant(self.i64_type, value)
+
+    def mul(self, lhs: ir.Value, rhs: ir.Value) -> ir.Value:
+        """Create a multiplication operation between two values."""
+        return llvm.mul(lhs, rhs, overflow_flags=llvm.IntegerOverflowFlags.none)
+
+    # expressions
+    def not_equal(self, lhs: ir.Value, rhs: ir.Value) -> ir.Value:
+        """Create a not-equal comparison between two values."""
+        return llvm.icmp(llvm.ICmpPredicate.ne, lhs, rhs)
+
+    def equal(self, lhs: ir.Value, rhs: ir.Value) -> ir.Value:
+        """Create an equal comparison between two values."""
+        return llvm.icmp(llvm.ICmpPredicate.eq, lhs, rhs)
+
+    def or_(self, lhs: ir.Value, rhs: ir.Value) -> ir.Value:
+        """Create a logical OR operation between two values."""
+        return llvm.or_(lhs, rhs)
+
+    def and_(self, lhs: ir.Value, rhs: ir.Value) -> ir.Value:
+        """Create a logical AND operation between two values."""
+        return llvm.and_(lhs, rhs)
+
+    def not_(self, value: ir.Value) -> ir.Value:
+        """Create a logical NOT operation."""
+        # Ensure we're working with i1 type for boolean operations
+        if value.type != self.i1_type:
+            value = llvm.trunc(res=self.i1_type, arg=value)
+        return llvm.xor(value, self.i1(1))
+
+    def i64_divisible_const(self, value: ir.Value, align_const: int) -> ir.Value:
+        """Check if i64 value is divisible by align_const.
+
+        Parameters
+        ----------
+        value : ir.Value
+            The i64 value to check.
+        align_const : int
+            The alignment constant to check divisibility against.
+
+        Returns
+        -------
+        ir.Value
+            A boolean value (i1) indicating if value is divisible by align_const.
+
+        Notes
+        -----
+        Uses fast path (bitwise AND) when align_const is a power of two,
+        otherwise uses modulo operation.
+        """
+        # Check if align_const is a power of two
+        is_power_of_two = (align_const > 0) and (align_const & (align_const - 1)) == 0
+
+        if is_power_of_two:
+            # Fast path: use bitwise AND
+            # value is divisible by align_const iff (value & (align_const - 1)) == 0
+            mask = self.i64(align_const - 1)
+            masked = llvm.and_(value, mask)
+            return self.equal(masked, self.i64(0))
+        else:
+            # Slow path: use modulo operation
+            align_val = self.i64(align_const)
+            remainder = llvm.urem(value, align_val)
+            return self.equal(remainder, self.i64(0))
+
+    def br(self, target_block: ir.Block, *, args: Optional[list[ir.Value]] = None) -> None:
+        """Create an unconditional branch.
+
+        Parameters
+        ----------
+        target_block : ir.Block
+            The target block to branch to.
+        args : list[ir.Value], optional
+            The values to pass as arguments to the target block. If None,
+            no arguments are passed. The target block must have the same
+            number of arguments as values in this list.
+
+        """
+        if args is None:
+            llvm.br(dest=target_block, dest_operands=[])
+        else:
+            llvm.br(dest_operands=args, dest=target_block)
+
+    def address_of(self, name: str, tp: ir.Type) -> ir.Value:
+        """Get the address of a global symbol."""
+        return llvm.AddressOfOp(tp, name).result
+
+    def getelementptr(
+        self, ptr: ir.Value, constant_indices: Sequence[int], elem_type: ir.Type
+    ) -> ir.Value:
+        """Create a getelementptr operation.
+
+        Parameters
+        ----------
+        ptr : ir.Value
+            The pointer to the element.
+        indices : Sequence[ir.Value]
+            The indices to the element.
+        elem_type : ir.Type
+            The type of the element.
+
+        Returns
+        -------
+        ir.Value
+            The resulting element pointer.
+        """
+        try:
+            return llvm.getelementptr(
+                self.ptr_type,
+                ptr,
+                [],
+                raw_constant_indices=ir.DenseI32ArrayAttr.get(constant_indices),
+                elem_type=elem_type,
+                **self.get_element_extra_kwargs,
+            )
+        except Exception:
+            # compatibility with different LLVM versions
+            self.get_element_extra_kwargs = {"no_wrap_flags": []}
+            return llvm.getelementptr(
+                self.ptr_type,
+                ptr,
+                [],
+                raw_constant_indices=ir.DenseI32ArrayAttr.get(constant_indices),
+                elem_type=elem_type,
+                **self.get_element_extra_kwargs,
+            )
+
+    def return_(self, ret: Optional[ir.Value] = None) -> None:
+        """Create a return statement."""
+        llvm.return_(arg=ret)
+
+    def cond_br(
+        self,
+        cond: ir.Value,
+        true_block: ir.Block,
+        false_block: ir.Block,
+        branch_weights=None
+    ) -> None:
+        """Create a conditional branch.
+
+        Parameters
+        ----------
+        cond : ir.Value
+            The condition value (i1 type).
+        true_block : ir.Block
+            The block to branch to if condition is true.
+        false_block : ir.Block
+            The block to branch to if condition is false.
+        branch_weights : Optional[tuple[int, int]]
+            Optional branch weights [true_weight, false_weight] for optimization hints.
+            Higher values indicate higher probability. For example, (99, 1) indicates
+            the true branch is much more likely than the false branch.
+        """
+        if branch_weights is not None:
+            # Branch weights should be a tuple/list of two integers [true_weight, false_weight]
+            if len(branch_weights) != 2:
+                raise ValueError("branch_weights must have exactly 2 elements")
+            llvm.cond_br(
+                cond,
+                true_dest_operands=[],
+                false_dest_operands=[],
+                true_dest=true_block,
+                false_dest=false_block,
+                branch_weights=ir.DenseI32ArrayAttr.get(list(branch_weights))
+            )
+        else:
+            llvm.cond_br(
+                cond,
+                true_dest_operands=[],
+                false_dest_operands=[],
+                true_dest=true_block,
+                false_dest=false_block,
+            )
+
+    def define_global_string(self, content: str) -> str:
+        """Define a global string symbol with the given content using standard MLIR APIs."""
+        if content in self.const_str_table:
+            return self.const_str_table[content]
+        symbol = f"__tvm_ffi__str_{len(self.const_str_table)}"
+
+        # The string_attr.value gives us the original string, but we need to escape it for
+        # MLIR parsing. Let's use a simple approach: escape quotes and backslashes
+        escaped_content = content.replace("\\", "\\\\").replace('"', '\\"')
+
+        # Parse the MLIR string with proper escaping using the standard API
+        module_body = self.module.body  # type: ignore[union-attr]
+        with ir.InsertionPoint(module_body):
+            parsed_op = ir.Operation.parse(
+                f'llvm.mlir.global private constant @{symbol}("{escaped_content}\\00")'
+            )
+            module_body.append(parsed_op)
+            self.const_str_table[content] = symbol
+        return symbol
+
+    # function
+    def function(
+        self,
+        name: str,
+        params_type: Sequence[ir.Type],
+        ret_type: ir.Type,
+        internal: bool = False,
+    ) -> tuple[list[ir.Value], ir.Block]:
+        """Create a function with the given signature."""
+        func_op = llvm.func(
+            name,
+            function_type=self.as_attr(
+                self.func_type(ret=ret_type, params=params_type)
+            ),
+        )
+        if internal:
+            func_op.attributes["llvm.linkage"] = ir.StringAttr.get("private")
+        else:
+            func_op.attributes["llvm.emit_c_interface"] = ir.UnitAttr.get()
+
+        params = []
+        func_body: Any = func_op.body
+        if func_body is not None:
+            entry_block = ir.Block.create_at_start(func_body)
+        else:
+            raise RuntimeError("Function body is None")
+        for param_type in params_type:
+            params.append(entry_block.add_argument(param_type, ir.Location.unknown()))
+        return params, entry_block
+
+    def declare_extern_func(
+        self, name: str, params: Sequence[ir.Type], ret: ir.Type
+    ) -> None:
+        """Define extern function."""
+        func_type = self.func_type(params=params, ret=ret)
+        func_op = llvm.func(
+            name,
+            function_type=self.as_attr(func_type),
+        )
+        func_op.attributes["llvm.linkage"] = ir.StringAttr.get("external")
+
+    def pack_values_to_alloca(
+        self,
+        current_block: ir.Block,
+        entry_block: ir.Block,
+        values: Sequence[ir.Value],
+    ) -> tuple[ir.Type, ir.Value]:
+        """Pack values to an alloca that lays out in the order of the values.
+
+        Parameters
+        ----------
+        current_block : ir.Block
+            The current block.
+        entry_block : ir.Block
+            The entry block to create an alloca for the struct.
+        values : Sequence[ir.Value]
+            The values to pack.
+
+        Returns
+        -------
+        tuple[ir.Type, ir.Value]
+            The struct type and the alloca.
+        """
+        with ir.InsertionPoint(entry_block.operations[0]):
+            # declare the struct type
+            struct_type = self.struct_type(fields=[value.type for value in values])
+            alloca = llvm.alloca(
+                res=self.ptr_type,
+                elem_type=struct_type,
+                array_size=self.i32(1),
+            )
+
+        with ir.InsertionPoint(current_block):
+            for index, value in enumerate(values):
+                # Get pointer to the field at the given index
+                field_ptr = self.getelementptr(
+                    alloca,
+                    [0, index],
+                    struct_type,
+                )
+                # Store the value into the field
+                llvm.store(value, field_ptr)
+        return (struct_type, alloca)
+
+    def find_operations_in_module(
+        self, module: ir.Module, name: str
+    ) -> list[ir.Operation]:
+        """Find operations in the module by the operation name."""
+        operations = []
+        for op in module.body:  # type: ignore[union-attr]
+            if op.name == name:
+                operations.append(op)
+        return operations
+
+    def find_func_in_module(
+        self, module: ir.Module, name: str
+    ) -> Optional[ir.Operation]:
+        """Find a function in the module."""
+        for op in module.body:  # type: ignore[union-attr]
+            if op.name == "llvm.func":
+                # Get the function name from the sym_name attribute
+                if "sym_name" in op.attributes:
+                    func_name = str(op.attributes["sym_name"]).strip('"')
+                    if func_name == name:
+                        return op
+        return None

python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/spec.py

@@ -0,0 +1,431 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 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/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.
+
+"""Kernel specification classes for TVM-FFI function parameters."""
+from abc import ABC
+
+from collections.abc import Sequence
+from typing import Optional, Union
+
+try:
+    import tvm_ffi
+except ModuleNotFoundError:
+    pass
+
+class DefaultConfig:
+    """Default configuration with context manager support."""
+
+    _current: Optional["DefaultConfig"] = None
+    _old_current: Optional["DefaultConfig"] = None
+    device_type: str
+
+    def __init__(self, *, device_type: Optional[str] = None) -> None:
+        """Initialize a default configuration.
+
+        Parameters
+        ----------
+        device_type : str, optional
+            The device type (e.g., "cpu", "cuda", "metal"). If None, copies from current config.
+        """
+        if device_type is None:
+            device_type = DefaultConfig.current().device_type  # type: ignore[union-attr]
+        self.device_type = device_type
+
+    def __enter__(self) -> "DefaultConfig":
+        """Enter the context manager."""
+        self._old_current = DefaultConfig._current
+        DefaultConfig._current = self
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Optional[type],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[object],
+    ) -> None:
+        """Exit the context manager."""
+        DefaultConfig._current = self._old_current
+
+    @classmethod
+    def current(cls) -> Optional["DefaultConfig"]:
+        """Get the current default configuration.
+
+        Returns
+        -------
+        Optional[DefaultConfig]
+            The current default configuration.
+        """
+        return cls._current
+
+    @classmethod
+    def _set_init_default_config(cls) -> None:
+        """Set the initial default configuration."""
+        current = cls.__new__(cls)
+        current.device_type = "cuda"
+        cls._current = current
+
+
+# Initialize the default config to cuda
+DefaultConfig._set_init_default_config()
+
+
+class Param(ABC):
+    """Base class for all parameters."""
+
+
+class Var(Param):
+    """variables: pointer, integer, floating-point, boolean, etc.
+
+    Parameters
+    ----------
+    name : str
+        The parameter name.
+    dtype : str | tvm_ffi.dtype
+        The data type of the parameter.
+    divisibility: Optional[int]
+        The divisibility of the parameter, by default None.
+
+    """
+
+    name: str
+    dtype: "tvm_ffi.dtype"
+    divisibility: Optional[int]
+
+    def __init__(
+        self,
+        name: str,
+        dtype: Union[str, "tvm_ffi.dtype"],
+        *,
+        divisibility: Optional[int] = None,
+    ) ->  None:
+        """Initialize a Var parameter.
+
+        Parameters
+        ----------
+        name : str
+            The parameter name.
+        dtype : str | tvm_ffi.dtype
+            The data type of the parameter.
+
+        """
+        self.name = name
+        self.dtype = tvm_ffi.dtype(dtype)
+        self.divisibility = divisibility
+
+
+class Shape(Param):
+    """Shape parameter.
+
+    Parameters
+    ----------
+    name : str
+        The parameter name.
+    shape: list[int | Var]
+        The shape of the parameter.
+
+    """
+
+    name: str
+    shape: list[Union[int, Var]]
+
+    def __init__(self, name: str, shape: list[Union[int, Var]]) -> None:
+        """Initialize a Shape parameter.
+
+        Parameters
+        ----------
+        name : str
+            The parameter name.
+        shape : list[int | Var]
+            The shape of the parameter.
+
+        """
+        self.name = name
+        self.shape = shape
+
+
+class Tensor(Param):
+    """Tensor parameter.
+
+    Parameters
+    ----------
+    name : str
+        The parameter name.
+    dtype : str | tvm_ffi.dtype
+        The data type of the parameter.
+    shape : Sequence[int | Var]
+        The shape of the parameter.
+    device_type : int
+        The device type of the parameter.
+    device_id : Var
+        The device id of the parameter.
+    strides : Optional[Sequence[Var]], optional
+        The strides of the parameter, by default None.
+    """
+
+    name: str
+    shape: list[Union[int, Var]]
+    dtype: "tvm_ffi.dtype"
+    strides: Optional[list[Var]]
+    dlpack_device_type: int
+    device_id: Var
+    map_tensor_dtype_f4x2_to_f4: bool
+    data_alignment: Optional[int]
+
+    def __init__(
+        self,
+        name: str,
+        shape: Sequence[Union[int, Var]],
+        dtype: Union[str, "tvm_ffi.dtype"],
+        *,
+        device_type: Optional[str] = None,
+        strides: Optional[Sequence[Var]] = None,
+        map_tensor_dtype_f4x2_to_f4: bool = False,
+        data_alignment: Optional[int] = None,
+    ) -> None:
+        """Initialize a Tensor parameter.
+
+        Parameters
+        ----------
+        name : str
+            The parameter name.
+        device_type : str
+            The device type of the parameter.
+        shape : Sequence[int | Var]
+            The shape of the parameter.
+        dtype : str | tvm_ffi.dtype
+            The data type of the parameter.
+        strides : Optional[Sequence[Var]], optional
+            The strides of the parameter, by default None.
+        map_tensor_dtype_f4x2_to_f4: bool
+            Whether to map tensor dtype float4x2 to float4 for internal use.
+        data_alignment: Optional[int], optional
+            The data alignment of the parameter, by default None.
+        """
+        self.name = name
+        self.data = Var(name + ".data", tvm_ffi.dtype("handle"))
+        self.shape: list[Union[int, Var]] = list(shape)
+        self.dtype = tvm_ffi.dtype(dtype)
+        self.strides: Optional[list[Var]] = list(strides) if strides is not None else None
+        self.data_alignment = data_alignment
+
+        # Use default device type if none specified
+        if device_type is None:
+            device_type = DefaultConfig.current().device_type  # type: ignore[union-attr]
+
+        example_device = tvm_ffi.device(device_type, 0)
+        self.dlpack_device_type = example_device.dlpack_device_type()
+        self.device_type_name = example_device.type
+        self.device_id = Var(name + ".device_id", tvm_ffi.dtype("int32"))
+        self.map_tensor_dtype_f4x2_to_f4 = map_tensor_dtype_f4x2_to_f4
+
+
+class Stream(Param):
+    """Stream parameter."""
+
+    name: str
+    var: Var
+
+    def __init__(self, name: str) -> None:
+        """Initialize a Stream parameter.
+
+        Parameters
+        ----------
+        name : str
+            The parameter name.
+        """
+        self.name = name
+        self.var = Var(name, tvm_ffi.dtype("handle"))
+
+
+class EnvStream(Param):
+    """EnvStream parameter.
+
+    Note
+    ----
+    This parameter spec indicates that we expect
+    to call the function with TVMFFIEnvGetStream to get the stream
+    and it is not part of the FFI function signature.
+    """
+
+    name: str
+    var: Var
+
+    def __init__(self, name: str) -> None:
+        """Initialize a EnvStream parameter.
+
+        Parameters
+        ----------
+        name : str
+            The parameter name.
+        """
+        self.name = name
+        self.var = Var(name, tvm_ffi.dtype("handle"))
+
+
+class DataPointer(Param):
+    """Data pointer parameter.
+
+    The main difference between DataPointer and Var with handle type
+    is that DataPointer can contain address space information.
+    We also allow passing in int as data pointer, so it can
+    conveniently support torch.Tensor.data_ptr() as input.
+
+    Parameters
+    ----------
+    name : str
+        The parameter name.
+    """
+
+    name: str
+    var: Var
+    address_space: Optional[int]
+
+    def __init__(self, name: str, address_space: Optional[int] = None) -> None:
+        """Initialize a DataPointer parameter.
+
+        Parameters
+        ----------
+        name : str
+            The parameter name.
+        """
+        self.name = name
+        self.var = Var(name, tvm_ffi.dtype("handle"))
+        self.address_space = address_space
+
+
+def signature(name: str, params: list[Param]) -> str:
+    """Generate a function signature string from name and parameters.
+
+    Parameters
+    ----------
+    name : str
+        The function name.
+    params : list[Param]
+        List of parameter objects (Var or Tensor).
+
+    Returns
+    -------
+    str
+        The formatted function signature.
+
+    Raises
+    ------
+    ValueError
+        If an unknown parameter type is encountered.
+
+    """
+    param_strs = []
+
+    for param in params:
+        if isinstance(param, Var):
+            param_str = f"{param.name}: {param.dtype}"
+        elif isinstance(param, Tensor):
+            # Format tensor shape
+            shape_strs = []
+            for dim in param.shape:
+                if isinstance(dim, Var):
+                    shape_strs.append(dim.name)
+                else:
+                    shape_strs.append(str(dim))
+            shape_str = "[" + ", ".join(shape_strs) + "]"
+            param_str = f"{param.name}: Tensor({shape_str}, {param.dtype})"
+        elif isinstance(param, Shape):
+            # Format shape parameter
+            shape_strs = []
+            for dim in param.shape:
+                if isinstance(dim, Var):
+                    shape_strs.append(dim.name)
+                else:
+                    shape_strs.append(str(dim))
+            shape_str = "[" + ", ".join(shape_strs) + "]"
+            param_str = f"{param.name}: Shape({shape_str})"
+        elif isinstance(param, Stream):
+            param_str = f"{param.name}: Stream"
+        elif isinstance(param, DataPointer):
+            param_str = f"{param.name}: DataPointer"
+        elif isinstance(param, EnvStream):
+            # env stream is not part of the FFI function signature
+            # continue to skip append
+            continue
+        else:
+            raise TypeError(f"Unsupported parameter type: {type(param)}")
+
+        param_strs.append(param_str)
+
+    return f"{name}({', '.join(param_strs)})"
+
+
+def create_map_tensor_dtype_f4x2_to_f4_spec(f4_tensor_spec: Tensor) -> Tensor:
+    """
+    Create a new Tensor spec that can be translated to f4 via f4x2->f4 conversion.
+
+    Parameters
+    ----------
+    f4_tensor_spec : Tensor
+        The original Tensor spec that declares a f4 tensor.
+
+    Returns
+    -------
+    Tensor
+        The new Tensor spec that can be translated to f4 via f4x2->f4 conversion.
+    """
+    if f4_tensor_spec.dtype != tvm_ffi.dtype("float4_e2m1fn"):
+        raise ValueError("f4_tensor_spec must be a float4 tensor")
+
+    def find_stride_one_index() -> int:
+        if f4_tensor_spec.strides is None:
+            return len(f4_tensor_spec.shape) - 1
+        for i, stride in enumerate(f4_tensor_spec.strides):
+            if isinstance(stride, int) and stride == 1:
+                return i
+        raise ValueError("Cannot find dimension with stride=1")
+    stride_one_index = find_stride_one_index()
+
+    def divisibility_divide_by_2(value: Var) -> Optional[int]:
+        if value.divisibility is not None:
+            if value.divisibility % 2 != 0:
+                raise ValueError(f"Dimension with stride=1 must be divisible by 2")
+            return value.divisibility // 2
+        return None
+
+    def map_shape(index: int, value: Union[int, Var]) -> Union[int, Var]:
+        if index == stride_one_index:
+            if isinstance(value, int):
+                if value % 2 != 0:
+                    raise ValueError(f"Dimension {index} with stride=1 must be even")
+                return value // 2
+            # create a new var with the same name and dtype
+            return Var(value.name, value.dtype, divisibility=divisibility_divide_by_2(value))
+        return value
+
+    def map_stride(index: int, value: Union[int, Var]) -> Union[int, Var]:
+        if index != stride_one_index:
+            if isinstance(value, int):
+                if value % 2 != 0:
+                    raise ValueError(f"Dimension {index} with stride != 1 must be even")
+                return value // 2
+            # create a new var with the same name and dtype
+            return Var(value.name, value.dtype, divisibility=divisibility_divide_by_2(value))
+        return value
+
+    new_shape = [map_shape(i, x) for i, x in enumerate(f4_tensor_spec.shape)]
+    if f4_tensor_spec.strides is not None:
+        new_strides = [map_stride(i, x) for i, x in enumerate(f4_tensor_spec.strides)]
+    else:
+        new_strides = None
+
+    return Tensor(
+        f4_tensor_spec.name,
+        new_shape,
+        dtype=tvm_ffi.dtype("float4_e2m1fnx2"),
+        strides=new_strides,
+        map_tensor_dtype_f4x2_to_f4=True,
+    )

python/CuTeDSL/cutlass/base_dsl/typing.py

@@ -239,6 +239,19 @@ def get_c_pointers(obj):
     return []
 
 
+def arg_compatible_with_tvm_ffi(arg):
+    """
+    Given the `arg`, check if it is a compatible argument for TVM FFI
+    """
+    import tvm_ffi
+
+    return (
+        hasattr(arg, "__tvm_ffi_object__")
+        or isinstance(arg, (int, float, bool))
+        or isinstance(arg, tvm_ffi.Shape)
+    )
+
+
 def get_mlir_types(obj):
     """
     Given the `obj`, recursively go through it to extract all contained MLIR types
@@ -1416,6 +1429,9 @@ class Integer(Numeric, metaclass=IntegerMeta, mlir_type=T.i32, is_abstract=True)
     def __rxor__(self, other, *, loc=None, ip=None):
         return self.__xor__(other, loc=loc, ip=ip)
 
+    def __tvm_ffi_int__(self):
+        return self.value
+
 
 class Float(Numeric, metaclass=FloatMeta, mlir_type=T.f32, is_abstract=True):
     """A class representing floating-point values.
@@ -1496,6 +1512,9 @@ class Float(Numeric, metaclass=FloatMeta, mlir_type=T.f32, is_abstract=True):
         else:
             raise DSLRuntimeError(f"{x} to Float conversion is not supported")
 
+    def __tvm_ffi_float__(self):
+        return self.value
+
 
 class Boolean(Integer, metaclass=IntegerMeta, width=1, signed=True, mlir_type=T.bool):
     """Boolean type representation in the DSL.

python/CuTeDSL/cutlass/cute/__init__.py

@@ -10,7 +10,7 @@
 # is strictly prohibited.
 
 # Use the auto-generated enum AddressSpace
-from cutlass._mlir.dialects.cute import AddressSpace
+from cutlass._mlir.dialects.cute import AddressSpace, CacheEvictionPriority
 
 # Explicitly import types that might be directly used by other modules.
 # This is a fix for using Sphinx to generate documentation
@@ -29,6 +29,7 @@ from .typing import (
     ComposedLayout,
     Pointer,
     Tensor,
+    SymInt,
 )
 
 # Import everything else
@@ -114,6 +115,9 @@ from .core import (
     ScaledBasis,
     get_divisibility,
     Ratio,
+    # FastDivmod operations
+    FastDivmodDivisor,
+    fast_divmod_create_divisor,
 )
 
 from .tuple import (
@@ -175,6 +179,8 @@ from .atom import (
 from .algorithm import gemm, copy, basic_copy, basic_copy_if, autovec_copy, prefetch
 
 from . import arch
+
+from . import export
 from . import nvgpu
 from . import testing
 from . import runtime
@@ -200,10 +206,16 @@ KeepPTX = _dsl.KeepPTX
 GPUArch = _dsl.GPUArch
 LinkLibraries = _dsl.LinkLibraries
 
+# attach the TVM FFI ABI interface postprocessor to the DSL
+from . import _tvm_ffi_args_spec_converter
+
+_tvm_ffi_args_spec_converter.attach_args_spec_converter()
+
 # Explicitly export all symbols for documentation generation
 __all__ = [
     # Core types
     "AddressSpace",
+    "CacheEvictionPriority",
     "Tensor",
     "Layout",
     "ComposedLayout",
@@ -219,6 +231,7 @@ __all__ = [
     "ThrCopy",
     "TensorSSA",
     "ReductionOp",
+    "SymInt",
     # Basic utility functions
     "assume",
     "is_integer",
@@ -344,8 +357,12 @@ __all__ = [
     "repeat_like",
     # User defined struct
     "struct",
+    # FastDivmod operations
+    "FastDivmodDivisor",
+    "fast_divmod_create_divisor",
     # Modules
     "arch",
+    "export",
     "nvgpu",
     "testing",
     "runtime",

python/CuTeDSL/cutlass/cute/_tvm_ffi_args_spec_converter.py

@@ -0,0 +1,225 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 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/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.
+
+from cutlass.base_dsl.tvm_ffi_builder import spec
+from cutlass.base_dsl.jit_executor import ExecutionArgs
+from cutlass.base_dsl.common import DSLRuntimeError
+from cutlass.cutlass_dsl import is_cute_algebra_type
+from cutlass._mlir.dialects import cute as _cute_ir
+from .runtime import _FakeStream
+from .typing import Tensor, Pointer, SymInt
+from .typing import (
+    Numeric,
+    Boolean,
+    Int4,
+    Int8,
+    Uint8,
+    Int16,
+    Uint16,
+    Int32,
+    Uint32,
+    Int64,
+    Uint64,
+    Float16,
+    BFloat16,
+    Float32,
+    TFloat32,
+    Float64,
+    Float6E2M3FN,
+    Float6E3M2FN,
+    Float8E5M2,
+    Float8E4M3FN,
+    Float8E8M0FNU,
+    Float4E2M1FN,
+)
+import cuda.bindings.driver as cuda
+
+from typing import List, Dict, Any, Optional
+import inspect
+
+NumericToTVMFFIDtype = {
+    Boolean: "bool",
+    Int4: "int4",
+    Int8: "int8",
+    Uint8: "uint8",
+    Int16: "int16",
+    Uint16: "uint16",
+    Int32: "int32",
+    Uint32: "uint32",
+    Int64: "int64",
+    Uint64: "uint64",
+    Float16: "float16",
+    BFloat16: "bfloat16",
+    Float32: "float32",
+    TFloat32: "float32",
+    Float64: "float64",
+    Float8E5M2: "float8_e5m2",
+    Float8E4M3FN: "float8_e4m3fn",
+    Float8E8M0FNU: "float8_e8m0fnu",
+    Float4E2M1FN: "float4_e2m1fn",
+    Float6E2M3FN: "float6_e2m3fn",
+    Float6E3M2FN: "float6_e3m2fn",
+}
+
+AcceptableNumericTypesForScalar = [
+    Boolean,
+    Int8,
+    Uint8,
+    Int16,
+    Uint16,
+    Int32,
+    Uint32,
+    Int64,
+    Uint64,
+    Float32,
+    Float64,
+]
+
+
+def _get_llvm_address_space_from_memspace(
+    memspace: _cute_ir.AddressSpace,
+) -> Optional[int]:
+    if memspace == _cute_ir.AddressSpace.gmem:
+        return 1
+    return None
+
+
+class SymIntId:
+    def __init__(self, sym_int: SymInt):
+        self.sym_int = sym_int
+
+    def __hash__(self):
+        return id(self.sym_int)
+
+    def __eq__(self, other) -> bool:
+        return self.sym_int is other.sym_int
+
+
+def _tvm_ffi_args_spec_converter(
+    function_name: str,
+    args_spec: inspect.FullArgSpec,
+    dynamic_args: List[Any],
+    dynamic_kwargs: Dict[str, Any],
+):
+    """Convert cute algebra args to tvm ffi spec params.
+
+    This function converts the cute arguments specs to tvm ffi spec params.
+    """
+    exec_args = ExecutionArgs(args_spec, function_name)
+    rectified_args = exec_args.get_rectified_args(dynamic_args, dynamic_kwargs)
+    arg_names = exec_args.args_spec.args + exec_args.args_spec.kwonlyargs
+
+    params = []
+    num_dyn_shape_vars = 0
+    num_dyn_stride_vars = 0
+    sym_int_id_mapping = {}
+
+
+    def alloc_shape_name():
+        nonlocal num_dyn_shape_vars
+        name = f"n{num_dyn_shape_vars}"
+        num_dyn_shape_vars += 1
+        return name
+
+    def alloc_stride_name():
+        nonlocal num_dyn_stride_vars
+        name = f"s{num_dyn_stride_vars}"
+        num_dyn_stride_vars += 1
+        return name
+
+    def alloc_or_reuse_symint_var(value, name_alloc_func):
+        nonlocal sym_int_id_mapping
+        sym_int_id = SymIntId(value)
+        if sym_int_id in sym_int_id_mapping:
+            return sym_int_id_mapping[sym_int_id]
+        name = name_alloc_func()
+        if value.width == 32:
+            dtype = NumericToTVMFFIDtype[Int32]
+        else:
+            dtype = NumericToTVMFFIDtype[Int64]
+        var = spec.Var(name, dtype, divisibility=value.divisibility)
+        sym_int_id_mapping[sym_int_id] = var
+        return var
+
+    for arg, arg_name in zip(rectified_args, arg_names):
+        arg_type = args_spec.annotations.get(arg_name, None)
+        if isinstance(arg, Numeric) and arg.dtype in AcceptableNumericTypesForScalar:
+            params.append(spec.Var(arg_name, NumericToTVMFFIDtype[arg.dtype]))
+        elif is_cute_algebra_type(arg_type):
+            shape = []
+            for i in range(len(arg)):
+                if isinstance(arg[i], int):
+                    shape.append(arg[i])
+                elif isinstance(arg[i], SymInt):
+                    shape.append(alloc_or_reuse_symint_var(arg[i], alloc_shape_name))
+                else:
+                    shape.append(spec.Var(alloc_shape_name(), NumericToTVMFFIDtype[arg[i].dtype]))
+            params.append(spec.Shape(arg_name, shape))
+        elif isinstance(arg, Tensor):
+            shapes = []
+            for i, dyn_mask in enumerate(arg.dynamic_shapes_mask):
+                if not dyn_mask:
+                    shapes.append(arg.shape[i])
+                elif isinstance(arg.shape[i], SymInt):
+                    shapes.append(alloc_or_reuse_symint_var(arg.shape[i], alloc_shape_name))
+                else:
+                    shapes.append(spec.Var(alloc_shape_name(), NumericToTVMFFIDtype[Int32]))
+            strides = []
+
+            for i, dyn_mask in enumerate(arg.dynamic_strides_mask):
+                if not dyn_mask:
+                    strides.append(arg.stride[i])
+                elif isinstance(arg.stride[i], SymInt):
+                    strides.append(alloc_or_reuse_symint_var(arg.stride[i], alloc_stride_name))
+                else:
+                    if hasattr(arg, "_use_32bit_stride") and arg._use_32bit_stride:
+                        dtype = NumericToTVMFFIDtype[Int32]
+                    else:
+                        dtype = NumericToTVMFFIDtype[Int64]
+                    strides.append(spec.Var(alloc_stride_name(), dtype))
+
+            tvm_ffi_cute_tensor = spec.Tensor(
+                arg_name,
+                shapes,
+                NumericToTVMFFIDtype[arg.element_type],
+                strides=strides,
+                data_alignment=arg._assumed_align,
+            )
+            if arg.element_type == Float4E2M1FN:
+                tvm_ffi_cute_tensor = spec.create_map_tensor_dtype_f4x2_to_f4_spec(
+                    tvm_ffi_cute_tensor
+                )
+            params.append(tvm_ffi_cute_tensor)
+        elif isinstance(arg, Pointer):
+            address_space = None
+            if hasattr(arg, "memspace"):
+                address_space = _get_llvm_address_space_from_memspace(arg.memspace)
+            params.append(spec.DataPointer(arg_name, address_space=address_space))
+        elif isinstance(arg, _FakeStream):
+            if arg.use_tvm_ffi_env_stream:
+                params.append(spec.EnvStream(arg_name))
+            else:
+                params.append(spec.Stream(arg_name))
+        elif isinstance(arg, cuda.CUstream):
+            params.append(spec.Stream(arg_name))
+        else:
+            raise DSLRuntimeError(f"Unsupported argument type: {type(arg)}")
+    # The following code can obtain signature of the function
+    # that maybe useful for future debugging and usecases.
+    # signature = spec.signature(function_name, params)
+    return params
+
+
+def attach_args_spec_converter():
+    """Attach TVM FFI ABI interface postprocessor to the DSL."""
+    from .. import cutlass_dsl as _dsl
+
+    _dsl.CuTeDSL._get_dsl()._tvm_ffi_args_spec_converter = _tvm_ffi_args_spec_converter

python/CuTeDSL/cutlass/cute/algorithm.py

@@ -221,8 +221,8 @@ def autovec_copy(src: Tensor, dst: Tensor, *, loc=None, ip=None) -> None:
     upper_bound = math.gcd(upper_bound, src.iterator.max_alignment * 8)
     upper_bound = math.gcd(upper_bound, dst.iterator.max_alignment * 8)
 
-    # Finally, we put a cap at 128b
-    num_bits_per_copy = math.gcd(upper_bound, 128)
+    # Finally, we put a cap at 256b
+    num_bits_per_copy = math.gcd(upper_bound, 256)
 
     if (num_common_elements > 1) and (num_bits_per_copy % 8 == 0):
         num_common_elements = num_bits_per_copy // src.element_type.width
@@ -355,7 +355,7 @@ def copy(
 
     .. code-block:: python
 
-        cute.copy(tma_atom, src, dst, tma_bar_ptr=mbar_ptr, mcast_mask=mask)
+        cute.copy(tma_atom, src, dst, tma_bar_ptr=mbar_ptr, mcast_mask=mask, cache_policy=policy)
 
     Optional predication is supported through an additional tensor parameter. For partitioned tensors with
     logical profile ``((ATOM_V,ATOM_REST),REST,...)``, the predication tensor must maintain profile

python/CuTeDSL/cutlass/cute/arch/numeric_conversion.py

@@ -10,7 +10,9 @@
 # is strictly prohibited.
 
 
-from cutlass.cutlass_dsl import dsl_user_op
+from cutlass.base_dsl.arch import Arch
+from cutlass.base_dsl.common import DSLRuntimeError
+from cutlass.cutlass_dsl import CuTeDSL, dsl_user_op
 from cutlass._mlir import ir
 from cutlass._mlir.dialects import builtin, arith, llvm, vector
 
@@ -34,15 +36,15 @@ from ..typing import (
     Int8,
     Int32,
     Float16,
+    Float32,
     BFloat16,
     Float32,
 )
 
-
 @dsl_user_op
 def cvt_i8_bf16_intrinsic(vec_i8, length, *, loc=None, ip=None):
     """
-    Convert a vector of int8 to a vector of bfloat16.
+    Fast conversion from int8 to bfloat16. It converts a vector of int8 to a vector of bfloat16.
 
     :param vec_i8: The input vector of int8.
     :type vec_i8: 1D vector of int8
@@ -51,6 +53,9 @@ def cvt_i8_bf16_intrinsic(vec_i8, length, *, loc=None, ip=None):
     :return: The output 1D vector of bfloat16 with the same length as the input vector.
     :rtype: 1D vector of bfloat16
     """
+    arch = CuTeDSL._get_dsl().get_arch_enum()
+    if not arch in cvt_i8_bf16_intrinsic.supported_archs:
+        raise DSLRuntimeError(f"cvt_i8_bf16_intrinsic is not supported on {arch}")
     src_pos = 0
     vec_i8x4_type = ir.VectorType.get([4], Int8.mlir_type, loc=loc)
     vec_i8x2_type = ir.VectorType.get([2], Int8.mlir_type, loc=loc)
@@ -116,7 +121,7 @@ def cvt_i8_bf16_intrinsic(vec_i8, length, *, loc=None, ip=None):
 @dsl_user_op
 def cvt_i4_bf16_intrinsic(vec_i4, length, *, loc=None, ip=None):
     """
-    Convert a vector of int4 to a vector of bfloat16.
+    Fast conversion from int4 to bfloat16. It converts a vector of int4 to a vector of bfloat16.
 
     :param vec_i4: The input vector of int4.
     :type vec_i4: 1D vector of int4
@@ -125,6 +130,9 @@ def cvt_i4_bf16_intrinsic(vec_i4, length, *, loc=None, ip=None):
     :return: The output 1D vector of bfloat16 with the same length as the input vector.
     :rtype: 1D vector of bfloat16
     """
+    arch = CuTeDSL._get_dsl().get_arch_enum()
+    if not arch in cvt_i4_bf16_intrinsic.supported_archs:
+        raise DSLRuntimeError(f"cvt_i4_bf16_intrinsic is not supported on {arch}")
     src_pos = 0
     vec_i4x8_type = ir.VectorType.get([8], Int4.mlir_type, loc=loc)
     vec_i4x4_type = ir.VectorType.get([4], Int4.mlir_type, loc=loc)
@@ -261,3 +269,18 @@ def cvt_f4e2m1_f16_intrinsic(vec_f4e2m1, length, *, loc=None, ip=None):
             ip=ip,
         )
     return vec_dst
+
+
+# Expose supported architectures via the intrinsic symbol
+cvt_i8_bf16_intrinsic.supported_archs = (
+    *Arch.AmpereArchs(),
+    *Arch.AdaArchs(),
+    *Arch.HopperArchs(),
+    *Arch.BlackwellArchs(),
+)
+cvt_i4_bf16_intrinsic.supported_archs = (
+    Arch.sm_100a,
+    Arch.sm_110a,
+    Arch.sm_120a,
+    Arch.sm_121a,
+)

python/CuTeDSL/cutlass/cute/arch/nvvm_wrappers.py

@@ -10,14 +10,19 @@
 # is strictly prohibited.
 
 from functools import partial
-from typing import Optional, Tuple, Union, Callable
+from typing import Optional, Tuple, Union, Callable, TYPE_CHECKING
 from typing_extensions import deprecated
 
-from cutlass.cutlass_dsl import T, dsl_user_op
+from cutlass.cutlass_dsl import T, dsl_user_op, cutlass_arith
+
+import cutlass.cutlass_dsl as cutlass_dsl
 
 from cutlass._mlir import ir
 from cutlass._mlir.dialects import arith, llvm, nvvm, vector
 
+if TYPE_CHECKING:
+    from cutlass.tensor import TensorSSA
+
 # Forward nvvm enums
 from cutlass._mlir.dialects.nvvm import (
     ProxyKind,
@@ -27,6 +32,8 @@ from cutlass._mlir.dialects.nvvm import (
     RoundingModeKind,
 )
 
+from ..core import size
+
 from ..typing import (
     Int,
     Boolean,
@@ -184,7 +191,7 @@ def block_idx_in_cluster(*, loc=None, ip=None) -> Int32:
 
 @dsl_user_op
 def shuffle_sync_op(
-    value: Numeric,
+    value: Union[Numeric, "TensorSSA"],
     offset: Int,
     mask: Int = FULL_MASK,
     mask_and_clamp: Int = WARP_SIZE - 1,
@@ -192,12 +199,12 @@ def shuffle_sync_op(
     *,
     loc=None,
     ip=None,
-) -> Numeric:
+) -> Union[Numeric, "TensorSSA"]:
     """
     Shuffles a value within the threads of a warp.
 
     :param value:          The value to shuffle
-    :type value:           Numeric
+    :type value:           Numeric or TensorSSA
     :param mask:           A mask describing the threads participating in this operation
     :type mask:            Int
     :param offset:         A source lane or a source lane offset depending on kind
@@ -211,12 +218,37 @@ def shuffle_sync_op(
     :return:               The shuffled value
     :rtype:                Numeric
     """
+    from ..tensor import TensorSSA
+
+    if isinstance(value, TensorSSA):
+        bit_width = value.dtype.width * size(value.shape)
+        if bit_width == 32:
+            i32_val = llvm.bitcast(
+                T.i32(), value.ir_value(loc=loc, ip=ip), loc=loc, ip=ip
+            )
+            i32_res = nvvm.shfl_sync(
+                T.i32(),
+                Int32(mask).ir_value(loc=loc, ip=ip),
+                i32_val,
+                Int32(offset).ir_value(loc=loc, ip=ip),
+                Int32(mask_and_clamp).ir_value(loc=loc, ip=ip),
+                kind,
+                loc=loc,
+                ip=ip,
+            )
+            result_vec = llvm.bitcast(value.type, i32_res, loc=loc, ip=ip)
+            return TensorSSA(result_vec, value.shape, value.dtype)
+        else:
+            raise ValueError(f"shuffle_sync only supports 32 bit, but got {value.type}")
+
     if not isinstance(value, Numeric):
         value = as_numeric(value)
+
     if value.width > 64:
         raise ValueError("shuffle_sync only supports values up to 64 bits")
 
     orig_type = type(value)
+
     if value.width < 32:
         if value.dtype.is_float:
             value = value.to(Float32)
@@ -331,6 +363,7 @@ def warp_reduction(
     :rtype: cutlass.Numeric
     """
     offset = threads_in_group // 2
+
     while offset > 0:
         val = op(
             val,
@@ -343,7 +376,8 @@ def warp_reduction(
 
 
 warp_reduction_max = partial(
-    warp_reduction, op=lambda x, y: fmax(x, y) if isinstance(x, Float32) else max(x, y)
+    warp_reduction,
+    op=lambda x, y: fmax(x, y) if isinstance(x, Float32) else cutlass_dsl.max(x, y),
 )
 warp_reduction_sum = partial(warp_reduction, op=lambda x, y: x + y)
 
@@ -796,11 +830,11 @@ add_packed_f32x2 = partial(
     calc_packed_f32x2_op, src_c=None, calc_func=nvvm.add_packed_f32x2
 )
 
-
 @dsl_user_op
 def fmax(
     a: Union[float, Float32], b: Union[float, Float32], *, loc=None, ip=None
 ) -> Float32:
+
     return Float32(
         nvvm.fmax(
             T.f32(),
@@ -811,7 +845,6 @@ def fmax(
         )
     )
 
-
 @dsl_user_op
 def rcp_approx(a: Union[float, Float32], *, loc=None, ip=None):
     return Float32(

python/CuTeDSL/cutlass/cute/atom.py

@@ -99,6 +99,9 @@ class Trait(ABC):
     def unpack(self, *, loc=None, ip=None, **kwargs) -> ir.Value:
         return self.value
 
+    def with_(self, *, loc=None, ip=None, **kwargs) -> "Trait":
+        return self.__class__(self.unpack(loc=loc, ip=ip, **kwargs))
+
 
 def make_atom(ty, values=None, *, loc=None, ip=None):
     """
@@ -185,6 +188,21 @@ class Atom(ABC):
         """
         return self._trait.get(field, loc=loc, ip=ip)
 
+    def with_(self, *, loc=None, ip=None, **kwargs) -> "Atom":
+        """
+        Returns a new Atom with the new Operation and Trait with the given runtime state. The runtime state
+        is provided as keyword arguments and it is Atom-specific.
+
+        .. code-block:: python
+
+            tiled_copy = cute.make_tiled_copy(tma_copy_op)
+            new_tiled_copy = tiled_copy.with_(tma_bar_ptr=tma_bar_ptr, cache_policy=cute.CacheEvictionPriority.EVICT_LAST)
+
+        The ``with_`` method provides a way to the user to modify such runtime state or create an executable Atom
+        (e.g. an Executable TMA Load Atom).
+        """
+        return self.__class__(self.op, self._trait.with_(loc=loc, ip=ip, **kwargs))
+
     def _unpack(self, *, loc=None, ip=None, **kwargs) -> ir.Value:
         return self._trait.unpack(loc=loc, ip=ip, **kwargs)
 
@@ -215,23 +233,28 @@ class MmaAtom(Atom):
     #
 
     @property
-    def thr_id(self) -> Layout:
-        return static(self._trait.value.type.thr_id)
+    @dsl_user_op
+    def thr_id(self, *, loc=None, ip=None) -> Layout:
+        return static(self._trait.value.type.thr_id, loc=loc, ip=ip)
 
     @property
-    def shape_mnk(self) -> Shape:
-        return _unpack_x_tuple(self._trait.value.type.shape_mnk)
+    @dsl_user_op
+    def shape_mnk(self, *, loc=None, ip=None) -> Shape:
+        return _unpack_x_tuple(self._trait.value.type.shape_mnk, loc=loc, ip=ip)
 
     @property
-    def tv_layout_A(self) -> Layout:
-        return static(self._trait.value.type.layout_a_tv)
+    @dsl_user_op
+    def tv_layout_A(self, *, loc=None, ip=None) -> Layout:
+        return static(self._trait.value.type.layout_a_tv, loc=loc, ip=ip)
 
     @property
-    def tv_layout_B(self) -> Layout:
-        return static(self._trait.value.type.layout_b_tv)
+    @dsl_user_op
+    def tv_layout_B(self, *, loc=None, ip=None) -> Layout:
+        return static(self._trait.value.type.layout_b_tv, loc=loc, ip=ip)
 
     @property
-    def tv_layout_C(self) -> Layout:
+    @dsl_user_op
+    def tv_layout_C(self, *, loc=None, ip=None) -> Layout:
         return static(self._trait.value.type.layout_c_tv)
 
     #

python/CuTeDSL/cutlass/cute/core.py

@@ -10,6 +10,7 @@
 # is strictly prohibited.
 
 from functools import partial, reduce
+import inspect
 from inspect import isclass
 from typing import Any, Dict, List, Optional, Tuple, Type, Union, overload
 
@@ -304,9 +305,10 @@ class IntValue(cutlass_arith.ArithValue):
     def __str__(self):
         if self.divisibility == 1:
             return "?"
-        else:
+        elif self.type.width == 32:
             return f"?{{div={self.divisibility}}}"
-
+        else:
+            return f"?{{i{self.type.width} div={self.divisibility}}}"
     def __repr__(self):
         parent_name = cutlass_arith.ArithValue.__name__
         return super().__str__().replace(parent_name, IntValue.__name__)
@@ -913,8 +915,8 @@ class _Layout(Layout):
 
             layout = make_layout((4, 8), stride=(8, 1))
 
-            # map linear index back to coordinate: 5 -> (1, 1)
-            coord = get_hier_coord(5, layout)
+            # map linear index back to coordinate: 9 -> (1, 1)
+            coord = layout.get_hier_coord(9)
         """
         idx_val = Int32(idx).ir_value(loc=loc, ip=ip)
         crd = _cute_ir.get_hier_coord(idx_val, self, loc=loc, ip=ip)
@@ -3116,12 +3118,6 @@ def make_ptr(
     if isinstance(value, ir.Value) and llvm.PointerType.isinstance(value.type):
         value = llvm.ptrtoint(T.i64(), value)
 
-    if not isinstance(mem_space, AddressSpace):
-        raise TypeError(f"expects mem_space to be an AddressSpace, but got {mem_space}")
-
-    if isinstance(value, ir.Value) and llvm.PointerType.isinstance(value.type):
-        value = llvm.ptrtoint(T.i64(), value)
-
     if not is_integer(value):
         raise TypeError(f"expects integer value, but got {type(value)}")
     value = Int32(value) if mem_space == AddressSpace.tmem else Int64(value)
@@ -3611,9 +3607,9 @@ def make_layout_image_mask(
     sliced_lay, offset = slice_and_offset(slicer, lay, loc=loc, ip=ip)
     # Given that we replace only one mode with _, the rank of the slice should be 1
     assert rank(sliced_lay) == 1
-
-    if not is_static(sliced_lay):
-        raise ValueError("make_layout_image_mask requires the layout to be static")
+    assert is_static(
+        sliced_lay
+    ), "make_layout_image_mask requires the layout to be static"
 
     # Create the mask of the image
     mcast_mask = Int16(0)
@@ -4125,3 +4121,166 @@ class ThrMma(_ThrMma):
 @deprecated("cute.core.ThrCopy is deprecated, use cute.ThrCopy instead")
 class ThrCopy(_ThrCopy):
     pass
+
+
+#
+# FastDivmod operations for optimized division and modulus
+#
+class FastDivmodDivisor:
+    """
+    First-class FastDivmod divisor with operator overloading support.
+
+    This class wraps a FastDivmod divisor and enables natural Python operator syntax:
+        quotient, remainder = divmod(dividend, divisor)
+        quotient = dividend // divisor
+        remainder = dividend % divisor
+
+    :ivar _divisor: The FastDivmod divisor MLIR value
+    """
+
+    @dsl_user_op
+    def __init__(
+        self,
+        divisor: Integer,
+        is_power_of_2: bool = None,
+        *,
+        loc=None,
+        ip=None,
+    ):
+        """
+        Create a FastDivmod divisor for optimized division operations.
+
+        :param divisor: The divisor value (should be runtime-dynamic value)
+        :param is_power_of_2: Whether divisor is known to be a power of 2.
+                              Defaults to False.
+        """
+        # Convert divisor to ir.Value for MLIR operation
+        if isinstance(divisor, ir.Value):
+            divisor_val = divisor
+        else:
+            divisor_val = Int32(divisor).ir_value()
+
+        # Use user-provided flag or default to False
+        # Power-of-2 optimization should be handled by compiler passes at IR level
+        if is_power_of_2 is None:
+            is_power_of_2 = False
+
+        # Create FastDivmod divisor
+        fast_divmod_divisor_type = _cute_ir.FastDivmodDivisorType.get(32, is_power_of_2)
+        self._divisor = _cute_ir.fast_divmod_create_divisor(
+            fast_divmod_divisor_type, divisor_val, loc=loc, ip=ip
+        )
+
+    @dsl_user_op
+    def __rdivmod__(
+        self, dividend: Integer, *, loc=None, ip=None
+    ) -> Tuple[Integer, Integer]:
+        """
+        Overload for: divmod(dividend, self)
+        Returns (quotient, remainder).
+
+        :param dividend: The dividend value
+        :param loc: Source location for MLIR
+        :param ip: Insertion point for MLIR
+        :return: Tuple of (quotient, remainder)
+        """
+        # Convert dividend to ir.Value for MLIR operation
+        if isinstance(dividend, ir.Value):
+            dividend_val = dividend
+        else:
+            dividend_val = Int32(dividend).ir_value()
+
+        quotient_type = dividend_val.type
+        remainder_type = dividend_val.type
+
+        results = _cute_ir.fast_divmod_compute(
+            quotient_type,
+            remainder_type,
+            dividend_val,
+            self._divisor,
+            loc=loc,
+            ip=ip,
+        )
+        return (IntValue(results[0]), IntValue(results[1]))
+
+    @dsl_user_op
+    def __rfloordiv__(self, dividend: Integer, *, loc=None, ip=None) -> Integer:
+        """
+        Overload for: dividend // self
+        Returns quotient only.
+
+        :param dividend: The dividend value
+        :param loc: Source location for MLIR
+        :param ip: Insertion point for MLIR
+        :return: The quotient
+        """
+        quotient, _ = self.__rdivmod__(dividend, loc=loc, ip=ip)
+        return quotient
+
+    @dsl_user_op
+    def __rmod__(self, dividend: Integer, *, loc=None, ip=None) -> Integer:
+        """
+        Overload for: dividend % self
+        Returns remainder only.
+
+        :param dividend: The dividend value
+        :param loc: Source location for MLIR
+        :param ip: Insertion point for MLIR
+        :return: The remainder
+        """
+        _, remainder = self.__rdivmod__(dividend, loc=loc, ip=ip)
+        return remainder
+
+    def __extract_mlir_values__(self):
+        """Extract MLIR values for Host->Device transfer."""
+        return [self._divisor]
+
+    def __new_from_mlir_values__(self, values):
+        """Reconstruct FastDivmodDivisor from MLIR values."""
+        new_obj = object.__new__(FastDivmodDivisor)
+        new_obj._divisor = values[0]
+        return new_obj
+    def __repr__(self):
+        return f"FastDivmodDivisor({self._divisor.type})"
+
+
+# Set explicit signature for Sphinx documentation to avoid issues with @dsl_user_op decorator
+FastDivmodDivisor.__init__.__signature__ = inspect.Signature(
+    [
+        inspect.Parameter("self", inspect.Parameter.POSITIONAL_OR_KEYWORD),
+        inspect.Parameter(
+            "divisor",
+            inspect.Parameter.POSITIONAL_OR_KEYWORD,
+            annotation=Integer,
+        ),
+        inspect.Parameter(
+            "is_power_of_2",
+            inspect.Parameter.POSITIONAL_OR_KEYWORD,
+            default=None,
+            annotation=bool,
+        ),
+    ]
+)
+
+
+@dsl_user_op
+def fast_divmod_create_divisor(
+    divisor: Integer, *, loc=None, ip=None
+) -> FastDivmodDivisor:
+    """Create a FastDivmod divisor for optimized division operations.
+
+    This function creates a FastDivmod divisor that precomputes auxiliary values
+    to enable fast division and modulus operations without using division instructions.
+
+    The returned FastDivmodDivisor object supports natural Python operator syntax:
+        divisor = fast_divmod_create_divisor(batch_size)
+        quotient, remainder = divmod(linear_idx, divisor)
+        quotient = linear_idx // divisor
+        remainder = linear_idx % divisor
+
+    :param divisor: The divisor value (should be runtime-dynamic value)
+    :type divisor: Integer
+    :return: FastDivmodDivisor object with operator overloading support
+    :rtype: FastDivmodDivisor
+    """
+    return FastDivmodDivisor(divisor, loc=loc, ip=ip)

python/CuTeDSL/cutlass/cute/export/__init__.py

@@ -0,0 +1,38 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 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/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.
+
+from .c_header_generator import CuteCHeaderGenerator
+
+from ...base_dsl.export import (
+    get_export_module,
+    dump_to_object as _dump_to_object,
+    export_to_c as _export_to_c,
+)
+from ...cutlass_dsl import CuTeDSL
+from functools import partial as _partial
+from ...cutlass_dsl.cuda_jit_executor import CudaDialectJitCompiledFunction
+
+dump_to_object = _partial(
+    _dump_to_object,
+    dsl=CuTeDSL._get_dsl(),
+)
+export_to_c = _partial(
+    _export_to_c,
+    dsl=CuTeDSL._get_dsl(),
+    c_header_generator=CuteCHeaderGenerator(),
+    use_gpu_dialect=False,
+)
+__all__ = [
+    "CuteCHeaderGenerator",
+    "get_export_module",
+    "dump_to_object",
+    "export_to_c",
+]

python/CuTeDSL/cutlass/cute/export/c_header_generator.py

@@ -0,0 +1,241 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 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/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.
+
+from cutlass.cute.typing import NumericMeta, Integer
+from cutlass.base_dsl.export import CHeaderGenerator
+from cutlass.base_dsl.dsl import is_dynamic_expression
+from cutlass.base_dsl.common import DSLRuntimeError
+from cutlass.base_dsl.jit_executor import ExecutionArgs
+from cutlass.cutlass_dsl.cutlass import is_cute_algebra_type
+from ..runtime import Tensor, Pointer
+
+from typing import List, Any, Dict
+from inspect import isclass
+import cuda.bindings.driver as cuda
+
+# =============================================================================
+# Cute DSL C Header Generator for c/cpp AOT support
+# =============================================================================
+
+
+class CuteCHeaderGenerator(CHeaderGenerator):
+    """This class provides a Export C Header Generator for cute c/cpp AOT support."""
+
+    def _get_cute_algebra_type(self, arg_type: Any, arg: Any) -> str:
+        """Judge if the dynamic elements of the cute algebra type are same(Int32 or Int64).
+        If so, generate the corresponding C type. Otherwise, refuse to generate the argument
+        in C header and raise an error.
+        """
+        if not isinstance(arg, (list, tuple, Integer)):
+            raise DSLRuntimeError(
+                f"Unsupported argument for c function argument generation: {arg} with type {arg_type}"
+            )
+        if isinstance(arg, Integer):
+            return self.numeric_to_c_type[arg.dtype]
+        dyn_type = None
+        for elem in arg:
+            if not isinstance(elem, (Integer, int)):
+                raise DSLRuntimeError(
+                    f"Unsupported argument for c function argument generation: {arg} with type {arg_type}"
+                )
+            if isinstance(elem, int):
+                continue
+            if dyn_type is not None:
+                if elem.dtype != dyn_type:
+                    raise DSLRuntimeError(
+                        f"Expects all dynamic elements of the cute algebra type to be of the same type, but got {elem.dtype} and {dyn_type}"
+                    )
+            else:
+                dyn_type = elem.dtype
+        if dyn_type is None:
+            return "int32_t "
+        return self.numeric_to_c_type[dyn_type]
+
+    def _generate_binary_declaration(self, symbol_prefix: str):
+        """
+        Generate the binary of the compiled function.
+        """
+        return ""
+
+    def _generate_kernel_metadata(
+        self, symbol_prefix: str, kernel_info: Dict[str, List], dsl_name: str
+    ):
+        """
+        Generate the kernel metadata for the compiled function.
+        """
+        kernel_metadata_struct = f"""
+typedef struct {{
+    CUlibrary module;
+}} {symbol_prefix}_Kernel_Metadata_t;
+"""
+        kernel_metadata_load = f"""
+#ifdef __cplusplus
+extern "C" {{
+#endif
+void _mlir_{symbol_prefix}_cuda_init(void **);
+void _mlir_{symbol_prefix}_cuda_load(void **);
+static inline void {symbol_prefix}_Kernel_Metadata_Load({symbol_prefix}_Kernel_Metadata_t *metadata) {{
+    CUlibrary *libraryPtr = &(metadata->module);
+    int32_t ret;
+    struct {{
+        CUlibrary **libraryPtr;
+        int32_t *ret;
+    }} initArgs = {{&libraryPtr, &ret}};
+    _mlir_{symbol_prefix}_cuda_init((void **)(&initArgs));
+    {dsl_name}_CUDA_ERROR_CHECK((CUresult)(ret));
+    struct {{
+        CUlibrary *library;
+        int32_t *ret;
+    }} loadArgs = {{libraryPtr, &ret}};
+    _mlir_{symbol_prefix}_cuda_load((void **)(&loadArgs));
+    {dsl_name}_CUDA_ERROR_CHECK((CUresult)(ret));
+    CUdevice device;
+    {dsl_name}_CUDA_ERROR_CHECK(cuCtxGetDevice(&device));
+    int max_shared_memory_per_block_optin;
+    {dsl_name}_CUDA_ERROR_CHECK(cuDeviceGetAttribute(&max_shared_memory_per_block_optin, CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN, device));
+    unsigned int num_kernels;
+    {dsl_name}_CUDA_ERROR_CHECK(cuLibraryGetKernelCount(&num_kernels, metadata->module));
+    CUkernel *kernels = (CUkernel *)malloc(num_kernels * sizeof(CUkernel));
+    {dsl_name}_CUDA_ERROR_CHECK(cuLibraryEnumerateKernels(kernels, num_kernels, metadata->module));
+    for (unsigned int i = 0; i < num_kernels; i++) {{
+        {dsl_name}_CUDA_ERROR_CHECK(cuKernelSetAttribute(CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, max_shared_memory_per_block_optin, kernels[i], device));
+    }}
+    free(kernels);
+}}
+"""
+        kernel_metadata_unload = f"""
+static inline void {symbol_prefix}_Kernel_Metadata_Unload({symbol_prefix}_Kernel_Metadata_t *metadata) {{
+    {dsl_name}_CUDA_ERROR_CHECK(cuLibraryUnload(metadata->module));
+}}
+
+#ifdef __cplusplus
+}}
+#endif
+"""
+
+        return kernel_metadata_struct + kernel_metadata_load + kernel_metadata_unload
+
+    def _generate_arguments(
+        self,
+        symbol_prefix: str,
+        args_spec: ExecutionArgs,
+        args: List[Any],
+        kwargs: Dict[str, Any],
+    ):
+        """
+        Generate the arguments of the wrapper function.
+        """
+        arguments = []
+        packed_args = []
+        declarations = []
+        # traverse the runtime args_spec and generate the arguments
+        rectified_args = args_spec.get_rectified_args(args, kwargs)
+        input_arg_names = args_spec.args_spec.args + args_spec.args_spec.kwonlyargs
+        for arg_name, arg in zip(input_arg_names, rectified_args):
+            arg_type = args_spec.args_spec.annotations.get(arg_name, None)
+
+            # process optional argument
+            if arg is None:
+                continue
+
+            if isinstance(arg, Pointer):
+                arguments.append(f"void *{arg_name}")
+                packed_args.append("&" + arg_name)
+            elif isinstance(arg, Tensor):
+                dynamic_shapes = (
+                    f"\n    int32_t dynamic_shapes[{sum(arg.dynamic_shapes_mask)}];"
+                    if sum(arg.dynamic_shapes_mask) > 0
+                    else ""
+                )
+                stride_type = "int32_t" if arg._use_32bit_stride else "int64_t"
+                dynamic_strides = (
+                    f"\n    {stride_type} dynamic_strides[{sum(arg.dynamic_strides_mask)}];"
+                    if sum(arg.dynamic_strides_mask) > 0
+                    else ""
+                )
+                declarations.append(
+                    f"""
+typedef struct {{
+    void *data;{dynamic_shapes}{dynamic_strides}
+}} {symbol_prefix}_Tensor_{arg_name}_t;
+"""
+                )
+                arguments.append(f"{symbol_prefix}_Tensor_{arg_name}_t *{arg_name}")
+                packed_args.append(arg_name)
+            # Generate basic numeric types
+            elif isinstance(arg_type, NumericMeta):
+                arguments.append(self._generate_numeric_argument(arg_name, arg_type))
+                packed_args.append("&" + arg_name)
+            elif is_cute_algebra_type(arg_type):
+                c_type = self._get_cute_algebra_type(arg_type, arg)
+                arguments.append(f"{c_type}*{arg_name}")
+                for i in range(self._count_dynamic_expression(arg)):
+                    packed_args.append("&" + arg_name + "[" + str(i) + "]")
+            elif isclass(arg_type) and issubclass(arg_type, cuda.CUstream):
+                arguments.append("CUstream " + arg_name)
+                packed_args.append("&" + arg_name)
+            else:
+                raise DSLRuntimeError(
+                    f"Unsupported argument for c function argument generation: {arg} with type {arg_type}"
+                )
+
+        return arguments, packed_args, declarations
+
+    def _generate_wrapper_function(
+        self,
+        symbol_prefix: str,
+        args_spec: ExecutionArgs,
+        function_name: str,
+        kernel_info: Dict[str, List],
+        dynamic_args: list,
+        dynamic_kwargs: dict,
+    ):
+        """
+        Generate the wrapper function for the compiled function which is provided to users as the entry point.
+        It uses the `symbol_prefix` as the function name for identification. The host/device symbols are hidden under the bytecode.
+        """
+        # 1. Get the name of the function wrapper
+        wrapper_function_name = f"{symbol_prefix}_wrapper"
+        capi_function_name = f"_mlir_{symbol_prefix}__mlir_ciface_{function_name}"
+        # 2. Generate the signature of the wrapper function
+        arguments, packed_args, declarations = self._generate_arguments(
+            symbol_prefix, args_spec, dynamic_args, dynamic_kwargs
+        )
+        # 3. Get the return type of the wrapper function.
+        # Note that this requires the return type to be properly annotated in python.
+        return_type = args_spec.args_spec.annotations.get("return", None)
+        if return_type is None:
+            return_type = "void"
+        else:
+            return_type = self.numeric_to_c_type[return_type][:-1]
+        declarations = "\n".join(declarations)
+
+        # 4. Generate the wrapper function
+        function = (
+            declarations
+            + f"""
+#ifdef __cplusplus
+extern "C"
+#endif
+void {capi_function_name}(void **args, int32_t num_args);
+
+static inline {return_type} {wrapper_function_name}({symbol_prefix}_Kernel_Metadata_t *metadata, {", ".join(arguments)}) {{
+    {return_type} ret;
+    void *args[{len(packed_args) + 1}] = {{
+        {", ".join(packed_args)},
+        &ret
+    }};
+    {capi_function_name}(args, {len(packed_args) + 1});
+    return ret;
+}}
+"""
+        )
+        return function

python/CuTeDSL/cutlass/cute/nvgpu/cpasync/copy.py

@@ -22,7 +22,7 @@ from cutlass._mlir import ir
 
 from ...atom import CopyOp, Trait
 from ...tensor import ReductionOp
-from ...typing import Int16, Pointer, Integer, Numeric
+from ...typing import Int16, Int64, Pointer, Integer, Numeric
 from ..common import OpError
 from ..tcgen05.mma import CtaGroup
 
@@ -113,6 +113,7 @@ TMA_MBAR_PTR_FIELD_NAME = "tma_bar"
 TMA_MCAST_MASK_FIELD_NAME = "mcast_mask"
 TMA_DESC_PTR_FIELD_NAME = "tma_descriptor_ptr"
 TMA_BYTE_MASK_FIELD_NAME = "byte_mask"
+TMA_CACHE_POLICY_FIELD_NAME = "cache_policy"
 
 
 class TmaCopyOp(CopyOp):
@@ -184,6 +185,10 @@ class CopyBulkTensorTileG2SOp(TmaCopyOp):
 class CopyBulkTensorTileG2SNonExecTrait(Trait):
     # We allow kw args to be dropped so that the user can write common code for non-multicast
     # and multicast loads.
+
+    def with_(self, *, loc=None, ip=None, **kwargs) -> "CopyBulkTensorTileG2STrait":
+        return CopyBulkTensorTileG2STrait(self.unpack(loc=loc, ip=ip, **kwargs))
+
     def unpack(
         self,
         *,
@@ -191,13 +196,15 @@ class CopyBulkTensorTileG2SNonExecTrait(Trait):
         ip=None,
         tma_bar_ptr: Optional[Pointer] = None,
         tma_desc_ptr: Optional[Pointer] = None,
+        cache_policy: Optional[Int64] = None,
         **kwargs,
     ):
         """
         Custom implementation of unpack for non-executable TMAs.
 
         The non-multicast TMA load requires a `tma_bar_ptr` keyword argument to be provided when
-        using `cute.copy`. Any other kw arguments will be ignored instead of triggering an error.
+        using `cute.copy`. `cache_policy` keyword argument to be provided to set the l2 cache eviction priority.
+        Any other kw arguments will be ignored instead of triggering an error.
         """
         if not isinstance(tma_bar_ptr, Pointer):
             raise ValueError(
@@ -215,9 +222,25 @@ class CopyBulkTensorTileG2SNonExecTrait(Trait):
             exec_value = _cute_nvgpu_ir.atom_set_value(
                 exec_value, attr, tma_desc_ptr.value, loc=loc, ip=ip
             )
+        if cache_policy is not None:
+            if not isinstance(cache_policy, Int64):
+                raise ValueError(
+                    "expects `Int64` value to be provided via the cache_policy kw argument"
+                )
+
+            attr_str = (
+                f"#cute_nvgpu.atom_copy_field_tmaload<{TMA_CACHE_POLICY_FIELD_NAME}>"
+            )
+            attr = ir.Attribute.parse(attr_str)
+            exec_value = _cute_nvgpu_ir.atom_set_value(
+                exec_value, attr, cache_policy.value, loc=loc, ip=ip
+            )
         return exec_value
 
 
+class CopyBulkTensorTileG2STrait(Trait):
+    pass
+
 #
 # TMA GMEM -> SMEM multicast copies
 #
@@ -277,6 +300,13 @@ class CopyBulkTensorTileG2SMulticastOp(TmaCopyOp):
 
 
 class CopyBulkTensorTileG2SMulticastNonExecTrait(Trait):
+    def with_(
+        self, *, loc=None, ip=None, **kwargs
+    ) -> "CopyBulkTensorTileG2SMulticastTrait":
+        return CopyBulkTensorTileG2SMulticastTrait(
+            self.unpack(loc=loc, ip=ip, **kwargs)
+        )
+
     def unpack(
         self,
         *,
@@ -285,12 +315,14 @@ class CopyBulkTensorTileG2SMulticastNonExecTrait(Trait):
         tma_bar_ptr: Optional[Pointer] = None,
         mcast_mask=None,
         tma_desc_ptr=None,
+        cache_policy: Optional[Int64] = None,
     ):
         """
         Custom implementation of unpack for non-executable TMAs.
 
         The multicast TMA load requires a `tma_bar_ptr`  and a `mcast_mask` keyword arguments to be
-        provided when using `cute.copy`.
+        provided when using `cute.copy`. `cache_policy` keyword argument to be provided to set the
+        l2 cache eviction priority.
         """
         if not isinstance(tma_bar_ptr, Pointer):
             raise ValueError(
@@ -312,6 +344,19 @@ class CopyBulkTensorTileG2SMulticastNonExecTrait(Trait):
             exec_value = _cute_nvgpu_ir.atom_set_value(
                 exec_value, attr, tma_desc_ptr.value, loc=loc, ip=ip
             )
+        if cache_policy is not None:
+            if not isinstance(cache_policy, Int64):
+                raise ValueError(
+                    "expects `Int64` value to be provided via the cache_policy kw argument"
+                )
+
+            attr_str = (
+                f"#cute_nvgpu.atom_copy_field_tmaload<{TMA_CACHE_POLICY_FIELD_NAME}>"
+            )
+            attr = ir.Attribute.parse(attr_str)
+            exec_value = _cute_nvgpu_ir.atom_set_value(
+                exec_value, attr, cache_policy.value, loc=loc, ip=ip
+            )
         # Set the tma_bar_ptr at last to ensure that the atom creation and setting
         # operations above can be moved outside the loop
         attr_str = "#cute_nvgpu.atom_copy_field_tmaload<tma_bar>"
@@ -321,6 +366,9 @@ class CopyBulkTensorTileG2SMulticastNonExecTrait(Trait):
         )
         return exec_value
 
+class CopyBulkTensorTileG2SMulticastTrait(Trait):
+    pass
+
 
 #
 # TMA SMEM -> GMEM copies
@@ -351,14 +399,24 @@ class CopyBulkTensorTileS2GOp(TmaCopyOp):
 
     def _make_trait(
         self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs
-    ) -> "CopyBulkTensorTileS2GTrait":
+    ) -> "CopyBulkTensorTileS2GNonExecTrait":
         raise NotImplementedError(
             "Use cpasync.make_tiled_tma_atom to obtain a copy Atom for TMA"
         )
 
 
-class CopyBulkTensorTileS2GTrait(Trait):
-    def unpack(self, *, loc=None, ip=None, tma_desc_ptr: Optional[Pointer] = None):
+class CopyBulkTensorTileS2GNonExecTrait(Trait):
+    def with_(self, *, loc=None, ip=None, **kwargs) -> "CopyBulkTensorTileS2GTrait":
+        return CopyBulkTensorTileS2GTrait(self.unpack(loc=loc, ip=ip, **kwargs))
+
+    def unpack(
+        self,
+        *,
+        loc=None,
+        ip=None,
+        tma_desc_ptr: Optional[Pointer] = None,
+        cache_policy: Optional[Int64] = None,
+    ):
         """
         Custom implementation of unpack for non-executable TMAs.
         """
@@ -371,9 +429,26 @@ class CopyBulkTensorTileS2GTrait(Trait):
             exec_value = _cute_nvgpu_ir.atom_set_value(
                 exec_value, attr, tma_desc_ptr.value, loc=loc, ip=ip
             )
+        if cache_policy is not None:
+            if not isinstance(cache_policy, Int64):
+                raise ValueError(
+                    "expects `Int64` value to be provided via the cache_policy kw argument"
+                )
+
+            attr_str = (
+                f"#cute_nvgpu.atom_copy_field_tmastore<{TMA_CACHE_POLICY_FIELD_NAME}>"
+            )
+            attr = ir.Attribute.parse(attr_str)
+            exec_value = _cute_nvgpu_ir.atom_set_value(
+                exec_value, attr, cache_policy.value, loc=loc, ip=ip
+            )
         return exec_value
 
 
+class CopyBulkTensorTileS2GTrait(Trait):
+    pass
+
+
 @dataclass(frozen=True)
 class CopyReduceBulkTensorTileS2GOp(TmaCopyOp):
     """
@@ -400,7 +475,7 @@ class CopyReduceBulkTensorTileS2GOp(TmaCopyOp):
 
     def _make_trait(
         self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs
-    ) -> "CopyReduceBulkTensorTileS2GTrait":
+    ) -> "CopyReduceBulkTensorTileS2GNonExecTrait":
         raise NotImplementedError(
             "Use cpasync.make_tiled_tma_atom to obtain a copy Atom for TMA"
         )
@@ -426,8 +501,20 @@ class CopyReduceBulkTensorTileS2GOp(TmaCopyOp):
             assert False, "unrecognized self.reduction_kind"
 
 
-class CopyReduceBulkTensorTileS2GTrait(Trait):
-    def unpack(self, *, loc=None, ip=None, tma_desc_ptr: Optional[Pointer] = None):
+class CopyReduceBulkTensorTileS2GNonExecTrait(Trait):
+    def with_(
+        self, *, loc=None, ip=None, **kwargs
+    ) -> "CopyReduceBulkTensorTileS2GTrait":
+        return CopyReduceBulkTensorTileS2GTrait(self.unpack(loc=loc, ip=ip, **kwargs))
+
+    def unpack(
+        self,
+        *,
+        loc=None,
+        ip=None,
+        tma_desc_ptr: Optional[Pointer] = None,
+        cache_policy: Optional[Int64] = None,
+    ):
         """
         Custom implementation of unpack for non-executable TMAs.
         """
@@ -440,9 +527,25 @@ class CopyReduceBulkTensorTileS2GTrait(Trait):
             exec_value = _cute_nvgpu_ir.atom_set_value(
                 exec_value, attr, tma_desc_ptr.value, loc=loc, ip=ip
             )
+        if cache_policy is not None:
+            if not isinstance(cache_policy, Int64):
+                raise ValueError(
+                    "expects `Int64` value to be provided via the cache_policy kw argument"
+                )
+            attr_str = (
+                f"#cute_nvgpu.atom_copy_field_tmareduce<{TMA_CACHE_POLICY_FIELD_NAME}>"
+            )
+            attr = ir.Attribute.parse(attr_str)
+            exec_value = _cute_nvgpu_ir.atom_set_value(
+                exec_value, attr, cache_policy.value, loc=loc, ip=ip
+            )
         return exec_value
 
 
+class CopyReduceBulkTensorTileS2GTrait(Trait):
+    pass
+
+
 #
 # Bulk GMEM -> SMEM copies
 #
@@ -494,13 +597,15 @@ class CopyBulkG2STrait(Trait):
         loc=None,
         ip=None,
         mbar_ptr: Optional[Pointer] = None,
+        cache_policy: Optional[Int64] = None,
         **kwargs,
     ):
         """
         Custom implementation of unpack for bulk copy load.
 
         The non-multicast bulk load requires a `mbar_ptr` keyword argument to be provided when
-        using `cute.copy`. Any other kw arguments will be ignored instead of triggering an error.
+        using `cute.copy`. `cache_policy` keyword argument to be provided to set the l2 cache eviction priority.
+        Any other kw arguments will be ignored instead of triggering an error.
         """
         if not isinstance(mbar_ptr, Pointer):
             raise ValueError(
@@ -511,6 +616,18 @@ class CopyBulkG2STrait(Trait):
         val = _cute_nvgpu_ir.atom_set_value(
             self.value, attr, mbar_ptr.value, loc=loc, ip=ip
         )
+        if cache_policy is not None:
+            if not isinstance(cache_policy, Int64):
+                raise ValueError(
+                    "expects `Int64` value to be provided via the cache_policy kw argument"
+                )
+            attr_str = (
+                f"#cute_nvgpu.atom_copy_field_bulkg2s<{TMA_CACHE_POLICY_FIELD_NAME}>"
+            )
+            attr = ir.Attribute.parse(attr_str)
+            val = _cute_nvgpu_ir.atom_set_value(
+                val, attr, cache_policy.value, loc=loc, ip=ip
+            )
         return val
 
 
@@ -566,6 +683,7 @@ class CopyBulkG2SMulticastTrait(Trait):
         ip=None,
         mbar_ptr: Optional[Pointer] = None,
         mcast_mask: Optional[Integer] = None,
+        cache_policy: Optional[Int64] = None,
         **kwargs,
     ):
         """
@@ -592,6 +710,18 @@ class CopyBulkG2SMulticastTrait(Trait):
         val = _cute_nvgpu_ir.atom_set_value(
             val, attr, Int16(mcast_mask).ir_value(loc=loc, ip=ip), loc=loc, ip=ip
         )
+        if cache_policy is not None:
+            if not isinstance(cache_policy, Int64):
+                raise ValueError(
+                    "expects `Int64` value to be provided via the cache_policy kw argument"
+                )
+            attr_str = (
+                f"#cute_nvgpu.atom_copy_field_bulkg2s<{TMA_CACHE_POLICY_FIELD_NAME}>"
+            )
+            attr = ir.Attribute.parse(attr_str)
+            val = _cute_nvgpu_ir.atom_set_value(
+                val, attr, cache_policy.value, loc=loc, ip=ip
+            )
         return val
 
 

python/CuTeDSL/cutlass/cute/nvgpu/cpasync/helpers.py

@@ -35,8 +35,8 @@ from .copy import (
     CopyReduceBulkTensorTileS2GOp,
     CopyBulkTensorTileG2SNonExecTrait,
     CopyBulkTensorTileG2SMulticastNonExecTrait,
-    CopyBulkTensorTileS2GTrait,
-    CopyReduceBulkTensorTileS2GTrait,
+    CopyBulkTensorTileS2GNonExecTrait,
+    CopyReduceBulkTensorTileS2GNonExecTrait,
 )
 
 
@@ -156,7 +156,7 @@ def make_tiled_tma_atom(
             loc=loc,
             ip=ip,
         )
-        return atom.CopyAtom(op, CopyBulkTensorTileS2GTrait(res[0])), res[1]
+        return atom.CopyAtom(op, CopyBulkTensorTileS2GNonExecTrait(res[0])), res[1]
     elif isinstance(op, CopyReduceBulkTensorTileS2GOp):
         res = _cute_nvgpu_ir.atom_make_non_exec_tiled_tma_reduce(
             gmem_tensor.value,
@@ -167,7 +167,10 @@ def make_tiled_tma_atom(
             loc=loc,
             ip=ip,
         )
-        return atom.CopyAtom(op, CopyReduceBulkTensorTileS2GTrait(res[0])), res[1]
+        return (
+            atom.CopyAtom(op, CopyReduceBulkTensorTileS2GNonExecTrait(res[0])),
+            res[1],
+        )
     else:
         raise ValueError(f"expects a bulk tensor (TMA) Copy Op, but got {op}")
 

python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/helpers.py

@@ -285,7 +285,9 @@ def find_tmem_tensor_col_offset(tmem_tensor: Tensor, *, loc=None, ip=None) -> In
     """
     tmem_col_mask = 0x0000FFFF
     offset = (
-        core.cosize(recast_tensor(tmem_tensor, Int32).layout, loc=loc, ip=ip)
+        core.cosize(
+            recast_tensor(tmem_tensor, Int32, loc=loc, ip=ip).layout, loc=loc, ip=ip
+        )
         & tmem_col_mask
     )
     if isinstance(offset, int):

python/CuTeDSL/cutlass/cute/runtime.py

@@ -10,19 +10,24 @@
 # is strictly prohibited.
 
 import ctypes
+import os
+import sys
+from pathlib import Path
 from functools import lru_cache
 import itertools
 import operator
-from typing import Union, Optional
+from typing import Union, Optional, Type, List
 
 # MLIR modules imports
 from cutlass._mlir import ir
 import cutlass._mlir.dialects.cute as _cute_ir
+import cutlass._mlir.dialects.cuda as _cuda_dialect
 
-from cutlass.cutlass_dsl import JitArgAdapterRegistry, DSLRuntimeError
+from cutlass.cutlass_dsl import JitArgAdapterRegistry, CuTeDSL as _CuTeDSL
+from cutlass.base_dsl.common import DSLRuntimeError
 
 # Local modules imports
-from .typing import AddressSpace, Tensor, Type, Pointer, Numeric
+from .typing import AddressSpace, Layout, Tensor, Pointer, Numeric, SymInt
 from . import core
 from .tensor import _Tensor as CoreTensor
 
@@ -76,6 +81,9 @@ class _Pointer(Pointer):
     def __get_mlir_types__(self):
         return [self.mlir_type]
 
+    def __tvm_ffi_opaque_ptr__(self):
+        return self._pointer
+
     def __c_pointers__(self):
         if self._c_pointer is None:
             self._desc = ctypes.c_void_p(int(self._pointer))
@@ -104,14 +112,6 @@ class _Pointer(Pointer):
     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
-
-        return False
-
     def __str__(self) -> str:
         return f"Ptr<0x{int(self._pointer):016x}@{self._addr_space}>"
 
@@ -120,12 +120,30 @@ class _Pointer(Pointer):
 
 
 class _Tensor(Tensor):
-    def __init__(self, tensor, assumed_align=None, use_32bit_stride=False):
+    def __init__(
+        self,
+        tensor,
+        assumed_align=None,
+        use_32bit_stride=False,
+        *,
+        enable_tvm_ffi=False,
+    ):
         # If tensor is already a DLPack object, use it directly
         if hasattr(tensor, "__dlpack_device__") and not hasattr(tensor, "__dlpack__"):
-            self._dlpack_data = tensor
+            self._dlpack_data = tensor.__dlpack_device__()
         else:
-            self._dlpack_data = tensor.__dlpack__()
+            try:
+                # we expect no stream sync. Because torch has different default behavior
+                # for stream parameter on different version.
+                # we need to explicitly pass -1 to achieve no sync effects.
+                self._dlpack_data = tensor.__dlpack__(stream=-1)
+            except Exception:
+                self._dlpack_data = tensor.__dlpack__()
+        if enable_tvm_ffi:
+            import tvm_ffi
+
+            self._tvm_ffi_tensor = tvm_ffi.from_dlpack(tensor)
+            self._dlpack_data = self._tvm_ffi_tensor.__dlpack__()
         self._dltensor_wrapper = None
         self._assumed_align = assumed_align
         self._is_dynamic = False
@@ -217,7 +235,6 @@ class _Tensor(Tensor):
             mode, stride_order, divisibility
         )
         return self
-
     @property
     @lazily_load_dltensor
     def element_type(self) -> Type[Numeric]:
@@ -342,6 +359,18 @@ class _Tensor(Tensor):
     def data_ptr(self):
         return self._dltensor_wrapper.data_ptr
 
+    @property
+    @lazily_load_dltensor
+    def dynamic_shapes_mask(self):
+        """Get the mask of dynamic shapes in the tensor."""
+        return self._dltensor_wrapper.get_dynamic_shapes_mask()
+
+    @property
+    @lazily_load_dltensor
+    def dynamic_strides_mask(self):
+        """Get the mask of dynamic strides in the tensor."""
+        return self._dltensor_wrapper.get_dynamic_strides_mask()
+
     @lazily_load_dltensor
     def __c_pointers__(self):
         self._memref_desc = self._dltensor_wrapper.build_memref_desc(
@@ -357,11 +386,336 @@ class _Tensor(Tensor):
         assert isinstance(values[0], CoreTensor)
         return CoreTensor(values[0].value, self._dtype)
 
+    def __tvm_ffi_object__(self):
+        return self._tvm_ffi_tensor
+
+
+def _get_cute_type_str(inp):
+    def _convert_dyn_elem(e):
+        return f"?{{i{e.width} div={e.divisibility}}}"
+
+    elems = [_convert_dyn_elem(e) if isinstance(e, SymInt) else str(e) for e in inp]
+    return "(" + ",".join(elems) + ")"
+
+
+class _FakeCompactTensor(Tensor):
+    def __init__(
+        self,
+        dtype,
+        shape,
+        stride_order,
+        memspace=None,
+        assumed_align=None,
+        use_32bit_stride=False,
+    ):
+        self._dtype = dtype
+        self._shape = shape
+        self._stride_order = stride_order or tuple(range(len(shape)))
+        self._memspace = memspace or AddressSpace.gmem
+        self._assumed_align = assumed_align or -(-dtype.width // 8)
+        self._use_32bit_stride = use_32bit_stride
+
+    def __str__(self) -> str:
+        return f"FakeTensorOrdered<{self._dtype}, {self._shape}, {self._stride_order}>"
+
+    def __repr__(self):
+        return self.__str__()
+
+    @property
+    def mlir_type(self) -> ir.Type:
+        shape_ty = ir.Type.parse(
+            '!cute.shape<"' + _get_cute_type_str(self._shape) + '">'
+        )
+        layout_ty = _cute_ir.LayoutType.get_ordered(
+            shape_ty, self._stride_order, self._use_32bit_stride
+        )
+        self._stride = layout_ty.stride
+        ptr_ty = _cute_ir.PtrType.get(
+            self._dtype.mlir_type, self._memspace, self._assumed_align
+        )
+        return _cute_ir.MemRefType.get(ptr_ty, layout_ty)
+
+    def __get_mlir_types__(self):
+        return [self.mlir_type]
+
+    def __new_from_mlir_values__(self, values):
+        assert len(values) == 1
+        assert isinstance(values[0], CoreTensor)
+        return CoreTensor(values[0].value, self._dtype)
+
+    def __setitem__(self, crd, value):
+        raise DSLRuntimeError("runtime._FakeCompactTensor is not indexable")
+
+    def __getitem__(self, crd):
+        raise DSLRuntimeError("runtime._FakeCompactTensor is not indexable")
+
+    @property
+    def element_type(self) -> Type[Numeric]:
+        return self._dtype
+
+    @property
+    def memspace(self):
+        return self._memspace
+
+    @property
+    def iterator(self):
+        raise DSLRuntimeError("runtime._FakeTensor has dummy iterator")
+
+    @property
+    def shape(self):
+        return self._shape
+
+    @property
+    def stride(self):
+        return self._stride
+
+    @property
+    def dynamic_shapes_mask(self):
+        return tuple(1 if isinstance(e, SymInt) else 0 for e in self._shape)
+
+    @property
+    def dynamic_strides_mask(self):
+        return tuple(1 if isinstance(e, SymInt) else 0 for e in self._stride)
+
+    def fill(self, value: Numeric):
+        raise DSLRuntimeError("runtime._FakeCompactTensor is not writable")
+
+
+class _FakeTensor(Tensor):
+    """Fake Tensor implementation as a placeholder.
+    It mimics the interface of Tensor, but does not hold real data or allow indexing.
+    Used for compilation or testing situations where only shape/type/layout information is needed.
+    All attempts to access or mutate data will raise errors.
+    """
+
+    """
+    Create a fake tensor with the given shape, type, and layout.
+
+    :param dtype: Data type of the tensor elements
+    :type dtype: Type[Numeric]
+    :param shape: Shape of the tensor, consists of int (static) or SymInt (dynamic)
+    :type shape: tuple[int, ...]
+    :param stride: Stride of the tensor, defaults to None, consists of int (static) or SymInt (dynamic)
+    :type stride: tuple[int, ...], optional
+    :param assumed_align: Assumed alignment of the tensor (bytes), defaults to None. If None, uses the element size bytes as the assumed alignment.
+    :type assumed_align: int, optional
+    :param use_32bit_stride: Whether to use 32-bit stride. Defaults to False. When True, the dynamic stride bitwidth
+        will be set to 32 for small problem sizes (cosize(layout) <= Int32_max) for better performance. This is only applied
+        when the dimension is dynamic.
+    :type use_32bit_stride: bool, optional
+
+    """
+
+    def __init__(self, dtype, shape, *, stride, memspace=None, assumed_align=None):
+        self._dtype = dtype
+        self._shape = shape
+        self._stride = stride
+        self._memspace = memspace or AddressSpace.generic
+        self._assumed_align = assumed_align
+        if assumed_align is None:
+            # use the bytes width of the element dtype. The alignment is at least one byte align.
+            self._assumed_align = (self._dtype.width + 7) // 8
+
+        if not isinstance(shape, (tuple, list)):
+            raise ValueError(f"Expected tuple or list but got {type(shape)}")
+
+        if not all(isinstance(s, (int, SymInt)) for s in self._shape):
+            raise ValueError("All shape elements must be int or SymInt")
+
+        if stride is not None and not all(
+            isinstance(s, (int, SymInt)) for s in self._stride
+        ):
+            raise ValueError("All stride elements must be int or SymInt")
+    @property
+    def mlir_type(self) -> ir.Type:
+        shape_str = _get_cute_type_str(self._shape)
+        stride_str = _get_cute_type_str(self._stride)
+        layout_ty = ir.Type.parse(f'!cute.layout<"{shape_str}:{stride_str}">')
+        ptr_ty = _cute_ir.PtrType.get(
+            self._dtype.mlir_type, self._memspace, self._assumed_align
+        )
+        return _cute_ir.MemRefType.get(ptr_ty, layout_ty)
+
+    def __get_mlir_types__(self):
+        return [self.mlir_type]
+
+    def __new_from_mlir_values__(self, values):
+        assert len(values) == 1
+        assert isinstance(values[0], CoreTensor)
+        return CoreTensor(values[0].value, self._dtype)
+
+    def __str__(self) -> str:
+        return f"FakeTensor<{self._dtype}, {self._shape}, {self._stride}>"
+
+    def __repr__(self):
+        return self.__str__()
+
+    def __setitem__(self, crd, value):
+        raise DSLRuntimeError("runtime._FakeTensor is not indexable")
+
+    def __getitem__(self, crd):
+        raise DSLRuntimeError("runtime._FakeTensor is not indexable")
+
+    @property
+    def element_type(self) -> Type[Numeric]:
+        return self._dtype
+
+    @property
+    def memspace(self):
+        return self._memspace
+
+    @property
+    def iterator(self):
+        raise DSLRuntimeError("runtime._FakeTensor has dummy iterator")
+
+    @property
+    def shape(self):
+        return self._shape
+
+    @property
+    def stride(self):
+        return self._stride
+
+    @property
+    def dynamic_shapes_mask(self):
+        return tuple(1 if isinstance(e, SymInt) else 0 for e in self._shape)
+
+    @property
+    def dynamic_strides_mask(self):
+        return tuple(1 if isinstance(e, SymInt) else 0 for e in self._stride)
+
+    def fill(self, value: Numeric):
+        raise DSLRuntimeError("runtime._FakeTensor is not writable")
+
+
+def make_fake_compact_tensor(
+    dtype,
+    shape,
+    *,
+    stride_order=None,
+    memspace=None,
+    assumed_align=None,
+    use_32bit_stride=False,
+):
+    """
+    Create a fake tensor with the specified shape, element type, and a compact memory layout.
+
+    :param dtype: Data type of the tensor elements.
+    :type dtype: Type[Numeric]
+    :param shape: Shape of the tensor.
+    :type shape: tuple[int, ...]
+    :param stride_order: Order in which strides (memory layout) are assigned to the tensor dimensions.
+        If None, the default layout is row-major. Otherwise, it should be a permutation of the dimension indices.
+    :type stride_order: tuple[int, ...], optional
+    :param memspace: Memory space where the fake tensor resides. Optional.
+    :type memspace: str, optional
+    :param assumed_align: Assumed byte alignment for the tensor data. If None, the default alignment is used.
+    :type assumed_align: int, optional
+    :param use_32bit_stride: Whether to use 32-bit stride for dynamic dimensions. If True and the total size of the
+        layout (cosize(layout)) fits within int32, then dynamic strides will use 32-bit integers for improved performance.
+        Only applies when dimensions are dynamic. Defaults to False.
+    :type use_32bit_stride: bool, optional
+    :return: An instance of a fake tensor with the given properties and compact layout.
+    :rtype: _FakeCompactTensor
+
+    **Examples:**
+
+    .. code-block:: python
+
+        @cute.jit
+        def foo(x: cute.Tensor):
+            ...
+
+        x = make_fake_compact_tensor(
+            cutlass.Float32, (100, cute.sym_int32(divisibility=8)), stride_order=(1, 0)
+        )
+
+        # Compiled function will take a tensor with the type:
+        #   tensor<ptr<f32, generic> o (100,?{div=8}):(?{i32 div=8},1)>
+        compiled_foo = cute.compile(foo, x)
+    """
+
+    return _FakeCompactTensor(
+        dtype,
+        shape,
+        stride_order=stride_order,
+        memspace=memspace,
+        assumed_align=assumed_align,
+        use_32bit_stride=use_32bit_stride,
+    )
+
+def make_fake_tensor(dtype, shape, stride, *, memspace=None, assumed_align=None):
+    """
+    Create a fake tensor with the specified element type, shape, and stride.
+
+    :param dtype: Data type of the tensor elements.
+    :type dtype: Type[Numeric]
+    :param shape: Shape of the tensor.
+    :type shape: tuple[int, ...]
+    :param stride: Stride of the tensor.
+    :type stride: tuple[int, ...]
+    :param assumed_align: Assumed byte alignment for the tensor data. If None, the default alignment is used. Defaults to None.
+    :type assumed_align: int, optional
+    :return: An instance of a fake tensor with the given properties.
+    :rtype: _FakeTensor
+    """
+    return _FakeTensor(
+        dtype, shape, stride=stride, memspace=memspace, assumed_align=assumed_align
+    )
+
+
+class _FakeStream:
+    """A fake stream that can be used as a placeholder for a stream in compilation.
+
+    When use_tvm_ffi_env_stream is True and the function is compiled with TVM-FFI,
+    the argument will be skipped from the function signature and we pass in
+    this value through the environment stream obtained from caller context
+    (e.g. torch.cuda.current_stream()).
+    """
+
+    use_tvm_ffi_env_stream: bool
+
+    def __init__(self, *, use_tvm_ffi_env_stream: bool = False):
+        self.use_tvm_ffi_env_stream = use_tvm_ffi_env_stream
+
+    def __str__(self) -> str:
+        return f"FakeStream"
+
+    def __repr__(self):
+        return self.__str__()
+
+    def __new_from_mlir_values__(self, values):
+        assert len(values) == 1
+        return values[0]
+
+    def __c_pointers__(self):
+        return [0]
+
+    def __get_mlir_types__(self):
+        return [_cuda_dialect.StreamType.get()]
+
+
+def make_fake_stream(*, use_tvm_ffi_env_stream: bool = False):
+    """Create a fake stream that can be used as a placeholder for a stream in compilation.
+
+    When use_tvm_ffi_env_stream is True and the function is compiled with TVM-FFI,
+    the argument will be skipped from the function signature and we pass in
+    this value through the environment stream obtained from caller context
+    (e.g. torch.cuda.current_stream()). This can speedup the calling process
+    since we no longer need to do stream query in python.
+
+    :param use_tvm_ffi_env_stream: Whether to skip this parameter use environment stream instead.
+    :type use_tvm_ffi_env_stream: bool
+    """
+    return _FakeStream(use_tvm_ffi_env_stream=use_tvm_ffi_env_stream)
+
 
 def from_dlpack(
     tensor_dlpack,
     assumed_align=None,
     use_32bit_stride=False,
+    *,
+    enable_tvm_ffi=False,
 ) -> Tensor:
     """Convert from tensor object supporting __dlpack__() to a CuTe Tensor.
 
@@ -374,25 +728,33 @@ def from_dlpack(
       stride bitwidth will be set to 32 for small problem size (cosize(layout) <= Int32_max) for better performance.
       This is only applied when the dimension is dynamic.
     :type use_32bit_stride: bool, optional
+    :param enable_tvm_ffi: Whether to enable TVM-FFI, defaults to False. When True, the tensor will be converted to
+      a TVM-FFI function compatible tensor.
+    :type enable_tvm_ffi: bool, optional
     :return: A CuTe Tensor object
     :rtype: Tensor
 
-    Examples:
-        .. code-block:: python
+    **Examples:**
 
-            import torch
-            from cutlass.cute.runtime import from_dlpack
-            x = torch.randn(100, 100)
-            y = from_dlpack(x)
-            y.shape
-            # (100, 100)
-            type(y)
-            # <class 'cutlass.cute.Tensor'>
+    .. code-block:: python
+
+        import torch
+        from cutlass.cute.runtime import from_dlpack
+        x = torch.randn(100, 100)
+        y = from_dlpack(x)
+        y.shape
+        # (100, 100)
+        type(y)
+        # <class 'cutlass.cute.Tensor'>
     """
+    # If the environment variable `CUTE_DSL_ENABLE_TVM_FFI` is set to True, the tensor will be converted to
+    # a TVM-FFI function compatible tensor.
+    enable_tvm_ffi = enable_tvm_ffi or _CuTeDSL._get_dsl().envar.enable_tvm_ffi
     return _Tensor(
         tensor_dlpack,
         assumed_align=assumed_align,
         use_32bit_stride=use_32bit_stride,
+        enable_tvm_ffi=enable_tvm_ffi,
     )
 
 
@@ -451,6 +813,23 @@ def make_ptr(
     return _Pointer(address_value, dtype, mem_space, assumed_align=assumed_align)
 
 
+def nullptr(
+    dtype: Type[Numeric],
+    mem_space: AddressSpace = AddressSpace.generic,
+    assumed_align=None,
+) -> Pointer:
+    """Create a null pointer which is useful for compilation
+
+    :param dtype: Data type of the pointer elements
+    :type dtype: Type[Numeric]
+    :param mem_space: Memory address space, defaults to AddressSpace.generic
+    :type mem_space: AddressSpace, optional
+    :return: A null pointer object
+    :rtype: Pointer
+    """
+    return make_ptr(dtype, 0, mem_space, assumed_align=assumed_align)
+
+
 class TensorAdapter:
     """
     Convert a DLPack protocol supported tensor/array to a cute tensor.
@@ -469,6 +848,74 @@ class TensorAdapter:
         return self._arg.__get_mlir_types__()
 
 
+def find_runtime_libraries(*, enable_tvm_ffi: bool = True) -> List[str]:
+    """
+    Find the runtime libraries that needs to be available for loading modules.
+
+    :param enable_tvm_ffi: Whether to enable TVM-FFI.
+    :type enable_tvm_ffi: bool, optional
+    :return: A list of runtime libraries that needs to be available for loading modules.
+    :rtype: list
+    """
+
+    def _get_cuda_dialect_runtime_path():
+        libs = os.environ.get("CUTE_DSL_LIBS")
+        if libs:
+            sep = ";" if sys.platform.startswith("win32") else ":"
+            for path in libs.split(sep):
+                if path.endswith("libcuda_dialect_runtime.so"):
+                    return path
+        try:
+            # find package library from wheel package
+            pkg_base = Path(__file__).resolve().parent.parent
+            lib_path = pkg_base / "lib" / "libcuda_dialect_runtime.so"
+            if lib_path.is_file():
+                return str(lib_path)
+        except OSError:
+            return None
+
+    libs = []
+    cuda_dialect_runtime_path = _get_cuda_dialect_runtime_path()
+    if cuda_dialect_runtime_path:
+        libs.append(cuda_dialect_runtime_path)
+
+    if enable_tvm_ffi:
+        import tvm_ffi
+
+        libs.append(tvm_ffi.libinfo.find_libtvm_ffi())
+
+    return libs
+
+# cache to load runtime libraries so they can be found by the DSO loader
+_LOAD_MODULE_LIBS_CACHE = []
+
+
+def load_module(file_path: str, *, enable_tvm_ffi: bool = True):
+    """Load a module from a file path. Today only support TVM-FFI module.
+
+    :param file_path: The path to the module file
+    :type file_path: str
+    :param enable_tvm_ffi: Whether to enable TVM-FFI, defaults to True. When True, the module will be loaded as a TVM-FFI module.
+    :type enable_tvm_ffi: bool, optional
+    :return: A module object
+    :rtype: module
+    """
+    if len(_LOAD_MODULE_LIBS_CACHE) == 0:
+        # ensure the runtime libraries are loaded so they can be found by the DSO loader
+        # no need to load tvm_ffi library here since it will be loaded by tvm_ffi package.
+        for path in find_runtime_libraries(enable_tvm_ffi=False):
+            if Path(path).exists():
+                _LOAD_MODULE_LIBS_CACHE.append(ctypes.CDLL(path))
+
+    if enable_tvm_ffi:
+        import tvm_ffi
+
+        return tvm_ffi.load_module(file_path)
+    else:
+        raise DSLRuntimeError(
+            "Unimplemented, please load the module with enable_tvm_ffi=True."
+        )
+
 # -------------------------------------------------------------------------
 # Try to register_jit_arg_adapter for TensorAdapter
 # -------------------------------------------------------------------------
@@ -479,7 +926,6 @@ try:  # Register for numpy.ndarray
     JitArgAdapterRegistry.register_jit_arg_adapter(numpy.ndarray)(TensorAdapter)
 except ImportError:
     pass  # silent attempt, suppress error
-
 try:  # Register for torch.Tensor
     import torch
 

python/CuTeDSL/cutlass/cute/tensor.py

@@ -20,6 +20,7 @@ from cutlass.cutlass_dsl import (
     T,
     cutlass_arith,
     _binary_op_type_promote,
+    CuTeDSL,
 )
 from cutlass._mlir import ir
 import cutlass._mlir.dialects.cute as _cute_ir
@@ -33,6 +34,7 @@ from .core import (
     _pack_coord,
     _pack_shape,
     _ComposedLayout,
+    _ComposedLayoutWithInnerFunc,
     is_static,
     is_weakly_congruent,
     rank,
@@ -250,7 +252,7 @@ class _Tensor(Tensor):
             val = tensor[0]  # Error: sub-byte scalar dereference not supported
         """
         if has_underscore(crd):
-            return slice_(self.value, crd, loc=loc, ip=ip)
+            return slice_(self, crd, loc=loc, ip=ip)
         elif isinstance(self.type, _cute_ir.CoordTensorType):
             res = _cute_ir.get_iter(
                 slice_(self, crd, loc=loc, ip=ip).value, loc=loc, ip=ip
@@ -448,7 +450,7 @@ class _Tensor(Tensor):
 
         self._check_can_load_store()
 
-        res_vect = _cute_ir.memref_load_vec(self.value, row_major=True, loc=loc, ip=ip)
+        res_vect = _cute_ir.memref_load_vec(self.value, loc=loc, ip=ip)
         if self.element_type is Boolean:
             assert res_vect.type.element_type == T.i8(), (
                 f"Boolean tensor must be stored as i8 in memory, but got {res_vect.type.element_type}"
@@ -501,9 +503,7 @@ class _Tensor(Tensor):
         # Implicit upcast to wider type
         new_data = self._cvt_to_dest(data, loc=loc, ip=ip)
 
-        return _cute_ir.memref_store_vec(
-            new_data, self.value, row_major=True, loc=loc, ip=ip
-        )
+        return _cute_ir.memref_store_vec(new_data, self.value, loc=loc, ip=ip)
 
     @dsl_user_op
     def fill(self, value: Numeric, *, loc=None, ip=None) -> None:
@@ -555,6 +555,11 @@ class _Tensor(Tensor):
         ):
             raise ValueError(f"{self} doesn't support load and store")
 
+        if self.type.is_swizzled:
+            raise NotImplementedError(
+                f"load & store swizzled memory is not supported yet: {self}"
+            )
+
     def _check_can_dereference(self):
         # Check for sub-byte types and raise error if needed
         if self.element_type.width % 8 != 0 and self.element_type is not Boolean:
@@ -579,11 +584,16 @@ def make_tensor(
     evaluates coordinates by applying the layout mapping and dereferencing the engine
     at the resulting offset.
 
-    :param iterator: Engine component (pointer, iterator, or counting iterator) that provides
-                    data access capabilities
-    :type iterator: Union[Pointer, IntTuple]
+    :param iterator: Engine component that provides data access capabilities. Can be:
+                    - A pointer (Pointer type)
+                    - An integer or integer tuple for coordinate tensors
+                    - A shared memory descriptor (SmemDescType)
+    :type iterator: Union[Pointer, IntTuple, ir.Value]
     :param layout: Layout component that defines the mapping from logical coordinates to
-                  physical offsets
+                  physical offsets. Can be:
+                  - A shape tuple that will be converted to a layout
+                  - A Layout object
+                  - A ComposedLayout object (must be a normal layout)
     :type layout: Union[Shape, Layout, ComposedLayout]
     :param loc: Source location for MLIR operation tracking, defaults to None
     :type loc: Optional[Location]
@@ -592,30 +602,40 @@ def make_tensor(
     :return: A tensor object representing the composition E ∘ L
     :rtype: Tensor
 
-    :raises ValueError: If iterator type is not supported
+    :raises TypeError: If iterator type is not a supported type
+    :raises ValueError: If layout is a composed layout with customized inner functions
 
     **Examples:**
 
     .. code-block:: python
 
-        # Create a tensor with row-major layout
+        # Create a tensor with row-major layout from a pointer
+        ptr = make_ptr(Float32, base_ptr, AddressSpace.gmem)
         layout = make_layout((64, 128), stride=(128, 1))
         tensor = make_tensor(ptr, layout)
 
-        # Create a tensor with hierarchical layout
+        # Create a tensor with hierarchical layout in shared memory
+        smem_ptr = make_ptr(Float16, base_ptr, AddressSpace.smem)
         layout = make_layout(((128, 8), (1, 4, 1)), stride=((32, 1), (0, 8, 4096)))
         tensor = make_tensor(smem_ptr, layout)
 
-        # Create a coord tensor
+        # Create a coordinate tensor
         layout = make_layout(2, stride=16 * E(0))
-        tensor = make_tensor(5, layout)
+        tensor = make_tensor(5, layout)  # coordinate tensor with iterator starting at 5
 
     Notes:
         - The engine (iterator) must support random access operations
         - Common engine types include raw pointers, arrays, and random-access iterators
         - The layout defines both the shape (logical dimensions) and stride (physical mapping)
         - Supports both direct coordinate evaluation T(c) and partial evaluation (slicing)
+        - ComposedLayouts must be "normal" layouts (no inner functions)
+        - For coordinate tensors, the iterator is converted to a counting sequence
     """
+    if isinstance(layout, _ComposedLayoutWithInnerFunc):
+        raise ValueError(
+            "CuTe DSL tensor does not support composed layouts with inner functions: {layout}"
+        )
+
     if not isinstance(layout, (Layout, ComposedLayout)):
         layout = make_layout(layout, loc=loc, ip=ip)
     # Automatic decay to normal layout
@@ -964,6 +984,65 @@ def print_tensor(
     _cute_ir.print_view(tensor.value, verbose=verbose, is_signed=signed, loc=loc, ip=ip)
 
 
+def _get_row_and_col_map(col_maj_shape_1d: tuple, is_row_to_col: bool):
+    """
+    Create an index mapping mask for converting between row-major and column-major vector ordering.
+    """
+
+    # create row-major layout with compact row-major stride
+    if col_maj_shape_1d is None:
+        raise ValueError("vector shape for row/col map cannot be None")
+    if isinstance(col_maj_shape_1d, tuple):
+        row_maj_shape_1d = tuple(reversed(flatten(col_maj_shape_1d)))
+    else:
+        # Single dimension
+        row_maj_shape_1d = col_maj_shape_1d
+
+    flat_shape = flatten(row_maj_shape_1d)
+    if isinstance(flat_shape, tuple):
+        strides = []
+        current_stride = 1
+        flat_shape = tuple(reversed(flat_shape))
+        for dim in flat_shape:
+            strides.append(current_stride)
+            current_stride *= dim
+        row_maj_stride = tuple(reversed(strides))
+    else:
+        # Single dimension
+        row_maj_stride = 1
+
+    row_maj_lay_1d = make_layout(row_maj_shape_1d, stride=row_maj_stride)
+
+    # get idx map
+    shape_size = size(row_maj_shape_1d)
+    mask = [0] * shape_size
+
+    for col_index in range(shape_size):
+        row_index = crd2idx(col_index, row_maj_lay_1d)
+        if is_row_to_col:
+            mask[row_index] = col_index
+        else:
+            mask[col_index] = row_index
+
+    return mask
+
+
+def _row2col(vec: ir.Value, *, shape, loc=None, ip=None) -> ir.Value:
+    """
+    Convert a vector or tensor from row-major order to column-major order.
+    """
+    row_and_col_map = _get_row_and_col_map(shape, is_row_to_col=True)
+    return vector.shuffle(vec, vec, row_and_col_map, loc=loc, ip=ip)
+
+
+def _col2row(vec: ir.Value, *, shape, loc=None, ip=None) -> ir.Value:
+    """
+    Convert a vector or tensor from column-major order to row-major order.
+    """
+    row_and_col_map = _get_row_and_col_map(shape, is_row_to_col=False)
+    return vector.shuffle(vec, vec, row_and_col_map, loc=loc, ip=ip)
+
+
 def _infer_broadcast_shape(*shapes: Shape) -> Shape:
     """
     Infer the broadcasted shape from multiple input shapes according to broadcasting rules.
@@ -1129,6 +1208,32 @@ class TensorSSA(cutlass_arith.ArithValue):
 
         return res
 
+    @dsl_user_op
+    def apply_op(self, op, other, flip=False, *, loc=None, ip=None) -> "TensorSSA":
+        """
+        Apply a binary operation to this tensor and another operand.
+
+        This is a public interface to the internal _apply_op method, providing
+        a stable API for external users who need to apply custom operations.
+
+        Args:
+            op: The operation function (e.g., operator.add, operator.mul, etc.)
+            other: The other operand (TensorSSA, ArithValue, or scalar)
+            flip: Whether to flip the operands (for right-hand operations)
+            loc: MLIR location (optional)
+            ip: MLIR insertion point (optional)
+
+        Returns:
+            TensorSSA: The result of the operation
+
+        Example:
+            >>> tensor1 = cute.Tensor(...)
+            >>> tensor2 = cute.Tensor(...)
+            >>> result = tensor1.apply_op(operator.add, tensor2)
+            >>> # Equivalent to: tensor1 + tensor2
+        """
+        return self._apply_op(op, other, flip=flip, loc=loc, ip=ip)
+
     @dsl_user_op
     def broadcast_to(self, target_shape: Shape, *, loc=None, ip=None) -> "TensorSSA":
         """
@@ -1147,20 +1252,21 @@ class TensorSSA(cutlass_arith.ArithValue):
 
         transform_leaf(_check_broadcast, shape, target_shape)
 
+        # convert TensorSSA col-major vec to row-m to be compatible with mlir vector ops
+        row_major_vec = _col2row(self, shape=shape, loc=loc, ip=ip)
         # reshape to flatten N-D vector
         flat_shp = flatten_to_tuple(shape)
         temp_ty = ir.VectorType.get(list(flat_shp), self.dtype.mlir_type)
-        temp_vect = vector.shape_cast(temp_ty, self, loc=loc, ip=ip)
+        temp_vect = vector.shape_cast(temp_ty, row_major_vec, loc=loc, ip=ip)
 
         # broadcast to result N-D vector
         flat_tgt_shp = flatten_to_tuple(target_shape)
         temp_tgt_ty = ir.VectorType.get(list(flat_tgt_shp), self.dtype.mlir_type)
         temp_tgt_vect = vector.broadcast(temp_tgt_ty, temp_vect, loc=loc, ip=ip)
 
-        res_1d_ty = ir.VectorType.get([size(target_shape)], self.dtype.mlir_type)  # type: ignore
-        res_1d_vect = vector.shape_cast(res_1d_ty, temp_tgt_vect, loc=loc, ip=ip)
-
-        return TensorSSA(res_1d_vect, target_shape, self.dtype)
+        return self._build_result(
+            temp_tgt_vect, target_shape, row_major=True, loc=loc, ip=ip
+        )
 
     @dsl_user_op
     def __pow__(self, other, *, loc=None, ip=None) -> "TensorSSA":
@@ -1499,15 +1605,43 @@ class TensorSSA(cutlass_arith.ArithValue):
         assert isinstance(flat_crd, tuple) and is_static(flat_crd)
         return flat_shp, flat_crd
 
-    def _build_result(self, res_vect, res_shp, *, loc=None, ip=None):
+    def _build_result(self, res_vect, res_shp, *, row_major=False, loc=None, ip=None):
         if isinstance(res_shp, ir.Value):
             raise ValueError(f"Expected static shape, but got {self._shape}")
 
         # cast back to 1D vector
         res_1d_ty = ir.VectorType.get([size(res_shp)], self.type.element_type)
         res_1d_vect = vector.shape_cast(res_1d_ty, res_vect, loc=loc, ip=ip)
+
+        if row_major:
+            res_1d_vect = _row2col(res_1d_vect, shape=res_shp, loc=loc, ip=ip)
+
         return TensorSSA(res_1d_vect, res_shp, self.dtype)
 
+    @dsl_user_op
+    def reshape(self, shape: Shape, *, loc=None, ip=None) -> "TensorSSA":
+        """Reshape the tensor to a new shape.
+
+        :param shape: The new shape to reshape to.
+        :type shape: Shape
+        :return: A new tensor with the same elements but with the new shape.
+        :rtype: TensorSSA
+        :raises NotImplementedError: If dynamic size is not supported
+        :raises ValueError: If the new shape is not compatible with the current shape
+        """
+
+        cur_size = size(self.shape, loc=loc, ip=ip)
+        shp_size = size(shape, loc=loc, ip=ip)
+
+        if type(shp_size) is not int:
+            raise NotImplementedError(f"dynamic shape is not supported: {shape}")
+        if cur_size != shp_size:
+            raise ValueError(
+                f"expected reshaped size to be the same: {self.shape} -> {shape}"
+            )
+
+        return TensorSSA(self, shape, self.dtype)
+
     @dsl_user_op
     def __getitem__(
         self, crd: Coord, *, loc=None, ip=None
@@ -1562,9 +1696,11 @@ class TensorSSA(cutlass_arith.ArithValue):
 
         flat_shp, flat_crd = self._flatten_shape_and_coord(crd, loc=loc, ip=ip)
 
+        # convert TensorSSA col-major vec to row-m to be compatible with mlir vector ops
+        row_major_vec = _col2row(self, shape=self._shape, loc=loc, ip=ip)
         multi_dim_ty = ir.VectorType.get(list(flat_shp), self.type.element_type)
         # vector<NxTy> -> vector<N_0xN_1x...xTy>
-        tmp_vect = vector.shape_cast(multi_dim_ty, self, loc=loc, ip=ip)
+        tmp_vect = vector.shape_cast(multi_dim_ty, row_major_vec, loc=loc, ip=ip)
 
         # Slice and keep dims matching `_` or None
         res_shp = slice_(self._shape, crd, loc=loc, ip=ip)
@@ -1593,7 +1729,7 @@ class TensorSSA(cutlass_arith.ArithValue):
 
         # Slice and keep dims matching `_` or None
         res_shp = slice_(self._shape, crd, loc=loc, ip=ip)
-        return self._build_result(res_vect, res_shp, loc=loc, ip=ip)
+        return self._build_result(res_vect, res_shp, row_major=True, loc=loc, ip=ip)
 
     @dsl_user_op
     def to(self, dtype: Type[Numeric], *, loc=None, ip=None):
@@ -1634,11 +1770,15 @@ class TensorSSA(cutlass_arith.ArithValue):
                 src, dtype.signed, dtype.mlir_type, loc=loc, ip=ip
             )
         elif issubclass(src_dtype, Integer) and dtype.is_float:
-            # fast conversion path for supported combinations
+            # check if there is a fast conversion path for given data types and arch
+            fast_cvt_func = None
             if src_dtype in (Int8, Uint8) and dtype == BFloat16:
-                res_vect = cvt_i8_bf16_intrinsic(src, size(self.shape), loc=loc, ip=ip)
+                fast_cvt_func = cvt_i8_bf16_intrinsic
             elif src_dtype == Int4 and dtype == BFloat16:
-                res_vect = cvt_i4_bf16_intrinsic(src, size(self.shape), loc=loc, ip=ip)
+                fast_cvt_func = cvt_i4_bf16_intrinsic
+            arch = CuTeDSL._get_dsl().get_arch_enum()
+            if fast_cvt_func is not None and arch in fast_cvt_func.supported_archs:
+                res_vect = fast_cvt_func(src, size(self.shape), loc=loc, ip=ip)
             else:
                 res_vect = cutlass_arith.itofp(
                     src, src_dtype.signed, dtype.mlir_type, loc=loc, ip=ip
@@ -1748,8 +1888,10 @@ class TensorSSA(cutlass_arith.ArithValue):
         assert depth(flat_shp) == 1 and depth(flat_prof) == 1
         assert rank(flat_shp) == rank(flat_prof)
 
+        # convert TensorSSA col-major vec to row-m to be compatible with mlir vector ops
+        row_major_vec = _col2row(self, shape=self._shape, loc=loc, ip=ip)
         temp_ty = ir.VectorType.get(list(flat_shp), elem_type.mlir_type)
-        temp_vect = vector.shape_cast(temp_ty, self, loc=loc, ip=ip)
+        temp_vect = vector.shape_cast(temp_ty, row_major_vec, loc=loc, ip=ip)
 
         red_dims = [i for i, x in enumerate(flat_prof) if x is not None]
 
@@ -1763,7 +1905,7 @@ class TensorSSA(cutlass_arith.ArithValue):
 
         # Slice and keep dims matching `_` or None
         res_shp = slice_(self.shape, reduction_profile, loc=loc, ip=ip)
-        return self._build_result(res_vect, res_shp, loc=loc, ip=ip)
+        return self._build_result(res_vect, res_shp, row_major=True, loc=loc, ip=ip)
 
 
 @dsl_user_op

python/CuTeDSL/cutlass/cute/testing.py

@@ -1057,7 +1057,7 @@ def tune(
         except NotImplementedError as e:
             logger.info(f"   Encountered unimplemented error, abort execution: {e}")
             raise e
-        except (ValueError, TypeError) as e:
+        except (ValueError, TypeError, CantImplementError) as e:
             logger.info(f"   Configuration parameter skipping: {e}")
             continue
         except Exception as e:
@@ -1073,3 +1073,17 @@ def tune(
     logger.info(f"Best configuration: {best_config}, execution time: {min_time} us")
     logger.info(f"Total tuning time: {tuning_time} s")
     return best_config
+
+
+class CantImplementError(Exception):
+    """Exception raised when a function is not implemented."""
+
+    def __init__(self, message=None):
+        self.message = message or "The current config is invalid/unsupported"
+        super().__init__(self.message)
+
+    def __str__(self):
+        return self.message
+
+    def __repr__(self):
+        return self.message

python/CuTeDSL/cutlass/cute/typing.py

@@ -10,17 +10,79 @@
 # is strictly prohibited.
 
 from abc import ABC, abstractmethod
-from typing import ForwardRef, Tuple, Union, Any, Type, List, Optional
+import ctypes
+from typing import ForwardRef, Tuple, Union, Any, Type, List, Optional, Literal
+from functools import lru_cache
 
 from cutlass.base_dsl.typing import *
 
 from cutlass._mlir import ir
 from cutlass._mlir.dialects.cute import AddressSpace, ConstrainedIntType
+from cutlass.base_dsl.typing import JitArgument
 
 
 Int = Union[int, Integer]
 
 
+class SymInt:
+    def __init__(self, width: Literal[32, 64] = 32, *, divisibility=1):
+        if width not in [32, 64]:
+            raise ValueError(f"Unsupported width: {width}")
+        self._width = width
+        self._divisibility = divisibility
+
+    @property
+    def width(self):
+        return self._width
+
+    @property
+    def divisibility(self):
+        return self._divisibility
+
+    def __str__(self) -> str:
+        return f"?{{i{self._width} div={self._divisibility}}}"
+
+    def __repr__(self) -> str:
+        return self.__str__()
+
+    def __eq__(self, other) -> bool:
+        if not isinstance(other, SymInt):
+            return False
+        return all(
+            [self._width == other._width, self._divisibility == other._divisibility]
+        )
+
+    def __c_pointers__(self):
+        return [ctypes.c_void_p(0).value]
+
+    def __get_mlir_types__(self) -> List[ir.Type]:
+        res_ty = ir.Type.parse(
+            f'!cute.int_tuple<"?{{i{self.width} div={self.divisibility}}}">'
+        )
+        return [res_ty]
+
+    def __new_from_mlir_values__(self, values) -> "SymInt":
+        from .core import IntValue
+
+        if self.width == 32:
+            return Int32(IntValue(values[0]))
+        elif self.width == 64:
+            return Int64(IntValue(values[0]))
+        else:
+            assert False, f"Unsupported width: {self.width}"
+            return self
+def sym_int(width: Literal[32, 64] = 32, *, divisibility=1) -> SymInt:
+    return SymInt(width, divisibility=divisibility)
+
+
+def sym_int32(divisibility=1) -> SymInt:
+    return sym_int(32, divisibility=divisibility)
+
+
+def sym_int64(divisibility=1) -> SymInt:
+    return sym_int(64, divisibility=divisibility)
+
+
 ScaledBasis = ForwardRef("ScaledBasis")
 
 IntTuple = Union[Int, Tuple["IntTuple", ...]]
@@ -288,7 +350,10 @@ def is_int_tuple(a) -> bool:
 
 
 __all__ = [
-    "Coord",
+    "SymInt",
+    "sym_int",
+    "sym_int32",
+    "sym_int64",
     "Numeric",
     "Integer",
     "Boolean",
@@ -317,13 +382,14 @@ __all__ = [
     "Float6E3M2FN",
     "IntTuple",
     "Layout",
-    "Pointer",
+    "Coord",
     "Shape",
     "Stride",
-    "Tensor",
     "Tile",
     "Tiler",
     "XTuple",
     "is_integer",
     "is_int_tuple",
+    "Pointer",
+    "Tensor",
 ]

python/CuTeDSL/cutlass/cutlass_dsl/cuda_jit_executor.py

@@ -0,0 +1,333 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 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/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.
+
+"""
+This module provides jit executor related classes for CUTLASS.
+"""
+import ctypes
+import functools
+import weakref
+import threading
+
+import cuda.bindings.runtime as cuda_runtime
+import cuda.bindings.driver as cuda_driver
+
+# Local modules imports
+from ..base_dsl.jit_executor import JitExecutor, ExecutionArgs, JitFunctionArtifacts
+from ..base_dsl.utils.logger import log
+from ..base_dsl.common import DSLCudaRuntimeError, DSLRuntimeError
+from ..base_dsl.typing import Int32
+
+class CudaDialectJitModule:
+    """Holds the execution engine and cuda libraries."""
+
+    def __init__(
+        self,
+        engine,
+        capi_func,
+        args_spec: ExecutionArgs,
+        cuda_library: list["cuda_runtime.cudaLibrary_t"],
+    ):
+        self.engine = engine
+        self.capi_func = capi_func
+        self.args_spec = args_spec
+        self.cuda_library = cuda_library
+        self._unloaded = False
+
+    def is_unloaded(self):
+        return self._unloaded
+
+    def unload(self):
+        try:
+            for library in self.cuda_library:
+                cuda_runtime.cudaLibraryUnload(library)
+            self.cuda_library.clear()
+        finally:
+            self._unloaded = True
+
+    def __del__(self):
+        self.unload()
+
+
+class CudaDialectJitCompiledFunction:
+    """Holds a compiled function and its module."""
+
+    def __init__(
+        self,
+        ir_module,
+        engine,
+        capi_func,
+        args_spec,
+        function_name,
+        kernel_info,
+        jit_time_profiling,
+        jit_function_artifacts,
+        prefix=None,
+        load_from_binary=False,
+    ):
+        self.ir_module = ir_module
+        self.engine = engine
+        self.capi_func = capi_func
+        self.function_name = function_name
+        self.kernel_info = kernel_info
+        if args_spec is not None:
+            self.args_spec = ExecutionArgs(args_spec, self.function_name)
+        self.jit_time_profiling = jit_time_profiling
+        self.prefix = prefix
+        assert (
+            isinstance(jit_function_artifacts, JitFunctionArtifacts)
+            or jit_function_artifacts is None
+        )
+        self.artifacts = jit_function_artifacts
+        self.load_from_binary = load_from_binary
+
+        # Set cuda result return type.
+        # When execution engine/capi function is None, do not set the return type.
+        if self.capi_func:
+            self.capi_func.restype = ctypes.c_int32
+        if self.args_spec:
+            self.args_spec.args_spec.annotations["return"] = Int32
+
+        # This runtime state is stored here so that we can preserve the module
+        # in the compiler cache. Callers can extend the lifetime of the module
+        # by creating and retaining the executor.
+        self.jit_module = None
+        self._executor_lock = threading.RLock()
+        self._default_executor = None
+
+    @property
+    def __ptx__(self):
+        """Returns the PTX code of the JIT-compiled function."""
+        return self.artifacts.PTX if self.artifacts is not None else None
+
+    @property
+    def __cubin__(self):
+        """Returns the CUBIN data of the JIT-compiled function."""
+        return self.artifacts.CUBIN if self.artifacts is not None else None
+
+    @property
+    def __mlir__(self):
+        """Returns the MLIR code of the JIT-compiled function."""
+        return self.artifacts.MLIR if self.artifacts is not None else None
+
+    @functools.cached_property
+    def num_devices(self):
+        """Returns the number of CUDA devices available."""
+        dev_err, devs = cuda_runtime.cudaGetDeviceCount()
+        if dev_err != cuda_runtime.cudaError_t.cudaSuccess:
+            raise DSLCudaRuntimeError(dev_err, cuda_runtime.cudaGetErrorName(dev_err))
+        return devs
+
+    def _deserializer(self):
+        """Load the cuda library from the binary execution engine.
+        @return: The list of cuda kernels.
+        """
+        library = ctypes.c_void_p()
+        pointer_to_library = ctypes.pointer(library)
+        pointer_to_pointer_to_library = ctypes.pointer(pointer_to_library)
+        err = ctypes.c_int32(0)
+        pointer_to_err = ctypes.pointer(err)
+
+        # cuda init takes in a pointer to a cudaLibrary_t and returns
+        # a i32 cudaError_t. It initialized (lazy loads) our cudaLibrary_t
+        cuda_init = self.engine.lookup(f"_mlir_{self.prefix}_cuda_init")
+        if cuda_init is None:
+            raise DSLRuntimeError("cuda_init not found")
+        cuda_init = ctypes.CFUNCTYPE(None, ctypes.c_void_p)(cuda_init)
+        # cuda load takes in a pointer to a cudaLibrary_t and returns
+        # a i32 cudaError_t. It loads the functions from the cuda library,
+        # sets function attributes, and returns an error if encountered.
+        cuda_load = self.engine.lookup(f"_mlir_{self.prefix}_cuda_load")
+        if cuda_load is None:
+            raise DSLRuntimeError("cuda_load not found")
+        cuda_load = ctypes.CFUNCTYPE(None, ctypes.c_void_p)(cuda_load)
+
+        cuda_init_args = [pointer_to_pointer_to_library, pointer_to_err]
+        packed_args = (ctypes.c_void_p * len(cuda_init_args))()
+        for i in range(len(cuda_init_args)):
+            packed_args[i] = ctypes.cast(cuda_init_args[i], ctypes.c_void_p)
+        cuda_init(packed_args)
+
+        if err.value != 0:
+            error_code = err.value
+            error_name = cuda_runtime.cudaGetErrorName(
+                cuda_runtime.cudaError_t(error_code)
+            )
+            raise DSLCudaRuntimeError(error_code, error_name)
+
+        cuda_load_args = [pointer_to_library, pointer_to_err]
+        packed_args = (ctypes.c_void_p * len(cuda_load_args))()
+        for i in range(len(cuda_load_args)):
+            packed_args[i] = ctypes.cast(cuda_load_args[i], ctypes.c_void_p)
+        cuda_load(packed_args)
+
+        if err.value != 0:
+            error_code = err.value
+            error_name = cuda_runtime.cudaGetErrorName(
+                cuda_runtime.cudaError_t(error_code)
+            )
+            raise DSLCudaRuntimeError(error_code, error_name)
+
+        return [cuda_runtime.cudaLibrary_t(library.value)]
+
+    def _get_cuda_init_and_load(self):
+        """Returns the cuda init and load functions from the engine."""
+        # cuda init takes in a pointer to a cudaLibrary_t and returns
+        # a i32 cudaError_t. It initialized (lazy loads) our cudaLibrary_t
+
+        cuda_init = None
+
+        # cuda load for device takes in a pointer to a cudaLibrary_t and a device index and returns
+        # a i32 cudaError_t. It resolves each kernel from the cuda library,
+        # and applies device scoped attributes for the given device id, and returns an error if encountered.
+        cuda_load_to_device = None
+
+        # When load_from_binary is true, the symbols are prefixed by _mlir_<prefix>_ and are looked
+        # up from the JIT engine. Otherwise we look for the unprefixed forms. Looking up cuda_init
+        # and cuda_load_to_device from the engine which are defined in CudaToLLVM.
+        if self.load_from_binary:
+            if self.prefix is None:
+                raise DSLRuntimeError("prefix is required to be set for binary loading")
+            cuda_init = self.engine.lookup(f"_mlir_{self.prefix}_cuda_init")
+            if cuda_init is None:
+                raise DSLRuntimeError(f"cuda_init not found for prefix {self.prefix}")
+            cuda_load_to_device = self.engine.lookup(
+                f"_mlir_{self.prefix}_cuda_load_to_device"
+            )
+            if cuda_load_to_device is None:
+                raise DSLRuntimeError(
+                    f"cuda_load_to_device not found for prefix {self.prefix}"
+                )
+        else:
+            cuda_init = self.engine.raw_lookup("cuda_init")
+            if cuda_init is None:
+                raise DSLRuntimeError("cuda_init not found")
+            cuda_load_to_device = self.engine.raw_lookup("cuda_load_to_device")
+            if cuda_load_to_device is None:
+                raise DSLRuntimeError("cuda_load_to_device not found")
+        cuda_init = ctypes.CFUNCTYPE(ctypes.c_int32, ctypes.c_void_p)(cuda_init)
+        cuda_load_to_device = ctypes.CFUNCTYPE(ctypes.c_int32, ctypes.c_void_p)(
+            cuda_load_to_device
+        )
+
+        return cuda_init, cuda_load_to_device
+
+    def _load_cuda_library(self):
+        """Loads the CUDA library from the engine."""
+
+        cuda_init, cuda_load_to_device = self._get_cuda_init_and_load()
+
+        library = ctypes.c_void_p()
+        pointer_to_library = ctypes.pointer(library)
+        pointer_to_pointer_to_library = ctypes.pointer(pointer_to_library)
+        err = ctypes.c_int32(0)
+        pointer_to_err = ctypes.pointer(err)
+
+        cuda_init_args = [pointer_to_pointer_to_library, pointer_to_err]
+        packed_args = (ctypes.c_void_p * len(cuda_init_args))()
+        for i in range(len(cuda_init_args)):
+            packed_args[i] = ctypes.cast(cuda_init_args[i], ctypes.c_void_p)
+        cuda_init(packed_args)
+
+        if err.value != 0:
+            error_code = err.value
+            error_name = cuda_runtime.cudaGetErrorName(
+                cuda_runtime.cudaError_t(error_code)
+            )
+            raise DSLCudaRuntimeError(error_code, error_name)
+
+        device_id = ctypes.c_int32(0)
+        pointer_to_device_id = ctypes.pointer(device_id)
+
+        cuda_load_args = [pointer_to_library, pointer_to_device_id, pointer_to_err]
+        packed_args = (ctypes.c_void_p * len(cuda_load_args))()
+        for i, arg in enumerate(cuda_load_args):
+            packed_args[i] = ctypes.cast(arg, ctypes.c_void_p)
+
+        for dev in range(self.num_devices):
+            device_id.value = dev
+            cuda_load_to_device(packed_args)
+            if err.value != 0:
+                raise DSLCudaRuntimeError(
+                    err.value,
+                    cuda_runtime.cudaGetErrorName(cuda_runtime.cudaError_t(err.value)),
+                )
+
+        if err.value != 0:
+            error_code = err.value
+            error_name = cuda_runtime.cudaGetErrorName(
+                cuda_runtime.cudaError_t(error_code)
+            )
+            raise DSLCudaRuntimeError(error_code, error_name)
+
+        return [cuda_runtime.cudaLibrary_t(library.value)]
+
+    def to(self, device=None) -> JitExecutor:
+        """Returns an executable function bound to the given device.
+
+        For multi-device execution this method can be called for each device where
+        the kernel will run.
+
+        Since CudaJitCompiledFunction uses CUDA libraries, which are context free,
+        binding to a device is not necessary and the device is ignored. Device is
+        kept in for compatibility with the JitCompiledFunction.
+
+        :param device: Specifies the device for the executor. If None the current device is used.
+        :type device: Optional[Union[int, CUdevice]]
+        :return: A callable executor function.
+        :rtype: JitExecutor
+        """
+        with self._executor_lock:
+            # We need to ensure that the modules are loaded if not already
+            if self.jit_module is None or self.jit_module.is_unloaded():
+                cuda_library = self._load_cuda_library()
+                self.jit_module = CudaDialectJitModule(
+                    self.engine,
+                    self.capi_func,
+                    self.args_spec,
+                    cuda_library,
+                )
+
+            return JitExecutor(self.jit_module, None, self.jit_time_profiling)
+
+    def generate_execution_args(self, *args, **kwargs):
+        return self.args_spec.generate_execution_args(args, kwargs)
+
+    def set_dynamic_args(self, dynamic_args, dynamic_kwargs):
+        """Sets the dynamic argument information required for export to c code generation."""
+        self.dynamic_args = dynamic_args
+        self.dynamic_kwargs = dynamic_kwargs
+
+    def __call__(self, *args, **kwargs):
+        """Executes the jit-compiled function under the currently active CUDA context.
+
+        Calling this method multiple devices is not allowed and will result in unexpected
+        CUDA errors. If you need to call the kernel on multiple devices use `to`
+        to return a per-device function.
+        """
+        exe_args, adapted_args = self.generate_execution_args(*args, **kwargs)
+        return self.run_compiled_program(exe_args)
+
+    def run_compiled_program(self, exe_args):
+        """Executes the jit-compiled function under the currently active CUDA context.
+
+        Calling this method multiple devices is not allowed and will result in unexpected
+        CUDA errors. If you need to call the kernel on multiple devices use `to`
+        to return a per-device function.
+        """
+        with self._executor_lock:
+            if self._default_executor is None:
+                log().debug("Creating default executor.")
+                # We use a weak reference here so that this instance does not keep this
+                # object alive as it hold a reference to self.
+                proxy_self = weakref.proxy(self)
+                self._default_executor = proxy_self.to(None)
+        return self._default_executor.run_compiled_program(exe_args)

python/CuTeDSL/cutlass/cutlass_dsl/cuda_stream_adapter.py

@@ -0,0 +1,44 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 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/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.
+
+"""
+This module provides CUDA Python helper functions
+"""
+
+import cuda.bindings.driver as cuda_driver
+
+# MLIR imports
+from .._mlir import ir
+from .._mlir.dialects import cuda
+
+# Local module imports
+from ..base_dsl.runtime.jit_arg_adapters import JitArgAdapterRegistry
+
+
+@JitArgAdapterRegistry.register_jit_arg_adapter(cuda_driver.CUstream)
+class CudaDialectStreamAdapter:
+    """
+    Convert a CUDA stream to a stream representation for JIT arg generation.
+    """
+
+    def __init__(self, arg):
+        self._arg = arg
+        self._c_pointer = self._arg.getPtr()
+
+    def __new_from_mlir_values__(self, values):
+        assert len(values) == 1
+        return values[0]
+
+    def __c_pointers__(self):
+        return [self._c_pointer]
+
+    def __get_mlir_types__(self):
+        return [cuda.StreamType.get()]

python/CuTeDSL/cutlass/cutlass_dsl/cutlass.py

@@ -43,11 +43,21 @@ from ..base_dsl.typing import *
 from ..base_dsl.typing import DynamicExpression, get_mlir_types
 from ..base_dsl.runtime.jit_arg_adapters import is_arg_spec_constexpr
 from ..base_dsl.runtime import cuda as cuda_helpers
+from .cuda_stream_adapter import CudaDialectStreamAdapter
 
+from .cuda_jit_executor import CudaDialectJitCompiledFunction
 
 # MLIR Imports
 from cutlass._mlir import ir, execution_engine, passmanager
-from cutlass._mlir.dialects import arith, func, gpu, scf, cute, gpu as cutlass_gpu
+from cutlass._mlir.dialects import (
+    arith,
+    func,
+    gpu,
+    scf,
+    cute,
+    gpu as cutlass_gpu,
+    cuda as cuda_dialect,
+)
 from cutlass._mlir.dialects._ods_common import (
     get_op_result_or_op_results as _get_op_result_or_op_results,
 )
@@ -220,6 +230,10 @@ class CutlassBaseDSL(BaseDSL):
             preprocess=preprocess,
         )
         self._smem_usage_tracker: tuple = None
+        # extra function to convert cute arguments to tvm ffi spec params
+        # this needs to be reverse registered because the arg convention
+        # depends on the runtime type of the DSL arguments
+        self._tvm_ffi_args_spec_converter = None
 
     # this method is not useful for cutlass_dsl, so we only provide a dummy implementation.
     def _is_tensor_descriptor(self, maybe_tensor_descriptor) -> bool:
@@ -241,7 +255,7 @@ class CutlassBaseDSL(BaseDSL):
 
     def _get_pipeline(self, pipeline):
         pipeline = super()._get_pipeline(pipeline)
-        if pipeline == None:
+        if pipeline is None:
             # cubin format is required to be cubin as we launch cuda module at python level.
             return (
                 "builtin.module(cute-to-nvvm{cubin-format=bin "
@@ -253,7 +267,10 @@ class CutlassBaseDSL(BaseDSL):
 
     def preprocess_pipeline(self, pipeline, arch) -> str:
         pipeline = super().preprocess_pipeline(pipeline, arch)
-        pipeline = pipeline.rstrip(")") + ",external-kernel-for-gpu-launch)"
+        pipeline = (
+            pipeline.rstrip("})")
+            + " enable-cuda-dialect=true cuda-dialect-external-module=true})"
+        )
         return pipeline
 
     def _enter_gpu_module(self):
@@ -265,13 +282,35 @@ class CutlassBaseDSL(BaseDSL):
         )
 
         ret = {}
-        # generate launch bound attr from LaunchConfig
-        max_threads = ", ".join(map(str, config.block))
-        ret["nvvm.reqntid"] = ir.Attribute.parse(f"array<i32 : {max_threads}>")
-        # min_blocks_per_mp is optional for kernel
-        min_blocks = config.min_blocks_per_mp
-        if min_blocks > 0:
-            ret["nvvm.minctasm"] = ir.Attribute.parse(f"{min_blocks} : i32")
+        if config.has_max_number_threads():
+            block_str = ", ".join(map(str, config.block))
+            has_dynamic = any(is_dynamic_expression(dim) for dim in config.block)
+            if not has_dynamic:
+                # Case 1. Auto-generate reqntid from block size
+                ret["nvvm.reqntid"] = ir.Attribute.parse(f"array<i32 : {block_str}>")
+            else:
+                # Case 2. Don't auto-generate reqntid from block size
+                self.print_warning_once(
+                    f"Dynamic variable in block size {block_str}, cannot auto-generate `nvvm.reqntid`"
+                )
+        else:
+            # Case 3. Use user-defined value for maxntid
+            max_ntid_str = ", ".join(map(str, config.max_number_threads))
+            for value in config.max_number_threads:
+                if is_dynamic_expression(value):
+                    raise DSLRuntimeError(
+                        f"User defined max number threads `{max_ntid_str}` contains dynamic expression, cannot generate `nvvm.maxntid`",
+                        suggestion="Consider using `Constexpr` annotation or Python constant",
+                    )
+            # Use user-defined maxntid
+            ret["nvvm.maxntid"] = ir.Attribute.parse(f"array<i32 : {max_ntid_str}>")
+
+        # Add optional minimum blocks per multiprocessor
+        if config.min_blocks_per_mp > 0:
+            ret["nvvm.minctasm"] = ir.Attribute.parse(
+                f"{config.min_blocks_per_mp} : i32"
+            )
+
         return ret
 
     @lru_cache(maxsize=1)
@@ -313,6 +352,116 @@ class CutlassBaseDSL(BaseDSL):
 
         return version_hash
 
+    def get_return_types(self) -> List[ir.Type]:
+        """
+        Get the return types of the function.
+        With cuda dialect, the return type is i32 to indicate cuda result.
+        """
+        i32_ty = ir.IntegerType.get_signless(32)
+        return [i32_ty]
+
+    def generate_default_return_values(self, ip=None) -> List[ir.Value]:
+        """
+        Generate the default return values of the function.
+        With cuda dialect, the default return value is 0 to indicate success.
+        """
+        if ip is None:
+            raise DSLRuntimeError("ip is required to generate default return values")
+        return_types = self.get_return_types()
+        if len(return_types) != 1:
+            raise DSLRuntimeError(
+                f"Expected exactly one return type, but got {len(return_types)}"
+            )
+        return_type = return_types[0]
+        with ip:
+            zero = arith.ConstantOp(return_type, 0).result
+        return [zero]
+
+    def compile_and_cache(
+        self,
+        module,
+        module_hash,
+        function_name,
+        pipeline,
+        args_spec,
+        no_cache,
+        *,
+        dynamic_args=None,
+        dynamic_kwargs=None,
+        original_function_name=None,
+    ):
+        """
+        Compile the module and cache the result.
+
+        :param module: The MLIR module to compile.
+        :param module_hash: The hash of the MLIR module.
+        :param function_name: The name of the function to compile.
+        :param pipeline: The pipeline to use for compilation.
+        :param args_spec: The args spec to use for compilation.
+        :param no_cache: Whether to cache the result.
+        :param dynamic_args: The dynamic arguments to use for compilation.
+        :param dynamic_kwargs: The dynamic keyword arguments to use for compilation.
+        :param original_function_name: The name of the original function without mangling.
+        :return: The compiled function.
+        """
+        if self.compile_options.enable_tvm_ffi:
+            # TVM FFI post compile logic
+            # attach extra ABI function to the MLIR module
+            from .tvm_ffi_provider import (
+                TVMFFIJitCompiledFunction,
+                TVMFFICuteCallProvider,
+            )
+            from cutlass.base_dsl.tvm_ffi_builder import attach_ffi_func
+
+            assert self._tvm_ffi_args_spec_converter is not None
+            tvm_ffi_spec_params = self._tvm_ffi_args_spec_converter(
+                function_name, args_spec, dynamic_args, dynamic_kwargs
+            )
+            tvm_ffi_provider = TVMFFICuteCallProvider(function_name)
+
+            # ensure we run the postprocessor hook after the compiler has run its passes
+            def post_compile_hook(module: ir.Module):
+                with module.context, module.operation.location:
+                    # attach the tvm ffi function to the mlir module
+                    attach_ffi_func(
+                        module,
+                        function_name,
+                        tvm_ffi_spec_params,
+                        tvm_ffi_provider,
+                        fn_display_name=original_function_name,
+                    )
+                module.operation.verify()
+
+            # ensure the compiler can run post-compile hook after its passes
+            # the context will restore the previous post-compile hook after it exits
+            with compiler.PostCompileHookContext(
+                self.compiler_provider, post_compile_hook
+            ):
+                return super().compile_and_cache(
+                    module,
+                    module_hash,
+                    function_name,
+                    pipeline,
+                    args_spec,
+                    no_cache,
+                    TVMFFIJitCompiledFunction,
+                    dynamic_args=dynamic_args,
+                    dynamic_kwargs=dynamic_kwargs,
+                )
+
+        return super().compile_and_cache(
+            module,
+            module_hash,
+            function_name,
+            pipeline,
+            args_spec,
+            no_cache,
+            CudaDialectJitCompiledFunction,
+            dynamic_args=dynamic_args,
+            dynamic_kwargs=dynamic_kwargs,
+            original_function_name=original_function_name,
+        )
+
     @staticmethod
     def track_smem_allocator(allocator, callback):
         """
@@ -343,6 +492,103 @@ class CutlassBaseDSL(BaseDSL):
         allocator, callback = self._smem_usage_tracker
         return callback(allocator)
 
+    @staticmethod
+    def cuda_launch_func(
+        stream: Union[list, tuple],
+        kernel,
+        grid_size_x,
+        grid_size_y,
+        grid_size_z,
+        block_size_x,
+        block_size_y,
+        block_size_z,
+        kernel_operands,
+        *,
+        cluster_size_x=None,
+        cluster_size_y=None,
+        cluster_size_z=None,
+        dynamic_shared_memory_size=None,
+        use_pdl=False,
+        loc=None,
+        ip=None,
+    ):
+
+        # Max number of attributes in the launch config is set to 2 for PDL and cluster size
+        max_num_attributes = 2
+        launch_config_type = cuda_dialect.LaunchConfigType.get(max_num_attributes)
+
+        if len(stream) == 0:
+            stream = cuda_dialect.CastOp(
+                src=Int64(0).ir_value(),
+                dst=cuda_dialect.StreamType.get(),
+                loc=loc,
+                ip=ip,
+            )
+        else:
+            stream = stream[0]
+
+        dynamic_shared_memory_size = Int64(dynamic_shared_memory_size).ir_value()
+        block_size_x = Int32(block_size_x).ir_value()
+        block_size_y = Int32(block_size_y).ir_value()
+        block_size_z = Int32(block_size_z).ir_value()
+        grid_size_x = Int32(grid_size_x).ir_value()
+        grid_size_y = Int32(grid_size_y).ir_value()
+        grid_size_z = Int32(grid_size_z).ir_value()
+
+        cfg = cuda_dialect.launch_cfg_create(
+            # Launch config type
+            launch_config_type,
+            # Max num of attributes the launch config can hold
+            # set to 2 for PDL and cluster size
+            ir.IntegerAttr.get(ir.IntegerType.get_signless(32), max_num_attributes),
+            block_size_x,
+            block_size_y,
+            block_size_z,
+            dynamic_shared_memory_size,
+            grid_size_x,
+            grid_size_y,
+            grid_size_z,
+            stream,
+            loc=loc,
+            ip=ip,
+        )
+
+        if use_pdl:
+            cuda_dialect.launch_cfg_programmatic_stream_serialization_allowed(
+                cfg, Int32(use_pdl).ir_value(), loc=loc, ip=ip
+            )
+
+        if cluster_size_x is not None:
+            if cluster_size_y is None:
+                cluster_size_y = 1
+            if cluster_size_z is None:
+                cluster_size_z = 1
+            cluster_size_x = Int32(cluster_size_x).ir_value(loc=loc, ip=ip)
+            cluster_size_y = Int32(cluster_size_y).ir_value(loc=loc, ip=ip)
+            cluster_size_z = Int32(cluster_size_z).ir_value(loc=loc, ip=ip)
+            cuda_dialect.launch_cfg_cluster_dim(
+                cfg, cluster_size_x, cluster_size_y, cluster_size_z, loc=loc, ip=ip
+            )
+
+        op = cuda_dialect.launch_ex(
+            cuda_dialect.ResultType.get(),
+            kernel,
+            cfg,
+            kernel_operands,
+            loc=loc,
+            ip=ip,
+        )
+
+        res_i32 = cuda_dialect.CastOp(
+            src=op,
+            dst=cuda_dialect.IntegerType.get_signless(32),
+            loc=loc,
+            ip=ip,
+        )
+
+        cuda_dialect.return_if_error([res_i32], loc=loc, ip=ip)
+        return None
+
     @staticmethod
     def gpu_launch_func(
         async_token,
@@ -396,20 +642,32 @@ class CutlassBaseDSL(BaseDSL):
 
             def generate_func_op(self, arg_types, arg_attrs, kernel_name, loc=None):
                 super().generate_func_op(arg_types, arg_attrs, kernel_name)
-                self.func_op = func.FuncOp(
+                self.func_op = cuda_dialect.KernelOp(
                     kernel_name, ir.FunctionType.get(arg_types, []), loc=loc
                 )
+                self.arg_types = arg_types
+                self.func_op.attributes["cu_attrs"] = ir.DictAttr.get(
+                    {
+                        str(
+                            cuda_dialect.CUFunctionAttribute.non_portable_cluster_size_allowed
+                        ): ir.IntegerAttr.get(ir.IntegerType.get_signless(32), 1),
+                        str(
+                            cuda_dialect.CUFunctionAttribute.max_dynamic_shared_size_bytes
+                        ): cuda_dialect.DevMaxSharedMemoryOptinAttr.get(),
+                    }
+                )
                 if arg_attrs is not None:
                     log().debug(arg_attrs)
                     self.func_op.arg_attrs = arg_attrs
                 return self.func_op
 
-            def generate_func_ret_op(self):
-                return func.ReturnOp([])
+            def generate_func_ret_op(self, loc=None, ip=None):
+                return cuda_dialect.ReturnOp([], loc=loc, ip=ip)
 
             def get_func_body_start(self):
                 assert self.func_op is not None, "Invalid func_op is not expected!"
-                return self.func_op.add_entry_block()
+                arg_locs = [self.func_op.operation.location for _ in self.arg_types]
+                return self.func_op.add_entry_block(arg_locs=arg_locs)
 
             def generate_launch_op(self, *args, **kwargs):
                 # Extract args and do validation
@@ -430,12 +688,6 @@ class CutlassBaseDSL(BaseDSL):
 
                 cfg = requiredArgs.config
 
-                # Apply to grid, block, and cluster if present
-                cfg.grid = [to_index(size) for size in cfg.grid]
-                cfg.block = [to_index(size) for size in cfg.block]
-                if cfg.has_cluster:
-                    cfg.cluster = [to_index(size) for size in cfg.cluster]
-
                 smem_usage = self.dsl._get_smem_usage()
                 if any(not isinstance(x, int) for x in [cfg.smem, smem_usage]):
                     pass  # cannot compare dynamic value inside kernel to launch op in py
@@ -449,12 +701,11 @@ class CutlassBaseDSL(BaseDSL):
                     )
                 cfg.smem = const(cfg.smem)
 
+                async_deps = cfg.async_deps
                 if not isinstance(cfg.async_deps, (list, tuple)):
-                    cfg.async_deps = [cfg.async_deps]
-                is_async = len(cfg.async_deps) > 0
-                token = CutlassBaseDSL.gpu_launch_func(
-                    gpu.AsyncTokenType.get() if is_async else None,
-                    cfg.async_deps,
+                    async_deps = [cfg.async_deps]
+                CutlassBaseDSL.cuda_launch_func(
+                    async_deps,
                     kernelSym,
                     *cfg.grid,
                     *cfg.block,
@@ -469,7 +720,7 @@ class CutlassBaseDSL(BaseDSL):
                     use_pdl=cfg.use_pdl,
                     loc=loc,
                 )
-                return token if is_async else None
+                return None
 
         return KernelLauncher(
             self,
@@ -526,6 +777,10 @@ class CutlassBaseDSL(BaseDSL):
                 # skip generic types such as List[int], Tuple[int, int], etc. for performance consideration?
                 pass
 
+            elif getattr(arg_annotation, "__name__", None) == "CUstream" and arg.__class__.__name__ == "_FakeStream":
+                # allow FakeStream to be passed as CUDA stream
+                pass
+
             elif isinstance(arg_annotation, type):
                 # Handle simple type annotations
                 if not isinstance(arg, arg_annotation) and arg is not None:
@@ -564,7 +819,23 @@ class CutlassBaseDSL(BaseDSL):
                     # Handle dynamic types
                     jit_arg_type.extend([v.type for v in dyn_vals])
                     jit_arg_attr.extend([default_attr] * len(dyn_vals))
-                    jit_exec_arg.extend(get_c_pointers(arg) if is_host else dyn_vals)
+                    if is_host and self.envar.enable_tvm_ffi:
+                        import tvm_ffi
+
+                        jit_exec_arg.extend(
+                            [
+                                tvm_ffi.Shape(
+                                    [
+                                        v.value if isinstance(v, Numeric) else v
+                                        for v in arg
+                                    ]
+                                )
+                            ]
+                        )
+                    else:
+                        jit_exec_arg.extend(
+                            get_c_pointers(arg) if is_host else dyn_vals
+                        )
                 else:
                     jit_exec_arg = jit_arg_type = jit_arg_attr = None
             elif not hasattr(arg, "__extract_mlir_values__") and not hasattr(
@@ -879,7 +1150,6 @@ def to_index(value):
 
     return res
 
-
 def _validate_iter_args_structure(iter_args, ir_values):
     """
     Validates that iter_args structure contains the same number of atomic values
@@ -938,6 +1208,15 @@ def _minmax(op, *args, loc=None, ip=None):
             emitter = getattr(cutlass_arith, f"_{op.__name__}")
             if not (is_dynamic_expression(res) or is_dynamic_expression(x)):
                 res = emitter(op(res), op(x))
+            elif (
+                hasattr(res, "type")
+                and hasattr(x, "type")
+                and isinstance(res.type, T.VectorType)
+                and isinstance(x.type, T.VectorType)
+            ):
+                lhs = op(res, loc=loc, ip=ip)
+                rhs = op(x, loc=loc, ip=ip)
+                return lhs.apply_op(emitter, rhs, loc=loc, ip=ip)
             else:
                 lhs = as_numeric(op(res, loc=loc, ip=ip))
                 rhs = as_numeric(op(x, loc=loc, ip=ip))

python/CuTeDSL/cutlass/cutlass_dsl/tvm_ffi_provider.py

@@ -0,0 +1,348 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 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/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.
+
+from cutlass.base_dsl.tvm_ffi_builder import (
+    DynamicParamPackCallProvider,
+    CallContext,
+    rename_tvm_ffi_function,
+    spec,
+)
+from cutlass._mlir import ir
+from cutlass._mlir.dialects import llvm
+from cutlass._mlir._mlir_libs import _execution_engine_extra
+from cutlass.cutlass_dsl.cuda_jit_executor import CudaDialectJitCompiledFunction
+from cutlass.base_dsl.common import DSLRuntimeError
+import tvm_ffi
+
+
+class TVMFFICuteCallProvider(DynamicParamPackCallProvider):
+    """Cute call provider that uses cute call convention."""
+
+    def __init__(self, target_func: str):
+        super().__init__(target_func, struct_call=True)
+        self.cuda_global_state_symbol = f"__{target_func}_cuda_state"
+
+    def get_callee_struct_for_param_tensor(
+        self,
+        param: spec.Tensor,
+        current_block: ir.Block,
+        data: ir.Value,
+        shape: list[ir.Value],
+        strides: list[ir.Value],
+        flatten_struct: ir.Type,
+    ) -> ir.Type:
+        """Routine used to override the tensor passing struct convention"""
+        with ir.InsertionPoint(current_block):
+            data_type = self.gpu_ptr_type
+            strides_type = (
+                self.struct_type(fields=[x.type for x in strides])
+                if len(strides) != 1
+                else strides[0].type
+            )
+            shape_type = (
+                self.struct_type(fields=[x.type for x in shape])
+                if len(shape) != 1
+                else shape[0].type
+            )
+            shape_stride_tuple_type = self.struct_type(
+                fields=[shape_type, strides_type]
+            )
+            tensor_type = self.struct_type(fields=[data_type, shape_stride_tuple_type])
+            return tensor_type
+
+    def pack_param_shape(
+        self, current_block: ir.Block, context: CallContext, param: spec.Shape
+    ) -> tuple[tuple[ir.Type], tuple[ir.Value]]:
+        """Pack a shape parameter to a struct."""
+        allocas: list[ir.Value] = []
+        arg_types: list[ir.Type] = []
+        for dim in param.shape:
+            if isinstance(dim, spec.Var):
+                allocas.append(
+                    self.pack_values_to_alloca(
+                        current_block,
+                        context.entry_block,
+                        [context.matched_var_binding[dim]],
+                    )[1]
+                )
+                arg_types.append(context.matched_var_binding[dim].type)
+        return tuple(arg_types), tuple(allocas)
+
+    def declare_extern_funcs(self, current_block: ir.Block, context: CallContext):
+        """Append the error handling function to the current block."""
+        with ir.InsertionPoint(context.module.body):
+            self.declare_extern_func(
+                "cuda_dialect_get_error_name",
+                [self.i32_type],
+                self.ptr_type,
+            )
+            self.declare_extern_func(
+                "_cudaSetDevice",
+                [self.i32_type],
+                self.i32_type,
+            )
+            self.declare_extern_func(
+                "cuda_dialect_init_library_once",
+                [self.ptr_type, self.ptr_type, self.ptr_type, self.ptr_type],
+                self.i32_type,
+            )
+            self.declare_extern_func(
+                "cuda_dialect_unload_library_once",
+                [self.ptr_type],
+                self.void_type,
+            )
+        return current_block
+
+    def insert_lazy_init_cuda(self, current_block: ir.Block, context: CallContext):
+        """Insert the lazy init cuda function."""
+        # create global private static that is initialized to nullptr
+        with ir.InsertionPoint(context.module.body):
+            parsed_op = ir.Operation.parse(
+                f"llvm.mlir.global private @{self.cuda_global_state_symbol}(0 : i64) : i64"
+            )
+            context.module.body.append(parsed_op)
+
+
+        with ir.InsertionPoint(current_block):
+            cuda_global_state_ptr = self.address_of(
+                self.cuda_global_state_symbol, self.ptr_type
+            )
+            cuda_init_ptr = self.address_of("cuda_init", self.ptr_type)
+            cuda_load_ptr = self.address_of("cuda_load", self.ptr_type)
+            set_error_ptr = self.address_of(
+                "TVMFFIErrorSetRaisedFromCStr", self.ptr_type
+            )
+            # Call the callback function with the loaded ptr value
+            init_result = llvm.call(
+                result=self.i32_type,  # function returns i32
+                callee="cuda_dialect_init_library_once",
+                callee_operands=[
+                    cuda_global_state_ptr,
+                    cuda_init_ptr,
+                    cuda_load_ptr,
+                    set_error_ptr,
+                ],
+                op_bundle_sizes=[],
+                op_bundle_operands=[],
+            )
+            # Create blocks for conditional branching
+            error_block = current_block.create_after()
+            success_block = error_block.create_after()
+            # Check if initialization failed (non-zero return code)
+            llvm.cond_br(
+                self.equal(init_result, self.i32(0)),
+                true_dest_operands=[],
+                false_dest_operands=[],
+                true_dest=success_block,
+                false_dest=error_block,
+            )
+        # Error block: return the error code
+        # error is already set by cuda_dialect_init_library_once
+        with ir.InsertionPoint(error_block):
+            llvm.return_(arg=self.i32(-1))
+
+        # Continue with success block
+        return success_block
+
+    def append_unload_to_global_dtors(
+        self, current_block: ir.Block, context: CallContext
+    ) -> ir.Block:
+        """Append the cuda_dialect_unload_library_once function to the global destructor list."""
+        unload_func_symbol = "cuda_dialect_unload_library_once"
+        # define a private function to call the extern function, we need this wrapper function
+        # since llvm.mlir.global_dtors require the dtor defined in the module
+        unload_func_wrapper_symbol = f"__dtor_{unload_func_symbol}"
+        with ir.InsertionPoint(context.module.body):
+            params, entry_block = self.function(
+                name=unload_func_wrapper_symbol,
+                params_type=[],
+                ret_type=self.void_type,
+                internal=True,
+            )
+            with ir.InsertionPoint(entry_block):
+                llvm.call(
+                    result=None,
+                    callee=unload_func_symbol,
+                    callee_operands=[
+                        self.address_of(self.cuda_global_state_symbol, self.ptr_type)
+                    ],
+                    op_bundle_sizes=[],
+                    op_bundle_operands=[],
+                )
+                llvm.return_()
+
+        # find or create the global destructors
+        global_dtors_list: list[ir.Operation] = self.find_operations_in_module(
+            context.module, "llvm.mlir.global_dtors"
+        )
+        if len(global_dtors_list) == 0:
+            # create the global destructors
+            with ir.InsertionPoint(context.module.body):
+                global_dtors = llvm.mlir_global_dtors(
+                    dtors=[],
+                    priorities=[],
+                )
+        else:
+            # use the existing global destructors
+            global_dtors = global_dtors_list[0]
+
+        # append the unload function to the global destructors
+        global_dtors.attributes["dtors"] += [
+            ir.FlatSymbolRefAttr.get(unload_func_wrapper_symbol)
+        ]
+        global_dtors.attributes["priorities"] += [
+            ir.IntegerAttr.get(self.i32_type, 65535)
+        ]  # the default priority
+
+        return current_block
+
+    def check_cuda_error(
+        self, code: ir.Value, current_block: ir.Block, context: CallContext
+    ):
+        """Check if the CUDA error is raised and return the error string if so."""
+        with ir.InsertionPoint(current_block):
+            # check if the call is successful
+            error_block = current_block.create_after()
+            success_block = error_block.create_after()
+            # Check if call is successful (non-zero return code)
+            self.cond_br(
+                cond=self.equal(code, self.i32(0)),
+                true_block=success_block,
+                false_block=error_block,
+                branch_weights=self.BRANCH_WEIGHTS_LIKELY,
+            )
+
+        # Error block: raise error and return
+        with ir.InsertionPoint(error_block):
+            error_str = llvm.call(
+                result=self.ptr_type,
+                callee="cuda_dialect_get_error_name",
+                callee_operands=[code],
+                op_bundle_sizes=[],
+                op_bundle_operands=[],
+            )
+            # Raise error and return -1
+            context.builder.raise_error_and_return(
+                error_kind="RuntimeError",
+                error_message_parts=["CUDA Error: ", error_str],
+            )
+        return success_block
+
+    def generate_llvm_call(
+        self,
+        current_block: ir.Block,
+        call_operands: list[ir.Value],
+        context: CallContext,
+    ) -> ir.Block:
+        """Generate the LLVM call operation and check if the call is successful."""
+        with ir.InsertionPoint(current_block):
+            result = llvm.call(
+                result=self.i32_type,
+                callee=self.target_func,
+                callee_operands=call_operands,
+                op_bundle_sizes=[],
+                op_bundle_operands=[],
+            )
+        return self.check_cuda_error(result, current_block, context)
+
+    def insert_set_cuda_device(self, current_block: ir.Block, context: CallContext):
+        """Call the _cudaSetDevice function if we can find device id from tensor parameters."""
+
+        def find_cuda_device_index_from_params():
+            for param in context.params:
+                if (
+                    isinstance(param, spec.Tensor)
+                    and param.dlpack_device_type != tvm_ffi.DLDeviceType.kDLCPU
+                ):
+                    return context.matched_var_binding[param.device_id]
+            return None
+
+        cuda_device_index = find_cuda_device_index_from_params()
+
+        if cuda_device_index is None:
+            return current_block
+
+        with ir.InsertionPoint(current_block):
+            result = llvm.call(
+                result=self.i32_type,
+                callee="_cudaSetDevice",
+                callee_operands=[cuda_device_index],
+                op_bundle_sizes=[],
+                op_bundle_operands=[],
+            )
+
+        return self.check_cuda_error(result, current_block, context)
+
+    def __call__(self, current_block: ir.Block, context: CallContext) -> ir.Block:
+        current_block = self.declare_extern_funcs(current_block, context)
+        current_block = self.insert_lazy_init_cuda(current_block, context)
+        current_block = self.append_unload_to_global_dtors(current_block, context)
+        current_block = self.insert_set_cuda_device(current_block, context)
+        current_block = super().__call__(current_block, context)
+        return current_block
+
+
+class TVMFFIJitCompiledFunction(tvm_ffi.Function, CudaDialectJitCompiledFunction):
+    """TVM FFI Function that contains metadata of the compiled function and interface to the FFI layer.
+
+    This function should not be directly used after
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        # initialize the tvm_ffi.Function from the current execution engine
+        self._init_ffi_function()
+
+    # use direct call to the tvm_ffi.Function.__call__
+    # to avoid most of python overhead
+    __call__ = tvm_ffi.Function.__call__
+
+    def _init_ffi_function(self):
+        """Initialize the tvm_ffi.Function from the current execution engine.
+
+        This function must be called at once during compilation time.
+        The reason why it is not called during init is because the original
+        flow may already created an execution engine and the function is not
+        guaranteed to be initialized at that time.
+        """
+        if self.__chandle__() != 0:
+            raise DSLRuntimeError("TVM FFI function is already initialized")
+        # get the MLIR function pointer from the execution engine
+        tvm_ffi_function_ptr = self.engine.raw_lookup("__tvm_ffi_" + self.function_name)
+        tvm_ffi_function = tvm_ffi.Function.__from_mlir_packed_safe_call__(
+            tvm_ffi_function_ptr
+        )
+        # move the handle from the tvm_ffi.Function to the current instance
+        self.__move_handle_from__(tvm_ffi_function)
+
+    def to(self, device=None):
+        """TVM FFI function itself is already support all devices."""
+        return self
+
+    def run_compiled_program(self, exe_args: list[ir.Value]):
+        """Run the compiled program. This override is needed for implicit compile and execution."""
+        return self.__call__(*exe_args)
+
+    def export_to_c(self, object_file_path: str, function_name: str = None):
+        """Export the TVM FFI function to an object file.
+
+        :param object_file_path: The path to the object file.
+        :param function_name: The name of the function to export.
+        """
+        if function_name is not None and function_name != self.function_name:
+            mod = self.ir_module
+            rename_tvm_ffi_function(mod, self.function_name, function_name)
+        else:
+            mod = self.ir_module
+
+        _execution_engine_extra.dump_object_file_pic(
+            mod, object_file_path, "__tvm_ffi_" + function_name, 2
+        )

python/CuTeDSL/cutlass/pipeline/__init__.py

@@ -20,7 +20,9 @@ from .helpers import (
     PipelineUserType,
     PipelineState,
     make_pipeline_state,
+    pipeline_init_arrive,
     pipeline_init_wait,
+    agent_sync,
     arrive,
     arrive_unaligned,
     wait,
@@ -68,7 +70,9 @@ __all__ = [
     "PipelineProducer",
     "PipelineConsumer",
     "make_pipeline_state",
+    "pipeline_init_arrive",
     "pipeline_init_wait",
+    "agent_sync",
     "arrive",
     "arrive_unaligned",
     "wait",

python/CuTeDSL/cutlass/pipeline/helpers.py

@@ -16,7 +16,14 @@ from typing import Optional, Union
 import warnings
 
 import cutlass.cute as cute
-from cutlass.cutlass_dsl import Boolean, Int32, Int64, if_generate
+from cutlass.cutlass_dsl import (
+    Boolean,
+    Int32,
+    Int64,
+    if_generate,
+    dsl_user_op,
+    dsl_user_op,
+)
 from cutlass._mlir.dialects import llvm
 import cutlass._mlir.dialects.cute as _cute_ir
 
@@ -39,6 +46,11 @@ class Agent(enum.Enum):
     ThreadBlockCluster = enum.auto()
 
 
+##############################################################################
+# CooperativeGroup class
+##############################################################################
+
+
 class CooperativeGroup:
     """
     CooperativeGroup contains size and alignment restrictions for an Agent.
@@ -74,6 +86,11 @@ class CooperativeGroup:
         self.agent = agent
 
 
+##############################################################################
+# PipelineOp class
+##############################################################################
+
+
 class PipelineOp(enum.Enum):
     """
     PipelineOp assigns an operation to an agent corresponding to a specific hardware feature.
@@ -134,6 +151,11 @@ class SyncObject(ABC):
         pass
 
 
+##############################################################################
+# MbarrierArray class
+##############################################################################
+
+
 class MbarrierArray(SyncObject):
     """
     MbarrierArray implements an abstraction for an array of smem barriers.
@@ -185,25 +207,33 @@ class MbarrierArray(SyncObject):
         return new_mbarrier_array
 
     # Mbarrier initialization
-    def mbarrier_init(self) -> None:
+    def mbarrier_init(self, *, loc=None, ip=None) -> None:
         """
         Initializes an array of mbarriers using warp 0.
         """
 
         def then_body():
             for index in range(self.num_stages):
-                cute.arch.mbarrier_init(self.get_barrier(index), self.arrive_count)
+                cute.arch.mbarrier_init(
+                    self.get_barrier(index, loc=loc, ip=ip),
+                    self.arrive_count,
+                    loc=loc,
+                    ip=ip,
+                )
 
-        warp_idx = cute.arch.warp_idx()
-        warp_idx = cute.arch.make_warp_uniform(warp_idx)
+        warp_idx = cute.arch.warp_idx(loc=loc, ip=ip)
+        warp_idx = cute.arch.make_warp_uniform(warp_idx, loc=loc, ip=ip)
 
-        if_generate(warp_idx == 0, then_body)
+        if_generate(warp_idx == 0, then_body, loc=loc, ip=ip)
 
     def arrive(
         self,
         index: int,
         dst: int,
         cta_group: Optional[cute.nvgpu.tcgen05.CtaGroup] = None,
+        *,
+        loc=None,
+        ip=None,
     ) -> None:
         """Select the arrive corresponding to this MbarrierArray's PipelineOp.
 
@@ -220,49 +250,75 @@ class MbarrierArray(SyncObject):
         :type cta_group: ``cute.nvgpu.tcgen05.CtaGroup``, optional
         """
         if self.op_type is PipelineOp.AsyncThread:
-            self.arrive_mbarrier(index, dst)
+            self.arrive_mbarrier(index, dst, loc=loc, ip=ip)
         elif self.op_type is PipelineOp.TCGen05Mma:
             assert cta_group is not None, (
                 "Error: CTA group must be provided for TCGen05Mma."
             )
-            self.arrive_tcgen05mma(index, dst, cta_group)
+            self.arrive_tcgen05mma(index, dst, cta_group, loc=loc, ip=ip)
         elif self.op_type in [PipelineOp.TmaLoad]:
             self.arrive_and_expect_tx(index, self.tx_count)
         elif self.op_type is PipelineOp.AsyncLoad:
-            self.arrive_cp_async_mbarrier(index)
+            self.arrive_cp_async_mbarrier(index, loc=loc, ip=ip)
         else:
             assert False, (
                 f"Error: MbarrierArray is not supported for PipelineOp: {_get_pipeline_op(self.op_type)}."
             )
 
-    def arrive_mbarrier(self, index: int, dst_rank: Optional[int] = None) -> None:
+    def arrive_mbarrier(
+        self, index: int, dst_rank: Optional[int] = None, *, loc=None, ip=None
+    ) -> None:
         if dst_rank is None:
-            cute.arch.mbarrier_arrive(self.get_barrier(index))
+            cute.arch.mbarrier_arrive(
+                self.get_barrier(index, loc=loc, ip=ip), loc=loc, ip=ip
+            )
         else:
-            cute.arch.mbarrier_arrive(self.get_barrier(index), dst_rank)
+            cute.arch.mbarrier_arrive(
+                self.get_barrier(index, loc=loc, ip=ip), dst_rank, loc=loc, ip=ip
+            )
 
-    def arrive_cp_async_mbarrier(self, index: int):
-        cute.arch.cp_async_mbarrier_arrive_noinc(self.get_barrier(index))
+    def arrive_cp_async_mbarrier(self, index: int, *, loc=None, ip=None):
+        cute.arch.cp_async_mbarrier_arrive_noinc(
+            self.get_barrier(index, loc=loc, ip=ip), loc=loc, ip=ip
+        )
 
     def arrive_tcgen05mma(
-        self, index: int, mask: Optional[int], cta_group: cute.nvgpu.tcgen05.CtaGroup
+        self,
+        index: int,
+        mask: Optional[int],
+        cta_group: cute.nvgpu.tcgen05.CtaGroup,
+        *,
+        loc=None,
+        ip=None,
     ) -> None:
         if mask is None:
-            with cute.arch.elect_one():
-                cute.nvgpu.tcgen05.commit(self.get_barrier(index))
+            with cute.arch.elect_one(loc=loc, ip=ip):
+                cute.nvgpu.tcgen05.commit(
+                    self.get_barrier(index, loc=loc, ip=ip), loc=loc, ip=ip
+                )
         else:
-            with cute.arch.elect_one():
-                cute.nvgpu.tcgen05.commit(self.get_barrier(index), mask, cta_group)
+            with cute.arch.elect_one(loc=loc, ip=ip):
+                cute.nvgpu.tcgen05.commit(
+                    self.get_barrier(index, loc=loc, ip=ip),
+                    mask,
+                    cta_group,
+                    loc=loc,
+                    ip=ip,
+                )
 
     def arrive_and_expect_tx(self, index: int, tx_count: int) -> None:
         with cute.arch.elect_one():
             cute.arch.mbarrier_arrive_and_expect_tx(self.get_barrier(index), tx_count)
 
-    def try_wait(self, index: int, phase: int) -> Boolean:
-        return cute.arch.mbarrier_try_wait(self.get_barrier(index), phase)
+    def try_wait(self, index: int, phase: int, *, loc=None, ip=None) -> Boolean:
+        return cute.arch.mbarrier_try_wait(
+            self.get_barrier(index, loc=loc, ip=ip), phase, loc=loc, ip=ip
+        )
 
-    def wait(self, index: int, phase: int) -> None:
-        cute.arch.mbarrier_wait(self.get_barrier(index), phase)
+    def wait(self, index: int, phase: int, *, loc=None, ip=None) -> None:
+        cute.arch.mbarrier_wait(
+            self.get_barrier(index, loc=loc, ip=ip), phase, loc=loc, ip=ip
+        )
 
     def arrive_and_wait(
         self,
@@ -270,14 +326,17 @@ class MbarrierArray(SyncObject):
         phase: int,
         dst: int,
         cta_group: Optional[cute.nvgpu.tcgen05.CtaGroup] = None,
+        *,
+        loc=None,
+        ip=None,
     ) -> None:
-        arrive(index, dst, cta_group)
-        wait(index, phase)
+        arrive(index, dst, cta_group, loc=loc, ip=ip)
+        wait(index, phase, loc=loc, ip=ip)
 
-    def arrive_and_drop(self) -> None:
+    def arrive_and_drop(self, *, loc=None, ip=None) -> None:
         raise NotImplementedError("Error: Not yet supported.")
 
-    def get_barrier(self, index: int) -> cute.Pointer:
+    def get_barrier(self, index: int, *, loc=None, ip=None) -> cute.Pointer:
         return self.mbarrier_base + index
 
     def max(self) -> int:
@@ -294,6 +353,11 @@ class MbarrierArray(SyncObject):
         )
 
 
+##############################################################################
+# NamedBarrier class
+##############################################################################
+
+
 @dataclass(frozen=True)
 class NamedBarrier(SyncObject):
     """
@@ -314,16 +378,21 @@ class NamedBarrier(SyncObject):
                 "NamedBarrier ID 0 is by other driver APIs (i.e. sync_threads()) and should not be used."
             )
 
-    def arrive(self) -> None:
+    @dsl_user_op
+    def arrive(self, *, loc=None, ip=None) -> None:
         """
         The aligned flavor of arrive is used when all threads in the CTA will execute the
         same instruction. See PTX documentation.
         """
         cute.arch.barrier_arrive(
-            barrier_id=self.barrier_id, number_of_threads=self.num_threads
+            barrier_id=self.barrier_id,
+            number_of_threads=self.num_threads,
+            loc=loc,
+            ip=ip,
         )
 
-    def arrive_unaligned(self) -> None:
+    @dsl_user_op
+    def arrive_unaligned(self, *, loc=None, ip=None) -> None:
         """
         The unaligned flavor of arrive can be used with an arbitrary number of threads in the CTA.
         """
@@ -337,7 +406,8 @@ class NamedBarrier(SyncObject):
             asm_dialect=llvm.AsmDialect.AD_ATT,
         )
 
-    def wait(self) -> None:
+    @dsl_user_op
+    def wait(self, *, loc=None, ip=None) -> None:
         """
         NamedBarriers do not have a standalone wait like mbarriers, only an arrive_and_wait.
         If synchronizing two warps in a producer/consumer pairing, the arrive count would be
@@ -348,7 +418,7 @@ class NamedBarrier(SyncObject):
         warnings.warn(
             "NamedBarrier wait also arrives on the barrier. Routing call to NamedBarrier.arrive_and_wait()."
         )
-        self.arrive_and_wait()
+        self.arrive_and_wait(loc=loc, ip=ip)
 
     def wait_unaligned(self) -> None:
         llvm.inline_asm(
@@ -361,18 +431,22 @@ class NamedBarrier(SyncObject):
             asm_dialect=llvm.AsmDialect.AD_ATT,
         )
 
-    def arrive_and_wait(self) -> None:
+    @dsl_user_op
+    def arrive_and_wait(self, *, loc=None, ip=None) -> None:
         cute.arch.barrier(
-            barrier_id=self.barrier_id, number_of_threads=self.num_threads
+            barrier_id=self.barrier_id,
+            number_of_threads=self.num_threads,
+            loc=loc,
+            ip=ip,
         )
 
-    def arrive_and_drop(self) -> None:
+    def arrive_and_drop(self, *, loc=None, ip=None) -> None:
         raise NotImplementedError("Error: Not supported.")
 
-    def sync(self) -> None:
+    def sync(self, *, loc=None, ip=None) -> None:
         self.arrive_and_wait()
 
-    def get_barrier(self) -> int:
+    def get_barrier(self, *, loc=None, ip=None) -> int:
         return self.barrier_id
 
     def max(self) -> int:
@@ -380,6 +454,11 @@ class NamedBarrier(SyncObject):
         return 4095
 
 
+##############################################################################
+# TmaStoreFence class
+##############################################################################
+
+
 class TmaStoreFence(SyncObject):
     """
     TmaStoreFence is used for a multi-stage epilogue buffer.
@@ -391,21 +470,27 @@ class TmaStoreFence(SyncObject):
 
         self.num_stages = num_stages
 
-    def arrive(self) -> None:
-        cute.arch.cp_async_bulk_commit_group()
+    @dsl_user_op
+    def arrive(self, *, loc=None, ip=None) -> None:
+        cute.arch.cp_async_bulk_commit_group(loc=loc, ip=ip)
 
-    def wait(self) -> None:
-        cute.arch.cp_async_bulk_wait_group(self.num_stages - 1, read=True)
+    @dsl_user_op
+    def wait(self, *, loc=None, ip=None) -> None:
+        cute.arch.cp_async_bulk_wait_group(
+            self.num_stages - 1, read=True, loc=loc, ip=ip
+        )
 
-    def arrive_and_wait(self) -> None:
-        self.arrive()
-        self.wait()
+    @dsl_user_op
+    def arrive_and_wait(self, *, loc=None, ip=None) -> None:
+        self.arrive(loc=loc, ip=ip)
+        self.wait(loc=loc, ip=ip)
 
-    def arrive_and_drop(self) -> None:
+    @dsl_user_op
+    def arrive_and_drop(self, *, loc=None, ip=None) -> None:
         raise NotImplementedError("Error: Not supported.")
 
     # TmaStoreFence doesn't have mbarriers
-    def get_barrier(self) -> None:
+    def get_barrier(self, *, loc=None, ip=None) -> None:
         assert False, (
             "Error: TmaStoreFence doesn't use mbarriers and cannot return a barrier."
         )
@@ -413,12 +498,13 @@ class TmaStoreFence(SyncObject):
     def max(self) -> None:
         raise NotImplementedError("Error: Not supported.")
 
-    def tail(self) -> None:
-        cute.arch.cp_async_bulk_wait_group(0, read=True)
+    @dsl_user_op
+    def tail(self, *, loc=None, ip=None) -> None:
+        cute.arch.cp_async_bulk_wait_group(0, read=True, loc=loc, ip=ip)
 
 
 ##############################################################################
-# PipelineState class
+# PipelineUserType class
 ##############################################################################
 
 
@@ -427,6 +513,11 @@ class PipelineUserType(enum.Enum):
     Consumer = enum.auto()
 
 
+##############################################################################
+# PipelineState class
+##############################################################################
+
+
 class PipelineState:
     """
     Pipeline state contains an index and phase bit corresponding to the current position in the circular buffer.
@@ -460,12 +551,13 @@ class PipelineState:
     def reset_count(self):
         self._count = Int32(0)
 
-    def advance(self):
+    @dsl_user_op
+    def advance(self, *, loc=None, ip=None):
         self._index += 1
         self._count += 1
 
         def then_body(index, phase):
-            new_index = Int32(0)
+            new_index = Int32(0, loc=loc, ip=ip)
             new_phase = phase ^ 1
             return new_index, new_phase
 
@@ -478,14 +570,17 @@ class PipelineState:
             else_body,
             [self.index, self.phase],
             [Int32, Int32],
+            loc=loc,
+            ip=ip,
         )
 
-    def reverse(self):
+    @dsl_user_op
+    def reverse(self, *, loc=None, ip=None):
         self._index -= 1
         self._count -= 1
 
         def then_body(index, phase):
-            new_index = Int32(self.stages - 1)
+            new_index = Int32(self.stages - 1, loc=loc, ip=ip)
             new_phase = phase ^ 1
             return new_index, new_phase
 
@@ -498,6 +593,8 @@ class PipelineState:
             else_body,
             [self.index, self.phase],
             [Int32, Int32],
+            loc=loc,
+            ip=ip,
         )
 
     def __get_mlir_types__(self):
@@ -516,23 +613,23 @@ class PipelineState:
         )
 
 
-def make_pipeline_state(type: PipelineUserType, stages: int):
+def make_pipeline_state(type: PipelineUserType, stages: int, *, loc=None, ip=None):
     """
     Creates a pipeline state. Producers are assumed to start with an empty buffer and have a flipped phase bit of 1.
     """
     if type is PipelineUserType.Producer:
         return PipelineState(
             stages,
-            Int32(0),
-            Int32(0),
-            Int32(1),
+            Int32(0, loc=loc, ip=ip),
+            Int32(0, loc=loc, ip=ip),
+            Int32(1, loc=loc, ip=ip),
         )
     elif type is PipelineUserType.Consumer:
         return PipelineState(
             stages,
-            Int32(0),
-            Int32(0),
-            Int32(0),
+            Int32(0, loc=loc, ip=ip),
+            Int32(0, loc=loc, ip=ip),
+            Int32(0, loc=loc, ip=ip),
         )
     else:
         assert False, (
@@ -545,31 +642,65 @@ def make_pipeline_state(type: PipelineUserType, stages: int):
 ##############################################################################
 
 
-def pipeline_init_wait(cta_layout_vmnk: Optional[cute.Layout] = None):
+def pipeline_init_arrive(
+    cluster_shape_mn: Optional[cute.Layout] = None,
+    is_relaxed: bool = False,
+    *,
+    loc=None,
+    ip=None,
+):
     """
-    Fences the mbarrier init and syncs the threadblock or cluster
+    Fences the mbarrier_init and sends an arrive if using clusters.
     """
-    cute.arch.mbarrier_init_fence()
 
-    if cta_layout_vmnk is None or cute.size(cta_layout_vmnk) == 1:
+    # If using clusters, send nonblocking arrives. Otherwise, do nothing
+    # because sync_threads() doesn't have a nonblocking arrive.
+
+    cute.arch.mbarrier_init_fence(loc=loc, ip=ip)
+
+    if cluster_shape_mn is not None and cute.size(cluster_shape_mn, loc=loc, ip=ip) > 1:
+        if is_relaxed:
+            # Fences memory operations issued before the arrive
+            cute.arch.cluster_arrive_relaxed(loc=loc, ip=ip)
+        else:
+            # Skips the memory barrier
+            cute.arch.cluster_arrive(loc=loc, ip=ip)
+
+
+def pipeline_init_wait(
+    cluster_shape_mn: Optional[cute.Layout] = None, *, loc=None, ip=None
+):
+    """
+    Syncs the threadblock or cluster
+    """
+
+    if cluster_shape_mn is None or cute.size(cluster_shape_mn, loc=loc, ip=ip) == 1:
         # If not using clusters, sync the threadblock
-        _sync(Agent.ThreadBlock)
+        agent_sync(Agent.ThreadBlock, loc=loc, ip=ip)
     else:
-        # If using clusters, sync the cluster
-        _sync(Agent.ThreadBlockCluster)
+        # If using clusters, wait on the cluster
+        cute.arch.cluster_wait(loc=loc, ip=ip)
 
 
-def _sync(group: Agent):
+def _sync(group: Agent, is_relaxed: bool = False, *, loc=None, ip=None):
+    warnings.warn("_sync is deprecated. Please use agent_sync instead.")
+    agent_sync(group, is_relaxed, loc=loc, ip=ip)
+
+
+def agent_sync(group: Agent, is_relaxed: bool = False, *, loc=None, ip=None):
     """
     Syncs all threads within an agent.
     """
     if group is Agent.Thread:
         raise NotImplementedError("Error: Not supported.")
     elif group is Agent.ThreadBlock:
-        cute.arch.sync_threads()
+        cute.arch.sync_threads(loc=loc, ip=ip)
     elif group is Agent.ThreadBlockCluster:
-        cute.arch.cluster_arrive()
-        cute.arch.cluster_wait()
+        if is_relaxed:
+            cute.arch.cluster_arrive_relaxed(loc=loc, ip=ip)
+        else:
+            cute.arch.cluster_arrive(loc=loc, ip=ip)
+        cute.arch.cluster_wait(loc=loc, ip=ip)
     else:
         assert False, (
             "Error: No explicit sync instruction exists. Please use barriers (named / mbarrier) instead."
@@ -589,15 +720,17 @@ def _mbarrier_i64_to_ptr(val: Int64) -> cute.Pointer:
 
 
 # NamedBarrier free functions
-def arrive(barrier_id: int, num_threads: int):
+def arrive(barrier_id: int, num_threads: int, *, loc=None, ip=None):
     """
     The aligned flavor of arrive is used when all threads in the CTA will execute the
     same instruction. See PTX documentation.
     """
-    cute.arch.barrier_arrive(barrier_id=barrier_id, number_of_threads=num_threads)
+    cute.arch.barrier_arrive(
+        barrier_id=barrier_id, number_of_threads=num_threads, loc=loc, ip=ip
+    )
 
 
-def arrive_unaligned(barrier_id: int, num_threads: int):
+def arrive_unaligned(barrier_id: int, num_threads: int, *, loc=None, ip=None):
     """
     The unaligned flavor of arrive can be used with an arbitrary number of threads in the CTA.
     """
@@ -623,10 +756,10 @@ def wait(barrier_id: int, num_threads: int):
     warnings.warn(
         "NamedBarrier wait also arrives on the barrier. Routing call to NamedBarrier.arrive_and_wait()."
     )
-    arrive_and_wait()
+    arrive_and_wait(loc=loc, ip=ip)
 
 
-def wait_unaligned(barrier_id: int, num_threads: int):
+def wait_unaligned(barrier_id: int, num_threads: int, *, loc=None, ip=None):
     warnings.warn(
         "NamedBarrier wait also arrives on the barrier. Routing call to NamedBarrier.arrive_and_wait()."
     )
@@ -641,9 +774,11 @@ def wait_unaligned(barrier_id: int, num_threads: int):
     )
 
 
-def arrive_and_wait(barrier_id: int, num_threads: int):
-    cute.arch.barrier(barrier_id=barrier_id, number_of_threads=num_threads)
+def arrive_and_wait(barrier_id: int, num_threads: int, *, loc=None, ip=None):
+    cute.arch.barrier(
+        barrier_id=barrier_id, number_of_threads=num_threads, loc=loc, ip=ip
+    )
 
 
-def sync(barrier_id: int = 0):
-    cute.arch.barrier(barrier_id=barrier_id)
+def sync(barrier_id: int = 0, *, loc=None, ip=None):
+    cute.arch.barrier(barrier_id=barrier_id, loc=loc, ip=ip)

python/CuTeDSL/cutlass/pipeline/sm100.py

@@ -17,11 +17,12 @@ import cutlass.cute as cute
 from cutlass.cutlass_dsl import Boolean, if_generate
 
 from cutlass.pipeline import (
+    Agent,
     CooperativeGroup,
     PipelineOp,
     PipelineState,
-    pipeline_init_wait,
     PipelineAsync,
+    agent_sync,
 )
 
 ##############################################################################
@@ -106,23 +107,27 @@ class PipelineTmaUmma(PipelineAsync):
         barrier_storage: cute.Pointer = None,
         cta_layout_vmnk: Optional[cute.Layout] = None,
         mcast_mode_mn: tuple[int, int] = (1, 1),
+        defer_sync: bool = False,
     ):
-        """
-        This helper function computes any necessary attributes and returns an instance of PipelineTmaUmma.
-        :param barrier_storage: Pointer to the smem address for this pipeline's mbarriers
-        :type barrier_storage: cute.Pointer
+        """Creates and initializes a new PipelineTmaUmma instance.
+
         :param num_stages: Number of buffer stages for this pipeline
-        :type num_stages: Int32
-        :param producer_group: `CooperativeGroup` for the producer agent
+        :type num_stages: int
+        :param producer_group: CooperativeGroup for the producer agent
         :type producer_group: CooperativeGroup
-        :param consumer_group: `CooperativeGroup` for the consumer agent
+        :param consumer_group: CooperativeGroup for the consumer agent
         :type consumer_group: CooperativeGroup
         :param tx_count: Number of bytes expected to be written to the transaction barrier for one stage
         :type tx_count: int
+        :param barrier_storage: Pointer to the shared memory address for this pipeline's mbarriers
+        :type barrier_storage: cute.Pointer, optional
         :param cta_layout_vmnk: Layout of the cluster shape
-        :type cta_layout_vmnk: cute.Layout | None
-        :param mcast_mode_mn: Tuple of two integers, specifying whether mcast is enabled for the m and n modes. At least one of the two integers must be 1.
-        :type mcast_mode_mn: tuple[int, int]
+        :type cta_layout_vmnk: cute.Layout, optional
+        :param mcast_mode_mn: Tuple specifying multicast modes for m and n dimensions (each 0 or 1)
+        :type mcast_mode_mn: tuple[int, int], optional
+        :raises ValueError: If barrier_storage is not a cute.Pointer instance
+        :return: A new PipelineTmaUmma instance configured with the provided parameters
+        :rtype: PipelineTmaUmma
         """
         if not isinstance(barrier_storage, cute.Pointer):
             raise ValueError(
@@ -161,7 +166,12 @@ class PipelineTmaUmma(PipelineAsync):
 
         consumer_mask = producer_mask
 
-        pipeline_init_wait(cta_layout_vmnk)
+        if not defer_sync:
+            cute.arch.mbarrier_init_fence()
+            if cta_layout_vmnk is None or cute.size(cta_layout_vmnk) == 1:
+                agent_sync(Agent.ThreadBlock)
+            else:
+                agent_sync(Agent.ThreadBlockCluster, is_relaxed=True)
 
         return PipelineTmaUmma(
             sync_object_full,
@@ -262,19 +272,23 @@ class PipelineAsyncUmma(PipelineAsync):
         consumer_group: CooperativeGroup,
         barrier_storage: cute.Pointer = None,
         cta_layout_vmnk: Optional[cute.Layout] = None,
+        defer_sync: bool = False,
     ):
-        """
-        This helper function computes any necessary attributes and returns an instance of PipelineAsyncUmma.
-        :param barrier_storage: Pointer to the smem address for this pipeline's mbarriers
-        :type barrier_storage: cute.Pointer
+        """Creates and initializes a new PipelineAsyncUmma instance.
+
         :param num_stages: Number of buffer stages for this pipeline
-        :type num_stages: Int32
-        :param producer_group: `CooperativeGroup` for the producer agent
+        :type num_stages: int
+        :param producer_group: CooperativeGroup for the producer agent
         :type producer_group: CooperativeGroup
-        :param consumer_group: `CooperativeGroup` for the consumer agent
+        :param consumer_group: CooperativeGroup for the consumer agent
         :type consumer_group: CooperativeGroup
+        :param barrier_storage: Pointer to the shared memory address for this pipeline's mbarriers
+        :type barrier_storage: cute.Pointer, optional
         :param cta_layout_vmnk: Layout of the cluster shape
-        :type cta_layout_vmnk: cute.Layout | None
+        :type cta_layout_vmnk: cute.Layout, optional
+        :raises ValueError: If barrier_storage is not a cute.Pointer instance
+        :return: A new PipelineAsyncUmma instance configured with the provided parameters
+        :rtype: PipelineAsyncUmma
         """
         if not isinstance(barrier_storage, cute.Pointer):
             raise ValueError(
@@ -315,7 +329,12 @@ class PipelineAsyncUmma(PipelineAsync):
             # consumer needs to get the mask to signal
             consumer_mask = PipelineAsyncUmma._compute_peer_cta_mask(cta_layout_vmnk)
 
-        pipeline_init_wait(cta_layout_vmnk)
+        if not defer_sync:
+            cute.arch.mbarrier_init_fence()
+            if cta_layout_vmnk is None or cute.size(cta_layout_vmnk) == 1:
+                agent_sync(Agent.ThreadBlock)
+            else:
+                agent_sync(Agent.ThreadBlockCluster, is_relaxed=True)
 
         return PipelineAsyncUmma(
             sync_object_full,
@@ -372,19 +391,23 @@ class PipelineUmmaAsync(PipelineAsync):
         consumer_group: CooperativeGroup,
         barrier_storage: cute.Pointer = None,
         cta_layout_vmnk: Optional[cute.Layout] = None,
+        defer_sync: bool = False,
     ):
-        """
-        This helper function computes any necessary attributes and returns an instance of PipelineUmmaAsync.
-        :param barrier_storage: Pointer to the smem address for this pipeline's mbarriers
-        :type barrier_storage: cute.Pointer
+        """Creates an instance of PipelineUmmaAsync with computed attributes.
+
         :param num_stages: Number of buffer stages for this pipeline
-        :type num_stages: Int32
-        :param producer_group: `CooperativeGroup` for the producer agent
+        :type num_stages: int
+        :param producer_group: ``CooperativeGroup`` for the producer agent
         :type producer_group: CooperativeGroup
-        :param consumer_group: `CooperativeGroup` for the consumer agent
+        :param consumer_group: ``CooperativeGroup`` for the consumer agent
         :type consumer_group: CooperativeGroup
+        :param barrier_storage: Pointer to the shared memory address for this pipeline's mbarriers
+        :type barrier_storage: cute.Pointer, optional
         :param cta_layout_vmnk: Layout of the cluster shape
-        :type cta_layout_vmnk: cute.Layout | None
+        :type cta_layout_vmnk: cute.Layout, optional
+        :raises ValueError: If barrier_storage is not a cute.Pointer instance
+        :return: New instance of ``PipelineUmmaAsync``
+        :rtype: PipelineUmmaAsync
         """
         if not isinstance(barrier_storage, cute.Pointer):
             raise ValueError(
@@ -411,7 +434,7 @@ class PipelineUmmaAsync(PipelineAsync):
             producer_mask = PipelineUmmaAsync._compute_tmem_sync_mask(cta_layout_vmnk)
 
         if cta_layout_vmnk is None or cute.size(cta_layout_vmnk, mode=[0]) == 1:
-            # Set mask to None if not using 2CTA intructions
+            # Set mask to None if not using 2CTA instructions
             consumer_mask = None
         else:
             consumer_mask = PipelineUmmaAsync._compute_peer_cta_rank()
@@ -422,7 +445,12 @@ class PipelineUmmaAsync(PipelineAsync):
             else cute.nvgpu.tcgen05.CtaGroup.TWO
         )
 
-        pipeline_init_wait(cta_layout_vmnk)
+        if not defer_sync:
+            cute.arch.mbarrier_init_fence()
+            if cta_layout_vmnk is None or cute.size(cta_layout_vmnk) == 1:
+                agent_sync(Agent.ThreadBlock)
+            else:
+                agent_sync(Agent.ThreadBlockCluster, is_relaxed=True)
 
         return PipelineUmmaAsync(
             sync_object_full,

python/CuTeDSL/cutlass/pipeline/sm90.py

@@ -13,8 +13,10 @@ from dataclasses import dataclass
 from typing import Optional
 
 import cutlass.cute as cute
-from cutlass.cutlass_dsl import Boolean, Int32, if_generate
+from cutlass.cutlass_dsl import Boolean, Int32, if_generate, dsl_user_op
+
 from cutlass.pipeline import (
+    Agent,
     CooperativeGroup,
     MbarrierArray,
     PipelineOp,
@@ -23,7 +25,7 @@ from cutlass.pipeline import (
     SyncObject,
     TmaStoreFence,
     make_pipeline_state,
-    pipeline_init_wait,
+    agent_sync,
 )
 
 ##############################################################################
@@ -157,6 +159,7 @@ class PipelineAsync:
         barrier_storage: cute.Pointer = None,
         producer_mask: Int32 = None,
         consumer_mask: Int32 = None,
+        defer_sync: bool = False,
     ):
         """Creates and initializes a new PipelineAsync instance.
 
@@ -167,17 +170,17 @@ class PipelineAsync:
         :type barrier_storage: cute.Pointer
         :param num_stages: Number of buffer stages for this pipeline
         :type num_stages: int
-        :param producer_group: `CooperativeGroup` for the producer agent
+        :param producer_group: ``CooperativeGroup`` for the producer agent
         :type producer_group: CooperativeGroup
-        :param consumer_group: `CooperativeGroup` for the consumer agent
+        :param consumer_group: ``CooperativeGroup`` for the consumer agent
         :type consumer_group: CooperativeGroup
-        :param producer_mask: Mask for signaling arrives for the producer agent, defaults to ``None``
+        :param producer_mask: Mask for signaling arrives for the producer agent
         :type producer_mask: Int32, optional
-        :param consumer_mask: Mask for signaling arrives for the consumer agent, defaults to ``None``
+        :param consumer_mask: Mask for signaling arrives for the consumer agent
         :type consumer_mask: Int32, optional
-        :return: A new PipelineAsync instance
-        :rtype: PipelineAsync
         :raises ValueError: If barrier_storage is not a cute.Pointer instance
+        :return: A new ``PipelineAsync`` instance
+        :rtype: PipelineAsync
         """
         if not isinstance(barrier_storage, cute.Pointer):
             raise ValueError(
@@ -197,7 +200,9 @@ class PipelineAsync:
             barrier_storage.align(min_align=8) + num_stages, num_stages, consumer
         )
 
-        pipeline_init_wait()
+        if not defer_sync:
+            cute.arch.mbarrier_init_fence()
+            agent_sync(Agent.ThreadBlock)
 
         return PipelineAsync(
             sync_object_full,
@@ -280,21 +285,24 @@ class PipelineCpAsync(PipelineAsync):
         consumer_group: CooperativeGroup,
         producer_mask: Int32 = None,
         consumer_mask: Int32 = None,
+        defer_sync: bool = False,
     ):
-        """
-        This helper function computes any necessary attributes and returns an instance of PipelineAsync.
-        :param barrier_storage: Pointer to the smem address for this pipeline's mbarriers
+        """Helper function that computes necessary attributes and returns a ``PipelineCpAsync`` instance.
+
+        :param barrier_storage: Pointer to the shared memory address for this pipeline's mbarriers
         :type barrier_storage: cute.Pointer
         :param num_stages: Number of buffer stages for this pipeline
         :type num_stages: Int32
-        :param producer_group: CooperativeGroup for the producer agent
+        :param producer_group: ``CooperativeGroup`` for the producer agent
         :type producer_group: CooperativeGroup
-        :param consumer_group: CooperativeGroup for the consumer agent
+        :param consumer_group: ``CooperativeGroup`` for the consumer agent
         :type consumer_group: CooperativeGroup
-        :param producer_mask: Mask for signaling arrives for the producer agent
-        :type producer_mask: Int32 | None
-        :param consumer_mask: Mask for signaling arrives for the consumer agent
-        :type consumer_mask: Int32 | None
+        :param producer_mask: Mask for signaling arrives for the producer agent, defaults to None
+        :type producer_mask: Int32, optional
+        :param consumer_mask: Mask for signaling arrives for the consumer agent, defaults to None
+        :type consumer_mask: Int32, optional
+        :return: A new ``PipelineCpAsync`` instance configured with the provided parameters
+        :rtype: PipelineCpAsync
         """
         producer_type = PipelineOp.AsyncLoad
         consumer_type = PipelineOp.AsyncThread
@@ -309,7 +317,9 @@ class PipelineCpAsync(PipelineAsync):
             barrier_storage.align(min_align=8) + num_stages, num_stages, consumer
         )
 
-        pipeline_init_wait()
+        if not defer_sync:
+            cute.arch.mbarrier_init_fence()
+            agent_sync(Agent.ThreadBlock)
 
         return PipelineCpAsync(
             sync_object_array_full,
@@ -331,11 +341,22 @@ class PipelineTmaAsync(PipelineAsync):
     @staticmethod
     @cute.jit
     def init_empty_barrier_arrive_signal(
-        cta_layout_vmnk: cute.Layout, tidx: Int32, mcast_mode_mn: tuple[int, int]
+        cta_layout_vmnk: cute.Layout, tidx: Int32, mcast_mode_mn: tuple[int, int] = (1, 1)
     ):
-        """
-        Initialize the empty barrier arrive signal
-        This function returns the destination cta rank and a boolean indicating if the signalling thread is the same as the current thread
+        """Initialize the empty barrier arrive signal.
+
+        This function determines which threads should signal empty barrier arrives based on the cluster layout
+        and multicast modes. It returns the destination CTA rank and whether the current thread should signal.
+
+        :param cta_layout_vmnk: Layout describing the cluster shape and CTA arrangement
+        :type cta_layout_vmnk: cute.Layout
+        :param tidx: Thread index within the warp
+        :type tidx: Int32
+        :param mcast_mode_mn: Tuple specifying multicast modes for m and n dimensions (each 0 or 1), defaults to (1,1)
+        :type mcast_mode_mn: tuple[int, int]
+        :raises ``AssertionError``: If both multicast modes are disabled (0,0)
+        :return: Tuple containing destination CTA rank and boolean indicating if current thread signals
+        :rtype: tuple[Int32, Boolean]
         """
         # Logic to optimally schedule Empty Arrives
         cluster_shape_vmnk = cta_layout_vmnk.shape
@@ -384,25 +405,29 @@ class PipelineTmaAsync(PipelineAsync):
         cta_layout_vmnk: Optional[cute.Layout] = None,
         tidx: Optional[Int32] = None,
         mcast_mode_mn: tuple[int, int] = (1, 1),
+        defer_sync: bool = False,
     ):
-        """
-        This helper function computes any necessary attributes and returns an instance of PipelineTmaAsync.
-        :param barrier_storage: Pointer to the smem address for this pipeline's mbarriers
-        :type barrier_storage: cute.Pointer
+        """Create a new ``PipelineTmaAsync`` instance.
+
         :param num_stages: Number of buffer stages for this pipeline
-        :type num_stages: Int32
-        :param producer_group: `CooperativeGroup` for the producer agent
+        :type num_stages: int
+        :param producer_group: ``CooperativeGroup`` for the producer agent
         :type producer_group: CooperativeGroup
-        :param consumer_group: `CooperativeGroup` for the consumer agent
+        :param consumer_group: ``CooperativeGroup`` for the consumer agent
         :type consumer_group: CooperativeGroup
         :param tx_count: Number of bytes expected to be written to the transaction barrier for one stage
         :type tx_count: int
-        :param cta_layout_vmnk: Layout of the cluster shape
-        :type cta_layout_vmnk: cute.Layout | None
-        :param tidx: thread index to consumer async threads
-        :type tidx: Int32 | None
-        :param mcast_mode_mn: Tuple of two integers, specifying whether mcast is enabled for the m and n modes. At least one of the two integers must be 1.
-        :type mcast_mode_mn: tuple[int, int]
+        :param barrier_storage: Pointer to the shared memory address for this pipeline's mbarriers, defaults to None
+        :type barrier_storage: cute.Pointer, optional
+        :param cta_layout_vmnk: Layout of the cluster shape, defaults to None
+        :type cta_layout_vmnk: cute.Layout, optional
+        :param tidx: Thread index to consumer async threads, defaults to None
+        :type tidx: Int32, optional
+        :param mcast_mode_mn: Tuple specifying multicast modes for m and n dimensions (each 0 or 1), defaults to (1,1)
+        :type mcast_mode_mn: tuple[int, int], optional
+        :raises ValueError: If barrier_storage is not a cute.Pointer instance
+        :return: New ``PipelineTmaAsync`` instance
+        :rtype: PipelineTmaAsync
         """
         if not isinstance(barrier_storage, cute.Pointer):
             raise ValueError(
@@ -438,7 +463,12 @@ class PipelineTmaAsync(PipelineAsync):
 
         producer_mask = None
 
-        pipeline_init_wait(cta_layout_vmnk)
+        if not defer_sync:
+            cute.arch.mbarrier_init_fence()
+            if cta_layout_vmnk is None or cute.size(cta_layout_vmnk) == 1:
+                agent_sync(Agent.ThreadBlock)
+            else:
+                agent_sync(Agent.ThreadBlockCluster, is_relaxed=True)
 
         return PipelineTmaAsync(
             sync_object_full,
@@ -449,17 +479,25 @@ class PipelineTmaAsync(PipelineAsync):
             is_signalling_thread,
         )
 
+    @dsl_user_op
     def producer_acquire(
-        self, state: PipelineState, try_acquire_token: Optional[Boolean] = None
+        self,
+        state: PipelineState,
+        try_acquire_token: Optional[Boolean] = None,
+        *,
+        loc=None,
+        ip=None,
     ):
         """
         TMA producer commit conditionally waits on buffer empty and sets the transaction barrier.
         """
         if_generate(
             try_acquire_token is None or try_acquire_token == 0,
-            lambda: self.sync_object_empty.wait(state.index, state.phase),
+            lambda: self.sync_object_empty.wait(
+                state.index, state.phase, loc=loc, ip=ip
+            ),
         )
-        self.sync_object_full.arrive(state.index, self.producer_mask)
+        self.sync_object_full.arrive(state.index, self.producer_mask, loc=loc, ip=ip)
 
     def producer_commit(self, state: PipelineState):
         """
@@ -467,13 +505,16 @@ class PipelineTmaAsync(PipelineAsync):
         """
         pass
 
-    def consumer_release(self, state: PipelineState):
+    @dsl_user_op
+    def consumer_release(self, state: PipelineState, *, loc=None, ip=None):
         """
         TMA consumer release conditionally signals the empty buffer to the producer.
         """
         if_generate(
             self.is_signalling_thread,
-            lambda: self.sync_object_empty.arrive(state.index, self.consumer_mask),
+            lambda: self.sync_object_empty.arrive(
+                state.index, self.consumer_mask, loc=loc, ip=ip
+            ),
         )
 
 
@@ -498,23 +539,29 @@ class PipelineTmaMultiConsumersAsync(PipelineAsync):
         tx_count: int,
         barrier_storage: cute.Pointer = None,
         cta_layout_vmnk: Optional[cute.Layout] = None,
+        defer_sync: bool = False,
     ):
-        """
-        This helper function computes any necessary attributes and returns an instance of PipelineTmaMultiConsumersAsync.
-        :param barrier_storage: Pointer to the smem address for this pipeline's mbarriers
+        """Creates an instance of PipelineTmaMultiConsumersAsync with computed attributes.
+
+        :param barrier_storage: Pointer to the shared memory address for this pipeline's mbarriers
         :type barrier_storage: cute.Pointer
         :param num_stages: Number of buffer stages for this pipeline
-        :type num_stages: Int32
-        :param producer_group: `CooperativeGroup` for the producer agent
+        :type num_stages: int
+        :param producer_group: ``CooperativeGroup`` for the producer agent
         :type producer_group: CooperativeGroup
-        :param consumer_group_umma: `CooperativeGroup` for the UMMA consumer agent
+        :param consumer_group_umma: ``CooperativeGroup`` for the UMMA consumer agent
         :type consumer_group_umma: CooperativeGroup
-        :param consumer_group_async: `CooperativeGroup` for the AsyncThread consumer agent
+        :param consumer_group_async: ``CooperativeGroup`` for the AsyncThread consumer agent
         :type consumer_group_async: CooperativeGroup
         :param tx_count: Number of bytes expected to be written to the transaction barrier for one stage
         :type tx_count: int
-        :param cta_layout_vmnk: Layout of the cluster shape
-        :type cta_layout_vmnk: cute.Layout | None
+        :param cta_layout_vmnk: Layout of the cluster shape, defaults to None
+        :type cta_layout_vmnk: Optional[cute.Layout]
+        :raises ValueError: If ``barrier_storage`` is not a ``cute.Pointer`` instance
+        :raises ValueError: If ``UMMA`` and ``AsyncThread`` consumer groups are not the same agent
+        :raises ValueError: If ``cta_layout_vmnk`` size is not 1
+        :return: New instance of ``PipelineTmaMultiConsumersAsync``
+        :rtype: PipelineTmaMultiConsumersAsync
         """
         if not isinstance(barrier_storage, cute.Pointer):
             raise ValueError(
@@ -560,7 +607,7 @@ class PipelineTmaMultiConsumersAsync(PipelineAsync):
         # No mcast mask if not using clusters
         producer_mask = None
         consumer_mask = None
-        # All threadblocks are leaders if not using clusters
+        # All thread-blocks are leaders if not using clusters
         is_leader_cta = True
         cta_group = (
             cute.nvgpu.tcgen05.CtaGroup.ONE
@@ -568,7 +615,11 @@ class PipelineTmaMultiConsumersAsync(PipelineAsync):
             else cute.nvgpu.tcgen05.CtaGroup.TWO
         )
 
-        pipeline_init_wait(cta_layout_vmnk)
+        if not defer_sync:
+            if cta_layout_vmnk is None or cute.size(cta_layout_vmnk) == 1:
+                agent_sync(Agent.ThreadBlock)
+            else:
+                agent_sync(Agent.ThreadBlockCluster, is_relaxed=True)
 
         return PipelineTmaMultiConsumersAsync(
             sync_object_full,
@@ -626,14 +677,15 @@ class PipelineTmaStore(PipelineAsync):
         num_stages: int,
         producer_group: CooperativeGroup,
     ):
-        """
-        This helper function computes any necessary attributes and returns an instance of PipelineTmaStore.
-        :param num_stages: Number of buffer stages for this pipeline
-        :type num_stages: Int32
-        :param producer_group: `CooperativeGroup` for the producer agent
-        :type producer_group: CooperativeGroup
-        """
+        """This helper function computes any necessary attributes and returns an instance of ``PipelineTmaStore``.
 
+        :param num_stages: Number of buffer stages for this pipeline
+        :type num_stages: int
+        :param producer_group: ``CooperativeGroup`` for the producer agent
+        :type producer_group: CooperativeGroup
+        :return: A new ``PipelineTmaStore`` instance
+        :rtype: PipelineTmaStore
+        """
         producer_type = PipelineOp.TmaStore
 
         producer = (producer_type, producer_group)
@@ -642,20 +694,25 @@ class PipelineTmaStore(PipelineAsync):
 
         return PipelineTmaStore(sync_object_full, None, num_stages, None, None)
 
-    def producer_acquire(self):
-        self.sync_object_full.wait()
+    @dsl_user_op
+    def producer_acquire(self, *, loc=None, ip=None):
+        self.sync_object_full.wait(loc=loc, ip=ip)
 
-    def producer_commit(self):
-        self.sync_object_full.arrive()
+    @dsl_user_op
+    def producer_commit(self, *, loc=None, ip=None):
+        self.sync_object_full.arrive(loc=loc, ip=ip)
 
-    def consumer_wait(self):
+    @dsl_user_op
+    def consumer_wait(self, *, loc=None, ip=None):
         assert False, "Error: PipelineTmaStore does not have a consumer agent."
 
-    def consumer_release(self):
+    @dsl_user_op
+    def consumer_release(self, *, loc=None, ip=None):
         assert False, "Error: PipelineTmaStore does not have a consumer agent."
 
-    def producer_tail(self):
-        self.sync_object_full.tail()
+    @dsl_user_op
+    def producer_tail(self, *, loc=None, ip=None):
+        self.sync_object_full.tail(loc=loc, ip=ip)
 
 
 @dataclass(frozen=True)
@@ -706,6 +763,7 @@ class PipelineOrder:
         length: int,
         group_id: int,
         producer_group: CooperativeGroup,
+        defer_sync: bool = False,
     ):
         if not isinstance(barrier_storage, cute.Pointer):
             raise ValueError(
@@ -722,7 +780,9 @@ class PipelineOrder:
             barrier_storage.align(min_align=8), num_stages, producer
         )
 
-        pipeline_init_wait()
+        if not defer_sync:
+            cute.arch.mbarrier_init_fence()
+            agent_sync(Agent.ThreadBlock)
 
         return PipelineOrder(
             sync_object_full,
@@ -740,16 +800,20 @@ class PipelineOrder:
     def get_barrier_for_current_stage_idx(self, group_id):
         return self.state.index * self.length + group_id
 
-    def arrive(self):
+    @dsl_user_op
+    def arrive(self, *, loc=None, ip=None):
         signalling_id = (self.group_id + 1) % self.length
         idx = self.get_barrier_for_current_stage_idx(signalling_id)
-        cute.arch.mbarrier_arrive(self.sync_object_full.get_barrier(idx))
-        self.state.advance()
+        cute.arch.mbarrier_arrive(
+            self.sync_object_full.get_barrier(idx), loc=loc, ip=ip
+        )
+        self.state.advance(loc=loc, ip=ip)
 
-    def wait(self):
+    @dsl_user_op
+    def wait(self, *, loc=None, ip=None):
         idx = self.get_barrier_for_current_stage_idx(self.group_id)
         cute.arch.mbarrier_wait(
-            self.sync_object_full.get_barrier(idx), self.state.phase
+            self.sync_object_full.get_barrier(idx), self.state.phase, loc=loc, ip=ip
         )
 
 
@@ -881,6 +945,10 @@ class PipelineProducer:
         self.__state = state
         self.__group = group
 
+    def clone(self):
+        """Create a new Producer instance with the same state."""
+        return PipelineProducer(self.__pipeline, self.__state.clone(), self.__group)
+
     def reset(self):
         """Reset the count of how many handles this producer has committed."""
         self.__state.reset_count()
@@ -1053,6 +1121,10 @@ class PipelineConsumer:
         self.__group = group
         self.__state = state
 
+    def clone(self):
+        """Create a new Consumer instance with the same state."""
+        return PipelineConsumer(self.__pipeline, self.__state.clone(), self.__group)
+
     def reset(self):
         """Reset the count of how many handles this consumer has consumed."""
         self.__state.reset_count()

python/CuTeDSL/cutlass/torch.py

@@ -47,6 +47,11 @@ def dtype(ty: Type[Numeric]):
         Float8E4M3FN: torch.float8_e4m3fn,
         Float8E4M3B11FNUZ: torch.float8_e4m3fnuz,
     }
+
+    # float8_e8m0fnu is introduced in latest version of torch
+    if hasattr(torch, "float8_e8m0fnu"):
+        torch_type_map[Float8E8M0FNU] = torch.float8_e8m0fnu
+
     if torch_dtype is None:
         torch_dtype = torch_type_map.get(ty)
 

python/CuTeDSL/cutlass/utils/__init__.py

@@ -13,6 +13,7 @@ from .static_persistent_tile_scheduler import (
     WorkTileInfo,
     PersistentTileSchedulerParams,
     StaticPersistentTileScheduler,
+    StaticPersistentRuntimeTileScheduler,
 )
 
 from .hardware_info import (
@@ -65,21 +66,27 @@ from .tmem_allocator import TmemAllocator
 
 from .layout import LayoutEnum
 
-from .distributed_helpers import (
-    spin_lock_wait,
-    spin_lock_multimem_arrive,
-    multimem_ld_reduce_8xf16,
-    multimem_ld_reduce_4xf32,
-    multimem_ld_reduce_8xbf16,
-    multimem_ld_reduce_16xe4m3,
-    multimem_ld_reduce_16xe5m2,
-    multimem_st_4xb32,
-    sm_wise_inter_gpu_multimem_barrier,
+from . import distributed
+
+from .mixed_input_helpers import (
+    TransformMode,
+    scale_tma_partition,
+    transform_partition,
+    scale_partition,
+    get_gmem_layout_scale,
+    get_smem_layout_scale,
+    compute_smem_layout,
+    get_transform_a_source,
+    get_tma_atom_kind,
+    get_copy_atom_a_transform,
+    is_valid_scale_granularity,
+    get_divisibility,
 )
 
 from . import hopper_helpers as sm90
 from . import blackwell_helpers as sm100
 
+from .print_latex import print_latex, print_latex_tv
 
 __all__ = [
     "get_smem_capacity_in_bytes",
@@ -89,6 +96,7 @@ __all__ = [
     "WorkTileInfo",
     "PersistentTileSchedulerParams",
     "StaticPersistentTileScheduler",
+    "StaticPersistentRuntimeTileScheduler",
     "TensorMapUpdateMode",
     "TensorMapManager",
     "GroupSearchResult",
@@ -96,6 +104,18 @@ __all__ = [
     "create_initial_search_state",
     "GroupedGemmTileSchedulerHelper",
     "HardwareInfo",
+    "TransformMode",
+    "scale_tma_partition",
+    "transform_partition",
+    "scale_partition",
+    "get_gmem_layout_scale",
+    "get_smem_layout_scale",
+    "compute_smem_layout",
+    "get_transform_a_source",
+    "get_tma_atom_kind",
+    "get_copy_atom_a_transform",
+    "is_valid_scale_granularity",
+    "get_divisibility",
     "compute_epilogue_tile_shape",
     "get_smem_store_op",
     "get_tmem_load_op",
@@ -107,4 +127,7 @@ __all__ = [
     "make_blockscaled_trivial_tiled_mma",
     "sm90",
     "sm100",
+    "print_latex",
+    "print_latex_tv",
+    "distributed",
 ]

python/CuTeDSL/cutlass/utils/blackwell_helpers.py

@@ -689,7 +689,7 @@ def make_smem_layout_a(
 
     is_k_major = tiled_mma.op.a_major_mode == OperandMajorMode.K
     a_smem_shape = tiled_mma.partition_shape_A(
-        cute.dice(mma_tiler_mnk, (1, None, 1), loc=loc, ip=ip)
+        cute.dice(mma_tiler_mnk, (1, None, 1), loc=loc, ip=ip), loc=loc, ip=ip
     )
     a_smem_shape_mn_k = (
         cute.size(a_smem_shape[0][0], loc=loc, ip=ip) * a_smem_shape[1],
@@ -741,7 +741,7 @@ def make_smem_layout_b(
 
     is_k_major = tiled_mma.op.b_major_mode == OperandMajorMode.K
     b_smem_shape = tiled_mma.partition_shape_B(
-        cute.dice(mma_tiler_mnk, (None, 1, 1), loc=loc, ip=ip)
+        cute.dice(mma_tiler_mnk, (None, 1, 1), loc=loc, ip=ip), loc=loc, ip=ip
     )
     b_smem_shape_nk = (
         cute.size(b_smem_shape[0][0], loc=loc, ip=ip) * b_smem_shape[1],
@@ -939,7 +939,9 @@ def make_trivial_tiled_mma(
     else:
         raise TypeError(f"unsupported ab_dtype, got {ab_dtype}")
 
-    return cute.make_tiled_mma(cute.make_mma_atom(mma_op), loc=loc, ip=ip)
+    return cute.make_tiled_mma(
+        cute.make_mma_atom(mma_op, loc=loc, ip=ip), loc=loc, ip=ip
+    )
 
 
 @dsl_user_op
@@ -1009,7 +1011,9 @@ def make_blockscaled_trivial_tiled_mma(
     else:
         raise TypeError(f"unsupported ab_dtype, got {ab_dtype}")
 
-    return cute.make_tiled_mma(cute.make_mma_atom(mma_op), loc=loc, ip=ip)
+    return cute.make_tiled_mma(
+        cute.make_mma_atom(mma_op, loc=loc, ip=ip), loc=loc, ip=ip
+    )
 
 
 @dsl_user_op

python/CuTeDSL/cutlass/utils/distributed.py

@@ -0,0 +1,386 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 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/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.
+
+from functools import partial
+from typing import Tuple
+
+import cutlass.cute as cute
+from cutlass.cute.typing import Pointer, Int32
+from cutlass.cutlass_dsl import T, dsl_user_op
+from cutlass._mlir import ir
+from cutlass._mlir.dialects import llvm, nvvm
+from cutlass._mlir.dialects.nvvm import (
+    MemOrderKind,
+    MemScopeKind,
+    AtomicOpKind,
+)
+
+
+__all__ = [
+    # misc
+    "atomicAdd",
+    "ld_bypass",
+    # Message Passing Lock & Unlock
+    "multimem_red_add1",
+    "red_add1",
+    "spin_lock_atom_cas_relaxed_wait",
+    # Load & Store
+    "multimem_ld_reduce_8xf16",
+    "multimem_ld_reduce_4xf32",
+    "multimem_ld_reduce_8xbf16",
+    "multimem_ld_reduce_16xe4m3",
+    "multimem_ld_reduce_16xe5m2",
+    "multimem_st_4xb32",
+]
+
+
+@dsl_user_op
+def atomicAdd(dst_ptr: Pointer, val: Int32, loc=None, ip=None) -> Int32:
+    return nvvm.atomicrmw(
+        T.i32(),
+        AtomicOpKind.ADD,
+        dst_ptr.llvm_ptr,
+        val.ir_value(loc=loc, ip=ip),
+        mem_order=MemOrderKind.RELAXED,
+        syncscope=MemScopeKind.SYS,
+        loc=loc,
+        ip=ip,
+    )
+
+
+@cute.jit
+def ld_bypass(input_tensor: cute.Tensor):
+    fragment = cute.make_rmem_tensor(input_tensor.layout, input_tensor.element_type)
+    copy_atom_load = cute.make_copy_atom(
+        cute.nvgpu.CopyUniversalOp(),
+        input_tensor.element_type,
+        memory_order=cute.nvgpu.common.MemoryOrder.VOLATILE,
+        memory_scope=cute.nvgpu.common.MemoryScope.SYS,
+    )
+    cute.copy_atom_call(copy_atom_load, input_tensor, fragment)
+    vals = fragment.load()
+    return vals
+
+
+########################################################
+# Message Passing Lock & Unlock
+########################################################
+
+
+@dsl_user_op
+def multimem_red_release_gpu_add1(
+    lock_ptr: Pointer,
+    loc=None,
+    ip=None,
+) -> None:
+    llvm.inline_asm(
+        None,
+        [lock_ptr.toint().ir_value(loc=loc, ip=ip)],
+        "multimem.red.release.gpu.global.add.u32 [$0], 1;",
+        "l",
+        has_side_effects=True,
+        asm_dialect=0,
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def multimem_red_release_sys_add1(
+    lock_ptr: Pointer,
+    loc=None,
+    ip=None,
+) -> None:
+    llvm.inline_asm(
+        None,
+        [lock_ptr.toint().ir_value(loc=loc, ip=ip)],
+        "multimem.red.release.sys.global.add.u32 [$0], 1;",
+        "l",
+        has_side_effects=True,
+        asm_dialect=0,
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def multimem_red_relaxed_gpu_add1(
+    lock_ptr: Pointer,
+    loc=None,
+    ip=None,
+) -> None:
+    llvm.inline_asm(
+        None,
+        [lock_ptr.toint().ir_value(loc=loc, ip=ip)],
+        "multimem.red.relaxed.gpu.global.add.u32 [$0], 1;",
+        "l",
+        has_side_effects=True,
+        asm_dialect=0,
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def multimem_red_relaxed_sys_add1(
+    lock_ptr: Pointer,
+    loc=None,
+    ip=None,
+) -> None:
+    llvm.inline_asm(
+        None,
+        [lock_ptr.toint().ir_value(loc=loc, ip=ip)],
+        "multimem.red.relaxed.sys.global.add.u32 [$0], 1;",
+        "l",
+        has_side_effects=True,
+        asm_dialect=0,
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def multimem_red_add1(
+    lock_ptr: Pointer,
+    *,
+    order: str,
+    scope: str,
+    loc=None,
+    ip=None,
+) -> None:
+    """
+    add 1 to multicast ptr
+    """
+    if scope == "gpu":
+        if order == "release":
+            multimem_red_release_gpu_add1(lock_ptr=lock_ptr, loc=loc, ip=ip)
+        elif order == "relaxed":
+            multimem_red_relaxed_gpu_add1(lock_ptr=lock_ptr, loc=loc, ip=ip)
+    elif scope == "sys":
+        if order == "release":
+            multimem_red_release_sys_add1(lock_ptr=lock_ptr, loc=loc, ip=ip)
+        elif order == "relaxed":
+            multimem_red_relaxed_sys_add1(lock_ptr=lock_ptr, loc=loc, ip=ip)
+
+
+@dsl_user_op
+def red_release_gpu_add1(
+    lock_ptr: Pointer,
+    loc=None,
+    ip=None,
+) -> None:
+    llvm.inline_asm(
+        None,
+        [lock_ptr.toint().ir_value(loc=loc, ip=ip)],
+        "red.release.gpu.global.add.u32 [$0], 1;",
+        "l",
+        has_side_effects=True,
+        asm_dialect=0,
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def red_release_sys_add1(
+    lock_ptr: Pointer,
+    loc=None,
+    ip=None,
+) -> None:
+    llvm.inline_asm(
+        None,
+        [lock_ptr.toint().ir_value(loc=loc, ip=ip)],
+        "red.release.sys.global.add.u32 [$0], 1;",
+        "l",
+        has_side_effects=True,
+        asm_dialect=0,
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def red_relaxed_gpu_add1(
+    lock_ptr: Pointer,
+    loc=None,
+    ip=None,
+) -> None:
+    llvm.inline_asm(
+        None,
+        [lock_ptr.toint().ir_value(loc=loc, ip=ip)],
+        "red.relaxed.gpu.global.add.u32 [$0], 1;",
+        "l",
+        has_side_effects=True,
+        asm_dialect=0,
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def red_relaxed_sys_add1(
+    lock_ptr: Pointer,
+    loc=None,
+    ip=None,
+) -> None:
+    llvm.inline_asm(
+        None,
+        [lock_ptr.toint().ir_value(loc=loc, ip=ip)],
+        "red.relaxed.sys.global.add.u32 [$0], 1;",
+        "l",
+        has_side_effects=True,
+        asm_dialect=0,
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def red_add1(
+    lock_ptr: Pointer,
+    *,
+    order: str,
+    scope: str,
+    loc=None,
+    ip=None,
+) -> None:
+    """
+    add 1 to multicast ptr
+    """
+    if scope == "gpu":
+        if order == "release":
+            red_release_gpu_add1(lock_ptr=lock_ptr, loc=loc, ip=ip)
+        elif order == "relaxed":
+            red_relaxed_gpu_add1(lock_ptr=lock_ptr, loc=loc, ip=ip)
+    elif scope == "sys":
+        if order == "release":
+            red_release_sys_add1(lock_ptr=lock_ptr, loc=loc, ip=ip)
+        elif order == "relaxed":
+            red_relaxed_sys_add1(lock_ptr=lock_ptr, loc=loc, ip=ip)
+
+
+@cute.jit
+def spin_lock_atom_cas_relaxed_wait(
+    lock_ptr: Pointer,
+    *,
+    expected_val: Int32,
+    reset_val: Int32,
+    scope: str,
+    loc=None,
+    ip=None,
+) -> None:
+    """
+    wait on a spin lock until the expected count is reached. Reset flag to reset_val if the expected count is reached.
+    """
+    if scope == "gpu":
+        result = 0
+        while result != expected_val:
+            result = nvvm.atomicrmw(
+                T.i32(),
+                AtomicOpKind.CAS,
+                lock_ptr.llvm_ptr,
+                Int32(reset_val).ir_value(loc=loc, ip=ip),
+                b=Int32(expected_val).ir_value(loc=loc, ip=ip),
+                mem_order=MemOrderKind.RELAXED,
+                syncscope=MemScopeKind.GPU,
+                loc=loc,
+                ip=ip,
+            )
+    elif scope == "sys":
+        result = 0
+        while result != expected_val:
+            result = nvvm.atomicrmw(
+                T.i32(),
+                AtomicOpKind.CAS,
+                lock_ptr.llvm_ptr,
+                Int32(reset_val).ir_value(loc=loc, ip=ip),
+                b=Int32(expected_val).ir_value(loc=loc, ip=ip),
+                mem_order=MemOrderKind.RELAXED,
+                syncscope=MemScopeKind.SYS,
+                loc=loc,
+                ip=ip,
+            )
+
+
+########################################################
+# Multimem Load & Store
+########################################################
+
+
+@dsl_user_op
+def multimem_ld_reduce_base(
+    mc_ptr: Pointer,
+    *,
+    ptx_string: str = "",
+    loc=None,
+    ip=None,
+) -> Tuple[Int32, Int32, Int32, Int32]:
+    # ld reduce 8xf16 elts
+    mc_ptr_int = mc_ptr.toint(loc=loc, ip=ip).ir_value(loc=loc, ip=ip)
+    return_struct = llvm.inline_asm(
+        ir.Type.parse("!llvm.struct<(i32,i32,i32,i32)>"),
+        [mc_ptr_int],
+        ptx_string,
+        "=r,=r,=r,=r,l",
+        has_side_effects=True,
+        asm_dialect=0,
+        loc=loc,
+        ip=ip,
+    )
+    return_regs = [llvm.extractvalue(T.i32(), return_struct, [i]) for i in range(4)]
+    return return_regs[0], return_regs[1], return_regs[2], return_regs[3]
+
+
+multimem_ld_reduce_8xf16 = partial(
+    multimem_ld_reduce_base,
+    ptx_string="multimem.ld_reduce.sys.relaxed.global.add.acc::f32.v4.f16x2 {$0, $1, $2, $3}, [$4];",
+)
+multimem_ld_reduce_4xf32 = partial(
+    multimem_ld_reduce_base,
+    ptx_string="multimem.ld_reduce.sys.relaxed.global.add.v4.f32 {$0, $1, $2, $3}, [$4];",
+)
+multimem_ld_reduce_8xbf16 = partial(
+    multimem_ld_reduce_base,
+    ptx_string="multimem.ld_reduce.sys.relaxed.global.add.acc::f32.v4.bf16x2 {$0, $1, $2, $3}, [$4];",
+)
+multimem_ld_reduce_16xe4m3 = partial(
+    multimem_ld_reduce_base,
+    ptx_string="multimem.ld_reduce.sys.relaxed.global.add.acc::f16.v4.e4m3x4 {$0, $1, $2, $3}, [$4];",
+)
+multimem_ld_reduce_16xe5m2 = partial(
+    multimem_ld_reduce_base,
+    ptx_string="multimem.ld_reduce.sys.relaxed.global.add.acc::f16.v4.e5m2x4 {$0, $1, $2, $3}, [$4];",
+)
+
+
+@dsl_user_op
+def multimem_st_4xb32(
+    mc_ptr: Pointer,
+    x: Int32,
+    y: Int32,
+    z: Int32,
+    w: Int32,
+    *,
+    loc=None,
+    ip=None,
+) -> None:
+    # st 4x32 bits of data
+    mc_ptr_int = mc_ptr.toint(loc=loc, ip=ip).ir_value(loc=loc, ip=ip)
+    llvm.inline_asm(
+        T.i32(),
+        [mc_ptr_int, x, y, z, w],
+        "multimem.st.sys.relaxed.global.v4.f32 [$1], {$2, $3, $4, $5};",
+        "=r,l,r,r,r,r",
+        has_side_effects=True,
+        asm_dialect=0,
+        loc=loc,
+        ip=ip,
+    )

python/CuTeDSL/cutlass/utils/grouped_gemm_tile_scheduler_helper.py

@@ -171,8 +171,18 @@ class GroupedGemmTileSchedulerHelper:
     def __new_from_mlir_values__(
         self, values: List[ir.Value]
     ) -> "GroupedGemmTileSchedulerHelper":
-        tile_sched_params = new_from_mlir_values(self.tile_sched_params, values)
-        search_state = new_from_mlir_values(self.search_state, values[1:])
+        # Reconstruct tile_sched_params and determine how many values it consumed.
+        # NOTE: tile_sched_params may contain FastDivmod divisors (when swizzle_size == 1),
+        # which adds extra MLIR values.
+        params_values = extract_mlir_values(self.tile_sched_params)
+        n_params_values = len(params_values)
+        tile_sched_params = new_from_mlir_values(
+            self.tile_sched_params, values[:n_params_values]
+        )
+
+        # Reconstruct search_state from remaining values
+        search_state = new_from_mlir_values(self.search_state, values[n_params_values:])
+
         return GroupedGemmTileSchedulerHelper(
             self.group_count,
             tile_sched_params,

python/CuTeDSL/cutlass/utils/hardware_info.py

@@ -9,7 +9,9 @@
 # and related documentation outside the scope permitted by the EULA
 # is strictly prohibited.
 
-from cuda.bindings import driver, nvrtc
+from cuda.bindings import driver, nvrtc, runtime
+from cutlass.cutlass_dsl.cuda_jit_executor import CudaDialectJitModule
+from cutlass.base_dsl.common import DSLRuntimeError
 
 import cutlass.cute as cute
 
@@ -51,7 +53,7 @@ class HardwareInfo:
                 f"Cluster size must be between 1 and 32, {cluster_size} is not supported"
             )
 
-        device_fn = self._get_device_function(self.device)
+        self._get_device_function(self.device)
 
         max_shared_memory_per_block = self._checkCudaErrors(
             driver.cuDeviceGetAttribute(
@@ -173,7 +175,14 @@ class HardwareInfo:
         )
 
     # get a empty kernel to compute occupancy
-    def _get_device_function(self, device) -> None:
+    def _get_device_function(self, device) -> driver.CUfunction:
         self.compiled_kernel = cute.compile(self._host_function).to(device)
-        self.kernel = self.compiled_kernel.exec_context.kernel_functions[0]
-        self.module = self.compiled_kernel.exec_context.module.cuda_modules[0]
+        assert isinstance(self.compiled_kernel.jit_module, CudaDialectJitModule)
+        err, kernels = runtime.cudaLibraryEnumerateKernels(
+            1, self.compiled_kernel.jit_module.cuda_library[0]
+        )
+        if err is not runtime.cudaError_t.cudaSuccess:
+            raise DSLRuntimeError(f"Failed to enumerate kernels: {err}")
+        self.kernel = kernels[0]
+        self.kernel = self._checkCudaErrors(driver.cuKernelGetFunction(self.kernel))
+        return self.kernel

python/CuTeDSL/cutlass/utils/mixed_input_helpers.py

@@ -0,0 +1,507 @@
+# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+from __future__ import annotations
+
+from enum import Enum, auto
+from math import log2
+from typing import Optional, Union
+
+import cutlass
+import cutlass.cute as cute
+import cutlass.utils.blackwell_helpers as sm100_utils
+from cutlass.cute.nvgpu import cpasync, tcgen05
+from cutlass.cute.runtime import from_dlpack
+
+"""
+This file contain common utilities for mixed-input GEMM.
+"""
+
+
+class TransformMode(Enum):
+    """
+    An enumeration for the possible transform modes of a mixed-input GEMM.
+    """
+
+    ConvertOnly = auto()
+    ConvertScale = auto()
+
+
+def scale_tma_partition(
+    tCsS: cute.Tensor,
+    tCgS: cute.Tensor,
+    tma_atom_s: cute.CopyAtom,
+    block_in_cluster_coord_vmnk: cute.Coord,
+    scale_cta_layout: cute.Layout,
+) -> tuple[cute.Tensor, cute.Tensor]:
+    """
+    Perform TMA partition for scale tensor.
+    This method partitions the gobal memory and shared memory buffer for scale tensor for TMA load.
+    :param tCsS: Input scale shared memory tensor
+    :type tCsS: cute.Tensor
+    :param tCgS: Input scale global memory tensor
+    :type tCgS: cute.Tensor
+    :param tma_atom_s: TMA copy atom for scale tensor
+    :type tma_atom_s: cute.CopyAtom
+    :param block_in_cluster_coord_vmnk: CTA coord in the cluster
+    :type block_in_cluster_coord_vmnk: cute.Coord
+    :param scale_cta_layout: Layout of CTA from the view of the scale tensor
+    :type scale_cta_layout: cute.Layout
+    :return: A tuple containing (tSsS, tSgS) where:
+
+        * tSsS: Partitioned scale tensor in shared memory
+        * tSgS: Partitioned scale tensor in global memory
+
+    :rtype: tuple[cute.Tensor, cute.Tensor]
+    """
+    tSsS, tSgS = cpasync.tma_partition(
+        tma_atom_s,
+        block_in_cluster_coord_vmnk[2],
+        scale_cta_layout,
+        cute.group_modes(tCsS, 0, 3),
+        cute.group_modes(tCgS, 0, 3),
+    )
+    # Add rest_v mode
+    # ((atom_v, rest_v), STAGE)
+    # ((atom_v, rest_v), loopM, loopK, loopL)
+    tSsS = cute.make_tensor(
+        tSsS.iterator,
+        cute.make_layout(
+            ((tSsS.layout.shape[0], 1), *tSsS.layout.shape[1:]),
+            stride=(
+                (tSsS.layout.stride[0], 0),
+                *tSsS.layout.stride[1:],
+            ),
+        ),
+    )
+    tSgS = cute.make_tensor(
+        tSgS.iterator,
+        cute.make_layout(
+            ((tSgS.layout.shape[0], 1), *tSgS.layout.shape[1:]),
+            stride=(
+                (tSgS.layout.stride[0], 0),
+                *tSgS.layout.stride[1:],
+            ),
+        ),
+    )
+    return tSsS, tSgS
+
+
+def transform_partition(
+    transform_a_source: tcgen05.OperandSource,
+    scale_mode: TransformMode,
+    copy_atom_a_input: cute.CopyAtom,
+    copy_atom_a_transform: cute.CopyAtom,
+    sA_input: cute.Tensor,
+    A_transform: cute.Tensor,
+    transform_local_tidx: cutlass.Int32,
+) -> tuple[cute.TiledCopy, cute.TiledCopy, cute.Tensor, cute.Tensor]:
+    """
+    Partition tensors for transform input and output.
+    This method sets up the copy atoms and partitions the shared/tensor memory
+    for the transformation of tensor A.
+    :param transform_a_source: Where the transformed tensor A is stored (TMEM or SMEM)
+    :type transform_a_source: tcgen05.OperandSource
+    :param scale_mode: The transform mode (ConvertOnly or ConvertScale)
+    :type scale_mode: TransformMode
+    :param copy_atom_a_input: Copy atom for loading A from shared memory
+    :type copy_atom_a_input: cute.CopyAtom
+    :param copy_atom_a_transform: Copy atom for storing transformed A
+    :type copy_atom_a_transform: cute.CopyAtom
+    :param sA_input: Input tensor A in shared memory
+    :type sA_input: cute.Tensor
+    :param A_transform: Transformed tensor A in tensor or shared memory
+    :type A_transform: cute.Tensor
+    :param transform_local_tidx: Local thread index for transformation warps
+    :type transform_local_tidx: cutlass.Int32
+    :return: A tuple containing (src_copy_a, dst_copy_a, tAsA_input, tA_transform) where:
+
+        * src_copy_a: Tiled copy for source tensor
+        * dst_copy_a: Tiled copy for destination tensor
+        * tAsA_input: Partitioned input tensor A
+        * tA_transform: Partitioned transformed tensor A
+
+    :rtype: tuple[cute.TiledCopy, cute.TiledCopy, cute.Tensor, cute.Tensor]
+    """
+    if cutlass.const_expr(transform_a_source == tcgen05.OperandSource.TMEM):
+        if cutlass.const_expr(
+            cute.size(A_transform, mode=[0, 0]) == 128
+            and cute.size(sA_input, mode=[0, 0]) == 64
+        ):
+            tensor_input = cute.make_tensor(
+                sA_input.iterator,
+                cute.logical_product(
+                    sA_input.layout,
+                    ((cute.make_layout(2, stride=0), None), None, None, None),
+                ),
+            )
+        else:
+            tensor_input = sA_input
+        reg2tmem_tiled_copy = tcgen05.make_tmem_copy(
+            copy_atom_a_transform, A_transform[(None, None, None, 0)]
+        )
+        thr_reg2tmem_tiled_copy = reg2tmem_tiled_copy.get_slice(transform_local_tidx)
+        partitioned_tensor_input = thr_reg2tmem_tiled_copy.partition_S(tensor_input)
+        partitioned_tensor_transform = thr_reg2tmem_tiled_copy.partition_D(A_transform)
+        src_copy_a = (
+            cute.make_tiled_copy_S(copy_atom_a_input, reg2tmem_tiled_copy)
+            if scale_mode is TransformMode.ConvertScale
+            else None
+        )
+        dst_copy_a = reg2tmem_tiled_copy
+        tAsA_input = partitioned_tensor_input
+        tA_transform = partitioned_tensor_transform
+    elif cutlass.const_expr(transform_a_source == tcgen05.OperandSource.SMEM):
+        # Construct tiled_copy satisfying 8 contiguous elts per copy atom
+        reg2smem_tiled_copy = cute.make_cotiled_copy(
+            copy_atom_a_transform,
+            cute.make_layout((128, 8), stride=(8, 1)),
+            A_transform[(None, None, None, 0)].layout,
+        )
+        thr_reg2smem_tiled_copy = reg2smem_tiled_copy.get_slice(transform_local_tidx)
+        partitioned_tensor_input = thr_reg2smem_tiled_copy.partition_S(sA_input)
+        partitioned_tensor_transform = thr_reg2smem_tiled_copy.partition_D(A_transform)
+        src_copy_a = (
+            cute.make_tiled_copy_S(copy_atom_a_input, reg2smem_tiled_copy)
+            if scale_mode is TransformMode.ConvertScale
+            else None
+        )
+        # auto-vec copy is enough for copy from register to shared memory here
+        dst_copy_a = None
+        tAsA_input = partitioned_tensor_input
+        tA_transform = partitioned_tensor_transform
+    return src_copy_a, dst_copy_a, tAsA_input, tA_transform
+
+
+def scale_partition(
+    src_copy_a: cute.TiledCopy,
+    tCsS: cute.Tensor,
+    transform_local_tidx: cutlass.Int32,
+    mma_dtype: type[cutlass.Numeric],
+) -> tuple[cute.TiledCopy, cute.Tensor, cute.Tensor, cute.Tensor]:
+    """
+    Partition the scale tensor for transformation.
+    This method prepares the copy atom and partitions the shared memory for the scale tensor.
+    :param src_copy_a: Tiled copy for the source tensor
+    :type src_copy_a: cute.TiledCopy
+    :param tCsS: Scale tensor in shared memory
+    :type tCsS: cute.Tensor
+    :param transform_local_tidx: Local thread index for transformation warps
+    :type transform_local_tidx: cutlass.Int32
+    :param mma_dtype: Data type for the MMA operation
+    :type mma_dtype: type[cutlass.Numeric]
+    :return: A tuple containing (smem_thr_copy_S, tSsS_trans, tSrS_copy, tSrS) where:
+
+        * smem_thr_copy_S: Tiled copy for the scale tensor
+        * tSsS_trans: Partitioned scale tensor for transformation
+        * tSrS_copy: Register fragment for the scale tensor
+        * tSrS: View of scale tensor used for transformation computation
+
+    :rtype: tuple[cute.TiledCopy, cute.Tensor, cute.Tensor, cute.Tensor]
+    """
+    smem_thr_copy_S = None
+    tSsS_trans = None
+    tSrS = None
+    # Partition scale tensor
+    smem_thr_copy_S = src_copy_a.get_slice(transform_local_tidx)
+    tSsS_trans = smem_thr_copy_S.partition_S(tCsS)
+    # Construct register fragment for scale tensor
+    tSsS_layout_per_stage = tSsS_trans[(None, None, None, None, 0)].layout
+    # tSrS for copy
+    tSrS_copy = cute.make_rmem_tensor(
+        cute.filter_zeros(tSsS_layout_per_stage).shape, mma_dtype
+    )
+    # tSrS view for transformation computation
+    tSrS = cute.make_tensor(
+        tSrS_copy.iterator,
+        cute.make_layout(tSsS_layout_per_stage.shape, stride=tSrS_copy.layout.stride),
+    )
+    return smem_thr_copy_S, tSsS_trans, tSrS_copy, tSrS
+
+
+def get_gmem_layout_scale(
+    scale_shape_mkl: tuple[int, int, int],
+    scale_granularity_m: int,
+    scale_granularity_k: int,
+    scale_major_mode: tcgen05.OperandMajorMode,
+) -> cute.Layout:
+    """
+    Get the layout of the scale tensor in global memory.
+    :param scale_shape_mkl: The shape of the scale tensor (M, K, L).
+    :type scale_shape_mkl: tuple[int, int, int]
+    :return: The layout of the scale tensor in global memory.
+    :rtype: cute.Layout
+    """
+    m, k, l = scale_shape_mkl
+    shape_scale = (
+        (scale_granularity_m, cute.ceil_div(m, scale_granularity_m)),
+        (scale_granularity_k, cute.ceil_div(k, scale_granularity_k)),
+    )
+    if cutlass.const_expr(scale_major_mode == tcgen05.OperandMajorMode.MN):
+        layout_mk = cute.make_layout(
+            shape_scale,
+            stride=(
+                (0, 1),
+                (0, cute.size(shape_scale[0][1])),
+            ),
+        )
+    else:
+        layout_mk = cute.make_layout(
+            shape_scale,
+            stride=(
+                (0, cute.size(shape_scale[1][1])),
+                (0, 1),
+            ),
+        )
+    return cute.make_layout(
+        (*layout_mk.shape, l),
+        stride=(*layout_mk.stride, cute.cosize(layout_mk)),
+    )
+
+
+def get_smem_layout_scale(
+    mma_tiler: tuple[int, int, int],
+    use_2cta_instrs: bool,
+    scale_granularity_m: int,
+    scale_granularity_k: int,
+    scale_major_mode: tcgen05.OperandMajorMode,
+    a_scale_dtype: type[cutlass.Numeric],
+    num_scale_load2trans_stage: int,
+) -> tuple[tuple[int, int], cute.ComposedLayout, cute.ComposedLayout]:
+    """
+    Get the layout of the scale tensor in shared memory.
+    :return: A tuple containing (scale_tile_shape, smem_layout_scale_per_stage, smem_layout_scale) where:
+
+        * scale_tile_shape: The tile shape
+        * smem_layout_scale_per_stage: Shared memory layout for scale tensor per stage
+        * smem_layout_scale: Shared memory layout for scale tensor
+
+    :rtype: tuple[tuple[int, int], cute.ComposedLayout, cute.ComposedLayout]
+    """
+    scale_tile_shape = (
+        (cute.size(mma_tiler[0]) // 2 if use_2cta_instrs else cute.size(mma_tiler[0])),
+        cute.size(mma_tiler[2]),
+    )
+    size_mn = scale_tile_shape[0]
+    size_k = scale_tile_shape[1]
+    smem_size_mn = scale_granularity_m if scale_granularity_m < size_mn else size_mn
+    smem_size_k = scale_granularity_k if scale_granularity_k < size_k else size_k
+    div_mn = cute.ceil_div(size_mn, smem_size_mn)
+    div_k = cute.ceil_div(size_k, smem_size_k)
+    smem_atom_shape = (
+        (smem_size_mn, div_mn),
+        (smem_size_k, div_k),
+    )
+    if cutlass.const_expr(scale_major_mode == tcgen05.OperandMajorMode.MN):
+        outer_layout = cute.make_layout(
+            smem_atom_shape,
+            stride=(
+                (0, 1),
+                (0, div_mn),
+            ),
+        )
+    else:
+        outer_layout = cute.make_layout(
+            smem_atom_shape,
+            stride=(
+                (0, div_k),
+                (0, 1),
+            ),
+        )
+    # Apply a trivial swizzle to make it a composed layout, which could be used to construct TMA atom
+    smem_layout_scale_per_stage = cute.make_composed_layout(
+        cute.make_swizzle(0, 4, 3), 0, outer_layout
+    )
+    assert cute.rank(smem_layout_scale_per_stage) == 2, "Scale layout must be rank 2"
+    assert (
+        cute.size(mma_tiler[0]) % cute.size(smem_layout_scale_per_stage.outer[0]) == 0
+    ), "smem_layout_scale_per_stage must equal the tile shape."
+    assert (
+        cute.size(mma_tiler[2]) % cute.size(smem_layout_scale_per_stage.outer[1]) == 0
+    ), "smem_layout_scale_per_stage must evenly divide tile k shape."
+    # Shared memory buffer for scale must be at least 128B to satisfy TMA requirement
+
+    assert cute.size_in_bytes(a_scale_dtype, smem_layout_scale_per_stage) >= 128, (
+        "smem size for scale must be at least 128B"
+    )
+    # Scale layout in smem with multiple stages
+    smem_layout_scale = cute.append(
+        smem_layout_scale_per_stage,
+        cute.make_layout(
+            (num_scale_load2trans_stage),
+            stride=(cute.cosize(smem_layout_scale_per_stage.outer)),
+        ),
+    )
+    return scale_tile_shape, smem_layout_scale_per_stage, smem_layout_scale
+
+
+def compute_smem_layout(
+    tiled_mma: cute.TiledMma,
+    mma_tiler_mnk: tuple[int, int, int],
+    a_dtype: type[cutlass.Numeric],
+    b_dtype: type[cutlass.Numeric],
+    load2trans_stage_count: int,
+    trans2mma_stage_count: int,
+) -> tuple[
+    cute.ComposedLayout,
+    cute.ComposedLayout,
+    cute.ComposedLayout,
+]:
+    """
+    Compute shared memory layouts for tensor A, transformed A and tensor B.
+    :param tiled_mma: The tiled MMA object defining the core computation.
+    :type tiled_mma: cute.TiledMma
+    :param mma_tiler_mnk: The shape (M, N, K) of the MMA tiler.
+    :type mma_tiler_mnk: tuple[int, int, int]
+    :param a_dtype: Data type of operand A.
+    :type a_dtype: type[cutlass.Numeric]
+    :param b_dtype: Data type of operand B.
+    :type b_dtype: type[cutlass.Numeric]
+    :param load2trans_stage_count: Number of stages for load-to-transform pipeline.
+    :type load2trans_stage_count: int
+    :param trans2mma_stage_count: Number of stages for transform-to-MMA pipeline.
+    :type trans2mma_stage_count: int
+    :return: A tuple containing (smem_layout_a, smem_layout_a_transform, smem_layout_b) where:
+
+        * smem_layout_a: Shared memory layout for tensor A
+        * smem_layout_a_transform: Shared memory layout for transformed tensor A
+        * smem_layout_b: Shared memory layout for tensor B
+
+    :rtype: tuple[cute.ComposedLayout, cute.ComposedLayout, cute.ComposedLayout]
+    """
+    smem_layout_a = sm100_utils.make_smem_layout_a(
+        tiled_mma,
+        mma_tiler_mnk,
+        a_dtype,
+        load2trans_stage_count,
+    )
+    smem_layout_a_transform = sm100_utils.make_smem_layout_a(
+        tiled_mma,
+        mma_tiler_mnk,
+        tiled_mma.op.a_dtype,
+        trans2mma_stage_count,
+    )
+    smem_layout_b = sm100_utils.make_smem_layout_b(
+        tiled_mma,
+        mma_tiler_mnk,
+        b_dtype,
+        load2trans_stage_count,
+    )
+    return (smem_layout_a, smem_layout_a_transform, smem_layout_b)
+
+
+def get_transform_a_source(
+    a_major_mode: tcgen05.OperandMajorMode,
+) -> tcgen05.OperandSource:
+    """
+    Determine the operand source for transformed A tensor based on the operand major mode.
+    """
+    if cutlass.const_expr(a_major_mode == tcgen05.OperandMajorMode.K):
+        return tcgen05.OperandSource.TMEM
+    else:
+        return tcgen05.OperandSource.SMEM
+
+
+def get_tma_atom_kind(
+    mcast: cutlass.Boolean,
+    use_2cta_instrs: bool,
+    is_b: bool,
+) -> Union[cpasync.CopyBulkTensorTileG2SMulticastOp, cpasync.CopyBulkTensorTileG2SOp]:
+    """
+    Get the TMA atom kind based on 1) whether it's a multicast operation,
+    2) whether 2CTA tcgen05.mma instruction is enabled, and
+    3) whether it's a B tensor
+    """
+    # Not using .2CTA instructions for tensor A as the consumer is threads on different CTAs
+    cta_group = (
+        tcgen05.CtaGroup.TWO if (use_2cta_instrs and is_b) else tcgen05.CtaGroup.ONE
+    )
+    if cutlass.const_expr(mcast):
+        return cpasync.CopyBulkTensorTileG2SMulticastOp(cta_group)
+    return cpasync.CopyBulkTensorTileG2SOp(cta_group)
+
+
+def get_copy_atom_a_transform(
+    mma_dtype: type[cutlass.Numeric],
+    use_2cta_instrs: bool,
+    transform_a_source: tcgen05.OperandSource,
+    a_smem_shape: cute.Shape,
+    a_dtype: type[cutlass.Numeric],
+) -> cute.CopyAtom:
+    """
+    Determine the copy atom for transformed A tensor based on the operand source and tile size.
+    """
+    if cutlass.const_expr(transform_a_source == tcgen05.OperandSource.TMEM):
+        if cutlass.const_expr(
+            cute.size(a_smem_shape[0][0]) == 64 and (not use_2cta_instrs)
+        ):
+            copy_op_r2t = tcgen05.St16x256bOp(
+                tcgen05.Repetition(1), tcgen05.Unpack.NONE
+            )
+        else:
+            copy_op_r2t = tcgen05.St32x32bOp(tcgen05.Repetition(8), tcgen05.Unpack.NONE)
+        return cute.make_copy_atom(copy_op_r2t, mma_dtype)
+    else:
+        return cute.make_copy_atom(
+            cute.nvgpu.CopyUniversalOp(), a_dtype, num_bits_per_copy=32
+        )
+
+
+def is_valid_scale_granularity(
+    scale_granularity_m: int,
+    scale_granularity_k: int,
+    a_dtype: type[cutlass.Numeric],
+    k: int,
+    mma_tiler_k: int,
+) -> bool:
+    """
+    Check if the scale granularity settings are valid for the given data type and problem size.
+    """
+    if a_dtype.width == 8:
+        # No scale tensor for 8bit data type A
+        if not (scale_granularity_m == 0 and scale_granularity_k == 0):
+            return False
+    elif a_dtype.width == 4:
+        if scale_granularity_m != 1 or (
+            scale_granularity_k == 0
+            or k % scale_granularity_k != 0
+            or scale_granularity_k % mma_tiler_k != 0
+        ):
+            return False
+    return True
+
+
+def get_divisibility(contiguous_dim_size: int, upper_bound: int = 128) -> int:
+    """
+    Calculate the largest power of 2 divisibility factor for memory alignment.
+    """
+    # Check the largest power of 2 factor of contiguous_dim_size
+    for i in range(int(log2(contiguous_dim_size)), 0, -1):
+        if contiguous_dim_size % (2**i) == 0:
+            return min(2**i, upper_bound)
+    return 1

python/CuTeDSL/cutlass/utils/print_latex.py

@@ -0,0 +1,165 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 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/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.
+
+from typing import Callable, Union
+
+from ..cute import (
+    Layout,
+    ComposedLayout,
+    append,
+    is_static,
+    make_layout,
+    size,
+    product_each,
+    rank,
+)
+from ..cute.typing import IntTuple
+
+__all__ = ["print_latex", "print_latex_tv"]
+
+
+def tikz_color_bwx8(idx: int):
+    color_map = [
+        "black!00",
+        "black!40",
+        "black!20",
+        "black!60",
+        "black!10",
+        "black!50",
+        "black!30",
+        "black!70",
+    ]
+    return color_map[idx % 8]
+
+
+def tikz_color_white(idx: int):
+    return "white"
+
+
+def tikz_color_tv(tid: int, vid: int):
+    color_map = [
+        "{rgb,255:red,175;green,175;blue,255}",
+        "{rgb,255:red,175;green,255;blue,175}",
+        "{rgb,255:red,255;green,255;blue,175}",
+        "{rgb,255:red,255;green,175;blue,175}",
+        "{rgb,255:red,210;green,210;blue,255}",
+        "{rgb,255:red,210;green,255;blue,210}",
+        "{rgb,255:red,255;green,255;blue,210}",
+        "{rgb,255:red,255;green,210;blue,210}",
+    ]
+    return color_map[tid % 8]
+
+
+def print_latex(x: Union[Layout, ComposedLayout], *, color: Callable = tikz_color_bwx8):
+    """
+    Prints a layout.
+    :param x: A layout
+    :type x: Union[Layout, ComposedLayout]
+    :param color: A function that returns TiKZ colors
+    :type color: Callable
+    """
+
+    if not is_static(x):
+        raise ValueError("Requires static input")
+    if rank(x) > 2:
+        raise ValueError("Requires rank <= 2 to print")
+
+    if rank(x) == 1:
+        layout = append(x, make_layout(1, stride=0))
+    else:
+        layout = x
+
+    print("%% Layout: {}", layout)
+    print("\\documentclass[convert]{standalone}")
+    print("\\usepackage{tikz}")
+    print("\\begin{document}")
+    print(
+        "\\begin{tikzpicture}[x={(0cm,-1cm)},y={(1cm,0cm)},every node/.style={minimum size=1cm, outer sep=0pt}]"
+    )
+
+    M, N = product_each(x.shape)
+
+    for m in range(M):
+        for n in range(N):
+            idx = layout((m, n))
+            print("\\node[fill=")
+            print(color(idx))
+            print("] at (%d,%d) {%d};\n" % (m, n, idx))
+    print(
+        "\\draw[color=black,thick,shift={(-0.5,-0.5)}] (0,0) grid (%d,%d);\n\n" % (M, N)
+    )
+    for m in range(M):
+        print("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n" % (m, -1, m))
+    for n in range(N):
+        print("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n" % (-1, n, n))
+
+    ## Footer
+    print("\\end{tikzpicture}")
+    print("\\end{document}")
+
+
+def print_latex_tv(
+    layout_tv: Union[Layout, ComposedLayout],
+    tile_mn: Union[IntTuple, Layout],
+    *,
+    color: Callable = tikz_color_tv,
+):
+    """
+    Prints a tv layout for a tile M N. Everything must be static.
+    :param layout_tv: A static thread value layout
+    :type layout_tv: Union[Layout, ComposedLayout]
+    :param tile_mn: A static M N tile
+    :type tile_mn: Union[IntTuple, Layout]
+    :param color: A function that returns TiKZ colors
+    :type color: Callable
+    """
+    if not is_static(layout_tv) or not is_static(tile_mn):
+        raise ValueError("Layout tv and tile_mn must be static")
+    if rank(layout_tv) != 2:
+        raise ValueError("Require layout_tv to be rank 2")
+
+    print("%% Layout TV: {}", layout_tv)
+    print("\\documentclass[convert]{standalone}")
+    print("\\usepackage{tikz}")
+    print("\\begin{document}")
+    print(
+        "\\begin{tikzpicture}[x={(0cm,-1cm)},y={(1cm,0cm)},every node/.style={minimum size=1cm, outer sep=0pt}]\n"
+    )
+
+    if not isinstance(tile_mn, Layout):
+        tile_mn = make_layout(tile_mn)
+
+    M, N = product_each(tile_mn.shape)
+    filled = [[False for n in range(N)] for m in range(M)]
+
+    for tid in range(size(layout_tv, mode=[0])):
+        for vid in range(size(layout_tv, mode=[1])):
+            idx = layout_tv((tid, vid))
+            m = (idx // tile_mn.stride[0]) % tile_mn.shape[0]
+            n = (idx // tile_mn.stride[1]) % tile_mn.shape[1]
+            if not filled[m][n]:
+                filled[m][n] = True
+                print(
+                    "\\node[fill=%s] at (%d,%d) {\\shortstack{T%d \\\\ V%d}};\n"
+                    % (color(tid, vid), m, n, tid, vid)
+                )
+
+    print(
+        "\\draw[color=black,thick,shift={(-0.5,-0.5)}] (0,0) grid (%d,%d);\n\n" % (M, N)
+    )
+    for m in range(M):
+        print("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n" % (m, -1, m))
+    for n in range(N):
+        print("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n" % (-1, n, n))
+
+    ## Footer
+    print("\\end{tikzpicture}")
+    print("\\end{document}")

python/CuTeDSL/cutlass/utils/smem_allocator.py

@@ -10,6 +10,7 @@
 # is strictly prohibited.
 
 from typing import Optional, Type, Union, overload
+import inspect
 
 import cutlass.cute as cute
 from cutlass.cute.arch import get_dyn_smem, get_dyn_smem_size
@@ -285,4 +286,11 @@ class SmemAllocator:
         return cute.make_tensor(ptr, layout, loc=loc, ip=ip)
 
 
+# Set explicit signature for Sphinx documentation to avoid issues with @dsl_user_op decorator
+SmemAllocator.__init__.__signature__ = inspect.Signature(
+    [
+        inspect.Parameter("self", inspect.Parameter.POSITIONAL_OR_KEYWORD),
+    ]
+)
+
 get_smem_capacity_in_bytes = SmemAllocator.capacity_in_bytes

python/CuTeDSL/cutlass/utils/static_persistent_tile_scheduler.py

@@ -19,6 +19,7 @@ from cutlass.cutlass_dsl import (
     extract_mlir_values,
     new_from_mlir_values,
     dsl_user_op,
+    const_expr,
 )
 from cutlass._mlir import ir
 import cutlass.cute as cute
@@ -135,6 +136,7 @@ class PersistentTileSchedulerParams:
             ip=ip,
         )
 
+        # Apply swizzle if swizzle_size > 1
         if swizzle_size > 1:
             problem_shape_ncluster_mnl = cute.round_up(
                 self.problem_layout_ncluster_mnl.shape,
@@ -172,6 +174,35 @@ class PersistentTileSchedulerParams:
                     ip=ip,
                 )
 
+        # Create FastDivmod divisors (only when swizzle_size == 1 for correctness)
+        # FastDivmod assumes simple col-major layout, incompatible with swizzled layouts
+        if swizzle_size == 1:
+            problem_layout_size = cute.size(
+                self.problem_layout_ncluster_mnl, loc=loc, ip=ip
+            )
+            cluster_count_m = self.problem_layout_ncluster_mnl.shape[0]
+            cluster_count_n = self.problem_layout_ncluster_mnl.shape[1]
+
+            # batch_fdd: Used to map linear_idx to work_unit_id (handles persistent scheduling)
+            self.batch_fdd = cute.fast_divmod_create_divisor(
+                problem_layout_size, loc=loc, ip=ip
+            )
+
+            # cluster_shape_m_fdd: Used to decode work_unit_id to cluster coordinates
+            self.cluster_shape_m_fdd = cute.fast_divmod_create_divisor(
+                cluster_count_m, loc=loc, ip=ip
+            )
+
+            # cluster_shape_n_fdd: Used for the second level decomposition
+            self.cluster_shape_n_fdd = cute.fast_divmod_create_divisor(
+                cluster_count_n, loc=loc, ip=ip
+            )
+        else:
+            # FastDivmod not applicable with swizzling, set to None
+            self.batch_fdd = None
+            self.cluster_shape_m_fdd = None
+            self.cluster_shape_n_fdd = None
+
     def __extract_mlir_values__(self):
         values, self._values_pos = [], []
         for obj in [
@@ -183,10 +214,38 @@ class PersistentTileSchedulerParams:
             obj_values = extract_mlir_values(obj)
             values += obj_values
             self._values_pos.append(len(obj_values))
+
+        # Add FastDivmod divisors to MLIR values for Host->Device transfer
+        # Only add non-None values to avoid MLIR type errors
+        fastdivmod_values = []
+        fastdivmod_indices = []  # Track which FastDivmod objects are present
+
+        for i, (fdd_name, fdd_obj) in enumerate(
+            [
+                ("batch_fdd", self.batch_fdd),
+                ("cluster_shape_m_fdd", self.cluster_shape_m_fdd),
+                ("cluster_shape_n_fdd", self.cluster_shape_n_fdd),
+            ]
+        ):
+            if fdd_obj is not None:
+                # Extract MLIR values from FastDivmodDivisor objects
+                fdd_values = extract_mlir_values(fdd_obj)
+                fastdivmod_values.extend(fdd_values)
+                fastdivmod_indices.append(i)
+
+        values += fastdivmod_values
+        self._values_pos.append(
+            len(fastdivmod_indices)
+        )  # Store count of FastDivmod objects, not values
+        self._fastdivmod_indices = fastdivmod_indices  # Store for reconstruction
+
         return values
 
     def __new_from_mlir_values__(self, values):
         obj_list = []
+        values_copy = list(values)  # Make a copy to avoid modifying original
+
+        # Reconstruct original objects from MLIR values
         for obj, n_items in zip(
             [
                 self.problem_shape_ntile_mnl,
@@ -194,11 +253,32 @@ class PersistentTileSchedulerParams:
                 self.swizzle_size,
                 self._raster_along_m,
             ],
-            self._values_pos,
+            self._values_pos[:-1],  # Exclude FastDivmod count
         ):
-            obj_list.append(new_from_mlir_values(obj, values[:n_items]))
-            values = values[n_items:]
-        return PersistentTileSchedulerParams(*(tuple(obj_list)), loc=self._loc)
+            obj_list.append(new_from_mlir_values(obj, values_copy[:n_items]))
+            values_copy = values_copy[n_items:]
+
+        # Create new params object by calling __init__ with reconstructed values
+        # This properly recreates layouts and other derived attributes in the device context
+        new_params = PersistentTileSchedulerParams(*(tuple(obj_list)), loc=self._loc)
+
+        # Restore FastDivmod divisors from remaining values
+        fdd_names = ["batch_fdd", "cluster_shape_m_fdd", "cluster_shape_n_fdd"]
+
+        if hasattr(self, "_fastdivmod_indices") and len(self._fastdivmod_indices) > 0:
+            # Override the FastDivmod divisors created by __init__ with reconstructed ones
+            for j, original_index in enumerate(self._fastdivmod_indices):
+                fdd_name = fdd_names[original_index]
+                # Get the original FastDivmodDivisor object
+                original_fdd = getattr(self, fdd_name)
+                if original_fdd is not None and j < len(values_copy):
+                    # Each FastDivmodDivisor has 1 MLIR value
+                    reconstructed_fdd = new_from_mlir_values(
+                        original_fdd, [values_copy[j]]
+                    )
+                    setattr(new_params, fdd_name, reconstructed_fdd)
+
+        return new_params
 
     @dsl_user_op
     def get_grid_shape(
@@ -285,12 +365,16 @@ class StaticPersistentTileScheduler:
         values.extend(extract_mlir_values(self._current_work_linear_idx))
         values.extend(extract_mlir_values(self.cta_id_in_cluster))
         values.extend(extract_mlir_values(self._num_tiles_executed))
+
+        # CRITICAL: Also extract FastDivmod divisors from params
+        values.extend(extract_mlir_values(self.params))
+
         return values
 
     def __new_from_mlir_values__(
         self, values: list[ir.Value]
     ) -> "StaticPersistentTileScheduler":
-        assert len(values) == 6
+        assert len(values) >= 6
         new_num_persistent_clusters = new_from_mlir_values(
             self.num_persistent_clusters, [values[0]]
         )
@@ -303,8 +387,13 @@ class StaticPersistentTileScheduler:
         new_num_tiles_executed = new_from_mlir_values(
             self._num_tiles_executed, [values[5]]
         )
+
+        # Reconstruct params with FastDivmod divisors
+        params_values = values[6:]  # Remaining values are from params
+        new_params = new_from_mlir_values(self.params, params_values)
+
         return StaticPersistentTileScheduler(
-            self.params,
+            new_params,  # Use reconstructed params with FastDivmod divisors
             new_num_persistent_clusters,
             new_current_work_linear_idx,
             new_cta_id_in_cluster,
@@ -403,11 +492,14 @@ class StaticPersistentTileScheduler:
             self.params.problem_layout_ncluster_mnl, loc=loc, ip=ip
         )
 
+        # Choose coordinate calculation method based on swizzle configuration
         if self.params.swizzle_size == 1:
-            cur_cluster_coord = self.params.problem_layout_ncluster_mnl.get_hier_coord(
+            # Use FastDivmod optimization for non-swizzled layouts
+            cur_cluster_coord = self._get_cluster_work_idx_with_fastdivmod(
                 current_work_linear_idx, loc=loc, ip=ip
             )
         else:
+            # Use get_flat_coord for swizzled layouts (FastDivmod doesn't support them)
             cur_cluster_coord = self.params.problem_layout_ncluster_mnl.get_flat_coord(
                 current_work_linear_idx, loc=loc, ip=ip
             )
@@ -424,6 +516,40 @@ class StaticPersistentTileScheduler:
 
         return WorkTileInfo(cur_tile_coord, is_valid)
 
+    def _get_cluster_work_idx_with_fastdivmod(
+        self, current_work_linear_idx: Int32, *, loc=None, ip=None
+    ) -> Tuple[Int32, Int32, Int32]:
+        """
+        FastDivmod optimized CLUSTER coordinate calculation.
+
+        CRITICAL: This should mimic problem_layout_ncluster_mnl.get_hier_coord()
+        which returns CLUSTER coordinates, not tile coordinates!
+
+        :param current_work_linear_idx: Linear index in the work space
+        :type current_work_linear_idx: Int32
+        :return: Cluster coordinates (m, n, l) or None if FastDivmod not available
+        :rtype: Tuple[Int32, Int32, Int32] or None
+        """
+
+        # Step 1: Handle persistent scheduling - map linear_idx to work_unit_id
+        work_iteration, work_unit_id = divmod(
+            current_work_linear_idx, self.params.batch_fdd
+        )
+
+        # Step 2: Decode work_unit_id using FastDivmod objects
+        # The layout structure is: problem_layout_ncluster_mnl has shape (cluster_count_m, cluster_count_n, batch_count)
+        # work_unit_id needs to be decomposed into (batch_l, cluster_n, cluster_m) in little-endian order
+
+        # First, get cluster_m using cluster_shape_m_fdd
+        cluster_n_batch, cluster_m = divmod(
+            work_unit_id, self.params.cluster_shape_m_fdd
+        )
+
+        # Then decode cluster_n_batch to get cluster_n and batch_l using FastDivmod
+        batch_l, cluster_n = divmod(cluster_n_batch, self.params.cluster_shape_n_fdd)
+
+        return (cluster_m, cluster_n, batch_l)
+
     @dsl_user_op
     def get_current_work(self, *, loc=None, ip=None) -> WorkTileInfo:
         return self._get_current_work_for_linear_idx(
@@ -446,3 +572,176 @@ class StaticPersistentTileScheduler:
         return self._num_tiles_executed
 
 
+class StaticPersistentRuntimeTileScheduler(StaticPersistentTileScheduler):
+    """A scheduler for static persistent runtime tile execution in CUTLASS/CuTe kernels.
+    This scheduler will always launch all the SMs and the scheduler will generate the real tile info for each SM.
+
+    :ivar params: Tile schedule related params, including cluster shape and problem_layout_ncluster_mnl
+    :type params: PersistentTileSchedulerParams
+    :ivar num_persistent_clusters: Number of persistent clusters that can be launched
+    :type num_persistent_clusters: Int32
+    :ivar cta_id_in_cluster: ID of the CTA within its cluster
+    :type cta_id_in_cluster: cute.Coord
+    :ivar _num_tiles_executed: Counter for executed tiles
+    :type _num_tiles_executed: Int32
+    :ivar _current_work_linear_idx: Current cluster index
+    :type _current_work_linear_idx: Int32
+    """
+
+    def __init__(
+        self,
+        params: PersistentTileSchedulerParams,
+        num_persistent_clusters: Int32,
+        current_work_linear_idx: Int32,
+        cta_id_in_cluster: cute.Coord,
+        num_tiles_executed: Int32,
+        inner_mode: int = 1,
+    ):
+        """
+        Initializes the StaticPersistentTileScheduler with the given parameters.
+
+        :param params: Tile schedule related params, including cluster shape and problem_layout_ncluster_mnl.
+        :type params: PersistentTileSchedulerParams
+        :param num_persistent_clusters: Number of persistent clusters that can be launched.
+        :type num_persistent_clusters: Int32
+        :param current_work_linear_idx: Current cluster index.
+        :type current_work_linear_idx: Int32
+        :param cta_id_in_cluster: ID of the CTA within its cluster.
+        :type cta_id_in_cluster: cute.Coord
+        :param num_tiles_executed: Counter for executed tiles.
+        :type num_tiles_executed: Int32
+        :param inner_mode: The inner mode along which the linear index will be decomposed first.
+        :type inner_mode: int
+        """
+        super().__init__(
+            params,
+            num_persistent_clusters,
+            current_work_linear_idx,
+            cta_id_in_cluster,
+            num_tiles_executed,
+        )
+        if inner_mode not in [0, 1]:
+            raise ValueError(
+                f"inner_mode must be 0(for M mode) or 1(for N mode), but got {inner_mode}"
+            )
+        self.inner_mode = inner_mode
+
+    def __new_from_mlir_values__(
+        self, values: list[ir.Value]
+    ) -> "StaticPersistentRuntimeTileScheduler":
+        assert len(values) >= 6
+        new_num_persistent_clusters = new_from_mlir_values(
+            self.num_persistent_clusters, [values[0]]
+        )
+        new_current_work_linear_idx = new_from_mlir_values(
+            self._current_work_linear_idx, [values[1]]
+        )
+        new_cta_id_in_cluster = new_from_mlir_values(
+            self.cta_id_in_cluster, values[2:5]
+        )
+        new_num_tiles_executed = new_from_mlir_values(
+            self._num_tiles_executed, [values[5]]
+        )
+
+        # Reconstruct params with FastDivmod divisors (same as parent class)
+        params_values = values[6:]  # Remaining values are from params
+        new_params = new_from_mlir_values(self.params, params_values)
+
+        return StaticPersistentRuntimeTileScheduler(
+            new_params,  # Use reconstructed params with FastDivmod divisors
+            new_num_persistent_clusters,
+            new_current_work_linear_idx,
+            new_cta_id_in_cluster,
+            new_num_tiles_executed,
+            self.inner_mode,
+        )
+
+    @staticmethod
+    @dsl_user_op
+    def create(
+        params: PersistentTileSchedulerParams,
+        block_idx: Tuple[Integer, Integer, Integer],
+        grid_dim: Tuple[Integer, Integer, Integer],
+        inner_mode: int = 1,
+        *,
+        loc=None,
+        ip=None,
+    ):
+        """Initialize the static persistent tile scheduler.
+
+        :param params: Parameters for the persistent
+            tile scheduler.
+        :type params: PersistentTileSchedulerParams
+        :param block_idx: The 3d block index in the format (bidx, bidy, bidz).
+        :type block_idx: Tuple[Integer, Integer, Integer]
+        :param grid_dim: The 3d grid dimensions for kernel launch.
+        :type grid_dim: Tuple[Integer, Integer, Integer]
+        :param inner_mode: The inner mode along which the linear index will be decomposed first.
+        :type inner_mode: int
+
+        :return: A StaticPersistentRuntimeTileScheduler object.
+        :rtype: StaticPersistentRuntimeTileScheduler
+        """
+
+        # Calculate the number of persistent clusters by dividing the total grid size
+        # by the number of CTAs per cluster
+        num_persistent_clusters = cute.size(grid_dim, loc=loc, ip=ip) // cute.size(
+            params.cluster_shape_mn, loc=loc, ip=ip
+        )
+
+        bidx, bidy, bidz = block_idx
+
+        # Initialize workload index equals to the cluster index in the grid
+        current_work_linear_idx = Int32(bidz)
+
+        # CTA id in the cluster
+        cta_id_in_cluster = (
+            Int32(bidx % params.cluster_shape_mn[0]),
+            Int32(bidy % params.cluster_shape_mn[1]),
+            Int32(0),
+        )
+        # Initialize number of tiles executed to zero
+        num_tiles_executed = Int32(0)
+        return StaticPersistentRuntimeTileScheduler(
+            params,
+            num_persistent_clusters,
+            current_work_linear_idx,
+            cta_id_in_cluster,
+            num_tiles_executed,
+            inner_mode,
+        )
+
+    # private method
+    def _get_current_work_for_linear_idx(
+        self, current_work_linear_idx: Int32, *, loc=None, ip=None
+    ) -> WorkTileInfo:
+        """Compute current tile coord given current_work_linear_idx and cta_id_in_cluster.
+
+        :param current_work_linear_idx: The linear index of the current work.
+        :type current_work_linear_idx: Int32
+
+        :return: An object containing information about the current tile coordinates
+            and validity status.
+        :rtype: WorkTileInfo
+        """
+        ntile_shape = self.params.problem_layout_ncluster_mnl.shape
+        int_max = 2147483647
+        if const_expr(self.inner_mode == 1):
+            ntile_layout = cute.make_layout(
+                (int_max, ntile_shape[1]), stride=(ntile_shape[1], 1)
+            )
+        else:
+            ntile_layout = cute.make_layout(
+                (ntile_shape[0], int_max), stride=(1, ntile_shape[0])
+            )
+        cluster_tile_coord_mn = ntile_layout.get_hier_coord(current_work_linear_idx)
+        cur_tile_coord = (
+            cluster_tile_coord_mn[0],
+            cluster_tile_coord_mn[1],
+            Int32(0),
+        )
+
+        # it is determined by kernel implementation
+        is_valid = True
+
+        return WorkTileInfo(cur_tile_coord, is_valid)

python/CuTeDSL/cutlass/utils/tmem_allocator.py

@@ -10,6 +10,7 @@
 # is strictly prohibited.
 
 from typing import Optional, Type
+import inspect
 
 from cutlass import const_expr
 from cutlass.cutlass_dsl import (
@@ -17,6 +18,7 @@ from cutlass.cutlass_dsl import (
     Float32,
     extract_mlir_values,
     new_from_mlir_values,
+    dsl_user_op,
 )
 import cutlass.pipeline as pipeline
 import cutlass.cute as cute
@@ -43,22 +45,27 @@ class TmemAllocator:
     :type _two_cta_tmem_dealloc_mbar_ptr: cute.Pointer
     """
 
+    @dsl_user_op
     @cute.jit
-    def _init_dealloc_mbarrier(self):
+    def _init_dealloc_mbarrier(self, *, loc=None, ip=None):
         assert self._two_cta_tmem_dealloc_mbar_ptr is not None, (
             "two_cta_tmem_dealloc_mbar_ptr is required for two cta"
         )
-        warp_idx = cute.arch.warp_idx()
-        warp_idx = cute.arch.make_warp_uniform(warp_idx)
+        warp_idx = cute.arch.warp_idx(loc=loc, ip=ip)
+        warp_idx = cute.arch.make_warp_uniform(warp_idx, loc=loc, ip=ip)
         _is_allocator_warp = warp_idx == self._allocator_warp_id
         if _is_allocator_warp:
             num_tmem_dealloc_threads = 32
-            with cute.arch.elect_one():
+            with cute.arch.elect_one(loc=loc, ip=ip):
                 cute.arch.mbarrier_init(
-                    self._two_cta_tmem_dealloc_mbar_ptr, num_tmem_dealloc_threads
+                    self._two_cta_tmem_dealloc_mbar_ptr,
+                    num_tmem_dealloc_threads,
+                    loc=loc,
+                    ip=ip,
                 )
         cute.arch.mbarrier_init_fence()
 
+    @dsl_user_op
     def __init__(
         self,
         alloc_result_dst_smem_ptr: cute.Pointer,
@@ -67,11 +74,40 @@ class TmemAllocator:
         is_two_cta: bool = False,
         num_allocated_columns: int = 0,
         two_cta_tmem_dealloc_mbar_ptr: Optional[cute.Pointer] = None,
+        *,
+        loc=None,
+        ip=None,
     ):
-        """Initialize the TmemAllocator instance.
+        """
+        Initialize a TmemAllocator instance for managing tensor memory on Blackwell GPUs.
 
-        Sets up the allocator state by initializing smem pointer that holds the base address of allocated tensor memory, allocator warp id, whether it is for two cta, number of allocated columns, and barrier for retrieving tensor memory ptr.
-        Meanwhile, it also initializes the mbarrier pointer for two cta deallocation case.
+        This initializer sets up the allocator's state, including the shared memory (smem) pointer
+        holding the base address of the allocated tensor memory, barrier synchronization for
+        retrieving the tensor memory pointer, allocator warp ID, whether the allocator is being used
+        for a 2-SM configuration, number of allocated columns in tensor
+        memory, and the optional mbarrier pointer for deallocation in the 2-SM case.
+
+        If `is_two_cta` is set to True, this will initialize the mbarrier pointer required for tensor
+        memory deallocation across two CTAs.
+
+        :param alloc_result_dst_smem_ptr: The shared memory pointer that holds the base address of allocated tensor memory.
+        :type alloc_result_dst_smem_ptr: cute.Pointer
+        :param barrier_for_retrieve: The named barrier for retrieving the tensor memory pointer.
+        :type barrier_for_retrieve: pipeline.NamedBarrier
+        :param allocator_warp_id: The warp ID of the allocator warp, defaults to 0.
+        :type allocator_warp_id: int, optional
+        :param is_two_cta: Whether the allocator should coordinate two CTAs, defaults to False.
+        :type is_two_cta: bool, optional
+        :param num_allocated_columns: The number of columns allocated in tensor memory, defaults to 0.
+        :type num_allocated_columns: int, optional
+        :param two_cta_tmem_dealloc_mbar_ptr: The mbarrier pointer required for two-CTA tensor memory deallocation, optional.
+        :type two_cta_tmem_dealloc_mbar_ptr: cute.Pointer, optional
+        :param loc: Optional codegen location for debugging and error reporting.
+        :type loc: Any, optional
+        :param ip: Optional insertion point for codegen.
+        :type ip: Any, optional
+
+        :raises AssertionError: If two_cta_tmem_dealloc_mbar_ptr is None while is_two_cta is True.
         """
         # TODO: automatically maintain a smem address
         self._alloc_result_dst_smem_ptr = alloc_result_dst_smem_ptr
@@ -83,7 +119,7 @@ class TmemAllocator:
 
         # Init tmem dealloc mbarrier if two cta
         if const_expr(self._is_two_cta):
-            self._init_dealloc_mbarrier()
+            self._init_dealloc_mbarrier(loc=loc, ip=ip)
 
     def __extract_mlir_values__(self) -> list[ir.Value]:
         values = extract_mlir_values(self._alloc_result_dst_smem_ptr)
@@ -137,8 +173,9 @@ class TmemAllocator:
             return False
         return True
 
+    @dsl_user_op
     @cute.jit
-    def allocate(self, num_columns: int):
+    def allocate(self, num_columns: int, *, loc=None, ip=None):
         """Allocate a block of tensor memory.
 
         This method allocates a block of tensor memory from allocator warp and returns a handle to retrieve
@@ -152,29 +189,34 @@ class TmemAllocator:
             "total allocated columns must be less than or equal to 512"
         )
 
-        warp_idx = cute.arch.warp_idx()
-        warp_idx = cute.arch.make_warp_uniform(warp_idx)
+        warp_idx = cute.arch.warp_idx(loc=loc, ip=ip)
+        warp_idx = cute.arch.make_warp_uniform(warp_idx, loc=loc, ip=ip)
         _is_allocator_warp = warp_idx == self._allocator_warp_id
         if _is_allocator_warp:
             cute.arch.alloc_tmem(
                 num_columns,
                 self._alloc_result_dst_smem_ptr,
                 is_two_cta=self._is_two_cta,
+                loc=loc,
+                ip=ip,
             )
         self._num_allocated_columns += num_columns
 
-    @cute.jit
-    def wait_for_alloc(self):
+    @dsl_user_op
+    def wait_for_alloc(self, *, loc=None, ip=None):
         """Wait for the allocator warp to finish allocation.
 
         This method is used to synchronize the allocator warp with the other warps before retrieving tmem ptr.
         """
-        self._barrier_for_retrieve.arrive_and_wait()
+        self._barrier_for_retrieve.arrive_and_wait(loc=loc, ip=ip)
 
-    @cute.jit
+    @dsl_user_op
     def retrieve_ptr(
         self,
         dtype: Type[Numeric] = Float32,
+        *,
+        loc=None,
+        ip=None,
     ) -> cute.Pointer:
         """Retrieve the pointer to the allocated tensor memory.
 
@@ -185,30 +227,36 @@ class TmemAllocator:
             dtype,
             alignment=16,
             ptr_to_buffer_holding_addr=self._alloc_result_dst_smem_ptr,
+            loc=loc,
+            ip=ip,
         )
 
+    @dsl_user_op
     @cute.jit
-    def relinquish_alloc_permit(self):
+    def relinquish_alloc_permit(self, *, loc=None, ip=None):
         """Relinquish the tensor memory allocation permit.
 
         This method relinquishes the tensor memory allocation permit for the allocator warp, promising
         the allocator warp will not allocate any more tensor memory.
         """
-        warp_idx = cute.arch.warp_idx()
-        warp_idx = cute.arch.make_warp_uniform(warp_idx)
+        warp_idx = cute.arch.warp_idx(loc=loc, ip=ip)
+        warp_idx = cute.arch.make_warp_uniform(warp_idx, loc=loc, ip=ip)
         _is_allocator_warp = warp_idx == self._allocator_warp_id
         if _is_allocator_warp:
-            cute.arch.relinquish_tmem_alloc_permit(is_two_cta=self._is_two_cta)
+            cute.arch.relinquish_tmem_alloc_permit(
+                is_two_cta=self._is_two_cta, loc=loc, ip=ip
+            )
 
+    @dsl_user_op
     @cute.jit
-    def free(self, tmem_ptr: cute.Pointer, num_columns: int = 0):
+    def free(self, tmem_ptr: cute.Pointer, num_columns: int = 0, *, loc=None, ip=None):
         """Deallocate the tensor memory.
 
         This method sync on mbarrier (for two cta use case) and deallocates the tensor memory from the allocator warp.
         User can optionally specify the number of columns to deallocate. If not specified, all allocated columns will be deallocated.
         """
-        warp_idx = cute.arch.warp_idx()
-        warp_idx = cute.arch.make_warp_uniform(warp_idx)
+        warp_idx = cute.arch.warp_idx(loc=loc, ip=ip)
+        warp_idx = cute.arch.make_warp_uniform(warp_idx, loc=loc, ip=ip)
         _is_allocator_warp = warp_idx == self._allocator_warp_id
 
         assert num_columns <= self._num_allocated_columns, (
@@ -223,15 +271,67 @@ class TmemAllocator:
         self._num_allocated_columns -= num_deallocate_columns
         if _is_allocator_warp:
             if const_expr(self._is_two_cta):
+                bid_in_cluster = cute.arch.block_idx_in_cluster(loc=loc, ip=ip)
                 _cta_rank_in_cluster = cute.arch.make_warp_uniform(
-                    cute.arch.block_idx_in_cluster()
+                    bid_in_cluster, loc=loc, ip=ip
                 )
                 # Arrive and wait for dealloc signal from peer cta
                 cute.arch.mbarrier_arrive(
-                    self._two_cta_tmem_dealloc_mbar_ptr, _cta_rank_in_cluster ^ 1
+                    self._two_cta_tmem_dealloc_mbar_ptr,
+                    _cta_rank_in_cluster ^ 1,
+                    loc=loc,
+                    ip=ip,
+                )
+                cute.arch.mbarrier_wait(
+                    self._two_cta_tmem_dealloc_mbar_ptr, 0, loc=loc, ip=ip
                 )
-                cute.arch.mbarrier_wait(self._two_cta_tmem_dealloc_mbar_ptr, 0)
             # Deallocate tmem
             cute.arch.dealloc_tmem(
-                tmem_ptr, num_deallocate_columns, is_two_cta=self._is_two_cta
+                tmem_ptr,
+                num_deallocate_columns,
+                is_two_cta=self._is_two_cta,
+                loc=loc,
+                ip=ip,
             )
+
+
+# Set explicit signature for Sphinx documentation to avoid issues with @dsl_user_op decorator
+TmemAllocator.__init__.__signature__ = inspect.Signature(
+    [
+        inspect.Parameter("self", inspect.Parameter.POSITIONAL_OR_KEYWORD),
+        inspect.Parameter(
+            "alloc_result_dst_smem_ptr",
+            inspect.Parameter.POSITIONAL_OR_KEYWORD,
+            annotation=cute.Pointer,
+        ),
+        inspect.Parameter(
+            "barrier_for_retrieve",
+            inspect.Parameter.POSITIONAL_OR_KEYWORD,
+            annotation=pipeline.NamedBarrier,
+        ),
+        inspect.Parameter(
+            "allocator_warp_id",
+            inspect.Parameter.POSITIONAL_OR_KEYWORD,
+            default=0,
+            annotation=int,
+        ),
+        inspect.Parameter(
+            "is_two_cta",
+            inspect.Parameter.POSITIONAL_OR_KEYWORD,
+            default=False,
+            annotation=bool,
+        ),
+        inspect.Parameter(
+            "num_allocated_columns",
+            inspect.Parameter.POSITIONAL_OR_KEYWORD,
+            default=0,
+            annotation=int,
+        ),
+        inspect.Parameter(
+            "two_cta_tmem_dealloc_mbar_ptr",
+            inspect.Parameter.POSITIONAL_OR_KEYWORD,
+            default=None,
+            annotation=Optional[cute.Pointer],
+        ),
+    ]
+)