nvmath-python 1.0.0__cp314-cp314t-win_amd64.whl

nvmath/device/cublasdx_numba.py

@@ -0,0 +1,1534 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import numbers
+import operator
+from collections.abc import Callable
+from functools import cached_property
+
+import llvmlite.ir as llvmir
+import numba
+import numpy
+from numba import cuda
+from numba.core import cgutils, typing
+from numba.core.base import BaseContext
+from numba.core.errors import TypingError
+from numba.core.extending import models
+from numba.cuda.cudaimpl import lower_constant
+from numba.cuda.extending import intrinsic, overload, overload_method, typeof_impl
+from numba.cuda.models import StructModel, register_model
+from numba.extending import types
+from numba.np import numpy_support
+
+from nvmath.device.common_cuda import get_default_code_type
+from nvmath.device.cublasdx_backend import generate_copy_wait_lto
+
+from .common import OpaqueTensor, axpby, clear, copy, copy_fragment, copy_wait, make_fragment_like, make_tensor
+from .common_numba import (
+    NUMBA_FE_TYPES_TO_NUMBA_IR,
+    EmptyStructModel,
+    MathdxOpaqueTensorType,
+    MathdxPipelineType,
+    cast_to_void_pointer,
+    declare_cabi_device,
+    get_array_ptr,
+    get_value_ptr,
+    overload_type_attribute,
+)
+from .common_opaque_tensor import (
+    LayoutModel,
+    LayoutType,
+    OpaqueTensorModel,
+    OpaqueTensorType,
+)
+from .cublasdx import (
+    DevicePipeline,
+    Matmul,
+    Partitioner,
+    TilePipeline,
+    _BlasMatmulLayout,
+    _BlasMatmulLikeLayout,
+    compile_blas_accumulator_init,
+    compile_blas_axpby,
+    compile_blas_clear,
+    compile_blas_copy,
+    compile_blas_device_pipeline_reset_tile,
+    compile_blas_execute,
+    compile_blas_is_index_in_bounds,
+    compile_blas_is_predicated,
+    compile_blas_is_thread_active,
+    compile_blas_map_idx2crd_partitioner,
+    compile_blas_tile_pipeline_destroy,
+    compile_blas_tile_pipeline_execute,
+    compile_blas_tile_pipeline_init,
+)
+from .cublasdx_backend import (
+    _BLAS_COMPILED_ARGS,
+    _BLAS_DEFINITION_ARGS,
+    _DEVICE_PIPELINE_COMPILED_ARGS,
+    _DEVICE_PIPELINE_DEFINITION_ARGS,
+    _TILE_PIPELINE_COMPILED_ARGS,
+    _TILE_PIPELINE_DEFINITION_ARGS,
+)
+
+
+class BlasType(types.Type):
+    """
+    Type class associated with the `cublasdx.Matmul`.
+    """
+
+    def __init__(self, blas: Matmul):
+        assert isinstance(blas, Matmul)
+        self._blas = blas
+        attributes = [f"{attr}={getattr(blas, attr)}" for attr in _BLAS_DEFINITION_ARGS if getattr(blas, attr)]
+        attributes.sort()
+
+        self.name = "BlasNumba(" + ",".join(attributes) + ")"
+
+    @property
+    def blas(self) -> Matmul:
+        return self._blas
+
+
+register_model(BlasType)(EmptyStructModel)
+
+
+@lower_constant(BlasType)
+def constant_dummy(context, builder, typ, pyval):
+    struct_ptr = cgutils.create_struct_proxy(typ)(context, builder)
+    return struct_ptr._getvalue()
+
+
+@typeof_impl.register(Matmul)
+def typeof_blas_numba(val: Matmul, c: typing.Context) -> BlasType:
+    return BlasType(val)
+
+
+for attribute in _BLAS_COMPILED_ARGS + _BLAS_DEFINITION_ARGS:
+    overload_type_attribute(BlasType, "blas", attribute)
+
+
+class DevicePipelineType(types.Type):
+    """
+    Type class associated with the `cublasdx.DevicePipeline`.
+    """
+
+    def __init__(self, pipeline: DevicePipeline):
+        assert isinstance(pipeline, DevicePipeline)
+        self._pipeline = pipeline
+        MM_type = BlasType(pipeline.mm)
+
+        attributes = [
+            f"a_layout=(dtype={pipeline._a_view.dtype},shape={pipeline._a_view.shape},strides={pipeline.a_strides})",
+            f"b_layout=(dtype={pipeline._b_view.dtype},shape={pipeline._b_view.shape},strides={pipeline.b_strides})",
+            f"pipeline_depth={pipeline.pipeline_depth}",
+            f"mm={MM_type}",
+        ]
+        attributes.sort()
+
+        self.name = "DevicePipeline(" + ",".join(attributes) + ")"
+
+    @property
+    def _buffer(self):
+        return self._pipeline._storage
+
+    @property
+    def pipeline(self) -> DevicePipeline:
+        return self._pipeline
+
+
+class PipelineModel(models.StructModel):
+    def __init__(self, dmm, fe_type: DevicePipelineType):
+        members = [
+            ("ptr", types.voidptr),
+        ]
+        models.StructModel.__init__(self, dmm, fe_type, members)
+
+
+register_model(DevicePipelineType)(PipelineModel)
+
+
+@typeof_impl.register(DevicePipeline)
+def typeof_device_pipeline_numba(val: DevicePipeline, c: typing.Context) -> DevicePipelineType:
+    return DevicePipelineType(val)
+
+
+for attribute in _DEVICE_PIPELINE_DEFINITION_ARGS + _DEVICE_PIPELINE_COMPILED_ARGS:
+    overload_type_attribute(DevicePipelineType, "pipeline", attribute)
+
+
+class TilePipelineType(types.Type):
+    """
+    Type class associated with the `cublasdx.DevicePipeline`.
+    """
+
+    def __init__(self, pipeline: TilePipeline):
+        assert isinstance(pipeline, TilePipeline)
+        self._pipeline = pipeline
+        device_pipeline_type = DevicePipelineType(pipeline.device_pipeline)
+        self.name = f"TilePipeline(device_pipeline={device_pipeline_type})"
+
+    @property
+    def pipeline(self) -> TilePipeline:
+        return self._pipeline
+
+
+register_model(TilePipelineType)(PipelineModel)
+
+
+@typeof_impl.register(TilePipeline)
+def typeof_tile_pipeline_numba(val: TilePipeline, c: typing.Context) -> TilePipelineType:
+    return TilePipelineType(val)
+
+
+@intrinsic
+def get_mathdx_tensor(typingctx: typing.Context, value):
+    """Cast value to MathdxOpaqueTensorType."""
+    if not isinstance(value, (OpaqueTensorType, BlasAccumulatorType)):
+        raise TypingError("get_mathdx_tensor expects an OpaqueTensorType or BlasAccumulatorType")
+
+    dynamic_strides_size = getattr(value.layout.layout, "dynamic_strides_size", 0)
+    dynamic_shape_size = getattr(value.layout.layout, "dynamic_shape_size", 0)
+
+    op_ty = MathdxOpaqueTensorType(dynamic_strides_size=dynamic_strides_size, dynamic_shape_size=dynamic_shape_size)
+    sig = typing.signature(op_ty, value)
+
+    def codegen(context: BaseContext, builder, sig, args):
+        input = cgutils.create_struct_proxy(sig.args[0])(context, builder, args[0])
+        output = cgutils.create_struct_proxy(sig.return_type)(context, builder)
+        output.ptr = input.ptr
+        if dynamic_strides_size > 0:
+            output.strides = input.strides
+        if dynamic_shape_size > 0:
+            output.shape = input.shape
+        return output._getvalue()
+
+    return sig, codegen
+
+
+@intrinsic
+def get_mathdx_pipeline(typingctx: typing.Context, value):
+    """Cast value to MathdxPipelineType."""
+    if not isinstance(value, (DevicePipelineType, TilePipelineType)):
+        raise TypingError("get_mathdx_pipeline expects a DevicePipelineType or TilePipelineType")
+
+    op_ty = MathdxPipelineType()
+    sig = typing.signature(op_ty, value)
+
+    def codegen(context: BaseContext, builder, sig, args):
+        input = cgutils.create_struct_proxy(sig.args[0])(context, builder, args[0])
+        output = cgutils.create_struct_proxy(sig.return_type)(context, builder)
+        output.ptr = input.ptr
+        return output._getvalue()
+
+    return sig, codegen
+
+
+for attribute in _TILE_PIPELINE_DEFINITION_ARGS + _TILE_PIPELINE_COMPILED_ARGS:
+    overload_type_attribute(TilePipelineType, "pipeline", attribute)
+
+
+@intrinsic
+def _create_tile_pipeline(typingctx: typing.Context, device_pipeline_type: DevicePipelineType):
+    """Create a new BLAS tile pipeline object."""
+
+    tile_pipeline = TilePipeline(device_pipeline_type.pipeline)
+    tile_pipeline_type = TilePipelineType(tile_pipeline)
+
+    return_ty = tile_pipeline_type
+    sig = typing.signature(return_ty, device_pipeline_type)
+
+    def codegen(context: BaseContext, builder, sig, args):
+        tile_pipeline_type = sig.return_type
+        struct_ptr = cgutils.create_struct_proxy(tile_pipeline_type)(context, builder)
+        return struct_ptr._getvalue()
+
+    return sig, codegen
+
+
+@intrinsic
+def _set_tile_pipeline_buffer(typingctx: typing.Context, tile_pipeline_type: TilePipelineType, buffer: types.Array):
+    """Create a new BLAS tile pipeline object."""
+
+    return_ty = tile_pipeline_type
+    sig = typing.signature(return_ty, tile_pipeline_type, buffer)
+
+    def codegen(context: BaseContext, builder, sig, args):
+        tile_pipeline_type, tile_pipeline = sig.args[0], args[0]
+        buffer_type, buffer = sig.args[1], args[1]
+        buffer_ptr = cgutils.create_struct_proxy(buffer_type)(context, builder, buffer).data
+        struct_ptr = cgutils.create_struct_proxy(tile_pipeline_type)(context, builder, tile_pipeline)
+        struct_ptr.ptr = buffer_ptr
+
+        return struct_ptr._getvalue()
+
+    return sig, codegen
+
+
+# Numba does not support method overload or variadic arguments, so we using
+# default values as a workaround
+# https://github.com/numba/numba/issues/9980
+# https://github.com/numba/numba/issues/9979
+# https://github.com/numba/numba/issues/10143
+@overload_method(BlasType, "execute", jit_options={"forceinline": True}, strict=False)
+def ol_blas_numba_execute(
+    blas_numba: BlasType, _arg1, _arg2, _arg3, _arg4=None, _arg5=None, _arg6=None, _arg7=None, _arg8=None
+):
+    none_set = {None, types.Omitted(None), types.none}
+
+    # execute(a, b, c) -> tensors API with rmem accumulator
+    if {_arg4, _arg5, _arg6, _arg7, _arg8} <= none_set:
+        return ol_blas_type_execute_tensors_rmem(blas_numba, _arg1, _arg2, _arg3)
+
+    # execute(alpha, a, b, beta, c) -> tensors API (smem) or plain arrays (static ld)
+    if {_arg6, _arg7, _arg8} <= none_set:
+        if isinstance(_arg2, OpaqueTensorType):
+            return ol_blas_type_execute_tensors_smem(blas_numba, _arg1, _arg2, _arg3, _arg4, _arg5)
+        return ol_blas_type_execute_basic(blas_numba, _arg1, _arg2, _arg3, _arg4, _arg5)
+
+    # execute(alpha, a, lda, b, ldb, beta, c, ldc) -> plain arrays (dynamic ld)
+    return ol_blas_type_execute_ldabc(blas_numba, _arg1, _arg2, _arg3, _arg4, _arg5, _arg6, _arg7, _arg8)
+
+
+def assert_suggested_tensors(tensor_types: tuple[types.Type, ...]):
+    """
+    Verify that all tensors are suggested or not suggested at the same time.
+    """
+    suggested_flags = []
+    for t in tensor_types:
+        assert isinstance(t, (OpaqueTensorType, BlasAccumulatorType))
+        layout_type = t.layout
+        assert isinstance(layout_type, BlasLayoutType)
+        suggested_flags.append(layout_type.layout.suggested)
+
+    if not all(suggested_flags) and any(suggested_flags):
+        raise TypingError("All tensors must be either suggested or not suggested at the same time.")
+
+
+def ol_blas_type_execute_tensors_rmem(
+    blas_numba: BlasType,
+    a: OpaqueTensorType,
+    b: OpaqueTensorType,
+    c: OpaqueTensorType,
+):
+    if not isinstance(blas_numba, BlasType):
+        return
+    MM = blas_numba.blas
+    if not all(isinstance(t, OpaqueTensorType) for t in (a, b)):
+        return
+    if not (isinstance(c, (OpaqueTensorType, BlasAccumulatorType))):
+        return
+
+    a_layout, b_layout, c_layout = a.layout.layout, b.layout.layout, c.layout.layout
+    assert isinstance(a_layout, _BlasMatmulLayout)
+    assert isinstance(b_layout, _BlasMatmulLayout)
+    assert isinstance(c_layout, _BlasMatmulLayout)
+
+    if c_layout.memory_space != "r" and not c_layout.accumulator:
+        return
+
+    assert_suggested_tensors((a, b, c))
+
+    return_type = types.void
+    mathdx_a = MathdxOpaqueTensorType(a_layout.dynamic_strides_size, a_layout.dynamic_shape_size)
+    mathdx_b = MathdxOpaqueTensorType(b_layout.dynamic_strides_size, b_layout.dynamic_shape_size)
+    mathdx_c = MathdxOpaqueTensorType(c_layout.dynamic_strides_size, c_layout.dynamic_shape_size)
+    sig = typing.signature(return_type, mathdx_a, mathdx_b, mathdx_c)
+
+    code, symbol = compile_blas_execute(
+        MM,
+        code_type=get_default_code_type(),
+        execute_api="tensors",
+        tensor_types=(a_layout.tensor_type, b_layout.tensor_type, c_layout.tensor_type),
+    )
+
+    lto = cuda.LTOIR(code.data)
+    blas_device_func = declare_cabi_device(symbol, sig, link=lto)
+
+    def impl(_, a, b, c, _arg4=None, _arg5=None, _arg6=None, _arg7=None, _arg8=None):
+        blas_device_func(get_mathdx_tensor(a), get_mathdx_tensor(b), get_mathdx_tensor(c))
+
+    return impl
+
+
+def ol_blas_type_execute_tensors_smem(
+    blas_numba: BlasType,
+    alpha: types.Number,
+    a: OpaqueTensorType,
+    b: OpaqueTensorType,
+    beta: types.Number,
+    c: OpaqueTensorType,
+):
+    if not isinstance(blas_numba, BlasType):
+        return
+    MM = blas_numba.blas
+    if not all(isinstance(a, types.Number) for a in (alpha, beta)):
+        return
+    if not all(isinstance(a, OpaqueTensorType) for a in (a, b, c)):
+        return
+    a_layout, b_layout, c_layout = a.layout.layout, b.layout.layout, c.layout.layout
+    assert isinstance(a_layout, _BlasMatmulLayout)
+    assert isinstance(b_layout, _BlasMatmulLayout)
+    assert isinstance(c_layout, _BlasMatmulLayout)
+    if c_layout.memory_space != "s":
+        return
+
+    return_type = types.void
+    mathdx_a = MathdxOpaqueTensorType(a_layout.dynamic_strides_size, a_layout.dynamic_shape_size)
+    mathdx_b = MathdxOpaqueTensorType(b_layout.dynamic_strides_size, b_layout.dynamic_shape_size)
+    mathdx_c = MathdxOpaqueTensorType(c_layout.dynamic_strides_size, c_layout.dynamic_shape_size)
+    c_ptr = types.CPointer(c.dtype)
+    sig = typing.signature(return_type, c_ptr, mathdx_a, mathdx_b, c_ptr, mathdx_c)
+
+    code, symbol = compile_blas_execute(
+        MM,
+        code_type=get_default_code_type(),
+        execute_api="tensors",
+        tensor_types=(a_layout.tensor_type, b_layout.tensor_type, c_layout.tensor_type),
+    )
+
+    lto = cuda.LTOIR(code.data)
+    blas_device_func = declare_cabi_device(symbol, sig, link=lto)
+
+    dtype = c.dtype
+
+    def impl(_, alpha, a, b, beta, c, _arg6=None, _arg7=None, _arg8=None):
+        alpha_ptr = get_value_ptr(dtype(alpha))
+        beta_ptr = get_value_ptr(dtype(beta))
+
+        blas_device_func(alpha_ptr, get_mathdx_tensor(a), get_mathdx_tensor(b), beta_ptr, get_mathdx_tensor(c))
+
+    return impl
+
+
+def ol_blas_type_execute_basic(
+    blas_numba: BlasType,
+    alpha: types.Number,
+    a: types.Array,
+    b: types.Array,
+    beta: types.Number,
+    c: types.Array,
+):
+    if not isinstance(blas_numba, BlasType):
+        return
+    MM = blas_numba.blas
+    if not all(isinstance(a, types.Number) for a in (alpha, beta)):
+        return
+    if not all(isinstance(a, types.Array) for a in (a, b, c)):
+        return
+    if (a.dtype, b.dtype, c.dtype) != tuple(NUMBA_FE_TYPES_TO_NUMBA_IR[vt] for vt in MM._traits.value_types):
+        return
+
+    # setting signature for intrinsic to much calling conventions. Numba will
+    # automatically cast to desired values.
+    return_type = types.void
+    c_ptr = types.CPointer(c.dtype)
+    sig = typing.signature(return_type, c_ptr, types.CPointer(a.dtype), types.CPointer(b.dtype), c_ptr, c_ptr)
+
+    code, symbol = compile_blas_execute(
+        MM,
+        code_type=get_default_code_type(),
+        execute_api="static_leading_dimensions",
+    )
+
+    lto = cuda.LTOIR(code.data)
+
+    blas_device_func = declare_cabi_device(symbol, sig, link=lto)
+
+    def impl(_, alpha, a, b, beta, c, _arg6=None, _arg7=None, _arg8=None):
+        aptr = get_array_ptr(a)
+        bptr = get_array_ptr(b)
+        cptr = get_array_ptr(c)
+        alpha = get_value_ptr(c.dtype.type(alpha))
+        beta = get_value_ptr(c.dtype.type(beta))
+
+        blas_device_func(alpha, aptr, bptr, beta, cptr)
+
+    return impl
+
+
+def ol_blas_type_execute_ldabc(
+    blas_numba: BlasType,
+    alpha: types.Number,
+    a: types.Array,
+    lda: types.Integer,
+    b: types.Array,
+    ldb: types.Integer,
+    beta: types.Number,
+    c: types.Array,
+    ldc: types.Integer,
+):
+    if not isinstance(blas_numba, BlasType):
+        return
+    MM = blas_numba.blas
+    if not all(isinstance(a, types.Number) for a in (alpha, beta)):
+        return
+    if not all(isinstance(a, types.Array) for a in (a, b, c)):
+        return
+    if (a.dtype, b.dtype, c.dtype) != tuple(NUMBA_FE_TYPES_TO_NUMBA_IR[vt] for vt in MM._traits.value_types):
+        return
+    if not all(isinstance(a, types.Integer) for a in (lda, ldb, ldc)):
+        return
+
+    ld_type = types.uint32
+    ld_ptr = types.CPointer(ld_type)
+    return_type = types.void
+    c_ptr = types.CPointer(c.dtype)
+    sig = typing.signature(
+        return_type, c_ptr, types.CPointer(a.dtype), ld_ptr, types.CPointer(b.dtype), ld_ptr, c_ptr, c_ptr, ld_ptr
+    )
+
+    code, symbol = compile_blas_execute(
+        MM,
+        code_type=get_default_code_type(),
+        execute_api="dynamic_leading_dimensions",
+    )
+
+    lto = cuda.LTOIR(code.data)
+    blas_device_func = declare_cabi_device(symbol, sig, link=lto)
+
+    def impl(_, alpha, a, lda, b, ldb, beta, c, ldc):
+        aptr = get_array_ptr(a)
+        bptr = get_array_ptr(b)
+        cptr = get_array_ptr(c)
+        alpha = get_value_ptr(c.dtype.type(alpha))
+        beta = get_value_ptr(c.dtype.type(beta))
+        lda = get_value_ptr(ld_type(lda))
+        ldb = get_value_ptr(ld_type(ldb))
+        ldc = get_value_ptr(ld_type(ldc))
+
+        blas_device_func(alpha, aptr, lda, bptr, ldb, beta, cptr, ldc)
+
+    return impl
+
+
+@overload(copy, jit_options={"forceinline": True}, strict=False)
+def ol_blas_copy(src: OpaqueTensorType, dst: OpaqueTensorType, alignment=None):
+    return ol_blas_copy_generic(src, dst, alignment, "copy")
+
+
+@overload(copy_fragment, jit_options={"forceinline": True}, strict=False)
+def ol_blas_copy_fragment(src: OpaqueTensorType, dst: OpaqueTensorType, alignment=None):
+    return ol_blas_copy_generic(src, dst, alignment, "copy_fragment")
+
+
+def ol_blas_copy_generic(src: OpaqueTensorType, dst: OpaqueTensorType, alignment_ty: types.Type | None, func: str):
+    assert isinstance(src, (OpaqueTensorType, BlasAccumulatorType))
+    assert isinstance(dst, OpaqueTensorType)
+    src_layout_ty = src.layout
+    dst_layout_ty = dst.layout
+    assert isinstance(src_layout_ty, BlasLayoutType)
+    assert isinstance(dst_layout_ty, BlasLayoutType)
+
+    alignment: int | None = None
+    if alignment_ty not in {None, types.Omitted(None)}:
+        if not isinstance(alignment_ty, types.Literal):
+            return lambda src, dst, alignment: numba.literally(alignment)
+        alignment = alignment_ty.literal_value
+        if alignment not in {1, 2, 4, 8, 16}:
+            raise TypingError(f"Alignment must be one of (1, 2, 4, 8, 16), got {alignment}")
+
+    rmem = dst_layout_ty.layout.memory_space == "r" or src_layout_ty.layout.memory_space == "r"
+
+    if func == "copy_fragment":
+        if not rmem:
+            raise TypingError("copy_fragment is only supported for rmem tensors. Please use copy instead.")
+    else:
+        assert func == "copy"
+        if rmem:
+            raise TypingError("copy is not supported for rmem tensors. Please use copy_fragment instead.")
+
+    if alignment is not None:
+        dtype = numpy_support.as_dtype(src.layout.dtype)
+        if alignment < dtype.itemsize:
+            raise TypingError(f"Alignment must be at least the size of the data type {dtype.itemsize}, got {alignment}")
+
+    code, symbol = compile_blas_copy(
+        src_tensor=src_layout_ty.layout,
+        dst_tensor=dst_layout_ty.layout,
+        code_type=get_default_code_type(),
+        alignment=alignment,
+    )
+
+    return_type = types.void
+    mathdx_src = MathdxOpaqueTensorType(src_layout_ty.layout.dynamic_strides_size, src_layout_ty.layout.dynamic_shape_size)
+    mathdx_dst = MathdxOpaqueTensorType(dst_layout_ty.layout.dynamic_strides_size, dst_layout_ty.layout.dynamic_shape_size)
+    sig = typing.signature(return_type, mathdx_src, mathdx_dst)
+
+    lto = cuda.LTOIR(code.data)
+    blas_device_func = declare_cabi_device(symbol, sig, link=lto)
+
+    def impl(src, dst, alignment=None):
+        return blas_device_func(get_mathdx_tensor(src), get_mathdx_tensor(dst))
+
+    return impl
+
+
+@overload(clear, jit_options={"forceinline": True}, strict=False)
+def ol_blas_clear(arr: OpaqueTensorType):
+    assert isinstance(arr, OpaqueTensorType)
+    assert isinstance(arr.layout, BlasLayoutType)
+
+    code, symbol = compile_blas_clear(
+        tensor=arr.layout.layout,
+        code_type=get_default_code_type(),
+    )
+
+    lto = cuda.LTOIR(code.data)
+
+    return_type = types.void
+    mathdx_arr = MathdxOpaqueTensorType(arr.layout.layout.dynamic_strides_size, arr.layout.layout.dynamic_shape_size)
+    sig = typing.signature(return_type, mathdx_arr)
+    blas_device_func = declare_cabi_device(symbol, sig, link=lto)
+
+    def impl(arr):
+        return blas_device_func(get_mathdx_tensor(arr))
+
+    return impl
+
+
+class BlasLayoutType(LayoutType):
+    """
+    Type class associated with opaque tensor layouts.
+    """
+
+    def __init__(self, layout: _BlasMatmulLayout):
+        assert isinstance(layout, _BlasMatmulLayout)
+        super().__init__(layout)
+        # Using handle descriptor in the type name to avoid symbol copy caching
+        # by numba.
+        MM_type = BlasType(layout.MM)
+        self.name = f"BlasLayout(uid={layout.uid},layout={layout.layout},dtype={layout.dtype},MM={MM_type})"
+
+    @property
+    def layout(self) -> _BlasMatmulLayout:
+        assert isinstance(self._layout, _BlasMatmulLayout)
+        return self._layout
+
+    def make_layout_like(self, dtype: numpy.number) -> "BlasLayoutType":
+        src_layout = self.layout
+        dst_layout = _BlasMatmulLikeLayout(
+            src_layout.MM,
+            src_layout.layout,
+            dtype,
+            leading_dimension=src_layout._default_ld,
+        )
+
+        return BlasLayoutType(dst_layout)
+
+
+register_model(BlasLayoutType)(LayoutModel)
+
+
+def lower_blas_layout_codegen(
+    context: BaseContext,
+    builder: llvmir.IRBuilder,
+    layout_type: BlasLayoutType,
+    ld_val: llvmir.NamedValue | numbers.Number | None,
+    ld_type: types.Type | None = None,
+):
+    layout = cgutils.create_struct_proxy(layout_type)(context, builder)
+    if layout_type.layout.dynamic_strides_size > 0:
+        assert layout_type.layout.dynamic_strides_size == 1
+        if ld_val is None:
+            # Use default MM ld when value is not provided
+            ld_val = layout_type.layout.MM.leading_dimension[layout_type.layout.tensor_index]
+        if isinstance(ld_val, numbers.Number):
+            ld = context.get_constant(types.int64, ld_val)
+        else:
+            ld = context.cast(builder, ld_val, ld_type, types.int64)
+        layout.strides = cgutils.pack_array(builder, [ld])
+    else:
+        assert ld_val is None
+    return layout._getvalue()
+
+
+@lower_constant(BlasLayoutType)
+def constant_blas_layout(context, builder, typ, pyval):
+    assert isinstance(typ, BlasLayoutType)
+    assert isinstance(pyval, _BlasMatmulLayout)
+    return lower_blas_layout_codegen(context, builder, typ, pyval.default_ld)
+
+
+@typeof_impl.register(_BlasMatmulLayout)
+def typeof_blas_mm_layout(layout: _BlasMatmulLayout, c: typing.Context) -> BlasLayoutType:
+    # We are clearing ld from the layout because it is a runtime value.
+    layout_no_ld = _BlasMatmulLayout(layout.MM, layout.layout)
+    return BlasLayoutType(layout_no_ld)
+
+
+@typeof_impl.register(_BlasMatmulLikeLayout)
+def typeof_blas_mm_like_layout(layout: _BlasMatmulLikeLayout, c: typing.Context) -> BlasLayoutType:
+    # We are clearing ld from the layout because it is a runtime value.
+    layout_no_ld = _BlasMatmulLikeLayout(layout.MM, layout.layout, layout.dtype)
+    return BlasLayoutType(layout_no_ld)
+
+
+for attribute in ["size", "cosize", "alignment"]:
+    overload_type_attribute(BlasLayoutType, "_layout", attribute)
+
+
+def ol_blas_layout(blas_numba: BlasType, method: str, leading_dimension: types.Number | None = None):
+    # leading_dimension is available only for global memory
+    if ("gmem" not in method) and (leading_dimension not in {None, types.Omitted(None)}):
+        return
+    MM = blas_numba.blas
+
+    layout = getattr(MM, method)()
+    return_type = BlasLayoutType(layout)
+
+    @intrinsic
+    def _intrinsic(typingctx, leading_dimension=None):
+        def codegen(context: BaseContext, builder: llvmir.IRBuilder, signature, args):
+            ld = args[0] if not isinstance(leading_dimension, types.NoneType) else None
+            ld_type = signature.args[0] if ld is not None else None
+            return lower_blas_layout_codegen(context, builder, return_type, ld, ld_type)
+
+        return typing.signature(return_type, leading_dimension), codegen
+
+    return lambda blas_numba, leading_dimension=None: _intrinsic(leading_dimension)
+
+
+def overload_blas_layout_method(method: str):
+    overload_method(
+        BlasType,
+        method,
+        jit_options={"forceinline": True},
+        strict=False,
+    )(lambda blas_numba, leading_dimension=None: ol_blas_layout(blas_numba, method, leading_dimension))
+
+
+for method in [
+    "get_layout_smem_a",
+    "get_layout_smem_b",
+    "get_layout_smem_c",
+    "get_layout_gmem_a",
+    "get_layout_gmem_b",
+    "get_layout_gmem_c",
+    "suggest_layout_smem_a",
+    "suggest_layout_smem_b",
+    "suggest_layout_smem_c",
+    "suggest_layout_rmem_c",
+    "get_layout_rmem_c",
+    "_get_accumulator_c",
+    "_suggest_accumulator_c",
+]:
+    overload_blas_layout_method(method)
+
+
+@overload(make_tensor, jit_options={"forceinline": True}, strict=False)
+def ol_make_tensor(array, layout):
+    assert isinstance(array, types.Array)
+    assert isinstance(layout, BlasLayoutType)
+
+    if array.dtype == layout.dtype:
+        return lambda array, layout: OpaqueTensor(array, layout)
+    else:
+
+        @intrinsic
+        def copy_strides(typingctx, array_type, layout_type):
+            def codegen(context: BaseContext, builder: llvmir.IRBuilder, signature, args):
+                array_type, layout_type = signature.args
+                array = cgutils.create_struct_proxy(array_type)(context, builder, args[0])
+                layout = cgutils.create_struct_proxy(layout_type)(context, builder, args[1])
+                if array_type.layout.layout.dynamic_strides_size > 0:
+                    array.strides = layout.strides
+                return array._getvalue()
+
+            return typing.signature(array_type, array_type, layout_type), codegen
+
+        np_dtype = numpy.dtype(numpy_support.as_dtype(array.dtype)).type
+        dst_layout_ty = layout.make_layout_like(np_dtype)
+        dst_layout = dst_layout_ty.layout
+
+        def impl(array, layout):
+            tensor = OpaqueTensor(array, dst_layout)
+            return copy_strides(tensor, layout)
+
+        return impl
+
+
+@overload(make_fragment_like, inline="always", strict=False)
+def ol_make_fragment_like(tensor, dtype):
+    assert isinstance(tensor, OpaqueTensorType)
+    layout = tensor.layout
+    assert isinstance(layout, BlasLayoutType)
+    assert layout.layout.memory_space == "r"
+
+    assert isinstance(dtype, types.NumberClass)
+    np_dtype = numpy.dtype(numpy_support.as_dtype(dtype)).type
+
+    dst_layout_ty = layout.make_layout_like(np_dtype)
+    dst_layout = dst_layout_ty.layout
+
+    def impl(tensor, dtype):
+        buff = cuda.local.array(dst_layout.cosize, dtype=dtype, alignment=dst_layout.alignment)
+        return OpaqueTensor(buff, dst_layout)
+
+    return impl
+
+
+@overload(copy_wait, jit_options={"forceinline": True}, strict=False)
+def ol_copy_wait():
+    # numba has cache per compute capability, so the function won't end up
+    # cached for the wrong compute capability.
+    return_type = types.void
+    sig = typing.signature(return_type)
+
+    ct = get_default_code_type()
+    symbol, code = generate_copy_wait_lto(ct.cc)
+
+    lto = cuda.LTOIR(code)
+    blas_device_func = declare_cabi_device(symbol, sig, link=lto)
+
+    return lambda: blas_device_func()
+
+
+@overload(axpby, jit_options={"forceinline": True}, strict=False)
+def ol_axpby(a, x, b, y):
+    if not isinstance(a, types.Number):
+        return
+    if not isinstance(x, OpaqueTensorType):
+        return
+    if not isinstance(b, types.Number):
+        return
+    if not isinstance(y, OpaqueTensorType):
+        return
+    if x.layout.layout.memory_space != "r":
+        raise TypeError("axpby is only supported for rmem tensors. x is not")
+    if y.layout.layout.memory_space != "r":
+        raise TypeError("axpby is only supported for rmem tensors. y is not")
+
+    code, symbol = compile_blas_axpby(
+        x_tensor=x.layout.layout,
+        y_tensor=y.layout.layout,
+        code_type=get_default_code_type(),
+    )
+
+    lto = cuda.LTOIR(code.data)
+
+    return_type = types.void
+    alpha_beta_ptr = types.CPointer(types.uint8)
+    mathdx_x = MathdxOpaqueTensorType(x.layout.layout.dynamic_strides_size, x.layout.layout.dynamic_shape_size)
+    mathdx_y = MathdxOpaqueTensorType(y.layout.layout.dynamic_strides_size, y.layout.layout.dynamic_shape_size)
+    sig = typing.signature(return_type, alpha_beta_ptr, mathdx_x, alpha_beta_ptr, mathdx_y)
+    blas_device_func = declare_cabi_device(symbol, sig, link=lto)
+
+    x_dtype = x.dtype
+    y_dtype = y.dtype
+
+    def impl(a, x, b, y):
+        a_ptr = cast_to_void_pointer(get_value_ptr(x_dtype(a)))
+        b_ptr = cast_to_void_pointer(get_value_ptr(y_dtype(b)))
+        return blas_device_func(a_ptr, get_mathdx_tensor(x), b_ptr, get_mathdx_tensor(y))
+
+    return impl
+
+
+class BlasPartitionerType(types.Type):
+    """
+    Type class for Blas partitioner.
+    """
+
+    def __init__(self, MM: Matmul, suggested: bool):
+        assert isinstance(MM, Matmul)
+        self._MM = MM
+        mm_type = BlasType(MM)
+        self._suggested = suggested
+        super().__init__(f"BlasPartitioner(MM={mm_type}, suggested={suggested})")
+
+    @property
+    def MM(self) -> Matmul:
+        return self._MM
+
+    @property
+    def suggested(self) -> bool:
+        return self._suggested
+
+    @property
+    def dtype(self) -> types.Number:
+        return NUMBA_FE_TYPES_TO_NUMBA_IR[self._MM._traits.value_types[2]]
+
+    @property
+    def fragment_layout(self) -> BlasLayoutType:
+        MM = self.MM
+        if self.suggested:
+            layout = MM.suggest_layout_rmem_c()
+        else:
+            layout = MM.get_layout_rmem_c()
+
+        assert isinstance(layout, _BlasMatmulLayout)
+
+        return BlasLayoutType(layout)
+
+
+@register_model(BlasPartitionerType)
+class PartitionerModel(StructModel):
+    def __init__(self, dmm, fe_type: BlasPartitionerType):
+        StructModel.__init__(self, dmm, fe_type, [])
+
+
+class BlasAccumulatorType(BlasPartitionerType):
+    """
+    Type class for Blas partitioner.
+    """
+
+    def __init__(self, MM: Matmul, suggested: bool):
+        super().__init__(MM, suggested)
+        mm = BlasType(MM)
+        self.name = f"BlasAccumulator(MM={mm}, suggested={suggested})"
+
+    @cached_property
+    def layout(self) -> BlasLayoutType:
+        if self._suggested:
+            layout = self._MM._suggest_accumulator_c()
+        else:
+            layout = self._MM._get_accumulator_c()
+
+        assert isinstance(layout, _BlasMatmulLayout)
+
+        return BlasLayoutType(layout)
+
+
+register_model(BlasAccumulatorType)(OpaqueTensorModel)
+
+
+def ol_blas_get_accumulator_generic(blas_numba: BlasType, suggested: bool):
+    assert isinstance(blas_numba, BlasType)
+
+    MM = blas_numba.blas
+    return_type = BlasAccumulatorType(MM, suggested=suggested)
+    opaque_tensor_type = OpaqueTensorType(return_type.layout)
+
+    @intrinsic
+    def _type_cast_ot_acc(typingctx, opaque_tensor):
+        def codegen(context: BaseContext, builder: llvmir.IRBuilder, signature, args):
+            return args[0]
+
+        return typing.signature(return_type, opaque_tensor_type), codegen
+
+    if suggested:
+        acc_layout = MM._suggest_accumulator_c()
+    else:
+        acc_layout = MM._get_accumulator_c()
+
+    alignment = max(acc_layout.alignment, 8)
+
+    def impl(mm: Matmul):
+        buff = cuda.local.array(acc_layout.cosize, dtype=mm.c_value_type, alignment=alignment)
+        tensor = make_tensor(buff, acc_layout)
+        clear(tensor)
+        acc = _type_cast_ot_acc(tensor)
+        acc._init()
+        return acc
+
+    return impl
+
+
+@overload_method(BlasType, "suggest_accumulator", inline="always", strict=False)
+def ol_blas_suggest_accumulator(blas_numba: BlasType):
+    return ol_blas_get_accumulator_generic(blas_numba, suggested=True)
+
+
+@overload_method(BlasType, "get_accumulator", inline="always", strict=False)
+def ol_blas_get_accumulator(blas_numba: BlasType):
+    return ol_blas_get_accumulator_generic(blas_numba, suggested=False)
+
+
+@overload_method(BlasPartitionerType, "partition_like_C", jit_options={"forceinline": True}, strict=False)
+def ol_blas_partition_like_C(partitioner: BlasPartitionerType, tensor: OpaqueTensorType):
+    assert isinstance(partitioner, BlasPartitionerType)
+    assert isinstance(tensor, OpaqueTensorType)
+    assert isinstance(tensor.layout, BlasLayoutType)
+    assert tensor.layout.layout.tensor_type == "gmem_c"
+
+    raise NotImplementedError("Blas partition_like_C is not yet implemented, please use map_fragment_index instead")
+
+
+def get_map_idx2crd_partitioner(symbol: str, lto: cuda.LTOIR):
+    assert isinstance(symbol, str)
+    return_type = types.Tuple((types.int32, types.int32))
+
+    @intrinsic
+    def map_idx2crd_partitioner(typingctx, index):
+        def codegen(context: BaseContext, builder: llvmir.IRBuilder, signature, args):
+            context.active_code_library.add_linking_file(lto)
+            idx: llvmir.AllocaInstr = args[0]
+            i_ptr = cgutils.alloca_once(builder, cgutils.int32_t)
+            j_ptr = cgutils.alloca_once(builder, cgutils.int32_t)
+            idx_ptr = cgutils.alloca_once_value(builder, idx)
+
+            int32_ptr = cgutils.int32_t.as_pointer()
+
+            fnTy = llvmir.FunctionType(cgutils.voidptr_t, [int32_ptr, int32_ptr, int32_ptr])
+            fn = cgutils.get_or_insert_function(builder.module, fnTy, symbol)
+            builder.call(fn, [idx_ptr, i_ptr, j_ptr])
+
+            return context.make_tuple(builder, return_type, (builder.load(i_ptr), builder.load(j_ptr)))
+
+        return typing.signature(return_type, types.int32), codegen
+
+    return map_idx2crd_partitioner
+
+
+@overload_method(
+    BlasPartitionerType,
+    "map_fragment_index",
+    jit_options={"forceinline": True},
+    strict=False,
+)
+def ol_blas_partitioner_map_fragment_index(
+    partitioner: BlasPartitionerType,
+    index: types.Integer,
+):
+    if not isinstance(partitioner, BlasPartitionerType):
+        return
+    if not isinstance(index, types.Integer):
+        return
+
+    code, symbol = compile_blas_map_idx2crd_partitioner(partitioner.MM, code_type=get_default_code_type())
+
+    lto = cuda.LTOIR(code.data)
+
+    map_idx2crd_partitioner = get_map_idx2crd_partitioner(symbol, lto=lto)
+
+    def map_fragment_index_impl(obj, idx):
+        i, j = map_idx2crd_partitioner(idx)
+        return (i, j)
+
+    return map_fragment_index_impl
+
+
+def get_bool_return_intrinsic(symbol: str, index: bool = False, lto=None):
+    return_type = types.bool
+
+    def codegen(context: BaseContext, builder: llvmir.IRBuilder, signature, args: list):
+        if lto is not None:
+            context.active_code_library.add_linking_file(lto)
+        active = cgutils.alloca_once(builder, cgutils.int32_t)
+        int32_ptr = cgutils.int32_t.as_pointer()
+
+        fn_args, fn_args_ty = [], []
+        if index:
+            idx_ptr = cgutils.alloca_once_value(builder, args[0])
+            fn_args += [idx_ptr]
+            fn_args_ty += [int32_ptr]
+        fn_args += [active]
+        fn_args_ty += [int32_ptr]
+
+        fnTy = llvmir.FunctionType(cgutils.voidptr_t, fn_args_ty)
+        fn = cgutils.get_or_insert_function(builder.module, fnTy, symbol)
+        builder.call(fn, fn_args)
+
+        res = builder.icmp_signed("!=", builder.load(active), llvmir.Constant(cgutils.int32_t, 0))
+        return res
+
+    _intrinsic: Callable = lambda typingctx: (typing.signature(return_type), codegen)
+    if index:
+        _intrinsic = lambda typingctx, index: (typing.signature(return_type, types.int32), codegen)
+
+    return intrinsic(_intrinsic)
+
+
+@overload_method(
+    BlasPartitionerType,
+    "is_thread_active",
+    jit_options={"forceinline": True},
+    strict=False,
+)
+def ol_blas_partitioner_is_thread_active(
+    partitioner: BlasPartitionerType,
+):
+    if not isinstance(partitioner, BlasPartitionerType):
+        return
+
+    code, symbol = compile_blas_is_thread_active(partitioner.MM, code_type=get_default_code_type())
+    lto = cuda.LTOIR(code.data)
+    is_thread_active = get_bool_return_intrinsic(symbol, lto=lto)
+
+    return lambda partitioner: is_thread_active()
+
+
+@overload_method(
+    BlasPartitionerType,
+    "is_predicated",
+    jit_options={"forceinline": True},
+    strict=False,
+)
+def ol_blas_partitioner_is_predicated(
+    partitioner: BlasPartitionerType,
+):
+    if not isinstance(partitioner, BlasPartitionerType):
+        return
+
+    code, symbol = compile_blas_is_predicated(partitioner.MM, code_type=get_default_code_type())
+    lto = cuda.LTOIR(code.data)
+    is_predicated = get_bool_return_intrinsic(symbol, lto=lto)
+
+    return lambda partitioner: is_predicated()
+
+
+@overload_method(
+    BlasPartitionerType,
+    "is_index_in_bounds",
+    jit_options={"forceinline": True},
+    strict=False,
+)
+def ol_blas_partition_is_index_in_bounds(
+    partitioner: BlasPartitionerType,
+    index: types.Integer,
+):
+    if not isinstance(partitioner, BlasPartitionerType):
+        return
+    if not isinstance(index, types.Integer):
+        return
+
+    code, symbol = compile_blas_is_index_in_bounds(partitioner.MM, code_type=get_default_code_type())
+    lto = cuda.LTOIR(code.data)
+    is_index_in_bounds = get_bool_return_intrinsic(symbol, index=True, lto=lto)
+
+    return lambda partitioner, idx: is_index_in_bounds(idx)
+
+
+@overload_method(
+    BlasPartitionerType,
+    "partition_and_copy",
+    inline="always",
+    strict=False,
+)
+def ol_blas_partitioner_partition_and_copy(
+    partitioner: BlasPartitionerType,
+    src: OpaqueTensorType,
+    dst: OpaqueTensorType,
+):
+    if not isinstance(partitioner, BlasPartitionerType):
+        return
+    if not isinstance(src, OpaqueTensorType):
+        return
+    if not isinstance(dst, OpaqueTensorType):
+        return
+    assert isinstance(src.layout, BlasLayoutType)
+    assert isinstance(dst.layout, BlasLayoutType)
+    src_mem = src.layout.layout.memory_space
+    dst_mem = dst.layout.layout.memory_space
+    assert src_mem == "r" and dst_mem == "g" or src_mem == "g" and dst_mem == "r"
+
+    alignment = max(partitioner.MM.alignment.c, src.layout.layout.alignment, dst.layout.layout.alignment)
+
+    def impl(partitioner: Partitioner, src, dst):
+        if partitioner.is_thread_active():
+            copy_fragment(src, dst, alignment=alignment)
+
+    return impl
+
+
+@overload_method(
+    BlasPartitionerType,
+    "make_empty_fragment",
+    inline="always",
+    strict=False,
+)
+def ol_blas_partitioner_make_empty_fragment(
+    partitioner: BlasPartitionerType,
+):
+    MM = partitioner.MM
+    layout = partitioner.fragment_layout.layout
+
+    alignment = max(8, layout.alignment, MM.alignment.c)
+
+    def impl(partitioner):
+        buff = cuda.local.array(layout.cosize, dtype=MM.c_value_type, alignment=alignment)
+        frag = make_tensor(buff, layout)
+        return frag
+
+    return impl
+
+
+@overload_method(
+    BlasPartitionerType,
+    "make_partition_and_copy",
+    inline="always",
+    strict=False,
+)
+def ol_blas_partitioner_make_partition_and_copy(
+    partitioner: BlasPartitionerType,
+    tensor: OpaqueTensorType,
+):
+    if not isinstance(tensor, OpaqueTensorType):
+        return
+
+    layout_ty = partitioner.fragment_layout
+    layout = layout_ty.layout
+    if layout_ty.dtype != tensor.layout.dtype:
+        layout = layout_ty.make_layout_like(tensor.layout.layout.dtype).layout
+    assert isinstance(layout, _BlasMatmulLayout)
+
+    MM = partitioner.MM
+    alignment = max(8, layout.alignment, MM.alignment.c)
+    dtype = layout.dtype
+
+    def impl(partitioner, tensor):
+        buff = cuda.local.array(layout.cosize, dtype=dtype, alignment=alignment)
+        frag = make_tensor(buff, layout)
+        partitioner.partition_and_copy(tensor, frag)
+        return frag
+
+    return impl
+
+
+@overload_method(
+    BlasAccumulatorType,
+    "_init",
+    target="cuda",
+    jit_options={"forceinline": True},
+    strict=False,
+)
+def ol_opaque_tensor_init(
+    accumulator: BlasAccumulatorType,
+):
+    if not isinstance(accumulator, BlasAccumulatorType):
+        return
+
+    layout = accumulator.layout
+    assert isinstance(layout, BlasLayoutType)
+    code, symbol = compile_blas_accumulator_init(layout.layout, code_type=get_default_code_type())
+    lto = cuda.LTOIR(code.data)
+
+    return_type = types.void
+    mathdx_accumulator = MathdxOpaqueTensorType(layout.layout.dynamic_strides_size, layout.layout.dynamic_shape_size)
+    sig = typing.signature(return_type, mathdx_accumulator)
+
+    cublasdx_init_accumulator = declare_cabi_device(symbol, sig, link=lto)
+
+    def impl(accumulator):
+        cublasdx_init_accumulator(get_mathdx_tensor(accumulator))
+
+    return impl
+
+
+@overload_method(
+    BlasAccumulatorType,
+    "get_results",
+    inline="always",
+    strict=False,
+)
+def ol_blas_accumulator_get_results(
+    accumulator: BlasAccumulatorType,
+):
+    if not isinstance(accumulator, BlasAccumulatorType):
+        return
+
+    def impl(accumulator):
+        frag = accumulator.make_empty_fragment()
+        copy_fragment(accumulator, frag)
+
+        return frag
+
+    return impl
+
+
+@overload_method(BlasAccumulatorType, "clear", jit_options={"forceinline": True}, strict=False)
+def ol_blas_accumulator_clear(accumulator: BlasAccumulatorType):
+    opaque_tensor_type = OpaqueTensorType(accumulator.layout)
+
+    @intrinsic
+    def _type_cast_ot_tensor(typingctx, opaque_tensor):
+        def codegen(context: BaseContext, builder: llvmir.IRBuilder, signature, args):
+            return args[0]
+
+        return_type = opaque_tensor_type
+        return typing.signature(return_type, accumulator), codegen
+
+    def impl(accumulator: BlasAccumulatorType):
+        tensor = _type_cast_ot_tensor(accumulator)
+        clear(tensor)
+
+    return impl
+
+
+@overload_method(DevicePipelineType, "get_tile", target="cuda", inline="always", strict=False)
+def ol_device_pipeline_get_tile(
+    device_pipeline: DevicePipelineType,
+    smem: types.Array,
+    idx: types.Number,
+    idy: types.Number,
+):
+    if not isinstance(device_pipeline, DevicePipelineType):
+        return
+    if not isinstance(smem, types.Array):
+        return
+    if not isinstance(idx, types.Number) and not isinstance(idx, types.UniTuple):
+        return
+    if not isinstance(idy, types.Number) and not isinstance(idy, types.UniTuple):
+        return
+
+    assert (
+        isinstance(idx, types.Number)
+        and isinstance(idy, types.Number)
+        or isinstance(idx, types.UniTuple)
+        and isinstance(idy, types.UniTuple)
+    )
+
+    def impl(device_pipeline, smem, idx, idy):
+        tile_pipeline: TilePipeline = _create_tile_pipeline(device_pipeline)
+        tile_pipeline_buffer = cuda.local.array(
+            shape=(tile_pipeline.storage_bytes,), dtype=numpy.byte, alignment=tile_pipeline.storage_alignment
+        )
+        tile_pipeline = _set_tile_pipeline_buffer(tile_pipeline, tile_pipeline_buffer)
+        tile_pipeline._init(device_pipeline, smem, idx, idy)
+
+        return tile_pipeline
+
+    return impl
+
+
+@overload_method(DevicePipelineType, "reset_tile", target="cuda", inline="always", strict=False)
+def ol_tile_pipeline_reset_tile(
+    device_pipeline: DevicePipelineType,
+    tile_pipeline: TilePipelineType,
+    idx: types.Number,
+    idy: types.Number,
+):
+    if not isinstance(device_pipeline, DevicePipelineType):
+        return
+    if not isinstance(tile_pipeline, TilePipelineType):
+        return
+    if not isinstance(idx, types.Number) and not isinstance(idx, types.UniTuple):
+        return
+    if not isinstance(idy, types.Number) and not isinstance(idy, types.UniTuple):
+        return
+
+    assert tile_pipeline.pipeline.device_pipeline == device_pipeline.pipeline
+
+    code, symbol = compile_blas_device_pipeline_reset_tile(
+        device_pipeline.pipeline, tile_pipeline.pipeline, code_type=get_default_code_type()
+    )
+    lto = cuda.LTOIR(code.data)
+
+    return_type = types.void
+    sig = typing.signature(
+        return_type,
+        MathdxPipelineType(),
+        MathdxPipelineType(),
+        types.CPointer(types.int32),
+        types.CPointer(types.int32),
+    )
+
+    cublasdx_reset_tile = declare_cabi_device(symbol, sig, link=lto)
+
+    if isinstance(idx, types.Number):
+
+        def impl(device_pipeline, tile_pipeline, idx, idy):
+            idx = get_value_ptr(types.int32(idx))
+            idy = get_value_ptr(types.int32(idy))
+            cublasdx_reset_tile(get_mathdx_pipeline(device_pipeline), get_mathdx_pipeline(tile_pipeline), idx, idy)
+
+        return impl
+    else:
+
+        def impl(device_pipeline, tile_pipeline, idx, idy):
+            idx_arr = cuda.local.array(shape=(2,), dtype=numpy.int32)
+            idx_arr[0], idx_arr[1] = idx
+            idy_arr = cuda.local.array(shape=(2,), dtype=numpy.int32)
+            idy_arr[0], idy_arr[1] = idy
+            idx_ptr = get_array_ptr(idx_arr)
+            idy_ptr = get_array_ptr(idy_arr)
+            cublasdx_reset_tile(get_mathdx_pipeline(device_pipeline), get_mathdx_pipeline(tile_pipeline), idx_ptr, idy_ptr)
+
+        return impl
+
+
+@overload_method(TilePipelineType, "_init", target="cuda", jit_options={"forceinline": True}, strict=False)
+def ol_tile_pipeline_init(
+    tile_pipeline: TilePipelineType,
+    device_pipeline: DevicePipelineType,
+    smem: types.Array,
+    idx: types.Number,
+    idy: types.Number,
+):
+    if not isinstance(tile_pipeline, TilePipelineType):
+        return
+    if not isinstance(device_pipeline, DevicePipelineType):
+        return
+    if not isinstance(smem, types.Array):
+        return
+    if not isinstance(idx, types.Number) and not isinstance(idx, types.UniTuple):
+        return
+    if not isinstance(idy, types.Number) and not isinstance(idy, types.UniTuple):
+        return
+
+    assert tile_pipeline.pipeline.device_pipeline == device_pipeline.pipeline
+
+    code, symbol = compile_blas_tile_pipeline_init(tile_pipeline.pipeline, code_type=get_default_code_type())
+    lto = cuda.LTOIR(code.data)
+
+    return_type = types.void
+    sig = typing.signature(
+        return_type,
+        MathdxPipelineType(),
+        MathdxPipelineType(),
+        types.CPointer(smem.dtype),
+        types.CPointer(types.int32),
+        types.CPointer(types.int32),
+    )
+
+    cublasdx_init_pipeline = declare_cabi_device(symbol, sig, link=lto)
+
+    if isinstance(idx, types.Number):
+
+        def impl(tile_pipeline, device_pipeline, smem, idx, idy):
+            smem_ptr = get_array_ptr(smem)
+            idx = get_value_ptr(types.int32(idx))
+            idy = get_value_ptr(types.int32(idy))
+
+            cublasdx_init_pipeline(get_mathdx_pipeline(device_pipeline), get_mathdx_pipeline(tile_pipeline), smem_ptr, idx, idy)
+
+        return impl
+    else:
+
+        def impl(tile_pipeline, device_pipeline, smem, idx, idy):
+            smem_ptr = get_array_ptr(smem)
+            idx_arr = cuda.local.array(shape=(2,), dtype=numpy.int32)
+            idx_arr[0], idx_arr[1] = idx
+            idy_arr = cuda.local.array(shape=(2,), dtype=numpy.int32)
+            idy_arr[0], idy_arr[1] = idy
+            idx_ptr = get_array_ptr(idx_arr)
+            idy_ptr = get_array_ptr(idy_arr)
+            cublasdx_init_pipeline(
+                get_mathdx_pipeline(device_pipeline), get_mathdx_pipeline(tile_pipeline), smem_ptr, idx_ptr, idy_ptr
+            )
+
+        return impl
+
+
+@overload_method(TilePipelineType, "_del", target="cuda", jit_options={"forceinline": True}, strict=False)
+def ol_tile_pipeline_destroy(tile_pipeline: TilePipelineType):
+    if not isinstance(tile_pipeline, TilePipelineType):
+        return
+
+    code, symbol = compile_blas_tile_pipeline_destroy(tile_pipeline.pipeline, code_type=get_default_code_type())
+    lto = cuda.LTOIR(code.data)
+
+    return_type = types.void
+    sig = typing.signature(return_type, tile_pipeline)
+
+    cublasdx_destroy_pipeline = declare_cabi_device(symbol, sig, link=lto)
+
+    def impl(tile_pipeline):
+        cublasdx_destroy_pipeline(tile_pipeline)
+
+    return impl
+
+
+@overload_method(TilePipelineType, "execute", target="cuda", jit_options={"forceinline": True}, strict=False)
+def ol_tile_pipeline_execute(
+    tile_pipeline: TilePipelineType,
+    accumulator: BlasAccumulatorType,
+):
+    assert isinstance(accumulator, BlasAccumulatorType)
+    acc_layout = accumulator.layout.layout
+    assert acc_layout.accumulator
+
+    code, symbol = compile_blas_tile_pipeline_execute(tile_pipeline.pipeline, acc_layout, code_type=get_default_code_type())
+    lto = cuda.LTOIR(code.data)
+
+    return_type = types.void
+    acc_ty = MathdxOpaqueTensorType(
+        dynamic_strides_size=getattr(getattr(acc_layout, "layout", None), "dynamic_strides_size", 0),
+        dynamic_shape_size=getattr(getattr(acc_layout, "layout", None), "dynamic_shape_size", 0),
+    )
+    sig = typing.signature(return_type, MathdxPipelineType(), acc_ty)
+
+    cublasdx_tile_pipeline_execute = declare_cabi_device(symbol, sig, link=lto)
+
+    def impl(tile_pipeline, accumulator):
+        cublasdx_tile_pipeline_execute(get_mathdx_pipeline(tile_pipeline), get_mathdx_tensor(accumulator))
+
+    return impl
+
+
+@intrinsic
+def _get_set_ot_item(typingctx, tensor: OpaqueTensorType, index: types.Integer, val=None):
+    def codegen(context: BaseContext, builder: llvmir.IRBuilder, signature, args):
+        tensor_ty = signature.args[0]
+        tensor, idx = args[0], args[1]
+
+        tensor_struct = cgutils.create_struct_proxy(tensor_ty)(context, builder, tensor)
+        buff_void_ptr = tensor_struct.ptr
+
+        assert isinstance(buff_void_ptr.type, llvmir.PointerType)
+
+        llvm_val_ty = context.get_value_type(tensor_ty.dtype)
+
+        data_ptr_ty = llvmir.PointerType(llvm_val_ty, buff_void_ptr.type.addrspace)
+        buff_ptr = builder.bitcast(buff_void_ptr, data_ptr_ty)
+
+        val_ptr = builder.gep(buff_ptr, [idx], inbounds=True)
+
+        if val is None:
+            return builder.load(val_ptr)
+        else:
+            val_arg = args[2]
+            builder.store(val_arg, val_ptr)
+
+    if val is None:
+        ret_type = tensor.dtype
+        sig = typing.signature(ret_type, tensor, index, types.none)
+    else:
+        ret_type = types.void
+        sig = typing.signature(ret_type, tensor, index, tensor.dtype)
+    return sig, codegen
+
+
+def _validate_fragment_tensor(tensor: OpaqueTensorType) -> bool:
+    if not isinstance(tensor, OpaqueTensorType):
+        return False
+
+    layout_ty = tensor.layout
+    if not isinstance(layout_ty, BlasLayoutType):
+        return False
+
+    layout = layout_ty.layout
+    if not isinstance(layout, _BlasMatmulLayout):
+        return False
+    return layout.memory_space == "r"
+
+
+@overload(operator.getitem, jit_options={"forceinline": True}, strict=False)
+def ol_blas_partition_getitem(tensor: OpaqueTensorType, index: types.Integer):
+    if not _validate_fragment_tensor(tensor):
+        return
+    if not isinstance(index, types.Integer):
+        return
+
+    return lambda tensor, index: _get_set_ot_item(tensor, index)
+
+
+@overload(operator.setitem, jit_options={"forceinline": True}, strict=False)
+def ol_blas_partition_setitem(tensor: OpaqueTensorType, index: types.Integer, value):
+    if not _validate_fragment_tensor(tensor):
+        return
+    if not isinstance(index, types.Integer):
+        return
+    if not isinstance(value, types.Number):
+        return
+
+    dtype = tensor.dtype
+
+    return lambda tensor, index, value: _get_set_ot_item(tensor, index, dtype(value))
+
+
+class _PipelineExtension:
+    def prepare_args(self, ty, val, **kwargs):
+        if isinstance(val, DevicePipeline):
+            assert isinstance(ty, DevicePipelineType)
+            c_ptr = types.CPointer(types.void)
+            return c_ptr, int(val._storage.handle)
+        else:
+            return ty, val
+
+
+# TODO: make implicit, once numba-cuda supports it
+#  https://github.com/NVIDIA/numba-cuda/pull/504
+pipeline_extensions = [_PipelineExtension()]