# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0

__all__ = ["MatmulComputeType", "Matmul", "matmul"]

import copy
from collections import namedtuple
from collections.abc import Sequence
from dataclasses import dataclass
import dataclasses
import functools
import logging
import operator
import typing
import random

try:
    from cuda.core import Device, EventOptions
except ImportError:
    from cuda.core.experimental import Device, EventOptions
import numpy as np

from nvmath import memory

from nvmath.linalg.advanced import _algorithmmod
from nvmath.linalg.advanced import _configuration
from nvmath.bindings import cublas
from nvmath.bindings import cublasLt as cublaslt  # type: ignore

from nvmath.internal import formatters
from nvmath.internal import tensor_wrapper
from nvmath.internal import typemaps
from nvmath.internal import utils

from nvmath.linalg._internal import matmul_desc_ifc, matmul_pref_ifc, matrix_layout_ifc
from nvmath.linalg._internal.typemaps import (
    NAMES_TO_DEFAULT_SCALE_TYPE,
    NAMES_TO_DEFAULT_COMPUTE_TYPE,
    COMPUTE_TYPE_TO_DEFAULT_SCALE_TYPE,
    SCALE_TYPE_TO_DEFAULT_COMPUTE_TYPE,
    SUPPORTED_TYPES,
)
from nvmath.linalg._internal.utils import (
    axis_order_in_memory,
    calculate_strides,
    check_batch_tileable,
    get_handle,
    pointer_aligned_to,
)
from nvmath.linalg._internal.epilog_protocol import (
    EPILOG_INPUT_HANDLERS_MAP,
    EPILOG_OUTPUT_HANDLERS_MAP,
    EPILOG_MINIMUM_VERSIONS_MAP,
    BATCHED_EPILOG_MINIMUM_VERSIONS_MAP,
    EpilogOutputHandler,
)
from nvmath._utils import CudaDataType

MatmulComputeType = cublas.ComputeType

EpilogInputTraits = namedtuple("EpilogInputTraits", ["dtype", "extents", "strides"])


@dataclass
class MatrixLayout:
    """An internal data class for capturing the tensor layout."""

    shape: Sequence[int]
    strides: Sequence[int]
    is_conjugate: bool = False  # Used to support is_conjugate via conjugate_transpose.


@dataclass
class LayoutTraits:
    """An internal data class for capturing the matrix multiplication traits."""

    order: cublaslt.Order
    ld: int
    batch_offset: int  # Based on strides
    is_conjugate: bool  # Used to support is_conjugate via conjugate_transpose.
    mm_shape: Sequence[int] | None = None
    mm_strides: Sequence[int] | None = None

    def get_mm_layout(self, transpose=False):
        if self.is_conjugate:
            transpose = True
        if not transpose:
            return *self.mm_shape, self.ld, self.order

        # Use of transpose is supported only for A and B for two specific use cases till the
        # C library directly supports these use cases:
        #
        #   1. When A or B has the conjugate qualifier, we transpose it internally and then
        #      use conjugate transpose in the MM (A @ B.conj() == A @ B.T.H).
        #
        #   2. When the epilog is BGRADB, we transpose B internally and use transpose in the
        #      MM since this epilog requires B to be transposed (A @ B == A @ B.T.T).
        #
        #  This requires that the layout order be ROW or COL (no special layouts such as
        #  structured or hierarchical).
        assert self.mm_shape is not None and self.mm_strides is not None, "Internal Error."
        assert self.ld != 0, "Internal Error."

        mm_shape = self.mm_shape[1], self.mm_shape[0]
        if self.order == cublaslt.Order.ROW:
            order = cublaslt.Order.COL
            ld = max(self.mm_shape[1], self.mm_strides[0])
        elif self.order == cublaslt.Order.COL:
            order = cublaslt.Order.ROW
            ld = max(self.mm_shape[0], self.mm_strides[1])
        else:
            raise AssertionError("Internal Error. Invalid layout order.")

        return *mm_shape, ld, order


@dataclass
class MMTraits:
    """An internal data class for capturing the matrix multiplication traits. The
    result traits are captured separately, because we need to wait for the
    epilog to be provided.
    """

    M: int
    N: int
    K: int
    d_mm_shape: Sequence[int]
    inplace: bool
    c_layout: MatrixLayout
    a_layout_traits: LayoutTraits
    b_layout_traits: LayoutTraits
    c_layout_traits: LayoutTraits
    batch_count: int
    batch_shape: Sequence[int]
    batch_axis_order: Sequence[int]


@dataclass
class ResultTraits:
    """An internal data class for capturing the result matrix's traits."""

    d_layout_traits: LayoutTraits
    result_shape: Sequence[int]
    result_strides: Sequence[int]


def get_matrix_layout_traits(
    mm_shape: Sequence[int],
    mm_strides: Sequence[int],
    batch_strides: Sequence[int],
    col_bcast: bool,
    ordering: cublaslt.Order | None = None,
    orientation: cublaslt.Order | None = None,
) -> tuple[cublaslt.Order, int, int]:
    """
    The 'ordering' option specifies the layout order, if it's not None, as in the case of
    the D matrix whose layout is determined by the other operands' layout. It is required if
    the matrix is degenerate (a vector or scalar: len(mm_shape) < 2).

    The 'orientation' option (ROW or COL) is required to infer the correct leading dimension
    for degenerate matrices (a vector or scalar: len(mm_shape) < 2).
    """
    if len(mm_shape) < 2:  # The result D can be a scalar or vector.
        assert ordering is not None, "Internal Error: 'ordering' must be specified for degenerate matrices."
        assert orientation is not None, "Internal Error: 'orientation' must be specified for degenerate matrices."
        batch_offset = min(batch_strides) if batch_strides else 0
        order = ordering
        if len(mm_shape) < 1:
            ld = 1
        elif order != orientation:
            # For a ROW vector in COL order, the LD should be 1 since we promote the row
            # vector to a matrix as (1, M). Similarly, for a COL vector in ROW order, the LD
            # should be 1 as well since we promote the column vector to a matrix as (M, 1).
            ld = 1
        else:
            ld = max(mm_strides[0], mm_shape[0])
        return order, ld, batch_offset

    M, N = mm_shape

    if ordering is not None:
        order = ordering
        message = f"Internal Error: incompatible ordering '{ordering}' and strides {mm_strides}"
        if order == cublaslt.Order.ROW:
            assert mm_strides[0] >= mm_strides[1] and mm_strides[1] == 1, message
        else:
            assert mm_strides[1] >= mm_strides[0] and mm_strides[0] == 1, message
    else:
        # Important: start with the first dimension so that cases like (M, 1) : (1, 1) or
        # (1, M) : (1, 1) in CuTe notation map to COL.
        if mm_strides[0] == 1:
            order = cublaslt.Order.COL
        elif mm_strides[1] == 1:
            order = cublaslt.Order.ROW
        else:
            if M == 1:
                order = cublaslt.Order.COL
            elif N == 1:
                order = cublaslt.Order.ROW
            else:
                raise ValueError("Unsupported layout.")

    # We need to handle broadcast dimensions with zero-stride for the c matrix.
    if col_bcast and N == 1:
        ld = 0
    else:
        ld = max(M, mm_strides[1]) if order == cublaslt.Order.COL else max(N, mm_strides[0])

    # Batch dimensions should be contiguous in memory, which we have already checked. The
    # batch_offset should be based on the lowest stride in the batch dimension to account
    # for embedded matrices.
    batch_offset = min(batch_strides) if batch_strides else 0

    return order, ld, batch_offset


def get_mm_traits(a_layout, b_layout, c_layout, inplace, logger):
    """
    First check A and B compatibility:

    1. Check MM compatibility (K):
        a. First pad A and/or B MM dimensions if 1-D according to NumPy convention.
        b. The padding is used to determine M, N, and K but should not appear in the output
           dimensions.
        c. If both A and B are N-D, the dimensions must match.
    2. Check batch dimensions:
        a. One of A or B can have missing batch extents, in which case it is broadcast,
           otherwise
        b. A and B must have the same batch ordering.
        c. In addition, the batch dimensions must be tileable (contiguous in memory).

    Then check C:

    C can be None. If C is passed in, it must be a matrix. The batching rule is the
    same as above.

    For the inplace option (result == C), C cannot be broadcast and so must have batch
    dimensions (if batched).
    """
    inplace_addendum = ""
    if inplace:
        inplace_addendum = "This is required for inplace operations."

    a_shape, a_strides = list(a_layout.shape), list(a_layout.strides)
    b_shape, b_strides = list(b_layout.shape), list(b_layout.strides)

    a_batch_shape, a_mm_shape = a_shape[:-2], a_shape[-2:]
    b_batch_shape, b_mm_shape = b_shape[:-2], b_shape[-2:]

    a_batch_strides, a_mm_strides = a_strides[:-2], a_strides[-2:]
    b_batch_strides, b_mm_strides = b_strides[:-2], b_strides[-2:]
    d_mm_shape = []
    if len(a_mm_shape) == 1:
        s, d = a_mm_shape[0], a_mm_strides[0]
        a_mm_shape = [1] + a_mm_shape
        a_mm_strides = [s * d] + a_mm_strides
    else:
        d_mm_shape.append(a_mm_shape[0])  # The first mode for d applies only when a is not a vector.

    if len(b_mm_shape) == 1:
        s, d = b_mm_shape[0], b_mm_strides[0]
        b_mm_shape = b_mm_shape + [1]
        b_mm_strides = b_mm_strides + [s * d]
    else:
        d_mm_shape.append(b_mm_shape[1])  # The second mode for d applies only when b is not a vector.

    logger.debug(f"The MM shape for operand A is {a_mm_shape} with strides {a_mm_strides}.")
    logger.debug(f"The MM shape for operand B is {b_mm_shape} with strides {b_mm_strides}.")
    logger.debug(f"The MM shape for operand D is {d_mm_shape}.")

    M0, K0 = a_mm_shape
    K1, N0 = b_mm_shape
    if K0 != K1:
        raise ValueError(
            f"The 'K' extent must match for the operands: K={K0} in operand A is not equal to K={K1} in operand B."
        )

    # Check if batch dimensions of A and B are tileable as well as compatible.
    batch_shape, batch_axis_order = [], ()
    if len(a_batch_shape) > 0:
        if not check_batch_tileable(a_batch_shape, a_batch_strides):
            message = (
                f"The batch layout for A corresponding to shape = {a_batch_shape} and strides = {a_batch_strides} "
                "is currently not supported because it is not tileable."
            )
            raise ValueError(message)
        logger.debug(
            f"The batch layout for A corresponding to shape = {a_batch_shape} and strides = {a_batch_strides} IS tileable."
        )
        batch_shape = a_batch_shape
        batch_axis_order = a_batch_axis_order = axis_order_in_memory(a_batch_strides)

    if len(b_batch_shape) > 0:
        if not check_batch_tileable(b_batch_shape, b_batch_strides):
            message = (
                f"The batch layout for B corresponding to shape = {b_batch_shape} and strides = {b_batch_strides} "
                "is currently not supported because it is not tileable."
            )
            raise ValueError(message)
        logger.debug(
            f"The batch layout for B corresponding to shape = {b_batch_shape} and strides = {b_batch_strides} IS tileable."
        )
        batch_shape = b_batch_shape
        batch_axis_order = b_batch_axis_order = axis_order_in_memory(b_batch_strides)

    if len(a_batch_shape) > 0 and len(b_batch_shape) > 0:
        if a_batch_shape != b_batch_shape:
            raise ValueError(f"The batch dimensions of operands A {a_batch_shape} and B {b_batch_shape} must match.")
        if a_batch_axis_order != b_batch_axis_order:
            raise ValueError(f"The batch order of operands A {a_batch_axis_order} and B {b_batch_axis_order} must match.")

    logger.debug(f"The batch shape is {batch_shape} with batch axis order {batch_axis_order}.")

    batch_count = functools.reduce(operator.mul, batch_shape, 1)

    # Create matrix layout traits.
    a_order, a_ld, a_batch_offset = get_matrix_layout_traits(a_mm_shape, a_mm_strides, a_batch_strides, col_bcast=False)
    a_layout_traits = LayoutTraits(
        order=a_order,
        ld=a_ld,
        batch_offset=a_batch_offset,
        is_conjugate=a_layout.is_conjugate,
        mm_shape=a_mm_shape,
        mm_strides=a_mm_strides,
    )
    logger.debug(f"The layout order for operand A is {a_order.name}, with LD {a_ld}, and batch offset {a_batch_offset}.")

    b_order, b_ld, b_batch_offset = get_matrix_layout_traits(b_mm_shape, b_mm_strides, b_batch_strides, col_bcast=False)
    b_layout_traits = LayoutTraits(
        order=b_order,
        ld=b_ld,
        batch_offset=b_batch_offset,
        is_conjugate=b_layout.is_conjugate,
        mm_shape=b_mm_shape,
        mm_strides=b_mm_strides,
    )
    logger.debug(f"The layout order for operand B is {b_order.name}, with LD {b_ld}, and batch offset {b_batch_offset}.")

    # Process matrix c, if provided.
    c_layout_traits = None
    if c_layout is not None:
        # 1. C cannot be a vector.
        # 2. C can be a matrix of dimension (M, N) or (M, 1), broadcast in the latter case
        #    and has to have contiguous strides.
        # 3. C can be batched matrices of dimension (..., M, N) or (..., M, 1), broadcast in
        #    the latter case and has to have contiguous strides.
        c_shape, c_strides = list(c_layout.shape), list(c_layout.strides)

        c_batch_shape, c_mm_shape = c_shape[:-2], c_shape[-2:]
        c_batch_strides, c_mm_strides = c_strides[:-2], c_strides[-2:]
        if len(c_mm_shape) == 1:
            raise ValueError(f"C cannot be a vector. C shape: {c_mm_shape}")
        logger.debug(f"The MM shape for operand C is {c_mm_shape} with strides {c_mm_strides}.")

        Mc, Nc = c_mm_shape
        if Mc != M0:
            raise ValueError(f"The M dimension of the C matrix ({Mc}) must match the M dimension of A.")

        if (inplace and Nc != N0) or (Nc != 1 and Nc != N0):
            raise ValueError(f"The N dimension of the C matrix ({Nc}) must match the N dimension of B. {inplace_addendum}")

        # For the inplace option, C must be batched if an operand is batched.
        if inplace or len(c_batch_shape) > 0:
            if c_batch_shape != batch_shape:
                raise ValueError(
                    f"The batch dimension of operand C {c_batch_shape} must match with that of the other operands "
                    f"{batch_shape}. {inplace_addendum}"
                )

            if batch_shape and (c_batch_axis_order := axis_order_in_memory(c_batch_strides)) != batch_axis_order:
                raise ValueError(
                    f"The batch axis order of operand C {c_batch_axis_order} must match with that of the other "
                    f"operands {batch_axis_order}."
                )

            if batch_shape and not check_batch_tileable(c_batch_shape, c_batch_strides):
                message = (
                    f"The batch layout for C corresponding to shape = {c_batch_shape} and strides = "
                    f"{c_batch_strides} is currently not supported because it is not tileable."
                )
                raise ValueError(message)

        c_order, c_ld, c_batch_offset = get_matrix_layout_traits(
            c_mm_shape, c_mm_strides, c_batch_strides, col_bcast=not inplace
        )
        c_layout_traits = LayoutTraits(order=c_order, ld=c_ld, batch_offset=c_batch_offset, is_conjugate=c_layout.is_conjugate)
        logger.debug(f"The layout order for operand C is {c_order.name}, with LD {c_ld}, and batch offset {c_batch_offset}.")

    return MMTraits(
        M=M0,
        N=N0,
        K=K0,
        d_mm_shape=d_mm_shape,
        inplace=inplace,
        c_layout=c_layout,
        batch_count=batch_count,
        batch_shape=batch_shape,
        batch_axis_order=batch_axis_order,
        a_layout_traits=a_layout_traits,
        b_layout_traits=b_layout_traits,
        c_layout_traits=c_layout_traits,
    )


def get_result_traits(mm_traits: MMTraits, epilog_ordering: cublaslt.Order | None, logger: logging.Logger) -> ResultTraits:
    """
    epilog_ordering = value of type cublaslt.Order or None.

    The result layout is determined from:
        - the ordering of operand c, if it is provided, or
        - the epilog requirement, if it exists, or
        - the ordering of operand a.

    The result batch dimensions must have the same extents and axis order as the inputs. The
    MM layout can be C or F.
    """

    # For inplace=True, the result traits are essentially the same as for operand C except
    # for conjugation.
    if mm_traits.inplace:
        result_shape, result_strides = mm_traits.c_layout.shape, mm_traits.c_layout.strides
        d_layout_traits = dataclasses.replace(mm_traits.c_layout_traits, is_conjugate=False)
        return ResultTraits(result_shape=result_shape, result_strides=result_strides, d_layout_traits=d_layout_traits)

    # The result shape is the batch shape + d_mm_shape.
    result_shape = (*mm_traits.batch_shape, *mm_traits.d_mm_shape)

    if mm_traits.c_layout_traits is not None:
        result_ordering = mm_traits.c_layout_traits.order
    elif epilog_ordering is not None:
        result_ordering = epilog_ordering
    else:
        result_ordering = mm_traits.a_layout_traits.order

    if result_ordering == cublaslt.Order.ROW:
        d_order = list(range(len(mm_traits.d_mm_shape) - 1, -1, -1))
    elif result_ordering == cublaslt.Order.COL:
        d_order = list(range(len(mm_traits.d_mm_shape)))
    else:
        raise AssertionError("Internal Error.")

    result_axis_order = [len(mm_traits.batch_axis_order) + a for a in d_order] + list(mm_traits.batch_axis_order)

    # Calculate the result strides.
    result_strides = calculate_strides(result_shape, result_axis_order)

    # For degenerate matrices, we need to specify the result orientation.
    result_orientation = None
    if len(mm_traits.d_mm_shape) < 2:
        if mm_traits.M == 1:
            result_orientation = cublaslt.Order.ROW
        elif mm_traits.N == 1:
            result_orientation = cublaslt.Order.COL

    # The result's traits.
    d_batch_strides, d_mm_strides = (
        result_strides[: len(mm_traits.batch_shape)],
        result_strides[len(mm_traits.batch_shape) :],
    )
    order, d_ld, d_batch_offset = get_matrix_layout_traits(
        mm_traits.d_mm_shape,
        d_mm_strides,
        d_batch_strides,
        col_bcast=False,
        ordering=result_ordering,
        orientation=result_orientation,
    )
    assert order == result_ordering, (
        f"Internal Error: d_order = {order.name}, result_ordering = {result_ordering.name}, mm_traits = {mm_traits}."
    )
    d_layout_traits = LayoutTraits(order=order, ld=d_ld, batch_offset=d_batch_offset, is_conjugate=False)
    logger.debug(f"The layout order for operand D is {order.name}, with LD {d_ld}, and batch offset {d_batch_offset}.")

    return ResultTraits(result_shape=result_shape, result_strides=result_strides, d_layout_traits=d_layout_traits)


SHARED_MM_DOCUMENTATION = utils.COMMON_SHARED_DOC_MAP.copy()
SHARED_MM_DOCUMENTATION.update(
    {
        "a": """\
A tensor representing the first operand to the matrix multiplication (see `Semantics`_). The currently supported types
are :class:`numpy.ndarray`, :class:`cupy.ndarray`, and :class:`torch.Tensor`.""".replace("\n", " "),
        #
        "b": """\
A tensor representing the second operand to the matrix multiplication (see `Semantics`_). The currently supported types
are :class:`numpy.ndarray`, :class:`cupy.ndarray`, and :class:`torch.Tensor`.""".replace("\n", " "),
        #
        "c": """\
(Optional) A tensor representing the operand to add to the matrix multiplication result (see `Semantics`_). The currently
supported types are :class:`numpy.ndarray`, :class:`cupy.ndarray`, and :class:`torch.Tensor`.""".replace("\n", " "),
        #
        "c_admonitions": """
                .. versionchanged:: 0.3.0
                    In order to avoid broadcasting behavior ambiguity, nvmath-python no longer
                    accepts a 1-D (vector) `c`. Use a singleton dimension to convert your input
                    array to 2-D.
""",
        #
        "alpha": """\
The scale factor for the matrix multiplication term as a real or complex number. The default is
:math:`1.0`.""".replace("\n", " "),
        #
        "beta": """\
The scale factor for the matrix addition term as a real or complex number. A value for `beta` must be provided if
operand `c` is specified.""".replace("\n", " "),
        #
        "quantization_scales": """\
Specify scale factors for the matrix multiplication as a :class:`~nvmath.linalg.advanced.MatmulQuantizationScales`
object. Alternatively, a `dict` containing the parameters for the
:class:`~nvmath.linalg.advanced.MatmulQuantizationScales`
constructor can also be provided. The scale factors can be provided as scalars or tensors.
If a scale factor is provided as a tensor, it must be from the same package and on the same
memory space (CPU or GPU device) as the operands of the matmul.
Allowed and required only for narrow-precision (FP8 and lower) operations.""".replace("\n", " "),
        #
        "algorithms": """\
A sequence of :class:`Algorithm` objects that can be directly provided to bypass planning. The algorithm objects must be
compatible with the matrix multiplication. A typical use for this option is to provide algorithms serialized (pickled)
from a previously planned and autotuned matrix multiplication.""".replace("\n", " "),
        #
        "epilog": """\
Specify an epilog :math:`F` as an object of type :class:`MatmulEpilog` to apply to the result of the matrix
multiplication: :math:`F(\\alpha A @ B + \\beta C`). The default is no epilog. See `cuBLASLt documentation
<https://docs.nvidia.com/cuda/cublas/#cublasltepilogue-t>`_ for the list of available epilogs.""".replace("\n", " "),
        #
        "epilog_inputs": """\
Specify the additional inputs needed for the selected epilog as a dictionary, where the key is the epilog input name and
the value is the epilog input. The epilog input must be a tensor with the same package and in the same memory space as
the operands (see the constructor for more information on the operands). If the required epilog inputs are not provided,
an exception is raised that lists the required epilog inputs. Some epilog inputs are generated by other epilogs. For
example, the epilog input for :class:`MatmulEpilog.DRELU` is generated by matrix multiplication with the same operands
using :class:`MatmulEpilog.RELU_AUX`. """.replace("\n", " "),
        #
        "qualifiers": """\
If desired, specify the matrix qualifiers as a :class:`numpy.ndarray` of
:class:`~nvmath.linalg.advanced.matrix_qualifiers_dtype` objects of length 3 corresponding to the operands `a`, `b`, and
`c`. See :ref:`matrix-tensor-qualifiers` for the motivation behind
qualifiers.""".replace("\n", " "),
        #
        "options": """\
Specify options for the matrix multiplication as a :class:`~nvmath.linalg.advanced.MatmulOptions` object. Alternatively,
a `dict` containing the parameters for the ``MatmulOptions`` constructor can also be provided. If not specified, the
value will be set to the default-constructed ``MatmulOptions`` object.""".replace("\n", " "),
        #
        "preferences": """\
This parameter specifies the preferences for planning as a :class:`MatmulPlanPreferences` object. Alternatively, a
dictionary containing the parameters for the :class:`MatmulPlanPreferences` constructor can also be provided. If not
specified, the value will be set to the default-constructed :class:`MatmulPlanPreferences` object.
""".replace("\n", " "),
        #
        "result": """\
The result of the specified matrix multiplication (epilog applied), which remains on the same device and belongs to the
same package as the input operands. If an epilog (like :attr:`nvmath.linalg.advanced.MatmulEpilog.RELU_AUX`) that
results in extra output is used, or an extra output is requested (for example by setting
:attr:`~nvmath.linalg.advanced.MatmulOptions.result_amax` option in ``options`` argument),
a tuple is returned with the first element being the matrix multiplication result (epilog applied) and the second element
being the auxiliary output provided as a `dict`. """.replace("\n", " "),
        #
        "narrow_precision": """\
        Matrix multiplication with narrow-precision operands is supported, in both FP8 and MXFP8 formats.

        .. note::

            Narrow-precision matrix multiplication in nvmath-python requires **CUDA Toolkit 12.8 or newer**.
            **FP8 requires a device with compute capability 8.9 or higher** (Ada, Hopper, Blackwell or newer architecture).
            **MXFP8 requires a device with compute capability 10.0 or higher** (Blackwell or newer architecture).
            Please refer to the `compute capability table <https://developer.nvidia.com/cuda-gpus>`_
            to check the compute capability of your device.

        For FP8 operations:

        * For each operand a scaling factor needs to be specified via ``quantization_scales`` argument.
        * Maximum absolute value of the result (amax) can be requested via
          :attr:`~nvmath.linalg.advanced.MatmulOptions.result_amax` option in ``options`` argument.
        * Custom result type (both FP8 and non-FP8) can be requested via
          :attr:`~nvmath.linalg.advanced.MatmulOptions.result_type` option in ``options`` argument.

        For MXFP8 operations:

        * To enable MXFP8 operations, :attr:`~nvmath.linalg.advanced.MatmulOptions.block_scaling` option
          must be set to ``True``.
        * Block scaling factors need to be specified via ``quantization_scales`` argument.
        * Utilities in :mod:`nvmath.linalg.advanced.helpers.matmul` can be used to create and modify
          block scaling factors.
        * When MXFP8 is used and the result type is a narrow-precision data type, the auxiliary output
          ``"d_out_scale"`` will be returned in the auxiliary output tensor. It will contain the scales
          that were used for the result quantization.

        Please refer to the examples and narrow-precision operations tutorial for more details.
        For more details on the FP8 and MXFP8 formats in cuBLAS,
        see the `cublasLtMatmul documentation <https://docs.nvidia.com/cuda/cublas/#cublasltmatmul>`_.
""".strip(),
        #
        "semantics": """\
        .. _semantics:

        The semantics of the matrix multiplication follows :external:py:data:`numpy.matmul` semantics, with some restrictions on
        broadcasting. In addition, the semantics for the fused matrix addition are described below:

        * For in-place matrix multiplication (where the result is written into `c`) the result has the same shape as `c`.
        * If arguments `a` and `b` are matrices, they are multiplied according to the rules of matrix multiplication.
        * If argument `a` is 1-D, it is promoted to a matrix by prefixing ``1`` to its dimensions. After matrix
          multiplication, the prefixed ``1`` is removed from the result's dimensions if the operation is not in-place.
        * If argument `b` is 1-D, it is promoted to a matrix by appending ``1`` to its dimensions. After matrix
          multiplication, the appended ``1`` is removed from the result's dimensions if the operation is not in-place.
        * If `a` or `b` is N-D (N > 2), then the operand is treated as a batch of matrices. If both `a` and `b` are N-D,
          their batch dimensions must match. If exactly one of `a` or `b` is N-D, the other operand is broadcast.
        * The operand for the matrix addition `c` may be a matrix of shape (M, 1) or (M, N), or the batched versions
          (..., M, 1) or (..., M, N). Here M and N are the dimensions of the result of the matrix multiplication. If N = 1, the
          columns of `c` are broadcast for the addition; the rows of `c` are never broadcast. If batch dimensions are not
          present, `c` is broadcast across batches as needed. If the operation is in-place, `c` cannot be broadcast since
          it must be large enough to hold the result.
        * Similarly, when operating on a batch, auxiliary outputs are 3-D for all epilogs. Therefore, epilogs that return 1-D
          vectors of length N in non-batched mode return 3-D matrices of size (batch, N, 1) in batched mode.
""".strip(),
    }
)


class InvalidMatmulState(Exception):
    pass


def _check_extents(shape: tuple, name: str):
    if any(e <= 0 for e in shape):
        message = f"The specified extents {shape} for operand {name} are not valid. The extents must be strictly positive. "
        raise ValueError(message)


@utils.docstring_decorator(SHARED_MM_DOCUMENTATION, skip_missing=False)
class Matmul:
    """
    Create a stateful object encapsulating the specified matrix multiplication computation
    :math:`\\alpha a @ b + \\beta c` and the required resources to perform the operation.  A
    stateful object can be used to amortize the cost of preparation (planning in the case of
    matrix multiplication) across multiple executions (also see the :ref:`Stateful APIs
    <host api types>` section).

    The function-form API :func:`matmul` is a convenient alternative to using stateful
    objects for *single* use (the user needs to perform just one matrix multiplication, for
    example), in which case there is no possibility of amortizing preparatory costs. The
    function-form APIs are just convenience wrappers around the stateful object APIs.

    Using the stateful object typically involves the following steps:

    1. **Problem Specification**: Initialize the object with a defined operation and
       options.
    2. **Preparation**: Use :meth:`plan` to determine the best algorithmic implementation
       for this specific matrix multiplication operation.
    3. **Execution**: Perform the matrix multiplication computation with :meth:`execute`.
    4. **Resource Management**: Ensure all resources are released either by explicitly
       calling :meth:`free` or by managing the stateful object within a context manager.

    Detailed information on what's happening in the various phases described above can be
    obtained by passing in a :class:`logging.Logger` object to :class:`MatmulOptions` or by
    setting the appropriate options in the root logger object, which is used by default:

        >>> import logging
        >>> logging.basicConfig(
        ...     level=logging.INFO,
        ...     format="%(asctime)s %(levelname)-8s %(message)s",
        ...     datefmt="%m-%d %H:%M:%S",
        ... )

    A user can select the desired logging level and, in general, take advantage of all of
    the functionality offered by the Python `logging` module.

    Args:
        a: {a}

        b: {b}

        c: {c}
            {c_admonitions}

        alpha: {alpha}

        beta: {beta}

        qualifiers: {qualifiers}

        options: {options}

        stream: {stream}

        quantization_scales: {quantization_scales}

    Semantics:
        {semantics}

    Narrow-precision support:
        {narrow_precision}

    .. seealso::
        :meth:`autotune`, :meth:`plan`, :meth:`reset_operands`, :meth:`execute`

    Examples:

        >>> import numpy as np
        >>> import nvmath

        Create two 2-D float64 ndarrays on the CPU:

        >>> M, N, K = 1024, 1024, 1024
        >>> a = np.random.rand(M, K)
        >>> b = np.random.rand(K, N)

        We will define a matrix multiplication operation followed by a RELU epilog function
        using the specialized matrix multiplication interface.

        Create a Matmul object encapsulating the problem specification above:

        >>> mm = nvmath.linalg.advanced.Matmul(a, b)

        Options can be provided above to control the behavior of the operation using the
        `options` argument (see :class:`MatmulOptions`).

        Next, plan the operation. The epilog is specified, and optionally, preferences can
        be specified for planning:

        >>> epilog = nvmath.linalg.advanced.MatmulEpilog.RELU
        >>> algorithms = mm.plan(epilog=epilog)

        Certain epilog choices (like :attr:`nvmath.linalg.advanced.MatmulEpilog.BIAS`)
        require additional input provided using the `epilog_inputs` argument to
        :meth:`plan`.

        Now execute the matrix multiplication, and obtain the result `r1` as a NumPy
        ndarray.

        >>> r1 = mm.execute()

        Finally, free the object's resources. To avoid having to explicitly make this
        call, it's recommended to use the Matmul object as a context manager as shown below,
        if possible.

        >>> mm.free()

        Note that all :class:`Matmul` methods execute on the current stream by default.
        Alternatively, the `stream` argument can be used to run a method on a specified
        stream.

        Let's now look at the same problem with CuPy ndarrays on the GPU.

        >>> import cupy as cp
        >>> a = cp.random.rand(M, K)
        >>> b = cp.random.rand(K, N)

        Create an Matmul object encapsulating the problem specification described earlier
        and use it as a context manager.

        >>> with nvmath.linalg.advanced.Matmul(a, b) as mm:
        ...     algorithms = mm.plan(epilog=epilog)
        ...
        ...     # Execute the operation to get the first result.
        ...     r1 = mm.execute()
        ...
        ...     # Update operands A and B in-place (see reset_operands() for an
        ...     # alternative).
        ...     a[:] = cp.random.rand(M, K)
        ...     b[:] = cp.random.rand(K, N)
        ...
        ...     # Execute the operation to get the new result.
        ...     r2 = mm.execute()


        All the resources used by the object are released at the end of the block.

        Further examples can be found in the `nvmath/examples/linalg/advanced/matmul
        <https://github.com/NVIDIA/nvmath-python/tree/main/examples/linalg/advanced/matmul>`_
        directory.
    """

    def __init__(
        self,
        a,
        b,
        /,
        c=None,
        *,
        alpha=None,
        beta=None,
        qualifiers=None,
        quantization_scales=None,
        options=None,
        stream: utils.AnyStream | int | None = None,
    ):
        options = utils.check_or_create_options(_configuration.MatmulOptions, options, "Matrix multiplication options")
        assert options is not None
        self.options = options

        if c is None and options.inplace:
            raise ValueError("The operation cannot be inplace if operand C is not provided.")

        self.logger = options.logger if options.logger is not None else logging.getLogger()

        if options.inplace and options.result_type is not None:
            self.logger.warning(
                f"Matmul: The provided result type {options.result_type} in options is ignored since \
the operation is in-place."
            )
        self.inplace = options.inplace

        def check_dtype(dtype, operand_name):
            if dtype not in SUPPORTED_TYPES:
                raise ValueError(f"The dtype of operand {operand_name} ({dtype}) is not supported.")

        # The matrix multiplication has two required operands 'a' and 'b', and one optional
        # operand 'c'.
        a = tensor_wrapper.wrap_operand(a)
        b = tensor_wrapper.wrap_operand(b)
        check_dtype(a.dtype, "A")
        check_dtype(b.dtype, "B")
        self.logger.info("= SPECIFICATION PHASE =")
        if self.inplace:
            self.logger.info("The MM operation will be performed in-place (the result will be written into operand C).")
        self.logger.info(f"The data type of operand A is '{a.dtype}', and that of operand B is '{b.dtype}'.")

        self.num_operands = 2
        if c is not None:
            self.num_operands = 3
            c = tensor_wrapper.wrap_operand(c)
            if len(c.shape) < 2:
                raise ValueError(
                    "In order to avoid broadcasting behavior ambiguity, `c` must be at least 2-D. "
                    "Use a singleton dimension to convert your input array to 2-D."
                )
            check_dtype(c.dtype, "C")
            self.logger.info(f"The data type of operand C is {c.dtype}.")

        if c is not None and beta is None:
            raise ValueError("A value for beta must be provided if operand C is provided.")

        if (a.dtype, b.dtype) not in NAMES_TO_DEFAULT_SCALE_TYPE:
            raise ValueError(f"Unsupported combination of dtypes for operands A {a.dtype} and B {b.dtype}.")

        # Currently, a.dtype != b.dtype is only supported for FP8 (different FP8 kinds are
        # allowed), so we assume that A and B have equal width.
        self.input_type_width = typemaps.NAME_TO_DATA_WIDTH[a.dtype]

        assert self.num_operands == 2 or self.num_operands == 3, "Internal Error."

        _check_extents(a.shape, "a")
        _check_extents(b.shape, "b")
        if c is not None:
            _check_extents(c.shape, "c")

        # Infer the library package & device ID the operands belong to.
        operands = [a, b]
        if self.num_operands == 3:
            operands.append(c)
        self.operands: None | Sequence[utils.TensorHolder | None] = operands

        self.package = utils.get_operands_package(operands)
        self.memory_space = "cuda"
        self.device_id = utils.get_operands_device_id(operands)
        if self.device_id == "cpu":
            if self.package == "numpy":
                self.package = "cuda"
            self.memory_space = "cpu"
            self.device_id = options.device_id
        self.logger.info(
            f"The input operands' memory space is {self.memory_space}, and the execution space is on device {self.device_id}."
        )

        # Allocate device memory (in stream context) if needed.
        stream_holder = utils.get_or_create_stream(self.device_id, stream, self.package)
        self.logger.info(f"The specified stream for the Matmul ctor is {stream_holder.obj}.")

        # Copy operands to device (and store reference to CPU operand), if needed.
        self.cpu_c_ref = None
        if self.memory_space == "cpu":
            if self.inplace:
                self.cpu_c_ref = self.operands[2]  # Hold reference, needed for inplace operations.
            self.operands = tensor_wrapper.to(self.operands, self.device_id, stream_holder)

        # Set qualifiers.
        self.qualifiers = qualifiers if qualifiers is not None else np.zeros((3,), dtype=_configuration.matrix_qualifiers_dtype)
        if self.qualifiers.dtype != _configuration.matrix_qualifiers_dtype:
            raise ValueError(
                "The qualifiers must be specified as a NumPy array of length 3 corresponding to the operands A, B, and "
                "C of type 'matrix_qualifiers_dtype'."
            )
        # Set qualifiers based on torch lazy conjugation flag if not provided.
        self.qualifiers[0]["is_conjugate"] = self.qualifiers[0]["is_conjugate"] ^ self.operands[0].is_conjugate
        self.qualifiers[1]["is_conjugate"] = self.qualifiers[1]["is_conjugate"] ^ self.operands[1].is_conjugate
        self.lazy_conjugation = (self.operands[0].is_conjugate, self.operands[1].is_conjugate, False)
        if self.num_operands == 3:
            self.qualifiers[2]["is_conjugate"] = self.qualifiers[2]["is_conjugate"] ^ self.operands[2].is_conjugate
        if self.qualifiers[2]["is_conjugate"]:
            raise ValueError("The conjugate flag is currently not supported for operand C.")

        # Set blocking or non-blocking behavior.
        self.blocking = self.options.blocking is True or self.memory_space == "cpu"
        if self.blocking:
            self.call_prologue = "This call is blocking and will return only after the operation is complete."
        else:
            self.call_prologue = (
                "This call is non-blocking and will return immediately after the operation is launched on the device."
            )

        # The result class is that of the first wrapped device operand.
        self.result_class = self.operands[0].__class__

        # Set memory allocator.
        self.allocator = (
            options.allocator
            if options.allocator is not None
            else memory._MEMORY_MANAGER[self.package](self.device_id, self.logger)
        )

        # Set memory limit.
        self.memory_limit = utils.get_memory_limit_from_device_id(self.options.memory_limit, self.device_id)
        self.logger.info(f"The memory limit is {formatters.MemoryStr(self.memory_limit)}.")

        # Set handle. We don't destroy handles we create.
        if options.handle is not None:
            self.handle = options.handle
        else:
            self.handle = get_handle(self.device_id)

        # Determine the data types for a and b.
        self.a_dtype = typemaps.NAME_TO_DATA_TYPE[a.dtype]
        self.b_dtype = typemaps.NAME_TO_DATA_TYPE[b.dtype]
        self.a_dtype_name = a.dtype
        self.b_dtype_name = b.dtype

        self.is_complex = "complex" in self.a_dtype_name or "complex" in self.b_dtype_name

        # Determine the data types for c and d.
        self.d_dtype = None if self.inplace else options.result_type
        if self.num_operands == 3:
            self.c_dtype = typemaps.NAME_TO_DATA_TYPE[c.dtype]
            if self.d_dtype is None:
                self.d_dtype = self.c_dtype
        elif self.num_operands == 2:
            if self.d_dtype is None:
                self.d_dtype = self.a_dtype
            if self.d_dtype in (CudaDataType.CUDA_R_8F_E5M2, CudaDataType.CUDA_R_8F_E4M3):
                self.c_dtype = CudaDataType.CUDA_R_16F
            else:
                self.c_dtype = self.d_dtype
        self.c_dtype_name = typemaps.DATA_TYPE_TO_NAME[self.c_dtype]
        self.d_dtype_name = typemaps.DATA_TYPE_TO_NAME[self.d_dtype]
        self.c_dtype_width = typemaps.NAME_TO_DATA_WIDTH[self.c_dtype_name]
        self.d_dtype_width = typemaps.NAME_TO_DATA_WIDTH[self.d_dtype_name]

        self.logger.info(f"The data type for the result D is '{self.d_dtype_name}'.")

        def assert_valid_compute_type(compute_type):
            if compute_type not in COMPUTE_TYPE_TO_DEFAULT_SCALE_TYPE["real"]:
                message = f"Unsupported compute type. The compute type '{repr(compute_type)}' is currently not supported."
                raise ValueError(message)

        # Determine the scale type.
        if options.scale_type is None:
            if options.compute_type is not None:
                assert_valid_compute_type(options.compute_type)
                if self.is_complex:
                    scale_type_map = COMPUTE_TYPE_TO_DEFAULT_SCALE_TYPE["complex"]
                else:
                    scale_type_map = COMPUTE_TYPE_TO_DEFAULT_SCALE_TYPE["real"]
                self.scale_type = scale_type_map[options.compute_type]
            else:
                self.scale_type = NAMES_TO_DEFAULT_SCALE_TYPE[(self.a_dtype_name, self.b_dtype_name)]
            self.scale_type_name = typemaps.DATA_TYPE_TO_NAME[self.scale_type]
        else:
            self.scale_type = options.scale_type
            if self.scale_type not in SCALE_TYPE_TO_DEFAULT_COMPUTE_TYPE:
                message = f"Unsupported scale type. The data type '{repr(self.scale_type)}' is currently not supported."
                raise ValueError(message)
            self.scale_type_name = typemaps.DATA_TYPE_TO_NAME[self.scale_type]
        self.logger.info(f"The scale type is '{self.scale_type_name}'.")

        # Determine the compute type.
        if options.compute_type is None:
            if options.scale_type is not None:
                self.compute_type = SCALE_TYPE_TO_DEFAULT_COMPUTE_TYPE[options.scale_type]
            else:
                self.compute_type = NAMES_TO_DEFAULT_COMPUTE_TYPE[(self.a_dtype_name, self.b_dtype_name)]
        else:
            self.compute_type = options.compute_type
        assert_valid_compute_type(self.compute_type)
        self.logger.info(f"The compute type is {self.compute_type.name}.")

        def is_supported(atype, btype, compute_type, scale_type):
            ct = cublas.ComputeType
            st = CudaDataType
            abtype = atype if atype == btype else (atype, btype)
            if compute_type in (ct.COMPUTE_16F, ct.COMPUTE_16F_PEDANTIC):
                return scale_type == st.CUDA_R_16F and abtype == "float16"
            elif compute_type == ct.COMPUTE_32F_PEDANTIC:
                if scale_type == st.CUDA_R_32F:
                    return abtype in ("float32", "bfloat16", "float16", "float8_e4m3fn", "float8_e5m2")
                elif scale_type == st.CUDA_C_32F:
                    return abtype == "complex64"
            elif compute_type == ct.COMPUTE_32F:
                if scale_type == st.CUDA_R_32F:
                    return abtype in (
                        "float32",
                        "bfloat16",
                        "float16",
                        "float8_e4m3fn",
                        "float8_e5m2",
                        ("float8_e4m3fn", "float8_e5m2"),
                        ("float8_e5m2", "float8_e4m3fn"),
                    )
                elif scale_type == st.CUDA_C_32F:
                    return abtype == "complex64"
            elif compute_type in (
                ct.COMPUTE_32F_FAST_16F,
                ct.COMPUTE_32F_FAST_16BF,
                ct.COMPUTE_32F_FAST_TF32,
                ct.COMPUTE_32F_EMULATED_16BFX9,
            ):
                if scale_type == st.CUDA_R_32F:
                    return abtype == "float32"
                if scale_type == st.CUDA_C_32F:
                    return abtype == "complex64"
            elif compute_type in (ct.COMPUTE_64F, ct.COMPUTE_64F_PEDANTIC, ct.COMPUTE_64F_EMULATED_FIXEDPOINT):
                if scale_type == st.CUDA_R_64F:
                    return abtype == "float64"
                if scale_type == st.CUDA_C_64F:
                    return abtype == "complex128"
            return False

        if not is_supported(self.a_dtype_name, self.b_dtype_name, self.compute_type, self.scale_type):
            raise ValueError(
                f"Selected scale_type={repr(self.scale_type)} compute_type={repr(self.compute_type)} "
                + f"are not supported for data types {self.a_dtype_name} (A) and {self.b_dtype_name} (B)."
            )

        # Set alpha and beta.
        self.alpha = np.zeros((1,), dtype=self.scale_type_name)
        try:
            self.alpha[0] = alpha if alpha is not None else 1
        except (ValueError, TypeError) as e:
            raise ValueError(f"The value provided for alpha {alpha} is not convertible to dtype '{self.alpha.dtype}'.") from e

        self.beta = np.zeros((1,), dtype=self.scale_type_name)
        if beta is not None and self.num_operands == 2:
            self.logger.warning(f"Matmul: The provided beta value {beta} is ignored since operand C is not specified.")
        try:
            self.beta[0] = beta if beta is not None and self.num_operands == 3 else 0
        except (ValueError, TypeError) as e:
            raise ValueError(f"The value provided for beta {beta} is not convertible to dtype '{self.beta.dtype}'.") from e

        # Set narrow-precision (FP8 and lower) quantization_scales.
        if self.input_type_width <= 8:
            self.quantization_scales = self._validate_operand_scales(quantization_scales, all_required=True)
        elif quantization_scales is not None:
            self.logger.warning(
                "Matmul: The provided scales are ignored, since they are only applicable to narrow-precision (FP8 and lower) "
                "operations."
            )

        if self.options.result_amax and self.d_dtype_width > 8:
            raise ValueError("result_amax=True is allowed only for narrow-precision (FP8 and lower) results")

        # Check operands alignment if needed
        if self.input_type_width <= 8:
            for operand, operand_name in zip(self.operands, "ABC", strict=False):
                if operand.data_ptr % 16 != 0:
                    raise ValueError(
                        f"For narrow-precision (FP8 and lower) multiplication, operand {operand_name} should be aligned to 16 "
                        "bytes."
                    )

        # Capture operand extents and strides for consistency check when resetting operands.
        self.operand_extents = tuple(o.shape for o in self.operands)
        self.operand_strides = tuple(o.strides for o in self.operands)

        # Create operand layouts.
        a_layout = MatrixLayout(self.operands[0].shape, self.operands[0].strides, self.qualifiers[0]["is_conjugate"])
        b_layout = MatrixLayout(self.operands[1].shape, self.operands[1].strides, self.qualifiers[1]["is_conjugate"])
        c_layout = MatrixLayout(self.operands[2].shape, self.operands[2].strides) if self.num_operands == 3 else None

        # Enforce equal batch shape for A and B if block_scaling=True.
        if self.options.block_scaling and a_layout.shape[:-2] != b_layout.shape[:-2]:
            raise ValueError(
                "When block_scaling=True, the batch dimensions of A and B must match (broadcasting is not supported)."
            )

        # Get the operation traits.
        self.mm_traits = get_mm_traits(a_layout, b_layout, c_layout, self.inplace, self.logger)
        self.result_traits = None  # Wait till planning to determine this based on the epilog.
        self.logger.info(
            f"The matrix multiplication attributes are M = {self.mm_traits.M}, N = {self.mm_traits.N}, and "
            f"K = {self.mm_traits.K}."
        )
        self.logger.info(
            f"The batch count is {self.mm_traits.batch_count}, and the batch shape is {self.mm_traits.batch_shape} "
            f"with batch axis order {self.mm_traits.batch_axis_order}."
        )

        # Create and set the operation descriptor.
        self.mm_desc = cublaslt.matmul_desc_create(self.compute_type, self.scale_type)

        self.mm_desc_ifc = matmul_desc_ifc.MatmulDescInterface(self.mm_desc)
        self.mm_desc_ifc.compute_type = self.compute_type
        self.mm_desc_ifc.scale_type = self.scale_type

        # Guard SM count target and fast accumulation flag.
        version = cublaslt.get_version()
        if options.sm_count_target > 0:
            if version < 111103:
                raise ValueError(f"The 'sm_count_target' option is not supported in cuBLASLt version {version}.")
            self.mm_desc_ifc.sm_count_target = options.sm_count_target
            self.logger.info(f"The SM count target is {options.sm_count_target}.")

        if options.fast_accumulation:
            if version < 111103:
                raise ValueError(f"The 'fast_accumulation' option is not supported in cuBLASLt version {version}.")
            self.mm_desc_ifc.fast_accum = options.fast_accumulation
            self.logger.info(f"The flag for fast accumulation mode is {options.fast_accumulation}.")

        if self.input_type_width == 8 and version < 120800:
            raise ValueError(
                f"FP8 is not supported for cuBLASLt version {version}. cuBLASLt version 12.8 or higher is required."
            )

        # Planning preferences
        self.preferences = None

        # Epilog attributes.
        self.epilog = None

        # Epilog attributes: name-to-operand.
        self.epilog_operands: dict[str, typing.Any] = {}

        # Epilog attributes: epilog input name-to-handler.
        self.epilog_input_name_to_handler: dict[str, typing.Any] = {}

        # Epilog attributes: name-to-output tensor.
        self.epilog_outputs: dict[str, typing.Any] = {}

        # Keep track of epilog input traits for resetting operands.
        self.epilog_inputs_traits: dict[str, typing.Any] = {}

        # Keep track of epilog output handlers to allocate output in execute().
        self.epilog_output_handlers: list[EpilogOutputHandler] = []

        # Non-epilog aux outputs. Currently, only used for quantization outputs (amax etc.)
        self.aux_outputs = None

        # Plan attributes.
        self.preference_ptr = None
        self.a_layout_ptr, self.b_layout_ptr, self.c_layout_ptr, self.d_layout_ptr = None, None, None, None
        self.flop_count = 0
        self.mm_planned = False

        # Algorithm attributes.
        self.algorithms_buffer = None
        self.algorithm_objects = None
        self.cached_best_algorithm_struct = None

        # Workspace attributes.
        self.workspace_ptr: None | memory.MemoryPointer = None
        self.workspace_size = 0
        self.workspace_allocated_size = 0
        self.workspace_allocated_here = False

        # Attributes to establish stream ordering.
        self.workspace_stream = None
        self.last_compute_event = None

        # Device-side array with the quantization_scales
        self.quantization_scales_device: dict[str, utils.TensorHolder] = {}

        self.valid_state = True
        self.logger.info("The Matmul operation has been created.")

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.free()

    def _check_valid_matmul(self, *args, **kwargs):
        """
        Check if the Matmul object is alive and well.
        """
        if not self.valid_state:
            raise InvalidMatmulState("The Matmul object cannot be used after resources are free'd")

    def _check_valid_operands(self, *args, **kwargs):
        """
        Check if the operands are available for the operation.
        """
        what = kwargs["what"]
        if self.operands is None:
            raise RuntimeError(
                f"{what} cannot be performed if the operands have been set to None. Use reset_operands() to set the "
                f"desired input before using performing the {what.lower()}."
            )

    def _free_plan_resources(self, exception: Exception | None = None) -> bool:
        """
        Free resources allocated in planning.
        """

        # Destroy matrix layouts.
        if self.a_layout_ptr is not None:
            cublaslt.matrix_layout_destroy(self.a_layout_ptr)
            self.a_layout_ptr = None
        if self.b_layout_ptr is not None:
            cublaslt.matrix_layout_destroy(self.b_layout_ptr)
            self.b_layout_ptr = None
        if self.c_layout_ptr != self.d_layout_ptr and self.c_layout_ptr is not None:
            cublaslt.matrix_layout_destroy(self.c_layout_ptr)
        self.c_layout_ptr = None
        if self.d_layout_ptr is not None:
            cublaslt.matrix_layout_destroy(self.d_layout_ptr)
            self.d_layout_ptr = None

        if self.preference_ptr is not None:
            cublaslt.matmul_preference_destroy(self.preference_ptr)
            self.preference_ptr = None

        self.mm_planned = False
        return True

    def _check_planned(self, *args, **kwargs):
        what = kwargs["what"]
        if not self.mm_planned:
            raise RuntimeError(f"{what} cannot be performed before plan() has been called.")

    def _free_workspace_memory(self, exception: Exception | None = None) -> bool:
        """
        Free workspace by releasing the MemoryPointer object.
        """
        if self.workspace_ptr is None:
            return True

        self.workspace_ptr = None
        self.workspace_allocated_size = 0
        self.logger.debug("[_free_workspace_memory] The workspace has been released.")

        return True

    def _reset_workspace_allocation_tracking(self):
        """
        Reset workspace allocation tracking attributes to False at the end of the methods
        where workspace memory is potentially allocated. This is necessary to prevent any
        exceptions raised before method entry from using stale tracking values.
        """
        self.workspace_allocated_here = False

    @utils.precondition(_check_valid_matmul)
    def _release_workspace_memory_perhaps(self, release_workspace):
        """
        Free workspace memory if it's larger than the specified limit.
        """
        if not release_workspace:
            return True

        # Establish ordering wrt the computation and free workspace if requested.
        if self.last_compute_event is not None:
            self.workspace_stream.wait(self.last_compute_event)
            self.logger.debug("Established ordering with respect to the computation before releasing the workspace.")
            self.last_compute_event = None

        self.logger.debug("[_release_workspace_memory_perhaps] The workspace memory will be released.")
        return self._free_workspace_memory()

    def _release_workspace_memory_perhaps_wrapper(self, exception: Exception | None = None) -> bool:
        """
        This is used in @atomic.
        """
        self._release_workspace_memory_perhaps(release_workspace=self.workspace_allocated_here)
        self._reset_workspace_allocation_tracking()
        return True

    @utils.precondition(_check_valid_matmul)
    @utils.precondition(_check_planned, "Workspace memory allocation")
    @utils.atomic(_free_workspace_memory, method=True)
    def _allocate_workspace_memory(self, stream_holder: utils.StreamHolder):
        """
        Allocate workspace memory using the specified allocator.
        """

        assert self.workspace_size is not None, "Internal Error."
        assert self.workspace_allocated_here is False, "Internal Error."

        if self.workspace_size == 0:  # For performance, bypass allocator for workspace size == 0.
            self.workspace_ptr = memory.MemoryPointer(0, 0, finalizer=None)
        else:
            self.logger.debug("Allocating workspace for performing the matrix multiplication...")
            with utils.device_ctx(self.device_id), stream_holder.ctx:
                try:
                    if isinstance(self.allocator, memory.BaseCUDAMemoryManagerAsync):
                        self.workspace_ptr = self.allocator.memalloc_async(self.workspace_size, stream_holder.obj)
                    else:
                        self.workspace_ptr = self.allocator.memalloc(self.workspace_size)
                    self.workspace_allocated_here = True
                except TypeError as e:
                    message = (
                        "The method 'memalloc' in the allocator object must conform to the interface in the "
                        "'BaseCUDAMemoryManager' protocol."
                    )
                    raise TypeError(message) from e

        self.workspace_allocated_size = self.workspace_size
        self.workspace_stream = stream_holder.obj
        self.logger.debug(
            f"Finished allocating device workspace of size {formatters.MemoryStr(self.workspace_size)} in the context "
            f"of stream {self.workspace_stream}."
        )

    def _allocate_workspace_memory_perhaps(self, stream_holder: utils.StreamHolder):
        """
        Allocate workspace memory using the specified allocator, if it hasn't already been
        done.
        """

        if self.workspace_ptr is not None and self.workspace_allocated_size >= self.workspace_size:
            return

        return self._allocate_workspace_memory(stream_holder)

    @utils.precondition(_check_valid_matmul)
    def applicable_algorithm_ids(self, limit=8):
        """Obtain the algorithm IDs that are applicable to this matrix multiplication.

        Args:
            limit: The maximum number of applicable algorithm IDs that is desired

        Returns:
            A sequence of algorithm IDs that are applicable to this matrix multiplication
            problem specification, in random order.
        """
        ...
        algo_ids = cublaslt.matmul_algo_get_ids(
            self.handle,
            self.compute_type,
            self.scale_type,
            self.a_dtype,
            self.b_dtype,
            self.c_dtype,
            self.d_dtype,
            limit,
        )
        return algo_ids

    def _validate_scalar_scale(self, operand: str):
        """
        Validates a scalar scale.
        """
        if self.options.block_scaling:
            raise ValueError(f"A scalar tensor-wide scale factor is not allowed for {operand.upper()} when block_scaling=True.")

    def _validate_tensor_scale(self, scale, operand: str, operand_size=None):
        """
        Validates a tensor scale.

        Args:
            scale: The tensor scale to validate
            operand: The operand name (a, b, c, d)
            operand_size: Size of the operand (needed for block scaling shape validation)
        """
        # Package validation: Normalize "numpy" to "cuda" to match the behavior in __init__.
        # When operands are NumPy on CPU, self.package is set to "cuda" (execution package),
        # so we must also normalize NumPy scales to "cuda" to allow the same input format.
        # This handles the NumPy <=> CuPy asymmetry where NumPy on CPU is accepted as input
        # but internally converted to CuPy for CUDA execution.
        scale_package = utils.infer_object_package(scale)
        scale_package = "cuda" if scale_package == "numpy" else scale_package
        if scale_package != self.package:
            raise TypeError(
                f"The quantization scaling tensor for {operand.upper()} must belong to the same package as the operands."
            )

        # Wrap temporarily since this is needed for validation
        scale_wrapped = tensor_wrapper.wrap_operand(scale)

        # Device/memory space validation
        expected_device_id = "cpu" if self.memory_space == "cpu" else self.device_id
        if expected_device_id != scale_wrapped.device_id:
            raise ValueError(
                f'The scale for {operand.upper()} is on device "{scale_wrapped.device_id}", '
                f'but it should be on device "{expected_device_id}" to match the operands memory space.'
            )

        # Shape and dtype validation for non-block-scaling
        if not self.options.block_scaling:
            if scale_wrapped.shape not in ((1,), ()):
                raise ValueError(
                    f"The scale for {operand.upper()} must be of shape (1,) or (). Got {scale_wrapped.shape} instead."
                )
            if scale_wrapped.dtype != "float32":
                raise ValueError(f"The scale for {operand.upper()} must be float32 type. Got {scale_wrapped.dtype} instead.")

        # Shape and dtype validation for block-scaling
        elif self.input_type_width == 8:
            # Dtype validation (always possible)
            if scale_wrapped.dtype != "uint8":
                raise ValueError(f"Block scales for {operand.upper()} must be uint8 tensor.")

            # Shape validation (only if operand_size is available)
            if operand_size is not None:
                expected_shape = (operand_size // 32,)
                if scale_wrapped.shape != expected_shape:
                    raise ValueError(
                        f"Scales for {operand.upper()} should have shape {expected_shape}. Got {scale_wrapped.shape}."
                    )
        else:
            raise ValueError("block_scaling == True is not supported for non-FP8 types.")

    def _validate_operand_scales(self, quantization_scales, all_required):
        """
        Validates quantization scales, wrapping them into a MatmulQuantizationScales
        object if needed.

        Args:
            quantization_scales: The quantization scales to validate.
            all_required: Whether all scales are required.

        Returns:
            A MatmulQuantizationScales object with the validated quantization scales.
        """
        if quantization_scales is None:
            raise ValueError(
                "Scales are required for narrow-precision (FP8 and lower) operations. "
                "Please set `quantization_scales` argument."
            )

        # wrap the quantization scales into a MatmulQuantizationScales object if needed
        # otherwise, return the quantization scales as is
        quantization_scales = utils.check_or_create_options(
            _configuration.MatmulQuantizationScales, quantization_scales, "Scale factors"
        )

        # Validate which scales are required/allowed
        expected_scales = "AB"
        if self.d_dtype_width <= 8 and not self.options.block_scaling:
            expected_scales += "D"
        elif quantization_scales.d is not None:
            if self.options.block_scaling:
                raise ValueError("Quantization scaling is not supported for D when `block_scaling` option is enabled.")
            if self.d_dtype_width > 8:
                raise ValueError(
                    "Quantization scaling is not supported for D when it is not a narrow-precision (FP8 and lower) type."
                )

        if self.num_operands == 3 and self.c_dtype_width <= 8:
            expected_scales += "C"
        elif quantization_scales.c is not None:
            raise ValueError(
                "Quantization scaling is not supported for C when it is not a narrow-precision (FP8 and lower) type."
            )

        if all_required:
            for operand in expected_scales:
                if getattr(quantization_scales, operand.lower()) is None:
                    raise ValueError(f"Scale for {operand.upper()} is not specified")

        # Validate each scale by delegating to scalar/tensor specific validators
        for operand in "abcd":
            scale = getattr(quantization_scales, operand)
            if scale is None:
                continue

            if isinstance(scale, (int, float)):
                self._validate_scalar_scale(operand)
            else:
                # For block scaling, pass operand size for shape validation
                if self.options.block_scaling and operand in ("a", "b"):
                    operand_idx = 0 if operand == "a" else 1
                    operand_size = self.operands[operand_idx].size  # type: ignore[union-attr,index]
                else:
                    operand_size = None
                self._validate_tensor_scale(scale, operand, operand_size)

        return quantization_scales

    def _validate_epilog_aux_scale(self, aux_quantization_scale, *, required):
        is_fp8_aux = (
            self.preferences.epilog.aux_type is not None
            and typemaps.NAME_TO_DATA_WIDTH[typemaps.DATA_TYPE_TO_NAME[self.preferences.epilog.aux_type]] <= 8
        )
        if aux_quantization_scale is not None and not is_fp8_aux:
            raise ValueError(
                "Scales for epilog auxiliary output are not supported when `preferences.epilog.aux_type` is not set to a "
                "narrow-precision type."
            )
        elif aux_quantization_scale is None and is_fp8_aux and required:
            raise ValueError(
                '"aux_quantization_scale" epilog input is required when `preferences.epilog.aux_type` is not set to a '
                "narrow-precision type."
            )

        # Validate scalar vs tensor scale (same as for operand scales)
        if aux_quantization_scale is not None:
            if isinstance(aux_quantization_scale, (int, float)):
                self._validate_scalar_scale("epilog_aux")
            else:
                # No operand_size for epilog_aux scales
                self._validate_tensor_scale(aux_quantization_scale, "epilog_aux", operand_size=None)

    def _prepare_validated_scalar_scale(self, scale: int | float, operand: str, stream_holder: utils.StreamHolder):
        """
        Converts validated scalar to float32 tensor and copies to GPU.
        Assumes validation already done in _validate_scalar_scale.
        """
        # If it's a scalar, copy to GPU. Float32 is the only type allowed by
        # cublasLtMatmulScale_t for tensor-wide scaling.
        self.logger.debug(f"Scale for {operand.upper()} will be copied to device {self.device_id}.")
        scale_op = tensor_wrapper.wrap_operand(np.asarray([scale], dtype="float32"))
        self.quantization_scales_device[operand] = scale_op.to(self.device_id, stream_holder)

    def _prepare_validated_tensor_scale(self, scale, operand: str, stream_holder: utils.StreamHolder):
        """
        Wraps validated tensor and copies to GPU.
        Assumes all validation already done in _validate_tensor_scale (called in __init__).
        This is pure preparation - no validation here.

        Note: We wrap the tensor a second time here (first wrap was for validation).
        This is acceptable because wrapping is cheap and we get early error detection.
        """
        # Wrap the scale (second time - first was for validation in __init__)
        self.quantization_scales_device[operand] = tensor_wrapper.wrap_operand(scale)

        # Copy to GPU if on CPU (no validation, just preparation)
        if self.quantization_scales_device[operand].device in (None, "cpu"):
            self.logger.debug(f"Scale for {operand.upper()} will be copied to device {self.device_id}.")
            self.quantization_scales_device[operand] = self.quantization_scales_device[operand].to(
                self.device_id, stream_holder
            )

    def _prepare_single_validated_scale(self, scale, operand: str, cublas_operand: str, stream_holder: utils.StreamHolder):
        """
        Prepares a single validated scale and sets its pointer/mode in mm_desc_ifc.
        Used for both operand scales (a,b,c,d) and epilog scales (epilog_aux).
        Assumes validation already done.
        """
        if scale is None:
            return

        # Delegate to specific preparer (validation already done)
        if isinstance(scale, (int, float)):
            self._prepare_validated_scalar_scale(scale, operand, stream_holder)
        else:
            self._prepare_validated_tensor_scale(scale, operand, stream_holder)

        # Set pointer and mode in descriptor
        setattr(self.mm_desc_ifc, f"{cublas_operand}_scale_pointer", self.quantization_scales_device[operand].data_ptr)

        if self.options.block_scaling:
            self.logger.debug(f"Using VEC32_UE8M0 scale mode for operand {operand.upper()}.")
            setattr(self.mm_desc_ifc, f"{cublas_operand}_scale_mode", cublaslt.MatmulMatrixScale.VEC32_UE8M0)
        else:
            self.logger.debug(f"Using SCALAR_32F scale mode for operand {operand.upper()}.")
            setattr(self.mm_desc_ifc, f"{cublas_operand}_scale_mode", cublaslt.MatmulMatrixScale.SCALAR_32F)

    def _prepare_operand_quantization_scales(self, scales, stream_holder: utils.StreamHolder):
        """
        Prepares validated operand scales (a,b,c,d).
        Assumes scales are validated and wrapped into a MatmulQuantizationScales object.
        """
        for operand in "abcd":
            scale = getattr(scales, operand)
            self._prepare_single_validated_scale(scale, operand, cublas_operand=operand, stream_holder=stream_holder)

    @utils.precondition(_check_valid_matmul)
    @utils.atomic(_free_plan_resources, method=True)
    def plan(
        self, *, preferences=None, algorithms=None, epilog=None, epilog_inputs=None, stream: utils.AnyStream | int | None = None
    ):  # Epilog inputs require as many inputs (with specific shapes etc) as required by the epilogue. It's a dict.
        """
        Plan the matrix multiplication operation, considering the epilog (if provided).

        Args:
            preferences: {preferences}

            algorithms: {algorithms}

            epilog: {epilog}

            epilog_inputs: {epilog_inputs}

            stream: {stream}

        Returns:
            A sequence of :class:`nvmath.linalg.advanced.Algorithm` objects that are
            applicable to this matrix multiplication problem specification, heuristically
            ordered from fastest to slowest.

        Notes:
            Epilogs that have ``BIAS`` in their name need an epilog input with the key
            ``'bias'``. Epilogs that have ``DRELU`` need an epilog input with the key
            ``'relu_aux'``, which is produced in a "forward pass" epilog like ``RELU_AUX``
            or ``RELU_AUX_BIAS``. Similarly, epilogs with ``DGELU`` in their name require an
            epilog input with the key ``'gelu_aux'``, produced in the corresponding forward
            pass operation.

        Examples:

            >>> import numpy as np
            >>> import nvmath

            Create two 3-D float64 ndarrays on the CPU representing batched matrices, along
            with a bias vector:

            >>> batch = 32
            >>> M, N, K = 1024, 1024, 1024
            >>> a = np.random.rand(batch, M, K)
            >>> b = np.random.rand(batch, K, N)
            >>> # The bias vector will be broadcast along the columns, as well as along the
            >>> # batch dimension.
            >>> bias = np.random.rand(M)

            We will define a matrix multiplication operation followed by a
            :attr:`nvmath.linalg.advanced.MatmulEpilog.RELU_BIAS` epilog function.

            >>> with nvmath.linalg.advanced.Matmul(a, b) as mm:
            ...     # Plan the operation with RELU_BIAS epilog and corresponding epilog
            ...     # input.
            ...     p = nvmath.linalg.advanced.MatmulPlanPreferences(limit=8)
            ...     epilog = nvmath.linalg.advanced.MatmulEpilog.RELU_BIAS
            ...     epilog_inputs = {{"bias": bias}}
            ...     # The preferences can also be provided as a dict: {{'limit': 8}}
            ...     algorithms = mm.plan(
            ...         preferences=p,
            ...         epilog=epilog,
            ...         epilog_inputs=epilog_inputs,
            ...     )
            ...
            ...     # Execute the matrix multiplication, and obtain the result `r` as a
            ...     # NumPy ndarray.
            ...     r = mm.execute()

            Some epilogs like :attr:`nvmath.linalg.advanced.MatmulEpilog.RELU_AUX` produce
            auxiliary output.

            >>> with nvmath.linalg.advanced.Matmul(a, b) as mm:
            ...     # Plan the operation with RELU_AUX epilog>
            ...     epilog = nvmath.linalg.advanced.MatmulEpilog.RELU_AUX
            ...     algorithms = mm.plan(epilog=epilog)
            ...
            ...     # Execute the matrix multiplication, and obtain the result `r` along
            ...     # with the auxiliary output.
            ...     r, auxiliary = mm.execute()

            The auxiliary output is a Python `dict` with the names of each auxiliary output
            as keys.

        Further examples can be found in the `nvmath/examples/linalg/advanced/matmul
        <https://github.com/NVIDIA/nvmath-python/tree/main/examples/linalg/advanced/matmul>`_
        directory.
        """
        log_info = self.logger.isEnabledFor(logging.INFO)

        self.logger.info("= PLANNING PHASE =")

        # Clear epilog operands, since different epilogs can be provided in different calls.
        # We don't need to worry about ordering, since it's the user's responsibility to
        # order calls that accept a stream argument. This applies to CPU operands as well,
        # even though we move them to the GPU, since the execution is blocking.
        self.epilog_operands = {}  # Clear operands in case of repeated planning.
        self.epilog_input_name_to_handler = {}  # Clear input name to handler map as well,
        self.epilog_inputs_traits = {}  # ... and the input traits as well.

        preferences = utils.check_or_create_options(
            _configuration.MatmulPlanPreferences, preferences, "Matrix multiplication plan preferences"
        )
        self.preferences = preferences

        mm_traits = self.mm_traits

        stream_holder = utils.get_or_create_stream(self.device_id, stream, self.package)
        self.logger.info(f"The specified stream for the matrix multiplication plan is {stream_holder.obj}.")

        # Base FLOP count.
        self.flop_count = 2 * mm_traits.M * mm_traits.N * mm_traits.K
        self.logger.info(f"The base matrix multiplication FLOP count is {formatters.FLOPSStr(self.flop_count, 'FLOP')}.")

        if epilog is None and epilog_inputs is not None:
            self.logger.warning(
                f"Matmul: The provided epilog inputs {epilog_inputs.keys()} are ignored since an epilog is not specified."
            )

        self.epilog = epilog
        epilog_ordering = None
        if epilog is not None:
            assert epilog in EPILOG_INPUT_HANDLERS_MAP, "Not supported."
            self.logger.info(f"The specified epilog is {epilog.name}.")

            epilog_minimum_versions = EPILOG_MINIMUM_VERSIONS_MAP[epilog]
            batched_epilog_minimum_versions = BATCHED_EPILOG_MINIMUM_VERSIONS_MAP[epilog]
            version = cublaslt.get_version()
            if version < epilog_minimum_versions["cublaslt"]:
                message = (
                    f"The epilog {epilog.name} requires cublaslt >= {epilog_minimum_versions['cublaslt']}; "
                    f"you have version {version}. Update to CUDA Toolkit >= {epilog_minimum_versions['ctk']}."
                )
                raise ValueError(message)

            if len(mm_traits.batch_shape) > 0 and version < batched_epilog_minimum_versions["cublaslt"]:
                message = (
                    f"The epilog {epilog.name} supports batching in "
                    f"cublaslt >= {batched_epilog_minimum_versions['cublaslt']}; "
                    f"you have version {version}. Update to CUDA Toolkit >= {epilog_minimum_versions['ctk']}."
                )
                raise ValueError(message)
            if (
                self.mm_traits.c_layout_traits is not None
                and self.mm_traits.c_layout_traits.order == cublaslt.Order.ROW
                and epilog
                in [
                    cublaslt.Epilogue.BGRADA,
                    cublaslt.Epilogue.BGRADB,
                ]
            ):
                msg = f"The epilog {epilog.name} requires input matrix 'c' to be F-contiguous (column-major)."
                raise ValueError(msg)
            if (
                version < 120804
                # A has one row
                and self.mm_traits.M == 1
                # C is broadcast
                and self.mm_traits.c_layout_traits is not None
                and self.mm_traits.c_layout_traits.ld == 0
                # Using both an bias epilog and C
                and self.epilog & _configuration.MatmulEpilog.BIAS > 0
            ):
                message = (
                    "When matrix 'a' has one row, "
                    "simultaneously broadcasting matrix 'c' and using a BIAS epilog requires cublaslt >= 120804; "
                    f"You have version {version}. Update to CUDA Toolkit >= 12.8.1."
                )
                raise ValueError(message)

            # Take a copy of the user-provided inputs.
            if epilog_inputs is not None:
                epilog_inputs = epilog_inputs.copy()
            else:
                epilog_inputs = {}

            # Get the dtype of auxiliary buffer
            aux_dtype_name = (
                typemaps.DATA_TYPE_TO_NAME[self.preferences.epilog.aux_type]  # type: ignore[attr-defined]
                if self.preferences.epilog.aux_type is not None  # type: ignore[attr-defined]
                else None
            )

            # Extract aux quantization scale from the inputs.
            aux_quantization_scale = (
                epilog_inputs.pop("aux_quantization_scale") if "aux_quantization_scale" in epilog_inputs else None
            )
            self._validate_epilog_aux_scale(aux_quantization_scale, required=True)
            self._prepare_single_validated_scale(
                aux_quantization_scale, "epilog_aux", cublas_operand="epilogue_aux", stream_holder=stream_holder
            )

            epilog_input_handler_types = EPILOG_INPUT_HANDLERS_MAP[epilog]
            if epilog_input_handler_types:
                epilog_input_handlers = [
                    handler_type(self.logger, mm_traits, epilog, self.c_dtype_name, self.d_dtype_name, aux_dtype_name)
                    for handler_type in epilog_input_handler_types
                ]

                # Check if the epilog requires a specific result layout, and if the
                # requirement is consistent for all the handlers.
                epilog_input_handlers_ordering = {h.order for h in epilog_input_handlers}
                assert len(epilog_input_handlers_ordering) == 1, "Internal error."
                epilog_ordering = epilog_input_handlers_ordering.pop()

                required_epilog_input_names = {h.name for h in epilog_input_handlers}

                self.logger.info(f"The epilog requires the following additional inputs: {required_epilog_input_names}.")
                if required_epilog_input_names != set(epilog_inputs.keys()):
                    raise ValueError(
                        f"The epilog {epilog.name} requires the following input tensors: "
                        f"{required_epilog_input_names}. The provided tensor names are: {epilog_inputs.keys()}"
                    )

                # Wrap epilog inputs.
                for name in epilog_inputs:
                    epilog_inputs[name] = tensor_wrapper.wrap_operand(epilog_inputs[name])

                # Check if epilog inputs all belong to the same package, which is the same
                # as the package of the MM operands.
                epilog_package = utils.get_operands_package(list(epilog_inputs.values()))
                epilog_package = "cuda" if epilog_package == "numpy" else epilog_package  # Handle the NumPy <=> CuPy asymmetry.
                if self.package != epilog_package:
                    message = f"Library package mismatch for epilog: '{self.package}' => '{epilog_package}'"
                    raise TypeError(message)

                # Check if all epilog inputs all are on the same device, which is the device
                # of the operands.
                device_id = utils.get_operands_device_id(list(epilog_inputs.values()))
                if device_id != "cpu" and self.device_id != device_id:
                    raise ValueError(
                        f"The epilog inputs must be on the same device ({device_id}) as the operands ({self.device_id})."
                    )

                # Move epilog inputs to the GPU, if needed.
                if device_id == "cpu":
                    for e in required_epilog_input_names:
                        self.logger.debug(f"The epilog input {e} will be copied to device{self.device_id}.")
                        self.epilog_operands[e] = epilog_inputs[e].to(self.device_id, stream_holder)
                else:
                    for e in required_epilog_input_names:
                        self.epilog_operands[e] = epilog_inputs[e]

                # First validate all epilog inputs. Use the GPU tensors in case metadata has
                # changed.
                for handler in epilog_input_handlers:
                    handler.validate(epilog_inputs[handler.name])

                # Finally, update the MM descriptor. Note that we pass in
                # self.epilog_operands (which are on the GPU).
                for handler in epilog_input_handlers:
                    handler.update(self.mm_desc_ifc, self.epilog_operands[handler.name])
                    self.epilog_input_name_to_handler[handler.name] = handler

                # Capture the epilog operands traits for consistency checks when resetting
                # operands.
                self.epilog_inputs_traits = {
                    name: EpilogInputTraits(
                        dtype=self.epilog_operands[name].dtype,
                        extents=self.epilog_operands[name].shape,
                        strides=self.epilog_operands[name].strides,
                    )
                    for name in self.epilog_operands
                }

            epilog_output_handler_types = EPILOG_OUTPUT_HANDLERS_MAP[epilog]
            if epilog_output_handler_types:
                self.epilog_output_handlers = epilog_output_handlers = [
                    handler_type(self.logger, mm_traits, epilog, self.c_dtype_name, self.d_dtype_name, aux_dtype_name)
                    for handler_type in epilog_output_handler_types
                ]
                # Check if the epilog requires a specific result layout, and if the
                # requirement is consistent for all the handlers.
                epilog_output_handlers_ordering = {h.order for h in epilog_output_handlers}
                assert len(epilog_output_handlers_ordering) == 1, "Internal error."
                op_epilog_ordering = epilog_output_handlers_ordering.pop()
                if epilog_ordering is None:
                    epilog_ordering = op_epilog_ordering
                else:
                    assert epilog_ordering == op_epilog_ordering, "Internal error."

                # Update the MM descriptor, except for the device pointer.
                for ohandler in epilog_output_handlers:
                    ohandler.update(self.mm_desc_ifc)

            # Set the epilog. At this point, we're sure that the epilog inputs, if any, are
            # valid and have been set.
            self.mm_desc_ifc.epilogue = epilog

        # Fill the result traits, now that we know the epilog.
        self.result_traits = result_traits = get_result_traits(mm_traits, epilog_ordering, self.logger)  # type: ignore[assignment]
        assert self.result_traits is not None, "Internal Error. self.result_traits should have been set by self.plan()."
        self.logger.info(
            f"The layout order for the result D is {self.result_traits.d_layout_traits.order.name}, with LD "
            f"{self.result_traits.d_layout_traits.ld}, and batch offset "
            f"{self.result_traits.d_layout_traits.batch_offset}."
        )

        # Internally transpose operand A if required (conjugate flag) and create layout.
        transpose = False
        if mm_traits.a_layout_traits.is_conjugate and self.is_complex:
            self.mm_desc_ifc.transa = cublas.Operation.C
            transpose = True
            self.logger.debug(
                "To conjugate A, the operand A will be internally transposed and the matrix multiplication will be "
                "performed with OP_C for operand A."
            )
        if self.input_type_width <= 8:
            # narrow-precision (FP8 and lower) data types are only supported for transa=OP_T
            self.mm_desc_ifc.transa = cublas.Operation.T
            transpose = True
            self.logger.debug(
                "For narrow-precision (FP8 and lower) multiplication, the operand A will be internally transposed and the "
                "matrix multiplication will be performed with OP_T for operand A."
            )
        m, n, ld, a_order = mm_traits.a_layout_traits.get_mm_layout(transpose=transpose)
        self.a_layout_ptr = cublaslt.matrix_layout_create(self.a_dtype, rows=m, cols=n, ld=ld)
        self.logger.debug(f"Layout for A: rows = {m}, cols = {n}, ld = {ld}.")

        # Internally transpose operand B if required (conjugate flag, or epilog is BGRADB)
        # and create layout.
        transpose = False
        if mm_traits.b_layout_traits.is_conjugate and self.is_complex:
            self.mm_desc_ifc.transb = cublas.Operation.C
            transpose = True
            self.logger.debug(
                "To conjugate B, the operand B will be internally transposed and the matrix multiplication will be "
                "performed with OP_C for operand B."
            )
        elif epilog == _configuration.MatmulEpilog.BGRADB:
            self.mm_desc_ifc.transb = cublas.Operation.T
            transpose = True
            self.logger.debug(
                "For BGRADB epilog, the operand B will be internally transposed and the matrix multiplication will be "
                "performed with OP_T for operand B."
            )
        m, n, ld, b_order = mm_traits.b_layout_traits.get_mm_layout(transpose=transpose)
        self.b_layout_ptr = cublaslt.matrix_layout_create(self.b_dtype, rows=m, cols=n, ld=ld)
        self.logger.debug(f"Layout for B: rows = {m}, cols = {n}, ld = {ld}.")

        self.d_layout_ptr = cublaslt.matrix_layout_create(
            self.d_dtype, rows=mm_traits.M, cols=mm_traits.N, ld=result_traits.d_layout_traits.ld
        )

        layout_a_ifc = matrix_layout_ifc.MatrixLayoutInterface(self.a_layout_ptr)
        layout_a_ifc.order = a_order
        layout_a_ifc.batch_count = mm_traits.batch_count
        layout_a_ifc.strided_batch_offset = mm_traits.a_layout_traits.batch_offset

        layout_b_ifc = matrix_layout_ifc.MatrixLayoutInterface(self.b_layout_ptr)
        layout_b_ifc.order = b_order
        layout_b_ifc.batch_count = mm_traits.batch_count
        layout_b_ifc.strided_batch_offset = mm_traits.b_layout_traits.batch_offset

        layout_d_ifc = matrix_layout_ifc.MatrixLayoutInterface(self.d_layout_ptr)
        layout_d_ifc.order = result_traits.d_layout_traits.order
        layout_d_ifc.batch_count = mm_traits.batch_count
        layout_d_ifc.strided_batch_offset = result_traits.d_layout_traits.batch_offset

        if self.num_operands == 2:  # By defn, this cannot be inplace.
            if self.c_dtype == self.d_dtype:
                # If C and D have equal types, reuse the layout.
                self.c_layout_ptr = self.d_layout_ptr
            else:
                # Otherwise, create a D-like layout, but with different type.
                self.c_layout_ptr = cublaslt.matrix_layout_create(
                    self.c_dtype, rows=mm_traits.M, cols=mm_traits.N, ld=result_traits.d_layout_traits.ld
                )
                layout_c_ifc = matrix_layout_ifc.MatrixLayoutInterface(self.c_layout_ptr)
                layout_c_ifc.order = result_traits.d_layout_traits.order
                layout_c_ifc.batch_count = mm_traits.batch_count
                layout_c_ifc.strided_batch_offset = result_traits.d_layout_traits.batch_offset
        else:
            # For inplace operation, use the same layout for C and D.
            if self.inplace:
                self.c_layout_ptr = self.d_layout_ptr
            else:
                self.c_layout_ptr = cublaslt.matrix_layout_create(
                    self.c_dtype, rows=mm_traits.M, cols=mm_traits.N, ld=mm_traits.c_layout_traits.ld
                )
                layout_c_ifc = matrix_layout_ifc.MatrixLayoutInterface(self.c_layout_ptr)
                layout_c_ifc.order = mm_traits.c_layout_traits.order
                layout_c_ifc.batch_count = mm_traits.batch_count
                layout_c_ifc.strided_batch_offset = mm_traits.c_layout_traits.batch_offset

        if (
            self.input_type_width == 8
            and self.options.block_scaling
            and (mm_traits.M % 128 != 0 or mm_traits.N % 128 != 0 or mm_traits.K % 128 != 0)
        ):
            raise ValueError(
                f"M={mm_traits.M} N={mm_traits.N} K={mm_traits.K} must be divisible by 128 when block_scaling=True."
            )

        if self.input_type_width == 8 and (mm_traits.M % 16 != 0 or mm_traits.N % 16 != 0 or mm_traits.K % 16 != 0):
            raise ValueError(f"M={mm_traits.M} N={mm_traits.N} K={mm_traits.K} must be divisible by 16 for FP8 operations")

        if self.options.block_scaling and self.d_dtype_width == 8:
            self.mm_desc_ifc.alpha_vector_batch_stride = 1  # Workaround for library caching issue

            # cublasLtMatmulAlgoGetHeuristic requires the scale pointer to be set.
            self.aux_outputs = {
                "d_out_scale": utils.create_empty_tensor(
                    self.result_class,
                    ((mm_traits.M * mm_traits.N) // 32 * self.mm_traits.batch_count),
                    "uint8",
                    self.device_id,
                    stream_holder,
                    verify_strides=False,
                )
            }
            self.mm_desc_ifc.d_out_scale_pointer = self.aux_outputs["d_out_scale"].data_ptr
            self.mm_desc_ifc.d_out_scale_mode = cublaslt.MatmulMatrixScale.VEC32_UE8M0

        limit = preferences.limit
        if algorithms is None:
            num_algorithms = np.empty((1,), dtype=np.int32)
            self.algorithms_buffer = cublaslt.MatmulHeuristicResult(limit)
        else:
            assert all(isinstance(algo, _algorithmmod.Algorithm) for algo in algorithms), (
                "The algorithms passed to plan() are of wrong type."
            )
            num_algorithms = len(algorithms)

        if self.preference_ptr is None:
            self.preference_ptr = cublaslt.matmul_preference_create()
        else:
            # We need to create a new preferences object to avoid preferences being set in a
            # cumulative manner if plan() is called multiple times.
            cublaslt.matmul_preference_destroy(self.preference_ptr)
            self.preference_ptr = cublaslt.matmul_preference_create()

        if self.input_type_width <= 8:
            self._prepare_operand_quantization_scales(self.quantization_scales, stream_holder)

        if algorithms is None:
            # Set preferences.
            preference_ifc = matmul_pref_ifc.MatmulPreferenceInterface(self.preference_ptr)
            preference_ifc.max_workspace_bytes = self.memory_limit
            preference_ifc.reduction_scheme_mask = preferences.reduction_scheme_mask
            preference_ifc.max_waves_count = preferences.max_waves_count
            preference_ifc.impl_mask = preferences.numerical_impl_mask

            # Set minimum alignments.
            a_ptr, b_ptr = self.operands[0].data_ptr, self.operands[1].data_ptr
            preference_ifc.min_alignment_a_bytes = min(256, pointer_aligned_to(a_ptr))
            preference_ifc.min_alignment_b_bytes = min(256, pointer_aligned_to(b_ptr))
            self.logger.debug(f"The minimum alignment for operand A is {preference_ifc.min_alignment_a_bytes} bytes.")
            self.logger.debug(f"The minimum alignment for operand B is {preference_ifc.min_alignment_b_bytes} bytes.")
            if self.num_operands == 3:
                c_ptr = self.operands[2].data_ptr
                preference_ifc.min_alignment_c_bytes = min(256, pointer_aligned_to(c_ptr))
                self.logger.debug(f"The minimum alignment for operand C is {preference_ifc.min_alignment_c_bytes} bytes.")
            # The result alignment should be 256 bytes.
            self.logger.debug("The minimum alignment for the result D is the default 256 bytes.")

            self.logger.info("Starting matrix multiplication planning...")
            assert isinstance(self.device_id, int), self.device_id
            assert stream_holder is not None
            with utils.cuda_call_ctx(stream_holder, blocking=True, timing=log_info) as (
                _,
                elapsed,
            ):
                cublaslt.matmul_algo_get_heuristic(
                    self.handle,
                    self.mm_desc,
                    self.a_layout_ptr,
                    self.b_layout_ptr,
                    self.c_layout_ptr,
                    self.d_layout_ptr,
                    self.preference_ptr,
                    limit,
                    self.algorithms_buffer.ptr,
                    num_algorithms.ctypes.data,
                )

            num_algorithms = num_algorithms[0]
            if num_algorithms == 0:
                raise RuntimeError("Planning failed to find any suitable algorithm.")
            assert self.algorithms_buffer is not None, (
                "Internal Error. self.algorithms_buffer should have been set by self.plan()."
            )
            self.algorithms_buffer = self.algorithms_buffer[:num_algorithms]

            # Create algorithm objects.
            self.algorithm_objects = tuple(_algorithmmod.Algorithm(a) for a in self.algorithms_buffer)
        else:
            self.algorithm_objects = tuple(algorithms)
            self.algorithms_buffer = cublaslt.MatmulHeuristicResult(len(algorithms))
            for i, algo in enumerate(algorithms):
                # we wrap it too well that it's hard to copy-construct...
                self.algorithms_buffer[i] = algo.algorithm._data

        # Cache the first (best) algorithm struct.
        self.cached_best_algorithm_struct = self.algorithms_buffer[0]["algo"]

        # Create the map from object to buffer.
        self.algorithm_object_to_buffer = dict(zip(self.algorithm_objects, self.algorithms_buffer, strict=True))

        self.workspace_size = int(np.max(self.algorithms_buffer["workspace_size"]))
        if self.workspace_size > 0 and self.epilog:
            self.workspace_size += 16  # Workaround for library issue

        if algorithms is None:
            self.logger.info(
                f"The plan found {num_algorithms} suitable algorithms within the requested limit of {limit} "
                f"algorithms, with a workspace requirement of {formatters.MemoryStr(self.workspace_size)}."
            )
        else:
            self.logger.info(
                f"The plan is using {num_algorithms} algorithm passed through the algorithms argument, with a "
                f"workspace requirement of {formatters.MemoryStr(self.workspace_size)}."
            )

        self.mm_planned = True
        if algorithms is None and elapsed.data is not None:
            self.logger.info(f"The matrix multiplication planning phase took {elapsed.data:.3f} ms to complete.")

        return self.algorithm_objects

    @property
    def algorithms(self):
        """
        After planning using :meth:`plan()`, get the sequence of algorithm objects to
        inquire their capabilities, configure them, or serialize them for later use.

        Returns:
            A sequence of :class:`nvmath.linalg.advanced.Algorithm` objects that are
            applicable to this matrix multiplication problem specification.
        """
        return self.algorithm_objects

    def _check_and_set_operand(
        self,
        operand,
        operand_name,
        mm_desc_ifc,
        stream_holder,
        *,
        operand_index=None,
        epilog_name=None,
        package=None,
        dtype=None,
        extents=None,
        strides=None,
    ):
        """
        Check to make sure that the provided operand is consistent with the one it's
        updating, and update it.
        """
        assert (operand_index is None) ^ (epilog_name is None), "Internal Error."
        assert self.operands is not None, "Internal Error."

        # Make sure that the data type and extents match.
        utils.check_attribute_match(dtype, operand.dtype, "data type")
        utils.check_attribute_match(extents, operand.shape, "extents")

        package = utils.infer_object_package(operand.tensor)

        # Conjugate flag of the provided operands must match the original qualifiers
        if operand_index is not None and self.lazy_conjugation[operand_index] != operand.is_conjugate:
            raise ValueError(f"The provided operand {operand_name} has different conjugate flag than the original operand")

        device_id = operand.device_id
        # When memory_space is "cpu", operands should be on CPU even though
        # self.device_id is the execution device.
        expected_device_id = "cpu" if self.memory_space == "cpu" else self.device_id
        if expected_device_id != device_id:
            raise ValueError(
                f'The operand {operand_name} is on device "{device_id}", but it should be on device '
                f'"{expected_device_id}" to match the original operand.'
            )

        if device_id == "cpu":
            package = "cuda" if package == "numpy" else package  # Handle the NumPy <=> CuPy asymmetry.
            if self.package != package:
                message = f"Library package mismatch: '{self.package}' => '{package}'"
                raise TypeError(message)

            # Check if we have a GPU buffer to update into.
            if operand_index is not None:
                o = self.operands[operand_index]
            else:
                o = self.epilog_operands[epilog_name]
            if o is None:  # No buffer, create one.
                # Copy operand across memory spaces (CPU to GPU).
                o = operand.to(self.device_id, stream_holder)
                if operand_index is not None:
                    self.operands[operand_index] = o
                else:
                    self.epilog_operands[epilog_name] = o
                    # Update the epilog pointer, since we're starting afresh.
                    self.epilog_input_name_to_handler[epilog_name].update(mm_desc_ifc, o)
            else:
                # In-place copy to existing device pointer because the new operand is on the
                # CPU.
                o.copy_(operand, stream_holder=stream_holder)
            if self.memory_space == "cpu" and operand_index == 2:
                # Hold references, needed for inplace operations.
                self.cpu_c_ref = operand
        else:
            if self.package != package:
                message = f"Library package mismatch: '{self.package}' => '{package}'"
                raise TypeError(message)

            utils.check_attribute_match(strides, operand.strides, "strides")

            # Finally, replace the original operand by the new one.
            if operand_index is not None:
                self.operands[operand_index] = operand
            else:
                self.epilog_operands[epilog_name] = operand
                # Update the epilog pointer, since we're starting afresh.
                self.epilog_input_name_to_handler[epilog_name].update(mm_desc_ifc, operand)

        self.logger.info(f"Operand '{operand_name}' has been reset to the new value.")

        return

    @utils.precondition(_check_valid_matmul)
    def reset_operands(
        self,
        a=None,
        b=None,
        c=None,
        *,
        alpha=None,
        beta=None,
        quantization_scales=None,
        epilog_inputs=None,
        stream: utils.AnyStream | int | None = None,
    ):
        """
        Reset the operands held by this :class:`Matmul` instance.

        This method has two use cases:
            (1) it can be used to provide new operands for execution when the original
                operands are on the CPU
            (2) it can be used to release the internal reference to the previous operands
                and make their memory available for other use by passing ``None`` for *all*
                arguments. In this case, this method must be called again to provide the
                desired operands before another call to execution APIs like :meth:`autotune`
                or :meth:`execute`.

        This method is not needed when the operands reside on the GPU and in-place
        operations are used to update the operand values.

        This method will perform various checks on the new operands to make sure:

            - The shapes, strides, datatypes match those of the old ones.
            - The packages that the operands belong to match those of the old ones.
            - If input tensors are on GPU, the device must match.

        Args:
            a: {a}

            b: {b}

            c: {c}
                {c_admonitions}

            alpha: {alpha}

            beta: {beta}

            epilog_inputs: {epilog_inputs}

            stream: {stream}

            quantization_scales: {quantization_scales}

        Examples:

            >>> import cupy as cp
            >>> import nvmath

            Create two 3-D float64 ndarrays on the GPU:

            >>> M, N, K = 128, 128, 256
            >>> a = cp.random.rand(M, K)
            >>> b = cp.random.rand(K, N)

            Create an matrix multiplication object as a context manager

            >>> with nvmath.linalg.advanced.Matmul(a, b) as mm:
            ...     # Plan the operation.
            ...     algorithms = mm.plan()
            ...
            ...     # Execute the MM to get the first result.
            ...     r1 = mm.execute()
            ...
            ...     # Reset the operands to new CuPy ndarrays.
            ...     c = cp.random.rand(M, K)
            ...     d = cp.random.rand(K, N)
            ...     mm.reset_operands(c, d)
            ...
            ...     # Execute to get the new result corresponding to the updated operands.
            ...     r2 = mm.execute()

            Note that if only a subset of operands are reset, the operands that are not
            reset hold their original values.

            With :meth:`reset_operands`, minimal overhead is achieved as problem
            specification and planning are only performed once.

            For the particular example above, explicitly calling :meth:`reset_operands` is
            equivalent to updating the operands in-place, i.e, replacing
            ``mm.reset_operand(c, d)`` with ``a[:]=c`` and ``b[:]=d``. Note that updating
            the operand in-place should be adopted with caution as it can only yield the
            expected result under the additional constraint below:

                - The operand is on the GPU (more precisely, the operand memory space should
                  be accessible from the execution space).

            For more details, please refer to `inplace update example
            <https://github.com/NVIDIA/nvmath-python/tree/main/examples/linalg/advanced/matmul/example05_stateful_inplace.py>`_.
        """

        if c is not None and self.num_operands == 2:
            raise ValueError(
                "The matrix multiplication problem specification does not include operand C, so it cannot be reset."
            )

        if a is None and b is None and c is None and epilog_inputs is None and alpha is None and beta is None:
            self.cpu_c_ref = self.operands = None
            self.epilog_operands = {}
            self.logger.info("The operands have been reset to None.")
            return

        # If the operands have been reset to None, then all required operands (a, b, c, and
        # epilog_inputs need to be provided).
        if self.operands is None:
            if a is None or b is None or (c is None and self.num_operands == 3):
                op_names = "A, B"
                if c is None and self.num_operands == 3:
                    op_names += ", C"
                raise ValueError(f"Operands {op_names} must be provided.")
            epilog_names = self.epilog_inputs_traits.keys()
            if epilog_inputs is None:
                if epilog_names:
                    raise ValueError(f"The epilog inputs {epilog_names} must be provided.")
            else:
                # Check that all required epilog inputs names are provided.
                if epilog_names != epilog_inputs.keys():
                    raise ValueError(
                        f"The epilog inputs {epilog_names} are required. The provided epilog input names are "
                        f"{epilog_inputs.keys()}."
                    )
            self.operands = [None] * self.num_operands
            self.epilog_operands = dict.fromkeys(epilog_names)

        # Future operations on the workspace stream should be ordered after the computation.
        if self.last_compute_event is not None:
            self.workspace_stream.wait(self.last_compute_event)
            self.last_compute_event = None

        # Update alpha.
        if alpha is not None:
            try:
                self.alpha[0] = alpha
            except (ValueError, TypeError) as e:
                raise ValueError(
                    f"The value provided for alpha {alpha} is not convertible to dtype '{self.alpha.dtype}'."
                ) from e

        # Update beta.
        if beta is not None:
            if self.num_operands == 2:
                self.logger.warning(f"Matmul: The provided beta value {beta} is ignored since operand C is not specified.")
            else:
                try:
                    self.beta[0] = beta
                except (ValueError, TypeError) as e:
                    raise ValueError(
                        f"The value provided for beta {beta} is not convertible to dtype '{self.beta.dtype}'."
                    ) from e

        stream_holder = utils.get_or_create_stream(self.device_id, stream, self.package)

        # Update quantization_scales.
        if quantization_scales is not None:
            quantization_scales = self._validate_operand_scales(quantization_scales, all_required=False)
            if quantization_scales.a is not None:
                self.quantization_scales.a = quantization_scales.a
            if quantization_scales.b is not None:
                self.quantization_scales.b = quantization_scales.b
            if quantization_scales.c is not None:
                self.quantization_scales.c = quantization_scales.c
            if quantization_scales.d is not None:
                self.quantization_scales.d = quantization_scales.d
            self._prepare_operand_quantization_scales(self.quantization_scales, stream_holder)

        if epilog_inputs is not None and "aux_quantization_scale" in epilog_inputs:
            epilog_inputs = epilog_inputs.copy()
            aux_quantization_scale = epilog_inputs.pop("aux_quantization_scale")
            self._validate_epilog_aux_scale(aux_quantization_scale, required=False)
            self._prepare_single_validated_scale(
                aux_quantization_scale, "epilog_aux", cublas_operand="epilogue_aux", stream_holder=stream_holder
            )

        # Reset the provided operands.
        if a is not None:
            a = tensor_wrapper.wrap_operand(a)
            index = 0
            self._check_and_set_operand(
                a,
                "A",
                self.mm_desc_ifc,
                stream_holder,
                operand_index=index,
                dtype=self.a_dtype_name,
                extents=self.operand_extents[index],
                strides=self.operand_strides[index],
            )

        if b is not None:
            b = tensor_wrapper.wrap_operand(b)
            index = 1
            self._check_and_set_operand(
                b,
                "B",
                self.mm_desc_ifc,
                stream_holder,
                operand_index=index,
                dtype=self.b_dtype_name,
                extents=self.operand_extents[index],
                strides=self.operand_strides[index],
            )

        if c is not None:  # If we get here, we know that C is one of the operands in the problem specification.
            c = tensor_wrapper.wrap_operand(c)
            index = 2
            self._check_and_set_operand(
                c,
                "C",
                self.mm_desc_ifc,
                stream_holder,
                operand_index=index,
                dtype=self.c_dtype_name,
                extents=self.operand_extents[index],
                strides=self.operand_strides[index],
            )

        # Reset the provided epilog inputs.
        if epilog_inputs is not None:
            for name in epilog_inputs:
                epilog_input = tensor_wrapper.wrap_operand(epilog_inputs[name])
                self._check_and_set_operand(
                    epilog_input,
                    name,
                    self.mm_desc_ifc,
                    stream_holder,
                    epilog_name=name,
                    dtype=self.epilog_inputs_traits[name].dtype,
                    extents=self.epilog_inputs_traits[name].extents,
                    strides=self.epilog_inputs_traits[name].strides,
                )

    @utils.precondition(_check_valid_matmul)
    @utils.precondition(_check_planned, "Autotuning")
    @utils.precondition(_check_valid_operands, "Autotuning")
    @utils.atomic(_release_workspace_memory_perhaps_wrapper, method=True)
    def autotune(
        self, iterations=10, prune=None, release_workspace=False, stream: utils.AnyStream | int | None = None
    ):  # Prune means keep top N of the algorithms only.
        """
        Autotune the matrix multiplication to order the algorithms from the fastest measured
        execution time to the slowest. Once autotuned, the optimally-ordered algorithm
        sequence can be accessed using :py:attr:`algorithms`.

        .. note::
            This function will benchmark each of the algorithms and order the algorithms
            based on the benchmark results. The measurements can be impacted by factors
            such as GPU temperature, clock settings, or power consumption. Autotuning
            in an unstable environment can result in a suboptimal algorithm ordering.
            If you experience performance problems, consider omitting the autotuning.

        Args:
            iterations: The number of autotuning iterations to perform.

            prune: An integer N, specifying the top N fastest algorithms to retain after
                autotuning. The default is to retain all algorithms.

            release_workspace: {release_workspace}

            stream: {stream}
        """
        self.logger.info("= AUTOTUNING PHASE =")
        # Measure time taken for autotuning.
        from timeit import default_timer as timer

        self.logger.info("Starting autotuning...")
        start = timer()

        assert self.algorithm_objects is not None, "Internal Error. self.algorithm_objects should have been set by self.plan()."
        num_algorithms = len(self.algorithm_objects)
        if num_algorithms == 1:
            self.logger.info("Skipping the autotuning, because only one algorithm has been returned in the planning phase.")
            return
        limit = min(prune, num_algorithms) if prune is not None else num_algorithms
        self.logger.info(
            f"The number of algorithms in the plan is {num_algorithms}, from which the top {limit} will be retained."
        )
        self.logger.info(f"The requested number of iterations is {iterations}.")

        # Autotune setup.
        stream_holder = utils.get_or_create_stream(self.device_id, stream, self.package)

        # Allocate workspace if needed.
        self._allocate_workspace_memory_perhaps(stream_holder)

        # Create empty tensors for auxiliary output.
        epilog_outputs = {}
        for handler in self.epilog_output_handlers:
            name = handler.name
            shape, strides, dtype_name = handler.attributes()
            epilog_outputs[name] = aux = utils.create_empty_tensor(
                self.result_class,
                shape,
                dtype_name,
                self.device_id,
                stream_holder,
                verify_strides=False,
                strides=strides,
            )

            # Update the data pointer in the MM descriptor.
            handler.update_ptr(self.mm_desc_ifc, aux.data_ptr)

        # Take a copy of `c` for autotuning to avoid clobbering it if inplace=True.
        c = None
        if self.inplace:
            # `c` is guaranteed to exist. Clone `c` and copy into it, since `tensor.to`
            # returns the original tensor when the device ID is the same.
            c = self.operands[2]
            c = utils.create_empty_tensor(
                c.__class__, c.shape, c.dtype, self.device_id, stream_holder, verify_strides=False, strides=c.strides
            )
            c.copy_(self.operands[2], stream_holder)
            assert c.data_ptr != self.operands[2].data_ptr, "Internal error."
        elif self.num_operands == 3:
            c = self.operands[2]

        # Create empty tensor for the result, if the operation is not in-place.
        assert self.result_traits is not None, "Internal Error. self.result_traits should have been set by self.plan()."
        if self.inplace:
            result = c
        else:
            result = utils.create_empty_tensor(
                self.result_class,
                self.result_traits.result_shape,
                self.d_dtype_name,
                self.device_id,
                stream_holder,
                verify_strides=False,
                strides=self.result_traits.result_strides,
            )
        result_ptr = result.data_ptr

        a, b = self.operands[0], self.operands[1]
        raw_workspace_ptr = utils.get_ptr_from_memory_pointer(self.workspace_ptr)
        alpha_ptr, a_ptr, b_ptr, beta_ptr = self.alpha.ctypes.data, a.data_ptr, b.data_ptr, self.beta.ctypes.data
        # If `c` is not provided, set c_ptr to nullptr.
        if self.num_operands == 3:
            c_ptr = c.data_ptr
        else:
            assert self.beta[0] == 0.0, "Internal error: beta must be zero if `c` is not specified."
            c_ptr = 0

        def execute_matmul(algorithm_ptr):
            cublaslt.matmul(
                self.handle,
                self.mm_desc,
                alpha_ptr,
                a_ptr,
                self.a_layout_ptr,
                b_ptr,
                self.b_layout_ptr,
                beta_ptr,
                c_ptr,
                self.c_layout_ptr,
                result_ptr,
                self.d_layout_ptr,
                algorithm_ptr,
                raw_workspace_ptr,
                self.workspace_size,
                stream_holder.ptr,
            )

        def flush_cache():
            """
            Write data to a temporary buffer to flush the L2 cache.
            """

            @functools.cache
            def get_l2_cache_size(device_id):
                device = Device(device_id)
                return device.properties.l2_cache_size

            l2_cache_size = get_l2_cache_size(self.device_id)
            cpu_buffer = np.zeros(l2_cache_size, dtype=np.uint8)
            tensor_wrapper.wrap_operand(cpu_buffer).to(device_id=self.device_id, stream_holder=stream_holder)

        # Tune.
        with utils.cuda_call_ctx(stream_holder, blocking=False, timing=False) as (
            self.last_compute_event,
            _,
        ):
            gpu_times = np.empty(shape=(len(self.algorithms_buffer), iterations), dtype=float)
            algorithm_idxs = list(range(len(self.algorithms_buffer)))
            timing_enabled_options = EventOptions(enable_timing=True)
            start0 = stream_holder.obj.device.create_event(options=timing_enabled_options)
            end0 = stream_holder.obj.device.create_event(options=timing_enabled_options)
            for i in range(iterations):
                random.shuffle(algorithm_idxs)
                for algorithm_idx in algorithm_idxs:
                    algorithm_ptr = self.algorithms_buffer[algorithm_idx]["algo"].ctypes.data
                    stream_holder.obj.record(start0)
                    execute_matmul(algorithm_ptr=algorithm_ptr)
                    stream_holder.obj.record(end0)
                    # FIXME: @dching Calling sync() here could slow down tuning by forcing
                    # the device to wait for the host to record the elapsed time. It may be
                    # faster to store references to 2 * len(iterations) *
                    # len(algorithms_buffer) Events and compute the elapsed time at the end.
                    end0.sync()
                    gpu_times[algorithm_idx, i] = end0 - start0
                    # Avoid potential overflow for inplace operations.
                    if self.inplace:
                        c.copy_(self.operands[2], stream_holder)
                    flush_cache()

        gpu_times = np.median(gpu_times, axis=1)

        # Establish ordering wrt the computation and free workspace if requested.
        self._release_workspace_memory_perhaps(release_workspace=release_workspace)
        self._reset_workspace_allocation_tracking()

        # Get the sort order based on the GPU times.
        sorted_gpu_times, sort_order = zip(*sorted(zip(gpu_times, range(num_algorithms), strict=True)), strict=True)

        # Reorder the algorithms buffer and algorithm objects according to the sort order,
        # and prune it.
        sort_order = sort_order[:limit]
        self.algorithms_buffer = self.algorithms_buffer[list(sort_order)]
        self.algorithm_objects = tuple(self.algorithm_objects[i] for i in sort_order)

        # Update cached first (best) algorithm struct after tuning.
        self.cached_best_algorithm_struct = self.algorithms_buffer[0]["algo"]

        # Create the map from object to buffer.
        self.algorithm_object_to_buffer = dict(zip(self.algorithm_objects, self.algorithms_buffer, strict=True))

        gpu_times_str = ", ".join(f"{t:0.3f}" for t in gpu_times)
        self.logger.info(f"The autotuned GPU times (in milliseconds) are: {gpu_times_str}.")
        self.logger.info(f"The corresponding sort order is: {sort_order}.")
        orig_flop_rate = self.flop_count / gpu_times[0] * 1000
        if sort_order[0] != 0:
            self.logger.info(
                f"Autotuning found that the algorithm originally ranked {sort_order[0]} is the best out of the "
                f"{num_algorithms} in the plan, and moved it to rank 0."
            )
            new_flop_rate = self.flop_count / sorted_gpu_times[0] * 1000
            self.logger.info(
                f"Autotuning has improved performance from {formatters.FLOPSStr(orig_flop_rate, 'FLOP/s')} to "
                f"{formatters.FLOPSStr(new_flop_rate, 'FLOP/s')}."
            )
        else:
            self.logger.info(
                f"Autotuning found that the algorithm ranked best by the plan heuristics remains the best out of the "
                f"{num_algorithms} algorithms in the plan."
            )
            self.logger.info(f"The best performance remains at {formatters.FLOPSStr(orig_flop_rate, 'FLOP/s')}.")

        end = timer()
        self.logger.info(f"The autotuning took {(end - start) * 1000.0:.3f} ms to complete.")

    @utils.precondition(_check_valid_matmul)
    @utils.precondition(_check_planned, "Execution")
    @utils.precondition(_check_valid_operands, "Execution")
    @utils.atomic(_release_workspace_memory_perhaps_wrapper, method=True)
    def execute(self, *, algorithm=None, release_workspace=False, stream: utils.AnyStream | int | None = None):
        """
        Execute a prepared (planned and possibly autotuned) matrix multiplication.

        Args:
            algorithm: (Experimental) An algorithm chosen from the sequence returned by
                :meth:`plan` or :py:attr:`algorithms`. By default, the first algorithm in
                the sequence is used.

            release_workspace: {release_workspace}

            stream: {stream}

        Returns:
           {result}
        """
        log_info = self.logger.isEnabledFor(logging.INFO)
        log_debug = self.logger.isEnabledFor(logging.DEBUG)

        if log_info:
            self.logger.info("= EXECUTION PHASE =")
        stream_holder = utils.get_or_create_stream(self.device_id, stream, self.package)
        if log_info:
            self.logger.info(f"The specified stream for execute() is {stream_holder.obj}.")

        # Allocate workspace if needed.
        self._allocate_workspace_memory_perhaps(stream_holder)

        # Create empty tensors for auxiliary output.
        for handler in self.epilog_output_handlers:
            name = handler.name
            shape, strides, dtype_name = handler.attributes()
            if log_debug:
                self.logger.debug(f"Beginning auxiliary output tensor '{name}' creation...")
                self.logger.debug(f"The '{name}' tensor shape = {shape} with strides = {strides} and data type '{dtype_name}'.")
            self.epilog_outputs[name] = aux = utils.create_empty_tensor(
                self.result_class,
                shape,
                dtype_name,
                self.device_id,
                stream_holder,
                verify_strides=False,
                strides=strides,
            )
            if log_debug:
                self.logger.debug(f"The auxiliary output tensor '{name}' has been created.")
            if self.preferences.epilog.aux_amax:  # type: ignore[attr-defined]
                if "float8" not in dtype_name:
                    raise ValueError("epilog.aux_amax=True is not supported when epilog output type is not FP8.")
                self.epilog_outputs[f"{name}_amax"] = utils.create_empty_tensor(
                    self.result_class,
                    (1,),
                    "float32",  # This is the only type allowed by cuBLAS for AMAX.
                    self.device_id,
                    stream_holder,
                    verify_strides=False,
                )
                self.mm_desc_ifc.epilogue_aux_amax_pointer = self.epilog_outputs[f"{name}_amax"].data_ptr

            # Update the data pointer in the MM descriptor.
            handler.update_ptr(self.mm_desc_ifc, aux.data_ptr)

        # Create empty tensor for the result, if the operation is not in-place.
        assert self.result_traits is not None, "Internal Error. self.result_traits should have been set by self.plan()"
        if self.inplace:
            if log_debug:
                self.logger.debug("The operation is in-place (operand C will be overwritten).")
            self.result = self.operands[2]
        else:
            if log_debug:
                self.logger.debug("Beginning output (empty) tensor creation...")
                self.logger.debug(
                    f"The output tensor shape = {self.result_traits.result_shape} with strides = "
                    f"{self.result_traits.result_strides} and data type '{self.d_dtype_name}'."
                )
            self.result = utils.create_empty_tensor(
                self.result_class,
                self.result_traits.result_shape,
                self.d_dtype_name,
                self.device_id,
                stream_holder,
                verify_strides=False,
                strides=self.result_traits.result_strides,
            )
            if log_debug:
                self.logger.debug("The output (empty) tensor has been created.")

        self.aux_outputs = {}

        if self.options.result_amax:
            self.aux_outputs["result_amax"] = utils.create_empty_tensor(
                self.result_class,
                (1,),
                "float32",  # This is the only type allowed by cuBLAS for AMAX.
                self.device_id,
                stream_holder,
                verify_strides=False,
            )
            self.mm_desc_ifc.amax_d_pointer = self.aux_outputs["result_amax"].data_ptr

        if self.options.block_scaling and self.d_dtype_width == 8:
            self.aux_outputs["d_out_scale"] = utils.create_empty_tensor(
                self.result_class,
                (self.mm_traits.batch_count * self.result_traits.result_shape[-1] * self.result_traits.result_shape[-2] // 32),
                "uint8",
                self.device_id,
                stream_holder,
                verify_strides=False,
            )
            self.mm_desc_ifc.d_out_scale_pointer = self.aux_outputs["d_out_scale"].data_ptr

        # Select the first (best) algorithm if one is not provided.
        if algorithm is None:
            algorithm_struct = self.cached_best_algorithm_struct
            if log_info:
                self.logger.info(
                    "The highest ranked algorithm in the plan (algorithm id = "
                    f"{self.algorithm_objects[0].algorithm_id}) will be used."
                )
        else:
            if algorithm not in self.algorithm_objects:
                raise ValueError("Algorithm passed to execute() has to be included in the plan() algorithms")
            algorithm_struct = self.algorithm_object_to_buffer[algorithm]["algo"]
            if log_info:
                self.logger.info(f"The specified algorithm (algorithm id = {algorithm.algorithm_id}) will be used.")

        a, b = self.operands[0], self.operands[1]

        # If `c` is not provided, set c_ptr to nullptr.
        if self.num_operands == 3:
            c_ptr = self.operands[2].data_ptr
        else:
            assert self.beta[0] == 0.0, "Internal error: beta must be zero if `c` is not specified."
            c_ptr = 0

        raw_workspace_ptr = utils.get_ptr_from_memory_pointer(self.workspace_ptr)
        if log_info:
            self.logger.info("Starting matrix multiplication...")
            self.logger.info(f"{self.call_prologue}")
        with utils.cuda_call_ctx(stream_holder, self.blocking, timing=log_info) as (
            self.last_compute_event,
            elapsed,
        ):
            cublaslt.matmul(
                self.handle,
                self.mm_desc,
                self.alpha.ctypes.data,
                a.data_ptr,
                self.a_layout_ptr,
                b.data_ptr,
                self.b_layout_ptr,
                self.beta.ctypes.data,
                c_ptr,
                self.c_layout_ptr,
                self.result.data_ptr,
                self.d_layout_ptr,
                algorithm_struct.ctypes.data,
                raw_workspace_ptr,
                self.workspace_size,
                stream_holder.ptr,
            )

        if log_info and elapsed.data is not None:
            self.logger.info(f"The matrix multiplication calculation took {elapsed.data:.3f} ms to complete.")

        # Establish ordering wrt the computation and free workspace if requested.
        if release_workspace:
            self._release_workspace_memory_perhaps(True)

        # Return the result and auxiliary outputs, if present.
        all_outputs = self.epilog_outputs | self.aux_outputs
        if self.memory_space == "cpu":
            if self.inplace:
                # Overwrite operand C.
                assert self.cpu_c_ref is not None, "Internal error."
                c = self.cpu_c_ref
                c.copy_(self.result, stream_holder=stream_holder)
                out = c.tensor
            else:
                out = self.result.to("cpu", stream_holder=stream_holder).tensor
            # Copy auxiliary output to CPU.
            aux = {name: all_outputs[name].to("cpu", stream_holder=stream_holder).tensor for name in all_outputs}
        else:
            out = self.result.tensor
            # Return the unwrapped epilog output tensor(s).
            aux = {name: all_outputs[name].tensor for name in all_outputs}

        # Release internal reference to the result to permit recycling of memory.
        self.result = None
        self.aux_outputs = {}
        self.epilog_outputs = {}
        self._reset_workspace_allocation_tracking()

        if aux:
            return out, aux

        return out

    def free(self):
        """Free Matmul resources.

        It is recommended that the :class:`Matmul` object be used within a context, but if
        it is not possible then this method must be called explicitly to ensure that the
        matrix multiplication resources (especially internal library objects) are properly
        cleaned up.
        """

        if not self.valid_state:
            return

        try:
            # Future operations on the workspace stream should be ordered after the
            # computation.
            if self.last_compute_event is not None:
                self.workspace_stream.wait(self.last_compute_event)
                self.last_compute_event = None

            self._free_workspace_memory()

            self._free_plan_resources()

            # We won't destroy the handle.

            # Release references to operands and results to ensure proper reference counting
            self.operands = None
            self.cpu_c_ref = None
            self.epilog_operands = {}
            self.result = None

        except Exception as e:
            self.logger.critical("Internal error: only part of the Matmul object's resources have been released.")
            self.logger.critical(str(e))
            raise e
        finally:
            self.valid_state = False

        self.logger.info("The Matmul object's resources have been released.")


@utils.docstring_decorator(SHARED_MM_DOCUMENTATION, skip_missing=False)
def matmul(
    a,
    b,
    /,
    c=None,
    *,
    alpha=None,
    beta=None,
    epilog=None,
    epilog_inputs=None,
    qualifiers=None,
    quantization_scales=None,
    options=None,
    preferences=None,
    algorithm=None,
    stream: utils.AnyStream | int | None = None,
):
    """
    Perform the specified matrix multiplication computation :math:`F(\\alpha a @ b + \\beta
    c)`, where :math:`F` is the epilog. This function-form is a wrapper around the stateful
    :class:`Matmul` object APIs and is meant for *single* use (the user needs to perform
    just one matrix multiplication, for example), in which case there is no possibility of
    amortizing preparatory costs.

    Detailed information on what's happening within this function can be obtained by passing
    in a :class:`logging.Logger` object to :class:`MatmulOptions` or by setting the
    appropriate options in the root logger object, which is used by default:

        >>> import logging
        >>> logging.basicConfig(
        ...     level=logging.INFO,
        ...     format="%(asctime)s %(levelname)-8s %(message)s",
        ...     datefmt="%m-%d %H:%M:%S",
        ... )

    A user can select the desired logging level and, in general, take advantage of all of
    the functionality offered by the Python `logging` module.

    Args:
        a: {a}

        b: {b}

        c: {c}
            {c_admonitions}

        alpha: {alpha}

        beta: {beta} from a previously planned and autotuned matrix multiplication.

        epilog: {epilog}

        epilog_inputs: {epilog_inputs}

        qualifiers: {qualifiers}

        options: {options}

        preferences: {preferences}

        algorithm: An object of type :class:`Algorithm` objects can be directly provided to
            bypass planning, if desired. The algorithm object must be compatible with the
            matrix multiplication. A typical use for this option is to provide an algorithm
            that has been serialized (pickled) from a previously planned and autotuned
            matrix multiplication.

        stream: {stream}

        quantization_scales: {quantization_scales}

    Returns:
        {result}

    Semantics:
        {semantics}

    Narrow-precision support:
        {narrow_precision}

    .. seealso::
        :class:`Matmul`, :class:`MatmulOptions`, :class:`MatmulEpilog`,
        :class:`MatmulPlanPreferences`

    Examples:

        >>> import cupy as cp
        >>> import nvmath

        Create three float32 ndarrays on the GPU:

        >>> M, N, K = 128, 64, 256
        >>> a = cp.random.rand(M, K, dtype=cp.float32)
        >>> b = cp.random.rand(K, N, dtype=cp.float32)
        >>> c = cp.random.rand(M, N, dtype=cp.float32)

        Perform the operation :math:`\\alpha A @ B + \\beta C` using :func:`matmul`. The
        result `r` is also a CuPy float32 ndarray:

        >>> r = nvmath.linalg.advanced.matmul(a, b, c, alpha=1.23, beta=0.74)

        An epilog can be used as well. Here we perform
        :math:`RELU(\\alpha A @ B + \\beta C)`:

        >>> epilog = nvmath.linalg.advanced.MatmulEpilog.RELU
        >>> r = nvmath.linalg.advanced.matmul(a, b, c, alpha=1.23, beta=0.74, epilog=epilog)

        Options can be provided to customize the operation:

        >>> compute_type = nvmath.linalg.advanced.MatmulComputeType.COMPUTE_32F_FAST_TF32
        >>> o = nvmath.linalg.advanced.MatmulOptions(compute_type=compute_type)
        >>> r = nvmath.linalg.advanced.matmul(a, b, options=o)

        See `MatmulOptions` for the complete list of available options.

        The package current stream is used by default, but a stream can be explicitly
        provided to the Matmul operation. This can be done if the operands are computed on a
        different stream, for example:

        >>> s = cp.cuda.Stream()
        >>> with s:
        ...     a = cp.random.rand(M, K)
        ...     b = cp.random.rand(K, N)
        >>> r = nvmath.linalg.advanced.matmul(a, b, stream=s)

        The operation above runs on stream `s` and is ordered with respect to the input
        computation.

        Create  NumPy ndarrays on the CPU.

        >>> import numpy as np
        >>> a = np.random.rand(M, K)
        >>> b = np.random.rand(K, N)

        Provide the NumPy ndarrays to :func:`matmul`, with the result also being a NumPy
        ndarray:

        >>> r = nvmath.linalg.advanced.matmul(a, b)

    Notes:
        - This function is a convenience wrapper around :class:`Matmul` and is
          specifically meant for *single* use.

    Further examples can be found in the `nvmath/examples/linalg/advanced/matmul
    <https://github.com/NVIDIA/nvmath-python/tree/main/examples/linalg/advanced/matmul>`_
    directory.
    """

    # Set algorithm limit to 1, but take a copy first if needed.
    if isinstance(preferences, _configuration.MatmulPlanPreferences):
        preferences = copy.copy(preferences)

    preferences = utils.check_or_create_options(
        _configuration.MatmulPlanPreferences, preferences, "Matrix multiplication plan preferences"
    )
    preferences.limit = 1

    if algorithm is None:
        algorithms = None
    else:
        algorithms = [algorithm]  # The type of algorithm should be algorithm.Algorithm and will be checked in plan()

    with Matmul(
        a,
        b,
        c=c,
        alpha=alpha,
        beta=beta,
        qualifiers=qualifiers,
        options=options,
        stream=stream,
        quantization_scales=quantization_scales,
    ) as mm:
        mm.plan(preferences=preferences, epilog=epilog, epilog_inputs=epilog_inputs, stream=stream, algorithms=algorithms)

        r = mm.execute(stream=stream)

    return r
