v4.5 dev update. (#3153)

python/CuTeDSL/cutlass/__init__.py

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

python/CuTeDSL/cutlass/base_dsl/__init__.py

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

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

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

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

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

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

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

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

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

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

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA
@@ -14,10 +14,48 @@ This module provides MLIR's OP helper functions
 """
 
 import inspect
+import os
+import types
 from functools import wraps
 
 from ..._mlir import ir
 from ..common import DSLRuntimeError
+from ..utils.stacktrace import walk_to_top_module
+
+
+# The DSL package root is empty by default.
+_DSL_PACKAGE_ROOT = ""
+
+
+def _is_framework_frame(filename: str) -> bool:
+    """Check if a frame's filename belongs to DSL library code."""
+    global _DSL_PACKAGE_ROOT
+    if _DSL_PACKAGE_ROOT == "":
+        # Compute the DSL package root once
+        # Any frame whose file starts with this prefix is considered DSL library code.
+        _DSL_PACKAGE_ROOT = walk_to_top_module(
+            os.path.dirname(os.path.abspath(__file__))
+        )
+
+    if _DSL_PACKAGE_ROOT is None:
+        return False
+
+    return os.path.abspath(filename).startswith(_DSL_PACKAGE_ROOT)
+
+
+def _find_user_frame(start_frame: types.FrameType | None) -> types.FrameType | None:
+    """Walk up the call stack from start_frame to find the first user (non-library) frame.
+
+    Returns the first frame whose file is not under the DSL package root.
+    Falls back to start_frame if no user frame is found (e.g. all frames are library code).
+    """
+    frame = start_frame
+    while frame is not None:
+        if not _is_framework_frame(frame.f_code.co_filename):
+            return frame
+        frame = frame.f_back
+    # Fallback: if everything is framework code, use the original caller
+    return start_frame
 
 
 def dsl_user_op(opFunc):
@@ -34,30 +72,39 @@ def dsl_user_op(opFunc):
     @wraps(opFunc)
     def wrapper(*args, **kwargs):
         loc = kwargs.pop("loc", None)
-        if loc is None:
-            frame = inspect.currentframe().f_back
+        frameInfo = None
+        verifier_error = False
+
+        if loc is None and ir.Context.current is not None:
+            frame = _find_user_frame(inspect.currentframe().f_back)
             frameInfo = inspect.getframeinfo(frame)
-            # In Python < 3.11, getframeinfo returns a NamedTuple without positions
-            if not hasattr(frameInfo, "positions"):
-                file_loc = ir.Location.file(
-                    frameInfo.filename,
-                    frameInfo.lineno,
-                    0,
+            try:
+                # In Python < 3.11, getframeinfo returns a NamedTuple without positions
+                if not hasattr(frameInfo, "positions"):
+                    file_loc = ir.Location.file(
+                        frameInfo.filename,
+                        frameInfo.lineno,
+                        0,
+                    )
+                else:
+                    file_loc = ir.Location.file(
+                        frameInfo.filename,
+                        frameInfo.positions.lineno,
+                        frameInfo.positions.col_offset or 0,
+                    )
+                loc = ir.Location.name(
+                    (
+                        "".join([c.strip() for c in frameInfo.code_context])
+                        if frameInfo.code_context
+                        else frameInfo.function
+                    ),
+                    childLoc=file_loc,
                 )
-            else:
-                file_loc = ir.Location.file(
-                    frameInfo.filename,
-                    frameInfo.positions.lineno,
-                    frameInfo.positions.col_offset,
-                )
-            loc = ir.Location.name(
-                (
-                    "".join([c.strip() for c in frameInfo.code_context])
-                    if frameInfo.code_context
-                    else frameInfo.function
-                ),
-                childLoc=file_loc,
-            )
+            except RuntimeError:
+                # No MLIR context available (e.g. validation-only call
+                # outside a kernel).  Proceed with loc=None so that the
+                # wrapped function's own validation can still fire.
+                pass
 
         try:
             res_or_list = opFunc(*args, **kwargs, loc=loc)

python/CuTeDSL/cutlass/base_dsl/arch.py

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

python/CuTeDSL/cutlass/base_dsl/ast_helpers.py

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

python/CuTeDSL/cutlass/base_dsl/ast_preprocessor.py

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

python/CuTeDSL/cutlass/base_dsl/cache_helpers.py

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

python/CuTeDSL/cutlass/base_dsl/common.py

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

python/CuTeDSL/cutlass/base_dsl/compiler.py

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

python/CuTeDSL/cutlass/base_dsl/dsl.py

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA
@@ -1341,7 +1341,11 @@ class BaseDSL(metaclass=DSLSingletonMeta):
         location=None,
     ):
         """Generate MLIR module and compile iself.T_provider."""
-        with ir.Context(), self.get_ir_location(location):
+        with ir.Context() as ctx, self.get_ir_location(location):
+            # If threading is enabled, each MLIR context will keep alive a thread pool.
+            # When we cache MLIR compilation results, we also cache its context thus accumulating #(compilations) * thread_pool_size threads.
+            # Disable threading to avoid such excessive number of threads.
+            ctx.enable_multithreading(False)
             try:
                 # Convert input arguments to MLIR arguments
                 exe_args, func_types, adapted_args = self.generate_mlir_function_types(
@@ -1491,6 +1495,11 @@ class BaseDSL(metaclass=DSLSingletonMeta):
 
         # Check if all non-default arguments are provided
         for param in sig.parameters.values():
+            if param.kind in (
+                inspect.Parameter.VAR_POSITIONAL,
+                inspect.Parameter.VAR_KEYWORD,
+            ):
+                continue
             if (
                 param.default is inspect.Parameter.empty
                 and param.name not in bound_args.arguments
@@ -1501,6 +1510,95 @@ class BaseDSL(metaclass=DSLSingletonMeta):
 
         return sig
 
+    def _get_full_arg_spec(self, funcBody):
+        """
+        Returns the full argument specification for a given function, handling PEP-563
+        (postponed evaluation of type annotations) if necessary.
+
+        If the function's annotations are provided as strings (which occurs when PEP-563
+        is enabled), this method evaluates those annotations so they are returned as objects
+        instead of strings.
+
+        Parameters
+        ----------
+        funcBody : function
+            The function whose argument specification is to be retrieved.
+
+        Returns
+        -------
+        inspect.FullArgSpec
+            The complete argument specification of the function, with its annotations
+            properly evaluated and resolved where relevant.
+        """
+        args_spec = inspect.getfullargspec(funcBody)
+        # Set `eval_str = True` to make it work when PEP-563 is enabled
+        if args_spec.annotations and all(
+            type(arg_type) is str for arg_type in args_spec.annotations.values()
+        ):
+            eval_annotations = inspect.get_annotations(funcBody, eval_str=True)
+            args_spec = inspect.FullArgSpec(
+                args_spec.args,
+                args_spec.varargs,
+                args_spec.varkw,
+                args_spec.defaults,
+                args_spec.kwonlyargs,
+                args_spec.kwonlydefaults,
+                eval_annotations,
+            )
+        return args_spec
+
+    @staticmethod
+    def _expand_varargs_varkw(
+        canonicalized_args: tuple,
+        canonicalized_kwargs: dict,
+        args_spec: inspect.FullArgSpec,
+    ) -> inspect.FullArgSpec:
+        """Expand *args and **kwargs into concrete named parameters in the FullArgSpec.
+
+        When a JIT function uses *args or **kwargs, the concrete call-site values
+        are known. This method synthesizes named parameters for them so the rest
+        of the pipeline (which expects fixed-arity signatures) works unchanged.
+
+        For *args: extra positional arguments beyond ``args_spec.args`` get
+        synthetic names ``_vararg_0``, ``_vararg_1``, etc.
+
+        For **kwargs: extra keyword arguments beyond ``args_spec.kwonlyargs``
+        are appended as keyword-only parameters.
+        """
+        if not args_spec.varargs and not args_spec.varkw:
+            return args_spec
+
+        expanded_args = list(args_spec.args)
+        expanded_annotations = dict(args_spec.annotations)
+        expanded_defaults = list(args_spec.defaults) if args_spec.defaults else []
+
+        if args_spec.varargs:
+            n_regular = len(args_spec.args)
+            n_extra = len(canonicalized_args) - n_regular
+            for i in range(n_extra):
+                expanded_args.append(f"varargs_{i}")
+
+        expanded_kwonlyargs = list(args_spec.kwonlyargs)
+        expanded_kwonlydefaults = (
+            dict(args_spec.kwonlydefaults) if args_spec.kwonlydefaults else {}
+        )
+
+        if args_spec.varkw:
+            existing_kwonly = set(args_spec.kwonlyargs)
+            for key in canonicalized_kwargs:
+                if key not in existing_kwonly:
+                    expanded_kwonlyargs.append(key)
+
+        return inspect.FullArgSpec(
+            args=expanded_args,
+            varargs=None,
+            varkw=None,
+            defaults=tuple(expanded_defaults) if expanded_defaults else None,
+            kwonlyargs=expanded_kwonlyargs,
+            kwonlydefaults=expanded_kwonlydefaults if expanded_kwonlydefaults else None,
+            annotations=expanded_annotations,
+        )
+
     def _func(self, funcBody, *args, **kwargs):
         """Decorator for MLIR functions.
         It cuts the boilerplate code, does the following:
@@ -1553,6 +1651,10 @@ class BaseDSL(metaclass=DSLSingletonMeta):
         canonicalized_args, canonicalized_kwargs = self._canonicalize_args(
             sig, *args, **kwargs
         )
+        # Expand *args/**kwargs into concrete named parameters
+        args_spec = self._expand_varargs_varkw(
+            canonicalized_args, canonicalized_kwargs, args_spec
+        )
         # Simple name mangling
         function_name = self.mangle_name(function_name, canonicalized_args, args_spec)
         if func_name_prefix:

python/CuTeDSL/cutlass/base_dsl/env_manager.py

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

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

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

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

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

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

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

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

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

python/CuTeDSL/cutlass/base_dsl/jit_executor.py

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

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

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

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

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

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

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

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

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

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

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

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

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

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

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

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

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

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

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

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

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

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

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

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

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

python/CuTeDSL/cutlass/base_dsl/typing.py

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA
@@ -35,6 +35,7 @@ from .common import *
 from .ast_helpers import const_expr
 from ._mlir_helpers import arith as arith_helper, lru_cache_ir
 from ._mlir_helpers.arith import ArithValue
+from ._mlir_helpers.op import dsl_user_op
 
 from .._mlir import ir
 from .._mlir.extras import types as T
@@ -843,7 +844,6 @@ def _binary_op(op, promote_operand=True, promote_bool=False, flip=False):
         if flip:
             lhs_val, rhs_val = rhs_val, lhs_val
 
-        # Check if the operation is supported by the operands
         res_val = op(lhs_val, rhs_val)
         return res_type(res_val, loc=loc, ip=ip)
 
@@ -1152,72 +1152,91 @@ class Numeric(metaclass=NumericMeta, is_abstract=True):
             )
         return res_type(value)
 
+    @dsl_user_op
     def __add__(self, other, *, loc=None, ip=None) -> "Numeric":
         return _binary_op(operator.add, promote_bool=True)(self, other, loc=loc, ip=ip)
 
+    @dsl_user_op
     def __sub__(self, other, *, loc=None, ip=None) -> "Numeric":
         return _binary_op(operator.sub, promote_bool=True)(self, other, loc=loc, ip=ip)
 
+    @dsl_user_op
     def __mul__(self, other, *, loc=None, ip=None) -> "Numeric":
         return _binary_op(operator.mul, promote_bool=True)(self, other, loc=loc, ip=ip)
 
+    @dsl_user_op
     def __floordiv__(self, other, *, loc=None, ip=None) -> "Numeric":
         return _binary_op(operator.floordiv, promote_bool=True)(
             self, other, loc=loc, ip=ip
         )
 
+    @dsl_user_op
     def __truediv__(self, other, *, loc=None, ip=None) -> "Numeric":
         return _binary_op(operator.truediv, promote_bool=True)(
             self, other, loc=loc, ip=ip
         )
 
+    @dsl_user_op
     def __mod__(self, other, *, loc=None, ip=None) -> "Numeric":
         return _binary_op(operator.mod, promote_bool=True)(self, other, loc=loc, ip=ip)
 
+    @dsl_user_op
     def __radd__(self, other, *, loc=None, ip=None) -> "Numeric":
         return self.__add__(other, loc=loc, ip=ip)
 
+    @dsl_user_op
     def __rsub__(self, other, *, loc=None, ip=None) -> "Numeric":
         return _binary_op(operator.sub, promote_bool=True, flip=True)(
             self, other, loc=loc, ip=ip
         )
 
+    @dsl_user_op
     def __rmul__(self, other, *, loc=None, ip=None) -> "Numeric":
         return self.__mul__(other, loc=loc, ip=ip)
 
+    @dsl_user_op
     def __rfloordiv__(self, other, *, loc=None, ip=None) -> "Numeric":
         return _binary_op(operator.floordiv, promote_bool=True, flip=True)(
             self, other, loc=loc, ip=ip
         )
 
+    @dsl_user_op
     def __rtruediv__(self, other, *, loc=None, ip=None) -> "Numeric":
         return _binary_op(operator.truediv, promote_bool=True, flip=True)(
             self, other, loc=loc, ip=ip
         )
 
+    @dsl_user_op
     def __rmod__(self, other, *, loc=None, ip=None) -> "Numeric":
         return _binary_op(operator.mod, promote_bool=True, flip=True)(
             self, other, loc=loc, ip=ip
         )
 
+    @dsl_user_op
     def __eq__(self, other, *, loc=None, ip=None) -> "Boolean":
         return _binary_op(operator.eq)(self, other, loc=loc, ip=ip)  # type: ignore
 
+    @dsl_user_op
     def __ne__(self, other, *, loc=None, ip=None) -> "Boolean":
         return _binary_op(operator.ne)(self, other, loc=loc, ip=ip)  # type: ignore
 
+    @dsl_user_op
     def __lt__(self, other, *, loc=None, ip=None) -> "Boolean":
         return _binary_op(operator.lt)(self, other, loc=loc, ip=ip)  # type: ignore
 
+    @dsl_user_op
     def __le__(self, other, *, loc=None, ip=None) -> "Boolean":
         return _binary_op(operator.le)(self, other, loc=loc, ip=ip)  # type: ignore
 
+    @dsl_user_op
     def __gt__(self, other, *, loc=None, ip=None) -> "Boolean":
         return _binary_op(operator.gt)(self, other, loc=loc, ip=ip)  # type: ignore
 
+    @dsl_user_op
     def __ge__(self, other, *, loc=None, ip=None) -> "Boolean":
         return _binary_op(operator.ge)(self, other, loc=loc, ip=ip)  # type: ignore
 
+    @dsl_user_op
     def __pow__(self, other, *, loc=None, ip=None) -> "Numeric":
         return _binary_op(operator.pow)(self, other, loc=loc, ip=ip)  # type: ignore
 

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

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

python/CuTeDSL/cutlass/base_dsl/utils/logger.py

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

python/CuTeDSL/cutlass/base_dsl/utils/stacktrace.py

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

python/CuTeDSL/cutlass/base_dsl/utils/timer.py

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

python/CuTeDSL/cutlass/base_dsl/utils/tree_utils.py

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

python/CuTeDSL/cutlass/base_dsl/version_info.py

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

python/CuTeDSL/cutlass/cute/__init__.py

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA
@@ -105,6 +105,7 @@ from .core import (
     E,
     # User defined struct
     struct,
+    union,
     pretty_str,
     make_layout_image_mask,
     repeat,

python/CuTeDSL/cutlass/cute/_tvm_ffi_args_spec_converter.py

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

python/CuTeDSL/cutlass/cute/algorithm.py

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

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

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA
@@ -50,6 +50,7 @@ __all__ = [
     "block_in_cluster_idx",
     "block_in_cluster_dim",
     "block_idx_in_cluster",
+    "dynamic_smem_size",
     "shuffle_sync",
     "shuffle_sync_up",
     "shuffle_sync_down",

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

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

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

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

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

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA
@@ -35,7 +35,10 @@ def mbarrier_init(mbar_ptr: Pointer, cnt: Int, *, loc=None, ip=None) -> None:
     :type cnt:       Int
     """
     nvvm.mbarrier_init_shared(
-        mbar_ptr.llvm_ptr, Int32(cnt).ir_value(loc=loc, ip=ip), loc=loc, ip=ip
+        mbar_ptr.to_llvm_ptr(loc=loc, ip=ip),
+        Int32(cnt).ir_value(loc=loc, ip=ip),
+        loc=loc,
+        ip=ip,
     )
 
 
@@ -65,7 +68,7 @@ def mbarrier_arrive_and_expect_tx(
     """
     BaseDSL._get_dsl().check_arch(lambda arch: arch >= Arch.sm_90)
 
-    mbar_llvm_ptr = mbar_ptr.llvm_ptr
+    mbar_llvm_ptr = mbar_ptr.to_llvm_ptr(loc=loc, ip=ip)
     if peer_cta_rank_in_cluster is not None:
         mbar_cluster_type = llvm.PointerType.get(AddressSpace.dsmem)
         mbar_llvm_ptr = nvvm.mapa(
@@ -108,7 +111,7 @@ def mbarrier_expect_tx(
     """
     BaseDSL._get_dsl().check_arch(lambda arch: arch >= Arch.sm_90)
 
-    mbar_llvm_ptr = mbar_ptr.llvm_ptr
+    mbar_llvm_ptr = mbar_ptr.to_llvm_ptr(loc=loc, ip=ip)
     if peer_cta_rank_in_cluster is not None:
         mbar_cluster_type = llvm.PointerType.get(AddressSpace.dsmem)
         mbar_llvm_ptr = nvvm.mapa(
@@ -150,7 +153,7 @@ def mbarrier_wait(mbar_ptr: Pointer, phase: Int, *, loc=None, ip=None) -> None:
     # This NVVM Op is a spin-loop wrapping the mbarrier.try_wait.parity.shared.b64 PTX
     # The timeout in ns only applies to the latter and this call is truly blocking
     nvvm.mbarrier_try_wait_parity_shared(
-        mbar_ptr.llvm_ptr,
+        mbar_ptr.to_llvm_ptr(loc=loc, ip=ip),
         Int32(phase).ir_value(loc=loc, ip=ip),
         Int32(timeout_ns).ir_value(loc=loc, ip=ip),
         loc=loc,
@@ -174,7 +177,7 @@ def mbarrier_try_wait(mbar_ptr: Pointer, phase: Int, *, loc=None, ip=None) -> Bo
 
     return Boolean(
         nvvm.mbarrier_wait_parity(
-            mbar_ptr.llvm_ptr,
+            mbar_ptr.to_llvm_ptr(loc=loc, ip=ip),
             Int32(phase).ir_value(loc=loc, ip=ip),
             nvvm.MBarrierWaitKind.TRY,
             loc=loc,
@@ -228,7 +231,7 @@ def mbarrier_arrive(
                                      the mbarrier is converted to a remote address in the peer CTA's
                                      SMEM.
     """
-    mbar_llvm_ptr = mbar_ptr.llvm_ptr
+    mbar_llvm_ptr = mbar_ptr.to_llvm_ptr(loc=loc, ip=ip)
     if peer_cta_rank_in_cluster is not None:
         BaseDSL._get_dsl().check_arch(lambda arch: arch >= Arch.sm_90)
 
@@ -269,10 +272,5 @@ def cp_async_mbarrier_arrive_noinc(mbar_ptr: Pointer, *, loc=None, ip=None) -> N
     """
     BaseDSL._get_dsl().check_arch(lambda arch: arch >= Arch.sm_90)
 
-    mbar_llvm_ptr = mbar_ptr.llvm_ptr
-    nvvm.cp_async_mbarrier_arrive_shared(
-        mbar_llvm_ptr,
-        noinc=True,
-        loc=loc,
-        ip=ip,
-    )
+    mbar_llvm_ptr = mbar_ptr.to_llvm_ptr(loc=loc, ip=ip)
+    nvvm.cp_async_mbarrier_arrive_shared(mbar_llvm_ptr, noinc=True, loc=loc, ip=ip)

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

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

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

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA
@@ -250,6 +250,26 @@ def block_idx_in_cluster(*, loc=None, ip=None) -> Int32:
     return Int32(nvvm.read_ptx_sreg_cluster_ctarank(T.i32(), loc=loc, ip=ip))
 
 
+@dsl_user_op
+def dynamic_smem_size(*, loc=None, ip=None) -> Int32:
+    """
+    Returns the launch dynamic smem size.
+    """
+    return Int32(
+        llvm.inline_asm(
+            Int32.mlir_type,
+            [],
+            "mov.u32 $0, %dynamic_smem_size;\n",
+            "=r",
+            has_side_effects=True,
+            is_align_stack=False,
+            asm_dialect=llvm.AsmDialect.AD_ATT,
+            loc=loc,
+            ip=ip,
+        )
+    )
+
+
 @dsl_user_op
 def shuffle_sync_op(
     value: Union[Numeric, "TensorSSA"],

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

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

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

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

python/CuTeDSL/cutlass/cute/atom.py

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

python/CuTeDSL/cutlass/cute/core.py

@@ -3,15 +3,17 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA
 # is strictly prohibited.
 
 from functools import partial, reduce
+import inspect
 from inspect import isclass
 from typing import Any, Dict, List, Optional, Tuple, Type, Union, overload
+from types import MethodType
 
 from cutlass import const_expr
 from typing_extensions import deprecated
@@ -31,6 +33,8 @@ from cutlass._mlir.dialects.cute import (
 from cutlass.cutlass_dsl import (
     T,
     const,
+    and_,
+    as_numeric,
     cutlass_arith,
     dsl_user_op,
     extract_mlir_values,
@@ -704,7 +708,7 @@ class ScaledBasis:
 
     def __eq__(self, other):
         if isinstance(other, ScaledBasis):
-            return self.value == other.value and self.mode == other.mode
+            return and_(self.mode == other.mode, self.value == other.value)
         else:
             return False
 
@@ -1212,6 +1216,10 @@ class _ComposedLayout(ComposedLayout):
     @property
     @dsl_user_op
     def shape(self, *, loc=None, ip=None) -> Shape:
+        return self.shape_method(loc=loc, ip=ip)
+
+    @dsl_user_op
+    def shape_method(self, *, loc=None, ip=None) -> Shape:
         return _unpack_x_tuple(
             _cute_ir.get_shape(self.value, loc=loc, ip=ip), loc=loc, ip=ip
         )
@@ -1352,6 +1360,60 @@ class _Pointer(Pointer):
     def type(self) -> ir.Type:
         return self.value.type
 
+    @dsl_user_op
+    def load(self, *, loc=None, ip=None) -> Numeric:
+        # LLVM doesn't support load/store narrow precision per element
+        tmp_ty = self.dtype.mlir_type
+        if self.dtype is Boolean or self.dtype.width == 8:
+            tmp_ty = T.i8()
+        elif self.dtype.width < 8:
+            raise ValueError(
+                f"Loading narrow precision type {self.dtype} is not supported"
+            )
+
+        llvm_ptr = self.to_llvm_ptr(loc=loc, ip=ip)
+        tmp_val = llvm.load(tmp_ty, llvm_ptr, loc=loc, ip=ip)
+        if self.dtype.width == 8:
+            tmp_val = arith.bitcast(self.dtype.mlir_type, tmp_val, loc=loc, ip=ip)
+
+        return self.dtype(tmp_val, loc=loc, ip=ip)
+
+    @dsl_user_op
+    def store(
+        self,
+        value: Union[Numeric, cutlass_arith.ArithValue, int, float, bool],
+        *,
+        loc=None,
+        ip=None,
+    ):
+        if isinstance(value, (int, float, bool, cutlass_arith.ArithValue)):
+            value = self.dtype(value, loc=loc, ip=ip)
+        elif isinstance(value, Numeric):
+            if value.dtype is not self.dtype:
+                value = value.to(self.dtype, loc=loc, ip=ip)
+        else:
+            raise ValueError(f"Unsupported value type: {type(value)}")
+        # LLVM doesn't support load/store narrow precision per element
+        tmp_val = value.ir_value(loc=loc, ip=ip)
+        if self.dtype.width == 8:
+            tmp_val = arith.bitcast(T.i8(), tmp_val, loc=loc, ip=ip)
+        elif self.dtype is not Boolean and self.dtype.width < 8:
+            raise ValueError(
+                f"Storing narrow precision type {self.dtype} is not supported"
+            )
+
+        llvm_ptr = self.to_llvm_ptr(loc=loc, ip=ip)
+        return llvm.store(tmp_val, llvm_ptr, loc=loc, ip=ip)
+
+    @dsl_user_op
+    def __getitem__(self, idx: Int, *, loc=None, ip=None) -> Pointer:
+        return (self + idx).load()
+
+    @dsl_user_op
+    def __setitem__(self, idx: Int, value: Numeric, *, loc=None, ip=None) -> Pointer:
+        (self + idx).store(value, loc=loc, ip=ip)
+        return value
+
     # Only use if you absolutely need to get the LLVM pointer Value
     @property
     @dsl_user_op
@@ -1360,6 +1422,25 @@ class _Pointer(Pointer):
         """
         Get the LLVM pointer representation of this pointer.
 
+        :param loc: Source location for MLIR, defaults to None
+        :type loc: Optional[Location]
+        :param ip: Insertion point for MLIR, defaults to None
+        :type ip: Optional[InsertionPoint]
+        :return: The LLVM pointer representation
+        :rtype: ir.Value
+        """
+        return self.to_llvm_ptr(loc=loc, ip=ip)
+
+    @dsl_user_op
+    @lru_cache_ir()
+    def to_llvm_ptr(self, *, loc=None, ip=None) -> ir.Value:
+        """
+        Get the LLVM pointer representation of this pointer. (Used by internal API to propagate loc and ip)
+
+        :param loc: Source location for MLIR, defaults to None
+        :type loc: Optional[Location]
+        :param ip: Insertion point for MLIR, defaults to None
+        :type ip: Optional[InsertionPoint]
         :return: The LLVM pointer representation
         :rtype: ir.Value
         """
@@ -1587,7 +1668,7 @@ def pretty_str(arg) -> str:
 
 
 @dsl_user_op
-def printf(*args, loc=None, ip=None) -> None:
+def printf(*args, loc=None, ip=None, end="\n") -> None:
     """
     Print one or more values with optional formatting.
 
@@ -1607,6 +1688,8 @@ def printf(*args, loc=None, ip=None) -> None:
     :type loc: Optional[Location]
     :param ip: Insertion point for code generation, defaults to None
     :type ip: Optional[InsertionPoint]
+    :param end: Suffix for the printed value, defaults to newline
+    :type end: Optional[str]
     :raises ValueError: If no arguments are provided
     :raises TypeError: If an unsupported argument type is passed
 
@@ -1636,10 +1719,10 @@ def printf(*args, loc=None, ip=None) -> None:
         raise ValueError("expects at least one argument to print")
 
     if isinstance(args[0], str):
-        fmt = args[0] + "\n"
+        fmt = args[0] + end
         args = args[1:]
     else:
-        fmt = "{}" + ", {}" * (len(args) - 1) + "\n"
+        fmt = "{}" + ", {}" * (len(args) - 1) + end
 
     def process_arg(arg):
         arg0 = arg.value if isinstance(arg, Numeric) else arg
@@ -3384,7 +3467,7 @@ def recast_ptr(
         if cvt_type is None:
             if not isclass(dtype) or not issubclass(dtype, Numeric):
                 raise TypeError(f"dtype must be a type of Numeric, but got {dtype}")
-            cvt_type = dtype.mlir_type
+            cvt_type = T.i8() if dtype is Boolean else dtype.mlir_type
 
     dtype = cvt_type
     value_type = ptr.type.value_type if dtype is None else dtype
@@ -4287,8 +4370,8 @@ class struct:
         storage = allocator.allocate(StorageB)
 
         storage.a[0] ...
-        storage.x ...
-        storage.compA.real ...
+        storage.x.ptr ...
+        storage.compA.real.ptr ...
 
     :param cls: The struct class with annotations.
     :return: The decorated struct class.
@@ -4306,8 +4389,8 @@ class struct:
         :ivar _size: The size of the MemRange.
         """
 
-        _dtype = None
-        _size = None
+        _dtype: Optional[Numeric] = None
+        _size: Optional[int] = None
 
         def __new__(cls, name, bases, dct):
             new_cls = super().__new__(cls, name, bases, dct)
@@ -4337,7 +4420,7 @@ class struct:
 
         @property
         def elem_width(cls):
-            return cls._dtype.width
+            return cls._dtype.width if cls._dtype is not Boolean else 8
 
         @property
         def size_in_bytes(cls):
@@ -4368,12 +4451,15 @@ class struct:
                          case the range can only be used for its address (e.g. as a partition marker).
             :param base: The base address of the memory range.
             """
-            self._dtype = dtype
-            self._size = size
-            self._base = base
+            self._dtype: Optional[Numeric] = dtype
+            self._size: Optional[int] = size
+            self._base: Optional[Pointer] = base
+
+        def __repr__(self):
+            return f"{object.__repr__(self)} <struct.MemRange[{self._dtype}, {self._size}]> <data_ptr = {self.data_ptr()}>"
 
         @dsl_user_op
-        def data_ptr(self, *, loc=None, ip=None):
+        def data_ptr(self, *, loc=None, ip=None) -> Pointer:
             """
             Returns start pointer to the data in this memory range.
 
@@ -4384,7 +4470,9 @@ class struct:
             return recast_ptr(self._base, dtype=self._dtype, loc=loc, ip=ip)
 
         @dsl_user_op
-        def get_tensor(self, layout, swizzle=None, dtype=None, *, loc=None, ip=None):
+        def get_tensor(
+            self, layout, swizzle=None, dtype=None, *, loc=None, ip=None
+        ) -> Tensor:
             """
             Creates a tensor from the memory range.
 
@@ -4404,9 +4492,10 @@ class struct:
             elem_type = self._dtype if dtype is None else dtype
             ptr = recast_ptr(self._base, swizzle, dtype=elem_type, loc=loc, ip=ip)
             res = make_tensor(ptr, layout, loc=loc, ip=ip)
-            return res
+            return type(res)(res, dtype=elem_type, loc=loc, ip=ip)
 
-        def __getitem__(self, index: int) -> Any:
+        @dsl_user_op
+        def __getitem__(self, index: int, *, loc=None, ip=None) -> Any:
             """
             Returns the element at the specified index in the memory range.
 
@@ -4415,7 +4504,21 @@ class struct:
             :raises AssertionError: If the index is out of range.
             """
             assert (index >= 0) and (index < self._size)
-            return self.data_ptr() + index
+            ptr = self.data_ptr() + index
+            return ptr.load(loc=loc, ip=ip)
+
+        @dsl_user_op
+        def __setitem__(self, index: int, val, *, loc=None, ip=None):
+            """
+            Set element value at the specified index in the memory range.
+
+            :param index: The index of the element to retrieve.
+            :val: The element value at the specified index.
+            :raises AssertionError: If the index is out of range.
+            """
+            assert (index >= 0) and (index < self._size)
+            ptr = self.data_ptr() + index
+            ptr.store(as_numeric(val).to(self._dtype), loc=loc, ip=ip)
 
     # inner class for aligning a member type
     class _AlignMeta(type):
@@ -4430,8 +4533,8 @@ class struct:
         :ivar _align: The alignment of the data type.
         """
 
-        _dtype = None
-        _align = None
+        _dtype: Optional[Any] = None
+        _align: Optional[int] = None
 
         def __new__(cls, name, bases, dct):
             return super().__new__(cls, name, bases, dct)
@@ -4473,6 +4576,88 @@ class struct:
 
         pass
 
+    class _ScalarData(_Pointer):
+        """
+        Represents a scalar value at a given pointer location in memory.
+
+        This class provides utility methods to get a scalar pointer.
+        It wraps a pointer to a scalar element and enables element-wise memory operations.
+
+        :ivar _ptr: The underlying pointer to the scalar value.
+        """
+
+        def __init__(self, ptr):
+            self._ptr: Optional[_Pointer] = ptr
+
+        def __repr__(self):
+            return f"{object.__repr__(self)} <{self.dtype}> <ptr = {self._ptr}>"
+
+        def __get_mlir_types__(self) -> List[ir.Type]:
+            return [self.value.type]
+
+        def __extract_mlir_values__(self) -> List[ir.Value]:
+            return [self.value]
+
+        def __new_from_mlir_values__(self, values) -> Pointer:
+            ptr = _Pointer(
+                values[0] if isinstance(values[0], ir.Value) else values[0].value
+            )
+            return self.__class__(ptr)
+
+        @dsl_user_op
+        def to_llvm_ptr(self, *, loc=None, ip=None) -> ir.Value:
+            llvm_ptr_ty = llvm.PointerType.get(
+                self._ptr.memspace.value
+                if self._ptr.memspace != AddressSpace.rmem
+                else 0
+            )
+            return builtin.unrealized_conversion_cast(
+                [llvm_ptr_ty], [self.value], loc=loc, ip=ip
+            )
+
+        @property
+        def ptr(self) -> Pointer:
+            """
+            Get the underlying pointer.
+
+            :return: The pointer to the scalar value.
+            :rtype: Pointer
+            """
+            return self._ptr
+
+        @property
+        def dtype(self) -> Numeric:
+            """
+            Get the data type of the scalar value.
+
+            :return: The numeric data type of the underlying pointer.
+            :rtype: Numeric
+            """
+            return self._ptr.dtype
+
+        @property
+        @deprecated("Using `struct.scalar` as pointer is deprecated.")
+        def value(self):
+            """
+            Get the raw MLIR value of the underlying pointer.
+
+            .. deprecated::
+                Using ``struct.scalar`` as pointer is deprecated.
+                Use explicit ``struct.scalar.ptr`` for pointer instead.
+
+            :return: The MLIR value of the underlying pointer.
+            :rtype: ir.Value
+            """
+            import warnings
+
+            with warnings.catch_warnings():
+                warnings.simplefilter("always")
+                warnings.warn(
+                    "Use explicit `struct.scalar.ptr` for pointer instead.",
+                    DeprecationWarning,
+                )
+            return self._ptr.value
+
     # util func for base dsl scalar types
     @staticmethod
     def _is_scalar_type(dtype):
@@ -4493,55 +4678,72 @@ class struct:
         :raises TypeError: If the struct is empty.
         """
         self._cls = cls
-        self.__name__ = f"struct::{cls.__name__}"
+        self.__name__ = f"cute.struct::{cls.__name__}"
         # Get the class annotations
         self._annotations = getattr(cls, "__annotations__", {})
         # Create a dictionary to store the offsets
         self._offsets: Dict[str, int] = {}
 
+        # Override `setattr` function for struct to assign scalar properly
+        def struct_setattr(self, name, value):
+            attr = getattr(self, name, None)
+            if isinstance(attr, struct._ScalarData):
+                value = as_numeric(value).to(attr.dtype)
+                attr.ptr.store(value)
+            else:
+                raise ValueError(f"cannot assign value to `{name}` in {self.__name__}")
+
+        type.__setattr__(self._cls, "__setattr__", struct_setattr)
+
+        # Override `__repr__` function for struct info
+        def struct_repr(self):
+            return f"{object.__repr__(self)} <{self.__name__}> <base = {self.base}>"
+
+        self._cls.__repr__ = struct_repr
+
         # Calculate the offsets and alignment
         offset = 0
         alignment = 1
         if len(self._annotations) == 0:
             raise TypeError("Empty struct is not supported!")
-        for name, object in self._annotations.items():
-            # get alignment of object
+        for name, member in self._annotations.items():
+            # get alignment of member
             sub_align = 1
-            if isinstance(object, struct._AlignMeta):
-                sub_align = object.align
-                object = object.dtype
+            if isinstance(member, struct._AlignMeta):
+                sub_align = member.align
+                member = member.dtype
 
             # switch addition order to support dynamic size
             def add_offset(val):
                 return val + offset if isinstance(val, ir.Value) else offset + val
 
             # size of scalar
-            if struct._is_scalar_type(object):
-                dtype_size = max(1, object.width // 8)
+            if struct._is_scalar_type(member):
+                dtype_size = max(1, member.width // 8)
                 sub_align = max(dtype_size, sub_align)
                 offset = self.align_offset(offset, sub_align)
                 self._offsets[name] = offset
                 offset = add_offset(dtype_size)
             # size of array is size_in_bytes, alignment is elem_size
-            elif isinstance(object, struct._MemRangeMeta):
+            elif isinstance(member, struct._MemRangeMeta):
                 # Allow empty array as a free marker-only struct member.
                 # Use max(sub_align, ) because we might have in the future some
-                # object.elem_width less than 8, such as fp4, bit and others,
+                # member.elem_width less than 8, such as fp4, bit and others,
                 # and align_offset() does not support an alignment of 0.
-                sub_align = max(object.elem_width // 8, sub_align)
+                sub_align = max(member.elem_width // 8, sub_align)
                 offset = self.align_offset(offset, sub_align)
                 self._offsets[name] = offset
-                offset = add_offset(object.size_in_bytes)
+                offset = add_offset(member.size_in_bytes)
             # size of struct
-            elif isinstance(object, struct):
-                sub_align = max(object.__alignof__(), sub_align)
+            elif isinstance(member, struct):
+                sub_align = max(member.__alignof__(), sub_align)
                 offset = self.align_offset(offset, sub_align)
                 self._offsets[name] = offset
-                offset = add_offset(object.__sizeof__())
+                offset = add_offset(member.__sizeof__())
             else:
                 raise TypeError(
                     f"Struct element only support struct/array/base_dsl scalar, "
-                    f"but got {object}"
+                    f"but got {member}"
                 )
             # Total alignment determined by the strictest requirement
             alignment = max(alignment, sub_align)
@@ -4564,20 +4766,22 @@ class struct:
         # make an new object of user-defined decorated struct
         # otherwise it will override same self._cls when new instance created
         cls = self._cls()
-        setattr(cls, "_base", base)
+        object.__setattr__(cls, "base", base)
+        object.__setattr__(cls, "__name__", self.__name__)
         for name, off in self._offsets.items():
             obj = self._annotations[name]
             if isinstance(obj, struct._AlignMeta):
                 obj = obj.dtype
             if struct._is_scalar_type(obj):
-                new_obj = recast_ptr(base + off, dtype=obj, loc=loc, ip=ip)
-                setattr(cls, name, new_obj)
+                ptr = recast_ptr(base + off, dtype=obj, loc=loc, ip=ip)
+                new_obj = struct._ScalarData(ptr)
+                object.__setattr__(cls, name, new_obj)
             elif isinstance(obj, struct._MemRangeMeta):
                 new_obj = struct._MemRangeData(obj._dtype, obj._size, base + off)
-                setattr(cls, name, new_obj)
+                object.__setattr__(cls, name, new_obj)
             elif isinstance(obj, struct):
                 new_obj = obj(base + off)
-                setattr(cls, name, new_obj)
+                object.__setattr__(cls, name, new_obj)
             else:
                 raise TypeError(
                     f"Struct element only support struct/array/base_dsl scalar, "
@@ -4614,6 +4818,196 @@ class struct:
         return (offset + (align - 1)) & ~(align - 1)
 
 
+##############################################################################
+# User defined struct
+##############################################################################
+
+
+class union(struct):
+    """
+    Decorator to abstract C union in Python DSL.
+
+    Similar to cute.struct, but lays out objects as a union:
+    - All objects start at offset 0
+    - The alignment is the maximum alignment of all objects
+    - The size is the maximum size of all objects
+
+    **Usage:**Expand commentComment on line R4131
+
+    .. code-block:: python
+
+        # Define a union with scalar int/float elements:
+        @cute.union
+        class value_union:
+            as_int : cutlass.Int32
+            as_float : cutlass.Float32
+
+
+        @cute.union
+        class data_union:
+            small : cutlass.Int16
+            medium : cutlass.Int32
+            large : cutlass.Int64
+
+
+        # Supports alignment for its elements:
+        @cute.union
+        class aligned_union:
+            a: cute.struct.Align[cutlass.Float32, 16]
+            b: cute.struct.Align[cutlass.Int32, 8]
+
+
+        # Statically get size and alignment:
+        size = data_union.__sizeof__()
+        align = data_union.__alignof__()
+
+        # Allocate and reference elements:
+        allocator = cutlass.utils.SmemAllocator()
+        value = allocator.allocate(data_union)
+
+        # Access union members (all at the same offset):
+        value.small.ptr ...
+        value.medium.ptr ...
+        value.large.ptr ...
+
+    :param cls: The union class with annotations.
+    :return: The decorated union class.
+    """
+
+    def __init__(self, cls):
+        """
+        Initializes a new cute.union decorator instance.
+
+        :param cls: The class representing the union data type.
+        :raises TypeError: If the union is empty.
+        """
+        object.__setattr__(self, "_cls", cls)
+        object.__setattr__(self, "__name__", f"cute.union::{cls.__name__}")
+        # Get the class annotations
+        object.__setattr__(self, "_annotations", getattr(cls, "__annotations__", {}))
+        # Create a dictionary to store the offsets (all zeros for union)
+        object.__setattr__(self, "_offsets", {})
+
+        # Override `setattr` function for struct to assign scalar properly
+        def union_setattr(self, name, value):
+            attr = getattr(self, name, None)
+            if isinstance(attr, struct._ScalarData):
+                value = as_numeric(value).to(attr.dtype)
+                attr.ptr.store(value)
+            else:
+                raise ValueError(f"cannot assign value to `{name}` in {self.__name__}")
+
+        type.__setattr__(self._cls, "__setattr__", union_setattr)
+
+        # Override `__repr__` function for struct info
+        def union_repr(self):
+            return f"{object.__repr__(self)} <{self.__name__}> <base = {self.base}>"
+
+        type.__setattr__(self._cls, "__repr__", union_repr)
+
+        # Calculate the maximum size and alignment
+        max_size = 0
+        max_alignment = 1
+        if len(self._annotations) == 0:
+            raise TypeError("Empty union is not supported!")
+        for name, item in self._annotations.items():
+            # All offsets are 0 for a union
+            self._offsets[name] = 0
+
+            # Get alignment of object
+            sub_align = 1
+            if isinstance(item, struct._AlignMeta):
+                sub_align = item.align
+                item = item.dtype
+
+            # Calculate size and alignment based on object type
+            if struct._is_scalar_type(item):
+                dtype_size = max(1, item.width // 8)
+                sub_align = max(dtype_size, sub_align)
+                max_size = max(max_size, dtype_size)
+            elif isinstance(item, struct._MemRangeMeta):
+                sub_align = max(item.elem_width // 8, sub_align)
+                max_size = max(max_size, item.size_in_bytes)
+            elif isinstance(item, struct):
+                sub_align = max(item.__alignof__(), sub_align)
+                max_size = max(max_size, item.__sizeof__())
+            else:
+                raise TypeError(
+                    f"Union element only support struct/array/DSL scalar, "
+                    f"but got `{item.__qualname__}`"
+                )
+            # Union alignment is the maximum alignment of all members
+            max_alignment = max(max_alignment, sub_align)
+
+        # Union size is the maximum size, aligned to the maximum alignment
+        object.__setattr__(self, "_align_of", max_alignment)
+        object.__setattr__(
+            self, "_size_of", struct.align_offset(max_size, max_alignment)
+        )
+
+    @dsl_user_op
+    def __call__(self, base: Any, *, loc=None, ip=None) -> None:
+        """
+        Creates a new instance of the decorated union.
+
+        :param base: The base address of the union.
+        :return: An instance of the decorated union.
+        :raises TypeError: If the base pointer is not byte-sized.
+        """
+        if base.type.value_type.width != 8:
+            raise TypeError("union base ptr value type must be byte sized.")
+        # Make a new object of user-defined decorated union
+        cls = self._cls()
+        object.__setattr__(cls, "base", base)
+        object.__setattr__(cls, "__name__", self.__name__)
+        for name, off in self._offsets.items():
+            obj = self._annotations[name]
+            if isinstance(obj, struct._AlignMeta):
+                obj = obj.dtype
+            if struct._is_scalar_type(obj):
+                ptr = recast_ptr(base + off, dtype=obj, loc=loc, ip=ip)
+                new_obj = struct._ScalarData(ptr)
+                object.__setattr__(cls, name, new_obj)
+            elif isinstance(obj, struct._MemRangeMeta):
+                new_obj = struct._MemRangeData(obj._dtype, obj._size, base + off)
+                object.__setattr__(cls, name, new_obj)
+            elif isinstance(obj, struct):
+                new_obj = obj(base + off)
+                object.__setattr__(cls, name, new_obj)
+            else:
+                raise TypeError(
+                    f"Union element only support struct/array/DSL scalar, "
+                    f"but got `{obj.__qualname__}`"
+                )
+        return cls
+
+    def __setattr__(self, name, value):
+        raise TypeError("Cannot add a new field after initialization")
+    def size_in_bytes(self) -> int:
+        """
+        Returns the size of the union in bytes.
+
+        :return: The size of the union.
+        """
+        return self._size_of
+
+    def __sizeof__(self) -> int:
+        """
+        Returns the size of the union in bytes.
+
+        :return: The size of the union.
+        """
+        return self._size_of
+
+    def __alignof__(self) -> int:
+        """
+        Returns the alignment of the union in bytes.
+
+        :return: The alignment of the union.
+        """
+        return self._align_of
+
+
 # Deprecated usage but keep them to avoid breaking some examples uses `cute.core.ThrMma`
 
 from .atom import ThrCopy as _ThrCopy
@@ -4752,6 +5146,23 @@ class FastDivmodDivisor:
         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

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

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

python/CuTeDSL/cutlass/cute/experimental/algorithm.py

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

python/CuTeDSL/cutlass/cute/experimental/math.py

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

python/CuTeDSL/cutlass/cute/experimental/memory.py

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

python/CuTeDSL/cutlass/cute/experimental/utils.py

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

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

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

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

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

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

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA
@@ -195,7 +195,7 @@ typedef struct {{
                 packed_args.append("&" + arg_name)
             else:
                 raise DSLRuntimeError(
-                    f"Unsupported argument for c function argument generation: {arg} with type {arg_type}"
+                    f"Unsupported argument for c function argument generation: {arg_name} = {arg} with type annotation {arg_type}"
                 )
 
         return arguments, packed_args, declarations

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

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

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

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

python/CuTeDSL/cutlass/cute/ffi.py

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

python/CuTeDSL/cutlass/cute/math.py

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA
@@ -15,6 +15,7 @@ from .typing import Numeric
 from .tensor import TensorSSA
 
 from cutlass._mlir.dialects import math, arith
+from cutlass.cutlass_dsl import dsl_user_op
 
 
 def _math_op(func: Callable, fastmath: bool, *args, **kwargs):
@@ -22,7 +23,7 @@ def _math_op(func: Callable, fastmath: bool, *args, **kwargs):
 
     :param func: The function to dispatch
     :param args: The input tensor or scalar
-    :param kwargs: The input tensor or scalar
+    :param kwargs: Extra keyword arguments (loc, ip) forwarded to the MLIR op
     """
     arg_type = type(args[0])
     for arg in args:
@@ -40,15 +41,16 @@ def _math_op(func: Callable, fastmath: bool, *args, **kwargs):
     fastmath_flag = arith.FastMathFlags.fast if fastmath else arith.FastMathFlags.none
     if isinstance(args[0], TensorSSA):
         return TensorSSA(
-            func(*args, fastmath=fastmath_flag), args[0].shape, args[0].dtype
+            func(*args, fastmath=fastmath_flag, **kwargs), args[0].shape, args[0].dtype
         )
     else:
         args = [a.ir_value() for a in args]
-        return func(*args, fastmath=fastmath_flag)
+        return func(*args, fastmath=fastmath_flag, **kwargs)
 
 
+@dsl_user_op
 def acos(
-    a: Union[TensorSSA, Numeric], fastmath: bool = False
+    a: Union[TensorSSA, Numeric], fastmath: bool = False, *, loc=None, ip=None
 ) -> Union[TensorSSA, Numeric]:
     """Compute element-wise arc cosine of the input tensor.
 
@@ -56,6 +58,10 @@ def acos(
     :type a: Union[TensorSSA, Numeric]
     :param fastmath: Enable fast math optimizations, defaults to False
     :type fastmath: bool, optional
+    :param loc: Source location information, defaults to None
+    :type loc: Optional[Location]
+    :param ip: Insertion point for IR generation, defaults to None
+    :type ip: Optional[InsertionPoint]
     :return: Tensor containing the arc cosine of each element in input tensor
     :rtype: Union[TensorSSA, Numeric]
 
@@ -67,11 +73,12 @@ def acos(
         y = x.load()  # Load values
         z = acos(y)  # Compute arc cosine
     """
-    return _math_op(math.acos, fastmath, a)
+    return _math_op(math.acos, fastmath, a, loc=loc, ip=ip)
 
 
+@dsl_user_op
 def asin(
-    a: Union[TensorSSA, Numeric], fastmath: bool = False
+    a: Union[TensorSSA, Numeric], fastmath: bool = False, *, loc=None, ip=None
 ) -> Union[TensorSSA, Numeric]:
     """Compute element-wise arc sine of the input tensor.
 
@@ -79,6 +86,10 @@ def asin(
     :type a: Union[TensorSSA, Numeric]
     :param fastmath: Enable fast math optimizations, defaults to False
     :type fastmath: bool, optional
+    :param loc: Source location information, defaults to None
+    :type loc: Optional[Location]
+    :param ip: Insertion point for IR generation, defaults to None
+    :type ip: Optional[InsertionPoint]
     :return: Tensor containing the arc sine of each element in input tensor
     :rtype: Union[TensorSSA, Numeric]
 
@@ -90,11 +101,12 @@ def asin(
         y = x.load()  # Load values
         z = asin(y)  # Compute arc sine
     """
-    return _math_op(math.asin, fastmath, a)
+    return _math_op(math.asin, fastmath, a, loc=loc, ip=ip)
 
 
+@dsl_user_op
 def atan(
-    a: Union[TensorSSA, Numeric], fastmath: bool = False
+    a: Union[TensorSSA, Numeric], fastmath: bool = False, *, loc=None, ip=None
 ) -> Union[TensorSSA, Numeric]:
     """Compute element-wise arc tangent of the input tensor.
 
@@ -102,6 +114,10 @@ def atan(
     :type a: Union[TensorSSA, Numeric]
     :param fastmath: Enable fast math optimizations, defaults to False
     :type fastmath: bool, optional
+    :param loc: Source location information, defaults to None
+    :type loc: Optional[Location]
+    :param ip: Insertion point for IR generation, defaults to None
+    :type ip: Optional[InsertionPoint]
     :return: Tensor containing the arc tangent of each element in input tensor
     :rtype: Union[TensorSSA, Numeric]
 
@@ -113,11 +129,13 @@ def atan(
         y = x.load()  # Load values
         z = atan(y)  # Compute arc tangent
     """
-    return _math_op(math.atan, fastmath, a)
+    return _math_op(math.atan, fastmath, a, loc=loc, ip=ip)
 
 
+@dsl_user_op
 def atan2(
-    a: Union[TensorSSA, Numeric], b: Union[TensorSSA, Numeric], fastmath: bool = False
+    a: Union[TensorSSA, Numeric], b: Union[TensorSSA, Numeric], fastmath: bool = False,
+    *, loc=None, ip=None
 ) -> Union[TensorSSA, Numeric]:
     """Compute element-wise arc tangent of two tensors.
 
@@ -130,6 +148,10 @@ def atan2(
     :type b: Union[TensorSSA, Numeric]
     :param fastmath: Enable fast math optimizations, defaults to False
     :type fastmath: bool, optional
+    :param loc: Source location information, defaults to None
+    :type loc: Optional[Location]
+    :param ip: Insertion point for IR generation, defaults to None
+    :type ip: Optional[InsertionPoint]
     :return: Tensor containing the arc tangent of a/b element-wise
     :rtype: Union[TensorSSA, Numeric]
 
@@ -141,11 +163,12 @@ def atan2(
         x = cute.make_rmem_tensor(ptr2, layout).load()  # x coordinates
         theta = atan2(y, x)  # Compute angles
     """
-    return _math_op(math.atan2, fastmath, a, b)
+    return _math_op(math.atan2, fastmath, a, b, loc=loc, ip=ip)
 
 
+@dsl_user_op
 def cos(
-    a: Union[TensorSSA, Numeric], fastmath: bool = False
+    a: Union[TensorSSA, Numeric], fastmath: bool = False, *, loc=None, ip=None
 ) -> Union[TensorSSA, Numeric]:
     """Compute element-wise cosine of the input tensor.
 
@@ -153,6 +176,10 @@ def cos(
     :type a: Union[TensorSSA, Numeric]
     :param fastmath: Enable fast math optimizations, defaults to False
     :type fastmath: bool, optional
+    :param loc: Source location information, defaults to None
+    :type loc: Optional[Location]
+    :param ip: Insertion point for IR generation, defaults to None
+    :type ip: Optional[InsertionPoint]
     :return: Tensor containing the cosine of each element
     :rtype: Union[TensorSSA, Numeric]
 
@@ -164,11 +191,12 @@ def cos(
         y = x.load()  # Load values
         z = cos(y)  # Compute cosine
     """
-    return _math_op(math.cos, fastmath, a)
+    return _math_op(math.cos, fastmath, a, loc=loc, ip=ip)
 
 
+@dsl_user_op
 def erf(
-    a: Union[TensorSSA, Numeric], fastmath: bool = False
+    a: Union[TensorSSA, Numeric], fastmath: bool = False, *, loc=None, ip=None
 ) -> Union[TensorSSA, Numeric]:
     """Compute element-wise error function of the input tensor.
 
@@ -179,6 +207,10 @@ def erf(
     :type a: Union[TensorSSA, Numeric]
     :param fastmath: Enable fast math optimizations, defaults to False
     :type fastmath: bool, optional
+    :param loc: Source location information, defaults to None
+    :type loc: Optional[Location]
+    :param ip: Insertion point for IR generation, defaults to None
+    :type ip: Optional[InsertionPoint]
     :return: Tensor containing the error function value for each element
     :rtype: Union[TensorSSA, Numeric]
 
@@ -190,11 +222,12 @@ def erf(
         y = x.load()  # Load values
         z = erf(y)  # Compute error function
     """
-    return _math_op(math.erf, fastmath, a)
+    return _math_op(math.erf, fastmath, a, loc=loc, ip=ip)
 
 
+@dsl_user_op
 def exp(
-    a: Union[TensorSSA, Numeric], fastmath: bool = False
+    a: Union[TensorSSA, Numeric], fastmath: bool = False, *, loc=None, ip=None
 ) -> Union[TensorSSA, Numeric]:
     """Compute element-wise exponential of the input tensor.
 
@@ -202,6 +235,10 @@ def exp(
     :type a: Union[TensorSSA, Numeric]
     :param fastmath: Enable fast math optimizations, defaults to False
     :type fastmath: bool, optional
+    :param loc: Source location information, defaults to None
+    :type loc: Optional[Location]
+    :param ip: Insertion point for IR generation, defaults to None
+    :type ip: Optional[InsertionPoint]
     :return: Tensor containing the exponential of each element
     :rtype: Union[TensorSSA, Numeric]
 
@@ -213,11 +250,12 @@ def exp(
         y = x.load()  # Load values
         z = exp(y)  # Compute exponential
     """
-    return _math_op(math.exp, fastmath, a)
+    return _math_op(math.exp, fastmath, a, loc=loc, ip=ip)
 
 
+@dsl_user_op
 def exp2(
-    a: Union[TensorSSA, Numeric], fastmath: bool = False
+    a: Union[TensorSSA, Numeric], fastmath: bool = False, *, loc=None, ip=None
 ) -> Union[TensorSSA, Numeric]:
     """Compute element-wise base-2 exponential of the input tensor.
 
@@ -225,6 +263,10 @@ def exp2(
     :type a: Union[TensorSSA, Numeric]
     :param fastmath: Enable fast math optimizations, defaults to False
     :type fastmath: bool, optional
+    :param loc: Source location information, defaults to None
+    :type loc: Optional[Location]
+    :param ip: Insertion point for IR generation, defaults to None
+    :type ip: Optional[InsertionPoint]
     :return: Tensor containing 2 raised to the power of each element
     :rtype: Union[TensorSSA, Numeric]
 
@@ -236,11 +278,12 @@ def exp2(
         y = x.load()  # Load values
         z = exp2(y)  # Compute 2^x
     """
-    return _math_op(math.exp2, fastmath, a)
+    return _math_op(math.exp2, fastmath, a, loc=loc, ip=ip)
 
 
+@dsl_user_op
 def log(
-    a: Union[TensorSSA, Numeric], fastmath: bool = False
+    a: Union[TensorSSA, Numeric], fastmath: bool = False, *, loc=None, ip=None
 ) -> Union[TensorSSA, Numeric]:
     """Compute element-wise natural logarithm of the input tensor.
 
@@ -248,6 +291,10 @@ def log(
     :type a: Union[TensorSSA, Numeric]
     :param fastmath: Enable fast math optimizations, defaults to False
     :type fastmath: bool, optional
+    :param loc: Source location information, defaults to None
+    :type loc: Optional[Location]
+    :param ip: Insertion point for IR generation, defaults to None
+    :type ip: Optional[InsertionPoint]
     :return: Tensor containing the natural logarithm of each element
     :rtype: Union[TensorSSA, Numeric]
 
@@ -259,11 +306,12 @@ def log(
         y = x.load()  # Load values
         z = log(y)  # Compute natural logarithm
     """
-    return _math_op(math.log, fastmath, a)
+    return _math_op(math.log, fastmath, a, loc=loc, ip=ip)
 
 
+@dsl_user_op
 def log2(
-    a: Union[TensorSSA, Numeric], fastmath: bool = False
+    a: Union[TensorSSA, Numeric], fastmath: bool = False, *, loc=None, ip=None
 ) -> Union[TensorSSA, Numeric]:
     """Compute element-wise base-2 logarithm of the input tensor.
 
@@ -271,6 +319,10 @@ def log2(
     :type a: Union[TensorSSA, Numeric]
     :param fastmath: Enable fast math optimizations, defaults to False
     :type fastmath: bool, optional
+    :param loc: Source location information, defaults to None
+    :type loc: Optional[Location]
+    :param ip: Insertion point for IR generation, defaults to None
+    :type ip: Optional[InsertionPoint]
     :return: Tensor containing the base-2 logarithm of each element
     :rtype: Union[TensorSSA, Numeric]
 
@@ -282,11 +334,12 @@ def log2(
         y = x.load()  # Load values
         z = log2(y)  # Compute log base 2
     """
-    return _math_op(math.log2, fastmath, a)
+    return _math_op(math.log2, fastmath, a, loc=loc, ip=ip)
 
 
+@dsl_user_op
 def log10(
-    a: Union[TensorSSA, Numeric], fastmath: bool = False
+    a: Union[TensorSSA, Numeric], fastmath: bool = False, *, loc=None, ip=None
 ) -> Union[TensorSSA, Numeric]:
     """Compute element-wise base-10 logarithm of the input tensor.
 
@@ -294,6 +347,10 @@ def log10(
     :type a: Union[TensorSSA, Numeric]
     :param fastmath: Enable fast math optimizations, defaults to False
     :type fastmath: bool, optional
+    :param loc: Source location information, defaults to None
+    :type loc: Optional[Location]
+    :param ip: Insertion point for IR generation, defaults to None
+    :type ip: Optional[InsertionPoint]
     :return: Tensor containing the base-10 logarithm of each element
     :rtype: Union[TensorSSA, Numeric]
 
@@ -305,11 +362,12 @@ def log10(
         y = x.load()  # Load values
         z = log10(y)  # Compute log base 10
     """
-    return _math_op(math.log10, fastmath, a)
+    return _math_op(math.log10, fastmath, a, loc=loc, ip=ip)
 
 
+@dsl_user_op
 def rsqrt(
-    a: Union[TensorSSA, Numeric], fastmath: bool = False
+    a: Union[TensorSSA, Numeric], fastmath: bool = False, *, loc=None, ip=None
 ) -> Union[TensorSSA, Numeric]:
     """Compute element-wise reciprocal square root of the input tensor.
 
@@ -319,6 +377,10 @@ def rsqrt(
     :type a: Union[TensorSSA, Numeric]
     :param fastmath: Enable fast math optimizations, defaults to False
     :type fastmath: bool, optional
+    :param loc: Source location information, defaults to None
+    :type loc: Optional[Location]
+    :param ip: Insertion point for IR generation, defaults to None
+    :type ip: Optional[InsertionPoint]
     :return: Tensor containing the reciprocal square root of each element
     :rtype: Union[TensorSSA, Numeric]
 
@@ -330,11 +392,12 @@ def rsqrt(
         y = x.load()  # Load values
         z = rsqrt(y)  # Compute 1/√x
     """
-    return _math_op(math.rsqrt, fastmath, a)
+    return _math_op(math.rsqrt, fastmath, a, loc=loc, ip=ip)
 
 
+@dsl_user_op
 def sin(
-    a: Union[TensorSSA, Numeric], fastmath: bool = False
+    a: Union[TensorSSA, Numeric], fastmath: bool = False, *, loc=None, ip=None
 ) -> Union[TensorSSA, Numeric]:
     """Compute element-wise sine of the input tensor.
 
@@ -342,6 +405,10 @@ def sin(
     :type a: Union[TensorSSA, Numeric]
     :param fastmath: Enable fast math optimizations, defaults to False
     :type fastmath: bool, optional
+    :param loc: Source location information, defaults to None
+    :type loc: Optional[Location]
+    :param ip: Insertion point for IR generation, defaults to None
+    :type ip: Optional[InsertionPoint]
     :return: Tensor containing the sine of each element
     :rtype: Union[TensorSSA, Numeric]
 
@@ -353,11 +420,12 @@ def sin(
         y = x.load()  # Load values
         z = sin(y)  # Compute sine
     """
-    return _math_op(math.sin, fastmath, a)
+    return _math_op(math.sin, fastmath, a, loc=loc, ip=ip)
 
 
+@dsl_user_op
 def sqrt(
-    a: Union[TensorSSA, Numeric], fastmath: bool = False
+    a: Union[TensorSSA, Numeric], fastmath: bool = False, *, loc=None, ip=None
 ) -> Union[TensorSSA, Numeric]:
     """Compute element-wise square root of the input tensor.
 
@@ -365,6 +433,10 @@ def sqrt(
     :type a: Union[TensorSSA, Numeric]
     :param fastmath: Enable fast math optimizations, defaults to False
     :type fastmath: bool, optional
+    :param loc: Source location information, defaults to None
+    :type loc: Optional[Location]
+    :param ip: Insertion point for IR generation, defaults to None
+    :type ip: Optional[InsertionPoint]
     :return: Tensor containing the square root of each element
     :rtype: Union[TensorSSA, Numeric]
 
@@ -376,11 +448,12 @@ def sqrt(
         y = x.load()  # Load values
         z = sqrt(y)  # Compute square root
     """
-    return _math_op(math.sqrt, fastmath, a)
+    return _math_op(math.sqrt, fastmath, a, loc=loc, ip=ip)
 
 
+@dsl_user_op
 def tan(
-    a: Union[TensorSSA, Numeric], fastmath: bool = False
+    a: Union[TensorSSA, Numeric], fastmath: bool = False, *, loc=None, ip=None
 ) -> Union[TensorSSA, Numeric]:
     """Compute element-wise tangent of the input tensor.
 
@@ -388,6 +461,10 @@ def tan(
     :type a: Union[TensorSSA, Numeric]
     :param fastmath: Enable fast math optimizations, defaults to False
     :type fastmath: bool, optional
+    :param loc: Source location information, defaults to None
+    :type loc: Optional[Location]
+    :param ip: Insertion point for IR generation, defaults to None
+    :type ip: Optional[InsertionPoint]
     :return: Tensor containing the tangent of each element
     :rtype: Union[TensorSSA, Numeric]
 
@@ -399,11 +476,12 @@ def tan(
         y = x.load()  # Load values
         z = tan(y)  # Compute tangent
     """
-    return _math_op(math.tan, fastmath, a)
+    return _math_op(math.tan, fastmath, a, loc=loc, ip=ip)
 
 
+@dsl_user_op
 def tanh(
-    a: Union[TensorSSA, Numeric], fastmath: bool = False
+    a: Union[TensorSSA, Numeric], fastmath: bool = False, *, loc=None, ip=None
 ) -> Union[TensorSSA, Numeric]:
     """Compute element-wise hyperbolic tangent of the input tensor.
 
@@ -411,6 +489,10 @@ def tanh(
     :type a: Union[TensorSSA, Numeric]
     :param fastmath: Enable fast math optimizations, defaults to False
     :type fastmath: bool, optional
+    :param loc: Source location information, defaults to None
+    :type loc: Optional[Location]
+    :param ip: Insertion point for IR generation, defaults to None
+    :type ip: Optional[InsertionPoint]
     :return: Tensor containing the hyperbolic tangent of each element
     :rtype: Union[TensorSSA, Numeric]
 
@@ -422,7 +504,7 @@ def tanh(
         y = x.load()  # Load values
         z = tanh(y)  # Compute hyperbolic tangent
     """
-    return _math_op(math.tanh, fastmath, a)
+    return _math_op(math.tanh, fastmath, a, loc=loc, ip=ip)
 
 
 __all__ = [

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

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

python/CuTeDSL/cutlass/cute/nvgpu/common.py

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

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

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

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

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

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

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

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

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

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

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

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

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

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

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA
@@ -105,27 +105,17 @@ def make_smem_layout_atom(
         SmemLayoutAtomKind.MN_SW128_32B,
     ):
         # M/N-major layout
-        return core.make_composed_layout(
-            sw,
-            0,
-            core.make_layout(
-                (num_contiguous_elems, 8), stride=(1, num_contiguous_elems)
-            ),
-            loc=loc,
-            ip=ip,
+        outer = core.make_layout(
+            (num_contiguous_elems, 8), stride=(1, num_contiguous_elems), loc=loc, ip=ip
         )
     else:
         # K-major layout
-        return core.make_composed_layout(
-            sw,
-            0,
-            core.make_layout(
-                (8, num_contiguous_elems), stride=(num_contiguous_elems, 1)
-            ),
-            loc=loc,
-            ip=ip,
+        outer = core.make_layout(
+            (8, num_contiguous_elems), stride=(num_contiguous_elems, 1), loc=loc, ip=ip
         )
 
+    return core.make_composed_layout(sw, 0, outer, loc=loc, ip=ip)
+
 
 @overload
 def tile_to_mma_shape(

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

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

python/CuTeDSL/cutlass/cute/nvgpu/warp/__init__.py

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

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

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

python/CuTeDSL/cutlass/cute/nvgpu/warp/mma.py

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

python/CuTeDSL/cutlass/cute/nvgpu/warpgroup/__init__.py

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

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

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

python/CuTeDSL/cutlass/cute/nvgpu/warpgroup/mma.py

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

python/CuTeDSL/cutlass/cute/runtime.py

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

python/CuTeDSL/cutlass/cute/tensor.py

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA
@@ -426,9 +426,10 @@ class _Tensor(Tensor):
         return _cute_ir.get_layout(self.value, loc=loc, ip=ip)
 
     @property
+    @dsl_user_op
     @lru_cache_ir()
-    def shape(self) -> Shape:
-        return self.layout.shape
+    def shape(self, *, loc=None, ip=None) -> Shape:
+        return self.layout.shape_method(loc=loc, ip=ip)
 
     @property
     @lru_cache_ir()
@@ -1011,7 +1012,9 @@ def recast_tensor(
 
     src_iter = recast_ptr(src.iterator, dtype=dtype, loc=loc, ip=ip)
     src_layout = recast_layout(dst_width, src_width, src.layout, loc=loc, ip=ip)
-    return make_tensor(src_iter, src_layout, loc=loc, ip=ip)
+    return type(src)(
+        make_tensor(src_iter, src_layout, loc=loc, ip=ip), dtype=dtype, loc=loc, ip=ip
+    )
 
 
 @dsl_user_op
@@ -1344,7 +1347,23 @@ class TensorSSA(cutlass_arith.ArithValue):
             if issubclass(rhs.dtype, Integer):
                 rhs_val = rhs_val.with_signedness(rhs.dtype.signed)
 
-            res_vect = op(lhs_val, rhs_val)
+            # Use ArithValue's operator method directly to avoid recursion
+            # through TensorSSA's __add__/__sub__/etc. when op() dispatches
+            # back to the subclass method
+            if op.__name__ == "_min":
+                arith_op = cutlass_arith._min
+            elif op.__name__ == "_max":
+                arith_op = cutlass_arith._max
+            elif op in (operator.and_, operator.or_):
+                arith_op_name = f"__{op.__name__}_"
+                arith_op = getattr(cutlass_arith.ArithValue, arith_op_name)
+            else:
+                arith_op_name = f"__{op.__name__}__"
+                arith_op = getattr(cutlass_arith.ArithValue, arith_op_name, None)
+            if arith_op:
+                res_vect = arith_op(lhs_val, rhs_val, loc=loc, ip=ip)
+            else:
+                res_vect = op(lhs_val, rhs_val)
             res = TensorSSA(res_vect, lhs._shape, res_type)
 
         return res

python/CuTeDSL/cutlass/cute/testing.py

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

python/CuTeDSL/cutlass/cute/tuple.py

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

python/CuTeDSL/cutlass/cute/typing.py

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

python/CuTeDSL/cutlass/cutlass_dsl/__init__.py

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

python/CuTeDSL/cutlass/cutlass_dsl/cuda_jit_executor.py

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

python/CuTeDSL/cutlass/cutlass_dsl/cuda_stream_adapter.py

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

python/CuTeDSL/cutlass/cutlass_dsl/cutlass.py

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA
@@ -16,6 +16,7 @@ regarding to that dialect.
 
 # Local module imports
 from types import GenericAlias, SimpleNamespace, UnionType
+from typing_extensions import deprecated
 from typing import (
     Callable,
     Union,
@@ -103,6 +104,25 @@ from .cutlass_ast_decorators import (
 
 from ..base_dsl.runtime.jit_arg_adapters import JitArgAdapterRegistry
 
+# =============================================================================
+# Cutlass DSL Device Info
+# =============================================================================
+
+# Contains a map of SM architecture to shared memory capacity in bytes
+SMEM_CAPACITY_MAP = {
+    "sm_121": (100 - 1) * 1024,
+    "sm_120": (100 - 1) * 1024,
+    "sm_110": (228 - 1) * 1024,
+    "sm_103": (228 - 1) * 1024,
+    "sm_101": (228 - 1) * 1024,
+    "sm_100": (228 - 1) * 1024,
+    "sm_90": (228 - 1) * 1024,
+    "sm_89": (100 - 1) * 1024,
+    "sm_86": (100 - 1) * 1024,
+    "sm_87": (164 - 1) * 1024,
+    "sm_80": (164 - 1) * 1024,
+}
+
 # =============================================================================
 # Cutlass DSL Base Abstract Class
 # =============================================================================
@@ -812,6 +832,27 @@ class CutlassBaseDSL(BaseDSL):
                     )
                 cfg.smem = const(cfg.smem)
 
+                # Warn user if shared memory exceed arch max
+                # Currently runtime only show 'CUDA_ERROR_INVALID_VALUE' error which is not useful
+                arch = self.dsl.get_arch_enum()
+                arch_str = f"sm_{arch.major}{arch.minor}"
+                if arch_str in SMEM_CAPACITY_MAP:
+                    arch_smem = SMEM_CAPACITY_MAP[arch_str]
+                    smem_msg = (
+                        f"\nError: kernel '{kernelSym}' launch shared memory "
+                        f"exceeds current GPU arch {arch} allowed. "
+                        f"Allocated: {{}} bytes. Max: {arch_smem} bytes.\n\n"
+                    )
+                    if_generate(
+                        arch_smem < cfg.smem,
+                        lambda: cute.print_([cfg.smem], fmt=smem_msg),
+                        loc=loc,
+                    )
+                else:
+                    raise DSLRuntimeError(
+                        f"Lack smem capacity info for GPU arch {arch}."
+                    )
+
                 async_deps = cfg.async_deps
                 if not isinstance(cfg.async_deps, (list, tuple)):
                     async_deps = [cfg.async_deps]
@@ -1175,6 +1216,7 @@ class KernelLauncher:
         self.func_kwargs = func_kwargs
 
         self._name_prefix = func_kwargs.pop("_name_prefix", None)
+        self._launch_name = None
 
         self._check_func_args(funcBody, *func_args, **func_kwargs)
 
@@ -1192,7 +1234,10 @@ class KernelLauncher:
                 cause=e,
             )
 
-    def smem_usage(self) -> int:
+    @deprecated(
+        "`smem_usage()` is deprecated, use public API `arch.dynamic_smem_size()` instead."
+    )
+    def smem_usage(self) -> Int32:
         """
         Check smem usage for this kernel, only available after `launch`
         """
@@ -1215,6 +1260,7 @@ class KernelLauncher:
 
         ret, name = kernel_generator(*self.func_args, **self.func_kwargs, config=config)
         self.dsl.kernel_info[name] = kernel_attrs
+        self._launch_name = name
         return ret.launch_op_ret
 
     def __call__(self, *args, **kwargs):
@@ -1443,7 +1489,7 @@ def _minmax(op, *args, loc=None, ip=None):
         for x in xs:
             emitter = getattr(cutlass_arith, f"_{op.__name__}")
             if not (is_dynamic_expression(res) or is_dynamic_expression(x)):
-                res = emitter(op(res), op(x))
+                res = emitter(op(res), op(x), loc=loc, ip=ip)
             elif (
                 hasattr(res, "type")
                 and hasattr(x, "type")
@@ -1473,7 +1519,7 @@ def _minmax(op, *args, loc=None, ip=None):
                     rhs_val = rhs.value.with_signedness(rhs.signed)
                 else:
                     rhs_val = rhs.value
-                res = res_type(emitter(lhs_val, rhs_val), loc=loc, ip=ip)
+                res = res_type(emitter(lhs_val, rhs_val, loc=loc, ip=ip), loc=loc, ip=ip)
         x = res
     else:
         raise DSLNotImplemented(f"{type(args)} is not supported")

python/CuTeDSL/cutlass/cutlass_dsl/cutlass_ast_decorators.py

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

python/CuTeDSL/cutlass/cutlass_dsl/tvm_ffi_provider.py

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

python/CuTeDSL/cutlass/impl_utils.py

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

python/CuTeDSL/cutlass/jax/__init__.py

@@ -3,25 +3,48 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA
 # is strictly prohibited.
 
 from functools import cache
+import logging
+
+logger = logging.getLogger(__name__)
+
+# This is the minimum JAX version that will work with CuTeDSL JAX extensions.
+#
+# See the following pages for details on JAX versioning:
+#   - https://docs.jax.dev/en/latest/jep/25516-effver.html
+#   - https://docs.jax.dev/en/latest/jep/9419-jax-versioning.html
+CUTE_DSL_MIN_SUPPORTED_JAX_VERSION = (0, 5, 0)
+
 
 @cache
 def is_available():
-    """Returns true of Jax support is enabled."""
+    """Returns true if JAX extensions are supported and available."""
     try:
         import jax
+        import jax.numpy  # Also verify jax.numpy is available
+    except ImportError:
+        logger.debug(
+            "CuTeDSL JAX extensions are not available because JAX was not found or could not be imported."
+        )
+        return False
 
-        _HAVE_JAX = True
-    except ImportError as e:
-        _HAVE_JAX = False
+    if not (
+        hasattr(jax.version, "__version_info__")
+        and jax.version.__version_info__ >= CUTE_DSL_MIN_SUPPORTED_JAX_VERSION
+    ):
+        logger.debug(
+            f"Your installed JAX v{jax.__version__} too old and not supported by CuTeDSL JAX extensions.\n"
+            "Please upgrade to the latest version."
+        )
+        return False
 
-    return _HAVE_JAX
+    return True
 
 
 if is_available():
@@ -29,6 +52,8 @@ if is_available():
     from .types import (
         jax_to_cutlass_dtype,
         cutlass_to_jax_dtype,
+        jax_to_cutlass_layout_order,
+        cutlass_to_jax_layout_order,
         from_dlpack,
         JaxArray,
         TensorSpec,
@@ -41,6 +66,7 @@ if is_available():
         find_cute_dsl_runtime_library,
         register_ffi,
         is_ffi_registered,
+        get_cutlass_call_ffi_version,
     )
     from . import testing
 
@@ -48,18 +74,24 @@ if is_available():
     TensorMode = TensorSpec
 
     __all__ = [
+        "CUTE_DSL_MIN_SUPPORTED_JAX_VERSION",
         "cutlass_call",
         "jax_to_cutlass_dtype",
         "cutlass_to_jax_dtype",
+        "jax_to_cutlass_layout_order",
+        "cutlass_to_jax_layout_order",
         "from_dlpack",
         "JaxArray",
         "TensorSpec",
         "TensorMode",
         "release_compile_cache",
         "get_export_disabled_safety_checks",
+        "is_ffi_registered",
+        "register_ffi",
+        "get_cutlass_call_ffi_version",
         "is_available",
         "testing",
     ]
 else:
     # export is_available check for callers or tests.
-    __all__ = ["is_available"]
+    __all__ = ["CUTE_DSL_MIN_SUPPORTED_JAX_VERSION", "is_available"]

python/CuTeDSL/cutlass/jax/compile.py

@@ -3,20 +3,16 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA
 # is strictly prohibited.
 
-import os
 import gc
-import ctypes
-import inspect
-from typing import Any, Callable, Optional, Sequence
+from typing import Any
 from dataclasses import dataclass
 from functools import partial
-from pathlib import Path
 
 import time
 import logging
@@ -27,17 +23,13 @@ import cuda.bindings.driver as cuda
 
 import jax
 import jax.numpy as jnp
-import jaxlib
 
 from .types import (
     jax_to_cutlass_dtype,
-    from_dlpack,
-    JaxArray,
     JaxArrayList,
     TensorSpec,
     JaxTracedArray,
     DEFAULT_CUTLASS_DEVICE_MEMSPACE,
-    DEFAULT_CUTLASS_DEVICE_BUFFER_ALIGNMENT,
 )
 
 import cutlass
@@ -90,6 +82,7 @@ class FunctionSpec:
                 leaf.spec.layout,
                 leaf.spec.mode,
                 leaf.get_static_flag(self.use_static_tensors),
+                leaf.spec.divisibility,
             )
             for leaf in self.in_args
         ]
@@ -102,6 +95,7 @@ class FunctionSpec:
                 leaf.spec.layout,
                 leaf.spec.mode,
                 leaf.get_static_flag(self.use_static_tensors),
+                leaf.spec.divisibility,
             )
             for leaf in self.out_args
         ]

python/CuTeDSL/cutlass/jax/ffi.py

@@ -3,16 +3,15 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA
 # is strictly prohibited.
 
-from typing import Sequence
+from typing import Sequence, Optional
 from pathlib import Path
 from functools import cache
-import os
 import logging
 import ctypes
 
@@ -27,22 +26,51 @@ logger = logging.getLogger(__name__)
 
 _CUTE_DSL_RUNTIME_LIBRARY_NAME = "cute_dsl_runtime"
 
-_CUTLASS_CALL_TARGETS = {
+# V1 targets for older jax clients
+_CUTLASS_CALL_TARGETS_V1 = {
     "CuteDSLRT_NvJaxCutlassCall": {
-        "execute": "CuteDSLRT_NvJaxCutlassCallExecute",
-        "prepare": "CuteDSLRT_NvJaxCutlassCallPrepare",
+        "prepare": "CuteDSLRT_NvJaxCutlassCallPrepare_v1",
+        "execute": "CuteDSLRT_NvJaxCutlassCallExecute_v1",
     },
     "CuteDSLRT_NvJaxCutlassCallNoCudaGraph": {
-        "execute": "CuteDSLRT_NvJaxCutlassCallExecuteNoCudaGraph",
-        "prepare": "CuteDSLRT_NvJaxCutlassCallPrepare",
+        "prepare": "CuteDSLRT_NvJaxCutlassCallPrepare_v1",
+        "execute": "CuteDSLRT_NvJaxCutlassCallExecuteNoCudaGraph_v1",
     },
 }
 
+# V2 targets for newer jax clients supporting stateful FFI calls.
+_JAX_FFI_V2_MIN_VERSION = (0, 9, 1)
+_CUTLASS_CALL_TARGETS_V2 = {
+    "CuteDSLRT_NvJaxCutlassCall": {
+        "execute": "CuteDSLRT_NvJaxCutlassCallExecute_v2",
+        "instantiate": "CuteDSLRT_NvJaxCutlassCallInstantiate_v2",
+        "prepare": "CuteDSLRT_NvJaxCutlassCallPrepare_v2",
+    },
+    "CuteDSLRT_NvJaxCutlassCallNoCudaGraph": {
+        "execute": "CuteDSLRT_NvJaxCutlassCallExecuteNoCudaGraph_v2",
+        "instantiate": "CuteDSLRT_NvJaxCutlassCallInstantiate_v2",
+        "prepare": "CuteDSLRT_NvJaxCutlassCallPrepare_v2",
+    },
+}
+_CUTLASS_CALL_TYPES_V2 = {
+    "CuteDSLRT_NvJaxCutlassCallTypes": {
+        "type_id": "CuteDSLRT_NvJaxCutlassCallStateTypeId_v2",
+        "type_info": "CuteDSLRT_NvJaxCutlassCallStateTypeInfo_v2",
+    }
+}
 
-def get_cutlass_call_ffi_name(allow_cuda_graph):
+
+def get_cutlass_call_ffi_version() -> int:
+    """Returns the FFI API version based on JAX version."""
+    if jax.version.__version_info__ >= _JAX_FFI_V2_MIN_VERSION:
+        return 2
+    else:
+        return 1
+
+
+def get_cutlass_call_ffi_name(allow_cuda_graph: bool) -> str:
     """Returns the FFI target to call when running cutlass_call functions."""
-    disable_cuda_graph = not allow_cuda_graph
-    if not disable_cuda_graph:
+    if allow_cuda_graph:
         return "CuteDSLRT_NvJaxCutlassCall"
     else:
         return "CuteDSLRT_NvJaxCutlassCallNoCudaGraph"
@@ -50,14 +78,14 @@ def get_cutlass_call_ffi_name(allow_cuda_graph):
 
 def get_export_disabled_safety_checks() -> Sequence[jax.export.DisabledSafetyCheck]:
     """Returns jax.export.DisabledSafetyCheck to allow cutlass_call kernels."""
-    checks = []
-    for target in _CUTLASS_CALL_TARGETS:
-        checks.append(jax.export.DisabledSafetyCheck.custom_call(target))
-    return tuple(checks)
+    targets = set(_CUTLASS_CALL_TARGETS_V1.keys()) | set(
+        _CUTLASS_CALL_TARGETS_V2.keys()
+    )
+    return tuple([jax.export.DisabledSafetyCheck.custom_call(t) for t in targets])
 
 
 @cache
-def find_cute_dsl_runtime_library():
+def find_cute_dsl_runtime_library() -> Optional[str]:
     """Searches for the CuTeDSL runtime library."""
     dsl = CuTeDSL._get_dsl()
     candidate_libs = []
@@ -85,7 +113,10 @@ def find_cute_dsl_runtime_library():
         candidate_libs.extend(dsl_libs)
 
     except Exception as e:
-        logger.debug(f"Failed to locate libraries due to an exception:", e)
+        logger.debug(
+            f"Failed to locate {_CUTE_DSL_RUNTIME_LIBRARY_NAME} library: {e}",
+            exc_info=True,
+        )
 
     for lib in candidate_libs:
         if lib.endswith(f"{_CUTE_DSL_RUNTIME_LIBRARY_NAME}.so"):
@@ -97,8 +128,12 @@ def find_cute_dsl_runtime_library():
 _FFI_CALLS_REGISTERED = False
 
 
-def register_ffi():
-    """Registers custom calls with Jax/XLA runtime."""
+def register_ffi(ffi_version: int = get_cutlass_call_ffi_version()):
+    """Registers custom calls with Jax/XLA runtime.
+
+    A specific version can be requested using `ffi_version` argument. Attempting
+    to register non default FFI versions may not work with your specific JAX.
+    """
     global _FFI_CALLS_REGISTERED
     if _FFI_CALLS_REGISTERED:
         return
@@ -112,27 +147,36 @@ def register_ffi():
 
     lib = ctypes.CDLL(runtime_library)
 
-    def _capsule(funcptr):
-        destructor = ctypes.CFUNCTYPE(None, ctypes.py_object)
-        builder = ctypes.pythonapi.PyCapsule_New
-        builder.restype = ctypes.py_object
-        builder.argtypes = (ctypes.c_void_p, ctypes.c_char_p, destructor)
-        return builder(funcptr, None, destructor(0))
-
     def _register_ffi_targets(lib, targets):
         for target_name, target in targets.items():
             handler = {}
             for stage, fn_name in target.items():
                 fn = getattr(lib, fn_name)
                 fn.restype = ctypes.c_void_p
-                handler[stage] = _capsule(fn)
+                handler[stage] = jax.ffi.pycapsule(fn)
             logger.debug(f"Registering ffi handler: {target_name}, {handler}")
-            jax.ffi.register_ffi_target(
-                target_name, handler["execute"], platform="CUDA"
-            )
+            jax.ffi.register_ffi_target(target_name, handler, platform="CUDA")
 
-    # Register the custom FFI targets
-    _register_ffi_targets(lib, _CUTLASS_CALL_TARGETS)
+    def _register_ffi_types(lib, types):
+        for type_name, type_dict_targets in types.items():
+            type_dict = {}
+            for field, fn_name in type_dict_targets.items():
+                fn = getattr(lib, fn_name)
+                fn.restype = ctypes.c_void_p
+                type_dict[field] = jax.ffi.pycapsule(fn())
+            logger.debug(f"Registering ffi type: {type_name}, {type_dict}")
+            jax.ffi.register_ffi_type(type_name, type_dict, platform="CUDA")
+
+    # Register the custom FFI targets.
+    match ffi_version:
+        case 1:
+            _register_ffi_targets(lib, _CUTLASS_CALL_TARGETS_V1)
+            # no types for v1
+        case 2:
+            _register_ffi_types(lib, _CUTLASS_CALL_TYPES_V2)
+            _register_ffi_targets(lib, _CUTLASS_CALL_TARGETS_V2)
+        case _:
+            raise ValueError(f"Invalid FFI version {ffi_version}")
 
     _FFI_CALLS_REGISTERED = True
 

python/CuTeDSL/cutlass/jax/primitive.py

@@ -3,33 +3,28 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA
 # is strictly prohibited.
 
-from typing import Any, Union, Sequence, Callable
-from functools import partial
+from typing import Any, Sequence, Callable
 import logging
-import os
 
-import cuda.bindings.driver as cuda
 
-import jax, jax.numpy as jnp
+import jax
 import jax.extend
 from jax.interpreters import mlir
 from jax._src.interpreters import ad
 from jax._src.interpreters import batching
-from jax._src import ffi
-from jax.tree import flatten, unflatten
 
-import cutlass
 
 from .compile import get_or_compile_kernel, build_function_spec
-from .types import row_major_layout, default_tensor_spec, TensorSpec
+from .types import cutlass_to_jax_layout_order, default_tensor_spec, TensorSpec
 from .ffi import get_cutlass_call_ffi_name, is_ffi_registered, register_ffi
 
+
 logger = logging.getLogger(__name__)
 
 cutlass_call_inner_p = jax.extend.core.Primitive("cutlass_call_inner")
@@ -39,7 +34,7 @@ cutlass_call_inner_p.multiple_results = True
 def cutlass_call(
     fn: Callable[..., None],
     *,
-    output_shape_dtype: Any,
+    output_shape_dtype: Any = None,
     input_spec: Any = None,
     output_spec: Any = None,
     input_mode: Any = None,
@@ -50,29 +45,64 @@ def cutlass_call(
     use_static_tensors=False,
     **kwargs,
 ):
-    """Creates a callable that invokes a @cute.jit function.
+    """Create a callable that invokes a ``@cute.jit`` function from JAX.
+
+    Returns a callable that accepts JAX arrays and dispatches to *fn* as part
+    of a ``jax.jit``-compiled computation.  The kernel is compiled once on the
+    first call and cached for subsequent invocations with the same shapes and
+    specs.
+
+    Example::
+
+        @cute.jit
+        def my_kernel(stream, A, B, C, D):
+            ...
+
+        @jax.jit
+        def run(a, b):
+            return cutlass_call(
+                my_kernel,
+                output_shape_dtype=(
+                    jax.ShapeDtypeStruct(a.shape, a.dtype),
+                    jax.ShapeDtypeStruct(b.shape, b.dtype),
+                ),
+            )(a, b)
+
+        c, d = run(a, b)
 
     Args:
-        fn: A @cute.jit decorated function that launches a cutlass kernel.
-        output_shape_dtype: A pytree representing the shape and dtype of the output buffers.
-        input_output_aliases: Optional mapping of input to output aliases. Positions are specified assuming
-            a flattened input and output pytree.
-        input_spec: Specifies a cute.Tensor dimension order for input tensors. If None then the order
-            will assume the corresponding layout order.
-        output_spec: Specifies a cute.Tensor dimension order for output tensors. If None then the order
-            will assume the corresponding layout order.
-        input_mode: Legacy alias for input_spec. This parameter may be removed in future versions.
-        output_spec: Legacy alias for output_spec. This parameter may be removed in future versions.
-        allow_cuda_graph: If false will prevent XLA from building a cuda graph of for this call.
-        compile_options: Optional compiler arguments to pass into cute.compile.
-        use_static_tensors: If True, tensor shapes and strides are treated as constexpr values by
-        default. This can improve performance through compiler specialization but may not work
-        properly with all kernels. Specific tensors may be marked static or dynamic using the mode
-        and override this flag.
-        kwargs: Optional constexpr parameters to pass into the kernel fn.
+        fn: A ``@cute.jit``-decorated function with the signature
+            ``(stream, *inputs, *outputs, **kwargs)``.
+        output_shape_dtype: A pytree of :class:`jax.ShapeDtypeStruct` (or
+            objects with ``.shape`` and ``.dtype`` attributes) describing each
+            output buffer.
+        input_spec: A :class:`TensorSpec` or list thereof providing
+            layout/mode/divisibility hints for input tensors.  ``None`` infers
+            defaults from each array.
+        output_spec: Same as *input_spec* but applied to output tensors.
+        input_output_aliases: ``{input_index: output_index}`` mapping that
+            allows an input buffer to alias an output, avoiding an extra copy.
+            Indices are into the flattened input and output pytrees.
+        allow_cuda_graph: If ``False``, prevents XLA from capturing this call
+            in a CUDA graph.  Defaults to ``True``.
+        compile_options: Optional dict of compiler flags forwarded to
+            ``cute.compile``.
+        use_static_tensors: If ``True``, tensor shapes and strides are baked in
+            as compile-time constants, improving performance when shapes are
+            fixed across calls.  Defaults to ``False``.
+        **kwargs: Additional keyword arguments forwarded to *fn* as compile-time
+            constants.
 
-    Note: This API is experimental and subject to change!
+    Returns:
+        A callable ``(*arrays) -> output_pytree`` that can be used inside
+        ``jax.jit``.
+
+    Note:
+        This API is experimental and subject to change.
     """
+    if output_shape_dtype is None:
+        raise ValueError("'output_shape_dtype' must be specified.")
+
     output_shape_dtype = jax.tree.map(
         lambda leaf: jax.ShapeDtypeStruct(leaf.shape, leaf.dtype), output_shape_dtype
     )
@@ -107,22 +137,84 @@ def cutlass_call(
     )
 
 
-def _normalize_tensor_spec(value: Any):
-    if value is None:
-        return [None]
-    elif isinstance(value, (tuple, list)):
-        if isinstance(value[0], int):  # single tuple of modes
-            return TensorSpec(mode=tuple(value))
+def _is_spec_leaf(x: Any) -> bool:
+    """Return True if *x* should be treated as a leaf when traversing a spec pytree.
+
+    Stops traversal at ``TensorSpec`` and ``None`` (both are valid leaf specs) and at
+    bare integer sequences (the legacy mode-spec shorthand).  Everything else is
+    treated as a pytree container and recursed into by ``jax.tree.leaves``.
+    """
+    if x is None or isinstance(x, TensorSpec):
+        return True
+    # Legacy: a bare sequence of ints represents a single TensorSpec(mode=...).
+    # Check *all* elements so that mixed sequences like (1, TensorSpec()) are NOT
+    # mistaken for a mode spec and instead cause a TypeError below.
+    if isinstance(x, (list, tuple)) and bool(x) and all(isinstance(i, int) for i in x):
+        return True
+    return False
+
+
+def _normalize_tensor_spec(value: Any) -> list[TensorSpec | None]:
+    """Normalize a spec pytree into a flat list of ``TensorSpec | None`` entries.
+
+    *value* may be any JAX pytree whose leaves are ``TensorSpec``, ``None``, or a
+    bare integer sequence (legacy shorthand for ``TensorSpec(mode=...)``).  Dict and
+    other non-list/tuple pytree containers are supported via ``jax.tree.leaves``.
+
+    Note: ``TensorSpec`` is itself a JAX-registered dataclass with all-static fields,
+    so traversal *must* use the ``is_leaf`` predicate to stop at ``TensorSpec`` nodes
+    rather than recursing into them (which would yield no children and silently drop
+    the spec).
+    """
+    leaves = jax.tree.leaves(value, is_leaf=_is_spec_leaf)
+    result = []
+    for leaf in leaves:
+        if leaf is None or isinstance(leaf, TensorSpec):
+            result.append(leaf)
+        elif isinstance(leaf, (list, tuple)):
+            # Legacy: bare int sequence → TensorSpec(mode=...)
+            result.append(TensorSpec(mode=tuple(leaf)))
         else:
-            flat, _ = jax.tree.flatten(
-                [_normalize_tensor_spec(x) for x in value],
-                is_leaf=lambda x: x is None or isinstance(x, TensorSpec),
+            raise TypeError(
+                f"Unexpected value for TensorSpec: {leaf!r} ({type(leaf).__name__})"
+            )
+    return result
+
+
+def _resolve_spec_flat(spec: Any, tensors: list) -> tuple[TensorSpec, ...]:
+    """Normalize *spec* and fill any ``None`` slots with defaults inferred from *tensors*."""
+    if spec is None:
+        return tuple(default_tensor_spec(t) for t in tensors)
+    specs = list(_normalize_tensor_spec(spec))
+    if len(specs) != len(tensors):
+        raise ValueError(
+            f"Must have the same number of specs ({len(specs)}) as tensors ({len(tensors)})."
+        )
+    return tuple(
+        default_tensor_spec(t) if s is None else s for s, t in zip(specs, tensors)
+    )
+
+
+def _validate_specs(label: str, tensors: list, specs: tuple[TensorSpec, ...]) -> None:
+    """Validate that each spec's rank-dependent fields match the corresponding tensor shape."""
+    for idx, (tensor, spec) in enumerate(zip(tensors, specs)):
+        ndim = len(tensor.shape)
+        if spec.layout is not None and len(spec.layout) != ndim:
+            raise ValueError(
+                f"{label} #{idx} has invalid layout {spec.layout} for shape {tensor.shape}."
+            )
+        if spec.mode is not None and len(spec.mode) != ndim:
+            raise ValueError(
+                f"{label} #{idx} has invalid mode {spec.mode} for shape {tensor.shape}."
+            )
+        if (
+            spec.divisibility is not None
+            and not isinstance(spec.divisibility, int)
+            and len(spec.divisibility) != ndim
+        ):
+            raise ValueError(
+                f"{label} #{idx} has invalid divisibility {spec.divisibility} for shape {tensor.shape}."
             )
-            return flat
-    elif isinstance(value, TensorSpec):
-        return [value]
-    else:
-        raise TypeError(f"Unexpected value for TensorMode {value} {type(value)}")
 
 
 def _cutlass_call_impl(
@@ -137,70 +229,21 @@ def _cutlass_call_impl(
     use_static_tensors,
     **kwargs,
 ):
+    # A single ShapeDtypeStruct means one output; a sequence means multiple.
     multiple_results = isinstance(output_shape_dtype, Sequence)
     if not multiple_results:
         output_shape_dtype = (output_shape_dtype,)
     output_shape_dtype_flat, output_tree = jax.tree.flatten(output_shape_dtype)
 
-    @partial(jax.jit, inline=True)
+    @jax.jit
     def call_wrapper(*args):
         args_flat, args_tree = jax.tree.flatten(args)
 
-        if input_spec is None:
-            input_spec_flat = tuple(default_tensor_spec(x) for x in args_flat)
-        else:
-            input_spec_flat = _normalize_tensor_spec(input_spec)
-            for idx, (spec, arg) in enumerate(zip(input_spec_flat, args_flat)):
-                if spec is None:
-                    input_spec_flat[idx] = default_tensor_spec(arg)
-            input_spec_flat = tuple(input_spec_flat)
+        input_spec_flat = _resolve_spec_flat(input_spec, args_flat)
+        output_spec_flat = _resolve_spec_flat(output_spec, output_shape_dtype_flat)
 
-        if output_spec is None:
-            output_spec_flat = tuple(
-                default_tensor_spec(x) for x in output_shape_dtype_flat
-            )
-        else:
-            output_spec_flat = _normalize_tensor_spec(output_spec)
-            for idx, (spec, arg) in enumerate(
-                zip(output_spec_flat, output_shape_dtype_flat)
-            ):
-                if spec is None:
-                    output_spec_flat[idx] = default_tensor_spec(arg)
-            output_spec_flat = tuple(output_spec_flat)
-        if len(input_spec_flat) != len(args_flat):
-            raise ValueError(
-                f"Must has same number of input modes ({len(input_spec_flat)}) as input arrays ({len(args_flat)})."
-            )
-
-        if len(output_spec_flat) != len(output_shape_dtype_flat):
-            raise ValueError(
-                f"Must has same number of output modes ({len(output_spec_flat)}) as output arrays ({len(output_shape_dtype_flat)})."
-            )
-
-        # Validate dynamic mode settings match whatever static shape
-        # information we got as input.
-        for idx, (arg, spec) in enumerate(zip(args_flat, input_spec_flat)):
-            if spec.layout is not None and len(spec.layout) != len(arg.shape):
-                raise ValueError(
-                    f"Input #{idx} has invalid layout {spec.layout} for shape {arg.shape}."
-                )
-            if spec.mode is not None and len(spec.mode) != len(arg.shape):
-                raise ValueError(
-                    f"Input #{idx} has invalid mode {spec.mode} for shape {arg.shape}."
-                )
-
-        for idx, (arg, spec) in enumerate(
-            zip(output_shape_dtype_flat, output_spec_flat)
-        ):
-            if spec.layout is not None and len(spec.layout) != len(arg.shape):
-                raise ValueError(
-                    f"Output #{idx} has invalid layout {spec.layout} for shape {arg.shape}."
-                )
-
-            if spec.mode is not None and len(spec.mode) != len(arg.shape):
-                raise ValueError(
-                    f"Output #{idx} has invalid mode {spec.mode} for shape {arg.shape}."
-                )
+        _validate_specs("Input", args_flat, input_spec_flat)
+        _validate_specs("Output", output_shape_dtype_flat, output_spec_flat)
 
         output_flat = cutlass_call_inner_p.bind(
             *args_flat,
@@ -208,8 +251,8 @@ def _cutlass_call_impl(
             args_tree=args_tree,
             output_shape_dtype_flat=tuple(output_shape_dtype_flat),
             output_tree=output_tree,
-            input_spec_flat=tuple(input_spec_flat),
-            output_spec_flat=tuple(output_spec_flat),
+            input_spec_flat=input_spec_flat,
+            output_spec_flat=output_spec_flat,
             input_output_aliases=tuple(input_output_aliases.items()),
             allow_cuda_graph=allow_cuda_graph,
             compile_options=compile_options,
@@ -264,10 +307,17 @@ def cutlass_call_inner_p_impl(
         register_ffi()
 
     call_name = get_cutlass_call_ffi_name(allow_cuda_graph)
+
+    # Convert layout from CuTeDSL to JAX order as ffi_call expects this.
+    input_layouts = [cutlass_to_jax_layout_order(s.layout) for s in input_spec_flat]
+    output_layouts = [cutlass_to_jax_layout_order(s.layout) for s in output_spec_flat]
+
     fun = jax.ffi.ffi_call(
         call_name,
         result_shape_dtypes=output_shape_dtype_flat,
         input_output_aliases=dict(spec.input_output_aliases),
+        input_layouts=input_layouts,
+        output_layouts=output_layouts,
     )
 
     return fun(*args_flat, module=kernel.module, key=kernel.fingerprint)

python/CuTeDSL/cutlass/jax/testing.py

@@ -3,13 +3,12 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA
 # is strictly prohibited.
 
-from functools import partial
 
 import jax
 import jax.numpy as jnp
@@ -17,6 +16,7 @@ import jax.numpy as jnp
 import cutlass.cute as cute
 from cutlass.cutlass_dsl import dsl_user_op
 
+
 def reorder_modes(src: str, target: str) -> tuple[int, ...]:
     """Computes the mode given a source and target order."""
     src = tuple(src)
@@ -88,6 +88,7 @@ def get_gemm_shape_from_tensors(
     n = b.shape[0]
     return (m, n, k, l)
 
+
 def create_tensor(
     shape, dtype, key, *, minval=-2.0, maxval=2.0, fill_value=None, fill_arange=False
 ):

python/CuTeDSL/cutlass/jax/types.py

@@ -3,24 +3,15 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA
 # is strictly prohibited.
 
-from typing import Type, Optional, Sequence, Union, Callable, Any, TypeVar
-import sys
-import ctypes
-import math
-import inspect
+from typing import Sequence
 from dataclasses import dataclass, field
-from functools import partial, reduce
-from operator import mul
-from itertools import chain
-from typing import Annotated
 
-import cuda.bindings.driver as cuda
 
 import jax
 import jax.numpy as jnp
@@ -28,13 +19,13 @@ import jax.numpy as jnp
 import cutlass
 import cutlass.cute as cute
 from cutlass.cute.runtime import from_dlpack as _from_dlpack
-from cutlass.cute import AddressSpace, Numeric, IntTuple
+from cutlass.cute import AddressSpace
 from cutlass._mlir import ir
 from cutlass._mlir.dialects import llvm, arith
-import cutlass._mlir.dialects.cute as _cute_ir
 
 JAX_DTYPE_TO_CUTLASS_DTYPE = {
     jnp.bool.dtype: cutlass.Boolean,
+    jnp.int4.dtype: cutlass.Int4,
     jnp.int8.dtype: cutlass.Int8,
     jnp.int16.dtype: cutlass.Int16,
     jnp.int32.dtype: cutlass.Int32,
@@ -65,38 +56,69 @@ DEFAULT_CUTLASS_DEVICE_BUFFER_ALIGNMENT = 256
 @jax.tree_util.register_dataclass
 @dataclass(frozen=True)
 class TensorSpec:
-    """Provides a specification of cute.Tensor modes and additional metadata about
-    dynamic/static shapes for compilation.
+    """Specifies the layout and metadata for a JAX array passed to a CuTe kernel.
 
-    Arguments:
-        layout : Specifies the position of stries as they relate to the framework
-                 tensor (S0, S1, ... SN)
-        mode : Specifies the position of each mode in the tensor (M0, M1, ... MN)
-        static : Specifies the tensor shape is represented as static constexpr.
-        ptr_assumed_align: Specifies the pointer alignment.
+    TensorSpec controls how a JAX array's dimensions are mapped to a cute.Tensor
+    during jit lowering, including stride ordering, mode permutation, and whether
+    shapes/strides are compiled as static constants.
+
+    Attributes:
+        layout: A minor-to-major stride ordering in CuTeDSL convention. ``layout[i]``
+            gives the stride rank of dimension ``i``, where rank 0 means the smallest
+            (innermost) stride. For example, row-major order for a 3-D tensor is
+            ``(2, 1, 0)``. If ``None``, row-major is assumed. Use
+            :func:`jax_to_cutlass_layout_order` to convert from JAX's major-to-minor
+            convention.
+        mode: A permutation that maps the stride-ordered dimensions to the mode
+            positions of the resulting ``cute.Layout``. For example, ``mode=(2, 0, 1)``
+            reorders an ``(M, K, L)`` layout into ``(K, L, M)`` mode order inside the
+            kernel. If ``None``, modes match the natural dimension order ``(0, 1, ..., N-1)``.
+        static: If ``True``, shapes and strides are compiled as static ``constexpr``
+            values, which may enable additional compiler optimisations. Kernels that
+            do not support static shapes will raise a compile error. Must be ``False``
+            when any dimension is symbolic (e.g. under ``jax.export``).
+        ptr_assumed_align: Assumed byte alignment of the tensor's data pointer.
+            Overrides the default of 256 bytes. Rarely needs to change.
+        divisibility: Optional per-mode divisibility hints. If a single int is passed
+            divisibility will be applied to the leading (stride=1) dimension only.
     """
 
-    # Specifies the layout of the Jax array. If not it will be assumed that the layout
-    # is row major.
+    # Minor-to-major stride ordering in CuTeDSL convention (layout[i] = stride rank
+    # of dimension i, 0 = innermost). Defaults to row-major if None.
     layout: tuple[int, ...] | None = field(metadata=dict(static=True), default=None)
-    # Indicates the order of modes. If unspecified the modes will match exactly with
-    # the layout of the Jax tensor (e.g. row-major). Typically used to map from the
-    # input layout to kernel layouts (e.g. MKL/NKL/MNL).
+    # Permutation from stride-ordered dimensions to cute.Layout mode positions.
+    # Defaults to identity (0, 1, ..., N-1) if None.
     mode: tuple[int, ...] | None = field(metadata=dict(static=True), default=None)
-    # Indicates the shape and strides will be defined statically. Setting True ay enable
-    # additional optimization. Kernels that do not support static shapes will generate
-    # compile errors if this is enabled so we leave it off by default.
+    # If True, shapes and strides are embedded as compile-time constants.
+    # Must be False for symbolic/dynamic shapes (e.g. jax.export).
     static: bool = field(metadata=dict(static=True), default=None)
-    # Overrides the default pointer alignment. Generally this should not be changed
-    # but is left here to provide a hook.
+    # Assumed alignment (bytes) of the data pointer. Default matches XLA's 256-byte alignment.
     ptr_assumed_align: int = field(
         metadata=dict(static=True), default=DEFAULT_CUTLASS_DEVICE_BUFFER_ALIGNMENT
     )
+    # Per-mode divisibility hints.
+    divisibility: tuple[int | None, ...] | int | None = field(
+        metadata=dict(static=True), default=None
+    )
+
 
 def row_major_layout(shaped):
-    """Returns a row major layout given a shaped value.
+    """Returns the CuTeDSL minor-to-major stride ordering for a row-major (C-contiguous) tensor.
 
-    Row major layout is (N-1, N-2, ... 1, 0) for an N-dimensional tensor.
+    In CuTeDSL convention, ``layout[i]`` is the stride rank of dimension ``i``,
+    where rank 0 denotes the innermost (stride-1) dimension.  Row-major means the
+    last dimension is innermost, so the result is ``(N-1, N-2, ..., 1, 0)`` for an
+    N-dimensional tensor.
+
+    Example::
+
+        row_major_layout((M, K, N))  # → (2, 1, 0)
+
+    Args:
+        shaped: An object with a ``.shape`` attribute, or a shape tuple/sequence.
+
+    Returns:
+        A tuple of length N representing the minor-to-major ordering.
     """
     if hasattr(shaped, "shape"):
         shaped = shaped.shape
@@ -104,9 +126,17 @@ def row_major_layout(shaped):
 
 
 def default_tensor_mode(shaped):
-    """Returns a default tensor mode given a shaped value.
+    """Returns the identity mode permutation for an N-dimensional tensor.
 
-    Default mode is (0, 1, ... N-2, N-1) for an N_dimensional tensor.
+    The mode permutation maps stride-ordered dimensions to ``cute.Layout`` mode
+    positions.  The default identity ``(0, 1, ..., N-1)`` leaves the mode order
+    unchanged relative to the dimension order.
+
+    Args:
+        shaped: An object with a ``.shape`` attribute, or a shape tuple/sequence.
+
+    Returns:
+        A tuple ``(0, 1, ..., N-1)`` of length N.
     """
     if hasattr(shaped, "shape"):
         shaped = shaped.shape
@@ -114,14 +144,108 @@ def default_tensor_mode(shaped):
 
 
 def default_tensor_spec(shaped) -> TensorSpec:
-    """Returns a default tensor spec given a shaped value.
+    """Returns a :class:`TensorSpec` with row-major layout and identity mode ordering.
 
-    Default layout is (N-1, N-2, ... 1, 0) for an N-dimensional tensor.
-    Default mode is (0, 1, ... N-2, N-1) for an N_dimensional tensor.
+    Equivalent to::
+
+        TensorSpec(layout=(N-1, ..., 1, 0), mode=(0, 1, ..., N-1), divisibility=(D0, D1, ... DN-1))
+
+    This is appropriate for standard row-major (C-contiguous) JAX arrays that
+    do not require dimension reordering inside the kernel.
+
+    Divisibility hints are inferred only for concrete integer dimensions. 
+    Symbolic dimensions always produce ``None`` for their slot; pass an
+    explicit ``TensorSpec`` with ``divisibility`` set if you need alignment
+    hints for symbolic shapes.
+
+    Args:
+        shaped: An object with a ``.shape`` attribute, or a shape tuple/sequence.
+
+    Returns:
+        A :class:`TensorSpec` with ``layout`` set to row-major minor-to-major order
+        and ``mode`` set to the identity permutation.
     """
     if hasattr(shaped, "shape"):
         shaped = shaped.shape
-    return TensorSpec(layout=row_major_layout(shaped), mode=default_tensor_mode(shaped))
+    inferred = tuple(d if isinstance(d, int) else None for d in shaped)
+    divisibility = inferred if any(d is not None for d in inferred) else None
+    return TensorSpec(
+        layout=row_major_layout(shaped),
+        mode=default_tensor_mode(shaped),
+        divisibility=divisibility,
+    )
+
+
+def _expand_divisibility(
+    divisibility, order: tuple[int, ...], ndim: int
+) -> tuple[int | None, ...] | None:
+    """Expand a divisibility spec to a full per-dimension tuple.
+
+    A bare ``int`` is placed at the leading-dimension slot (where
+    ``order[i] == 0``, i.e. stride == 1) and ``None`` everywhere else.
+    A tuple is returned unchanged.  ``None`` returns ``None``.
+    """
+    if divisibility is None or isinstance(divisibility, tuple):
+        return divisibility
+    leading = order.index(0)
+    result = [None] * ndim
+    result[leading] = divisibility
+    return tuple(result)
+
+
+def cutlass_to_jax_layout_order(
+    layout: Sequence[int] | None,
+) -> Sequence[int] | None:
+    """Converts a CuTeDSL layout order (minor-to-major) to JAX layout order (major-to-minor).
+
+    CuTeDSL uses minor-to-major ordering: ``layout[i]`` is the stride rank of
+    dimension ``i`` (0 = innermost).  JAX uses major-to-minor ordering: position
+    ``j`` in the result is the dimension index of the ``j``-th outermost axis.
+
+    Example::
+
+        cutlass_to_jax_layout_order((2, 1, 0))  # row-major → (0, 1, 2)
+        cutlass_to_jax_layout_order((0, 1, 2))  # col-major → (2, 1, 0)
+
+    Args:
+        layout: Minor-to-major stride permutation, or ``None`` (returned unchanged).
+
+    Returns:
+        Major-to-minor axis permutation compatible with ``jax.Array.layout``, or ``None``.
+    """
+    if layout is None:
+        return None
+    return tuple(sorted(range(len(layout)), key=lambda i: layout[i], reverse=True))
+
+
+def jax_to_cutlass_layout_order(
+    layout: Sequence[int] | None,
+) -> Sequence[int] | None:
+    """Converts a JAX layout order (major-to-minor) to CuTeDSL layout order (minor-to-major).
+
+    JAX uses major-to-minor ordering: position ``j`` is the dimension index of the
+    ``j``-th outermost axis.  CuTeDSL uses minor-to-major ordering: ``layout[i]``
+    is the stride rank of dimension ``i`` (0 = innermost).
+
+    This is the inverse of :func:`cutlass_to_jax_layout_order`.
+
+    Example::
+
+        jax_to_cutlass_layout_order((0, 1, 2))  # row-major → (2, 1, 0)
+        jax_to_cutlass_layout_order((2, 1, 0))  # col-major → (0, 1, 2)
+
+    Args:
+        layout: Major-to-minor axis permutation, or ``None`` (returned unchanged).
+
+    Returns:
+        Minor-to-major stride permutation for use as :attr:`TensorSpec.layout`, or ``None``.
+    """
+    if layout is None:
+        return None
+    inv = [0] * len(layout)
+    for i, p in enumerate(layout):
+        inv[p] = len(layout) - 1 - i
+    return tuple(inv)
 
 
 def jax_to_cutlass_dtype(dtype):
@@ -144,6 +268,16 @@ def from_dlpack(array, assumed_align: int = DEFAULT_CUTLASS_DEVICE_BUFFER_ALIGNM
     return _from_dlpack(array, assumed_align=assumed_align)
 
 
+def _validate_permutation(name: str, perm, shape):
+    if len(perm) != len(shape):
+        raise ValueError(f"{name} must be same length as shape", perm, shape)
+    for s in perm:
+        if s < 0 or s >= len(shape):
+            raise ValueError(f"Invalid index {s} in {name}", perm, shape)
+    if len(set(perm)) != len(perm):
+        raise ValueError(f"{name} has duplicate indices", perm)
+
+
 class JaxArray:
     """Base class for JaxArray argument type.
 
@@ -172,32 +306,21 @@ class JaxArray:
         order=None,
         mode=None,
         static=False,
+        divisibility=None,
     ):
         self.dtype = dtype
         self.shape = tuple(shape)
         self.ndim = len(self.shape)
         self.mem_space = mem_space
         self.assumed_align = assumed_align
+
         if order is None:
             order = row_major_layout(shape)
         if mode is None:
             mode = default_tensor_mode(shape)
 
-        if len(order) != len(shape):
-            raise ValueError(f"layout must be same length as shape", order, shape)
-        for s in order:
-            if s < 0 or s >= len(shape):
-                raise ValueError(f"Invalid index {s} in stride order", order, shape)
-        if len(tuple(set(order))) != len(order):
-            raise ValueError(f"layout has duplicate indices", order)
-
-        if len(mode) != len(shape):
-            raise ValueError(f"mode must be same length as shape", mode, shape)
-        for s in mode:
-            if s < 0 or s >= len(shape):
-                raise ValueError(f"Invalid index {s} in stride order", mode, shape)
-        if len(tuple(set(mode))) != len(mode):
-            raise ValueError(f"mode has duplicate indices", mode)
+        _validate_permutation("order", order, shape)
+        _validate_permutation("mode", mode, shape)
 
         self.order = tuple(order)
         self.mode = tuple(mode)
@@ -208,6 +331,20 @@ class JaxArray:
             )
         self.static = static
 
+        if divisibility is not None:
+            divisibility = _expand_divisibility(divisibility, self.order, self.ndim)
+            divisibility = tuple(divisibility)
+            if len(divisibility) != len(shape):
+                raise ValueError(
+                    "divisibility must be same length as shape", divisibility, shape
+                )
+            for d in divisibility:
+                if not (d is None or isinstance(d, int)):
+                    raise ValueError(
+                        f"divisibility entries must be None or integer, got {d!r}"
+                    )
+        self.divisibility = divisibility
+
 
 class JaxArrayValue(JaxArray):
     """The IR representation of the JaxArray."""
@@ -222,12 +359,15 @@ class JaxArrayValue(JaxArray):
         order,
         mode,
         static,
+        divisibility=None,
     ):
-        super().__init__(dtype, shape, mem_space, assumed_align, order, mode, static)
+        super().__init__(
+            dtype, shape, mem_space, assumed_align, order, mode, static, divisibility
+        )
         self.value = ir_value
 
     def __str__(self):
-        return f"JaxArrayValue<{self.value}:{self.dtype}:{self.shape}:{self.order}:{self.mode}:{self.static}>"
+        return f"JaxArrayValue<{self.value}:{self.dtype}:{self.shape}:{self.order}:{self.mode}:{self.static}:{self.divisibility}>"
 
     def __repr__(self):
         return str(self)
@@ -236,38 +376,44 @@ class JaxArrayValue(JaxArray):
         self, shape, order: tuple[int, ...], *, loc=None, ip=None
     ):
         i32 = ir.IntegerType.get_signless(32)
-        i64 = ir.IntegerType.get_signless(64)
-        one = arith.constant(i64, 1)
-        zero = arith.constant(i64, 0)
         pairs = sorted(zip(shape, order), key=lambda x: x[1])
 
         # Compute strides for each element in order.
         strides = [1]  # static 1 for leading
         if len(shape) > 1:
             strides.append(pairs[0][0])
-        for i, idx in enumerate(range(len(pairs[:-2]))):
-            strides.append(arith.muli(pairs[i][0], strides[-1]))
+        for i in range(len(pairs) - 2):
+            strides.append(arith.muli(pairs[i + 1][0], strides[-1]))
 
         # Apply the order to strides
         strides_ordered = []
         for i in range(len(shape)):
             strides_ordered.append(strides[order[i]])
 
-        # zero out any stride for a shape of size 1 to align with make_ordered_layout
-        # We ignore the leading dimension of 1
-        final_stride = []
-        for i in range(len(shape)):
-            x = arith.cmpi(0, one, shape[i])
-            s = strides_ordered[i]
-            if isinstance(s, int) and s == 1:
-                final_stride.append(s)
-            else:
-                final_stride.append(arith.select(x, zero, s))
-
         # Shapes are expected to be int32 so truncate to that before creating layout
-        shape = tuple([arith.trunci(i32, s) for s in shape])
+        shape_i32 = tuple(arith.trunci(i32, s) for s in shape)
 
-        return cute.make_layout(shape, stride=tuple(final_stride))
+        # Apply per-mode divisibility assumptions so the compiler can exploit alignment.
+        if self.divisibility is not None:
+            assumed = []
+            for s32, div_spec, static_s in zip(
+                shape_i32, self.divisibility, self.shape
+            ):
+                if isinstance(static_s, int):
+                    # Pure static shape is known even though a dynamic shape is
+                    # used. We can assume the exact shape here. We keep the shape
+                    # as a dynamic value to avoid breaking code that may expect
+                    # a dynamic value.
+                    assumed.append(cute.assume(s32, divby=static_s))
+                elif div_spec is not None:
+                    # Using a dynamic value so apply the div_spec if its provided.
+                    assumed.append(cute.assume(s32, divby=div_spec))
+                else:
+                    # No divisibility specification for this shape
+                    assumed.append(s32)
+            shape_i32 = tuple(assumed)
+
+        return cute.make_layout(shape_i32, stride=tuple(strides_ordered))
 
     def _load_dynamic_shapes(self, ffi_buffer, *, loc=None, ip=None):
         i64 = ir.IntegerType.get_signless(64)
@@ -325,7 +471,9 @@ class JaxArrayValue(JaxArray):
             layout = cute.make_ordered_layout(shape, order=self.order, loc=loc, ip=ip)
         else:
             shape = self._load_dynamic_shapes(ffi_buffer)
-            layout = self._make_ordered_layout_dynamic_strides(shape, self.order)
+            layout = self._make_ordered_layout_dynamic_strides(
+                shape, self.order, loc=loc, ip=ip
+            )
 
         # Apply mode order
         if self.mode is not None:
@@ -346,6 +494,7 @@ class JaxArrayValue(JaxArray):
             self.order,
             self.mode,
             self.static,
+            self.divisibility,
         )
 
 
@@ -355,20 +504,8 @@ class JaxTracedArray(JaxArray):
     Traced values are not real tensors or allocated on the device.
     """
 
-    def __init__(
-        self,
-        dtype,
-        shape,
-        mem_space,
-        assumed_align,
-        order,
-        mode,
-        static,
-    ):
-        super().__init__(dtype, shape, mem_space, assumed_align, order, mode, static)
-
     def __str__(self):
-        return f"JaxTracedArray<{self.dtype}:{self.shape}:{self.order}:{self.mode}:{self.static}>"
+        return f"JaxTracedArray<{self.dtype}:{self.shape}:{self.order}:{self.mode}:{self.static}:{self.divisibility}>"
 
     def __repr__(self):
         return str(self)
@@ -387,6 +524,7 @@ class JaxTracedArray(JaxArray):
             self.order,
             self.mode,
             self.static,
+            self.divisibility,
         )
 
     def __c_pointers__(self):

python/CuTeDSL/cutlass/pipeline/__init__.py

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA

python/CuTeDSL/cutlass/pipeline/helpers.py

@@ -3,7 +3,7 @@
 #
 # 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
+# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
 #
 # Any use, reproduction, disclosure, or distribution of this software
 # and related documentation outside the scope permitted by the EULA