You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by tq...@apache.org on 2026/10/14 23:44:17 UTC
(tvm) branch main updated: [REFACTOR][TIR] Retire legacy CUDA tensorization and warp memory lowering (#20734)
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new eb2a8f7b8c [REFACTOR][TIR] Retire legacy CUDA tensorization and warp memory lowering (#20734)
eb2a8f7b8c is described below
commit eb2a8f7b8c9e621eea89a978a410da51c913a4ae
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Wed Oct 14 19:44:09 2026 -0400
[REFACTOR][TIR] Retire legacy CUDA tensorization and warp memory lowering (#20734)
Retire the legacy CUDA s_tir tensor-core intrinsic catalogue and its
MetaSchedule tensor-core presets, automatic selection, and dedicated
callers. CUDA scheduling uses the generic preset; new tensor-core
kernels use the PTX/tile paths.
Remove LowerWarpMemory and its public entry points and pipeline calls.
ROCm MFMA uses explicit per-lane local fragments, and CUDA/OpenCL reject
residual warp storage. Keep shared scheduling transforms and MemHammer
support, with the retained matrixC layout mapping defined directly in
its transform.
Move the unchanged definitions of `tirx.ptx_legacy.mma`,
`tirx.ptx_legacy.ldmatrix`, `tirx.cuda.mma_store`, and
`tirx.cuda.mma_fill` into CUDA `op/legacy.cc`. Preserve their
construction APIs for downstream code generators while removing TVM's
CUDA emitters and their exclusive helpers. Shared PTX dtype support
remains available.
---
docs/reference/api/python/s_tir/tensor_intrin.rst | 6 -
docs/tirx/api/cuda.rst | 4 +
docs/tirx/api/ptx.rst | 4 +-
docs/tirx/arch/lowering_pipeline.rst | 2 +-
include/tvm/s_tir/meta_schedule/mutator.h | 2 -
include/tvm/s_tir/meta_schedule/postproc.h | 2 -
include/tvm/s_tir/meta_schedule/schedule_rule.h | 30 -
include/tvm/tirx/transform.h | 6 -
python/tvm/backend/cuda/op.py | 2 +-
python/tvm/s_tir/meta_schedule/mutator/mutator.py | 4 +-
.../tvm/s_tir/meta_schedule/postproc/postproc.py | 5 +-
.../s_tir/meta_schedule/schedule_rule/__init__.py | 1 -
.../schedule_rule/multi_level_tiling.py | 58 +-
.../meta_schedule/schedule_rule/schedule_rule.py | 5 +-
.../space_generator/space_generator.py | 12 +-
.../meta_schedule/testing/space_generation.py | 4 +-
python/tvm/s_tir/pipeline.py | 1 -
python/tvm/s_tir/tensor_intrin/__init__.py | 1 -
python/tvm/s_tir/tensor_intrin/cuda.py | 1193 ------------------
python/tvm/s_tir/tensor_intrin/rocm.py | 107 +-
python/tvm/s_tir/testing/tir.py | 119 +-
python/tvm/tirx/compilation_pipeline.py | 1 -
python/tvm/tirx/transform/transform.py | 11 -
src/backend/cuda/codegen/codegen_cuda.cc | 131 +-
src/backend/cuda/codegen/ptx.cc | 494 +-------
src/backend/cuda/codegen/ptx.h | 76 --
src/backend/cuda/op/legacy.cc | 61 +
src/backend/cuda/op/target_builtin.cc | 25 -
src/backend/opencl/codegen/codegen_opencl.cc | 2 +
src/s_tir/meta_schedule/mutator/mutator.cc | 3 -
src/s_tir/meta_schedule/postproc/postproc.cc | 15 -
.../multi_level_tiling_tensor_core.cc | 984 ---------------
.../meta_schedule/schedule_rule/schedule_rule.cc | 43 -
.../space_generator/space_generator.cc | 22 -
src/s_tir/transform/inject_permuted_layout.cc | 116 +-
src/s_tir/transform/inject_software_pipeline.cc | 40 -
src/s_tir/transform/transform_mma_tensor_layout.cc | 23 +-
src/tirx/transform/lower_warp_memory.cc | 814 -------------
.../test_meta_schedule_mma_tensorize.py | 343 ------
...est_meta_schedule_postproc_rewrite_tensorize.py | 2 +-
...ule_space_cuda_async_multiple_initialization.py | 2 +-
.../test_meta_schedule_space_generator.py | 21 +
.../test_meta_schedule_trace_apply.py | 3 +-
.../s_tir/schedule/test_tir_schedule_analysis.py | 31 +-
..._tir_schedule_tensorize_ldmatrix_mma_numeric.py | 464 -------
.../test_tir_schedule_tensorize_mfma_numeric.py | 108 +-
.../test_s_tir_transform_inject_permuted_layout.py | 360 +-----
...est_s_tir_transform_inject_software_pipeline.py | 169 ---
.../test_s_tir_transform_mma_tensor_layout.py | 23 +
tests/python/tirx/base/test_tirx_ptx_ldmatrix.py | 101 --
tests/python/tirx/base/test_tirx_ptx_mma.py | 1264 --------------------
tests/python/tirx/codegen/test_codegen_cuda.py | 54 +
52 files changed, 432 insertions(+), 6942 deletions(-)
diff --git a/docs/reference/api/python/s_tir/tensor_intrin.rst b/docs/reference/api/python/s_tir/tensor_intrin.rst
index dfdfed5645..383c4c4e14 100644
--- a/docs/reference/api/python/s_tir/tensor_intrin.rst
+++ b/docs/reference/api/python/s_tir/tensor_intrin.rst
@@ -24,12 +24,6 @@ tvm.s_tir.tensor_intrin
:members:
:imported-members:
-tvm.s_tir.tensor_intrin.cuda
-****************************
-.. automodule:: tvm.s_tir.tensor_intrin.cuda
- :members:
- :noindex:
-
tvm.s_tir.tensor_intrin.arm_cpu
*******************************
.. automodule:: tvm.s_tir.tensor_intrin.arm_cpu
diff --git a/docs/tirx/api/cuda.rst b/docs/tirx/api/cuda.rst
index 1c89493d27..3762be0f6b 100644
--- a/docs/tirx/api/cuda.rst
+++ b/docs/tirx/api/cuda.rst
@@ -76,6 +76,10 @@ The current ``Tx.cuda`` surface includes:
``iket.official_event``, ``timer_init``, ``timer_start``, ``timer_end``,
``timer_finalize``, ``mma_store``, and ``mma_fill``
+``mma_store`` and ``mma_fill`` remain available for downstream emitters.
+TVM CUDA codegen does not emit them; express stores and initialization with
+ordinary tensor assignments.
+
NVSHMEM namespace
-----------------
diff --git a/docs/tirx/api/ptx.rst b/docs/tirx/api/ptx.rst
index 567b9cc6ee..c244389c77 100644
--- a/docs/tirx/api/ptx.rst
+++ b/docs/tirx/api/ptx.rst
@@ -182,8 +182,8 @@ PTX versus CUDA helpers
Use ``Tx.ptx`` for one table-described PTX instruction. Use ``Tx.cuda`` for
backend helpers that require multiple statements, C/C++ expressions, descriptor
packing, or other behavior that cannot be represented as one PTX instruction.
-``Tx.ptx_legacy`` only preserves historical spellings needed by compatibility
-passes and is not the API for new kernels.
+``Tx.ptx_legacy`` preserves operation construction for downstream emitters.
+TVM CUDA codegen does not emit these operations; use ``Tx.ptx`` in new kernels.
Namespace API
-------------
diff --git a/docs/tirx/arch/lowering_pipeline.rst b/docs/tirx/arch/lowering_pipeline.rst
index 55aa5e9a51..672f400c39 100644
--- a/docs/tirx/arch/lowering_pipeline.rst
+++ b/docs/tirx/arch/lowering_pipeline.rst
@@ -121,7 +121,7 @@ The ``tirx_pipeline`` module pass applies this sequence (a few passes are gated
- **host**: ``LowerTVMBuiltin`` (lower ``tvm_*`` builtins), ``LowerIntrin``
(target-specific intrinsics)
-- **device**: ``LowerWarpMemory`` (warp-scoped tensors → shuffles), ``StmtSimplify``,
+- **device**: ``StmtSimplify``,
``LowerIntrin``
Inside LowerTIRx
diff --git a/include/tvm/s_tir/meta_schedule/mutator.h b/include/tvm/s_tir/meta_schedule/mutator.h
index c394d57b8c..85205213d8 100644
--- a/include/tvm/s_tir/meta_schedule/mutator.h
+++ b/include/tvm/s_tir/meta_schedule/mutator.h
@@ -136,8 +136,6 @@ class Mutator : public ffi::ObjectRef {
TVM_DLL static ffi::Map<Mutator, FloatImm, void> DefaultLLVM();
/*! \brief Create default mutators for CUDA */
TVM_DLL static ffi::Map<Mutator, FloatImm, void> DefaultCUDA();
- /*! \brief Create default mutators for CUDA with TensorCore */
- TVM_DLL static ffi::Map<Mutator, FloatImm, void> DefaultCUDATensorCore();
/*! \brief Create default mutators for Hexagon */
TVM_DLL static ffi::Map<Mutator, FloatImm, void> DefaultHexagon();
diff --git a/include/tvm/s_tir/meta_schedule/postproc.h b/include/tvm/s_tir/meta_schedule/postproc.h
index 57a1c2019a..3f235f9a36 100644
--- a/include/tvm/s_tir/meta_schedule/postproc.h
+++ b/include/tvm/s_tir/meta_schedule/postproc.h
@@ -170,8 +170,6 @@ class Postproc : public ffi::ObjectRef {
TVM_DLL static ffi::Array<Postproc, void> DefaultRISCV();
/*! \brief Create default postprocessors for CUDA */
TVM_DLL static ffi::Array<Postproc, void> DefaultCUDA();
- /*! \brief Create default postprocessors for CUDA with TensorCore */
- TVM_DLL static ffi::Array<Postproc, void> DefaultCUDATensorCore();
/*! \brief Create default postprocessors for Hexagon */
TVM_DLL static ffi::Array<Postproc, void> DefaultHexagon();
diff --git a/include/tvm/s_tir/meta_schedule/schedule_rule.h b/include/tvm/s_tir/meta_schedule/schedule_rule.h
index 39b119b309..cf0f47dbf9 100644
--- a/include/tvm/s_tir/meta_schedule/schedule_rule.h
+++ b/include/tvm/s_tir/meta_schedule/schedule_rule.h
@@ -192,34 +192,6 @@ class ScheduleRule : public ffi::ObjectRef {
ffi::Optional<ffi::Map<ffi::String, ffi::Any>> reuse_read,
ffi::Optional<ffi::Map<ffi::String, ffi::Any>> reuse_write);
- /*!
- * \brief Extension of MultiLevelTiling for auto-tensorization with multiple groups of candidate
- * tensor core intrinsics
- * \param intrin_groups A list of groups of tensor core intrinsics. The map should contains key
- * "init", "load_a", "load_b", "compute", "store", which represent the tensor intrin for
- * initialization, loading operand A, loading operand B, tensor core computation, storing the
- * result. The value of the map should be names of tensor intrinsics, must be registered via
- * TensorIntrin.register(...) beforehand
- * \param structure The tiling structure. Recommended:
- * - 'SSSRRSRS' on GPU
- * \param tile_binds For each level of tiles, which thread axis it is bound to. Recommended:
- * - [blockIdx.y, blockIdx.x, threadIdx.y] on GPU
- * \param max_innermost_factor The maximum size of the innermost factor. std::nullopt means no
- * limit \param vector_load_lens The length of vector lane in vectorized cooperative fetching.
- * std::nullopt means disable vectorization
- * \param reuse_read Data reuse configuration for reading. std::nullopt means no reuse.
- * \param reuse_write Data reuse configuration for writing. std::nullopt means no reuse.
- * \param use_software_pipeline Whether use the software pipeline.
- * \return The schedule rule created
- */
- TVM_DLL static ScheduleRule MultiLevelTilingTensorCore(
- ffi::Array<ffi::Map<ffi::String, ffi::String>> intrin_groups, ffi::String structure,
- ffi::Optional<ffi::Array<ffi::String>> tile_binds,
- ffi::Optional<int64_t> max_innermost_factor,
- ffi::Optional<ffi::Array<int64_t>> vector_load_lens,
- ffi::Optional<ffi::Map<ffi::String, ffi::Any>> reuse_read,
- ffi::Optional<ffi::Map<ffi::String, ffi::Any>> reuse_write, bool use_software_pipeline);
-
/*!
* \brief Extension of MultiLevelTiling for backends with wide vectors.
* The loop over the innermost spatial axis of the output tensor is always vectorized with the
@@ -304,8 +276,6 @@ class ScheduleRule : public ffi::ObjectRef {
TVM_DLL static ffi::Array<ScheduleRule, void> DefaultX86(const ffi::String& type);
/*! \brief Create default schedule rules for CUDA */
TVM_DLL static ffi::Array<ScheduleRule, void> DefaultCUDA();
- /*! \brief Create default postprocessors for CUDA with TensorCore */
- TVM_DLL static ffi::Array<ScheduleRule, void> DefaultCUDATensorCore();
/*! \brief Create default schedule rules for Hexagon */
TVM_DLL static ffi::Array<ScheduleRule, void> DefaultHexagon();
/*! \brief Create default schedule rules for ARM CPU (NEON and DOTPROD) */
diff --git a/include/tvm/tirx/transform.h b/include/tvm/tirx/transform.h
index 3b6ede5942..7da082b7e3 100644
--- a/include/tvm/tirx/transform.h
+++ b/include/tvm/tirx/transform.h
@@ -209,12 +209,6 @@ TVM_TIRX_DLL Pass LowerTVMBuiltin();
*/
TVM_TIRX_DLL Pass LowerIntrin();
-/*!
- * \brief Lower warp memory access to low-level device related function calls.
- * \return The pass.
- */
-TVM_TIRX_DLL Pass LowerWarpMemory();
-
/*!
* \brief Narrow down PrimExpr datatype in stmt to target_bits.
*
diff --git a/python/tvm/backend/cuda/op.py b/python/tvm/backend/cuda/op.py
index 495d5269a1..b0b796ee7e 100644
--- a/python/tvm/backend/cuda/op.py
+++ b/python/tvm/backend/cuda/op.py
@@ -578,7 +578,7 @@ def ptx_legacy_ldmatrix(
):
"""Load a matrix with explicit pointer and element-offset operands.
- The legacy lowering uses the result ``ty`` for its element width, including
+ Downstream legacy emitters use the result ``ty`` for its element width, including
the transposed int8 gather form. ``dtype`` is the PTX instruction type token.
"""
return _legacy_ldmatrix(
diff --git a/python/tvm/s_tir/meta_schedule/mutator/mutator.py b/python/tvm/s_tir/meta_schedule/mutator/mutator.py
index 52d432d97c..59b79a27bb 100644
--- a/python/tvm/s_tir/meta_schedule/mutator/mutator.py
+++ b/python/tvm/s_tir/meta_schedule/mutator/mutator.py
@@ -81,7 +81,6 @@ class Mutator(Object):
kind: Literal[
"llvm",
"cuda",
- "cuda-tensorcore",
"hexagon",
],
) -> dict["Mutator", float]:
@@ -89,7 +88,7 @@ class Mutator(Object):
Parameters
----------
- kind : Literal["llvm", "cuda", "cuda-tensorcore", "hexagon"]
+ kind : Literal["llvm", "cuda", "hexagon"]
The kind of mutators.
Returns
@@ -101,7 +100,6 @@ class Mutator(Object):
# pylint: disable=no-member
"llvm": _ffi_api.MutatorDefaultLLVM, # type: ignore
"cuda": _ffi_api.MutatorDefaultCUDA, # type: ignore
- "cuda-tensorcore": _ffi_api.MutatorDefaultCUDATensorCore, # type: ignore
"hexagon": _ffi_api.MutatorDefaultHexagon, # type: ignore
# pylint: enable=no-member
}
diff --git a/python/tvm/s_tir/meta_schedule/postproc/postproc.py b/python/tvm/s_tir/meta_schedule/postproc/postproc.py
index e75fa13ccc..5e5a7f1374 100644
--- a/python/tvm/s_tir/meta_schedule/postproc/postproc.py
+++ b/python/tvm/s_tir/meta_schedule/postproc/postproc.py
@@ -78,12 +78,12 @@ class Postproc(Object):
return _ffi_api.PostprocClone(self) # type: ignore # pylint: disable=no-member
@staticmethod
- def create(kind: Literal["llvm", "cuda", "cuda-tensorcore", "hexagon"]) -> list["Postproc"]:
+ def create(kind: Literal["llvm", "cuda", "hexagon"]) -> list["Postproc"]:
"""Create a list of default postprocessors.
Parameters
----------
- kind : Literal["llvm", "cuda", "cuda-tensorcore", "hexagon"]
+ kind : Literal["llvm", "cuda", "hexagon"]
The kind of the postprocessors.
Returns
@@ -95,7 +95,6 @@ class Postproc(Object):
# pylint: disable=no-member
"llvm": _ffi_api.PostprocDefaultLLVM, # type: ignore
"cuda": _ffi_api.PostprocDefaultCUDA, # type: ignore
- "cuda-tensorcore": _ffi_api.PostprocDefaultCUDATensorCore, # type: ignore
"hexagon": _ffi_api.PostprocDefaultHexagon, # type: ignore
# pylint: enable=no-member
}
diff --git a/python/tvm/s_tir/meta_schedule/schedule_rule/__init__.py b/python/tvm/s_tir/meta_schedule/schedule_rule/__init__.py
index dcff4f1e07..b0ccdd3af4 100644
--- a/python/tvm/s_tir/meta_schedule/schedule_rule/__init__.py
+++ b/python/tvm/s_tir/meta_schedule/schedule_rule/__init__.py
@@ -28,7 +28,6 @@ from .auto_inline import AutoInline, InlineConstantScalars
from .cross_thread_reduction import CrossThreadReduction
from .multi_level_tiling import (
MultiLevelTiling,
- MultiLevelTilingTensorCore,
MultiLevelTilingWideVector,
MultiLevelTilingWithIntrin,
ReuseType,
diff --git a/python/tvm/s_tir/meta_schedule/schedule_rule/multi_level_tiling.py b/python/tvm/s_tir/meta_schedule/schedule_rule/multi_level_tiling.py
index 7c4b84581d..604a72babd 100644
--- a/python/tvm/s_tir/meta_schedule/schedule_rule/multi_level_tiling.py
+++ b/python/tvm/s_tir/meta_schedule/schedule_rule/multi_level_tiling.py
@@ -16,7 +16,7 @@
# under the License.
"""Multi-level tiling with reuse."""
-from collections.abc import Callable, Mapping
+from collections.abc import Callable
from typing import Any, NamedTuple
from tvm_ffi import register_object
@@ -144,62 +144,6 @@ class MultiLevelTilingWithIntrin(ScheduleRule):
)
-@register_object("s_tir.meta_schedule.MultiLevelTilingTensorCore")
-class MultiLevelTilingTensorCore(ScheduleRule):
- """Extension of MultiLevelTiling for auto-tensorizing with multiple groups of candidate tensor
- core intrinsics.
-
- Parameters
- ----------
- intrin_groups : List[Mapping[str, str]]
- A list of groups of tensor core intrinsics. The map should contains key "init", "load_a",
- "load_b", "compute", "store", which represent the tensor intrin for initialization,
- loading operand A, loading operand B, tensor core computation, storing the result.
- The value of the map should be names of tensor intrinsics, must be registerd via
- TensorIntrin.register(...) beforehand
- structure : str
- The tiling structure. Recommended:
- - 'SSSRRSRS' on GPU
- tile_bind : Optional[List[str]]
- For each level of tiles, which thread axis it is bound to. Recommended:
- - [blockIdx.y, vthread.x, threadIdx.y] on GPU
- max_innermost_factor : Optional[int]
- The maximum size of the innermost factor. None means no limit
- vector_load_lens : Optional[List[int]]
- The length of vector lane in vectorized cooperative fetching.
- None means disable vectorization
- reuse_read : Optional[ReuseType]
- Data reuse configuration for reading. None means no reuse.
- reuse_write : Optional[ReuseType]
- Data reuse configuration for writing. None means no reuse.
- use_software_pipeline : bool
- Whether to use the software pipeline.
- """
-
- def __init__(
- self,
- intrin_groups: list[Mapping[str, str]],
- structure: str,
- tile_binds: list[str] | None = None,
- max_innermost_factor: int | None = None,
- vector_load_lens: list[int] | None = None,
- reuse_read: ReuseType | None = None,
- reuse_write: ReuseType | None = None,
- use_software_pipeline: bool = False,
- ) -> None:
- self.__init_handle_by_constructor__(
- _ffi_api.ScheduleRuleMultiLevelTilingTensorCore, # type: ignore # pylint: disable=no-member
- intrin_groups,
- structure,
- tile_binds,
- max_innermost_factor,
- vector_load_lens,
- reuse_read.as_dict() if reuse_read is not None else None,
- reuse_write.as_dict() if reuse_write is not None else None,
- use_software_pipeline,
- )
-
-
@register_object("s_tir.meta_schedule.MultiLevelTilingWideVector")
class MultiLevelTilingWideVector(ScheduleRule):
"""Extension of MultiLevelTiling for backends with wide vectors. The loop over the innermost
diff --git a/python/tvm/s_tir/meta_schedule/schedule_rule/schedule_rule.py b/python/tvm/s_tir/meta_schedule/schedule_rule/schedule_rule.py
index 746f8c6af9..3ac3181dbe 100644
--- a/python/tvm/s_tir/meta_schedule/schedule_rule/schedule_rule.py
+++ b/python/tvm/s_tir/meta_schedule/schedule_rule/schedule_rule.py
@@ -85,12 +85,12 @@ class ScheduleRule(Object):
return _ffi_api.ScheduleRuleClone(self) # type: ignore # pylint: disable=no-member
@staticmethod
- def create(kind: Literal["llvm", "cuda", "cuda-tensorcore", "hexagon"]) -> list["ScheduleRule"]:
+ def create(kind: Literal["llvm", "cuda", "hexagon"]) -> list["ScheduleRule"]:
"""Create a list of schedule rules for the given kind.
Parameters
----------
- kind : Literal["llvm", "cuda", "cuda-tensorcore", "hexagon"]
+ kind : Literal["llvm", "cuda", "hexagon"]
The kind of the schedule rules.
Returns
@@ -102,7 +102,6 @@ class ScheduleRule(Object):
# pylint: disable=no-member
"llvm": _ffi_api.ScheduleRuleDefaultLLVM, # type: ignore
"cuda": _ffi_api.ScheduleRuleDefaultCUDA, # type: ignore
- "cuda-tensorcore": _ffi_api.ScheduleRuleDefaultCUDATensorCore, # type: ignore
"hexagon": _ffi_api.ScheduleRuleDefaultHexagon, # type: ignore
# pylint: enable=no-member
}
diff --git a/python/tvm/s_tir/meta_schedule/space_generator/space_generator.py b/python/tvm/s_tir/meta_schedule/space_generator/space_generator.py
index 73068a0d42..732b00351f 100644
--- a/python/tvm/s_tir/meta_schedule/space_generator/space_generator.py
+++ b/python/tvm/s_tir/meta_schedule/space_generator/space_generator.py
@@ -133,15 +133,9 @@ class SpaceGenerator(Object):
ScheduleFnType = SpaceGenerator.ScheduleFnType
-ScheduleRuleType = (
- list["ScheduleRule"] | Literal["llvm", "cuda", "cuda-tensorcore", "hexagon", "from-target"]
-)
-PostprocType = (
- list["Postproc"] | Literal["llvm", "cuda", "cuda-tensorcore", "hexagon", "from-target"]
-)
-MutatorProbType = (
- dict["Mutator", float] | Literal["llvm", "cuda", "cuda-tensorcore", "hexagon", "from-target"]
-)
+ScheduleRuleType = list["ScheduleRule"] | Literal["llvm", "cuda", "hexagon", "from-target"]
+PostprocType = list["Postproc"] | Literal["llvm", "cuda", "hexagon", "from-target"]
+MutatorProbType = dict["Mutator", float] | Literal["llvm", "cuda", "hexagon", "from-target"]
create = SpaceGenerator.create # pylint: disable=invalid-name
diff --git a/python/tvm/s_tir/meta_schedule/testing/space_generation.py b/python/tvm/s_tir/meta_schedule/testing/space_generation.py
index 877daf423b..13d49c8959 100644
--- a/python/tvm/s_tir/meta_schedule/testing/space_generation.py
+++ b/python/tvm/s_tir/meta_schedule/testing/space_generation.py
@@ -32,7 +32,7 @@ from tvm.target import Target
def get_rules(
- kind: Literal["llvm", "cuda", "cuda-tensorcore", "hexagon"],
+ kind: Literal["llvm", "cuda", "hexagon"],
types: type | tuple[type, ...],
) -> list[ms.ScheduleRule]:
"""Get default schedule rules"""
@@ -57,7 +57,7 @@ def structural_equal_no_gs(mod1: IRModule, mod2: IRModule) -> bool:
def generate_design_space(
- kind: Literal["llvm", "cuda", "cuda-tensorcore", "hexagon"],
+ kind: Literal["llvm", "cuda", "hexagon"],
mod: IRModule,
target: Target,
types: type | tuple[type, ...],
diff --git a/python/tvm/s_tir/pipeline.py b/python/tvm/s_tir/pipeline.py
index f02ccafbda..442aa3a8fd 100644
--- a/python/tvm/s_tir/pipeline.py
+++ b/python/tvm/s_tir/pipeline.py
@@ -231,7 +231,6 @@ def finalize_host_passes(): # pylint: disable=unused-argument
def finalize_device_passes(): # pylint: disable=unused-argument
"""The default finalization passes for TIR backend."""
device_pass_list = [
- tirx.transform.LowerWarpMemory(),
tirx.transform.StmtSimplify(),
tirx.transform.LowerIntrin(),
]
diff --git a/python/tvm/s_tir/tensor_intrin/__init__.py b/python/tvm/s_tir/tensor_intrin/__init__.py
index 0032db630c..c2384a65e4 100644
--- a/python/tvm/s_tir/tensor_intrin/__init__.py
+++ b/python/tvm/s_tir/tensor_intrin/__init__.py
@@ -19,7 +19,6 @@
"""Intrinsics for tensorization."""
from tvm.runtime import enabled
-from . import cuda
if enabled("llvm"):
from . import arm_cpu, x86, rocm, hexagon, riscv_cpu
diff --git a/python/tvm/s_tir/tensor_intrin/cuda.py b/python/tvm/s_tir/tensor_intrin/cuda.py
deleted file mode 100644
index 7929be294d..0000000000
--- a/python/tvm/s_tir/tensor_intrin/cuda.py
+++ /dev/null
@@ -1,1193 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-# pylint: disable=invalid-name,missing-function-docstring,unused-variable
-# ruff: noqa: E731
-"""Intrinsics for tensorization on NVIDIA GPU."""
-
-from typing import Literal
-
-from tvm_ffi import register_global_func
-
-from tvm.runtime import MemoryScope, convert
-from tvm.s_tir import TensorIntrin
-from tvm.script import s_tir as Ts
-from tvm.script import tirx as T
-from tvm.tirx import Cast, IntImm
-from tvm.tirx.function import Function
-
-_SCOPE_GLOBAL = MemoryScope("global")
-_SCOPE_SHARED = MemoryScope("shared")
-_SCOPE_SHARED_DYN = MemoryScope("shared.dyn")
-
-
-def shared_16x16_to_ldmatrix_32x8_layout(i, j):
- thread_id = 4 * (i % 8) + (j % 8) // 2
- return thread_id, 4 * (j // 8) + (i // 8) * 2 + (j % 2)
-
-
-def shared_16x32_to_ldmatrix_32x16_layout(i, j):
- thread_id = 4 * (i % 8) + (j % 16) // 4
- return thread_id, 8 * (j // 16) + (i // 8) * 4 + j % 4
-
-
-def shared_32x16_to_ldmatrix_32x16_layout(i, j):
- thread_id = (i % 16) // 4 + 4 * (j % 8)
- return thread_id, 8 * (j // 8) + (i // 16) * 4 + i % 4
-
-
-def ldmatrix_32x8_to_shared_16x16_layout(thread_id, local_id):
- row = 8 * (local_id % 4 // 2) + (thread_id // 4)
- col = 8 * (local_id // 4) + (thread_id % 4) * 2 + (local_id % 2)
- return row, col
-
-
-@register_global_func("tirx.index_map.shared_16x16_to_ldmatrix_32x8_layout")
-def index_map_shared_16x16_to_ldmatrix_32x8_layout(ind):
- i, j = ind[0], ind[1]
- thread_id, local_id = shared_16x16_to_ldmatrix_32x8_layout(i, j)
- return convert([thread_id, local_id])
-
-
-lift = convert
-
-M_DIM = 16
-N_DIM = 16
-WARP_SIZE = 32
-HALF_WARP = WARP_SIZE // 2
-HALF_WARP_expr = lift(HALF_WARP)
-
-
-def get_ldmatrix_intrin(
- k_dim: int,
- dtype: str,
- matrix_name: Literal["A", "B"],
- transposed: bool,
- shared_scope: str = "shared",
-):
- local_size = (M_DIM * k_dim) // WARP_SIZE
- smem_offset = None
- index_map = None
-
- if matrix_name == "A":
- transpose_in_ldmatrix = transposed
- # transpose_layout_for_ldmatrix_input: Every thread loads 8 bytes data. This determines
- # which 8 bytes every thread loads.
- # If transpose_layout_for_ldmatrix_input is False, the load pattern is
- # T0 T0 T0 T0 T16 T16 T16 T16
- # T1 T1 T1 T1 T17 T17 T17 T17
- # ...
- # T8 T8 T8 T8 T24 T24 T24 T24
- # T9 T9 T9 T9 T25 T25 T25 T25
- # ...
- # T15 T15 T15 T15 T31 T31 T31 T31
- # Otherwise, the load pattern is
- # T0 T0 T0 T0 T8 T8 T8 T8
- # T1 T1 T1 T1 T9 T9 T9 T9
- # ...
- # T7 T7 T7 T7 T15 T15 T15 T15
- # T16 T16 T16 T16 T24 T24 T24 T24
- # T17 T17 T17 T17 T25 T25 T25 T25
- # ...
- # T23 T23 T23 T23 T31 T31 T31 T31
- transpose_layout_for_ldmatrix_input = transposed
- smem_tile_row, smem_tile_col = (M_DIM, k_dim) if not transposed else (k_dim, M_DIM)
- else:
- assert matrix_name == "B"
- transpose_in_ldmatrix = not transposed
- transpose_layout_for_ldmatrix_input = transposed
- smem_tile_row, smem_tile_col = (k_dim, N_DIM) if not transposed else (N_DIM, k_dim)
-
- if k_dim == 16:
- assert dtype == "float16"
-
- index_map = shared_16x16_to_ldmatrix_32x8_layout
-
- if transpose_layout_for_ldmatrix_input:
- smem_offset = lambda tx, stride: (
- stride * 8 * (tx // HALF_WARP_expr)
- + stride * (tx % 8)
- + 8 * ((tx % HALF_WARP_expr) // 8)
- )
- else:
- smem_offset = lambda tx, stride: (
- stride * (tx % HALF_WARP_expr) + 8 * (tx // HALF_WARP_expr)
- )
- else:
- # TODO(yixin): Support TN and TT matmul for int8
- assert matrix_name == "B" or not transposed, (
- "Now only B matrix can be transposed for int8 matmul"
- )
- assert k_dim == 32 and (
- dtype == "int8" or dtype == "float8_e4m3fn" or dtype == "float8_e5m2"
- ), "Only k_dim == 16 (float16) or k_dim == 32 (int8) supported for now"
-
- if matrix_name == "B" and not transposed:
- index_map = shared_32x16_to_ldmatrix_32x16_layout
- # A dummy offset, ldmatrix cannot be used for int8 + trans case.
- # We still use the ldmatrix intrinsic, but lower it to a manual loop in the codegen.
- # Only the stride information is required.
- smem_offset = lambda _, stride: stride
- elif matrix_name == "B" and transposed:
- index_map = shared_16x32_to_ldmatrix_32x16_layout
- smem_offset = lambda tx, stride: (
- stride * 8 * (tx // HALF_WARP_expr)
- + (tx % 8) * stride
- + 16 * ((tx % HALF_WARP_expr) // 8)
- )
- else: # A, not transposed
- index_map = shared_16x32_to_ldmatrix_32x16_layout
- smem_offset = lambda tx, stride: stride * (tx % 16) + 16 * (tx // 16)
-
- offset_factor = smem_tile_col
-
- @Ts.function
- def ldmatrix_desc(
- warp: T.Tensor(
- (WARP_SIZE, local_size),
- dtype,
- align=64,
- offset_factor=offset_factor,
- elem_offset=T.int32(),
- scope="warp",
- ),
- shared: T.Tensor(
- (smem_tile_row, smem_tile_col),
- dtype,
- align=64,
- offset_factor=offset_factor,
- elem_offset=T.int32(),
- scope=shared_scope,
- ),
- ) -> None:
- with Ts.sblock("root"):
- Ts.reads(shared[0:smem_tile_row, 0:smem_tile_col])
- Ts.writes(warp[0:WARP_SIZE, 0:local_size])
-
- for ax0, ax1 in T.grid(smem_tile_row, smem_tile_col):
- with Ts.sblock("shared_warp"):
- v0, v1 = Ts.axis.remap("SS", [ax0, ax1])
- Ts.reads(shared[v0, v1])
-
- warp_indices = index_map(v0, v1)
- Ts.writes(warp[warp_indices[0], warp_indices[1]])
- warp[warp_indices[0], warp_indices[1]] = shared[v0, v1]
-
- s0 = T.dynamic("s0", "int32")
- s1 = T.dynamic("s1", "int32")
-
- @Ts.function
- def ldmatrix_impl(
- warp: T.Tensor(
- (WARP_SIZE, local_size),
- dtype,
- align=64,
- offset_factor=offset_factor,
- elem_offset=T.int32(),
- scope="warp",
- ),
- shared: T.Tensor(
- (smem_tile_row, smem_tile_col),
- dtype,
- align=64,
- offset_factor=offset_factor,
- elem_offset=T.int32(),
- scope=shared_scope,
- strides=[s0, s1],
- ),
- ) -> None:
- with Ts.sblock("root"):
- Ts.reads(shared[0:smem_tile_row, 0:smem_tile_col])
- Ts.writes(warp[0:WARP_SIZE, 0:local_size])
- for tx in T.parallel(0, WARP_SIZE, annotations={"thread_binding": "threadIdx.x"}):
- T.evaluate(
- T.ptx_legacy.ldmatrix(
- transpose_in_ldmatrix,
- 4, # Always load 4 matrices
- ".b16",
- warp.data,
- warp.elem_offset + lift(local_size) * tx,
- shared.ptr_to([0, 0]),
- smem_offset(tx, s0),
- ty=dtype,
- )
- )
-
- return ldmatrix_desc, ldmatrix_impl
-
-
-LDMATRIX_f16_A_INTRIN = "mma_ldmatrix_f16_a"
-TensorIntrin.register(LDMATRIX_f16_A_INTRIN, *get_ldmatrix_intrin(16, "float16", "A", False))
-
-LDMATRIX_f16_B_INTRIN = "mma_ldmatrix_f16_b"
-TensorIntrin.register(LDMATRIX_f16_B_INTRIN, *get_ldmatrix_intrin(16, "float16", "B", False))
-
-LDMATRIX_f16_A_TRANS_INTRIN = "mma_ldmatrix_f16_a_trans"
-TensorIntrin.register(LDMATRIX_f16_A_TRANS_INTRIN, *get_ldmatrix_intrin(16, "float16", "A", True))
-
-LDMATRIX_f16_B_TRANS_INTRIN = "mma_ldmatrix_f16_b_trans"
-TensorIntrin.register(LDMATRIX_f16_B_TRANS_INTRIN, *get_ldmatrix_intrin(16, "float16", "B", True))
-
-LDMATRIX_f16_A_DYN_INTRIN = "mma_ldmatrix_f16_a_dyn"
-TensorIntrin.register(
- LDMATRIX_f16_A_DYN_INTRIN, *get_ldmatrix_intrin(16, "float16", "A", False, "shared.dyn")
-)
-
-LDMATRIX_f16_B_DYN_INTRIN = "mma_ldmatrix_f16_b_dyn"
-TensorIntrin.register(
- LDMATRIX_f16_B_DYN_INTRIN, *get_ldmatrix_intrin(16, "float16", "B", False, "shared.dyn")
-)
-
-LDMATRIX_f16_A_TRANS_DYN_INTRIN = "mma_ldmatrix_f16_a_trans_dyn"
-TensorIntrin.register(
- LDMATRIX_f16_A_TRANS_DYN_INTRIN, *get_ldmatrix_intrin(16, "float16", "A", True, "shared.dyn")
-)
-
-LDMATRIX_f16_B_TRANS_DYN_INTRIN = "mma_ldmatrix_f16_b_trans_dyn"
-TensorIntrin.register(
- LDMATRIX_f16_B_TRANS_DYN_INTRIN, *get_ldmatrix_intrin(16, "float16", "B", True, "shared.dyn")
-)
-
-LDMATRIX_i8_A_INTRIN = "mma_ldmatrix_i8_a"
-TensorIntrin.register(LDMATRIX_i8_A_INTRIN, *get_ldmatrix_intrin(32, "int8", "A", False))
-
-LDMATRIX_i8_B_INTRIN = "mma_ldmatrix_i8_b"
-TensorIntrin.register(LDMATRIX_i8_B_INTRIN, *get_ldmatrix_intrin(32, "int8", "B", False))
-
-LDMATRIX_i8_B_TRANS_INTRIN = "mma_ldmatrix_i8_b_trans"
-TensorIntrin.register(LDMATRIX_i8_B_TRANS_INTRIN, *get_ldmatrix_intrin(32, "int8", "B", True))
-
-LDMATRIX_e4m3_A_INTRIN = "mma_ldmatrix_e4m3_a"
-TensorIntrin.register(LDMATRIX_e4m3_A_INTRIN, *get_ldmatrix_intrin(32, "float8_e4m3fn", "A", False))
-
-LDMATRIX_e4m3_B_INTRIN = "mma_ldmatrix_e4m3_b"
-TensorIntrin.register(LDMATRIX_e4m3_B_INTRIN, *get_ldmatrix_intrin(32, "float8_e4m3fn", "B", False))
-
-LDMATRIX_e4m3_B_TRANS_INTRIN = "mma_ldmatrix_e4m3_b_trans"
-TensorIntrin.register(
- LDMATRIX_e4m3_B_TRANS_INTRIN, *get_ldmatrix_intrin(32, "float8_e4m3fn", "B", True)
-)
-
-LDMATRIX_e5m2_A_INTRIN = "mma_ldmatrix_e5m2_a"
-TensorIntrin.register(LDMATRIX_e5m2_A_INTRIN, *get_ldmatrix_intrin(32, "float8_e5m2", "A", False))
-
-LDMATRIX_e5m2_B_INTRIN = "mma_ldmatrix_e5m2_b"
-TensorIntrin.register(LDMATRIX_e5m2_B_INTRIN, *get_ldmatrix_intrin(32, "float8_e5m2", "B", False))
-
-LDMATRIX_e5m2_B_TRANS_INTRIN = "mma_ldmatrix_e5m2_b_trans"
-TensorIntrin.register(
- LDMATRIX_e5m2_B_TRANS_INTRIN, *get_ldmatrix_intrin(32, "float8_e5m2", "B", True)
-)
-
-
-def get_mma_intrin(
- k_dim,
- a_dtype="float16",
- b_dtype="float16",
- out_dtype="float16",
- a_transposed=False,
- b_transposed=False,
-):
- local_size = (M_DIM * k_dim) // WARP_SIZE
- local_size_out = (M_DIM * N_DIM) // 32
-
- index_map_C = shared_16x16_to_ldmatrix_32x8_layout
-
- if k_dim == 16:
- index_map_A = shared_16x16_to_ldmatrix_32x8_layout
- index_map_B = shared_16x16_to_ldmatrix_32x8_layout
- mma_prefix = "m16n8k16"
- elif k_dim == 32 and b_transposed:
- index_map_A = index_map_B = shared_16x32_to_ldmatrix_32x16_layout
- mma_prefix = "m16n8k32"
- elif k_dim == 32 and not b_transposed:
- index_map_A = shared_16x32_to_ldmatrix_32x16_layout
- index_map_B = shared_32x16_to_ldmatrix_32x16_layout
- mma_prefix = "m16n8k32"
- else:
- assert False
-
- dtype_abbrv = {
- "float16": "fp16",
- "float32": "fp32",
- "int8": "int8",
- "int32": "int32",
- "float8_e4m3fn": "e4m3",
- "float8_e5m2": "e5m2",
- }
- a_dtype_abbrv = dtype_abbrv[a_dtype]
- b_dtype_abbrv = dtype_abbrv[b_dtype]
- out_dtype_abbrv = dtype_abbrv[out_dtype]
-
- def cast_to_out_dtype(v):
- if out_dtype in ["float32", "int32"]:
- return Cast(v, out_dtype)
- return v
-
- def swap_if_flag(i, j, flag):
- return (j, i) if flag else (i, j)
-
- A_offset_factor = M_DIM if a_transposed else k_dim
- B_offset_factor = k_dim if b_transposed else N_DIM
- out_offset_factor = N_DIM
-
- @Ts.function
- def mma_sync_desc(
- A: T.Tensor(
- (WARP_SIZE, local_size),
- a_dtype,
- align=64,
- offset_factor=A_offset_factor,
- elem_offset=T.int32(),
- scope="warp",
- ),
- B: T.Tensor(
- (WARP_SIZE, local_size),
- b_dtype,
- align=64,
- offset_factor=B_offset_factor,
- elem_offset=T.int32(),
- scope="warp",
- ),
- C: T.Tensor(
- (WARP_SIZE, local_size_out),
- out_dtype,
- align=64,
- offset_factor=out_offset_factor,
- elem_offset=T.int32(),
- scope="warp",
- ),
- ) -> None:
- with Ts.sblock("root"):
- Ts.reads(
- C[0:WARP_SIZE, 0:local_size_out],
- A[0:WARP_SIZE, 0:local_size],
- B[0:WARP_SIZE, 0:local_size],
- )
- Ts.writes(C[0:WARP_SIZE, 0:local_size_out])
-
- for i, j, k in T.grid(M_DIM, N_DIM, k_dim):
- with Ts.sblock("C"):
- vi, vj, vk = Ts.axis.remap("SSR", [i, j, k])
- a_indices = swap_if_flag(vi, vk, a_transposed)
- b_indices = swap_if_flag(vk, vj, b_transposed)
-
- c_warp_indices = index_map_C(vi, vj)
- a_warp_indices = index_map_A(a_indices[0], a_indices[1])
- b_warp_indices = index_map_B(b_indices[0], b_indices[1])
-
- Ts.reads(
- C[c_warp_indices[0], c_warp_indices[1]],
- A[a_warp_indices[0], a_warp_indices[1]],
- B[b_warp_indices[0], b_warp_indices[1]],
- )
- Ts.writes(C[c_warp_indices[0], c_warp_indices[1]])
-
- C[c_warp_indices[0], c_warp_indices[1]] += cast_to_out_dtype(
- A[a_warp_indices[0], a_warp_indices[1]]
- ) * cast_to_out_dtype(B[b_warp_indices[0], b_warp_indices[1]])
-
- @Ts.function
- def mma_sync_impl(
- A: T.Tensor(
- (WARP_SIZE, local_size),
- a_dtype,
- align=64,
- offset_factor=A_offset_factor,
- elem_offset=T.int32(),
- scope="warp",
- ),
- B: T.Tensor(
- (WARP_SIZE, local_size),
- b_dtype,
- align=64,
- offset_factor=B_offset_factor,
- elem_offset=T.int32(),
- scope="warp",
- ),
- C: T.Tensor(
- (WARP_SIZE, local_size_out),
- out_dtype,
- align=64,
- offset_factor=out_offset_factor,
- elem_offset=T.int32(),
- scope="warp",
- ),
- ) -> None:
- with Ts.sblock("root"):
- Ts.reads(
- C[0:WARP_SIZE, 0:local_size_out],
- A[0:WARP_SIZE, 0:local_size],
- B[0:WARP_SIZE, 0:local_size],
- )
- Ts.writes(C[0:WARP_SIZE, 0:local_size_out])
-
- for tx in T.parallel(0, WARP_SIZE, annotations={"thread_binding": "threadIdx.x"}):
- T.evaluate(
- T.ptx_legacy.mma(
- mma_prefix,
- "row",
- "col",
- a_dtype_abbrv,
- b_dtype_abbrv,
- out_dtype_abbrv,
- A.data,
- A.elem_offset + tx * lift(local_size),
- B.data,
- B.elem_offset + tx * lift(local_size),
- C.data,
- C.elem_offset + tx * lift(local_size_out),
- False,
- ty=out_dtype,
- )
- )
-
- T.evaluate(
- T.ptx_legacy.mma(
- mma_prefix,
- "row",
- "col",
- a_dtype_abbrv,
- b_dtype_abbrv,
- out_dtype_abbrv,
- A.data,
- A.elem_offset + tx * lift(local_size),
- B.data,
- B.elem_offset + tx * lift(local_size) + lift(local_size) // 2,
- C.data,
- C.elem_offset + tx * lift(local_size_out) + lift(local_size_out) // 2,
- False,
- ty=out_dtype,
- )
- )
-
- return mma_sync_desc, mma_sync_impl
-
-
-MMA_f16f16f32_INTRIN = "mma_f16f16f32"
-TensorIntrin.register(
- MMA_f16f16f32_INTRIN, *get_mma_intrin(16, "float16", "float16", "float32", False, False)
-)
-
-MMA_f16f16f32_TRANS_B_INTRIN = "mma_f16f16f32_trans_b"
-TensorIntrin.register(
- MMA_f16f16f32_TRANS_B_INTRIN, *get_mma_intrin(16, "float16", "float16", "float32", False, True)
-)
-
-MMA_f16f16f32_TRANS_A_INTRIN = "mma_f16f16f32_trans_a"
-TensorIntrin.register(
- MMA_f16f16f32_TRANS_A_INTRIN, *get_mma_intrin(16, "float16", "float16", "float32", True, False)
-)
-
-MMA_f16f16f32_TRANS_A_TRANS_B_INTRIN = "mma_f16f16f32_trans_a_trans_b"
-TensorIntrin.register(
- MMA_f16f16f32_TRANS_A_TRANS_B_INTRIN,
- *get_mma_intrin(16, "float16", "float16", "float32", True, True),
-)
-
-MMA_f16f16f16_INTRIN = "mma_f16f16f16"
-TensorIntrin.register(
- MMA_f16f16f16_INTRIN, *get_mma_intrin(16, "float16", "float16", "float16", False, False)
-)
-
-MMA_f16f16f16_TRANS_B_INTRIN = "mma_f16f16f16_trans_b"
-TensorIntrin.register(
- MMA_f16f16f16_TRANS_B_INTRIN, *get_mma_intrin(16, "float16", "float16", "float16", False, True)
-)
-
-MMA_f16f16f16_TRANS_A_INTRIN = "mma_f16f16f16_trans_a"
-TensorIntrin.register(
- MMA_f16f16f16_TRANS_A_INTRIN, *get_mma_intrin(16, "float16", "float16", "float16", True, False)
-)
-
-MMA_f16f16f16_TRANS_A_TRANS_B_INTRIN = "mma_f16f16f16_trans_a_trans_b"
-TensorIntrin.register(
- MMA_f16f16f16_TRANS_A_TRANS_B_INTRIN,
- *get_mma_intrin(16, "float16", "float16", "float16", True, True),
-)
-
-MMA_i8i8i32_INTRIN = "mma_i8i8i32"
-TensorIntrin.register(
- MMA_i8i8i32_INTRIN, *get_mma_intrin(32, "int8", "int8", "int32", False, False)
-)
-
-MMA_i8i8i32_TRANS_B_INTRIN = "mma_i8i8i32_trans_b"
-TensorIntrin.register(
- MMA_i8i8i32_TRANS_B_INTRIN, *get_mma_intrin(32, "int8", "int8", "int32", False, True)
-)
-
-MMA_e5m2e5m2f32_INTRIN = "mma_e5m2e5m2f32"
-TensorIntrin.register(
- MMA_e5m2e5m2f32_INTRIN,
- *get_mma_intrin(32, "float8_e5m2", "float8_e5m2", "float32", False, False),
-)
-
-MMA_e5m2e5m2f32_TRANS_B_INTRIN = "mma_e5m2e5m2f32_trans_b"
-TensorIntrin.register(
- MMA_e5m2e5m2f32_TRANS_B_INTRIN,
- *get_mma_intrin(32, "float8_e5m2", "float8_e5m2", "float32", False, True),
-)
-
-MMA_e4m3e4m3f32_INTRIN = "mma_e4m3e4m3f32"
-TensorIntrin.register(
- MMA_e4m3e4m3f32_INTRIN,
- *get_mma_intrin(32, "float8_e4m3fn", "float8_e4m3fn", "float32", False, False),
-)
-
-MMA_e4m3e4m3f32_TRANS_B_INTRIN = "mma_e4m3e4m3f32_trans_b"
-TensorIntrin.register(
- MMA_e4m3e4m3f32_TRANS_B_INTRIN,
- *get_mma_intrin(32, "float8_e4m3fn", "float8_e4m3fn", "float32", False, True),
-)
-
-
-def get_mma_fill_intrin(dtype, local_size):
- zero = IntImm("int32", 0).astype(dtype)
-
- # Assume M = N = 16
- index_map = shared_16x16_to_ldmatrix_32x8_layout
-
- @Ts.function
- def mma_fill_desc(C_warp: T.Tensor([WARP_SIZE, local_size], dtype=dtype, scope="warp")) -> None:
- with Ts.sblock("root"):
- Ts.reads()
- Ts.writes(C_warp[0:WARP_SIZE, 0:local_size])
- for i0, i1 in T.grid(M_DIM, N_DIM):
- with Ts.sblock("C_warp"):
- i, j = Ts.axis.remap("SS", [i0, i1])
- warp_indices = index_map(i, j)
- Ts.reads()
- Ts.writes(C_warp[warp_indices[0], warp_indices[1]])
- C_warp[warp_indices[0], warp_indices[1]] = zero
-
- @Ts.function
- def mma_fill_impl(
- C_warp: T.Tensor(
- [WARP_SIZE, local_size],
- dtype=dtype,
- scope="warp",
- offset_factor=1,
- elem_offset=T.int32(),
- ),
- ) -> None:
- with Ts.sblock("root"):
- Ts.reads()
- Ts.writes(C_warp[0:WARP_SIZE, 0:local_size])
-
- for tx in T.parallel(0, WARP_SIZE, annotations={"thread_binding": "threadIdx.x"}):
- T.evaluate(T.cuda.mma_fill(local_size, C_warp.data, C_warp.elem_offset, ty=dtype))
-
- return mma_fill_desc, mma_fill_impl
-
-
-MMA_fill_16x16_f32_INTRIN = "mma_fill_16x16_f32"
-TensorIntrin.register(MMA_fill_16x16_f32_INTRIN, *get_mma_fill_intrin("float32", 8))
-
-MMA_fill_16x16_f16_INTRIN = "mma_fill_16x16_f16"
-TensorIntrin.register(MMA_fill_16x16_f16_INTRIN, *get_mma_fill_intrin("float16", 8))
-
-MMA_fill_16x16_i32_INTRIN = "mma_fill_16x16_i32"
-TensorIntrin.register(MMA_fill_16x16_i32_INTRIN, *get_mma_fill_intrin("int32", 8))
-
-
-def get_mma_store_intrin(dtype, local_size, scope="global", use_mma_store_intrinic=True):
- # Assume M = N = 16
- index_map = shared_16x16_to_ldmatrix_32x8_layout
- index_map_rev = ldmatrix_32x8_to_shared_16x16_layout
-
- @Ts.function
- def mma_store_desc(
- C_warp: T.Tensor([WARP_SIZE, local_size], dtype=dtype, scope="warp"),
- C: T.Tensor([M_DIM, N_DIM], dtype=dtype, scope=scope),
- ) -> None:
- with Ts.sblock("root"):
- Ts.reads(C_warp[0:WARP_SIZE, 0:local_size])
- Ts.writes(C[0:M_DIM, 0:N_DIM])
- for i0, i1 in T.grid(M_DIM, N_DIM):
- with Ts.sblock("C_warp"):
- v0, v1 = Ts.axis.remap("SS", [i0, i1])
- warp_indices = index_map(v0, v1)
- Ts.reads(C_warp[warp_indices[0], warp_indices[1]])
- Ts.writes(C[v0, v1])
- C[v0, v1] = C_warp[warp_indices[0], warp_indices[1]]
-
- if use_mma_store_intrinic:
- s0 = T.dynamic("s0", "int32")
- s1 = T.dynamic("s1", "int32")
-
- @Ts.function
- def mma_store_impl(
- C_warp: T.Tensor(
- [WARP_SIZE, local_size],
- dtype=dtype,
- scope="warp",
- offset_factor=1,
- elem_offset=T.int32(),
- ),
- C: T.Tensor(
- [M_DIM, N_DIM],
- dtype=dtype,
- scope=scope,
- offset_factor=1,
- elem_offset=T.int32(),
- strides=[s0, s1],
- ),
- ) -> None:
- with Ts.sblock("root"):
- Ts.reads(C_warp[0:WARP_SIZE, 0:local_size])
- Ts.writes(C[0:M_DIM, 0:N_DIM])
-
- for tx in T.parallel(0, WARP_SIZE, annotations={"thread_binding": "threadIdx.x"}):
- T.evaluate(
- T.cuda.mma_store(
- M_DIM,
- N_DIM,
- C.ptr_to([0, 0]),
- C_warp.data,
- C_warp.elem_offset,
- s0,
- ty=dtype,
- )
- )
-
- else:
- s0 = T.dynamic("s0", "int32")
- s1 = T.dynamic("s1", "int32")
-
- @Ts.function
- def mma_store_impl(
- C_warp: T.Tensor(
- [WARP_SIZE, local_size],
- dtype=dtype,
- scope="warp",
- offset_factor=1,
- elem_offset=T.int32(),
- ),
- C: T.Tensor(
- [M_DIM, N_DIM],
- dtype=dtype,
- scope=scope,
- offset_factor=1,
- elem_offset=T.int32(),
- strides=[s0, s1],
- ),
- ) -> None:
- with Ts.sblock("root"):
- Ts.reads(C_warp[0:WARP_SIZE, 0:local_size])
- Ts.writes(C[0:M_DIM, 0:N_DIM])
-
- for tx in T.parallel(0, WARP_SIZE, annotations={"thread_binding": "threadIdx.x"}):
- for local_id in T.serial(local_size):
- matrix_indices = index_map_rev(tx, local_id)
- C[matrix_indices[0], matrix_indices[1]] = C_warp[tx, local_id]
-
- return mma_store_desc, mma_store_impl
-
-
-MMA_store_16x16_f32_global_INTRIN = "mma_store_16x16_f32_global_"
-TensorIntrin.register(
- MMA_store_16x16_f32_global_INTRIN, *get_mma_store_intrin("float32", 8, "global", True)
-)
-
-MMA_store_16x16_f32_shared_dyn_INTRIN = "mma_store_16x16_f32_shared_dyn_"
-TensorIntrin.register(
- MMA_store_16x16_f32_shared_dyn_INTRIN, *get_mma_store_intrin("float32", 8, "shared.dyn", True)
-)
-
-MMA_store_16x16_f32_shared_dyn_INTRIN_SIMPLE = "mma_store_16x16_f32_shared_dyn_simple_"
-TensorIntrin.register(
- MMA_store_16x16_f32_shared_dyn_INTRIN_SIMPLE,
- *get_mma_store_intrin("float32", 8, "shared.dyn", False),
-)
-
-MMA_store_16x16_f16_shared_dyn_INTRIN_SIMPLE = "mma_store_16x16_f16_shared_dyn_simple_"
-TensorIntrin.register(
- MMA_store_16x16_f16_shared_dyn_INTRIN_SIMPLE,
- *get_mma_store_intrin("float16", 8, "shared.dyn", False),
-)
-
-MMA_store_16x16_f16_global_INTRIN = "mma_store_16x16_f16_global_"
-TensorIntrin.register(
- MMA_store_16x16_f16_global_INTRIN, *get_mma_store_intrin("float16", 8, "global", True)
-)
-
-MMA_store_16x16_i32_global_INTRIN = "mma_store_16x16_i32_global_"
-TensorIntrin.register(
- MMA_store_16x16_i32_global_INTRIN, *get_mma_store_intrin("int32", 8, "global", True)
-)
-
-
-def get_mma_intrin_group(
- load_scope: Literal["shared", "shared.dyn"],
- store_scope: Literal["global", "shared", "shared.dyn"],
- in_dtype: Literal["float16", "int8", "float8_e4m3fn", "float8_e5m2"],
- out_dtype: Literal["float16", "float32", "int32"],
- trans_a: bool,
- trans_b: bool,
- not_use_mma_store_intrinic: bool = True,
- store_to_smem_dtype: Literal["float16", "float32", "int32"] | None = None,
-) -> dict[str, str]:
- """Get a group of intrinsics for mma tensor core with the given configurations
-
- Parameters
- ----------
- load_scope : Literal["shared", "shared.dyn"]
- The memory scope of the input tensor.
-
- store_scope : Literal["global", "shared", "shared.dyn"]
- The memory scope of the result tensor.
-
- in_dtype : str
- The input data type.
-
- out_dtype : str
- The output data dtype.
-
- trans_a : bool
- Whether the input matrix A is transposed.
-
- trans_b : bool
- Whether the input matrix B is transposed.
-
- not_use_mma_store_intrinic : bool
- Whether to not use the mma_store intrinsic. If True, use TensorStore stmts to store the
- result of mma. Otherwise, use mma_store intrinsic.
-
- This is because if we use mma_store intrinsic, during swizzling shared memory visits, our
- rearrangement scheme will involve areas accessed by different mma_store calls. This makes
- swizzling quite complex. But TensorStore will not face this problem.
-
- store_to_smem_dtype : Optional[Literal["float16", "float32", "int32"]]
- The dtype that we use to store from register to shared memory. By default it is out_dtype.
-
- Returns
- -------
- ret : Dict[str, str]
- A group of tensor intrinsics.
- """
- load_scope = MemoryScope(load_scope)
- store_scope = MemoryScope(store_scope)
- assert load_scope in [_SCOPE_SHARED, _SCOPE_SHARED_DYN]
- assert store_scope in [_SCOPE_GLOBAL, _SCOPE_SHARED, _SCOPE_SHARED_DYN]
- assert in_dtype in ["float16", "int8", "float8_e4m3fn", "float8_e5m2"]
- assert out_dtype in ["float16", "float32", "int32"]
-
- shape = "16x16"
-
- dtype_mapping = {
- "float16": "f16",
- "float32": "f32",
- "int8": "i8",
- "float8_e4m3fn": "e4m3",
- "float8_e5m2": "e5m2",
- "int32": "i32",
- }
- a_dtype = dtype_mapping[in_dtype]
- b_dtype = dtype_mapping[in_dtype]
- out_dtype = dtype_mapping[out_dtype]
-
- # e.g. mma_fill_16x16_f32
- init_intrin = f"mma_fill_{shape}_{out_dtype}"
-
- # e.g. mma_ldmatrix_f16_a_trans_dyn, mma_ldmatrix_f16_b_trans_dyn
- trans_a = "_trans" if trans_a else ""
- trans_b = "_trans" if trans_b else ""
- load_scope = "_dyn" if load_scope == _SCOPE_SHARED_DYN else ""
- load_a_intrin = f"mma_ldmatrix_{a_dtype}_a{trans_a}{load_scope}"
- load_b_intrin = f"mma_ldmatrix_{b_dtype}_b{trans_b}{load_scope}"
-
- # e.g. mma_f16f16f32_trans_a_trans_b
- trans_a_str = trans_a + "_a" if trans_a != "" else ""
- trans_b_str = trans_b + "_b" if trans_b != "" else ""
- compute_intrin = f"mma_{a_dtype}{b_dtype}{out_dtype}{trans_a_str}{trans_b_str}"
-
- # e.g. mma_store_16x16_f32_shared_dyn_simple_
- store_scope = store_scope.name.replace(".", "_")
- store_to_smem_dtype = dtype_mapping[store_to_smem_dtype] if store_to_smem_dtype else out_dtype
- suffix = "simple_" if not_use_mma_store_intrinic else ""
- store_intrin = f"mma_store_{shape}_{store_to_smem_dtype}_{store_scope}_{suffix}"
-
- return {
- "init": init_intrin,
- "load_a": load_a_intrin,
- "load_b": load_b_intrin,
- "compute": compute_intrin,
- "store": store_intrin,
- }
-
-
-######## MMA intrinsics ########
-
-
-def get_index_A(elem_offset, stride):
- i = elem_offset // stride
- j = elem_offset % stride
- stride_b = stride // 8
- bi = i // 32
- bj = j // 8
- no = bi * stride_b + bj
- return no * 8 + (i % 32) // 16 * 4
-
-
-def get_index_B(elem_offset, stride):
- i = elem_offset // stride
- j = elem_offset % stride
- stride_b = stride // 32
- bi = i // 8
- bj = j // 32
- no = bi * stride_b + bj
- return no * 8 + (j % 32) // 8 * 2
-
-
-def get_index_C(elem_offset, stride):
- i = elem_offset // stride
- j = elem_offset % stride
- stride_b = stride // 8
- bi = i // 8
- bj = j // 8
- return ((bi // 2) * 2 * stride_b + bi % 2 + bj * 2) * 2
-
-
-def get_mma_init_intrin(
- m_dim: int, n_dim: int, k_dim: int, dtype: str
-) -> tuple[Function, Function]:
- """Generator of mma init intrins"""
- del k_dim # unused
- zero = IntImm("int32", 0).astype(dtype)
- assert m_dim % 8 == 0 and n_dim % 4 == 0, "m_dim and n_dim must be multiple of 8 and 4"
- assert dtype in ["float16", "float32"]
- assert n_dim // 4 * int(dtype[-2:]) <= 128, "n_dim vectorize failed"
-
- @Ts.function
- def mma_init_desc(
- dst: T.Tensor(
- (m_dim, n_dim),
- dtype,
- align=64,
- offset_factor=1,
- elem_offset=T.int32(),
- scope="mma.matrixC",
- ),
- ) -> None:
- with Ts.sblock("root"):
- Ts.reads()
- Ts.writes(dst[0:m_dim, 0:n_dim])
- for i, j in T.grid(m_dim, n_dim):
- with Ts.sblock("init"):
- vi, vj = Ts.axis.remap("SS", [i, j])
- dst[vi, vj] = zero
-
- @Ts.function
- def mma_init_impl(
- dst: T.Tensor(
- (m_dim, n_dim),
- dtype,
- align=64,
- offset_factor=1,
- elem_offset=T.int32(),
- scope="mma.matrixC",
- ),
- ) -> None:
- with Ts.sblock("root"):
- Ts.reads()
- Ts.writes(dst[0:m_dim, 0:n_dim])
-
- for tx in T.parallel(0, WARP_SIZE, annotations={"thread_binding": "threadIdx.x"}):
- for b in range(m_dim // 8):
- for v in T.vectorized(n_dim // 4):
- dst[b * 8 + tx // 4, (tx % 4) * (n_dim // 4) + v] = zero
-
- return mma_init_desc, mma_init_impl
-
-
-def get_mma_load_intrin(
- m_dim: int,
- n_dim: int,
- k_dim: int,
- dtype: str,
- shared_scope: str,
- is_b: bool,
- is_col_major: bool,
-) -> tuple[Function, Function]:
- """Generator of mma ldmatrix intrins"""
- mma_fragment_scope = f"mma.matrix{'B' if is_b else 'A'}"
- frag_m, frag_n = (k_dim, n_dim) if is_b else (m_dim, k_dim)
- trans = (not is_col_major) if is_b else is_col_major
- if is_col_major:
- frag_m, frag_n = frag_n, frag_m
- get_index = get_index_B if is_b else get_index_A
- get_tx_index = (
- (lambda tx, s0: (tx % 8) * s0 + (tx // 8) * 8) if trans else (lambda tx, s0: tx * s0)
- )
-
- @Ts.function
- def mma_load_desc(
- src: T.Tensor(
- (frag_m, frag_n),
- dtype,
- align=64,
- offset_factor=1,
- elem_offset=T.int32(),
- scope=shared_scope,
- ),
- dst: T.Tensor(
- (frag_m, frag_n),
- dtype,
- align=64,
- offset_factor=1,
- elem_offset=T.int32(),
- scope=mma_fragment_scope,
- ),
- ) -> None:
- with Ts.sblock("root"):
- Ts.reads(src[0:frag_m, 0:frag_n])
- Ts.writes(dst[0:frag_m, 0:frag_n])
- for i, j in T.grid(frag_m, frag_n):
- with Ts.sblock("root"):
- vi, vj = Ts.axis.remap("SS", [i, j])
- dst[vi, vj] = src[vi, vj]
-
- s0 = T.dynamic("s0", "int32")
- s1 = T.dynamic("s1", "int32")
- d0 = T.dynamic("d0", "int32")
- d1 = T.dynamic("d1", "int32")
-
- @Ts.function
- def mma_load_impl(
- src: T.Tensor(
- (frag_m, frag_n),
- dtype,
- align=64,
- offset_factor=1,
- elem_offset=T.int32(),
- scope=shared_scope,
- strides=[s0, s1],
- ),
- dst: T.Tensor(
- (frag_m, frag_n),
- dtype,
- align=64,
- offset_factor=1,
- elem_offset=T.int32(),
- scope=mma_fragment_scope,
- strides=[d0, d1],
- ),
- ) -> None:
- with Ts.sblock("root"):
- Ts.reads(src[0:frag_m, 0:frag_n])
- Ts.writes(dst[0:frag_m, 0:frag_n])
-
- for tx in T.parallel(0, WARP_SIZE, annotations={"thread_binding": "threadIdx.x"}):
- T.evaluate(
- T.ptx_legacy.ldmatrix(
- trans,
- 4, # Always load 4 matrices
- ".b16",
- dst.data,
- get_index(dst.elem_offset, d0),
- src.ptr_to([0, 0]),
- get_tx_index(tx, s0),
- ty=dtype,
- )
- )
-
- return mma_load_desc, mma_load_impl
-
-
-def get_mma_sync_intrin(
- m_dim: int, n_dim: int, k_dim: int, in_dtype: str, out_dtype: str, b_transposed: bool
-) -> tuple[Function, Function]:
- """Generator of mma sync intrins"""
-
- def maybe_cast(v):
- if in_dtype != out_dtype:
- return Cast(v, out_dtype)
- return v
-
- def maybe_swap(i, j):
- if b_transposed:
- return j, i
- return i, j
-
- B_shape_0, B_shape_1 = maybe_swap(k_dim, n_dim)
-
- @Ts.function
- def mma_sync_desc(
- A: T.Tensor(
- (m_dim, k_dim),
- in_dtype,
- align=64,
- offset_factor=1,
- elem_offset=T.int32(),
- scope="mma.matrixA",
- ),
- B: T.Tensor(
- (B_shape_0, B_shape_1),
- in_dtype,
- align=64,
- offset_factor=1,
- elem_offset=T.int32(),
- scope="mma.matrixB",
- ),
- C: T.Tensor(
- (m_dim, n_dim),
- out_dtype,
- align=64,
- offset_factor=1,
- elem_offset=T.int32(),
- scope="mma.matrixC",
- ),
- ) -> None:
- with Ts.sblock("root"):
- Ts.reads(C[0:m_dim, 0:n_dim], A[0:m_dim, 0:k_dim], B[0:B_shape_0, 0:B_shape_1])
- Ts.writes(C[0:m_dim, 0:n_dim])
- for i, j, k in T.grid(m_dim, n_dim, k_dim):
- with Ts.sblock("m16n8k8_sync"):
- vi, vj, vk = Ts.axis.remap("SSR", [i, j, k])
- b_indices = maybe_swap(vk, vj)
- C[vi, vj] = C[vi, vj] + maybe_cast(A[vi, vk]) * maybe_cast(
- B[b_indices[0], b_indices[1]]
- )
-
- a0 = T.dynamic("a0", "int32")
- a1 = T.dynamic("a1", "int32")
- b0 = T.dynamic("b0", "int32")
- b1 = T.dynamic("b1", "int32")
- c0 = T.dynamic("c0", "int32")
- c1 = T.dynamic("c1", "int32")
-
- @Ts.function
- def mma_sync_impl(
- A: T.Tensor(
- (m_dim, k_dim),
- in_dtype,
- align=64,
- offset_factor=1,
- elem_offset=T.int32(),
- scope="mma.matrixA",
- strides=[a0, a1],
- ),
- B: T.Tensor(
- (B_shape_0, B_shape_1),
- in_dtype,
- align=64,
- offset_factor=1,
- elem_offset=T.int32(),
- scope="mma.matrixB",
- strides=[b0, b1],
- ),
- C: T.Tensor(
- (m_dim, n_dim),
- out_dtype,
- align=64,
- offset_factor=1,
- elem_offset=T.int32(),
- scope="mma.matrixC",
- strides=[c0, c1],
- ),
- ) -> None:
- with Ts.sblock("root"):
- Ts.reads(C[0:m_dim, 0:n_dim], A[0:m_dim, 0:k_dim], B[0:B_shape_0, 0:B_shape_1])
- Ts.writes(C[0:m_dim, 0:n_dim])
- T.evaluate(
- T.ptx_legacy.mma(
- f"m{m_dim}n{n_dim}k{k_dim}",
- "row",
- "col",
- in_dtype,
- in_dtype,
- out_dtype,
- A.data,
- get_index_A(A.elem_offset, a0),
- B.data,
- get_index_B(B.elem_offset, b0),
- C.data,
- get_index_C(C.elem_offset, c0),
- False,
- ty=out_dtype,
- )
- )
-
- return mma_sync_desc, mma_sync_impl
-
-
-def get_mma_store_dummy_intrin(
- m_dim: int, n_dim: int, k_dim: int, dtype: str
-) -> tuple[Function, Function]:
- """Disable mma store intrin for now."""
- del k_dim # unused
-
- @Ts.function
- def mma_store_desc(
- src: T.Tensor(
- (m_dim, n_dim),
- dtype,
- align=64,
- offset_factor=1,
- elem_offset=T.int32(),
- scope="mma.matrixC",
- ),
- dst: T.Tensor(
- (m_dim, n_dim),
- dtype,
- align=64,
- offset_factor=1,
- elem_offset=T.int32(),
- scope="shared.dyn",
- ),
- ) -> None:
- with Ts.sblock("root"):
- Ts.reads(src[0:m_dim, 0:n_dim])
- Ts.writes(dst[0:m_dim, 0:n_dim])
- for i, j in T.grid(m_dim, n_dim):
- with Ts.sblock("m16n8k8_store"):
- vi, vj = Ts.axis.remap("SS", [i, j])
- dst[vi, vj] = src[vi, vj]
-
- return mma_store_desc, mma_store_desc
-
-
-TensorIntrin.register("mma_init_m16n8k8_f16", *get_mma_init_intrin(16, 8, 8, "float16"))
-TensorIntrin.register("mma_init_m16n8k8_f32", *get_mma_init_intrin(16, 8, 8, "float32"))
-
-TensorIntrin.register(
- "mma_load_m16n8k8_f16_A_shared_dyn",
- *get_mma_load_intrin(32, 32, 8, "float16", "shared.dyn", False, False),
-)
-TensorIntrin.register(
- "mma_load_m16n8k8_f16_B_shared_dyn",
- *get_mma_load_intrin(32, 32, 8, "float16", "shared.dyn", True, False),
-)
-
-TensorIntrin.register(
- "mma_sync_m16n8k8_f16f16f16", *get_mma_sync_intrin(16, 8, 8, "float16", "float16", False)
-)
-TensorIntrin.register(
- "mma_sync_m16n8k8_f16f16f32", *get_mma_sync_intrin(16, 8, 8, "float16", "float32", False)
-)
-
-TensorIntrin.register(
- "mma_store_m16n8k8_f16_global", *get_mma_store_dummy_intrin(16, 8, 8, "float16")
-)
-TensorIntrin.register(
- "mma_store_m16n8k8_f32_global", *get_mma_store_dummy_intrin(16, 8, 8, "float32")
-)
-
-
-@register_global_func("tirx.index_map_m16n8k8.matrixC")
-def index_map_m16n8k8_matrixC(ind):
- i, j = ind[0], ind[1]
- return convert([(i // 8) // 2, j // 8, (i // 8) % 2, (j % 8) % 2])
diff --git a/python/tvm/s_tir/tensor_intrin/rocm.py b/python/tvm/s_tir/tensor_intrin/rocm.py
index ad4b11bf9f..318f1c0826 100644
--- a/python/tvm/s_tir/tensor_intrin/rocm.py
+++ b/python/tvm/s_tir/tensor_intrin/rocm.py
@@ -15,7 +15,16 @@
# specific language governing permissions and limitations
# under the License.
# pylint: disable=invalid-name,missing-function-docstring
-"""Intrinsics for AMDGPU tensorization."""
+"""Intrinsics for AMDGPU tensorization.
+
+Descriptors use a logical (lane, value) layout to match a whole wavefront's
+matrix operation. Implementations bind only (1, values_per_lane) in local
+storage: every lane owns its own fragment, indexed at lane zero. The hardware
+lane id selects the shared/global matrix coordinates for loads and stores;
+MFMA communicates the register operands across lanes. Binding a singleton
+lane dimension also lets CompactTensorAllocation remove the logical 64-lane
+extent without a warp-memory lowering pass.
+"""
from tvm.runtime import convert
from tvm.script import s_tir as Ts
@@ -60,7 +69,7 @@ N_DIM = 16
def shared_16x4_to_local_64x1_layout_A(i, j):
thread_id = j * 16 + i
- return thread_id, convert(0)
+ return thread_id, IntImm("int32", 0)
def thread_id_shared_access_64x1_to_16x4_layout_A(thread_id, local_id):
@@ -71,7 +80,7 @@ def thread_id_shared_access_64x1_to_16x4_layout_A(thread_id, local_id):
def shared_4x16_to_local_64x1_layout_B(i, j):
thread_id = i * 16 + j
- return thread_id, convert(0)
+ return thread_id, IntImm("int32", 0)
def thread_id_shared_access_64x1_to_4x16_layout_B(thread_id, local_id):
@@ -123,7 +132,9 @@ def get_mma_fill_intrin(dtype, local_size):
index_map = shared_16x16_to_local_64x4_layout_C
@Ts.function
- def mma_fill_desc(C_warp: T.Tensor([WARP_SIZE, local_size], dtype=dtype, scope="warp")) -> None:
+ def mma_fill_desc(
+ C_warp: T.Tensor([WARP_SIZE, local_size], dtype=dtype, scope="local"),
+ ) -> None:
with Ts.sblock("root"):
Ts.reads()
Ts.writes(C_warp[0:WARP_SIZE, 0:local_size])
@@ -138,19 +149,19 @@ def get_mma_fill_intrin(dtype, local_size):
@Ts.function
def mma_fill_impl(
C_warp: T.Tensor(
- [WARP_SIZE, local_size],
+ [1, local_size],
dtype=dtype,
- scope="warp",
+ scope="local",
offset_factor=1,
elem_offset=T.int32(),
),
) -> None:
with Ts.sblock("root"):
Ts.reads()
- Ts.writes(C_warp[0:WARP_SIZE, 0:local_size])
- tx = T.launch_thread("threadIdx.x", WARP_SIZE)
+ Ts.writes(C_warp[0:1, 0:local_size])
+ _tx = T.launch_thread("threadIdx.x", WARP_SIZE)
for local_id in T.serial(0, local_size):
- C_warp[tx, local_id] = zero
+ C_warp[0, local_id] = zero
return mma_fill_desc, mma_fill_impl
@@ -204,7 +215,7 @@ def get_mfma_load_intrin(
@Ts.function
def mfma_load_desc(
reg: T.Tensor(
- (WARP_SIZE, local_size), dtype, offset_factor=1, elem_offset=T.int32(), scope="warp"
+ (WARP_SIZE, local_size), dtype, offset_factor=1, elem_offset=T.int32(), scope="local"
),
memory: T.Tensor(memory_shape, dtype, offset_factor=1, elem_offset=T.int32(), scope=scope),
) -> None:
@@ -227,12 +238,12 @@ def get_mfma_load_intrin(
@Ts.function
def mfma_load_impl(
reg: T.Tensor(
- (WARP_SIZE, local_size),
+ (1, local_size),
dtype,
align=64,
offset_factor=1,
elem_offset=T.int32(),
- scope="warp",
+ scope="local",
),
memory: T.Tensor(
memory_shape,
@@ -246,11 +257,11 @@ def get_mfma_load_intrin(
) -> None:
with Ts.sblock("root"):
Ts.reads(memory[0:row_dim, 0:col_dim])
- Ts.writes(reg[0:WARP_SIZE, 0:local_size])
+ Ts.writes(reg[0:1, 0:local_size])
for local_id in T.serial(0, local_size):
tx = T.launch_thread("threadIdx.x", WARP_SIZE)
matrix_indices = reverse_index_map(tx, local_id)
- reg[tx, local_id] = memory[matrix_indices[0], matrix_indices[1]]
+ reg[0, local_id] = memory[matrix_indices[0], matrix_indices[1]]
return mfma_load_desc, mfma_load_impl
@@ -288,17 +299,17 @@ def get_mfma_intrin(k_dim, in_dtype="float32", out_dtype="float32", b_transposed
@Ts.function
def mfma_sync_desc(
A: T.Tensor(
- (WARP_SIZE, local_size), in_dtype, offset_factor=1, elem_offset=T.int32(), scope="warp"
+ (WARP_SIZE, local_size), in_dtype, offset_factor=1, elem_offset=T.int32(), scope="local"
),
B: T.Tensor(
- (WARP_SIZE, local_size), in_dtype, offset_factor=1, elem_offset=T.int32(), scope="warp"
+ (WARP_SIZE, local_size), in_dtype, offset_factor=1, elem_offset=T.int32(), scope="local"
),
C: T.Tensor(
(WARP_SIZE, local_size_out),
out_dtype,
offset_factor=1,
elem_offset=T.int32(),
- scope="warp",
+ scope="local",
),
) -> None:
with Ts.sblock("root"):
@@ -332,32 +343,32 @@ def get_mfma_intrin(k_dim, in_dtype="float32", out_dtype="float32", b_transposed
@Ts.function
def mfma_sync_impl_float(
A: T.Tensor(
- (WARP_SIZE, local_size), in_dtype, offset_factor=1, elem_offset=T.int32(), scope="warp"
+ (1, local_size), in_dtype, offset_factor=1, elem_offset=T.int32(), scope="local"
),
B: T.Tensor(
- (WARP_SIZE, local_size), in_dtype, offset_factor=1, elem_offset=T.int32(), scope="warp"
+ (1, local_size), in_dtype, offset_factor=1, elem_offset=T.int32(), scope="local"
),
C: T.Tensor(
- (WARP_SIZE, local_size_out),
+ (1, local_size_out),
out_dtype,
offset_factor=1,
elem_offset=T.int32(),
- scope="warp",
+ scope="local",
),
) -> None:
with Ts.sblock("root"):
Ts.reads(
- A[0:WARP_SIZE, 0:local_size],
- B[0:WARP_SIZE, 0:local_size],
- C[0:WARP_SIZE, 0:local_size_out],
+ A[0:1, 0:local_size],
+ B[0:1, 0:local_size],
+ C[0:1, 0:local_size_out],
)
- Ts.writes(C[0:WARP_SIZE, 0:local_size_out])
- tx = T.launch_thread("threadIdx.x", WARP_SIZE)
- C[tx, T.ramp(0, 1, local_size_out)] = T.call_llvm_pure_intrin(
+ Ts.writes(C[0:1, 0:local_size_out])
+ _tx = T.launch_thread("threadIdx.x", WARP_SIZE)
+ C[0, T.ramp(0, 1, local_size_out)] = T.call_llvm_pure_intrin(
T.llvm_lookup_intrinsic_id(mfma_intrin),
- A[tx, T.ramp(0, 1, local_size) if T.constexpr(local_size > 1) else 0],
- B[tx, T.ramp(0, 1, local_size) if T.constexpr(local_size > 1) else 0],
- C[tx, T.ramp(0, 1, local_size_out)],
+ A[0, T.ramp(0, 1, local_size) if T.constexpr(local_size > 1) else 0],
+ B[0, T.ramp(0, 1, local_size) if T.constexpr(local_size > 1) else 0],
+ C[0, T.ramp(0, 1, local_size_out)],
T.int32(0),
T.int32(0),
T.int32(0),
@@ -367,41 +378,41 @@ def get_mfma_intrin(k_dim, in_dtype="float32", out_dtype="float32", b_transposed
@Ts.function
def mfma_sync_impl_integer(
A: T.Tensor(
- (WARP_SIZE, local_size), in_dtype, offset_factor=1, elem_offset=T.int32(), scope="warp"
+ (1, local_size), in_dtype, offset_factor=1, elem_offset=T.int32(), scope="local"
),
B: T.Tensor(
- (WARP_SIZE, local_size), in_dtype, offset_factor=1, elem_offset=T.int32(), scope="warp"
+ (1, local_size), in_dtype, offset_factor=1, elem_offset=T.int32(), scope="local"
),
C: T.Tensor(
- (WARP_SIZE, local_size_out),
+ (1, local_size_out),
out_dtype,
offset_factor=1,
elem_offset=T.int32(),
- scope="warp",
+ scope="local",
),
) -> None:
with Ts.sblock("root"):
Ts.reads(
- A[0:WARP_SIZE, 0:local_size],
- B[0:WARP_SIZE, 0:local_size],
- C[0:WARP_SIZE, 0:local_size_out],
+ A[0:1, 0:local_size],
+ B[0:1, 0:local_size],
+ C[0:1, 0:local_size_out],
)
- Ts.writes(C[0:WARP_SIZE, 0:local_size_out])
- tx = T.launch_thread("threadIdx.x", WARP_SIZE)
+ Ts.writes(C[0:1, 0:local_size_out])
+ _tx = T.launch_thread("threadIdx.x", WARP_SIZE)
- C[tx, T.ramp(0, 1, local_size_out)] = T.call_llvm_pure_intrin(
+ C[0, T.ramp(0, 1, local_size_out)] = T.call_llvm_pure_intrin(
T.llvm_lookup_intrinsic_id(mfma_intrin),
T.call_intrin(
"tirx.reinterpret",
- A[tx, T.ramp(0, 1, local_size) if T.constexpr(local_size > 1) else 0],
+ A[0, T.ramp(0, 1, local_size) if T.constexpr(local_size > 1) else 0],
ty="int32",
),
T.call_intrin(
"tirx.reinterpret",
- B[tx, T.ramp(0, 1, local_size) if T.constexpr(local_size > 1) else 0],
+ B[0, T.ramp(0, 1, local_size) if T.constexpr(local_size > 1) else 0],
ty="int32",
),
- C[tx, T.ramp(0, 1, local_size_out)],
+ C[0, T.ramp(0, 1, local_size_out)],
T.int32(0),
T.int32(0),
T.int32(0),
@@ -420,7 +431,7 @@ def get_mfma_store_intrin(local_size=4, dtype="float32", scope="global"):
@Ts.function
def mfma_store_desc(
- C_warp: T.Tensor([WARP_SIZE, local_size], dtype=dtype, scope="warp"),
+ C_warp: T.Tensor([WARP_SIZE, local_size], dtype=dtype, scope="local"),
C: T.Tensor([M_DIM, N_DIM], dtype=dtype, scope=scope),
) -> None:
with Ts.sblock("root"):
@@ -440,9 +451,9 @@ def get_mfma_store_intrin(local_size=4, dtype="float32", scope="global"):
@Ts.function
def mfma_store_impl(
C_warp: T.Tensor(
- [WARP_SIZE, local_size],
+ [1, local_size],
dtype=dtype,
- scope="warp",
+ scope="local",
offset_factor=1,
elem_offset=T.int32(),
),
@@ -456,11 +467,11 @@ def get_mfma_store_intrin(local_size=4, dtype="float32", scope="global"):
),
) -> None:
with Ts.sblock("root"):
- Ts.reads(C_warp[0:WARP_SIZE, 0:local_size])
+ Ts.reads(C_warp[0:1, 0:local_size])
Ts.writes(C[0:M_DIM, 0:N_DIM])
tx = T.launch_thread("threadIdx.x", WARP_SIZE)
for i in range(local_size):
- C[((tx // 16) * 4) + i, (tx % 16)] = C_warp[tx, i]
+ C[((tx // 16) * 4) + i, (tx % 16)] = C_warp[0, i]
return mfma_store_desc, mfma_store_impl
diff --git a/python/tvm/s_tir/testing/tir.py b/python/tvm/s_tir/testing/tir.py
index af5c8839a2..89e2ff3901 100644
--- a/python/tvm/s_tir/testing/tir.py
+++ b/python/tvm/s_tir/testing/tir.py
@@ -18,118 +18,6 @@
"""Common utility functions in TVM tirx"""
-def mma_schedule(
- workload,
- k_inner,
- in_dtype,
- b_transposed,
- i_factors,
- j_factors,
- k_factors,
- index_map_A,
- index_map_B,
- index_map_C,
- ldmatrix_a_intrin,
- ldmatrix_b_intrin,
- mma_intrin,
- mma_fill_intrin,
- mma_store_intrin,
- shared_scope="shared",
-):
- """Create a tensorized schedule for GEMM with MMA intrinsics."""
- import tvm # pylint: disable=import-outside-toplevel
-
- ir_module = tvm.IRModule({"main": workload})
- sch = tvm.s_tir.Schedule(ir_module)
-
- block = sch.get_sblock("C")
- i, j, k = sch.get_loops(block)
- i, i_tc = sch.split(i, factors=[None, 16])
- j, j_tc = sch.split(j, factors=[None, 16])
- k, k_tc = sch.split(k, factors=[None, k_inner])
-
- sch.reorder(i, j, k, i_tc, j_tc, k_tc)
-
- block_inner = sch.blockize(i_tc)
- block_outer, block_inner = block_inner, block
-
- num_ty = i_factors[2] * j_factors[2]
-
- i0, i1, i2, i3, i4 = sch.split(i, factors=i_factors)
- j0, j1, j2, j3, j4 = sch.split(j, factors=j_factors)
- k0, k1, k2 = sch.split(k, k_factors)
-
- sch.reorder(i0, j0, i1, j1, j2, i2, k0, k1, i3, j3, k2, i4, j4)
-
- block_idx = sch.fuse(i0, j0)
- block_idy = sch.fuse(i1, j1)
- thread_idy = sch.fuse(j2, i2)
- sch.bind(block_idx, "blockIdx.x")
- sch.bind(block_idy, "blockIdx.y")
- sch.bind(thread_idy, "threadIdx.y")
-
- def fetch_to_shared(block, idx, ndim):
- block_read = sch.cache_read(block, idx, shared_scope)
- sch.compute_at(block_read, k0)
- vector_size = 16 if in_dtype == "int8" else 8
- warp_size = 32
- fused = sch.fuse(*sch.get_loops(block_read)[-ndim:])
- _, f_1, f_2, f_3 = sch.split(fused, factors=[None, num_ty, warp_size, vector_size])
- sch.bind(f_2, "threadIdx.x")
- sch.bind(f_1, "threadIdx.y")
- sch.vectorize(f_3)
- offset = 8 if in_dtype == "float16" else 16
- sch.storage_align(block_read, 0, axis=-2, factor=32, offset=offset)
-
- return block_read
-
- fetch_to_shared(block_outer, 0, 2)
- fetch_to_shared(block_outer, 1, 2)
-
- A_warp = sch.cache_read(block_outer, 0, "warp")
- B_warp = sch.cache_read(block_outer, 1, "warp")
-
- sch.compute_at(A_warp, k1)
- sch.compute_at(B_warp, k1)
-
- C_warp = sch.cache_write(block_outer, 0, "warp")
- sch.reverse_compute_at(C_warp, thread_idy)
-
- ii, jj = sch.get_loops(C_warp)[-2:]
- io, ii = sch.split(ii, factors=[None, 16])
- jo, ji = sch.split(jj, factors=[None, 16])
- sch.reorder(io, jo, ii, ji)
-
- sch.decompose_reduction(block_outer, sch.get_loops(block_outer)[3])
- block_init_c = sch.get_sblock("C_init")
-
- def tile_wmma_fragment(block_read, height, width):
- i, j = sch.get_loops(block_read)[-2:]
- i0, i1 = sch.split(i, factors=[None, height])
- j0, j1 = sch.split(j, factors=[None, width])
- sch.reorder(i0, j0, i1, j1)
- return i1
-
- loop_a = tile_wmma_fragment(A_warp, 16, k_inner)
-
- if b_transposed:
- loop_b = tile_wmma_fragment(B_warp, 16, k_inner)
- else:
- loop_b = tile_wmma_fragment(B_warp, k_inner, 16)
-
- sch.transform_layout(A_warp, ("write", 0), index_map_A)
- sch.transform_layout(B_warp, ("write", 0), index_map_B)
- sch.transform_layout(C_warp, ("read", 0), index_map_C)
-
- sch.tensorize(loop_a, ldmatrix_a_intrin)
- sch.tensorize(loop_b, ldmatrix_b_intrin)
- sch.tensorize(sch.get_loops(block_inner)[-3], mma_intrin)
- sch.tensorize(sch.get_loops(block_init_c)[-2], mma_fill_intrin)
- sch.tensorize(sch.get_loops(C_warp)[-2], mma_store_intrin)
-
- return sch
-
-
def mfma_schedule(
workload,
k_inner,
@@ -198,12 +86,13 @@ def mfma_schedule(
fetch_to_shared(block_outer, 0, 2)
fetch_to_shared(block_outer, 1, 2)
- A_warp = sch.cache_read(block_outer, 0, "warp")
- B_warp = sch.cache_read(block_outer, 1, "warp")
+ # Keep the logical lane axis until tensorization binds per-lane local fragments.
+ A_warp = sch.cache_read(block_outer, 0, "local")
+ B_warp = sch.cache_read(block_outer, 1, "local")
sch.compute_at(A_warp, k1)
sch.compute_at(B_warp, k1)
- C_warp = sch.cache_write(block_outer, 0, "warp")
+ C_warp = sch.cache_write(block_outer, 0, "local")
sch.reverse_compute_at(C_warp, thread_idy)
ii, jj = sch.get_loops(C_warp)[-2:]
diff --git a/python/tvm/tirx/compilation_pipeline.py b/python/tvm/tirx/compilation_pipeline.py
index 8f63fae2b9..b40d140de4 100644
--- a/python/tvm/tirx/compilation_pipeline.py
+++ b/python/tvm/tirx/compilation_pipeline.py
@@ -146,7 +146,6 @@ def finalize_host_passes(): # pylint: disable=unused-argument
def finalize_device_passes(): # pylint: disable=unused-argument
"""The default finalization passes for TIR backend."""
device_pass_list = [
- tirx.transform.LowerWarpMemory(),
tirx.transform.StmtSimplify(),
tirx.transform.LowerIntrin(),
]
diff --git a/python/tvm/tirx/transform/transform.py b/python/tvm/tirx/transform/transform.py
index ffe4d1ba56..2f9ca83ce6 100644
--- a/python/tvm/tirx/transform/transform.py
+++ b/python/tvm/tirx/transform/transform.py
@@ -316,17 +316,6 @@ def SkipAssert():
return _ffi_api.SkipAssert() # type: ignore
-def LowerWarpMemory():
- """Lower warp memory access to low-level device related function calls.
-
- Returns
- -------
- fpass : tvm.ir.transform.Pass
- The result pass
- """
- return _ffi_api.LowerWarpMemory() # type: ignore
-
-
def LowerTVMBuiltin():
"""Lower tvm builtin intrinsics.
diff --git a/src/backend/cuda/codegen/codegen_cuda.cc b/src/backend/cuda/codegen/codegen_cuda.cc
index 6ce6861f06..032c81a6ca 100644
--- a/src/backend/cuda/codegen/codegen_cuda.cc
+++ b/src/backend/cuda/codegen/codegen_cuda.cc
@@ -31,7 +31,6 @@
#include <tvm/runtime/memory_scope.h>
#include <tvm/support/gpu_thread_tag.h>
#include <tvm/sym/analyzer.h>
-#include <tvm/tirx/index_map.h>
#include <tvm/tirx/op/gpu.h>
#include <tvm/tirx/op/memory.h>
#include <tvm/tirx/op/region.h>
@@ -48,7 +47,6 @@
#include "cuda_fallback_module.h"
#include "literal/cuda_half_t.h"
#include "literal/cuda_int8_t.h"
-#include "ptx.h"
namespace tvm {
namespace codegen {
@@ -56,14 +54,6 @@ using namespace tvm::prim;
namespace {
-bool IsOp(const CallNode* call, const Op& compat_op, const char* canonical_name) {
- if (call->op.same_as(compat_op)) {
- return true;
- }
- const auto* op_node = call->op.as<OpNode>();
- return op_node != nullptr && op_node->name == canonical_name;
-}
-
TVM_FFI_INLINE bool IsFloat8(const PrimType& ty) {
return ty.MatchesCode(DLDataTypeCode::kDLFloat8_e3m4, DLDataTypeCode::kDLFloat8_e4m3,
DLDataTypeCode::kDLFloat8_e4m3b11fnuz, DLDataTypeCode::kDLFloat8_e4m3fn,
@@ -917,6 +907,8 @@ void CodeGenCUDA::PrintStorageScope(const runtime::MemoryScope& scope,
TVM_FFI_ICHECK(scope != runtime::MemoryScope::Global())
<< "Cannot allocate global memory when targeting CUDA. You must pass "
"all global arrays as input instead";
+ TVM_FFI_ICHECK(scope != runtime::MemoryScope::Get("warp"))
+ << "CUDA warp storage is unsupported; use explicit per-thread local storage and shuffles";
if (scope == scope_shared) {
os << "__shared__ ";
} else if (scope == scope_shared_dyn) {
@@ -1128,126 +1120,9 @@ void CodeGenCUDA::Dispatch_(const CallNode* op, std::ostream& os) {
}
}
- static const Op mma_store_op = Op::Get("tirx.cuda.mma_store");
- static const Op mma_fill_op = Op::Get("tirx.cuda.mma_fill");
- static const Op ptx_mma_legacy_op = Op::Get("tirx.ptx_legacy.mma");
- static const Op ptx_ldmatrix_legacy_op = Op::Get("tirx.ptx_legacy.ldmatrix");
static const Op cuda_func_call_op = Op::Get("tirx.cuda.func_call");
- if (op->op.same_as(mma_store_op)) {
- int m = op->args[0].as_or_throw<IntImm>()->value.as<int>().value();
- int n = op->args[1].as_or_throw<IntImm>()->value.as<int>().value();
- std::string dst = this->PrintExpr(op->args[2]);
- std::string src = this->PrintExpr(op->args[3]);
- std::string src_offset = this->PrintExpr(op->args[4]);
- PrimExpr stride = op->args[5].as_or_throw<PrimExpr>();
-
- TVM_FFI_ICHECK(m == 16 && n == 16) << "Only m == 16 && n == 16 case supported for now";
-
- // Each thread in a warp holds a certain number of elements of an MMA output.
- // For example, if we compute a 16x16 tile using MMA, each thread holds 8 elements
- // in its registers. So conceptually, a warp memory is organized as a 32x8 block.
- // A map from a 16x16 tile to a 32x8 block of memory is specified by the index map below.
-
- // To store the 32x8 output back to a 16x16 tile in shared or global memory, we invert this map
- // to determine the output location for each 8 element.
-
- const auto index_map_func =
- tvm::ffi::Function::GetGlobal("tirx.index_map.shared_16x16_to_ldmatrix_32x8_layout");
- TVM_FFI_ICHECK(index_map_func.has_value());
-
- sym::Analyzer analyzer;
- auto inverse_index_map =
- IndexMap::FromFunc(2, *index_map_func).Inverse({Range(0, m), Range(0, n)}, analyzer);
- auto indices_16x16 = inverse_index_map->final_indices;
-
- // "//" and "%" in the index map are translated to FloorDiv/Mod, but the plain Div/Mod are fine.
- // FloorDiv/Mod are supposed to be lowered before they reach codegen, so manually replace them
- // to the plain ones here.
- class LowerFloorDivMod : public tirx::StmtExprMutator {
- public:
- UnchangedOr<PrimExpr> Mutate_(const prim::FloorDivNode* op, InplaceMode inplace_mode) {
- return prim::Div(this->Mutate(op->a, inplace_mode).ValueOrUnchanged(op->a),
- this->Mutate(op->b, inplace_mode).ValueOrUnchanged(op->b));
- }
- UnchangedOr<PrimExpr> Mutate_(const prim::FloorModNode* op, InplaceMode inplace_mode) {
- return prim::Mod(this->Mutate(op->a, inplace_mode).ValueOrUnchanged(op->a),
- this->Mutate(op->b, inplace_mode).ValueOrUnchanged(op->b));
- }
- };
-
- PrimExpr dst_expr = indices_16x16[0] * stride + indices_16x16[1];
- auto dst_ind =
- ffi::make_object<LowerFloorDivMod>()->Mutate(dst_expr).ValueOrUnchanged(dst_expr);
-
- var_idmap_[inverse_index_map->initial_indices[0].get()] = "threadIdx.x";
- var_idmap_[inverse_index_map->initial_indices[1].get()] = "local_id";
-
- os << "for (int local_id = 0; local_id < 8; ++local_id) {\n";
- os << dst << "[" + this->PrintExpr(dst_ind) + "] = " << src << "[" << src_offset
- << " + local_id];\n";
- os << "}\n";
-
- } else if (op->op.same_as(mma_fill_op)) {
- std::string num_elem = this->PrintExpr(op->args[0]);
- std::string dst = this->PrintExpr(op->args[1]);
- std::string dst_offset = this->PrintExpr(op->args[2]);
-
- os << "for (int i = 0; i < " << num_elem << "; ++i) {\n";
- os << dst << "[" << dst_offset << " + i] = 0.0;";
- os << "}\n";
- } else if (IsOp(op, ptx_mma_legacy_op, "tirx.ptx_legacy.mma")) {
- // args: shape, A_layout, B_layout, A_dtype, B_dtype, C_dtype,
- // a_ptr_var, a_offset, b_ptr_var, b_offset,
- // c_ptr_var, c_offset, saturate, [bit_op]
- codegen_tags_.insert("mma");
- TVM_FFI_ICHECK(op->args.size() == 13U || op->args.size() == 14U);
- std::string shape = op->args[0].as_or_throw<StringImm>()->value;
- std::string A_layout = op->args[1].as_or_throw<StringImm>()->value;
- std::string B_layout = op->args[2].as_or_throw<StringImm>()->value;
- std::string A_dtype = op->args[3].as_or_throw<StringImm>()->value;
- std::string B_dtype = op->args[4].as_or_throw<StringImm>()->value;
- std::string C_dtype = op->args[5].as_or_throw<StringImm>()->value;
- std::string a_ref = this->PrintExpr(op->args[6]);
- std::string a_bias = this->PrintExpr(op->args[7]);
- std::string b_ref = this->PrintExpr(op->args[8]);
- std::string b_bias = this->PrintExpr(op->args[9]);
- std::string c_ref = this->PrintExpr(op->args[10]);
- std::string c_bias = this->PrintExpr(op->args[11]);
- bool saturate = op->args[12].as_or_throw<IntImm>()->value != 0;
- std::string bit_op = op->args.size() > 13 ? op->args[13].as_or_throw<StringImm>()->value : "";
- this->stream << PrintMMAAssembly(shape, A_layout, B_layout, A_dtype, B_dtype, C_dtype, a_ref,
- a_bias, b_ref, b_bias, c_ref, c_bias, bit_op, saturate);
- } else if (IsOp(op, ptx_ldmatrix_legacy_op, "tirx.ptx_legacy.ldmatrix")) {
- // args: trans, num, type, local_ptr_var, local_offset, smem_ptr_var, smem_offset
- codegen_tags_.insert("mma");
- TVM_FFI_ICHECK_EQ(op->args.size(), 7U);
- // `trans` and `num` may arrive as Bool/IntImm; both Downcastable
- // to PrimExpr whose IntImmNode value tells us the literal.
- bool trans = op->args[0].as_or_throw<IntImm>()->value != 0;
- int num = op->args[1].as_or_throw<IntImm>()->value.as<int>().value();
- std::string type_str = op->args[2].as_or_throw<StringImm>()->value;
- std::string local_ptr = this->PrintExpr(op->args[3]);
- std::string local_offset = this->PrintExpr(op->args[4]);
- std::string smem_ptr = this->PrintExpr(op->args[5]);
- PrimType res_ty = op->ty.as_or_throw<PrimType>();
- if (trans && res_ty.bits() == 8) {
- // ldmatrix can't transpose 8-bit elements (it assumes 16-bit), so
- // synthesize the equivalent manual gather loop. args[6] is the
- // shared-memory stride for this fallback.
- std::string smem_stride = this->PrintExpr(op->args[6]);
- TVM_FFI_ICHECK(num == 4);
- os << "for (int i = 0; i < 16; ++i) {\n";
- os << local_ptr << "[" + local_offset + " + i] = " << smem_ptr
- << "[(i % 8) / 4 * " + smem_stride + " * 16 + (threadIdx.x % 4) * 4 * " + smem_stride +
- "+ (i % 4) * " + smem_stride + " + threadIdx.x / 4 + (i / 8) * 8];\n";
- os << "}\n";
- } else {
- std::string smem_offset = this->PrintExpr(op->args[6]);
- this->stream << PrintLoadMatrixAssembly(trans, num, type_str, local_ptr, local_offset,
- smem_ptr, smem_offset);
- }
- } else if (op->op.same_as(tirx::reinterpret_op())) {
+ if (op->op.same_as(tirx::reinterpret_op())) {
// Compile-time pointer reinterpret of a literal (e.g. a tcgen05 descriptor
// template encoded at address 0): emit C++-style reinterpret_cast<T*>(...)
// to match the encoded template. Runtime pointer reinterprets fall through
diff --git a/src/backend/cuda/codegen/ptx.cc b/src/backend/cuda/codegen/ptx.cc
index 6ffb43b21e..d50b97789c 100644
--- a/src/backend/cuda/codegen/ptx.cc
+++ b/src/backend/cuda/codegen/ptx.cc
@@ -21,15 +21,10 @@
* \file ptx.cc
*/
-#include "ptx.h"
-
+#include <tvm/ffi/error.h>
#include <tvm/ffi/reflection/registry.h>
-#include <algorithm>
#include <string>
-#include <tuple>
-#include <utility>
-#include <vector>
namespace tvm {
namespace codegen {
@@ -80,8 +75,6 @@ static const char* dtype_str[] = {".s4", ".u4", ".s8", ".u8", ".s16",
".u32", ".s64", ".u64", ".e2m1", ".e2m3", ".e3m2", ".e4m3",
".ue4m3", ".e5m2", ".ue8m0", ".f16", ".bf16", ".f16x2", ".f32",
".tf32", ".f64", ".b1", ".b8", ".b16", ".b32", ".b64"};
-static const uint32_t num_bits[] = {4, 4, 8, 8, 16, 16, 32, 32, 64, 64, 4, 6, 6, 8,
- 7, 8, 8, 16, 16, 32, 32, 32, 64, 1, 8, 16, 32, 64};
/*!
* \brief Create PTX data type from string.
@@ -170,489 +163,6 @@ void RegisterCudaPTXHelpers() {
TVM_FFI_STATIC_INIT_BLOCK() { RegisterCudaPTXHelpers(); }
-/*!
- * \brief Get the number of bits of given PTX data type.
- */
-inline uint32_t DTypeBits(DataType dtype) { return num_bits[static_cast<int>(dtype)]; }
-
-/*!
- * \brief Extract the value m, n, k from string m*n*k*
- */
-inline std::tuple<int, int, int> ParseMMAShape(const std::string& str) {
- size_t pos_m = str.find("m"), pos_n = str.find("n"), pos_k = str.find("k");
- TVM_FFI_ICHECK(pos_m != str.npos && pos_n != str.npos && pos_k != str.npos)
- << "Cannot parse MMA shape " << str;
- int m = std::stoi(str.substr(pos_m + 1, pos_n - pos_m - 1)),
- n = std::stoi(str.substr(pos_n + 1, pos_k - pos_n - 1)), k = std::stoi(str.substr(pos_k + 1));
- return std::make_tuple(m, n, k);
-}
-
-/*!
- * \brief Layout Type
- */
-enum class LayoutType : int { kRowMajor = 0, kColumnMajor = 1 };
-
-/*!
- * \brief Parse layout type
- */
-LayoutType LayoutTypeFromString(const std::string& str) {
- if (str == "row") {
- return LayoutType::kRowMajor;
- } else if (str == "col") {
- return LayoutType::kColumnMajor;
- } else {
- TVM_FFI_THROW(InternalError) << "Unrecognized layout type " << str;
- }
-}
-
-static const char* layout_type_str[] = {"row", "col"};
-
-/*!
- * \brief Convert layout type to string.
- */
-inline std::string LayoutTypeToString(LayoutType layout) {
- return layout_type_str[static_cast<int>(layout)];
-}
-
-/*!
- * \brief MMA Configurations, used to determine validity.
- */
-struct MMAConfig {
- explicit MMAConfig(int m, int n, int k, DataType dtype_mul, bool use_bit_op)
- : m(m), n(n), k(k), dtype_mul(dtype_mul), use_bit_op(use_bit_op) {}
- int m, n, k;
- DataType dtype_mul;
- bool use_bit_op;
- inline bool operator==(const MMAConfig& other) {
- return m == other.m && n == other.n && k == other.k && dtype_mul == other.dtype_mul &&
- use_bit_op == other.use_bit_op;
- }
-};
-
-/*!
- * \brief Valid MMA configurations
- * \note Reference:
- * https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#warp-level-matrix-shape
- */
-const MMAConfig valid_mma_configs[] = {
- MMAConfig(8, 8, 4, DataType::kFloat64, false),
- MMAConfig(8, 8, 4, DataType::kFloat16, false),
- MMAConfig(16, 8, 8, DataType::kFloat16, false),
- MMAConfig(16, 8, 16, DataType::kFloat16, false),
- MMAConfig(16, 8, 8, DataType::kBFloat16, false),
- MMAConfig(16, 8, 16, DataType::kBFloat16, false),
- MMAConfig(16, 8, 4, DataType::kTensorFloat32, false),
- MMAConfig(16, 8, 8, DataType::kTensorFloat32, false),
- MMAConfig(8, 8, 16, DataType::kInt8, false),
- MMAConfig(16, 8, 16, DataType::kInt8, false),
- MMAConfig(16, 8, 32, DataType::kInt8, false),
- MMAConfig(8, 8, 16, DataType::kUInt8, false),
- MMAConfig(16, 8, 16, DataType::kUInt8, false),
- MMAConfig(16, 8, 32, DataType::kUInt8, false),
- MMAConfig(8, 8, 32, DataType::kInt4, false),
- MMAConfig(16, 8, 32, DataType::kInt4, false),
- MMAConfig(16, 8, 64, DataType::kInt4, false),
- MMAConfig(8, 8, 32, DataType::kUInt4, false),
- MMAConfig(16, 8, 32, DataType::kUInt4, false),
- MMAConfig(16, 8, 64, DataType::kUInt4, false),
- MMAConfig(8, 8, 128, DataType::kBit1, true),
- MMAConfig(16, 8, 128, DataType::kBit1, true),
- MMAConfig(16, 8, 256, DataType::kBit1, true),
- MMAConfig(16, 8, 32, DataType::kFloat8_e4m3fn, false),
- MMAConfig(16, 8, 32, DataType::kFloat8_e5m2, false),
-};
-
-/*!
- * \brief Check whether the multiplicand data type and accumulator data type is valid for MMA
- * computation.
- * \param dtype_a The data type of multiplicand a.
- * \param dtype_b The data type of multiplicand b.
- * \param dtype_c The data type of accumulator c.
- * \note Reference:
- * https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#warp-level-matrix-data-types
- */
-void CheckMMADTypeCompatible(DataType dtype_a, DataType dtype_b, DataType dtype_c) {
- std::string ab_not_match_err_str = "The multiplicands' data type " + DTypeToString(dtype_a) +
- DTypeToString(dtype_b) + " do not match.";
- // check a and b
- switch (dtype_a) {
- case DataType::kBit1:
- case DataType::kFloat16:
- case DataType::kBFloat16:
- case DataType::kTensorFloat32:
- case DataType::kFloat64:
- TVM_FFI_ICHECK(dtype_a == dtype_b) << ab_not_match_err_str;
- break;
- case DataType::kInt4:
- case DataType::kUInt4:
- TVM_FFI_ICHECK(dtype_b == DataType::kInt4 || dtype_b == DataType::kUInt4)
- << ab_not_match_err_str;
- break;
- case DataType::kInt8:
- case DataType::kUInt8:
- TVM_FFI_ICHECK(dtype_b == DataType::kInt8 || dtype_b == DataType::kUInt8)
- << ab_not_match_err_str;
- break;
- case DataType::kFloat8_e4m3fn:
- case DataType::kFloat8_e5m2:
- TVM_FFI_ICHECK(dtype_b == DataType::kFloat8_e4m3fn || dtype_b == DataType::kFloat8_e5m2)
- << ab_not_match_err_str;
- break;
- default:
- TVM_FFI_ICHECK(false) << "Invalid multiplicand data types: " << DTypeToString(dtype_a)
- << DTypeToString(dtype_b);
- }
- // check a,b and c
- switch (dtype_a) {
- case DataType::kBit1:
- case DataType::kInt4:
- case DataType::kUInt4:
- case DataType::kInt8:
- case DataType::kUInt8:
- TVM_FFI_ICHECK(dtype_c == DataType::kInt32)
- << "For multiplicand data type " << DTypeToString(dtype_a) << DTypeToString(dtype_b)
- << ", accumulator data type should be s32.";
- break;
- case DataType::kFloat16:
- TVM_FFI_ICHECK(dtype_c == DataType::kFloat16 || dtype_c == DataType::kFloat32)
- << "For multiplicand data type f16, accumulator data type should be f16/f32.";
- break;
- case DataType::kBFloat16:
- case DataType::kTensorFloat32:
- TVM_FFI_ICHECK(dtype_c == DataType::kFloat32)
- << "For multiplicand data type bf16/tf32, accumulator data type can only be f32.";
- break;
- case DataType::kFloat64:
- TVM_FFI_ICHECK(dtype_c == DataType::kFloat64)
- << "For multiplicand data type f64, accumulator data type can only be f64.";
- break;
- case DataType::kFloat8_e4m3fn:
- case DataType::kFloat8_e5m2:
- TVM_FFI_ICHECK(dtype_c == DataType::kFloat32)
- << "For multiplicand data type e4m3/e5m2, accumulator data type can only be f32.";
- break;
- default:
- TVM_FFI_ICHECK(false) << "Invalid multiplicand/accumulator data types: "
- << DTypeToString(dtype_a) << DTypeToString(dtype_b)
- << DTypeToString(dtype_c) << ".";
- }
-}
-
-/*!
- * \brief Check whether the given configuration is valid for MMA computation.
- * \param m The M in mMnNkK of MMA instructions.
- * \param n The N in mMnNkK of MMA instructions.
- * \param k The K in mMnNkK of MMA instructions.
- * \param layout_a The layout of multiplicand A (row/col).
- * \param layout_b The layout of multiplicand B (row/col).
- * \param dtype_a The data type of multiplicand A.
- * \param dtype_b The data type of multiplicand B.
- * \param dtype_c The data type of accumulator C.
- * \param bit_op The bit operator for 1-bit MMA computation, can be "xor"/"and" or ""(if it's not
- * 1-bit MMA).
- * \param saturate Whether saturate output or not.
- */
-void CheckMMAConfigValidity(int m, int n, int k, LayoutType layout_a, LayoutType layout_b,
- DataType dtype_a, DataType dtype_b, DataType dtype_c,
- const std::string& bit_op, bool saturate) {
- TVM_FFI_ICHECK(bit_op == "xor" || bit_op == "and" || bit_op == "")
- << "Unrecognized 1-bit operation " << bit_op << " , can only be xor/and.";
- bool use_bit_op = !bit_op.empty();
- if (use_bit_op) {
- TVM_FFI_ICHECK(dtype_a == DataType::kBit1)
- << "Bit operator is only compatible with 1-bit multiplicand.";
- }
- CheckMMADTypeCompatible(dtype_a, dtype_b, dtype_c);
- if (saturate) {
- TVM_FFI_ICHECK(dtype_a == DataType::kInt4 || dtype_a == DataType::kUInt4 ||
- dtype_a == DataType::kInt8 || dtype_a == DataType::kUInt8)
- << "Output saturation only applicable to multiplicand type s4/u4/s8/u8.";
- }
-
- if (!(m == 8 && n == 8 && k == 4 && dtype_a == ptx::DataType::kFloat16)) {
- // Only MMA on m8n8k4 for fp16 supports customized layouts.
- TVM_FFI_ICHECK(layout_a == LayoutType::kRowMajor && layout_b == LayoutType::kColumnMajor)
- << "Invalid layout combination " << LayoutTypeToString(layout_a) << ","
- << LayoutTypeToString(layout_b) << ".";
- }
-
- MMAConfig config(m, n, k, dtype_a, use_bit_op);
- bool match = false;
- for (const MMAConfig& valid_config : valid_mma_configs) {
- if (config == valid_config) {
- match = true;
- break;
- }
- }
- TVM_FFI_ICHECK(match) << "Cannot find matched MMA configurations.";
-}
-
-/*!
- * \brief Fragment attributes
- */
-class FragAttrs {
- public:
- explicit FragAttrs(char reg_type, uint32_t size, std::string ptr_type)
- : reg_type(reg_type), size(size), ptr_type(ptr_type) {}
- /*! \brief PTX register type */
- char reg_type;
- /*! \brief Fragment size */
- uint32_t size;
- /*! \brief Fragment pointer type */
- std::string ptr_type;
-};
-
-/*!
- * \brief Fragment attributes of given data type.
- */
-inline FragAttrs GetFragAttrs(DataType dtype) {
- switch (dtype) {
- case DataType::kBit1:
- case DataType::kInt4:
- case DataType::kUInt4:
- case DataType::kInt8:
- case DataType::kUInt8:
- case DataType::kFloat8_e4m3fn:
- case DataType::kFloat8_e5m2:
- case DataType::kBit16:
- case DataType::kFloat16: // .f16x2 register
- case DataType::kBFloat16:
- case DataType::kTensorFloat32:
- return FragAttrs('r', 32, "(unsigned *)");
- case DataType::kInt32:
- return FragAttrs('r', 32, "(int *)");
- case DataType::kFloat32:
- return FragAttrs('f', 32, "(float *)");
- case DataType::kFloat64:
- return FragAttrs('d', 64, "(double *)");
- default:
- TVM_FFI_ICHECK(false) << DTypeToString(dtype) << " is not matrix data type in MMA.";
- return FragAttrs('\0', 0, "");
- }
-}
-
-}; // namespace ptx
-
-/*!
- * \brief Replace patterns with replacement strings.
- * \note should use std::format instead when codebase is ported to C++20.
- */
-class Replacer {
- public:
- void register_rule(const std::string& pattern, const std::string& replacement) {
- _rules.emplace_back(pattern, replacement);
- }
- std::string rewrite(std::string str) {
- for (auto&& rule : _rules) {
- auto [pattern, replacement] = rule;
- size_t len = pattern.size();
- size_t new_len = replacement.size();
- size_t pos = str.find(pattern);
- while (pos != std::string::npos) {
- str = str.replace(pos, len, replacement);
- pos = str.find(pattern, pos + new_len);
- }
- }
- return str;
- }
- void empty_rules() { _rules.clear(); }
-
- private:
- std::vector<std::pair<std::string, std::string>> _rules;
-};
-
-/*!
- * \brief Get the number of MMA computations for given shape and datatype.
- */
-inline uint32_t GetNumMMAComputations(int m, int n, int k, ptx::DataType dtype) {
- if (m == 8 && n == 8 && k == 4 && dtype == ptx::DataType::kFloat16) {
- // MMA for m8n8k4 on fp16 would launch 4 MMA computations instead of one.
- return 4;
- } else {
- return 1;
- }
-}
-
-/*!
- * \brief Return template string, input operands string and output operands string.
- * \param m The M in mMnNkK of MMA instructions.
- * \param n The N in mMnNkK of MMA instructions.
- * \param k The K in mMnNkK of MMA instructions.
- * \param dtype_a The data type of multiplicand a.
- * \param dtype_b The data type of multiplicand b.
- * \param dtype_c The data type of accumulator c.
- */
-inline std::tuple<std::string, std::string, std::string> GetMMAOperands(int m, int n, int k,
- ptx::DataType dtype_a,
- ptx::DataType dtype_b,
- ptx::DataType dtype_c) {
- std::stringstream templates, inputs, outputs;
- const ptx::FragAttrs frag_attr_a = ptx::GetFragAttrs(dtype_a),
- frag_attr_b = ptx::GetFragAttrs(dtype_b),
- frag_attr_c = ptx::GetFragAttrs(dtype_c);
- constexpr uint32_t warp_size = 32;
- const uint32_t threads = warp_size / GetNumMMAComputations(m, n, k, dtype_a);
- const int num_operands_a = (m * k) * ptx::DTypeBits(dtype_a) / frag_attr_a.size / threads,
- num_operands_b = (k * n) * ptx::DTypeBits(dtype_b) / frag_attr_b.size / threads,
- num_operands_c = (m * n) * ptx::DTypeBits(dtype_c) / frag_attr_c.size / threads;
-
- // generate templates;
- int arg_counter = 0;
- templates << "{"
- << "%" << arg_counter++;
- for (int i = 1; i < num_operands_c; ++i) {
- templates << ", %" << arg_counter++;
- }
- templates << "}, {"
- << "%" << arg_counter++;
- for (int i = 1; i < num_operands_a; ++i) {
- templates << ", %" << arg_counter++;
- }
- templates << "}, {"
- << "%" << arg_counter++;
- for (int i = 1; i < num_operands_b; ++i) {
- templates << ", %" << arg_counter++;
- }
- templates << "}, {"
- << "%" << arg_counter++;
- for (int i = 1; i < num_operands_c; ++i) {
- templates << ", %" << arg_counter++;
- }
- templates << "}";
-
- // generate inputs
- for (int i = 0; i < num_operands_a; ++i) {
- if (i != 0) {
- inputs << ", ";
- }
- inputs << "\"" << frag_attr_a.reg_type << "\"((" << frag_attr_a.ptr_type << "(A))[" << i
- << "])";
- }
- for (int i = 0; i < num_operands_b; ++i) {
- inputs << ", \"" << frag_attr_b.reg_type << "\"((" << frag_attr_b.ptr_type << "(B))[" << i
- << "])";
- }
- for (int i = 0; i < num_operands_c; ++i) {
- inputs << ", \"" << frag_attr_c.reg_type << "\"((" << frag_attr_c.ptr_type << "(C))[" << i
- << "])";
- }
-
- // generate outputs
- for (int i = 0; i < num_operands_c; ++i) {
- if (i != 0) {
- outputs << ",";
- }
- outputs << " \"=" << frag_attr_c.reg_type << "\"((" << frag_attr_c.ptr_type << "(D))[" << i
- << "])";
- }
- return std::make_tuple(templates.str(), inputs.str(), outputs.str());
-}
-
-// ldmatrix assembly emitter.
-// `local_elem_offset` / `smem_elem_offset` are element offsets in the
-// respective buffer's dtype; the generated C expression `ptr + offset`
-// relies on C pointer arithmetic to scale them to bytes.
-inline std::tuple<std::string, std::string> GetLoadMatrixOperands(
- int num, const std::string& local_ptr, const std::string& local_elem_offset) {
- std::stringstream templates, outputs;
- int arg_counter = 0;
- templates << "{%" << arg_counter++;
- for (int i = 1; i < num; ++i) {
- templates << ", %" << arg_counter++;
- }
- templates << "}, [%" << arg_counter++ << "]";
- std::string ptr_type = "(unsigned *)";
- for (int i = 0; i < num; ++i) {
- if (i != 0) outputs << ", ";
- outputs << "\"=r\"((" << ptr_type << "(" << local_ptr << " + " << local_elem_offset << "))["
- << i << "])";
- }
- return std::make_tuple(templates.str(), outputs.str());
-}
-
-std::string PrintLoadMatrixAssembly(bool trans, int num, const std::string& type,
- const std::string& local_ptr,
- const std::string& local_elem_offset,
- const std::string& smem_ptr,
- const std::string& smem_elem_offset) {
- TVM_FFI_ICHECK(num == 1 || num == 2 || num == 4)
- << "ldmatrix only accept loading 1/2/4 matrices.";
- ptx::DataType data_type = ptx::DTypeFromString(type);
- TVM_FFI_ICHECK(data_type == ptx::DataType::kBit16)
- << "ldmatrix only accept matrix with type .b16.";
- std::string asm_code = R"(
- {
- unsigned int addr = __cvta_generic_to_shared({smem_addr});
- __asm__ __volatile__(
- "ldmatrix.sync.aligned{.shape}{.num}{.trans}{.ss}{.type}"
- "{templates};\n"
- : {outputs}
- : "r"(addr)
- );
- }
-)";
- auto [templates_str, outputs_str] = GetLoadMatrixOperands(num, local_ptr, local_elem_offset);
- Replacer replacer;
- replacer.register_rule("{.shape}", ".m8n8");
- replacer.register_rule("{.num}", ".x" + std::to_string(num));
- replacer.register_rule("{.trans}", trans ? ".trans" : "");
- replacer.register_rule("{.ss}", ".shared");
- replacer.register_rule("{.type}", ptx::DTypeToString(data_type));
- replacer.register_rule("{smem_addr}", smem_ptr + " + " + smem_elem_offset);
- replacer.register_rule("{templates}", templates_str);
- replacer.register_rule("{outputs}", outputs_str);
- asm_code = replacer.rewrite(asm_code);
- return asm_code;
-}
-
-std::string PrintMMAAssembly(const std::string& shape, const std::string& A_layout,
- const std::string& B_layout, const std::string& A_dtype,
- const std::string& B_dtype, const std::string& C_dtype,
- const std::string& a_ptr, const std::string& a_elem_offset,
- const std::string& b_ptr, const std::string& b_elem_offset,
- const std::string& c_ptr, const std::string& c_elem_offset,
- const std::string& bit_op, bool saturate) {
- ptx::DataType dtype_a = ptx::DTypeFromString(A_dtype), dtype_b = ptx::DTypeFromString(B_dtype),
- dtype_c = ptx::DTypeFromString(C_dtype);
- ptx::LayoutType layout_a = ptx::LayoutTypeFromString(A_layout),
- layout_b = ptx::LayoutTypeFromString(B_layout);
- auto [m, n, k] = ptx::ParseMMAShape(shape);
- CheckMMAConfigValidity(m, n, k, layout_a, layout_b, dtype_a, dtype_b, dtype_c, bit_op, saturate);
- std::string asm_code = R"(
- {
- __asm__ __volatile__(
- "mma.sync.aligned{.shape}{.alayout}{.blayout}{.saturate}{.dtype}{.atype}{.btype}{.ctype}{.bitop}"
- "{templates};\n"
- : {outputs}
- : {inputs});
- }
-)";
- auto [templates_str, inputs_str, outputs_str] =
- GetMMAOperands(m, n, k, dtype_a, dtype_b, dtype_c);
-
- // replace patterns
- Replacer replacer;
- replacer.register_rule("{.shape}", "." + shape);
- replacer.register_rule("{.saturate}", saturate ? ".satfinite" : "");
- replacer.register_rule("{.alayout}", "." + A_layout);
- replacer.register_rule("{.blayout}", "." + B_layout);
- replacer.register_rule("{.atype}", ptx::DTypeToString(dtype_a));
- replacer.register_rule("{.btype}", ptx::DTypeToString(dtype_b));
- replacer.register_rule("{.ctype}", ptx::DTypeToString(dtype_c));
- replacer.register_rule("{.dtype}", ptx::DTypeToString(dtype_c));
- replacer.register_rule("{.bitop}", bit_op.empty() ? "" : "." + bit_op + ".popc");
- replacer.register_rule("{templates}", templates_str);
- replacer.register_rule("{outputs}", outputs_str);
- replacer.register_rule("{inputs}", inputs_str);
- asm_code = replacer.rewrite(asm_code);
- replacer.empty_rules();
- replacer.register_rule("A", a_ptr + " + " + a_elem_offset);
- replacer.register_rule("B", b_ptr + " + " + b_elem_offset);
- replacer.register_rule("C", c_ptr + " + " + c_elem_offset);
- replacer.register_rule("D", c_ptr + " + " + c_elem_offset);
- asm_code = replacer.rewrite(asm_code);
- return asm_code;
-}
-
+} // namespace ptx
} // namespace codegen
} // namespace tvm
diff --git a/src/backend/cuda/codegen/ptx.h b/src/backend/cuda/codegen/ptx.h
deleted file mode 100644
index 3965c6f4de..0000000000
--- a/src/backend/cuda/codegen/ptx.h
+++ /dev/null
@@ -1,76 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied. See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-
-/*!
- * \file ptx.h
- * \brief Code generation with inlined PTX code.
- */
-#ifndef TVM_TARGET_SOURCE_PTX_H_
-#define TVM_TARGET_SOURCE_PTX_H_
-
-#include <tvm/ffi/error.h>
-
-#include <string>
-#include <tuple>
-
-#include "codegen_cuda.h"
-
-namespace tvm {
-namespace codegen {
-
-/*!
- * \brief Print MMA assembly string given parameters.
- * \param shape The shape string mMnNkK
- * \param A_layout The layout of multiplicand A, can be either "row" or "col".
- * \param B_layout The layout of multiplicand B, can be either "row" or "col".
- * \param A_dtype The data type of multiplicand A.
- * \param B_dtype The data type of multiplicand B.
- * \param C_dtype The data type of multiplicand C.
- * \param a_ptr Pointer to buffer A.
- * \param a_offset The offset of element in A.
- * \param b_ptr Pointer to buffer B.
- * \param b_offset The offset of element in B.
- * \param c_ptr Pointer to buffer C.
- * \param c_offset The offset of element in C.
- * \param bit_op The bit operator used in 1-bit mma, can be either "xor" or "and".
- * \param saturate Whether saturate output or not.
- */
-/*!
- * \brief ldmatrix assembly emitter. Offsets are element offsets in the
- * buffer's dtype; the generated C pointer arithmetic ``ptr + offset``
- * scales them to bytes.
- */
-std::string PrintLoadMatrixAssembly(bool trans, int num, const std::string& type,
- const std::string& local_ptr,
- const std::string& local_elem_offset,
- const std::string& smem_ptr,
- const std::string& smem_elem_offset);
-
-std::string PrintMMAAssembly(const std::string& shape, const std::string& A_layout,
- const std::string& B_layout, const std::string& A_dtype,
- const std::string& B_dtype, const std::string& C_dtype,
- const std::string& a_ptr, const std::string& a_offset,
- const std::string& b_ptr, const std::string& b_offset,
- const std::string& c_ptr, const std::string& c_offset,
- const std::string& bit_op, bool saturate);
-
-} // namespace codegen
-} // namespace tvm
-
-#endif // TVM_TARGET_SOURCE_PTX_H_
diff --git a/src/backend/cuda/op/legacy.cc b/src/backend/cuda/op/legacy.cc
new file mode 100644
index 0000000000..42841e8804
--- /dev/null
+++ b/src/backend/cuda/op/legacy.cc
@@ -0,0 +1,61 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file backend/cuda/op/legacy.cc
+ * \brief Legacy CUDA operations retained for downstream code generators.
+ */
+#include <tvm/ir/op.h>
+#include <tvm/ir/prim/op.h>
+
+namespace tvm {
+namespace tirx {
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ OpDef("tirx.cuda.mma_store")
+ .signature(sig::arg<IntExpr>("m", "The M dimension."),
+ sig::arg<IntExpr>("n", "The N dimension."),
+ sig::arg("dst_ptr", "The destination pointer."),
+ sig::arg("src_ptr", "The source pointer."),
+ sig::arg<IntExpr>("src_offset", "The source offset."),
+ sig::arg<IntExpr>("dst_stride", "The destination stride."))
+ .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque));
+
+ OpDef("tirx.cuda.mma_fill")
+ .signature(sig::arg<IntExpr>("local_size", "The local allocation size."),
+ sig::arg("local_ptr", "The local pointer."),
+ sig::arg<IntExpr>("offset", "The offset."))
+ .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque));
+
+ OpDef("tirx.ptx_legacy.ldmatrix")
+ .signature(sig::arg("trans"), sig::arg<IntExpr>("num"), sig::arg("dtype"),
+ sig::arg("local_ptr"), sig::arg<IntExpr>("local_offset"), sig::arg("smem_ptr"),
+ sig::arg<IntExpr>("smem_offset"))
+ .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque));
+ OpDef("tirx.ptx_legacy.mma")
+ .signature(sig::arg("shape"), sig::arg("a_layout"), sig::arg("b_layout"), sig::arg("a_dtype"),
+ sig::arg("b_dtype"), sig::arg("c_dtype"), sig::arg("a_ptr"),
+ sig::arg<IntExpr>("a_offset"), sig::arg("b_ptr"), sig::arg<IntExpr>("b_offset"),
+ sig::arg("acc_ptr"), sig::arg<IntExpr>("c_offset"), sig::arg("saturate"),
+ sig::var_args("args"))
+ .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque));
+}
+
+} // namespace tirx
+} // namespace tvm
diff --git a/src/backend/cuda/op/target_builtin.cc b/src/backend/cuda/op/target_builtin.cc
index f90b4e6041..d69f2131a1 100644
--- a/src/backend/cuda/op/target_builtin.cc
+++ b/src/backend/cuda/op/target_builtin.cc
@@ -129,21 +129,6 @@ void RegisterCudaTargetBuiltins() {
sig::arg<IntExpr>("cp_size", "The copy size."), sig::var_args("args"))
.set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque));
- OpDef("tirx.cuda.mma_store")
- .signature(sig::arg<IntExpr>("m", "The M dimension."),
- sig::arg<IntExpr>("n", "The N dimension."),
- sig::arg("dst_ptr", "The destination pointer."),
- sig::arg("src_ptr", "The source pointer."),
- sig::arg<IntExpr>("src_offset", "The source offset."),
- sig::arg<IntExpr>("dst_stride", "The destination stride."))
- .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque));
-
- OpDef("tirx.cuda.mma_fill")
- .signature(sig::arg<IntExpr>("local_size", "The local allocation size."),
- sig::arg("local_ptr", "The local pointer."),
- sig::arg<IntExpr>("offset", "The offset."))
- .set_attr(tvm::op_attr::kCallEffectKind, static_cast<int64_t>(CallEffectKind::kOpaque));
-
OpDef("tirx.cuda.timer_init")
.signature(sig::arg("profiler_buffer", "The profiler buffer."),
sig::arg("profiler_tag", "The profiler tag."),
@@ -412,16 +397,6 @@ void RegisterDeviceIntrinsics() {
sig::arg("ivar"), sig::arg("cmp"), sig::arg<IntExpr>("cmp_value"),
sig::arg("type"))
.set_attr(tvm::op_attr::kFixedReturnType, PrimType::Void());
- RegisterDeviceIntrinsic(OpDef("tirx.ptx_legacy.ldmatrix"), CallEffectKind::kOpaque,
- sig::arg("trans"), sig::arg<IntExpr>("num"), sig::arg("dtype"),
- sig::arg("local_ptr"), sig::arg<IntExpr>("local_offset"),
- sig::arg("smem_ptr"), sig::arg<IntExpr>("smem_offset"));
- RegisterDeviceIntrinsic(
- OpDef("tirx.ptx_legacy.mma"), CallEffectKind::kOpaque, sig::arg("shape"),
- sig::arg("a_layout"), sig::arg("b_layout"), sig::arg("a_dtype"), sig::arg("b_dtype"),
- sig::arg("c_dtype"), sig::arg("a_ptr"), sig::arg<IntExpr>("a_offset"), sig::arg("b_ptr"),
- sig::arg<IntExpr>("b_offset"), sig::arg("acc_ptr"), sig::arg<IntExpr>("c_offset"),
- sig::arg("saturate"), sig::var_args("args"));
}
} // namespace
diff --git a/src/backend/opencl/codegen/codegen_opencl.cc b/src/backend/opencl/codegen/codegen_opencl.cc
index 7037c04afe..365aaa1921 100644
--- a/src/backend/opencl/codegen/codegen_opencl.cc
+++ b/src/backend/opencl/codegen/codegen_opencl.cc
@@ -426,6 +426,8 @@ void CodeGenOpenCL::PrintStorageQualifier(const VarNode* var, std::ostream& os)
void CodeGenOpenCL::PrintStorageScope(const runtime::MemoryScope& scope,
std::ostream& os) { // NOLINT(*)
static const auto scope_shared = runtime::MemoryScope::Get("shared");
+ TVM_FFI_ICHECK(scope != runtime::MemoryScope::Get("warp"))
+ << "OpenCL warp storage is unsupported; use explicit per-thread local storage and shuffles";
if (scope == runtime::MemoryScope::Global()) {
os << "__global ";
} else if (scope == scope_shared) {
diff --git a/src/s_tir/meta_schedule/mutator/mutator.cc b/src/s_tir/meta_schedule/mutator/mutator.cc
index 7f7465b5bf..57d0cb0c0c 100644
--- a/src/s_tir/meta_schedule/mutator/mutator.cc
+++ b/src/s_tir/meta_schedule/mutator/mutator.cc
@@ -67,8 +67,6 @@ ffi::Map<Mutator, FloatImm> Mutator::DefaultCUDA() {
{Mutator::MutateThreadBinding(), FloatImm(PrimType::Float(64), 0.02)}};
}
-ffi::Map<Mutator, FloatImm> Mutator::DefaultCUDATensorCore() { return Mutator::DefaultCUDA(); }
-
ffi::Map<Mutator, FloatImm> Mutator::DefaultHexagon() {
return ffi::Map<Mutator, FloatImm>{
{Mutator::MutateTileSize(), FloatImm(PrimType::Float(64), 0.9)},
@@ -98,7 +96,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.def("s_tir.meta_schedule.MutatorPyMutator", Mutator::PyMutator)
.def("s_tir.meta_schedule.MutatorDefaultLLVM", Mutator::DefaultLLVM)
.def("s_tir.meta_schedule.MutatorDefaultCUDA", Mutator::DefaultCUDA)
- .def("s_tir.meta_schedule.MutatorDefaultCUDATensorCore", Mutator::DefaultCUDATensorCore)
.def("s_tir.meta_schedule.MutatorDefaultHexagon", Mutator::DefaultHexagon);
}
diff --git a/src/s_tir/meta_schedule/postproc/postproc.cc b/src/s_tir/meta_schedule/postproc/postproc.cc
index f1dec5c5e9..17d1dc5fb6 100644
--- a/src/s_tir/meta_schedule/postproc/postproc.cc
+++ b/src/s_tir/meta_schedule/postproc/postproc.cc
@@ -87,20 +87,6 @@ ffi::Array<Postproc> Postproc::DefaultCUDA() {
};
}
-ffi::Array<Postproc> Postproc::DefaultCUDATensorCore() {
- return ffi::Array<Postproc>{
- Postproc::DisallowDynamicLoop(),
- Postproc::RewriteCooperativeFetch(),
- Postproc::RewriteUnboundBlock(/*max_threadblocks=*/256),
- Postproc::RewriteParallelVectorizeUnroll(),
- Postproc::RewriteReductionBlock(),
- Postproc::VerifyGPUCode(),
- // RewriteTensorize is relatively expensive and it doesn't affect the validity of a sample, so
- // run it only on samples that have passed VerifyGPUCode.
- Postproc::RewriteTensorize(/*vectorize_init_loop=*/false),
- };
-}
-
ffi::Array<Postproc> Postproc::DefaultHexagon() {
return ffi::Array<Postproc>{
Postproc::DisallowDynamicLoop(), Postproc::RewriteParallelVectorizeUnroll(),
@@ -126,7 +112,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.def("s_tir.meta_schedule.PostprocPyPostproc", Postproc::PyPostproc)
.def("s_tir.meta_schedule.PostprocDefaultLLVM", Postproc::DefaultLLVM)
.def("s_tir.meta_schedule.PostprocDefaultCUDA", Postproc::DefaultCUDA)
- .def("s_tir.meta_schedule.PostprocDefaultCUDATensorCore", Postproc::DefaultCUDATensorCore)
.def("s_tir.meta_schedule.PostprocDefaultHexagon", Postproc::DefaultHexagon);
}
diff --git a/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling_tensor_core.cc b/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling_tensor_core.cc
deleted file mode 100644
index 5d0d462766..0000000000
--- a/src/s_tir/meta_schedule/schedule_rule/multi_level_tiling_tensor_core.cc
+++ /dev/null
@@ -1,984 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied. See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-#include <tvm/ffi/cast.h>
-#include <tvm/ffi/reflection/registry.h>
-#include <tvm/ir/prim/op.h>
-#include <tvm/runtime/memory_scope.h>
-#include <tvm/s_tir/meta_schedule/schedule_rule.h>
-#include <tvm/s_tir/stmt.h>
-#include <tvm/s_tir/tensor_intrin.h>
-
-#include <algorithm>
-#include <array>
-#include <utility>
-#include <vector>
-
-#include "../utils.h"
-#include "./multi_level_tiling.h"
-
-namespace tvm {
-namespace s_tir {
-namespace meta_schedule {
-
-using s_tir::GetSBlockVarTypes;
-using s_tir::IterVarType;
-using s_tir::LoopRV;
-using s_tir::SBlockRV;
-using s_tir::Schedule;
-
-struct TensorCoreIntrinGroup {
- ffi::String init_intrin;
- ffi::String load_a_intrin;
- ffi::String load_b_intrin;
- ffi::String compute_intrin;
- ffi::String store_intrin;
-
- /*! \brief Create TensorCoreIntrinGroup from config in a map. The map should contains the
- * following keys:
- * - init
- * - load_a
- * - load_b
- * - compute
- * - store
- * The values of the keys should be the names of the corresponding intrinsics and should be
- * registered via TensorIntrin.Register beforehand.
- */
- static TensorCoreIntrinGroup FromConfig(const ffi::Map<ffi::String, ffi::String>& config);
-};
-
-TensorCoreIntrinGroup TensorCoreIntrinGroup::FromConfig(
- const ffi::Map<ffi::String, ffi::String>& config) {
- auto f_initialize_intrin = [&config](ffi::String key_name, ffi::String* intrin_name) {
- TVM_FFI_CHECK(config.count(key_name), ValueError) << key_name << " is not set.";
- *intrin_name = config.at(key_name);
- // Check the existence of the intrin
- TensorIntrin::Get(*intrin_name);
- };
- TensorCoreIntrinGroup intrin_group;
- f_initialize_intrin("init", &intrin_group.init_intrin);
- f_initialize_intrin("load_a", &intrin_group.load_a_intrin);
- f_initialize_intrin("load_b", &intrin_group.load_b_intrin);
- f_initialize_intrin("compute", &intrin_group.compute_intrin);
- f_initialize_intrin("store", &intrin_group.store_intrin);
- return intrin_group;
-}
-
-class TensorCoreStateObj : public StateObj {
- public:
- explicit TensorCoreStateObj(Schedule sch) : StateObj(std::move(sch)) {}
- /*! \brief The tensor core intrinsic group. */
- TensorCoreIntrinGroup intrin_group;
- /*! \brief The auto tensorization maping info. */
- s_tir::AutoTensorizeMappingInfo mapping_info{ffi::UnsafeInit()};
- /*! \brief The Tensor Core reindex block A for Tensor Core computation */
- s_tir::SBlockRV tensor_core_reindex_A;
- /*! \brief The Tensor Core reindex block B for Tensor Core computation */
- s_tir::SBlockRV tensor_core_reindex_B;
- /*! \brief The Tensor Core reindex store block for Tensor Core computation */
- s_tir::SBlockRV tensor_core_reindex_store;
- /*! \brief Flag to indicate its a WMMA or MMA intrin group */
- bool is_mma;
- /*! \brief Flag to indicate whether to use async software pipeline */
- bool use_async;
-
- State Copy() const final;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<TensorCoreStateObj>();
- }
-
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.meta_schedule.TensorCoreState", TensorCoreStateObj,
- StateObj);
-};
-
-class TensorCoreState : public State {
- public:
- explicit TensorCoreState(TensorCoreIntrinGroup intrin_group,
- s_tir::AutoTensorizeMappingInfo mapping_info, Schedule sch,
- SBlockRV block_rv, bool use_async,
- ffi::Array<ffi::Array<s_tir::LoopRV>> tiles = {});
-
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TensorCoreState, State, TensorCoreStateObj);
-};
-
-TensorCoreState::TensorCoreState(TensorCoreIntrinGroup intrin_group,
- s_tir::AutoTensorizeMappingInfo mapping_info, Schedule sch,
- SBlockRV block_rv, bool use_async,
- ffi::Array<ffi::Array<LoopRV>> tiles)
- : State(ffi::UnsafeInit{}) {
- ffi::ObjectPtr<TensorCoreStateObj> node = ffi::make_object<TensorCoreStateObj>(std::move(sch));
- node->intrin_group = intrin_group;
- node->mapping_info = mapping_info;
- node->block_rv = std::move(block_rv);
- node->tiles = std::move(tiles);
- node->is_mma = intrin_group.compute_intrin.starts_with("mma_sync");
- node->use_async = use_async;
- data_ = std::move(node);
-}
-
-State TensorCoreStateObj::Copy() const {
- ffi::ObjectPtr<TensorCoreStateObj> node = ffi::make_object<TensorCoreStateObj>(*this);
- node->sch = sch->Copy();
- return State(ffi::UnsafeInit{}, node);
-}
-
-/*!
- * \brief Extension of MultiLevelTiling for auto-tensorizing with a single group of tensor core
- * intrinsics.
- */
-class MultiLevelTilingTensorCoreObj : public MultiLevelTilingObj {
- private:
- // SubRule: Add tensorization-related transformations
- inline std::vector<State> TransformForTensorization(TensorCoreState state) const;
- // Subrule: Transform the layout of the output. This is necessary for efficient cache write the
- // output in the shared memory.
- std::vector<State> TransformIntermediateOutputLayout(TensorCoreState state);
- // Subrule: Add read cache for mma
- // Basically same with MultiLevelTilingObj::AddReadReuse, but change CacheRead + ComputeAt to
- // ReadAt
- inline std::vector<State> MMAAddReadReuse(TensorCoreState state) const;
- // Subrule: Add tensorized load
- inline std::vector<State> AddReadReuseTensorCore(TensorCoreState state) const;
- // Subrule: Add tensorized store
- inline std::vector<State> AddWriteReuseTensorCore(TensorCoreState state) const;
- // Subrule: Add software pipeline
- inline std::vector<State> AddSoftwarePipeline(TensorCoreState state) const;
- // Subrule: split loop for mma using sample partitioned tile
- inline std::pair<ffi::Array<ffi::Optional<s_tir::ExprRV>>, ffi::Array<s_tir::LoopRV>>
- MMASplitLoop(const Schedule& sch, SBlockRV block, LoopRV loop, int n_tiles, int partition_pos,
- int innerpart_factor) const;
- // Subrule: tile loop nest for mma
- // Basically same with MultiLevelTilingObj::TileLoopNest, but change SamplePerfectTile to
- // SamplePartitionedTile
- inline std::vector<State> MMATileLoopNest(TensorCoreState state) const;
-
- // Override ApplySubRules to apply tensorization-specific sub-rules
- std::vector<State> ApplySubRules(std::vector<State> states) final;
-
- // Override Apply to apply tensorization-specific analysis before applying sub-rules
- ffi::Array<Schedule> Apply(const Schedule& sch, const SBlockRV& block_rv) final;
-
- // Inherited from ScheduleRuleObj
- ScheduleRule Clone() const final {
- ffi::ObjectPtr<MultiLevelTilingTensorCoreObj> n =
- ffi::make_object<MultiLevelTilingTensorCoreObj>(*this);
- return ScheduleRule(ffi::UnsafeInit{}, n);
- }
-
- /*!
- * \brief Transform and tensorize with the given tensor intrin
- * \param state The state of the meta schedule rule
- * \param intrin_name The name of the tensor intrin
- * \return The loop to be tensorized. std::nullopt if the workload can't be tensorized.
- */
- ffi::Optional<LoopRV> TransformWithTensorIntrin(TensorCoreStateObj* state,
- const ffi::String& intrin_name) const;
-
- /*!
- * \brief Tile, blockize and annotate for tensorization with the given intrin
- * \param block_rv The block to be tensorized
- * \param intrin_name The name of the tensor intrin
- */
- void TileAndAnnotateTensorize(Schedule* sch, const SBlockRV& block_rv,
- const ffi::String& intrin_name,
- const ffi::String& permuted_layout_annotate_value) const;
-
- public:
- /*! \brief The candidate tensor core intrin groups to apply */
- std::vector<TensorCoreIntrinGroup> intrin_groups;
- /*! \brief Whether to use software pipeline */
- bool use_software_pipeline = false;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<MultiLevelTilingTensorCoreObj>();
- }
-
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.meta_schedule.MultiLevelTilingTensorCore",
- MultiLevelTilingTensorCoreObj, MultiLevelTilingObj);
-
- private:
-};
-
-// Entry of the mega rule; Inherited from ScheduleRuleObj
-ffi::Array<Schedule> MultiLevelTilingTensorCoreObj::Apply(const Schedule& sch,
- const SBlockRV& block_rv) {
- if (!NeedsMultiLevelTiling(sch->state(), sch->GetSRef(block_rv))) {
- return {sch};
- }
-
- std::unordered_map<int, s_tir::AutoTensorizeMappingInfo> intrin_group_to_mapping_info;
- for (int i = 0, n = intrin_groups.size(); i < n; ++i) {
- TensorCoreIntrinGroup intrin_group = intrin_groups[i];
- ffi::Optional<s_tir::AutoTensorizeMappingInfo> mapping_info =
- s_tir::GetAutoTensorizeMappingInfo(
- sch->state(), sch->GetSRef(block_rv),
- TensorIntrin::Get(intrin_groups[i].compute_intrin).value()->desc);
- if (mapping_info.has_value()) {
- intrin_group_to_mapping_info.emplace(i, mapping_info.value());
- }
- }
-
- if (intrin_group_to_mapping_info.empty()) {
- // No tensor intrinsics can be applied.
- return {sch};
- }
-
- // Save the original schedule so that we can roll back transformations if tensorization
- // fail.
- Schedule original_sch = sch;
-
- std::vector<State> initial_states;
- for (const auto& kv : intrin_group_to_mapping_info) {
- const TensorCoreIntrinGroup& intrin_group = intrin_groups[kv.first];
- const s_tir::AutoTensorizeMappingInfo& mapping_info = kv.second;
- Schedule new_sch = sch->Copy();
- new_sch->Annotate(block_rv, tvm::s_tir::attr::kMetaScheduleTilingStructure, structure);
- initial_states.push_back(TensorCoreState(intrin_group, mapping_info, new_sch, block_rv, true));
- }
- ffi::Array<Schedule> results;
- for (auto&& state : ApplySubRules(initial_states)) {
- TVM_PY_LOG(INFO, logger) << "Sketch " << results.size() << ": tensorizing with "
- << state.as<TensorCoreStateObj>()->intrin_group.compute_intrin;
- results.push_back(std::move(state->sch));
- }
- if (results.empty()) {
- TVM_PY_LOG(INFO, logger) << "The workload cannot be tensorized.";
- return {original_sch};
- }
- return results;
-}
-
-std::vector<State> MultiLevelTilingTensorCoreObj::ApplySubRules(std::vector<State> states) {
- states = SubRule(std::move(states), [&](State state) {
- return TransformForTensorization(state.as_or_throw<TensorCoreState>());
- });
- states = SubRule(std::move(states), [&](State state) {
- TensorCoreState tc_state = state.as_or_throw<TensorCoreState>();
- return tc_state->is_mma ? MMATileLoopNest(tc_state) : TileLoopNest(state, 2);
- });
- states = SubRule(std::move(states), [&](State state) {
- return TransformIntermediateOutputLayout(state.as_or_throw<TensorCoreState>());
- });
- states = SubRule(std::move(states), [&](State state) { return AddWriteReuse(state); });
- states = SubRule(std::move(states), [&](State state) {
- return AddWriteReuseTensorCore(state.as_or_throw<TensorCoreState>());
- });
- states = SubRule(std::move(states), [&](State state) {
- TensorCoreState tc_state = state.as_or_throw<TensorCoreState>();
- return tc_state->is_mma ? MMAAddReadReuse(tc_state) : AddReadReuse(state);
- });
- states = SubRule(std::move(states), [&](State state) {
- return AddReadReuseTensorCore(state.as_or_throw<TensorCoreState>());
- });
- states = SubRule(std::move(states), [&](State state) {
- return AddSoftwarePipeline(state.as_or_throw<TensorCoreState>());
- });
- return states;
-}
-
-void MultiLevelTilingTensorCoreObj::TileAndAnnotateTensorize(
- Schedule* sch, const SBlockRV& block_rv, const ffi::String& intrin_name,
- const ffi::String& permuted_layout_annotate_value) const {
- ffi::Optional<LoopRV> loop = s_tir::TileWithTensorIntrin(*sch, block_rv, intrin_name).value();
- TVM_FFI_ICHECK(loop.has_value());
- SBlockRV blockized_outer = (*sch)->Blockize(loop.value());
- (*sch)->Annotate(blockized_outer, tvm::s_tir::attr::kMetaScheduleAutoTensorize, intrin_name);
- if (!permuted_layout_annotate_value.empty()) {
- (*sch)->Annotate(blockized_outer, tvm::s_tir::attr::kPermutedLayout,
- permuted_layout_annotate_value);
- }
-}
-
-std::vector<State> MultiLevelTilingTensorCoreObj::MMAAddReadReuse(TensorCoreState state) const {
- const ReuseConfig& config = this->reuse_read_;
- if (config.req == ReuseType::kNoReuse) {
- return {std::move(state)};
- }
- TVM_FFI_ICHECK(config.req != ReuseType::kMayReuse);
- const SBlockRV& block_rv = state->block_rv;
- std::vector<State> results;
- results.reserve(config.levels.size());
- for (int level : config.levels) {
- State new_state = state->Copy();
- Schedule& sch = new_state->sch;
- const LoopRV& loop_rv = state->tiles[level - 1].back();
- // Enumerate all tensors that are read but not written
- std::vector<int> read_tensor_ndims = s_tir::GetReadTensorNDims(sch->GetSRef(block_rv));
- for (int i = 0, n_reads = read_tensor_ndims.size(); i < n_reads; ++i) {
- int tensor_ndim = read_tensor_ndims[i];
- if (tensor_ndim == -1) {
- continue;
- }
- // Do cache_read
- SBlockRV cache_read_block = sch->ReadAt(loop_rv, block_rv, i, config.scope);
- new_state->read_reuse.emplace(i, cache_read_block);
- if (state->is_mma) {
- new_state->sch->Annotate(
- cache_read_block, tvm::s_tir::attr::kPermutedLayout,
- ffi::String(std::string("g2s_") + std::string(i == 0 ? "A" : "B")));
- }
- }
- results.push_back(std::move(new_state));
- }
- return results;
-}
-
-std::pair<ffi::Array<ffi::Optional<s_tir::ExprRV>>, ffi::Array<s_tir::LoopRV>>
-MultiLevelTilingTensorCoreObj::MMASplitLoop(const Schedule& sch, SBlockRV block, LoopRV loop,
- int n_tiles, int partition_pos,
- int innerpart_factor) const {
- auto sampled_factors = sch->SamplePartitionedTile(
- /*loop=*/loop,
- /*n=*/n_tiles,
- /*partition_pos=*/partition_pos,
- /*innerpart_factor=*/innerpart_factor);
- ffi::Array<ffi::Optional<s_tir::ExprRV>> factors(sampled_factors.begin(), sampled_factors.end());
- ffi::Array<s_tir::LoopRV> splits = sch->Split(/*loop=*/loop,
- /*factors=*/factors);
- return {factors, splits};
-}
-
-std::vector<State> MultiLevelTilingTensorCoreObj::MMATileLoopNest(TensorCoreState state) const {
- Schedule& sch = state->sch;
- const SBlockRV& block_rv = state->block_rv;
- // Step 1. Assuming trivial binding, pair the loops and their iter-var-types
- ffi::Array<LoopRV> loops = sch->GetLoops(block_rv);
- if (!(loops.size() == 3 || !state->is_mma)) {
- LOG(DEBUG) << "The MMA tensor core only supports SSR loops now";
- return {};
- }
- std::vector<IterVarType> iter_types = GetSBlockVarTypes(sch->GetSRef(state->block_rv));
- TVM_FFI_ICHECK_EQ(loops.size(), iter_types.size());
- // Step 2. For each loop axis, tile it
- int64_t spatial_loop_product = 1;
- std::vector<ffi::Array<LoopRV>> tiles(s_indices_.size() + r_indices_.size());
- state->tile_factors.resize(tiles.size());
- std::vector<ffi::Array<ffi::Optional<s_tir::ExprRV>>> tile_factors;
- tile_factors.resize(tiles.size());
- for (int i = 0, n = loops.size(); i < n; ++i) {
- LoopRV loop = loops[i];
- const std::vector<int>* idx = nullptr;
-
- if (iter_types[i] == IterVarType::kDataPar) {
- idx = &s_indices_;
- if (spatial_loop_product != -1) {
- const auto* extent_imm = sch->Get(loop)->extent.as<IntImmNode>();
- if (auto extent = extent_imm ? extent_imm->value.as<int64_t>() : std::nullopt;
- extent.has_value()) {
- spatial_loop_product *= *extent;
- } else {
- spatial_loop_product = -1;
- }
- }
- } else if (iter_types[i] == IterVarType::kCommReduce) {
- idx = &r_indices_;
- } else {
- continue;
- }
-
- const int n_tiles = idx->size();
-
- if (n_tiles == 1) {
- tiles[idx->at(0)].push_back(loop);
- } else {
- auto [factors, splits] =
- iter_types[i] == IterVarType::kDataPar
- ? MMASplitLoop(
- sch, block_rv, loop, n_tiles, 3,
- i == 0 ? 2 // 32 (load A intrin i shape) / 16 (sync intrin i shape) == 2
- : 4 // 32 (load B intrin j shape) / 8 (sync intrin j shape) == 4
- )
- : SplitLoop(sch, block_rv, loop, n_tiles);
-
- // Put every tile to its slot
- for (int j = 0; j < n_tiles; ++j) {
- tiles[idx->at(j)].push_back(splits[j]);
- tile_factors[idx->at(j)].push_back(factors[j]);
- }
- }
- }
- state->tile_factors = std::move(tile_factors);
- // Step 3. Reorder to organize the tiles
- sch->Reorder(support::ConcatArrayList<LoopRV>(tiles.begin(), tiles.end()));
- // Step 4. Bind the tiles to threads
- int n_binds = std::min(tile_binds.size(), tiles.size());
- for (int i = 0; i < n_binds; ++i) {
- LoopRV fused = sch->Fuse(tiles[i]);
- sch->Bind(fused, tile_binds[i]);
- tiles[i] = {fused};
- }
- state->tiles = ffi::Array<ffi::Array<LoopRV>>{tiles.begin(), tiles.end()};
- if (this->thread_warp_size_ != -1) {
- int64_t low_inclusive = 1;
- int64_t high_inclusive = this->max_threads_per_block_;
- if (spatial_loop_product > 2 * this->thread_warp_size_) {
- low_inclusive = this->thread_warp_size_;
- }
- sch->Annotate(block_rv, tvm::s_tir::attr::kMetaScheduleThreadExtentLowInclusive, low_inclusive);
- sch->Annotate(block_rv, tvm::s_tir::attr::kMetaScheduleThreadExtentHighInclusive,
- high_inclusive);
- }
- return {state};
-}
-
-std::vector<State> MultiLevelTilingTensorCoreObj::TransformIntermediateOutputLayout(
- TensorCoreState state) {
- if (state->is_mma) {
- return {state};
- }
- // Transform the intermediate output to packed layout
- // [..., warp_m, warp_n, accum_frag_m, accum_frag_n, accum_elem_m, accum_elem_n]
- // where warp_m, warp_n are thread indices bound to the warp id, accum_frag_m, accum_frag_n are
- // the index of the fragments in each warp, accum_elem_m, accum_elem_n are the index of the
- // elements in each accumulator fragment.
-
- // Get the shape of the wmma accumulator
- auto [frag_shape_m, frag_shape_n] = [&]() {
- tirx::Function desc = TensorIntrin::Get(state->intrin_group.init_intrin).value()->desc;
- TVM_FFI_ICHECK(desc->body.has_value());
- TVM_FFI_ICHECK_EQ(desc->body.value()->size(), 1);
- s_tir::SBlock intrin_block =
- desc->body.value()->seq[0].as_or_throw<s_tir::SBlockRealize>()->block;
- TVM_FFI_ICHECK_EQ(intrin_block->body->size(), 1);
- For loop_m = intrin_block->body->seq[0].as_or_throw<For>();
- TVM_FFI_ICHECK_EQ(loop_m->body->size(), 1);
- For loop_n = loop_m->body->seq[0].as_or_throw<For>();
- return std::make_tuple(loop_m->extent, loop_n->extent);
- }();
-
- // Get the tile index of the warp id (i.e. threadIdx.y)
- auto it = std::find(tile_binds.begin(), tile_binds.end(), "threadIdx.y");
- TVM_FFI_ICHECK(it != tile_binds.end());
- auto tile_index_warp_id = std::distance(tile_binds.begin(), it);
-
- // Get the extent of loop indicated by `loop_idx` inside the warp scope.
- // For example, after spatial loops i, j are tiled, we will have
- // tile_factors = ((i0, j0), (i1, j1), ..., (in, jn))
- // This function computes the product of tile_factors[i][loop_idx] for i > tile_index_warp_id.
- // `loop_idx` can be negative, in which case it is counted from the end.
- auto f_get_inner_tile_product = [&](int loop_idx) {
- ffi::Array<ffi::Optional<s_tir::ExprRV>> factors;
- for (int i = tile_index_warp_id + 1; i < static_cast<int>(s_indices_.size()); ++i) {
- auto s_factors = state->tile_factors[s_indices_[i]];
- if (loop_idx < 0) {
- loop_idx += s_factors.size();
- }
- factors.push_back(s_factors[loop_idx]);
- }
- TVM_FFI_ICHECK(!factors.empty());
- if (factors.size() == 1) {
- return factors[0].value();
- }
- auto result = factors[0].value();
- for (int i = 1; i < static_cast<int>(factors.size()); ++i) {
- result = result * factors[i].value();
- }
- return result;
- };
-
- // Compute the number of output fragment of each warp
- auto warp_num_frag_m = f_get_inner_tile_product(-2);
- auto warp_num_frag_n = f_get_inner_tile_product(-1);
-
- Schedule& sch = state->sch;
- int tensor_ndim = static_cast<int>(sch->Get(state->block_rv)
- ->writes[0]
- ->source.as_or_throw<tvm::tirx::TensorVar>()
- .shape()
- .size());
- // The dimension of the tensor should be larger or same as that of the tensor intrin.
- TVM_FFI_ICHECK_GE(tensor_ndim, 2);
- int num_higher_dims = tensor_ndim - 2;
-
- auto index_map =
- tirx::IndexMap::FromFunc(tensor_ndim,
- // frag_shape_m and frag_shape_n are structural bindings that cannot
- // not be automatically captured until c++20
- [&, frag_shape_m = frag_shape_m,
- frag_shape_n = frag_shape_n](const ffi::Array<tvm::Var>& indices) {
- ffi::Array<PrimExpr> result;
- result.reserve(indices.size() + 4);
- for (int i = 0; i < num_higher_dims; ++i) {
- result.push_back(indices[i].as_or_throw<PrimExpr>());
- }
- const auto& m = indices[num_higher_dims];
- const auto& n = indices[num_higher_dims + 1];
- auto accum_m = floormod(m.as_or_throw<PrimExpr>(), frag_shape_m);
- auto accum_n = floormod(n.as_or_throw<PrimExpr>(), frag_shape_n);
- auto outer_m = floordiv(m.as_or_throw<PrimExpr>(), frag_shape_m);
- auto outer_n = floordiv(n.as_or_throw<PrimExpr>(), frag_shape_n);
-
- result.push_back(floordiv(outer_m, warp_num_frag_m));
- result.push_back(floordiv(outer_n, warp_num_frag_n));
- result.push_back(floormod(outer_m, warp_num_frag_m));
- result.push_back(floormod(outer_n, warp_num_frag_n));
- result.push_back(accum_m);
- result.push_back(accum_n);
- return result;
- });
- sch->TransformLayout(state->block_rv, 0, s_tir::TensorIndexType::kWrite, index_map,
- /*pad_value=*/std::nullopt, /*assume_injective_transform=*/true);
-
- return {state};
-}
-
-std::vector<State> MultiLevelTilingTensorCoreObj::AddWriteReuseTensorCore(
- TensorCoreState state) const {
- if (state->is_mma) {
- state->sch->WriteAt(state->tiles[2].back(), state->block_rv, 0, "mma.matrixC");
- state->sch->ReverseComputeInline(state->tensor_core_reindex_store);
- return {state};
- }
- // Add the cache write stage for Tensor Core
- Schedule& sch = state->sch;
- auto cache_write = sch->CacheWrite(state->block_rv, 0, "wmma.accumulator");
-
- // The compute block has been tiled by the warp shape and the fragment shape.
- // We need to bind the cache write block (from the accumulator to the shared memory) to the warp
- // id. The schedule is as follows:
- //
- // After adding cache write for wmma.accumulator, we will have
- // for i0, j0, i1, j1, accum_m, accum_n:
- // shared_mem[i0, j0, i1, j1, accum_m, accum_n] = accum[i0, j0, i1, j1, accum_m, accum_n]
- // for i0', j0', i1', j1', accum_m', accum_n':
- // global_mem[i0', j0', i1', j1', accum_m', accum_n'] =
- // shared_mem[i0', j0', i1', j1', accum_m', accum_n']
- // where i0' and j0' are already bound to the block id and warp id.
- //
- // To reduce the shared memory usage and allow efficient data movement, we will apply
- // transformations to generate the following schedule:
- //
- // for i1':
- // for i0_j0 (fused and bound to threadIdx.y):
- // for j1, accum_m, accum_n:
- // shared_mem[i0, j0, i1, j1, accum_m, accum_n] = accum[i0, j0, i1, j1, accum_m, accum_n]
- // for i0', j0', j1', accum_m', accum_n':
- // global_mem[i0', j0', i1', j1', accum_m', accum_n'] =
- // shared_mem[i0', j0', i1', j1', accum_m', accum_n']
- //
- // i1' is reordered to the outermost. This effectively allows only a row (i.e. loop i1') of the
- // fragments are moved to the shared memory and then to the global memory each time.
- // As a result, shared memory for the output will only have shape of [j1, accum_m, accum_n]
- // instead of [i0 * i1 * accum_m, j0 * j1 * accum_n].
-
- // Get the loops other than the innermost two loops (accum_m and accum_n).
- auto f_get_loops = [&](const SBlockRV& block_rv) -> std::array<LoopRV, 4> {
- ffi::Array<LoopRV> tensor_loops = sch->GetLoops(block_rv);
- TVM_FFI_ICHECK_GT(tensor_loops.size(), 6);
- return {tensor_loops[tensor_loops.size() - 6], tensor_loops[tensor_loops.size() - 5],
- tensor_loops[tensor_loops.size() - 4], tensor_loops[tensor_loops.size() - 3]};
- };
- {
- const auto& [i0, j0, i1, j1] = f_get_loops(state->write_reuse[0]);
- sch->Reorder({i1, i0, j0, j1});
- sch->ComputeAt(cache_write, i1, true);
- }
- {
- auto loops = f_get_loops(cache_write);
- const auto& i0 = loops[0];
- const auto& j0 = loops[1];
- auto fused = sch->Fuse({i0, j0});
- sch->Bind(fused, "threadIdx.y");
- }
-
- sch->ReverseComputeInline(state->tensor_core_reindex_store);
- auto loops = sch->GetLoops(cache_write);
- auto blockized_store = sch->Blockize(loops[loops.size() - 2]);
- sch->Annotate(blockized_store, tvm::s_tir::attr::kMetaScheduleAutoTensorize,
- state->intrin_group.store_intrin);
-
- ffi::Array<LoopRV> tensor_loops = sch->GetLoops(state->write_reuse[0]);
- TVM_FFI_ICHECK_GT(tensor_loops.size(), 5);
- sch->Fuse(ffi::Array<LoopRV>{tensor_loops.end() - 5, // The src shmem is always 2D
- tensor_loops.end()});
- AnnotateCooperativeFetching(&sch, state->write_reuse[0]);
- return {state};
-}
-
-std::vector<State> MultiLevelTilingTensorCoreObj::AddReadReuseTensorCore(
- TensorCoreState state) const {
- const ffi::Array<LoopRV>& r_tiles = state->tiles[r_indices_[1]];
- Schedule& sch = state->sch;
- TVM_FFI_CHECK(!r_tiles.empty(), ValueError)
- << "Cannot find the suitable reduction loop in the block";
-
- auto f_tensorize_load = [&](int read_index, ffi::String scope, ffi::String intrin_name) {
- auto cache_read = sch->CacheRead(state->block_rv, read_index, scope);
- state->sch->ComputeAt(cache_read, r_tiles.back(), true);
- ffi::String permuted_layout_annotate_value =
- state->is_mma ? std::string("s2l_") + std::string(read_index == 0 ? "A" : "B") : "";
- TileAndAnnotateTensorize(&sch, cache_read, intrin_name, permuted_layout_annotate_value);
- };
- f_tensorize_load(0, state->is_mma ? "mma.matrixA" : "wmma.matrix_a",
- state->intrin_group.load_a_intrin);
- f_tensorize_load(1, state->is_mma ? "mma.matrixB" : "wmma.matrix_b",
- state->intrin_group.load_b_intrin);
-
- for (int i = 0; i < 2; ++i) {
- const s_tir::SBlockRV cache_read = state->read_reuse.at(i);
- // Inline the reindex / padding block
- sch->ComputeInline(sch->GetProducers(cache_read)[0]);
- const s_tir::SBlockNode* cache_read_block =
- sch->GetSRef(cache_read)->StmtAs<s_tir::SBlockNode>();
- tirx::TensorVar cache_read_tensor =
- s_tir::GetNthAccessTensor(sch->state(), ffi::GetRef<s_tir::SBlock>(cache_read_block), 0,
- s_tir::TensorIndexType::kWrite);
- const DLDataType dtype = cache_read_tensor.dtype()->dtype;
- // Storage alignment is chosen from element storage width; this schedule rule uses scalar
- // cache-read tensors, so the old element-type-only test is preserved.
- if ((((dtype).code == kDLFloat) && ((dtype).bits == 16))) {
- sch->StorageAlign(cache_read, 0, -2, 32, 8);
- } else if (((dtype).code == kDLInt) && dtype.bits == 8) {
- sch->StorageAlign(cache_read, 0, -2, 32, 16);
- } else {
- TVM_PY_LOG(WARNING, logger) << "StorageAlign is not applied for data type " << dtype
- << ", shared memory accesses might be inefficient.";
- }
- }
- return {state};
-}
-
-std::vector<State> MultiLevelTilingTensorCoreObj::AddSoftwarePipeline(TensorCoreState state) const {
- if (!use_software_pipeline) {
- return {state};
- }
- // The current config is not suitable for software pipelining.
- if (r_indices_.size() < 2) {
- return {state};
- }
-
- Schedule& sch = state->sch;
- // Check reduction length after blockize.
- ffi::BigInt reduction_length = 1;
- for (int r_index : r_indices_) {
- const ffi::Array<LoopRV>& tiles = state->tiles[r_index];
- for (const LoopRV& tile : tiles) {
- const auto* extent = sch->Get(tile)->extent.as<IntImmNode>();
- TVM_FFI_ICHECK(extent != nullptr) << "Dynamic extent is not supported.";
- reduction_length *= extent->value;
- }
- }
- if (reduction_length <= 1) {
- return {state};
- }
-
- for (int i = 0; i < 2; ++i) {
- const s_tir::SBlockRV cache_read = state->read_reuse.at(i);
- if (state->is_mma) {
- // Add vector bytes for memhammer
- sch->Annotate(cache_read, tvm::s_tir::attr::kVectorBytes, 16);
- if (!state->use_async) {
- sch->Annotate(cache_read, tvm::s_tir::attr::kLocalStage, 1);
- sch->Annotate(cache_read, tvm::s_tir::attr::kDoubleBufferScope, 0);
- }
- } else {
- // Add local stage and double buffering
- sch->Annotate(cache_read, tvm::s_tir::attr::kManifestSharedMemoryLocalStage, 1);
- sch->Annotate(cache_read, tvm::s_tir::attr::kDoubleBufferScope, 0);
- }
- }
-
- // Add annotations of software pipeline
- //
- // Before pipelining, the original loop can be expressed as the pseudo code below:
- //
- // for k0 in [0, K0):
- // load tile k0 to registers
- // load tile k0 from registers to shared memory
- //
- // for k1 in [0, K1):
- // load fragment k1 of tile k0
- // compute matmul with fragment k1
- //
-
- // Inner software pipeline: Prefetch to tensor core fragment by one iteration
- // The following annotation for the inner loop is equivalent the pesudo code below:
- //
- // Pipelined inner loop:
- //
- // prologue:
- // load fragment 0
- // body:
- // for k1 in [0, K1 - 1):
- // load fragment k1 + 1
- // compute matmul with fragment k1
- // epilogue:
- // compute matmul with fragment K1 - 1
- //
- sch->Annotate(state->tiles[r_indices_[1]].back(), tvm::s_tir::attr::kSoftwarePipelineStage,
- ffi::Array<int64_t>{0, 0, 1});
- sch->Annotate(state->tiles[r_indices_[1]].back(), tvm::s_tir::attr::kSoftwarePipelineOrder,
- ffi::Array<int64_t>{0, 1, 2});
- if (state->is_mma && state->use_async) {
- sch->Annotate(state->tiles[r_indices_[0]].back(),
- tvm::s_tir::attr::kSoftwarePipelineAsyncStages, ffi::Array<int64_t>{0});
- sch->Annotate(state->tiles[r_indices_[0]].back(), tvm::s_tir::attr::kSoftwarePipelineStage,
- ffi::Array<int64_t>{0, 0, 1, 2, 2});
- sch->Annotate(state->tiles[r_indices_[0]].back(), tvm::s_tir::attr::kSoftwarePipelineOrder,
- ffi::Array<int64_t>{0, 1, 3, 2, 4});
- } else {
- // Outer software pipeline: Interleave the outer loop with the (pipelined) inner loop.
- // The prefetching stage of the inner pipeline is executed by one iteration in the outer loop.
- // The following annotation for the outer loop is equivalent the pesudo code below:
- //
- // Pipelined outer loop with nested inner pipeline:
- //
- // prologue:
- // load tile 0 to registers
- // load tile 0 from registers to shared memory
- //
- // // prologue of the inner pipeline
- // load fragment 0 of tile 0
- //
- // body:
- // for k0 in [0, K0 - 1):
- // load tile k0 + 1 to registers
- //
- // // body of the inner pipeline
- // for k1 in [0, K1 - 1):
- // load fragment k1 + 1 of tile k0
- // compute matmul with fragment k1 of tile k0
- //
- // load tile k0 + 1 from registers to shared memory
- //
- // // prologue of the inner pipeline
- // load fragment 0 of tile k0 + 1
- //
- // // epilogue of the inner pipeline
- // compute matmul with fragment K1 - 1 of tile k0
- //
- // epilogue:
- //
- // // body of the inner pipeline
- // for k1 in [0, K1 - 1):
- // load fragment k1 + 1 of tile K0 - 1
- // compute matmul with fragment k1 of tile K0 - 1
- //
- // // epilogue of the inner pipeline
- // compute matmul with fragment K1 - 1 of tile K0 - 1
- //
- sch->Annotate(state->tiles[r_indices_[0]].back(), tvm::s_tir::attr::kSoftwarePipelineStage,
- ffi::Array<int64_t>{0, 0, 0, 0, 0, 1, 1});
- sch->Annotate(state->tiles[r_indices_[0]].back(), tvm::s_tir::attr::kSoftwarePipelineOrder,
- ffi::Array<int64_t>{0, 3, 1, 4, 5, 2, 6});
- }
-
- return {state};
-}
-
-ffi::Optional<LoopRV> MultiLevelTilingTensorCoreObj::TransformWithTensorIntrin(
- TensorCoreStateObj* state, const ffi::String& intrin_name) const {
- SBlockRV block_rv = state->block_rv;
- const s_tir::AutoTensorizeMappingInfo& mapping_info = state->mapping_info;
- tirx::StmtSRef block_sref = state->sch->GetSRef(state->block_rv);
-
- // Add reindex stages
- const s_tir::SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref);
- // Hold the reference of the block before reindex
- const s_tir::SBlock block_before_reindex = ffi::GetRef<s_tir::SBlock>(block);
- if (block->reads.size() != 2 || block->writes.size() != 1) {
- // only matmul-like computation is allowed
- return std::nullopt;
- }
- state->tensor_core_reindex_store =
- state->sch->ReIndex(state->block_rv, 0, s_tir::TensorIndexType::kWrite);
- state->tensor_core_reindex_A =
- state->sch->ReIndex(state->block_rv, 0, s_tir::TensorIndexType::kRead);
- state->tensor_core_reindex_B =
- state->sch->ReIndex(state->block_rv, 1, s_tir::TensorIndexType::kRead);
-
- // Transform the layout of reindex tensors accordingly.
- // The index map defines the mapping for the computation block. We need to extract the sub index
- // map to transform the load and store block.
- TVM_FFI_ICHECK_EQ(mapping_info->mappings.size(), 1U); // assume only one mapping is present
- const tirx::IndexMap& index_map = mapping_info->mappings[0];
-
- // Find the correspondence between block iters and the iters in the index map.
- std::unordered_map<PrimVar, PrimVar, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>
- lhs_to_index_map_src;
- std::unordered_map<PrimVar, PrimExpr, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>
- rhs_to_index_map_tgt;
- std::unordered_set<PrimVar, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> unmapped_index_map_src;
- TVM_FFI_ICHECK_EQ(mapping_info->lhs_iters.size(), index_map->initial_indices.size());
- for (int i = 0; i < static_cast<int>(mapping_info->lhs_iters.size()); ++i) {
- lhs_to_index_map_src.insert_or_assign(mapping_info->lhs_iters[i]->var,
- index_map->initial_indices[i]);
- }
- // The number of result iters in the index map is equal or more than the number of rhs (the
- // tensor intrin) iters. When there are extra iters, these iters represent unmapped iters from
- // the lhs. They will be skipped during pattern matching for tensorization. An example of such
- // case is batch matmul, the batch dimension is kept after layout transformations and it will be
- // kept as a outer loop after tensorization.
- int offset = static_cast<int>(index_map->final_indices.size()) -
- static_cast<int>(mapping_info->rhs_iters.size());
- TVM_FFI_ICHECK_GE(offset, 0);
- for (int i = 0; i < offset; ++i) {
- auto var = index_map->final_indices[i].as<PrimVar>();
- TVM_FFI_ICHECK(var.has_value());
- unmapped_index_map_src.insert(var.value());
- }
- for (int i = offset; i < static_cast<int>(index_map->final_indices.size()); ++i) {
- rhs_to_index_map_tgt.insert_or_assign(mapping_info->rhs_iters[i - offset]->var,
- index_map->final_indices[i]);
- }
-
- auto f_get_sub_index_map = [&](const tirx::TensorVar& lhs_tensor,
- const ffi::Array<Range>& lhs_region) {
- std::vector<PrimVar> sub_index_map_src;
- std::vector<PrimExpr> sub_index_map_tgt;
- const tirx::TensorVar& rhs_tensor = mapping_info->lhs_tensor_map[lhs_tensor];
- for (const Range& range : lhs_region) {
- TVM_FFI_ICHECK(tvm::prim::IsOne(range->extent));
- auto var = range->min.as<PrimVar>();
- TVM_FFI_ICHECK(var.has_value());
- const PrimVar& lhs_representer = lhs_to_index_map_src.at(var.value());
- sub_index_map_src.push_back(lhs_representer);
- if (unmapped_index_map_src.count(lhs_representer)) {
- sub_index_map_tgt.push_back(lhs_representer);
- }
- }
- for (size_t i = 0; i < mapping_info->rhs_tensor_indices[rhs_tensor].size(); ++i) {
- auto var = mapping_info->rhs_tensor_indices[rhs_tensor][i].as<PrimVar>();
- TVM_FFI_ICHECK(var.has_value());
- sub_index_map_tgt.push_back(rhs_to_index_map_tgt.at(var.value()));
- }
- return tirx::IndexMap(sub_index_map_src, sub_index_map_tgt);
- };
-
- std::unordered_set<tirx::TensorVar, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> visited_tensors;
-
- ffi::Map<tirx::TensorVar, tirx::IndexMap> tensor_sub_index_map; // cache of the sub index map
- // associated with each tensor
-
- auto f_transform_tensor_layout = [&](s_tir::TensorIndexType index_type, int tensor_index) {
- const tirx::TensorVar& lhs_tensor = s_tir::GetNthAccessTensor(
- state->sch->state(), block_before_reindex, tensor_index, index_type);
- if (visited_tensors.count(lhs_tensor)) {
- return;
- }
- visited_tensors.insert(lhs_tensor);
- // Refresh block pointer (block sref is not invalidated)
- block = TVM_SREF_TO_SBLOCK(block_sref);
- const tvm::TensorRegion& reindexed_tensor_region = s_tir::GetNthAccessTensorRegion(
- state->sch->state(), ffi::GetRef<s_tir::SBlock>(block), tensor_index, index_type);
- auto sub_index_map = f_get_sub_index_map(lhs_tensor, reindexed_tensor_region->region);
- tensor_sub_index_map.Set(lhs_tensor, sub_index_map);
- state->sch->TransformLayout(state->block_rv, tensor_index, index_type, sub_index_map,
- /*pad_value=*/std::nullopt, /*assume_injective_transform=*/true);
- };
-
- for (int i = 0, n = block_before_reindex->reads.size(); i < n; ++i) {
- f_transform_tensor_layout(s_tir::TensorIndexType::kRead, i);
- }
- for (int i = 0, n = block_before_reindex->writes.size(); i < n; ++i) {
- f_transform_tensor_layout(s_tir::TensorIndexType::kWrite, i);
- }
-
- // Transform the layout of current block and reindex blocks
- auto f_transform_reindex_block_layout = [&](const SBlockRV& block_rv,
- s_tir::TensorIndexType tensor_type) {
- tirx::TensorVar tensor =
- s_tir::GetNthAccessTensor(state->sch->state(), state->sch->Get(block_rv), 0, tensor_type);
- const auto& sub_index_map = tensor_sub_index_map.at(tensor);
- state->sch->TransformBlockLayout(block_rv, sub_index_map);
- };
- f_transform_reindex_block_layout(state->tensor_core_reindex_store,
- s_tir::TensorIndexType::kWrite);
- f_transform_reindex_block_layout(state->tensor_core_reindex_A, s_tir::TensorIndexType::kRead);
- f_transform_reindex_block_layout(state->tensor_core_reindex_B, s_tir::TensorIndexType::kRead);
- state->sch->TransformBlockLayout(state->block_rv, index_map);
- return s_tir::TileWithTensorIntrin(state->sch, state->block_rv, intrin_name,
- /*allow_padding=*/true);
-}
-
-inline std::vector<State> MultiLevelTilingTensorCoreObj::TransformForTensorization(
- TensorCoreState state) const {
- // Do reindex and layout transformations.
- ffi::Optional<LoopRV> transformed_loop_rv =
- TransformWithTensorIntrin(state.operator->(), state->intrin_group.compute_intrin);
- if (!transformed_loop_rv.has_value()) {
- // The workload can't be tensorized.
- return {};
- }
-
- // Do blockize
- state->block_rv = state->sch->Blockize(transformed_loop_rv.value());
-
- // Add annotations for post processors.
- state->sch->Annotate(state->block_rv, tvm::s_tir::attr::kMetaScheduleAutoTensorize,
- state->intrin_group.compute_intrin);
- state->sch->Annotate(state->block_rv, tvm::s_tir::attr::kMetaScheduleAutoTensorizeInit,
- state->intrin_group.init_intrin);
- state->sch->Annotate(state->block_rv, tvm::s_tir::attr::kWarpExecution, 1);
- return {std::move(state)};
-}
-
-ScheduleRule ScheduleRule::MultiLevelTilingTensorCore(
- ffi::Array<ffi::Map<ffi::String, ffi::String>> intrin_groups, ffi::String structure,
- ffi::Optional<ffi::Array<ffi::String>> tile_binds, ffi::Optional<int64_t> max_innermost_factor,
- ffi::Optional<ffi::Array<int64_t>> vector_load_lens,
- ffi::Optional<ffi::Map<ffi::String, ffi::Any>> reuse_read,
- ffi::Optional<ffi::Map<ffi::String, ffi::Any>> reuse_write, bool use_software_pipeline) {
- static const auto scope_shared = runtime::MemoryScope::Get("shared");
- static const auto scope_shared_dyn = runtime::MemoryScope::Get("shared.dyn");
-
- if (tile_binds.has_value()) {
- for (const ffi::String& tile_bind : tile_binds.value()) {
- TVM_FFI_ICHECK_NE(tile_bind, "threadIdx.x")
- << "Cannot bind to threadIdx.x when using tensor core.";
- }
- }
- auto node = MultiLevelTilingInitCommon<MultiLevelTilingTensorCoreObj>(
- structure, tile_binds, max_innermost_factor, vector_load_lens, reuse_read, reuse_write);
-
- bool have_wmma_intrin_group = false;
- node->intrin_groups.reserve(intrin_groups.size());
- for (const auto& intrin_group_config : intrin_groups) {
- TensorCoreIntrinGroup group = TensorCoreIntrinGroup::FromConfig(intrin_group_config);
- if (group.compute_intrin.starts_with("wmma")) {
- have_wmma_intrin_group = true;
- }
- node->intrin_groups.emplace_back(group);
- }
-
- if (have_wmma_intrin_group) {
- auto reuse_scope = runtime::MemoryScope::Get(node->reuse_write_.scope);
- TVM_FFI_CHECK(node->reuse_write_.req == ReuseType::kMustReuse &&
- (reuse_scope == scope_shared || reuse_scope == scope_shared_dyn),
- ValueError)
- << "Shared memory write reuse must be enabled for MultiLevelTilingTensorCore.";
- }
-
- node->use_software_pipeline = use_software_pipeline;
- return ScheduleRule(ffi::UnsafeInit{}, node);
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- MultiLevelTilingTensorCoreObj::RegisterReflection();
- TensorCoreStateObj::RegisterReflection();
- refl::GlobalDef().def("s_tir.meta_schedule.ScheduleRuleMultiLevelTilingTensorCore",
- ScheduleRule::MultiLevelTilingTensorCore);
-}
-
-} // namespace meta_schedule
-} // namespace s_tir
-} // namespace tvm
diff --git a/src/s_tir/meta_schedule/schedule_rule/schedule_rule.cc b/src/s_tir/meta_schedule/schedule_rule/schedule_rule.cc
index ec87f79272..68e45106fd 100644
--- a/src/s_tir/meta_schedule/schedule_rule/schedule_rule.cc
+++ b/src/s_tir/meta_schedule/schedule_rule/schedule_rule.cc
@@ -173,47 +173,6 @@ ffi::Array<ScheduleRule> ScheduleRule::DefaultCUDA() {
};
}
-ffi::Array<ScheduleRule> ScheduleRule::DefaultCUDATensorCore() {
- ffi::Array<ffi::Map<ffi::String, ffi::String>> mma_intrin_groups = {
- // Tensor Core MMA
- {
- {"init", "mma_init_m16n8k8_f16"},
- {"load_a", "mma_load_m16n8k8_f16_A_shared_dyn"},
- {"load_b", "mma_load_m16n8k8_f16_B_shared_dyn"},
- {"compute", "mma_sync_m16n8k8_f16f16f16"},
- {"store", "mma_store_m16n8k8_f16_global"},
- },
- {
- {"init", "mma_init_m16n8k8_f32"},
- {"load_a", "mma_load_m16n8k8_f16_A_shared_dyn"},
- {"load_b", "mma_load_m16n8k8_f16_B_shared_dyn"},
- {"compute", "mma_sync_m16n8k8_f16f16f32"},
- {"store", "mma_store_m16n8k8_f32_global"},
- },
- };
- ffi::Array<ScheduleRule> results{
- ScheduleRule::ApplyCustomRule(),
- ScheduleRule::MultiLevelTilingTensorCore(
- /*intrin_groups=*/mma_intrin_groups,
- /*structure=*/"SSSRRSRS",
- /*tile_binds=*/ffi::Array<ffi::String>{"blockIdx.y", "blockIdx.x", "threadIdx.y"},
- /*max_innermost_factor=*/static_cast<int64_t>(4),
- /*vector_load_lens=*/ffi::Array<int64_t>{1, 2, 3, 4, 8, 16},
- /*reuse_read=*/
- ffi::Map<ffi::String, ffi::Any>{{"req", ffi::String("must")},
- {"levels", ffi::Array<int64_t>{4}}, //
- {"scope", ffi::String("shared.dyn")}},
- /*reuse_write=*/
- ffi::Map<ffi::String, ffi::Any>{{"req", ffi::String("no")},
- {"levels", ffi::Array<int64_t>{2}}, //
- {"scope", ffi::String("shared.dyn")}},
- /*use_software_pipeline=*/true) //
- };
- ffi::Array<ScheduleRule> append = ScheduleRule::DefaultCUDA();
- results.insert(results.end(), append.begin() + 1, append.end());
- return results;
-}
-
ffi::Array<ScheduleRule> ScheduleRule::DefaultHexagon() {
return {
ScheduleRule::ApplyCustomRule(),
@@ -405,8 +364,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.def("s_tir.meta_schedule.ScheduleRulePyScheduleRule", ScheduleRule::PyScheduleRule)
.def("s_tir.meta_schedule.ScheduleRuleDefaultLLVM", ScheduleRule::DefaultLLVM)
.def("s_tir.meta_schedule.ScheduleRuleDefaultCUDA", ScheduleRule::DefaultCUDA)
- .def("s_tir.meta_schedule.ScheduleRuleDefaultCUDATensorCore",
- ScheduleRule::DefaultCUDATensorCore)
.def("s_tir.meta_schedule.ScheduleRuleDefaultHexagon", ScheduleRule::DefaultHexagon)
.def("s_tir.meta_schedule.ScheduleRuleDefaultARM", ScheduleRule::DefaultARM);
}
diff --git a/src/s_tir/meta_schedule/space_generator/space_generator.cc b/src/s_tir/meta_schedule/space_generator/space_generator.cc
index 0af2c75cea..47cbab1cb9 100644
--- a/src/s_tir/meta_schedule/space_generator/space_generator.cc
+++ b/src/s_tir/meta_schedule/space_generator/space_generator.cc
@@ -56,24 +56,6 @@ ffi::String GetRuleKindFromTarget(const Target& target) {
if (target->kind->name == "hexagon") {
return "hexagon";
}
- if (target->kind->name == "cuda") {
- if (ffi::Optional<ffi::String> opt_sm = target->GetAttr<ffi::String>("arch")) {
- std::string sm = opt_sm.value();
- if (sm.starts_with("sm_")) {
- sm = sm.substr(3);
- try {
- if (std::stoi(sm) >= 75) {
- return "cuda-tensorcore";
- }
- } catch (const std::invalid_argument& e) {
- LOG(WARNING) << "ValueError: Unable to parse `target.arch`: " << sm
- << ". Details: " << e.what();
- }
- }
- }
- return "cuda";
- }
-
if (IsGPUTarget(target->kind->name)) {
return "cuda";
}
@@ -103,10 +85,6 @@ void SpaceGeneratorObj::InitializeWithTuneContext(const TuneContext& context) {
default_sch_rules = ScheduleRule::DefaultCUDA();
default_postprocs = Postproc::DefaultCUDA();
default_mutator_probs = Mutator::DefaultCUDA();
- } else if (kind == "cuda-tensorcore") {
- default_sch_rules = ScheduleRule::DefaultCUDATensorCore();
- default_postprocs = Postproc::DefaultCUDATensorCore();
- default_mutator_probs = Mutator::DefaultCUDATensorCore();
} else if (kind == "hexagon") {
default_sch_rules = ScheduleRule::DefaultHexagon();
default_postprocs = Postproc::DefaultHexagon();
diff --git a/src/s_tir/transform/inject_permuted_layout.cc b/src/s_tir/transform/inject_permuted_layout.cc
index 979912ab5f..98661dccf6 100644
--- a/src/s_tir/transform/inject_permuted_layout.cc
+++ b/src/s_tir/transform/inject_permuted_layout.cc
@@ -31,7 +31,6 @@
#include <tvm/s_tir/transform.h>
#include <tvm/sym/analyzer.h>
#include <tvm/tirx/function.h>
-#include <tvm/tirx/op/memory.h>
namespace tvm {
namespace s_tir {
@@ -40,21 +39,6 @@ using namespace tvm::tirx;
using namespace sym;
using namespace runtime;
-namespace {
-
-ffi::Optional<Var> GetBufferDataVar(const ffi::Any& data) {
- if (auto var = data.as<Var>()) {
- return var;
- }
- if (const auto* call = data.as<CallNode>();
- call && call->op.same_as(tirx::tensor_data_ptr_op()) && call->args.size() == 1) {
- return call->args[0].as<Var>();
- }
- return std::nullopt;
-}
-
-} // namespace
-
class PermutedLayoutInjector : public StmtExprMutatorWithAnalyzer {
public:
using StmtExprMutatorWithAnalyzer::Mutate;
@@ -63,7 +47,7 @@ class PermutedLayoutInjector : public StmtExprMutatorWithAnalyzer {
static Function Transform(Function func) {
Analyzer analyzer;
- auto new_body = ffi::make_object<PermutedLayoutInjector>(func, analyzer)
+ auto new_body = ffi::make_object<PermutedLayoutInjector>(analyzer)
->Mutate(func->body)
.ValueOrUnchanged(func->body);
auto func_node = func.CopyOnWrite();
@@ -71,14 +55,8 @@ class PermutedLayoutInjector : public StmtExprMutatorWithAnalyzer {
return func;
}
- explicit PermutedLayoutInjector(Function func, const Analyzer& analyzer)
- : StmtExprMutatorWithAnalyzer(analyzer) {
- for (const Var& param : func->params) {
- if (auto tensor = param.as<TensorVar>()) {
- tensor_map_.insert({tensor.value(), tensor.value()});
- }
- }
- }
+ explicit PermutedLayoutInjector(const Analyzer& analyzer)
+ : StmtExprMutatorWithAnalyzer(analyzer) {}
private:
ffi::Array<PrimExpr> PermuteIndices(PrimExpr row_idx, PrimExpr col_idx, int row_size) {
@@ -139,14 +117,6 @@ class PermutedLayoutInjector : public StmtExprMutatorWithAnalyzer {
}
UnchangedOr<Stmt> Mutate_(const SBlockNode* op, InplaceMode inplace_mode) final {
- // Record the mapping from tensor identity to tensor for later lookup.
- for (auto tensor : op->alloc_tensors) {
- tensor_map_.insert({tensor, tensor});
- }
- for (auto match_tensor : op->match_tensors) {
- tensor_map_.insert({match_tensor->tensor, match_tensor->tensor});
- }
-
if (op->annotations.count(tvm::s_tir::attr::kPermutedLayout) == 0 ||
!CheckAnnotation(op->annotations.at(tvm::s_tir::attr::kPermutedLayout))) {
return StmtExprMutatorWithAnalyzer::Mutate_(op, inplace_mode);
@@ -254,88 +224,8 @@ class PermutedLayoutInjector : public StmtExprMutatorWithAnalyzer {
load->loc);
}
- // Decode physical byte additions around a logical tensor address. Keep the
- // logical load untouched here: the intrinsic's explicit offset must be added
- // before applying the permutation, and visiting the load first would swizzle twice.
- TensorVar DecodePointer(Expr pointer, PrimExpr* byte_offset) {
- if (const auto* call = pointer.as<CallNode>()) {
- if (call->op.same_as(tirx::ptr_byte_offset_op())) {
- *byte_offset =
- *byte_offset +
- Mutate(call->args[1]).ValueOrUnchanged(call->args[1]).as_or_throw<PrimExpr>();
- return DecodePointer(call->args[0], byte_offset);
- }
- if (call->op.same_as(tirx::reinterpret_op())) {
- return DecodePointer(call->args[0], byte_offset);
- }
- if (call->op.same_as(tirx::address_of_op())) {
- const auto* load = call->args[0].as<TensorLoadNode>();
- TVM_FFI_ICHECK(load) << "Expected a tensor address for permuted layout";
- TensorVar tensor = load->source.as_or_throw<TensorVar>();
- auto indices = Mutate(load->indices)
- .ValueOrUnchanged(load->indices)
- .as_or_throw<ffi::Array<PrimExpr>>();
- auto flat_indices = tensor.ty()->ElemOffset(indices);
- // S-TIR flattening uses ElemOffset, including explicit strides and the
- // element offset, even when the tensor carries a default layout.
- TVM_FFI_ICHECK_EQ(flat_indices.size(), 1U);
- PrimType dtype = tensor.dtype();
- int bytes = (dtype.bits() * dtype.lanes() + 7) / 8;
- *byte_offset = *byte_offset + flat_indices[0] * bytes;
- return tensor;
- }
- }
- auto data_var = GetBufferDataVar(pointer);
- TVM_FFI_ICHECK(data_var.has_value()) << "Expected a tensor pointer, received " << pointer;
- auto it = tensor_map_.find(data_var.value());
- TVM_FFI_ICHECK(it != tensor_map_.end()) << "Unknown tensor pointer: " << pointer;
- return it->second;
- }
-
- Expr PermutePointer(Expr pointer, ffi::Optional<PrimExpr> offset = std::nullopt) {
- PrimExpr byte_offset = PrimExpr(0);
- TensorVar tensor = DecodePointer(pointer, &byte_offset);
- PrimType dtype = tensor.dtype();
- int bytes = (dtype.bits() * dtype.lanes() + 7) / 8;
- int row_size = CheckAndGetTensorRowSize(tensor);
- PrimExpr smem_offset = floordiv(byte_offset, bytes) + offset.value_or(PrimExpr(0));
- auto indices =
- PermuteIndices(floordiv(smem_offset, row_size), floormod(smem_offset, row_size), row_size);
- PrimExpr new_offset = analyzer_->Simplify(indices[0] * row_size + indices[1]);
- PrimExpr new_bytes = new_offset * bytes;
- PrimExpr remainder = analyzer_->Simplify(floormod(byte_offset, bytes));
- if (!prim::IsZero(remainder)) new_bytes = new_bytes + remainder;
- return Call(pointer->ty, tirx::ptr_byte_offset_op(), {tensor.data(), new_bytes});
- }
-
- UnchangedOr<Expr> Mutate_(const CallNode* op, InplaceMode inplace_mode) final {
- static const Op ptx_ldmatrix_op = Op::Get("tirx.ptx_legacy.ldmatrix");
- static const Op mma_store_op = Op::Get("tirx.cuda.mma_store");
- if (!permute_ || (!op->op.same_as(ptx_ldmatrix_op) && !op->op.same_as(mma_store_op))) {
- return StmtExprMutatorWithAnalyzer::Mutate_(op, inplace_mode);
- }
- bool is_ldmatrix = op->op.same_as(ptx_ldmatrix_op);
- int pointer_index = is_ldmatrix ? 5 : 2;
- ffi::Array<Expr> args;
- for (int i = 0; i < static_cast<int>(op->args.size()); ++i) {
- args.push_back(i == pointer_index ? op->args[i]
- : Mutate(op->args[i]).ValueOrUnchanged(op->args[i]));
- }
- if (is_ldmatrix) {
- PrimExpr offset = args[6].as_or_throw<PrimExpr>();
- args.Set(pointer_index, PermutePointer(args[pointer_index], offset));
- args.Set(6, IntImm(offset.ty(), 0));
- } else {
- args.Set(pointer_index, PermutePointer(args[pointer_index]));
- }
- return Call(op->ty, op->op, args, op->attrs, op->ty_args, op->loc);
- }
-
static constexpr size_t VECTORIZE_FACTOR = 8;
- static constexpr size_t BANK_SIZE_BYTES = 128;
- // Mapping from data Var of a TensorVar to TensorVar, for lookup
- std::unordered_map<Var, TensorVar> tensor_map_;
bool permute_ = false;
};
diff --git a/src/s_tir/transform/inject_software_pipeline.cc b/src/s_tir/transform/inject_software_pipeline.cc
index 9b80b1e55d..090da95f51 100644
--- a/src/s_tir/transform/inject_software_pipeline.cc
+++ b/src/s_tir/transform/inject_software_pipeline.cc
@@ -131,8 +131,6 @@ class PipelineOpaqueAccessRewriter {
static const Op load_matrix_sync = Op::Get("tirx.gpu_load_matrix_sync");
static const Op store_matrix_sync = Op::Get("tirx.gpu_store_matrix_sync");
static const Op mma_sync = Op::Get("tirx.gpu_mma_sync");
- static const Op ptx_ldmatrix_legacy = Op::Get("tirx.ptx_legacy.ldmatrix");
- static const Op ptx_mma_legacy = Op::Get("tirx.ptx_legacy.mma");
if (call->op.same_as(load_matrix_sync) || call->op.same_as(store_matrix_sync)) {
const TensorVar& tensor = tensor_data_to_tensor_.at(GetBufferDataVar(call->args[0]).value());
auto it = tensor_remap_.find(tensor);
@@ -159,10 +157,6 @@ class PipelineOpaqueAccessRewriter {
}
}
return Call(call->ty, call->op, new_args, call->attrs, call->ty_args, call->loc);
- } else if (call->op.same_as(ptx_mma_legacy)) {
- return RewriteTensorAccess(call, {6, 8, 10});
- } else if (call->op.same_as(ptx_ldmatrix_legacy)) {
- return RewriteTensorAccess(call, {3});
}
return call;
}
@@ -189,40 +183,6 @@ class PipelineOpaqueAccessRewriter {
return new_tensor_offset;
}
- Expr RewriteTensorAccess(const Call& call, const std::vector<int> arg_indices) {
- const auto scope_mma_matrixA = GetRegisteredMemoryScope("mma.matrixA");
- const auto scope_mma_matrixB = GetRegisteredMemoryScope("mma.matrixB");
-
- auto product = [](const ffi::Array<PrimExpr>& input) {
- return foldl([](PrimExpr a, PrimExpr b, Location loc) { return mul(a, b, loc); },
- IntImm::Int32(1), input);
- };
- ffi::Array<Expr> new_args = call->args;
- for (int i : arg_indices) {
- const TensorVar& tensor = tensor_data_to_tensor_.at(GetBufferDataVar(call->args[i]).value());
- auto it = tensor_remap_.find(tensor);
- if (it != tensor_remap_.end()) {
- const TensorVar& new_tensor = (*it).second;
- new_args.Set(i, new_tensor.data());
- PrimExpr old_index = call->args[i + 1].as_or_throw<PrimExpr>();
- PrimExpr offset =
- new_tensor.strides().empty() ? product(tensor.shape()) : new_tensor.strides()[0];
- if ((scope_mma_matrixA && tensor.scope() == *scope_mma_matrixA) ||
- (scope_mma_matrixB && tensor.scope() == *scope_mma_matrixB)) {
- // mma scope size will shrink by warp size
- // @see transform_mma_tensor_layout
- TVM_FFI_ICHECK_EQ(floormod(offset, 32).as_or_throw<IntImm>()->value, 0)
- << "mma scope size should be multiple of warp size";
- offset = floordiv(offset, 32);
- }
- PrimExpr new_index =
- old_index + floormod(pipeline_loop_->loop_var, new_tensor.shape()[0]) * offset;
- new_args.Set(i + 1, new_index);
- }
- }
- return Call(call->ty, call->op, new_args, call->attrs, call->ty_args, call->loc);
- }
-
const ffi::Map<Var, TensorVar>& tensor_data_to_tensor_;
const ffi::Map<TensorVar, TensorVar>& tensor_remap_;
const For& pipeline_loop_;
diff --git a/src/s_tir/transform/transform_mma_tensor_layout.cc b/src/s_tir/transform/transform_mma_tensor_layout.cc
index 2c2e956914..2990d89099 100644
--- a/src/s_tir/transform/transform_mma_tensor_layout.cc
+++ b/src/s_tir/transform/transform_mma_tensor_layout.cc
@@ -19,6 +19,7 @@
#include <tvm/ffi/cast.h>
#include <tvm/ffi/reflection/registry.h>
+#include <tvm/ir/prim/op.h>
#include <tvm/runtime/memory_scope.h>
#include <tvm/s_tir/analysis.h>
#include <tvm/s_tir/stmt.h>
@@ -26,7 +27,6 @@
#include <tvm/s_tir/transform.h>
#include <tvm/sym/analyzer.h>
#include <tvm/tirx/analysis.h>
-#include <tvm/tirx/index_map.h>
#include "../memory_scope.h"
@@ -40,8 +40,7 @@ using namespace tvm::tirx;
* own part of the matrix;
* 2. Rewrite access of mma.matrixC so it can access the correct part of the matrix.
* The reason why access of m16n8k8.matrix[A/B] tensor doesn't need this kind of rewrite is
- * that their access is through opaque access inside ldmatrix and mma_sync. Please refer to
- * get_index_[A/B] in python/tvm/tirx/tensor_intrin/cuda.py.
+ * that their access is through opaque accesses to per-thread fragments.
* We cannot use this kind of opaque access in matrixC too since the ptx stmatrix is only
* supported for sm90 or higher. Therefore, writeback of matrixC is limited to the
* transparent way.
@@ -153,10 +152,7 @@ class MmaTensorLayoutTransformer : public StmtExprMutator {
if (auto replacement = VarRemapGet(original_tensor).as<TensorVar>()) {
auto* n = store.CopyOnWrite();
if ((scope_mma_matrixC && original_tensor.scope() == *scope_mma_matrixC)) {
- const auto index_map_func = tvm::ffi::Function::GetGlobal("tirx.index_map_m16n8k8.matrixC");
- TVM_FFI_ICHECK(index_map_func.has_value());
- auto index_map = IndexMap::FromFunc(2, *index_map_func);
- auto new_indices = index_map->MapIndices(store->indices, analyzer);
+ auto new_indices = MatrixCIndices(store->indices);
n->dest = replacement.value();
n->indices = std::move(new_indices);
} else if ((scope_mma_matrixA && original_tensor.scope() == *scope_mma_matrixA) ||
@@ -184,10 +180,7 @@ class MmaTensorLayoutTransformer : public StmtExprMutator {
if (auto replacement = VarRemapGet(tensor).as<TensorVar>()) {
ffi::Array<PrimExpr> indices = load->indices;
if ((scope_mma_matrixC && tensor.scope() == *scope_mma_matrixC)) {
- const auto index_map_func = tvm::ffi::Function::GetGlobal("tirx.index_map_m16n8k8.matrixC");
- TVM_FFI_ICHECK(index_map_func.has_value());
- auto index_map = IndexMap::FromFunc(2, *index_map_func);
- indices = index_map->MapIndices(load->indices, analyzer);
+ indices = MatrixCIndices(load->indices);
} else {
TVM_FFI_ICHECK(false)
<< "TransformMmaTensorLayout requires " << tensor.scope()
@@ -200,6 +193,14 @@ class MmaTensorLayoutTransformer : public StmtExprMutator {
}
private:
+ // Preserve the per-thread accumulator layout independently of tensor-intrinsic registration.
+ ffi::Array<PrimExpr> MatrixCIndices(const ffi::Array<PrimExpr>& indices) {
+ TVM_FFI_ICHECK_EQ(indices.size(), 2U);
+ PrimExpr i = indices[0], j = indices[1];
+ return {analyzer->Simplify(floordiv(i, 16)), analyzer->Simplify(floordiv(j, 8)),
+ analyzer->Simplify(floormod(floordiv(i, 8), 2)), analyzer->Simplify(floormod(j, 2))};
+ }
+
sym::Analyzer analyzer;
};
diff --git a/src/tirx/transform/lower_warp_memory.cc b/src/tirx/transform/lower_warp_memory.cc
deleted file mode 100644
index cb495ee04c..0000000000
--- a/src/tirx/transform/lower_warp_memory.cc
+++ /dev/null
@@ -1,814 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied. See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-
-/*!
- * Lower warp memory to use local memory
- * and shuffle intrinsics.
- *
- * \file lower_warp_memory.cc
- */
-// Thanks to Andrew Adams and Vinod Grover for
-// explaining the concept of warp shuffle.
-#include <tvm/ffi/cast.h>
-#include <tvm/ffi/extra/structural_visit.h>
-#include <tvm/ffi/function.h>
-#include <tvm/ffi/reflection/registry.h>
-#include <tvm/ir/function.h>
-#include <tvm/ir/op.h>
-#include <tvm/ir/prim/expr.h>
-#include <tvm/ir/prim/op.h>
-#include <tvm/runtime/memory_scope.h>
-#include <tvm/sym/analyzer.h>
-#include <tvm/sym/pattern.h>
-#include <tvm/target/target.h>
-#include <tvm/tirx/analysis.h>
-#include <tvm/tirx/op/gpu.h>
-#include <tvm/tirx/op/memory.h>
-#include <tvm/tirx/op/region.h>
-#include <tvm/tirx/stmt_functor.h>
-#include <tvm/tirx/transform.h>
-
-#include <unordered_map>
-#include <unordered_set>
-#include <utility>
-
-#include "../../sym/pattern_match.h"
-#include "ir_utils.h"
-
-namespace tvm {
-namespace tirx {
-using namespace tvm::prim;
-
-namespace {
-
-class UpdatePointerStorageScope : public StmtExprMutator {
- public:
- using StmtExprMutator::Mutate_;
-
- UnchangedOr<Stmt> Mutate_(const BindNode* op, InplaceMode inplace_mode) final;
-
- UnchangedOr<Expr> Mutate_(const CallNode* op, InplaceMode inplace_mode) final;
-
- explicit UpdatePointerStorageScope(
- const std::unordered_map<Var, runtime::MemoryScope, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>&
- new_storage_scopes);
-
- private:
- std::unordered_map<const CallNode*, runtime::MemoryScope> tensor_scopes_;
-};
-
-Var WithStorageScope(const VarNode* buffer_var, runtime::MemoryScope storage_scope) {
- auto* ptr_type = buffer_var->ty.as<PtrTypeNode>();
- TVM_FFI_ICHECK(ptr_type) << "The provided variable is not of pointer type";
- return Var(buffer_var->name, PtrType(ptr_type->element_type, storage_scope), buffer_var->loc);
-}
-
-UpdatePointerStorageScope::UpdatePointerStorageScope(
- const std::unordered_map<Var, runtime::MemoryScope, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>&
- new_storage_scopes) {
- for (auto& kv : new_storage_scopes) {
- if (kv.first->ty.as<TensorTypeNode>()) {
- TensorVar tensor = ffi::GetRef<Var>(kv.first.get()).as_or_throw<TensorVar>();
- auto type = ffi::make_object<TensorTypeNode>(*tensor.ty().get());
- type->scope = kv.second;
- TensorVar replacement = TensorVar(tensor->name, TensorType(std::move(type)), tensor->loc);
- VarRemapSet(kv.first, replacement);
- } else {
- VarRemapSet(kv.first, WithStorageScope(kv.first.get(), kv.second));
- }
- }
-}
-
-UnchangedOr<Stmt> UpdatePointerStorageScope::Mutate_(const BindNode* op, InplaceMode inplace_mode) {
- const auto* call = op->value.as<CallNode>();
- if (call &&
- (call->op.same_as(tirx::alloc_tensor_op()) || call->op.same_as(tirx::decl_tensor_op()))) {
- if (auto mapped = VarRemapGet(op->var); mapped != nullptr) {
- tensor_scopes_.emplace(call, mapped.as_or_throw<TensorVar>().scope());
- auto result = StmtExprMutator::Mutate_(op, inplace_mode);
- tensor_scopes_.erase(call);
- return result;
- }
- }
- return StmtExprMutator::Mutate_(op, inplace_mode);
-}
-
-UnchangedOr<Expr> UpdatePointerStorageScope::Mutate_(const CallNode* op, InplaceMode inplace_mode) {
- auto result = StmtExprMutator::Mutate_(op, inplace_mode);
- if (auto it = tensor_scopes_.find(op); it != tensor_scopes_.end()) {
- Expr value = std::move(result).ValueOrUnchanged(ffi::GetRef<Expr>(op));
- auto call = value.as_or_throw<Call>();
- auto type = call->ty_args[0].as_or_throw<TensorType>();
- if (type->scope != it->second) {
- auto updated = ffi::make_object<TensorTypeNode>(*type.get());
- updated->scope = it->second;
- auto copy = ffi::make_object<CallNode>(*call.get());
- copy->ty_args.Set(0, TensorType(std::move(updated)));
- return ReinferMutatedCallType(Expr(std::move(copy)), op, inplace_mode);
- }
- return value;
- }
- if (!op->op.same_as(tirx::tensor_data_ptr_op())) return result;
- return ReinferMutatedCallType(std::move(result), op, inplace_mode);
-}
-
-} // namespace
-
-// Rewrite Rule
-//
-// There is no special warp memory in most GPUs.
-// Instead, we can stripe the data into threads
-// and store the data into local memory.
-//
-// This requires us to do the following rewriting:
-// - Rewrite allocation to use local memory.
-// - Rewrite store of warp memory to local store.
-// - Rewrite load of warp memory to local plus a shuffle.
-//
-// Define a generic shuffle intrinsic warp_shuffle(data, warp_index).
-// We can use the following rewriting rule
-//
-// Before rewrite,
-//
-// alloc warp warp_mem[n * width * m]
-// store warp_mem[m * warp_index + (width * m) * y + x]
-// load warp_mem[m * z + (width * m) * y + x]
-// subject to x \in [0, m), y \in [0, n)
-//
-// where width equals to the extent of threadIdx.x, which should
-// be no larger than the warp size
-//
-// After rewrite:
-//
-// alloc local local_mem[n * m]
-// store warp_mem[m * y + x]
-// warp_shuffle(load warp_mem[m * y + x], z)
-// subject to (m * y + x) is invariant to warp_index
-//
-// If width == warp size, we are shuffling on full warps.
-// Otherwise, we are virtually shuffling on sub-warps,
-// whose size equals to width. In this case, you can imagine
-// a warp only consists of `width` threads. Width is passed
-// as an argument to the shuffle primitive, and will be
-// lowered to the device code if the target supports.
-//
-// A limitation of this sub-warp approach is that users
-// cannot shuffle across the sub-warp boundary (i.e. shuffle
-// with threadIdx.y or threadIdx.z indices). It can be solved
-// via fusing threadIdx.x to the warp size, or improving the
-// analyzer to detect both 3 thread axes, which is left for
-// future improvements.
-
-// Algorithm
-//
-// To implement this rewrite rule, we can do the follow step:
-// For each warp memory alloc
-// - Use linear pattern detector on load index to find m
-// - Deduce n given width and alloc size
-// - Now that we have m, n, width, we can proceed with the rewrite
-
-// Visitor to find m in pattern
-// store warp_mem[m * warp_index + (width * m) * y + x]
-const VarNode* GetTensorVar(const Expr& expr) {
- if (const auto* var = expr.as<VarNode>()) {
- return var;
- }
- if (const auto* call = expr.as<CallNode>();
- call && call->op.same_as(tensor_data_ptr_op()) && call->args.size() == 1) {
- return call->args[0].as<VarNode>();
- }
- return nullptr;
-}
-
-// Hardware axis identity belongs to a lexical definition, not to a reused Var.
-// Each pass instance processes one device Function after SplitHostDevice.
-struct WarpThreadBinding {
- ffi::String tag;
- PrimExpr extent;
-};
-using WarpThreadBindings = std::unordered_map<const VarNode*, WarpThreadBinding>;
-using WarpIndexAliases = std::unordered_map<const VarNode*, PrimExpr>;
-
-PrimExpr ExpandWarpIndexAliases(const PrimExpr& index, const WarpIndexAliases& aliases) {
- return SubstituteWithDataTypeLegalization(index, [&](const Var& var) -> ffi::Optional<PrimExpr> {
- auto it = aliases.find(var.get());
- if (it != aliases.end()) return it->second;
- return std::nullopt;
- });
-}
-
-// Normalize equivalent hardware bindings only for index analysis. The selected
-// representative is the innermost live definition, so emitted indexes stay in scope.
-PrimExpr NormalizeWarpIndex(const PrimExpr& index, const WarpThreadBindings& active_bindings,
- const Var& warp_index, const WarpIndexAliases& aliases) {
- PrimExpr expanded = ExpandWarpIndexAliases(index, aliases);
- return SubstituteWithDataTypeLegalization(
- expanded, [&](const Var& var) -> ffi::Optional<PrimExpr> {
- auto it = active_bindings.find(var.get());
- if (it != active_bindings.end() && it->second.tag == "threadIdx.x" &&
- !var.same_as(warp_index)) {
- return prim::cast(warp_index.as_or_throw<PrimVar>(), var->ty.as_or_throw<PrimType>());
- }
- return std::nullopt;
- });
-}
-
-class WarpStoreCoeffFinder : public StmtExprVisitor {
- public:
- WarpStoreCoeffFinder(const VarNode* tensor, const WarpThreadBindings& bindings,
- WarpThreadBindings active_bindings, Var warp_index,
- sym::AnalyzerObj* analyzer, const WarpIndexAliases& aliases)
- : tensor_(tensor),
- bindings_(bindings),
- active_bindings_(std::move(active_bindings)),
- warp_index_(warp_index),
- analyzer_(analyzer),
- aliases_(aliases) {}
- // find the warp co-efficient in the statement given the warp size
- int Find(const Stmt& stmt) {
- this->Visit(stmt);
- return warp_coeff_;
- }
-
- private:
- ffi::Optional<VisitInterrupt> Visit_(const RegionStmtNode* op) final {
- // Operands are evaluated before the new launch binding comes into scope.
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(op->args));
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(op->attrs));
- Var previous_index = warp_index_;
- auto previous_bindings = active_bindings_;
- if (op->op.same_as(tirx::launch_thread_op())) {
- PrimVar var = op->body_params[0].as_or_throw<PrimVar>();
- const auto& binding = bindings_.at(var.get());
- active_bindings_.insert_or_assign(var.get(), binding);
- if (binding.tag == "threadIdx.x") warp_index_ = var;
- }
- auto result = Visit(op->body);
- active_bindings_ = std::move(previous_bindings);
- warp_index_ = previous_index;
- return result;
- }
-
- void UpdateCoefficient(int64_t coefficient) {
- TVM_FFI_ICHECK(warp_index_.defined())
- << "Warp memory store must be inside a threadIdx.x launch";
- TVM_FFI_ICHECK_GT(coefficient, 0);
- if (warp_coeff_ != 0) {
- TVM_FFI_ICHECK_EQ(warp_coeff_, coefficient)
- << "LowerWarpMemory requires compatible store coefficients across launch bindings";
- } else {
- warp_coeff_ = coefficient;
- }
- }
-
- /// Visitor implementation
- ffi::Optional<VisitInterrupt> Visit_(const CallNode* op) final {
- static const Op mma_fill_op = Op::Get("tirx.cuda.mma_fill");
- static const Op ptx_ldmatrix_legacy_op = Op::Get("tirx.ptx_legacy.ldmatrix");
- if (op->op.same_as(mma_fill_op) && GetTensorVar(op->args[1]) == tensor_) {
- auto* local_size = op->args[0].as<IntImmNode>();
- TVM_FFI_ICHECK(local_size) << "Integer expected for the first argument of mma_fill";
- UpdateCoefficient(local_size->value.as<int>().value());
- } else if (op->op.same_as(ptx_ldmatrix_legacy_op) && GetTensorVar(op->args[3]) == tensor_) {
- // ldmatrix writes the warp tensor; its local_offset carries
- // ``... + lift(local_size) * tx`` from which the warp coefficient
- // is derived.
- UpdatePattern(op->args[4].as_or_throw<PrimExpr>());
- }
- // mma_store/ptx_mma_legacy only *use* the warp tensor
- // (read+rewrite); WarpStoreCoeffFinder relies on ldmatrix/mma_fill
- // (the actual stores) for the warp coefficient.
-
- return StmtExprVisitor::Visit_(op);
- }
-
- ffi::Optional<VisitInterrupt> Visit_(const TensorStoreNode* op) final {
- if (op->dest.as_or_throw<TensorVar>().get() != tensor_) {
- return StmtExprVisitor::Visit_(op);
- }
-
- TVM_FFI_ICHECK_EQ(op->indices.size(), 1) << "Expected flat memory to use as warp memory. "
- << "Has FlattenTensor been run?";
-
- PrimExpr index = op->indices[0];
- PrimType value_ty = op->value.ty();
- if (value_ty.lanes() != 1) {
- sym::PVar<PrimExpr> base;
- TVM_FFI_ICHECK(sym::ramp(base, 1, value_ty.lanes()).Match(index))
- << "LowerWarpMemory failed due to store index=" << index
- << ", can only handle continuous store";
- UpdatePattern(base.Eval());
-
- index = base.Eval();
- }
-
- UpdatePattern(index);
- return std::nullopt;
- }
-
- void UpdatePattern(const PrimExpr& index) {
- TVM_FFI_ICHECK(warp_index_.defined())
- << "Warp memory store must be inside a threadIdx.x launch";
- PrimExpr normalized_index = NormalizeWarpIndex(index, active_bindings_, warp_index_, aliases_);
- ffi::Array<PrimExpr> m =
- sym::DetectLinearEquation(normalized_index, {warp_index_.as_or_throw<PrimVar>()});
- TVM_FFI_ICHECK_EQ(m.size(), 2U)
- << "LowerWarpMemory failed. Could not simplify the store index `" << index
- << "` into the form ax + by + cz + ... Warp memory is approximated by storing values in "
- "thread local registers and shuffling values between these registers. Currently only "
- "linear equation indices are supported.";
- PrimExpr mcoeff = analyzer_->canonical_simplify(m[0]);
- const auto* mcoeff_as_int = mcoeff.as<IntImmNode>();
- TVM_FFI_ICHECK(mcoeff_as_int && mcoeff_as_int->value > 0)
- << "LowerWarpMemory failed due to store index=" << index
- << ", require positive constant coefficient on warp index " << warp_index_ << " but get "
- << mcoeff;
-
- UpdateCoefficient(mcoeff_as_int->value.as<int>().value());
- }
-
- // The tensor variable
- const VarNode* tensor_;
- const WarpThreadBindings& bindings_;
- WarpThreadBindings active_bindings_;
- // The active lexical warp index.
- Var warp_index_{ffi::UnsafeInit{}};
- // the coefficient
- int64_t warp_coeff_{0};
- // analyzer.
- sym::AnalyzerObj* analyzer_;
- const WarpIndexAliases& aliases_;
-};
-
-// Collect each lexical launch definition and validate compatible hardware widths.
-class WarpIndexFinder : public StmtExprVisitor {
- public:
- explicit WarpIndexFinder(int warp_size, WarpThreadBindings enclosing_bindings)
- : warp_size_(warp_size), bindings_(std::move(enclosing_bindings)) {
- for (const auto& [var, binding] : bindings_) CheckWidth(binding);
- }
-
- std::pair<WarpThreadBindings, int> Find(const Stmt& stmt) {
- this->Visit(stmt);
- TVM_FFI_ICHECK_GT(width_, 0) << "Cannot find threadIdx.x within the scope of warp memory";
- return {std::move(bindings_), width_};
- }
-
- private:
- ffi::Optional<VisitInterrupt> Visit_(const RegionStmtNode* op) final {
- if (op->op.same_as(tirx::launch_thread_op())) {
- WarpThreadBinding binding{op->args[0].as_or_throw<StringImm>()->value,
- op->args[1].as_or_throw<PrimExpr>()};
- CheckWidth(binding);
- bindings_.insert_or_assign(op->body_params[0].as_or_throw<PrimVar>().get(),
- std::move(binding));
- }
- return StmtExprVisitor::Visit_(op);
- }
-
- void CheckWidth(const WarpThreadBinding& binding) {
- if (binding.tag != "threadIdx.x") return;
- const auto* extent = binding.extent.as<IntImmNode>();
- TVM_FFI_ICHECK(extent && extent->value > 0 && extent->value <= warp_size_ &&
- warp_size_ % extent->value == 0)
- << "Expect threadIdx.x size to be a positive factor of warp size (" << warp_size_
- << "), but got " << binding.extent;
- int width = extent->value.as<int>().value();
- TVM_FFI_ICHECK(width_ == 0 || width_ == width)
- << "LowerWarpMemory requires compatible threadIdx.x extents across launch bindings";
- width_ = width;
- }
-
- int warp_size_;
- int width_{0};
- WarpThreadBindings bindings_;
-};
-// Mutator to change the read pattern
-class WarpAccessRewriter : public StmtExprMutator {
- public:
- using StmtExprMutator::Mutate;
- using StmtExprMutator::Mutate_;
- explicit WarpAccessRewriter(int warp_size, sym::AnalyzerObj* analyzer,
- WarpThreadBindings enclosing_bindings, Var warp_index,
- const WarpIndexAliases& aliases)
- : warp_size_(warp_size),
- warp_index_(warp_index),
- analyzer_(analyzer),
- bindings_(enclosing_bindings),
- active_bindings_(std::move(enclosing_bindings)),
- aliases_(aliases) {}
- // Rewrite the AllocTensor statement which transforms
- // warp memory to local memory.
- // \param op The allocation binding for warp memory.
- // \param tensor_call The matched allocation Call.
- // \param body The remaining statements (siblings) that use this tensor.
- Stmt Rewrite(const BindNode* op, const CallNode* tensor_call, Stmt body) {
- tvm::Tuple shape = tvm::Tuple(tensor_call->ty_args[0].as_or_throw<TensorType>()->shape);
- DLDataType dtype = tensor_call->ty_args[0].as_or_throw<TensorType>()->dtype->dtype;
- PrimType element_type(dtype);
- tensor_ = op->var.get();
- int64_t alloc_size = 1;
- for (const auto& dim : shape->fields) {
- if (const IntImmNode* int_size = dim.as<IntImmNode>()) {
- alloc_size = static_cast<int64_t>(alloc_size * int_size->value);
- } else {
- alloc_size = 0;
- }
- }
- TVM_FFI_ICHECK_GT(alloc_size, 0) << "warp memory only support constant alloc size";
- alloc_size *= element_type.lanes();
- std::tie(bindings_, width_) =
- ffi::make_object<WarpIndexFinder>(warp_size_, bindings_)->Find(body);
- warp_coeff_ = ffi::make_object<WarpStoreCoeffFinder>(tensor_, bindings_, active_bindings_,
- warp_index_, analyzer_, aliases_)
- ->Find(body);
-
- // Align the local memory size. The number of elements may not
- // be a multiple of width_ * warp_coeff_; round it up.
- int factor = width_ * warp_coeff_;
- TVM_FFI_ICHECK_NE(factor, 0) << "Divide by zero";
- warp_group_ = (alloc_size + (factor - 1)) / factor;
- alloc_size = warp_group_ * factor;
-
- auto type = ffi::make_object<TensorTypeNode>(*op->var.as_or_throw<TensorVar>().ty().get());
- static const auto local = runtime::MemoryScope::Get("local");
- type->scope = local;
- type->shape = {IntImm::Int32(alloc_size / width_)};
- type->strides = {};
- type->elem_offset = IntImm(op->var.as_or_throw<TensorVar>().elem_offset().ty(), 0);
- TensorVar new_buf =
- TensorVar(op->var.as_or_throw<TensorVar>()->name, TensorType(std::move(type)),
- op->var.as_or_throw<TensorVar>()->loc);
- new_tensor_ = new_buf;
- Stmt rewritten_body = this->Mutate(body, InplaceMode::kDisallow).ValueOrUnchanged(body);
- ffi::Array<Expr> args = tensor_call->args;
- if (!args.empty()) {
- args.Set(0, Mutate(args[0], InplaceMode::kDisallow).ValueOrUnchanged(args[0]));
- }
-
- return SeqStmt({Bind(new_buf,
- Call(new_buf.ty(), tirx::alloc_tensor_op(), args, tensor_call->attrs,
- {new_buf.ty()}, tensor_call->loc),
- op->loc),
- rewritten_body});
- }
-
- protected:
- UnchangedOr<Stmt> Mutate_(const RegionStmtNode* op, InplaceMode inplace_mode) final {
- if (!op->op.same_as(tirx::launch_thread_op()))
- return StmtExprMutator::Mutate_(op, inplace_mode);
- PrimExpr old_extent = op->args[1].as_or_throw<PrimExpr>();
- PrimExpr extent = Mutate(old_extent, inplace_mode).ValueOrUnchanged(old_extent);
- PrimVar var = op->body_params[0].as_or_throw<PrimVar>();
- Var previous_index = warp_index_;
- auto previous_bindings = active_bindings_;
- const auto& binding = bindings_.at(var.get());
- active_bindings_.insert_or_assign(var.get(), binding);
- if (binding.tag == "threadIdx.x") warp_index_ = var;
- Stmt body = Mutate(op->body, inplace_mode).ValueOrUnchanged(op->body);
- active_bindings_ = std::move(previous_bindings);
- warp_index_ = previous_index;
- return RegionStmt(op->op, {op->args[0], extent}, op->body_params, op->attrs, body,
- op->result_vars, op->loc);
- }
-
- Expr RewriteIndicesAt(const CallNode* op, const std::vector<int>& indices) {
- ffi::Array<Expr> new_args = op->args;
- for (int i : indices) {
- // Preserve the pointer operand as an Expr and narrow only its scalar index.
- if (GetTensorVar(op->args[i]) == tensor_) {
- PrimExpr local_index = SplitIndexByGroup(op->args[i + 1].as_or_throw<PrimExpr>()).first;
- new_args.Set(i, new_tensor_.data());
- new_args.Set(i + 1, local_index);
- }
- }
- return Call(op->ty, op->op, new_args, op->attrs, {}, op->loc);
- }
-
- UnchangedOr<Expr> Mutate_(const CallNode* op, InplaceMode inplace_mode) override {
- static const Op mma_store_op = Op::Get("tirx.cuda.mma_store");
- static const Op mma_fill_op = Op::Get("tirx.cuda.mma_fill");
- static const Op ptx_mma_legacy_op = Op::Get("tirx.ptx_legacy.mma");
- static const Op ptx_ldmatrix_legacy_op = Op::Get("tirx.ptx_legacy.ldmatrix");
- if (op->op.same_as(mma_store_op)) {
- return RewriteIndicesAt(op, {3});
- }
-
- if (op->op.same_as(mma_fill_op)) {
- return RewriteIndicesAt(op, {1});
- }
-
- // Legacy variants: (ptr_var, offset) pairs in apache positions.
- if (op->op.same_as(ptx_mma_legacy_op)) {
- return RewriteIndicesAt(op, {6, 8, 10});
- }
- if (op->op.same_as(ptx_ldmatrix_legacy_op)) {
- // args: trans, num, type, local_ptr, local_offset, smem_ptr_call, smem_offset
- // Only local_ptr is a raw warp tensor Var; smem_ptr is an
- // pointer expression referencing shared memory.
- return RewriteIndicesAt(op, {3});
- }
-
- return StmtExprMutator::Mutate_(op, inplace_mode);
- }
-
- UnchangedOr<Expr> Mutate_(const VarNode* op, InplaceMode inplace_mode) override {
- if (def_region_kind() == kTVMFFIDefRegionKindNone) {
- TVM_FFI_ICHECK(op != tensor_) << "Cannot access address of warp memory directly";
- }
- return StmtExprMutator::Mutate_(op, inplace_mode);
- }
-
- UnchangedOr<Stmt> Mutate_(const TensorStoreNode* op, InplaceMode inplace_mode) override {
- // The source is a memory access, not a direct address use checked by the Var hook.
- auto value = Mutate(op->value, inplace_mode);
- auto indices =
- Mutate(op->indices, inplace_mode).as_or_throw<UnchangedOr<ffi::Array<PrimExpr>>>();
- TensorStore store = ffi::GetRef<TensorStore>(op);
- if (!value.UnchangedOrSameAs(op->value) || !indices.UnchangedOrSameAs(op->indices)) {
- auto* n = store.CopyOnWrite();
- n->value = std::move(value).ValueOrUnchanged(op->value);
- n->indices = std::move(indices).ValueOrUnchanged(op->indices);
- }
-
- if (store->dest.as_or_throw<TensorVar>().get() == tensor_) {
- TVM_FFI_ICHECK_EQ(store->indices.size(), 1) << "Expected flat memory to use as warp memory. "
- << "Has FlattenTensor been run?";
-
- auto [local_index, group] = SplitIndexByGroup(store->indices[0]);
- (void)group; // https://gcc.gnu.org/bugzilla/show_bug.cgi?id=81767
-
- auto writer = store.CopyOnWrite();
- writer->dest = new_tensor_;
- writer->indices = {local_index};
- }
-
- return store;
- }
-
- UnchangedOr<PrimExpr> Mutate_(const TensorLoadNode* op, InplaceMode inplace_mode) override {
- // Keep the memory source opaque to the direct-address check in the Var hook.
- auto indices =
- Mutate(op->indices, inplace_mode).as_or_throw<UnchangedOr<ffi::Array<PrimExpr>>>();
- TensorLoad load = ffi::GetRef<TensorLoad>(op);
- if (!indices.UnchangedOrSameAs(op->indices)) {
- load.CopyOnWrite()->indices = std::move(indices).ValueUnchecked();
- }
-
- if (load->source.as_or_throw<tvm::tirx::TensorVar>().get() != tensor_) {
- return load;
- }
-
- TVM_FFI_ICHECK_EQ(op->indices.size(), 1) << "Expected flat memory to use as warp memory. "
- << "Has FlattenTensor been run?";
-
- auto [local_index, group] = SplitIndexByGroup(op->indices[0]);
- // invariance: local index must do not contain warp id
- auto walkfn = [this](const Var& var) -> ffi::Expected<ffi::WalkResult> {
- return var.get() == warp_index_.get() ? ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
- : ffi::WalkResult::Advance();
- };
- TVM_FFI_ICHECK(!ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(local_index, walkfn).has_value())
- << "LowerWarpMemory failed to rewrite load to shuffle for index " << op->indices[0]
- << " local_index=" << local_index;
-
- load = TensorLoad(new_tensor_, {local_index}, load->loc);
-
- if (analyzer_->CanProveEqual(group, warp_index_.as_or_throw<PrimExpr>())) {
- return load;
- }
-
- PrimExpr mask = Call(PrimType::UInt(32), gpu_warp_activemask_op(), {}).as_or_throw<PrimExpr>();
- return Call(load.ty(), gpu_warp_shuffle_op(),
- ffi::Array<PrimExpr>{mask, load, group, width_, warp_size_})
- .as_or_throw<PrimExpr>();
- }
-
- // Split the index to the two component
- // <local_index, source_index>
- // local index is the index in the local
- // source index is the corresponding source index
- // in this access pattern.
- std::pair<PrimExpr, PrimExpr> SplitIndexByGroup(const PrimExpr& original_index) {
- TVM_FFI_ICHECK(warp_index_.defined())
- << "Warp memory access must be inside a threadIdx.x launch";
- PrimExpr index = NormalizeWarpIndex(original_index, active_bindings_, warp_index_, aliases_);
- PrimType index_ty = index.ty();
- if (index_ty.lanes() != 1) {
- sym::PVar<PrimExpr> base;
- TVM_FFI_ICHECK(sym::ramp(base, 1, index_ty.lanes()).Match(index));
-
- auto [local_index, group] = SplitIndexByGroup(base.Eval());
- local_index = prim::Ramp(local_index, IntImm(local_index.ty(), 1), index_ty.lanes());
- return std::make_pair(local_index, group);
- }
- PrimExpr m = IntImm(index_ty, warp_coeff_);
-
- // simple case, warp index is on the highest.
- if (warp_group_ == 1) {
- PrimExpr x = analyzer_->canonical_simplify(indexmod(index, m));
- PrimExpr z = analyzer_->canonical_simplify(indexdiv(index, m));
- return std::make_pair(x, z);
- } else {
- PrimExpr x = analyzer_->canonical_simplify(indexmod(index, m));
- PrimExpr y = index / MakeConst(index_ty, warp_coeff_ * width_);
- y = y * m + x;
- PrimExpr z = indexdiv(indexmod(index, IntImm(index_ty, warp_coeff_ * width_)), m);
- return std::make_pair(analyzer_->canonical_simplify(y), analyzer_->canonical_simplify(z));
- }
- }
-
- private:
- // the warp size
- int warp_size_{0};
- // The tensor variable
- const VarNode* tensor_;
- // The fresh local tensor replacing the warp-scoped definition.
- TensorVar new_tensor_{ffi::UnsafeInit{}};
- // number of threads involved in one shuffle
- int width_{0};
- // Warp index
- Var warp_index_{ffi::UnsafeInit{}};
- // the coefficient m
- int warp_coeff_{0};
- // the coefficient n
- int warp_group_{0};
- // Internal analyzer
- sym::AnalyzerObj* analyzer_;
- WarpThreadBindings bindings_;
- WarpThreadBindings active_bindings_;
- const WarpIndexAliases& aliases_;
-};
-
-// Bind bound information of variables to make analyzer more effective
-// TODO(tqchen): consider a pass to inline the bound info into the expr
-// so analysis can be context independent.
-class BindVarBoundInfo : public StmtExprVisitor {
- public:
- ffi::Optional<VisitInterrupt> Visit(ffi::AnyView value) override {
- if (value.as<ExprNode>()) return std::nullopt;
- return StmtExprVisitor::Visit(value);
- }
- explicit BindVarBoundInfo(sym::AnalyzerObj* analyzer, WarpIndexAliases* aliases)
- : analyzer_(analyzer), aliases_(aliases) {}
-
- ffi::Optional<VisitInterrupt> Visit_(const BindNode* op) final {
- if (auto value = op->value.as<PrimExpr>();
- value && op->var.as<PrimVar>() && value->ty().IsScalar() &&
- value->ty().MatchesCode(DLDataTypeCode::kDLInt, DLDataTypeCode::kDLUInt) &&
- SideEffect(*value) <= CallEffectKind::kPure) {
- // Preserve thread dependence even for unit extents. Analyzer range
- // simplification would erase that dependence before coefficient matching.
- aliases_->insert_or_assign(op->var.get(), ExpandWarpIndexAliases(*value, *aliases_));
- }
- return StmtExprVisitor::Visit_(op);
- }
-
- ffi::Optional<VisitInterrupt> Visit_(const ForNode* op) final {
- const Var& loop_var = op->loop_var;
- analyzer_->Bind(loop_var, Range::FromMinExtent(op->min, op->extent));
- return StmtExprVisitor::Visit_(op);
- }
-
- ffi::Optional<VisitInterrupt> Visit_(const RegionStmtNode* op) final {
- if (op->op.same_as(tirx::launch_thread_op())) {
- PrimVar var = op->body_params[0].as_or_throw<PrimVar>();
- PrimExpr extent = op->args[1].as_or_throw<PrimExpr>();
- Range dom = Range::FromMinExtent(IntImm(extent.ty(), 0), extent);
- analyzer_->Bind(var, dom);
- }
- return StmtExprVisitor::Visit_(op);
- }
-
- protected:
- // internal analyzer.
- sym::AnalyzerObj* analyzer_;
- WarpIndexAliases* aliases_;
-};
-
-// Mutator to change the read pattern
-class WarpMemoryRewriter : public StmtExprMutator {
- public:
- using StmtExprMutator::Mutate;
- using StmtExprMutator::Mutate_;
- UnchangedOr<ffi::Any> Mutate(ffi::AnyView input, InplaceMode inplace_mode) override {
- if (input.as<ExprNode>()) return ffi::Unchanged();
- return StmtExprMutator::Mutate(input, inplace_mode);
- }
- explicit WarpMemoryRewriter(int warp_size) : warp_size_(warp_size) {}
-
- Stmt Rewrite(Stmt stmt) {
- if (warp_size_ == 1) return stmt;
- auto binder = ffi::make_object<BindVarBoundInfo>(analyzer_.get(), &aliases_);
- binder->Visit(stmt);
- stmt = Mutate(stmt, InplaceMode::kAllow).ValueOrUnchanged(stmt);
- return stmt;
- }
-
- // Keep the old variables alive until UpdatePointerStorageScope reads their types.
- std::unordered_map<Var, runtime::MemoryScope, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>
- new_storage_scopes_;
-
- private:
- UnchangedOr<Stmt> Mutate_(const RegionStmtNode* op, InplaceMode inplace_mode) final {
- if (!op->op.same_as(tirx::launch_thread_op()))
- return StmtExprMutator::Mutate_(op, inplace_mode);
- PrimVar var = op->body_params[0].as_or_throw<PrimVar>();
- auto previous_bindings = active_bindings_;
- Var previous_index = warp_index_;
- ffi::String tag = op->args[0].as_or_throw<StringImm>()->value;
- active_bindings_.insert_or_assign(var.get(),
- WarpThreadBinding{tag, op->args[1].as_or_throw<PrimExpr>()});
- if (tag == "threadIdx.x") warp_index_ = var;
- Stmt body = Mutate(op->body, inplace_mode).ValueOrUnchanged(op->body);
- active_bindings_ = std::move(previous_bindings);
- warp_index_ = previous_index;
- return RegionStmt(op->op, op->args, op->body_params, op->attrs, body, op->result_vars, op->loc);
- }
-
- UnchangedOr<Stmt> Mutate_(const SeqStmtNode* op, InplaceMode inplace_mode) {
- // Process SeqStmt to find warp AllocTensor and gather remaining siblings as body.
- ffi::Array<Stmt> new_seq;
- bool changed = false;
- for (size_t i = 0; i < op->seq.size(); ++i) {
- const auto* alloc = op->seq[i].as<BindNode>();
- static const auto warp = runtime::MemoryScope::Get("warp");
- static const auto local = runtime::MemoryScope::Get("local");
- if (const auto* call = alloc ? alloc->value.as<CallNode>() : nullptr;
- call && call->op.same_as(alloc_tensor_op()) &&
- call->ty_args[0].as_or_throw<TensorType>()->scope == warp) {
- new_storage_scopes_[alloc->var] = local;
- // Gather remaining siblings as the "body" for rewriting.
- ffi::Array<Stmt> remaining;
- for (size_t j = i + 1; j < op->seq.size(); ++j) {
- remaining.push_back(op->seq[j]);
- }
- Stmt body = SeqStmt(remaining);
- auto rewriter = ffi::make_object<WarpAccessRewriter>(
- warp_size_, analyzer_.get(), active_bindings_, warp_index_, aliases_);
- Stmt rewritten = rewriter->Rewrite(alloc, call, body);
- // Continue through the remaining allocation bindings in the same scope.
- new_seq.push_back(Mutate(rewritten, inplace_mode).ValueOrUnchanged(rewritten));
- changed = true;
- break;
- } else {
- auto result = this->Mutate(op->seq[i]);
- changed |= !result.UnchangedOrSameAs(op->seq[i]);
- new_seq.push_back(std::move(result).ValueOrUnchanged(op->seq[i]).as_or_throw<Stmt>());
- }
- }
- if (!changed) return ffi::Unchanged();
- return SeqStmt(new_seq, op->loc);
- }
-
- int warp_size_{0};
- sym::Analyzer analyzer_;
- WarpThreadBindings active_bindings_;
- Var warp_index_{ffi::UnsafeInit{}};
- WarpIndexAliases aliases_;
-};
-
-namespace transform {
-
-Pass LowerWarpMemory() {
- auto pass_func = [](Function f, IRModule m, PassContext ctx) {
- if (!f->body.has_value()) return f;
- auto* n = f.CopyOnWrite();
- auto target = f->GetAttr<Target>(tvm::attr::kTarget);
- TVM_FFI_ICHECK(target.has_value()) << "LowerWarpMemory: Require the target attribute";
- int warp_size = target.value()->GetAttr<int64_t>("thread_warp_size", 1).value();
- auto warp_memory_rewriter = ffi::make_object<WarpMemoryRewriter>(warp_size);
- auto stmt = warp_memory_rewriter->Rewrite(std::move(n->body).value());
- n->body = ffi::make_object<UpdatePointerStorageScope>(warp_memory_rewriter->new_storage_scopes_)
- ->Mutate(stmt, InplaceMode::kAllow)
- .ValueOrUnchanged(stmt);
- return f;
- };
- return CreateFunctionPass(pass_func, 0, "tirx.LowerWarpMemory");
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def("tirx.transform.LowerWarpMemory", LowerWarpMemory);
-}
-
-} // namespace transform
-
-} // namespace tirx
-} // namespace tvm
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_mma_tensorize.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_mma_tensorize.py
deleted file mode 100644
index 75201cf84b..0000000000
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_mma_tensorize.py
+++ /dev/null
@@ -1,343 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-# ruff: noqa: E501, F841
-
-import numpy as np
-import pytest
-
-import tvm
-import tvm.s_tir.tensor_intrin # pylint: disable=unused-import
-import tvm.testing
-from tvm.s_tir.schedule import Schedule
-from tvm.script import s_tir as Ts
-from tvm.script import tirx as T
-from tvm.testing import env
-
-torch = pytest.importorskip("torch")
-
-M, N, K = 4096, 4096, 4096
-np.random.seed(0)
-
-
-@tvm.script.ir_module
-class Gemm_F16F16F16:
- # fmt: off
- @Ts.function
- def main(
- A: T.Tensor((M, K), "float16"), # type: ignore
- B: T.Tensor((K, N), "float16"), # type: ignore
- C: T.Tensor((M, N), "float16"), # type: ignore
- ):
- for i, j, k in T.grid(M, N, K):
- with Ts.sblock("C"):
- vi, vj, vk = Ts.axis.remap("SSR", [i, j, k])
- with Ts.init():
- C[vi, vj] = T.float32(0)
- C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]
-
-
-@tvm.script.ir_module
-class Gemm_F16F16F32:
- # fmt: off
- @Ts.function
- def main(
- A: T.Tensor((M, K), "float16"), # type: ignore
- B: T.Tensor((K, N), "float16"), # type: ignore
- C: T.Tensor((M, N), "float32"), # type: ignore
- ):
- for i, j, k in T.grid(M, N, K):
- with Ts.sblock("C"):
- vi, vj, vk = Ts.axis.remap("SSR", [i, j, k])
- with Ts.init():
- C[vi, vj] = T.float32(0)
- C[vi, vj] = C[vi, vj] + T.cast(A[vi, vk], "float32") * T.cast(B[vk, vj], "float32")
-
-
-@pytest.mark.gpu
-@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
-def test_run_target(mod=None, tgt_str=None, in_dtype="float16", out_dtype="float16"):
- if mod is None:
- return
- tgt_str = tgt_str or "cuda"
- target = tvm.target.Target(target=tgt_str)
- with tvm.ir.transform.PassContext(opt_level=3):
- lib: tvm.runtime.Module = tvm.compile(mod, target=target)
-
- a_np = np.random.rand(M, K).astype(in_dtype)
- b_np = np.random.rand(K, N).astype(in_dtype)
- c_np = np.ones((M, N), dtype=out_dtype)
- f = lib["main"]
-
- def run_and_check():
- dev = tvm.cuda(0)
- a = tvm.runtime.tensor(a_np, dev)
- b = tvm.runtime.tensor(b_np, dev)
- c = tvm.runtime.tensor(c_np, dev)
- f(a, b, c)
- c_th = torch.matmul(torch.tensor(a_np).to(tgt_str), torch.tensor(b_np).to(tgt_str)).to(
- torch.float32 if out_dtype == "float32" else torch.float16
- )
- c_f = torch.tensor(c.numpy()).to(tgt_str)
- assert torch.allclose(c_th, c_f, rtol=0.05, atol=0.05)
-
- tvm.testing.run_with_gpu_lock(run_and_check)
-
-
-@pytest.mark.gpu
-@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
-def test_f16f16f16_mma_gemm():
- # fmt: off
- mod = Gemm_F16F16F16
- sch = Schedule(mod)
- b0 = sch.get_sblock(name="C", func_name="main")
- b1 = sch.get_sblock(name="root", func_name="main")
- sch.annotate(block_or_loop=b0, ann_key="meta_schedule.tiling_structure", ann_val="SSSRRSRS")
- b2 = sch.reindex(block=b0, tensor=("write", 0))
- b3 = sch.reindex(block=b0, tensor=("read", 0))
- b4 = sch.reindex(block=b0, tensor=("read", 1))
- sch.transform_layout(block=b0, tensor=("read", 0), index_map=lambda vi, vk: (vi, vk,), pad_value=None, assume_injective_transform=True)
- sch.transform_layout(block=b0, tensor=("read", 1), index_map=lambda vj, vk: (vk, vj,), pad_value=None, assume_injective_transform=True)
- sch.transform_layout(block=b0, tensor=("write", 0), index_map=lambda vi, vj: (vi, vj,), pad_value=None, assume_injective_transform=True)
- sch.transform_block_layout(block=b2, index_map=lambda vi, vj: (vi, vj,))
- sch.transform_block_layout(block=b3, index_map=lambda vi, vk: (vi, vk,))
- sch.transform_block_layout(block=b4, index_map=lambda vj, vk: (vk, vj,))
- sch.transform_block_layout(block=b0, index_map=lambda vi, vj, vk: (vi, vj, vk,))
- l5, l6, l7 = sch.get_loops(block=b0)
- l8, l9 = sch.split(loop=l7, factors=[None, 8], preserve_unit_iters=True, disable_predication=False)
- l10, l11 = sch.split(loop=l6, factors=[None, 8], preserve_unit_iters=True, disable_predication=False)
- l12, l13 = sch.split(loop=l5, factors=[None, 16], preserve_unit_iters=True, disable_predication=False)
- l14, l15, l16, l17, l18, l19 = sch.get_loops(block=b0)
- sch.reorder(l16, l18, l13, l11, l9)
- b20 = sch.blockize(target=l13, preserve_unit_iters=True)
- sch.annotate(block_or_loop=b20, ann_key="meta_schedule.auto_tensorize", ann_val="mma_sync_m16n8k8_f16f16f16")
- sch.annotate(block_or_loop=b20, ann_key="meta_schedule.auto_tensorize_init", ann_val="mma_init_m16n8k8_f16")
- sch.annotate(block_or_loop=b20, ann_key="warp_execution", ann_val=1)
- l21, l22, l23 = sch.get_loops(block=b20)
- v24, v25, v26, v27, v28 = sch.sample_partitioned_tile(loop=l21, n=5, partition_pos=3, innerpart_factor=2, decision=[2, 16, 4, 1, 2])
- l29, l30, l31, l32, l33 = sch.split(loop=l21, factors=[v24, v25, v26, v27, v28], preserve_unit_iters=True, disable_predication=False)
- v34, v35, v36, v37, v38 = sch.sample_partitioned_tile(loop=l22, n=5, partition_pos=3, innerpart_factor=4, decision=[2, 16, 4, 1, 4])
- l39, l40, l41, l42, l43 = sch.split(loop=l22, factors=[v34, v35, v36, v37, v38], preserve_unit_iters=True, disable_predication=False)
- v44, v45, v46 = sch.sample_perfect_tile(loop=l23, n=3, max_innermost_factor=4, decision=[128, 1, 4])
- l47, l48, l49 = sch.split(loop=l23, factors=[v44, v45, v46], preserve_unit_iters=True, disable_predication=False)
- sch.reorder(l29, l39, l30, l40, l31, l41, l47, l48, l32, l42, l49, l33, l43)
- l50 = sch.fuse(l29, l39, preserve_unit_iters=True)
- sch.bind(loop=l50, thread_axis="blockIdx.y")
- l51 = sch.fuse(l30, l40, preserve_unit_iters=True)
- sch.bind(loop=l51, thread_axis="blockIdx.x")
- l52 = sch.fuse(l31, l41, preserve_unit_iters=True)
- sch.bind(loop=l52, thread_axis="threadIdx.y")
- sch.annotate(block_or_loop=b20, ann_key="meta_schedule.thread_extent_low_inclusive", ann_val=32)
- sch.annotate(block_or_loop=b20, ann_key="meta_schedule.thread_extent_high_inclusive", ann_val=1024)
- b53 = sch.write_at(loop=l52, block=b20, write_tensor_index=0, storage_scope="mma.matrixC")
- sch.reverse_compute_inline(block=b2)
- b54 = sch.read_at(loop=l47, block=b20, read_tensor_index=0, storage_scope="shared.dyn")
- sch.annotate(block_or_loop=b54, ann_key="permuted_layout", ann_val="g2s_A")
- b55 = sch.read_at(loop=l47, block=b20, read_tensor_index=1, storage_scope="shared.dyn")
- sch.annotate(block_or_loop=b55, ann_key="permuted_layout", ann_val="g2s_B")
- b56 = sch.cache_read(block=b20, read_tensor_index=0, storage_scope="mma.matrixA")
- sch.compute_at(block=b56, loop=l48, preserve_unit_loops=True, index=-1)
- l57, l58, l59, l60, l61, l62, l63 = sch.get_loops(block=b56)
- l64, l65 = sch.split(loop=l63, factors=[None, 8], preserve_unit_iters=True, disable_predication=False)
- l66, l67 = sch.split(loop=l62, factors=[None, 32], preserve_unit_iters=True, disable_predication=False)
- l68, l69, l70, l71, l72, l73, l74, l75, l76 = sch.get_loops(block=b56)
- sch.reorder(l75, l67, l65)
- b77 = sch.blockize(target=l67, preserve_unit_iters=True)
- sch.annotate(block_or_loop=b77, ann_key="meta_schedule.auto_tensorize", ann_val="mma_load_m16n8k8_f16_A_shared_dyn")
- sch.annotate(block_or_loop=b77, ann_key="permuted_layout", ann_val="s2l_A")
- b78 = sch.cache_read(block=b20, read_tensor_index=1, storage_scope="mma.matrixB")
- sch.compute_at(block=b78, loop=l48, preserve_unit_loops=True, index=-1)
- l79, l80, l81, l82, l83, l84, l85 = sch.get_loops(block=b78)
- l86, l87 = sch.split(loop=l85, factors=[None, 32], preserve_unit_iters=True, disable_predication=False)
- l88, l89 = sch.split(loop=l84, factors=[None, 8], preserve_unit_iters=True, disable_predication=False)
- l90, l91, l92, l93, l94, l95, l96, l97, l98 = sch.get_loops(block=b78)
- sch.reorder(l97, l89, l87)
- b99 = sch.blockize(target=l89, preserve_unit_iters=True)
- sch.annotate(block_or_loop=b99, ann_key="meta_schedule.auto_tensorize", ann_val="mma_load_m16n8k8_f16_B_shared_dyn")
- sch.annotate(block_or_loop=b99, ann_key="permuted_layout", ann_val="s2l_B")
- b100, = sch.get_producers(block=b54)
- sch.compute_inline(block=b100)
- sch.storage_align(block=b54, tensor_index=0, axis=-2, factor=32, offset=8)
- b101, = sch.get_producers(block=b55)
- sch.compute_inline(block=b101)
- sch.storage_align(block=b55, tensor_index=0, axis=-2, factor=32, offset=8)
- sch.annotate(block_or_loop=b54, ann_key="vector_bytes", ann_val=16)
- sch.annotate(block_or_loop=b55, ann_key="vector_bytes", ann_val=16)
- sch.annotate(block_or_loop=l48, ann_key="software_pipeline_stage", ann_val=[0, 0, 1])
- sch.annotate(block_or_loop=l48, ann_key="software_pipeline_order", ann_val=[0, 1, 2])
- sch.annotate(block_or_loop=l47, ann_key="software_pipeline_async_stages", ann_val=[0])
- sch.annotate(block_or_loop=l47, ann_key="software_pipeline_stage", ann_val=[0, 0, 1, 2, 2])
- sch.annotate(block_or_loop=l47, ann_key="software_pipeline_order", ann_val=[0, 1, 3, 2, 4])
- v102 = sch.sample_categorical(candidates=[0, 16, 64, 512, 1024], probs=[0.20000000000000001, 0.20000000000000001, 0.20000000000000001, 0.20000000000000001, 0.20000000000000001], decision=0)
- sch.annotate(block_or_loop=b1, ann_key="meta_schedule.unroll_explicit", ann_val=v102)
- sch.enter_postproc()
- b103 = sch.get_sblock(name="root", func_name="main")
- sch.unannotate(block_or_loop=b103, ann_key="meta_schedule.unroll_explicit")
- b104, b105, b106, b107, b108, b109 = sch.get_child_blocks(b103)
- l110, l111, l112, l113 = sch.get_loops(block=b104)
- l114, l115, l116, l117 = sch.get_loops(block=b105)
- l118, l119, l120, l121, l122, l123, l124 = sch.get_loops(block=b106)
- l125, l126, l127, l128, l129, l130, l131 = sch.get_loops(block=b107)
- l132, l133, l134, l135, l136, l137, l138, l139, l140, l141 = sch.get_loops(block=b108)
- l142, l143, l144 = sch.get_loops(block=b109)
- b145 = sch.get_sblock(name="C_o", func_name="main")
- l146, l147, l148, l149, l150, l151, l152, l153, l154, l155 = sch.get_loops(block=b145)
- b156 = sch.decompose_reduction(block=b145, loop=l149)
- sch.unannotate(block_or_loop=b156, ann_key="meta_schedule.auto_tensorize")
- sch.annotate(block_or_loop=b156, ann_key="meta_schedule.auto_tensorize", ann_val="mma_init_m16n8k8_f16")
- sch.unannotate(block_or_loop=b145, ann_key="meta_schedule.auto_tensorize_init")
- sch.unannotate(block_or_loop=b156, ann_key="meta_schedule.auto_tensorize_init")
- b157 = sch.get_sblock(name="C_o_init", func_name="main")
- sch.unannotate(block_or_loop=b157, ann_key="meta_schedule.auto_tensorize")
- sch.tensorize(block_or_loop=b157, tensor_intrin="mma_init_m16n8k8_f16", preserve_unit_iters=True)
- b158 = sch.get_sblock(name="A_reindex_shared.dyn_mma.matrixA_o", func_name="main")
- sch.unannotate(block_or_loop=b158, ann_key="meta_schedule.auto_tensorize")
- sch.tensorize(block_or_loop=b158, tensor_intrin="mma_load_m16n8k8_f16_A_shared_dyn", preserve_unit_iters=True)
- b159 = sch.get_sblock(name="B_reindex_shared.dyn_mma.matrixB_o", func_name="main")
- sch.unannotate(block_or_loop=b159, ann_key="meta_schedule.auto_tensorize")
- sch.tensorize(block_or_loop=b159, tensor_intrin="mma_load_m16n8k8_f16_B_shared_dyn", preserve_unit_iters=True)
- b160 = sch.get_sblock(name="C_o_update", func_name="main")
- sch.unannotate(block_or_loop=b160, ann_key="meta_schedule.auto_tensorize")
- sch.tensorize(block_or_loop=b160, tensor_intrin="mma_sync_m16n8k8_f16f16f16", preserve_unit_iters=True)
- mod = sch.mod
- test_run_target(mod)
-
-
-@pytest.mark.gpu
-@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
-def test_f16f16f32_mma_gemm():
- mod = Gemm_F16F16F32
- sch = Schedule(mod)
- # fmt: off
- sch = Schedule(mod)
- b0 = sch.get_sblock(name="C", func_name="main")
- b1 = sch.get_sblock(name="root", func_name="main")
- sch.annotate(block_or_loop=b0, ann_key="meta_schedule.tiling_structure", ann_val="SSSRRSRS")
- b2 = sch.reindex(block=b0, tensor=("write", 0))
- b3 = sch.reindex(block=b0, tensor=("read", 0))
- b4 = sch.reindex(block=b0, tensor=("read", 1))
- sch.transform_layout(block=b0, tensor=("read", 0), index_map=lambda vi, vk: (vi, vk,), pad_value=None, assume_injective_transform=True)
- sch.transform_layout(block=b0, tensor=("read", 1), index_map=lambda vj, vk: (vk, vj,), pad_value=None, assume_injective_transform=True)
- sch.transform_layout(block=b0, tensor=("write", 0), index_map=lambda vi, vj: (vi, vj,), pad_value=None, assume_injective_transform=True)
- sch.transform_block_layout(block=b2, index_map=lambda vi, vj: (vi, vj,))
- sch.transform_block_layout(block=b3, index_map=lambda vi, vk: (vi, vk,))
- sch.transform_block_layout(block=b4, index_map=lambda vj, vk: (vk, vj,))
- sch.transform_block_layout(block=b0, index_map=lambda vi, vj, vk: (vi, vj, vk,))
- l5, l6, l7 = sch.get_loops(block=b0)
- l8, l9 = sch.split(loop=l7, factors=[None, 8], preserve_unit_iters=True, disable_predication=False)
- l10, l11 = sch.split(loop=l6, factors=[None, 8], preserve_unit_iters=True, disable_predication=False)
- l12, l13 = sch.split(loop=l5, factors=[None, 16], preserve_unit_iters=True, disable_predication=False)
- l14, l15, l16, l17, l18, l19 = sch.get_loops(block=b0)
- sch.reorder(l16, l18, l13, l11, l9)
- b20 = sch.blockize(target=l13, preserve_unit_iters=True)
- sch.annotate(block_or_loop=b20, ann_key="meta_schedule.auto_tensorize", ann_val="mma_sync_m16n8k8_f16f16f32")
- sch.annotate(block_or_loop=b20, ann_key="meta_schedule.auto_tensorize_init", ann_val="mma_init_m16n8k8_f32")
- sch.annotate(block_or_loop=b20, ann_key="warp_execution", ann_val=1)
- l21, l22, l23 = sch.get_loops(block=b20)
- v24, v25, v26, v27, v28 = sch.sample_partitioned_tile(loop=l21, n=5, partition_pos=3, innerpart_factor=2, decision=[1, 16, 2, 2, 4])
- l29, l30, l31, l32, l33 = sch.split(loop=l21, factors=[v24, v25, v26, v27, v28], preserve_unit_iters=True, disable_predication=False)
- v34, v35, v36, v37, v38 = sch.sample_partitioned_tile(loop=l22, n=5, partition_pos=3, innerpart_factor=4, decision=[2, 16, 2, 4, 2])
- l39, l40, l41, l42, l43 = sch.split(loop=l22, factors=[v34, v35, v36, v37, v38], preserve_unit_iters=True, disable_predication=False)
- v44, v45, v46 = sch.sample_perfect_tile(loop=l23, n=3, max_innermost_factor=4, decision=[128, 1, 4])
- l47, l48, l49 = sch.split(loop=l23, factors=[v44, v45, v46], preserve_unit_iters=True, disable_predication=False)
- sch.reorder(l29, l39, l30, l40, l31, l41, l47, l48, l32, l42, l49, l33, l43)
- l50 = sch.fuse(l29, l39, preserve_unit_iters=True)
- sch.bind(loop=l50, thread_axis="blockIdx.y")
- l51 = sch.fuse(l30, l40, preserve_unit_iters=True)
- sch.bind(loop=l51, thread_axis="blockIdx.x")
- l52 = sch.fuse(l31, l41, preserve_unit_iters=True)
- sch.bind(loop=l52, thread_axis="threadIdx.y")
- sch.annotate(block_or_loop=b20, ann_key="meta_schedule.thread_extent_low_inclusive", ann_val=32)
- sch.annotate(block_or_loop=b20, ann_key="meta_schedule.thread_extent_high_inclusive", ann_val=1024)
- b53 = sch.write_at(loop=l52, block=b20, write_tensor_index=0, storage_scope="mma.matrixC")
- sch.reverse_compute_inline(block=b2)
- b54 = sch.read_at(loop=l47, block=b20, read_tensor_index=0, storage_scope="shared.dyn")
- sch.annotate(block_or_loop=b54, ann_key="permuted_layout", ann_val="g2s_A")
- b55 = sch.read_at(loop=l47, block=b20, read_tensor_index=1, storage_scope="shared.dyn")
- sch.annotate(block_or_loop=b55, ann_key="permuted_layout", ann_val="g2s_B")
- b56 = sch.cache_read(block=b20, read_tensor_index=0, storage_scope="mma.matrixA")
- sch.compute_at(block=b56, loop=l48, preserve_unit_loops=True, index=-1)
- l57, l58, l59, l60, l61, l62, l63 = sch.get_loops(block=b56)
- l64, l65 = sch.split(loop=l63, factors=[None, 8], preserve_unit_iters=True, disable_predication=False)
- l66, l67 = sch.split(loop=l62, factors=[None, 32], preserve_unit_iters=True, disable_predication=False)
- l68, l69, l70, l71, l72, l73, l74, l75, l76 = sch.get_loops(block=b56)
- sch.reorder(l75, l67, l65)
- b77 = sch.blockize(target=l67, preserve_unit_iters=True)
- sch.annotate(block_or_loop=b77, ann_key="meta_schedule.auto_tensorize", ann_val="mma_load_m16n8k8_f16_A_shared_dyn")
- sch.annotate(block_or_loop=b77, ann_key="permuted_layout", ann_val="s2l_A")
- b78 = sch.cache_read(block=b20, read_tensor_index=1, storage_scope="mma.matrixB")
- sch.compute_at(block=b78, loop=l48, preserve_unit_loops=True, index=-1)
- l79, l80, l81, l82, l83, l84, l85 = sch.get_loops(block=b78)
- l86, l87 = sch.split(loop=l85, factors=[None, 32], preserve_unit_iters=True, disable_predication=False)
- l88, l89 = sch.split(loop=l84, factors=[None, 8], preserve_unit_iters=True, disable_predication=False)
- l90, l91, l92, l93, l94, l95, l96, l97, l98 = sch.get_loops(block=b78)
- sch.reorder(l97, l89, l87)
- b99 = sch.blockize(target=l89, preserve_unit_iters=True)
- sch.annotate(block_or_loop=b99, ann_key="meta_schedule.auto_tensorize", ann_val="mma_load_m16n8k8_f16_B_shared_dyn")
- sch.annotate(block_or_loop=b99, ann_key="permuted_layout", ann_val="s2l_B")
- b100, = sch.get_producers(block=b54)
- sch.compute_inline(block=b100)
- sch.storage_align(block=b54, tensor_index=0, axis=-2, factor=32, offset=8)
- b101, = sch.get_producers(block=b55)
- sch.compute_inline(block=b101)
- sch.storage_align(block=b55, tensor_index=0, axis=-2, factor=32, offset=8)
- sch.annotate(block_or_loop=b54, ann_key="vector_bytes", ann_val=16)
- sch.annotate(block_or_loop=b55, ann_key="vector_bytes", ann_val=16)
- sch.annotate(block_or_loop=l48, ann_key="software_pipeline_stage", ann_val=[0, 0, 1])
- sch.annotate(block_or_loop=l48, ann_key="software_pipeline_order", ann_val=[0, 1, 2])
- sch.annotate(block_or_loop=l47, ann_key="software_pipeline_async_stages", ann_val=[0])
- sch.annotate(block_or_loop=l47, ann_key="software_pipeline_stage", ann_val=[0, 0, 1, 2, 2])
- sch.annotate(block_or_loop=l47, ann_key="software_pipeline_order", ann_val=[0, 1, 3, 2, 4])
- v102 = sch.sample_categorical(candidates=[0, 16, 64, 512, 1024], probs=[0.20000000000000001, 0.20000000000000001, 0.20000000000000001, 0.20000000000000001, 0.20000000000000001], decision=0)
- sch.annotate(block_or_loop=b1, ann_key="meta_schedule.unroll_explicit", ann_val=v102)
- sch.enter_postproc()
- b103 = sch.get_sblock(name="root", func_name="main")
- sch.unannotate(block_or_loop=b103, ann_key="meta_schedule.unroll_explicit")
- b104, b105, b106, b107, b108, b109 = sch.get_child_blocks(b103)
- l110, l111, l112, l113 = sch.get_loops(block=b104)
- l114, l115, l116, l117 = sch.get_loops(block=b105)
- l118, l119, l120, l121, l122, l123, l124 = sch.get_loops(block=b106)
- l125, l126, l127, l128, l129, l130, l131 = sch.get_loops(block=b107)
- l132, l133, l134, l135, l136, l137, l138, l139, l140, l141 = sch.get_loops(block=b108)
- sch.annotate(block_or_loop=l132, ann_key="auto_unroll_max_step", ann_val=0)
- sch.annotate(block_or_loop=l132, ann_key="unroll_explicit", ann_val=1)
- l142, l143, l144 = sch.get_loops(block=b109)
- b145 = sch.get_sblock(name="C_o", func_name="main")
- l146, l147, l148, l149, l150, l151, l152, l153, l154, l155 = sch.get_loops(block=b145)
- b156 = sch.decompose_reduction(block=b145, loop=l149)
- sch.unannotate(block_or_loop=b156, ann_key="meta_schedule.auto_tensorize")
- sch.annotate(block_or_loop=b156, ann_key="meta_schedule.auto_tensorize", ann_val="mma_init_m16n8k8_f32")
- sch.unannotate(block_or_loop=b145, ann_key="meta_schedule.auto_tensorize_init")
- sch.unannotate(block_or_loop=b156, ann_key="meta_schedule.auto_tensorize_init")
- b157 = sch.get_sblock(name="C_o_init", func_name="main")
- sch.unannotate(block_or_loop=b157, ann_key="meta_schedule.auto_tensorize")
- sch.tensorize(block_or_loop=b157, tensor_intrin="mma_init_m16n8k8_f32", preserve_unit_iters=True)
- b158 = sch.get_sblock(name="A_reindex_shared.dyn_mma.matrixA_o", func_name="main")
- sch.unannotate(block_or_loop=b158, ann_key="meta_schedule.auto_tensorize")
- sch.tensorize(block_or_loop=b158, tensor_intrin="mma_load_m16n8k8_f16_A_shared_dyn", preserve_unit_iters=True)
- b159 = sch.get_sblock(name="B_reindex_shared.dyn_mma.matrixB_o", func_name="main")
- sch.unannotate(block_or_loop=b159, ann_key="meta_schedule.auto_tensorize")
- sch.tensorize(block_or_loop=b159, tensor_intrin="mma_load_m16n8k8_f16_B_shared_dyn", preserve_unit_iters=True)
- b160 = sch.get_sblock(name="C_o_update", func_name="main")
- sch.unannotate(block_or_loop=b160, ann_key="meta_schedule.auto_tensorize")
- sch.tensorize(block_or_loop=b160, tensor_intrin="mma_sync_m16n8k8_f16f16f32", preserve_unit_iters=True)
- mod = sch.mod
- test_run_target(mod, out_dtype="float32")
-
-
-if __name__ == """__main__""":
- test_f16f16f16_mma_gemm()
- test_f16f16f32_mma_gemm()
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_tensorize.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_tensorize.py
index 809fcf9750..18c799fca0 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_tensorize.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_postproc_rewrite_tensorize.py
@@ -18,7 +18,7 @@
# ruff: noqa: F401, F841
import tvm
from tvm.s_tir import meta_schedule as ms
-from tvm.s_tir.tensor_intrin import cuda, rocm, x86
+from tvm.s_tir.tensor_intrin import dot_product_common, x86
from tvm.script import s_tir as Ts
from tvm.script import tirx as T
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_space_cuda_async_multiple_initialization.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_space_cuda_async_multiple_initialization.py
index dc16994657..008c98d575 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_space_cuda_async_multiple_initialization.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_space_cuda_async_multiple_initialization.py
@@ -29,7 +29,7 @@ from tvm.target import Target
def generate_design_space(
- kind: Literal["llvm", "cuda", "cuda-tensorcore", "hexagon"],
+ kind: Literal["llvm", "cuda", "hexagon"],
mod: IRModule,
target: Target,
types: type | tuple[type, ...],
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_space_generator.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_space_generator.py
index 3952ef4440..9e5448bc6d 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_space_generator.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_space_generator.py
@@ -24,6 +24,7 @@ import pytest
import tvm
import tvm.testing
from tvm.s_tir.meta_schedule.space_generator import (
+ PostOrderApply,
PySpaceGenerator,
ScheduleFn,
SpaceGeneratorUnion,
@@ -109,5 +110,25 @@ def test_meta_schedule_design_space_generator_NIE():
generator._initialize_with_tune_context(TuneContext())
+@pytest.mark.parametrize("arch", ["sm_75", "sm_80", "sm_100"])
+def test_cuda_default_rules_use_generic_scheduling(arch):
+ target = tvm.target.Target(
+ {
+ "kind": "cuda",
+ "arch": arch,
+ "max_threads_per_block": 1024,
+ "max_shared_memory_per_block": 49152,
+ "thread_warp_size": 32,
+ }
+ )
+ generator = PostOrderApply()
+ TuneContext(mod=Matmul, target=target, space_generator=generator, rand_state=0)
+ assert any(
+ isinstance(rule, tvm.s_tir.meta_schedule.schedule_rule.MultiLevelTiling)
+ for rule in generator.sch_rules
+ )
+ assert generator.generate_design_space(Matmul)
+
+
if __name__ == "__main__":
tvm.testing.main()
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_trace_apply.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_trace_apply.py
index 244c3f08d6..c77b0d0520 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_trace_apply.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_trace_apply.py
@@ -14,7 +14,7 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-# ruff: noqa: E501, F401, F403, F841
+# ruff: noqa: E501, F401, F841
import pytest
import tvm
@@ -22,7 +22,6 @@ import tvm.s_tir.meta_schedule as ms
import tvm.testing
from tvm.codegen import llvm_lookup_intrinsic_id
from tvm.s_tir import Schedule
-from tvm.s_tir.tensor_intrin.cuda import *
from tvm.s_tir.tensor_intrin.x86 import VNNI_DOT_16x4_INTRIN as VNNI_INTRIN
from tvm.script import s_tir as Ts
from tvm.script import tirx as T
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_analysis.py b/tests/python/s_tir/schedule/test_tir_schedule_analysis.py
index 57851cecd3..22288f00aa 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_analysis.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_analysis.py
@@ -298,6 +298,24 @@ def test_get_tensorize_loop_mapping_matmul_mma():
assert s.get(desc_loop_to_sref[desc_loops[2]]) == s.get(i2)
+def matmul_16x8x8_desc(out_dtype):
+ @Ts.function
+ def desc(
+ A: T.Tensor((16, 8), "float16", offset_factor=1, elem_offset=T.int32()),
+ B: T.Tensor((8, 8), "float16", offset_factor=1, elem_offset=T.int32()),
+ C: T.Tensor((16, 8), out_dtype, offset_factor=1, elem_offset=T.int32()),
+ ):
+ with Ts.sblock("root"):
+ Ts.reads(C[0:16, 0:8], A[0:16, 0:8], B[0:8, 0:8])
+ Ts.writes(C[0:16, 0:8])
+ for i, j, k in T.grid(16, 8, 8):
+ with Ts.sblock("update"):
+ vi, vj, vk = Ts.axis.remap("SSR", [i, j, k])
+ C[vi, vj] += T.cast(A[vi, vk], out_dtype) * T.cast(B[vk, vj], out_dtype)
+
+ return desc
+
+
def test_get_tensorize_loop_mapping_padding_matmul():
matmul = create_function(
te_workload.matmul_relu(
@@ -311,7 +329,7 @@ def test_get_tensorize_loop_mapping_padding_matmul():
s = Schedule(matmul)
block = s.get_sblock("C")
- desc = TensorIntrin.get("mma_sync_m16n8k8_f16f16f16").desc
+ desc = matmul_16x8x8_desc("float16")
info = get_tensorize_loop_mapping(s, block, desc, allow_padding=True)
assert info is not None
expected_padding = [16, 1, 8]
@@ -322,10 +340,9 @@ def test_get_tensorize_loop_mapping_padding_matmul():
assert actual == expected
-def check_index_map(workload, block_name, intrin_name, expected_index_map):
+def check_index_map(workload, block_name, desc_func, expected_index_map):
s = Schedule(workload)
block = s.get_sblock(block_name)
- desc_func = TensorIntrin.get(intrin_name).desc
info = get_auto_tensorize_mapping_info(s, block, desc_func)
if expected_index_map is None:
assert info is None
@@ -341,7 +358,7 @@ def test_get_auto_tensorize_mapping_info_conv2d():
check_index_map(
conv2d,
"conv2d_nhwc",
- "mma_sync_m16n8k8_f16f16f32",
+ matmul_16x8x8_desc("float32"),
lambda n, h, w, c, rh, rw, rc: (n * 256 + h * 16 + w, c, rh * 192 + rw * 64 + rc),
)
@@ -353,7 +370,7 @@ def test_get_auto_tensorize_mapping_info_conv2d_unit_batch():
check_index_map(
conv2d,
"conv2d_nhwc",
- "mma_sync_m16n8k8_f16f16f32",
+ matmul_16x8x8_desc("float32"),
lambda n, h, w, c, rh, rw, rc: (n * 256 + h * 16 + w, c, rh * 192 + rw * 64 + rc),
)
@@ -363,7 +380,7 @@ def test_get_auto_tensorize_mapping_info_batch_matmul(b, m, n, k):
matmul = create_function(
te_workload.batch_matmul_nkkm(b, m, n, k, in_dtype="float16", out_dtype="float32")
)
- check_index_map(matmul, "Z", "mma_sync_m16n8k8_f16f16f32", lambda b, m, n, k: (b, m, n, k))
+ check_index_map(matmul, "Z", matmul_16x8x8_desc("float32"), lambda b, m, n, k: (b, m, n, k))
@pytest.mark.parametrize(
@@ -384,7 +401,7 @@ def test_get_auto_tensorize_mapping_info_batch_matmul(b, m, n, k):
)
def test_get_auto_tensorize_mapping_info_matmul(n, m, k, expected):
matmul = create_function(te_workload.matmul(n, m, k, in_dtype="float16", out_dtype="float32"))
- check_index_map(matmul, "C", "mma_sync_m16n8k8_f16f16f32", expected)
+ check_index_map(matmul, "C", matmul_16x8x8_desc("float32"), expected)
def test_is_output_block():
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_tensorize_ldmatrix_mma_numeric.py b/tests/python/s_tir/schedule/test_tir_schedule_tensorize_ldmatrix_mma_numeric.py
deleted file mode 100644
index 6e44ef9400..0000000000
--- a/tests/python/s_tir/schedule/test_tir_schedule_tensorize_ldmatrix_mma_numeric.py
+++ /dev/null
@@ -1,464 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-# pylint: disable=missing-docstring
-import numpy as np
-import pytest
-
-import tvm
-import tvm.testing
-from tvm import te
-from tvm.s_tir.tensor_intrin.cuda import (
- LDMATRIX_e4m3_A_INTRIN,
- LDMATRIX_e4m3_B_TRANS_INTRIN,
- LDMATRIX_e5m2_A_INTRIN,
- LDMATRIX_e5m2_B_TRANS_INTRIN,
- LDMATRIX_f16_A_INTRIN,
- LDMATRIX_f16_B_INTRIN,
- LDMATRIX_f16_B_TRANS_INTRIN,
- LDMATRIX_i8_A_INTRIN,
- LDMATRIX_i8_B_INTRIN,
- LDMATRIX_i8_B_TRANS_INTRIN,
- MMA_e4m3e4m3f32_TRANS_B_INTRIN,
- MMA_e5m2e5m2f32_TRANS_B_INTRIN,
- MMA_f16f16f16_INTRIN,
- MMA_f16f16f16_TRANS_B_INTRIN,
- MMA_f16f16f32_INTRIN,
- MMA_f16f16f32_TRANS_B_INTRIN,
- MMA_fill_16x16_f16_INTRIN,
- MMA_fill_16x16_f32_INTRIN,
- MMA_fill_16x16_i32_INTRIN,
- MMA_i8i8i32_INTRIN,
- MMA_i8i8i32_TRANS_B_INTRIN,
- MMA_store_16x16_f16_global_INTRIN,
- MMA_store_16x16_f32_global_INTRIN,
- MMA_store_16x16_i32_global_INTRIN,
- shared_16x16_to_ldmatrix_32x8_layout,
- shared_16x32_to_ldmatrix_32x16_layout,
- shared_32x16_to_ldmatrix_32x16_layout,
-)
-from tvm.s_tir.testing.tir import mma_schedule
-from tvm.testing import env
-
-M = 4096
-N = 4096
-K = 4096
-measure_perf = False
-gflops = (N * M * K) * 2 / 1e9
-
-
-def matmul(m, n, k, in_dtype, out_dtype, b_transposed):
- b_shape = (n, k) if b_transposed else (k, n)
- a = te.placeholder((m, k), name="A", dtype=in_dtype)
- b = te.placeholder(b_shape, name="B", dtype=in_dtype)
- k = te.reduce_axis((0, k), name="k")
-
- def maybe_cast(v):
- if in_dtype != out_dtype:
- return tvm.tirx.Cast(v, out_dtype)
- return v
-
- def maybe_swap(i, j):
- if b_transposed:
- return j, i
- return i, j
-
- c = te.compute(
- (m, n),
- lambda i, j: te.sum(maybe_cast(a[i, k]) * maybe_cast(b[maybe_swap(k, j)]), axis=[k]),
- name="C",
- )
- return (a, b, c)
-
-
-def run_test(
- k_inner,
- in_dtype,
- out_dtype,
- b_transposed,
- i_factors,
- j_factors,
- k_factors,
- index_map_A,
- index_map_B,
- index_map_C,
- ldmatrix_a_intrin,
- ldmatrix_b_intrin,
- mma_intrin,
- mma_fill_intrin,
- mma_store_intrin,
-):
- sch = mma_schedule(
- te.create_function(matmul(M, N, K, in_dtype, out_dtype, b_transposed)),
- k_inner,
- in_dtype,
- b_transposed,
- i_factors,
- j_factors,
- k_factors,
- index_map_A,
- index_map_B,
- index_map_C,
- ldmatrix_a_intrin,
- ldmatrix_b_intrin,
- mma_intrin,
- mma_fill_intrin,
- mma_store_intrin,
- )
-
- f = tvm.compile(sch.mod["main"], target="cuda")
-
- if in_dtype == "float16":
- a_np = np.random.normal(size=(M, K)).astype("float16")
-
- if b_transposed:
- b_np = np.random.normal(size=(N, K)).astype("float16")
- c_np = np.dot(a_np.astype("float32"), b_np.astype("float32").transpose()).astype(
- out_dtype
- )
- else:
- b_np = np.random.normal(size=(K, N)).astype("float16")
- c_np = np.dot(a_np.astype("float32"), b_np.astype("float32")).astype(out_dtype)
- elif in_dtype in ["float8_e4m3fn", "float8_e5m2"]:
- typemap = {
- "float8_e4m3fn": "float8_e4m3fn",
- "float8_e5m2": "float8_e5m2",
- }
- a_np = (
- np.random.uniform(low=-5, high=5, size=(M * K))
- .reshape((M, K))
- .astype(typemap[in_dtype])
- )
- if b_transposed:
- b_np = (
- np.random.uniform(low=-5, high=5, size=(N * K))
- .reshape((N, K))
- .astype(typemap[in_dtype])
- )
- c_np = np.dot(a_np.astype("float32"), b_np.T.astype("float32")).astype(out_dtype)
- else:
- b_np = (
- np.random.uniform(low=-5, high=5, size=(N * K))
- .reshape((K, N))
- .astype(typemap[in_dtype])
- )
- c_np = np.dot(a_np.astype("float32"), b_np.astype("float32")).astype(out_dtype)
- else:
- a_np = np.random.randint(-128, 128, (M, K)).astype("int8")
-
- if b_transposed:
- b_np = np.random.randint(-128, 128, (N, K)).astype("int8")
- c_np = np.dot(a_np.astype("float32"), b_np.astype("float32").transpose()).astype(
- "int32"
- )
- else:
- b_np = np.random.randint(-128, 128, (K, N)).astype("int8")
- c_np = np.dot(a_np.astype("float32"), b_np.astype("float32")).astype("int32")
-
- def run_and_check(measure=False):
- dev = tvm.cuda(0)
- a = tvm.runtime.tensor(a_np, dev)
- b = tvm.runtime.tensor(b_np, dev)
- c = tvm.runtime.tensor(np.zeros((M, N), dtype=out_dtype), dev)
- if measure:
- return f.time_evaluator(f.entry_name, dev, number=500)(a, b, c)
- f(a, b, c)
- dev.sync()
- if out_dtype != "float16" and in_dtype not in ["float8_e4m3fn", "float8_e5m2"]:
- tvm.testing.assert_allclose(c.numpy(), c_np, rtol=1e-2, atol=1e-2)
-
- tvm.testing.run_with_gpu_lock(run_and_check)
- return lambda: tvm.testing.run_with_gpu_lock(run_and_check, True)
-
-
-@pytest.mark.gpu
-@pytest.mark.skipif(not env.has_cuda_compute(8), reason="need cuda compute >= 8.0")
-def test_f16f16f32_m16n16k16():
- def index_map(i, j):
- return (
- i // 16,
- j // 16,
- *shared_16x16_to_ldmatrix_32x8_layout(i % 16, j % 16),
- )
-
- k_inner = 16
- in_dtype = "float16"
- out_dtype = "float32"
- i_factors, j_factors, k_factors = [4, 8, 2, 4, 1], [1, 64, 2, 1, 2], [128, 2, 1]
-
- timer = run_test(
- k_inner,
- in_dtype,
- out_dtype,
- False, # b_transposed
- i_factors,
- j_factors,
- k_factors,
- index_map,
- index_map,
- index_map,
- LDMATRIX_f16_A_INTRIN,
- LDMATRIX_f16_B_INTRIN,
- MMA_f16f16f32_INTRIN,
- MMA_fill_16x16_f32_INTRIN,
- MMA_store_16x16_f32_global_INTRIN,
- )
-
- if measure_perf and timer:
- print("f16f16f32_m16n16k16: %f GFLOPS" % (gflops / (timer().mean)))
-
- timer = run_test(
- k_inner,
- in_dtype,
- out_dtype,
- True, # b_transposed
- i_factors,
- j_factors,
- k_factors,
- index_map,
- index_map,
- index_map,
- LDMATRIX_f16_A_INTRIN,
- LDMATRIX_f16_B_TRANS_INTRIN,
- MMA_f16f16f32_TRANS_B_INTRIN,
- MMA_fill_16x16_f32_INTRIN,
- MMA_store_16x16_f32_global_INTRIN,
- )
-
- if measure_perf and timer:
- print("f16f16f32_m16n16k16_trans: %f GFLOPS" % (gflops / (timer().mean)))
-
-
-@pytest.mark.gpu
-@pytest.mark.skipif(not env.has_cuda_compute(8), reason="need cuda compute >= 8.0")
-def test_f16f16f16_m16n16k16():
- def index_map(i, j):
- return (
- i // 16,
- j // 16,
- *shared_16x16_to_ldmatrix_32x8_layout(i % 16, j % 16),
- )
-
- k_inner = 16
- in_dtype = "float16"
- out_dtype = "float16"
- i_factors, j_factors, k_factors = [16, 2, 1, 4, 2], [16, 2, 2, 1, 4], [128, 2, 1]
-
- timer = run_test(
- k_inner,
- in_dtype,
- out_dtype,
- False, # b_transposed
- i_factors,
- j_factors,
- k_factors,
- index_map,
- index_map,
- index_map,
- LDMATRIX_f16_A_INTRIN,
- LDMATRIX_f16_B_INTRIN,
- MMA_f16f16f16_INTRIN,
- MMA_fill_16x16_f16_INTRIN,
- MMA_store_16x16_f16_global_INTRIN,
- )
-
- if measure_perf and timer:
- print("f16f16f16_m16n16k16: %f GFLOPS" % (gflops / (timer().mean)))
-
- timer = run_test(
- k_inner,
- in_dtype,
- out_dtype,
- True, # b_transposed
- i_factors,
- j_factors,
- k_factors,
- index_map,
- index_map,
- index_map,
- LDMATRIX_f16_A_INTRIN,
- LDMATRIX_f16_B_TRANS_INTRIN,
- MMA_f16f16f16_TRANS_B_INTRIN,
- MMA_fill_16x16_f16_INTRIN,
- MMA_store_16x16_f16_global_INTRIN,
- )
-
- if measure_perf and timer:
- print("f16f16f16_m16n16k16_trans: %f GFLOPS" % (gflops / (timer().mean)))
-
-
-@pytest.mark.gpu
-@pytest.mark.skipif(not env.has_cuda_compute(8), reason="need cuda compute >= 8.0")
-def test_i8i8i32_m16n16k32():
- def index_map_A(i, j):
- return (
- i // 16,
- j // 32,
- *shared_16x32_to_ldmatrix_32x16_layout(i % 16, j % 32),
- )
-
- def index_map_B(i, j):
- return (
- i // 32,
- j // 16,
- *shared_32x16_to_ldmatrix_32x16_layout(i % 32, j % 16),
- )
-
- def index_map_C(i, j):
- return (
- i // 16,
- j // 16,
- *shared_16x16_to_ldmatrix_32x8_layout(i % 16, j % 16),
- )
-
- k_inner = 32
- in_dtype = "int8"
- out_dtype = "int32"
- i_factors, j_factors, k_factors = [1, 32, 1, 4, 2], [8, 4, 4, 2, 1], [32, 2, 2]
-
- timer = run_test(
- k_inner,
- in_dtype,
- out_dtype,
- False, # b_transposed
- i_factors,
- j_factors,
- k_factors,
- index_map_A,
- index_map_B,
- index_map_C,
- LDMATRIX_i8_A_INTRIN,
- LDMATRIX_i8_B_INTRIN,
- MMA_i8i8i32_INTRIN,
- MMA_fill_16x16_i32_INTRIN,
- MMA_store_16x16_i32_global_INTRIN,
- )
-
- if measure_perf and timer:
- print("i8i8i32_m16n16k32: %f GOPS" % (gflops / (timer().mean)))
-
- timer = run_test(
- k_inner,
- in_dtype,
- out_dtype,
- True, # b_transposed
- i_factors,
- j_factors,
- k_factors,
- index_map_A,
- index_map_A,
- index_map_C,
- LDMATRIX_i8_A_INTRIN,
- LDMATRIX_i8_B_TRANS_INTRIN,
- MMA_i8i8i32_TRANS_B_INTRIN,
- MMA_fill_16x16_i32_INTRIN,
- MMA_store_16x16_i32_global_INTRIN,
- )
-
- if measure_perf and timer:
- print("i8i8i32_m16n16k32_trans: %f GOPS" % (gflops / (timer().mean)))
-
-
-@pytest.mark.gpu
-@pytest.mark.skipif(not env.has_cuda_compute(8, 9), reason="need cuda compute >= 8.9")
-def test_e4m3e4m3f32_m16n16k32():
- def index_map_A(i, j):
- return (
- i // 16,
- j // 32,
- *shared_16x32_to_ldmatrix_32x16_layout(i % 16, j % 32),
- )
-
- def index_map_C(i, j):
- return (
- i // 16,
- j // 16,
- *shared_16x16_to_ldmatrix_32x8_layout(i % 16, j % 16),
- )
-
- k_inner = 32
- in_dtype = "float8_e4m3fn"
- out_dtype = "float32"
- i_factors, j_factors, k_factors = [1, 32, 1, 4, 2], [8, 4, 4, 2, 1], [32, 2, 2]
-
- timer = run_test(
- k_inner,
- in_dtype,
- out_dtype,
- True, # b_transposed
- i_factors,
- j_factors,
- k_factors,
- index_map_A,
- index_map_A,
- index_map_C,
- LDMATRIX_e4m3_A_INTRIN,
- LDMATRIX_e4m3_B_TRANS_INTRIN,
- MMA_e4m3e4m3f32_TRANS_B_INTRIN,
- MMA_fill_16x16_f32_INTRIN,
- MMA_store_16x16_f32_global_INTRIN,
- )
-
- if measure_perf and timer:
- print("e4m3e4m3f32_m16n16k32_trans: %f GOPS" % (gflops / (timer().mean)))
-
-
-@pytest.mark.gpu
-@pytest.mark.skipif(not env.has_cuda_compute(8, 9), reason="need cuda compute >= 8.9")
-def test_e5m2e5m2f32_m16n16k32():
- def index_map_A(i, j):
- return (
- i // 16,
- j // 32,
- *shared_16x32_to_ldmatrix_32x16_layout(i % 16, j % 32),
- )
-
- def index_map_C(i, j):
- return (
- i // 16,
- j // 16,
- *shared_16x16_to_ldmatrix_32x8_layout(i % 16, j % 16),
- )
-
- k_inner = 32
- in_dtype = "float8_e5m2"
- out_dtype = "float32"
- i_factors, j_factors, k_factors = [1, 32, 1, 4, 2], [8, 4, 4, 2, 1], [32, 2, 2]
-
- timer = run_test(
- k_inner,
- in_dtype,
- out_dtype,
- True, # b_transposed
- i_factors,
- j_factors,
- k_factors,
- index_map_A,
- index_map_A,
- index_map_C,
- LDMATRIX_e5m2_A_INTRIN,
- LDMATRIX_e5m2_B_TRANS_INTRIN,
- MMA_e5m2e5m2f32_TRANS_B_INTRIN,
- MMA_fill_16x16_f32_INTRIN,
- MMA_store_16x16_f32_global_INTRIN,
- )
-
- if measure_perf and timer:
- print("e5m2e5m2f32_m16n16k32_trans: %f GOPS" % (gflops / (timer().mean)))
-
-
-if __name__ == "__main__":
- tvm.testing.main()
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_tensorize_mfma_numeric.py b/tests/python/s_tir/schedule/test_tir_schedule_tensorize_mfma_numeric.py
index 043683f276..c1dc1a069d 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_tensorize_mfma_numeric.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_tensorize_mfma_numeric.py
@@ -17,6 +17,7 @@
# pylint: disable=missing-docstring
import numpy as np
import pytest
+import tvm_ffi
import tvm
import tvm.testing
@@ -155,8 +156,7 @@ def run_test(
return f.time_evaluator(f.entry_name, dev, number=500)(a, b, c)
f(a, b, c)
dev.sync()
- if in_dtype != "float16":
- tvm.testing.assert_allclose(c.numpy(), c_np, rtol=1e-2, atol=1e-2)
+ tvm.testing.assert_allclose(c.numpy(), c_np, rtol=1e-2, atol=1e-2)
tvm.testing.run_with_gpu_lock(run_and_check)
return lambda: tvm.testing.run_with_gpu_lock(run_and_check, True)
@@ -270,15 +270,15 @@ def test_f32f32f32_m16n16k4():
def index_map_A(i, j):
return (
i // 16,
- j // 16,
- *shared_16x4_to_local_64x1_layout_A(i % 16, j % 16),
+ j // 4,
+ *shared_16x4_to_local_64x1_layout_A(i % 16, j % 4),
)
def index_map_B(i, j):
return (
- i // 16,
+ i // 4,
j // 16,
- *shared_4x16_to_local_64x1_layout_B(i % 16, j % 16),
+ *shared_4x16_to_local_64x1_layout_B(i % 4, j % 16),
)
def index_map_C(i, j):
@@ -315,5 +315,101 @@ def test_f32f32f32_m16n16k4():
print("test_f32f32f32_m16n16k4: %f GFLOPS" % (gflops / (timer().mean)))
+@pytest.mark.parametrize(
+ "k_inner,in_dtype,out_dtype,load_a,load_b,mfma,fill,store",
+ [
+ (
+ 4,
+ "float32",
+ "float32",
+ ROCM_MFMA_LOAD_16x4_A_SHARED_f32_INTRIN,
+ ROCM_MFMA_LOAD_16x4_B_SHARED_f32_INTRIN,
+ ROCM_MFMA_f32f32f32_INTRIN,
+ ROCM_MFMA_fill_16x16_f32_INTRIN,
+ ROCM_MFMA_STORE_16x16_f32_INTRIN,
+ ),
+ (
+ 16,
+ "float16",
+ "float32",
+ ROCM_MFMA_LOAD_16x16_A_SHARED_f16_INTRIN,
+ ROCM_MFMA_LOAD_16x16_B_SHARED_f16_INTRIN,
+ ROCM_MFMA_f16f16f32_INTRIN,
+ ROCM_MFMA_fill_16x16_f32_INTRIN,
+ ROCM_MFMA_STORE_16x16_f32_INTRIN,
+ ),
+ (
+ 16,
+ "int8",
+ "int32",
+ ROCM_MFMA_LOAD_16x16_A_SHARED_s8_INTRIN,
+ ROCM_MFMA_LOAD_16x16_B_SHARED_s8_INTRIN,
+ ROCM_MFMA_s8s8s32_INTRIN,
+ ROCM_MFMA_fill_16x16_i32_INTRIN,
+ ROCM_MFMA_STORE_16x16_s32_INTRIN,
+ ),
+ ],
+)
+def test_mfma_per_lane_storage(k_inner, in_dtype, out_dtype, load_a, load_b, mfma, fill, store):
+ """Tensorize and lower multi-tile, multi-wave GEMM without ROCm hardware."""
+ from tvm.s_tir.pipeline import _lower_s_tir
+
+ layout_a = (
+ shared_16x4_to_local_64x1_layout_A if k_inner == 4 else shared_16x16_to_local_64x4_layout_A
+ )
+ layout_b = (
+ shared_4x16_to_local_64x1_layout_B if k_inner == 4 else shared_16x16_to_local_64x4_layout_B
+ )
+
+ def index_map_a(i, j):
+ return i // 16, j // k_inner, *layout_a(i % 16, j % k_inner)
+
+ def index_map_b(i, j):
+ return i // k_inner, j // 16, *layout_b(i % k_inner, j % 16)
+
+ def index_map_c(i, j):
+ return i // 16, j // 16, *shared_16x16_to_local_64x4_layout_C(i % 16, j % 16)
+
+ sch = mfma_schedule(
+ te.create_function(matmul(64, 64, 2 * k_inner, in_dtype, out_dtype, False)),
+ k_inner,
+ in_dtype,
+ False,
+ [2, 1, 1, 2, 1],
+ [1, 1, 2, 1, 2],
+ [1, 2, 1],
+ index_map_a,
+ index_map_b,
+ index_map_c,
+ load_a,
+ load_b,
+ mfma,
+ fill,
+ store,
+ )
+ lowered = _lower_s_tir()(sch.mod)["main"]
+ allocations = []
+ accesses = []
+
+ def check(node):
+ if isinstance(node, tvm.ir.Bind) and isinstance(node.var.ty, tvm.tirx.TensorType):
+ assert str(node.var.ty.scope) != "warp"
+ if str(node.var.ty.scope) == "local":
+ allocations.append(tuple(int(extent) for extent in node.var.ty.shape))
+ if isinstance(node, tvm.ir.TensorLoad | tvm.ir.TensorStore):
+ tensor = node.source if isinstance(node, tvm.ir.TensorLoad) else node.dest
+ if str(tensor.ty.scope) == "local":
+ accesses.append(node)
+ # All load/fill/MFMA/store operations address this lane's fragment.
+ assert int(node.indices[-2]) == 0
+
+ tvm_ffi.structural_walk(lowered.body, check)
+ values_per_lane = k_inner // 4
+ assert sorted(allocations) == sorted(
+ [(2, 1, 1, values_per_lane), (1, 2, 1, values_per_lane), (2, 2, 1, 4)]
+ )
+ assert len(accesses) >= 8
+
+
if __name__ == "__main__":
tvm.testing.main()
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_inject_permuted_layout.py b/tests/python/s_tir/transform/test_s_tir_transform_inject_permuted_layout.py
index bbeebc3011..ba48398f68 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_inject_permuted_layout.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_inject_permuted_layout.py
@@ -14,335 +14,55 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-# ruff: noqa: E501, E741
+import pytest
+
import tvm
-import tvm.s_tir
import tvm.testing
-from tvm import IRModule
from tvm.script import s_tir as Ts
from tvm.script import tirx as T
-from tvm.tirx import Function
-
-
-def _check_function_transform(before: Function, expected: Function):
- before_module = IRModule.from_expr(before)
- after_module = tvm.s_tir.transform.InjectPermutedLayout()(before_module)
-
- after = after_module["before"].without_attr("global_symbol")
- expected = expected.without_attr("global_symbol")
-
- tvm.ir.assert_structural_equal(after, expected)
-
-
-# This pass is adapted from another previous pass, so we need to ensure backward compatibility here
-def test_backward_compatibility_shared_a():
- # fmt: off
- @Ts.function
- def before(X: T.Tensor((4096, 4096), "float16")):
- # with Ts.sblock("root"):
- for blockIdx_y in T.parallel(256, annotations={"thread_binding": "blockIdx.y"}):
- for threadIdx_y in T.parallel(4, annotations={"thread_binding": "threadIdx.y"}):
- for threadIdx_x in T.parallel(32, annotations={"thread_binding": "threadIdx.x"}):
- with Ts.sblock(""):
- Ts.reads(X[blockIdx_y // 8 * 128 + threadIdx_y * 8 + threadIdx_x // 4:blockIdx_y // 8 * 128 + threadIdx_y * 8 + threadIdx_x // 4 + 97, threadIdx_x % 4 * 8:threadIdx_x % 4 * 8 + 4072])
- Ts.writes()
- for ax2_0_0 in range(128):
- with Ts.sblock(""):
- Ts.reads(X[blockIdx_y // 8 * 128 + threadIdx_y * 8 + threadIdx_x // 4:blockIdx_y // 8 * 128 + threadIdx_y * 8 + threadIdx_x // 4 + 97, ax2_0_0 * 32 + threadIdx_x % 4 * 8:ax2_0_0 * 32 + threadIdx_x % 4 * 8 + 8])
- Ts.writes()
- X_reindex_shared_dyn = Ts.sblock_alloc_tensor((128, 32), "float16", strides=(32, 1), scope="shared.dyn")
- with Ts.sblock("X_reindex_shared.dyn"):
- Ts.reads(X[blockIdx_y // 8 * 128 + threadIdx_y * 8 + threadIdx_x // 4:blockIdx_y // 8 * 128 + threadIdx_y * 8 + threadIdx_x // 4 + 97, ax2_0_0 * 32 + threadIdx_x % 4 * 8:ax2_0_0 * 32 + threadIdx_x % 4 * 8 + 8])
- Ts.writes(X_reindex_shared_dyn[threadIdx_y * 8 + threadIdx_x // 4:threadIdx_y * 8 + threadIdx_x // 4 + 97, threadIdx_x % 4 * 8:threadIdx_x % 4 * 8 + 8])
- Ts.sblock_attr({"permuted_layout": "g2s_A"})
- for ax0_ax1_fused_0 in range(4):
- for ax0_ax1_fused_3 in T.vectorized(8):
- X_reindex_shared_dyn[ax0_ax1_fused_0 * 32 + threadIdx_y * 8 + threadIdx_x // 4, threadIdx_x % 4 * 8 + ax0_ax1_fused_3] = X[blockIdx_y // 8 * 128 + ax0_ax1_fused_0 * 32 + threadIdx_y * 8 + threadIdx_x // 4, ax2_0_0 * 32 + threadIdx_x % 4 * 8 + ax0_ax1_fused_3]
- for ax2_0_1 in range(4):
- with Ts.sblock(""):
- Ts.reads(X_reindex_shared_dyn[threadIdx_y // 2 * 64:threadIdx_y // 2 * 64 + 64, ax2_0_1 * 8:ax2_0_1 * 8 + 8])
- Ts.writes()
- X_reindex_shared_dyn_m16n8k8_matrixA = Ts.sblock_alloc_tensor((64, 8), "float16", scope="mma.matrixA")
- for ax0_0, ax1_0 in T.grid(2, 1):
- with Ts.sblock("X_reindex_shared.dyn_mma.matrixA_o"):
- Ts.reads(X_reindex_shared_dyn[threadIdx_y // 2 * 64 + ax0_0 * 32:threadIdx_y // 2 * 64 + ax0_0 * 32 + 32, ax2_0_1 * 8:ax2_0_1 * 8 + 8])
- Ts.writes(X_reindex_shared_dyn_m16n8k8_matrixA[ax0_0 * 32:ax0_0 * 32 + 32, 0:8])
- Ts.sblock_attr({"permuted_layout": "s2l_A"})
- T.ptx_legacy.ldmatrix(T.bool(False), 4, ".b16", X_reindex_shared_dyn_m16n8k8_matrixA.data, ax0_0 * 8, T.ptr_byte_offset(X_reindex_shared_dyn.data, (threadIdx_y // 2 * 2048 + ax0_0 * 1024 + ax2_0_1 * 8) * 2, ty=X_reindex_shared_dyn.data.ty), threadIdx_x * 32, ty="float16")
-
- @Ts.function
- def expected(X: T.Tensor((4096, 4096), "float16")):
- for blockIdx_y in T.parallel(256, annotations={"thread_binding": "blockIdx.y"}):
- for threadIdx_y in T.parallel(4, annotations={"thread_binding": "threadIdx.y"}):
- for threadIdx_x in T.parallel(32, annotations={"thread_binding": "threadIdx.x"}):
- with Ts.sblock(""):
- for ax2_0_0 in T.serial(128):
- with Ts.sblock(""):
- X_reindex_shared_dyn = Ts.sblock_alloc_tensor((128, 32), "float16", strides=(32, 1), scope="shared.dyn")
- with Ts.sblock("X_reindex_shared.dyn"):
- # annotate the reads and writes because they cannot be inferred from tirx.bitwise_xor
- Ts.reads(X[blockIdx_y // 8 * 128 + threadIdx_y * 8 + threadIdx_x // 4:blockIdx_y // 8 * 128 + threadIdx_y * 8 + threadIdx_x // 4 + 97, ax2_0_0 * 32 + threadIdx_x % 4 * 8:ax2_0_0 * 32 + threadIdx_x % 4 * 8 + 8])
- Ts.writes(X_reindex_shared_dyn[threadIdx_y * 8 + threadIdx_x // 4:threadIdx_y * 8 + threadIdx_x // 4 + 97, threadIdx_x % 4 * 8:threadIdx_x % 4 * 8 + 8])
- for ax0_ax1_fused_0 in range(4):
- for ax0_ax1_fused_3 in T.vectorized(8):
- X_reindex_shared_dyn[ax0_ax1_fused_0 * 32 + threadIdx_y * 8 + threadIdx_x // 4, T.bitwise_xor(threadIdx_x % 4, threadIdx_x // 8) * 8 + ax0_ax1_fused_3] = X[blockIdx_y // 8 * 128 + ax0_ax1_fused_0 * 32 + threadIdx_y * 8 + threadIdx_x // 4, ax2_0_0 * 32 + threadIdx_x % 4 * 8 + ax0_ax1_fused_3]
- for ax2_0_1 in T.serial(4):
- with Ts.sblock(""):
- X_reindex_shared_dyn_m16n8k8_matrixA = Ts.sblock_alloc_tensor((64, 8), "float16", scope="mma.matrixA")
- for ax0_0, ax1_0 in T.grid(2, 1):
- with Ts.sblock("X_reindex_shared.dyn_mma.matrixA_o"):
- Ts.reads(X_reindex_shared_dyn[threadIdx_y // 2 * 64 + ax0_0 * 32:threadIdx_y // 2 * 64 + ax0_0 * 32 + 32, ax2_0_1 * 8:ax2_0_1 * 8 + 8])
- Ts.writes(X_reindex_shared_dyn_m16n8k8_matrixA[ax0_0 * 32:ax0_0 * 32 + 32, 0:8])
- T.ptx_legacy.ldmatrix(T.bool(False), 4, ".b16", X_reindex_shared_dyn_m16n8k8_matrixA.data, ax0_0 * 8, T.ptr_byte_offset(X_reindex_shared_dyn.data, (threadIdx_y // 2 * 2048 + ax0_0 * 1024 + threadIdx_x * 32 + T.bitwise_xor(ax2_0_1, threadIdx_x % 8 // 2) * 8) * 2, ty=X_reindex_shared_dyn.data.ty), 0, ty="float16")
- # fmt: on
- _check_function_transform(before, expected)
-
-
-def test_backward_compatibility_shared_a_and_b():
- # fmt: off
- @Ts.function
- def before(X: T.Tensor((4096, 4096), "float16"), Y: T.Tensor((4096, 4096), "float16")):
- for blockIdx_x in T.parallel(4, annotations={"thread_binding": "blockIdx.x"}):
- for blockIdx_y in T.parallel(256, annotations={"thread_binding": "blockIdx.y"}):
- for threadIdx_y in T.parallel(4, annotations={"thread_binding": "threadIdx.y"}):
- for threadIdx_x in T.parallel(32, annotations={"thread_binding": "threadIdx.x"}):
- with Ts.sblock(""):
- for ax2_0_0 in T.serial(128):
- with Ts.sblock(""):
- X_reindex_shared_dyn = Ts.sblock_alloc_tensor((128, 32), "float16", strides=(32, 1), scope="shared.dyn")
- Y_reindex_shared_dyn = Ts.sblock_alloc_tensor((32, 128), "float16", strides=(128, 1), scope="shared.dyn")
- with Ts.sblock("X_reindex_shared.dyn"):
- Ts.sblock_attr({"permuted_layout": "g2s_A"})
- for ax0_ax1_fused_0 in range(4):
- for ax0_ax1_fused_3 in T.vectorized(8):
- X_reindex_shared_dyn[ax0_ax1_fused_0 * 32 + threadIdx_y * 8 + threadIdx_x // 4, threadIdx_x % 4 * 8 + ax0_ax1_fused_3] = X[blockIdx_y // 8 * 128 + ax0_ax1_fused_0 * 32 + threadIdx_y * 8 + threadIdx_x // 4, ax2_0_0 * 32 + threadIdx_x % 4 * 8 + ax0_ax1_fused_3]
- with Ts.sblock("Y_reindex_shared.dyn"):
- Ts.sblock_attr({"permuted_layout": "g2s_B"})
- for ax0_ax1_fused_0 in range(4):
- for ax0_ax1_fused_3 in T.vectorized(8):
- Y_reindex_shared_dyn[ax0_ax1_fused_0 * 8 + threadIdx_y * 2 + threadIdx_x // 16, threadIdx_x % 16 * 8 + ax0_ax1_fused_3] = Y[ax2_0_0 * 32 + ax0_ax1_fused_0 * 8 + threadIdx_y * 2 + threadIdx_x // 16, blockIdx_x * 1024 + blockIdx_y % 8 * 128 + threadIdx_x % 16 * 8 + ax0_ax1_fused_3]
- for ax2_0_1 in T.serial(4):
- with Ts.sblock(""):
- X_reindex_shared_dyn_m16n8k8_matrixA = Ts.sblock_alloc_tensor((64, 8), "float16", scope="mma.matrixA")
- Y_reindex_shared_dyn_m16n8k8_matrixB = Ts.sblock_alloc_tensor((8, 64), "float16", scope="mma.matrixB")
- for ax0_0, ax1_0 in T.grid(2, 1):
- with Ts.sblock("X_reindex_shared.dyn_mma.matrixA_o"):
- Ts.reads(X_reindex_shared_dyn[threadIdx_y // 2 * 64 + ax0_0 * 32:threadIdx_y // 2 * 64 + ax0_0 * 32 + 32, ax2_0_1 * 8:ax2_0_1 * 8 + 8])
- Ts.writes(X_reindex_shared_dyn_m16n8k8_matrixA[ax0_0 * 32:ax0_0 * 32 + 32, 0:8])
- Ts.sblock_attr({"permuted_layout": "s2l_A"})
- T.ptx_legacy.ldmatrix(T.bool(False), 4, ".b16", X_reindex_shared_dyn_m16n8k8_matrixA.data, ax0_0 * 8, T.ptr_byte_offset(X_reindex_shared_dyn.data, (threadIdx_y // 2 * 2048 + ax0_0 * 1024 + ax2_0_1 * 8) * 2, ty=X_reindex_shared_dyn.data.ty), threadIdx_x * 32, ty="float16")
- for ax0_0, ax1_0 in T.grid(1, 2):
- with Ts.sblock("Y_reindex_shared.dyn_mma.matrixB_o"):
- Ts.reads(Y_reindex_shared_dyn[ax2_0_1 * 8:ax2_0_1 * 8 + 8, threadIdx_y % 2 * 64 + ax1_0 * 32:threadIdx_y % 2 * 64 + ax1_0 * 32 + 32])
- Ts.writes(Y_reindex_shared_dyn_m16n8k8_matrixB[0:8, ax1_0 * 32:ax1_0 * 32 + 32])
- Ts.sblock_attr({"permuted_layout": "s2l_B"})
- T.ptx_legacy.ldmatrix(T.bool(True), 4, ".b16", Y_reindex_shared_dyn_m16n8k8_matrixB.data, ax1_0 * 8, T.ptr_byte_offset(Y_reindex_shared_dyn.data, (ax2_0_1 * 1024 + threadIdx_y % 2 * 64 + ax1_0 * 32) * 2, ty=Y_reindex_shared_dyn.data.ty), threadIdx_x % 8 * 128 + threadIdx_x // 8 * 8, ty="float16")
-
- @Ts.function
- def expected(X: T.Tensor((4096, 4096), "float16"), Y: T.Tensor((4096, 4096), "float16")):
- for blockIdx_x in T.parallel(4, annotations={"thread_binding": "blockIdx.x"}):
- for blockIdx_y in T.parallel(256, annotations={"thread_binding": "blockIdx.y"}):
- for threadIdx_y in T.parallel(4, annotations={"thread_binding": "threadIdx.y"}):
- for threadIdx_x in T.parallel(32, annotations={"thread_binding": "threadIdx.x"}):
- with Ts.sblock(""):
- Ts.reads(X[blockIdx_y // 8 * 128 + threadIdx_y * 8 + threadIdx_x // 4:blockIdx_y // 8 * 128 + threadIdx_y * 8 + threadIdx_x // 4 + 97, threadIdx_x % 4 * 8:threadIdx_x % 4 * 8 + 4072], Y[threadIdx_y * 2 + threadIdx_x // 16:threadIdx_y * 2 + threadIdx_x // 16 + 4089, blockIdx_x * 1024 + blockIdx_y % 8 * 128 + threadIdx_x % 16 * 8:blockIdx_x * 1024 + blockIdx_y % 8 * 128 + threadIdx_x % 16 * 8 + 8])
- Ts.writes()
- for ax2_0_0 in T.serial(128):
- with Ts.sblock(""):
- Ts.reads(X[blockIdx_y // 8 * 128 + threadIdx_y * 8 + threadIdx_x // 4:blockIdx_y // 8 * 128 + threadIdx_y * 8 + threadIdx_x // 4 + 97, ax2_0_0 * 32 + threadIdx_x % 4 * 8:ax2_0_0 * 32 + threadIdx_x % 4 * 8 + 8], Y[ax2_0_0 * 32 + threadIdx_y * 2 + threadIdx_x // 16:ax2_0_0 * 32 + threadIdx_y * 2 + threadIdx_x // 16 + 25, blockIdx_x * 1024 + blockIdx_y % 8 * 128 + threadIdx_x % 16 * 8:blockIdx_x * 1024 + blockIdx_y % 8 * 128 + threadIdx_x % 16 * 8 + 8])
- Ts.writes()
- X_reindex_shared_dyn = Ts.sblock_alloc_tensor((128, 32), "float16", strides=(32, 1), scope="shared.dyn")
- Y_reindex_shared_dyn = Ts.sblock_alloc_tensor((32, 128), "float16", strides=(128, 1), scope="shared.dyn")
- with Ts.sblock("X_reindex_shared.dyn"):
- Ts.reads(X[blockIdx_y // 8 * 128 + threadIdx_y * 8 + threadIdx_x // 4:blockIdx_y // 8 * 128 + threadIdx_y * 8 + threadIdx_x // 4 + 97, ax2_0_0 * 32 + threadIdx_x % 4 * 8:ax2_0_0 * 32 + threadIdx_x % 4 * 8 + 8])
- Ts.writes(X_reindex_shared_dyn[threadIdx_y * 8 + threadIdx_x // 4:threadIdx_y * 8 + threadIdx_x // 4 + 97, threadIdx_x % 4 * 8:threadIdx_x % 4 * 8 + 8])
- for ax0_ax1_fused_0 in range(4):
- for ax0_ax1_fused_3 in T.vectorized(8):
- X_reindex_shared_dyn[ax0_ax1_fused_0 * 32 + threadIdx_y * 8 + threadIdx_x // 4, T.bitwise_xor(threadIdx_x % 4, threadIdx_x // 8) * 8 + ax0_ax1_fused_3] = X[blockIdx_y // 8 * 128 + ax0_ax1_fused_0 * 32 + threadIdx_y * 8 + threadIdx_x // 4, ax2_0_0 * 32 + threadIdx_x % 4 * 8 + ax0_ax1_fused_3]
- with Ts.sblock("Y_reindex_shared.dyn"):
- Ts.reads(Y[ax2_0_0 * 32 + threadIdx_y * 2 + threadIdx_x // 16:ax2_0_0 * 32 + threadIdx_y * 2 + threadIdx_x // 16 + 25, blockIdx_x * 1024 + blockIdx_y % 8 * 128 + threadIdx_x % 16 * 8:blockIdx_x * 1024 + blockIdx_y % 8 * 128 + threadIdx_x % 16 * 8 + 8])
- Ts.writes(Y_reindex_shared_dyn[threadIdx_y * 2 + threadIdx_x // 16:threadIdx_y * 2 + threadIdx_x // 16 + 25, threadIdx_x % 16 * 8:threadIdx_x % 16 * 8 + 8])
- for ax0_ax1_fused_0 in range(4):
- for ax0_ax1_fused_3 in T.vectorized(8):
- Y_reindex_shared_dyn[ax0_ax1_fused_0 * 8 + threadIdx_y * 2 + threadIdx_x // 16, T.bitwise_xor(threadIdx_x % 16, threadIdx_y * 2 + threadIdx_x // 16) * 8 + ax0_ax1_fused_3] = Y[ax2_0_0 * 32 + ax0_ax1_fused_0 * 8 + threadIdx_y * 2 + threadIdx_x // 16, blockIdx_x * 1024 + blockIdx_y % 8 * 128 + threadIdx_x % 16 * 8 + ax0_ax1_fused_3]
- for ax2_0_1 in T.serial(4):
- with Ts.sblock(""):
- X_reindex_shared_dyn_m16n8k8_matrixA = Ts.sblock_alloc_tensor((64, 8), "float16", scope="mma.matrixA")
- Y_reindex_shared_dyn_m16n8k8_matrixB = Ts.sblock_alloc_tensor((8, 64), "float16", scope="mma.matrixB")
- for ax0_0, ax1_0 in T.grid(2, 1):
- with Ts.sblock("X_reindex_shared.dyn_mma.matrixA_o"):
- Ts.reads(X_reindex_shared_dyn[threadIdx_y // 2 * 64 + ax0_0 * 32:threadIdx_y // 2 * 64 + ax0_0 * 32 + 32, ax2_0_1 * 8:ax2_0_1 * 8 + 8])
- Ts.writes(X_reindex_shared_dyn_m16n8k8_matrixA[ax0_0 * 32:ax0_0 * 32 + 32, 0:8])
- T.ptx_legacy.ldmatrix(T.bool(False), 4, ".b16", X_reindex_shared_dyn_m16n8k8_matrixA.data, ax0_0 * 8, T.ptr_byte_offset(X_reindex_shared_dyn.data, (threadIdx_y // 2 * 2048 + ax0_0 * 1024 + threadIdx_x * 32 + T.bitwise_xor(ax2_0_1, threadIdx_x % 8 // 2) * 8) * 2, ty=X_reindex_shared_dyn.data.ty), 0, ty="float16")
- for ax0_0, ax1_0 in T.grid(1, 2):
- with Ts.sblock("Y_reindex_shared.dyn_mma.matrixB_o"):
- Ts.reads(Y_reindex_shared_dyn[ax2_0_1 * 8:ax2_0_1 * 8 + 8, threadIdx_y % 2 * 64 + ax1_0 * 32:threadIdx_y % 2 * 64 + ax1_0 * 32 + 32])
- Ts.writes(Y_reindex_shared_dyn_m16n8k8_matrixB[0:8, ax1_0 * 32:ax1_0 * 32 + 32])
- T.ptx_legacy.ldmatrix(T.bool(True), 4, ".b16", Y_reindex_shared_dyn_m16n8k8_matrixB.data, ax1_0 * 8, T.ptr_byte_offset(Y_reindex_shared_dyn.data, (ax2_0_1 * 1024 + threadIdx_x % 8 * 128 + T.bitwise_xor(threadIdx_y % 2 * 8 + ax1_0 * 4 + threadIdx_x // 8, threadIdx_x % 8) * 8) * 2, ty=Y_reindex_shared_dyn.data.ty), 0, ty="float16")
- # fmt: on
- _check_function_transform(before, expected)
-def test_tensor_a():
- # fmt: off
+@pytest.mark.parametrize("columns", [32, 64])
+@pytest.mark.parametrize("scope", ["shared", "shared.dyn"])
+@pytest.mark.parametrize("annotation", [1, "g2s_A"])
+def test_permute_shared_tensor_accesses(columns, scope, annotation):
@Ts.function
- def before(A: T.Tensor((T.int64(128), T.int64(32)), 'float16')):
-
- A_shared_dyn = Ts.sblock_alloc_tensor((T.int64(128), T.int64(32)), "float16", scope="shared.dyn")
- A_warp = Ts.sblock_alloc_tensor((T.int64(4), T.int64(1), T.int64(32), T.int64(8)), "float16", scope="warp")
- for threadIdx_z in T.parallel(T.int64(2), annotations={"thread_binding": "threadIdx.z"}):
- for threadIdx_y in T.parallel(T.int64(2), annotations={"thread_binding": "threadIdx.y"}):
- for threadIdx_x in T.parallel(T.int64(32), annotations={"thread_binding": "threadIdx.x"}):
- for v0 in range(T.int64(4)):
- for v1 in T.vectorized(T.int64(8)):
- with Ts.sblock("A_reindex_shared.dyn"):
- Ts.sblock_attr({"permuted_layout": 1})
- A_shared_dyn[
- v0 * T.int64(32) + threadIdx_z * T.int64(16) + threadIdx_y * T.int64(8) + threadIdx_x // T.int64(4),
- threadIdx_x % T.int64(4) * T.int64(8) + v1
- ] = A[
- (v0 * T.int64(32) + threadIdx_z * T.int64(16) + threadIdx_y * T.int64(8) + threadIdx_x // T.int64(4)) % T.int64(32),
- threadIdx_x % T.int64(4) * T.int64(8) + v1
- ]
- for v0, v1 in T.grid(T.int64(2), T.int64(4)):
- with Ts.sblock("A_reindex_shared.dyn_warp_o"):
- Ts.sblock_attr({"permuted_layout": 1})
- with Ts.sblock("A_reindex_shared.dyn_warp_o"):
- Ts.reads(A_shared_dyn[threadIdx_z * T.int64(64) + v1 * T.int64(16):threadIdx_z * T.int64(64) + v1 * T.int64(16) + T.int64(16), v0 * T.int64(16):v0 * T.int64(16) + T.int64(16)])
- Ts.writes(A_warp[v1, T.int64(0), T.int64(0):T.int64(32), T.int64(0):T.int64(8)])
- T.ptx_legacy.ldmatrix(T.bool(False), 4, ".b16",
- A_warp.data,
- v1 * T.int64(256) + threadIdx_x * T.int64(8),
- T.ptr_byte_offset(A_shared_dyn.data, (threadIdx_z * T.int64(2048) + v1 * T.int64(512) + v0 * T.int64(16)) * 2, ty=A_shared_dyn.data.ty),
- threadIdx_x % T.int64(16) * T.int64(32) + threadIdx_x // T.int64(16) * T.int64(8)
- , ty="float16")
-
- @Ts.function
- def expected(A: T.Tensor((T.int64(128), T.int64(32)), "float16")):
- A_shared_dyn = Ts.sblock_alloc_tensor((T.int64(128), T.int64(32)), "float16", scope="shared.dyn")
- A_warp = Ts.sblock_alloc_tensor((T.int64(4), T.int64(1), T.int64(32), T.int64(8)), "float16", scope="warp")
- for threadIdx_z in T.parallel(T.int64(2), annotations={"thread_binding": "threadIdx.z"}):
- for threadIdx_y in T.parallel(T.int64(2), annotations={"thread_binding": "threadIdx.y"}):
- for threadIdx_x in T.parallel(T.int64(32), annotations={"thread_binding": "threadIdx.x"}):
- for v0 in range(T.int64(4)):
- for v1 in T.vectorized(T.int64(8)):
- with Ts.sblock("A_reindex_shared.dyn"):
- Ts.reads(A[(v0 * T.int64(32) + threadIdx_z * T.int64(16) + threadIdx_y * T.int64(8) + threadIdx_x // T.int64(4)) % T.int64(32), threadIdx_x % T.int64(4) * T.int64(8) + v1])
- Ts.writes(A_shared_dyn[v0 * T.int64(32) + threadIdx_z * T.int64(16) + threadIdx_y * T.int64(8) + threadIdx_x // T.int64(4), threadIdx_x % T.int64(4) * T.int64(8) + v1])
- A_shared_dyn[v0 * T.int64(32) + threadIdx_z * T.int64(16) + threadIdx_y * T.int64(8) + threadIdx_x // T.int64(4), T.bitwise_xor(threadIdx_x % T.int64(4), threadIdx_x // T.int64(8)) * T.int64(8) + v1] = A[(v0 * T.int64(32) + threadIdx_z * T.int64(16) + threadIdx_y * T.int64(8) + threadIdx_x // T.int64(4)) % T.int64(32), threadIdx_x % T.int64(4) * T.int64(8) + v1]
- for v0, v1 in T.grid(T.int64(2), T.int64(4)):
- with Ts.sblock("A_reindex_shared.dyn_warp_o"):
- Ts.reads(A_shared_dyn[threadIdx_z * T.int64(64) + v1 * T.int64(16):threadIdx_z * T.int64(64) + v1 * T.int64(16) + T.int64(16), v0 * T.int64(16):v0 * T.int64(16) + T.int64(16)])
- Ts.writes(A_warp[v1, T.int64(0), T.int64(0):T.int64(32), T.int64(0):T.int64(8)])
- with Ts.sblock("A_reindex_shared.dyn_warp_o"):
- Ts.reads(A_shared_dyn[threadIdx_z * T.int64(64) + v1 * T.int64(16):threadIdx_z * T.int64(64) + v1 * T.int64(16) + T.int64(16), v0 * T.int64(16):v0 * T.int64(16) + T.int64(16)])
- Ts.writes(A_warp[v1, T.int64(0), T.int64(0):T.int64(32), T.int64(0):T.int64(8)])
- T.ptx_legacy.ldmatrix(T.bool(False), 4, ".b16", A_warp.data, v1 * T.int64(256) + threadIdx_x * T.int64(8), T.ptr_byte_offset(A_shared_dyn.data, (threadIdx_z * T.int64(2048) + v1 * T.int64(512) + threadIdx_x % T.int64(16) * T.int64(32) + T.bitwise_xor(v0 * T.int64(2) + threadIdx_x // T.int64(16), threadIdx_x % T.int64(8) // T.int64(2)) * T.int64(8)) * 2, ty=A_shared_dyn.data.ty), T.int64(0), ty="float16")
-
- # fmt: on
- _check_function_transform(before, expected)
-
-
-def test_tensor_b():
- # fmt: off
- @Ts.function
- def before(B: T.Tensor((T.int64(128), T.int64(32)), "float16")):
- B_shared_dyn = Ts.sblock_alloc_tensor((T.int64(128), T.int64(32)), "float16", scope="shared.dyn")
- for threadIdx_z in T.parallel(T.int64(2), annotations={"thread_binding": "threadIdx.z"}):
- for threadIdx_y in T.parallel(T.int64(2), annotations={"thread_binding": "threadIdx.y"}):
- for threadIdx_x in T.parallel(T.int64(32), annotations={"thread_binding": "threadIdx.x"}):
- for v0 in range(T.int64(4)):
- for v1 in T.vectorized(T.int64(8)):
- with Ts.sblock("B_reindex_shared.dyn"):
- Ts.sblock_attr({"permuted_layout": 1})
- B_shared_dyn[v0 * T.int64(32) + threadIdx_z * T.int64(16) + threadIdx_y * T.int64(8) + threadIdx_x // T.int64(4), threadIdx_x % T.int64(4) * T.int64(8) + v1] = B[v0 * T.int64(32) + threadIdx_z * T.int64(16) + threadIdx_y * T.int64(8) + threadIdx_x // T.int64(4), threadIdx_x % T.int64(4) * T.int64(8) + v1]
- for v0 in range(T.int64(2)):
- with Ts.sblock(""):
- B_warp = Ts.sblock_alloc_tensor((T.int64(4), T.int64(1), T.int64(32), T.int64(8)), "float16", scope="warp")
- for v1 in range(T.int64(4)):
- with Ts.sblock("B_reindex_shared.dyn_warp_o"):
- Ts.sblock_attr({"permuted_layout": 1})
- with Ts.sblock("B_reindex_shared.dyn_warp_o"):
- Ts.reads(B_shared_dyn[threadIdx_y * T.int64(64) + v1 * T.int64(16):threadIdx_y * T.int64(64) + v1 * T.int64(16) + T.int64(16), v0 * T.int64(16):v0 * T.int64(16) + T.int64(16)])
- Ts.writes(B_warp[v1, T.int64(0), T.int64(0):T.int64(32), T.int64(0):T.int64(8)])
- T.ptx_legacy.ldmatrix(T.bool(False), 4, ".b16", B_warp.data, v1 * T.int64(256) + threadIdx_x * T.int64(8), T.ptr_byte_offset(B_shared_dyn.data, (threadIdx_y * T.int64(2048) + v1 * T.int64(512) + v0 * T.int64(16)) * 2, ty=B_shared_dyn.data.ty), threadIdx_x // T.int64(16) * T.int64(256) + threadIdx_x % T.int64(8) * T.int64(32) + threadIdx_x % T.int64(16) // T.int64(8) * T.int64(8), ty="float16")
+ def before(A: T.Tensor((8, columns), "float16"), B: T.Tensor((8, columns), "float16")):
+ shared = Ts.sblock_alloc_tensor((8, columns), "float16", scope=scope)
+ for i, j in T.grid(8, columns):
+ with Ts.sblock("load"):
+ vi, vj = Ts.axis.remap("SS", [i, j])
+ Ts.sblock_attr({"permuted_layout": annotation})
+ shared[vi, vj] = A[vi, vj]
+ for i, j in T.grid(8, columns):
+ with Ts.sblock("store"):
+ vi, vj = Ts.axis.remap("SS", [i, j])
+ Ts.sblock_attr({"permuted_layout": annotation})
+ B[vi, vj] = shared[vi, vj]
+
+ def permuted_column(i, j):
+ return T.bitwise_xor(j // 8, i // 2 if columns == 32 else i) * 8 + j % 8
@Ts.function
- def expected(B: T.Tensor((T.int64(128), T.int64(32)), "float16")):
- B_shared_dyn = Ts.sblock_alloc_tensor((T.int64(128), T.int64(32)), "float16", scope="shared.dyn")
- for threadIdx_z in T.parallel(T.int64(2), annotations={"thread_binding": "threadIdx.z"}):
- for threadIdx_y in T.parallel(T.int64(2), annotations={"thread_binding": "threadIdx.y"}):
- for threadIdx_x in T.parallel(T.int64(32), annotations={"thread_binding": "threadIdx.x"}):
- for v0 in range(T.int64(4)):
- for v1 in T.vectorized(T.int64(8)):
- with Ts.sblock("B_reindex_shared.dyn"):
- Ts.reads(B[v0 * T.int64(32) + threadIdx_z * T.int64(16) + threadIdx_y * T.int64(8) + threadIdx_x // T.int64(4), threadIdx_x % T.int64(4) * T.int64(8) + v1])
- Ts.writes(B_shared_dyn[v0 * T.int64(32) + threadIdx_z * T.int64(16) + threadIdx_y * T.int64(8) + threadIdx_x // T.int64(4), threadIdx_x % T.int64(4) * T.int64(8) + v1])
- B_shared_dyn[v0 * T.int64(32) + threadIdx_z * T.int64(16) + threadIdx_y * T.int64(8) + threadIdx_x // T.int64(4), T.bitwise_xor(threadIdx_x % T.int64(4), threadIdx_x // T.int64(8)) * T.int64(8) + v1] = B[v0 * T.int64(32) + threadIdx_z * T.int64(16) + threadIdx_y * T.int64(8) + threadIdx_x // T.int64(4), threadIdx_x % T.int64(4) * T.int64(8) + v1]
- for v0 in range(T.int64(2)):
- with Ts.sblock(""):
- B_warp = Ts.sblock_alloc_tensor((T.int64(4), T.int64(1), T.int64(32), T.int64(8)), "float16", scope="warp")
- for v1 in range(T.int64(4)):
- with Ts.sblock("B_reindex_shared.dyn_warp_o"):
- Ts.reads(B_shared_dyn[threadIdx_y * T.int64(64) + v1 * T.int64(16):threadIdx_y * T.int64(64) + v1 * T.int64(16) + T.int64(16), v0 * T.int64(16):v0 * T.int64(16) + T.int64(16)])
- Ts.writes(B_warp[v1, T.int64(0), T.int64(0):T.int64(32), T.int64(0):T.int64(8)])
- with Ts.sblock("B_reindex_shared.dyn_warp_o"):
- Ts.reads(B_shared_dyn[threadIdx_y * T.int64(64) + v1 * T.int64(16):threadIdx_y * T.int64(64) + v1 * T.int64(16) + T.int64(16), v0 * T.int64(16):v0 * T.int64(16) + T.int64(16)])
- Ts.writes(B_warp[v1, T.int64(0), T.int64(0):T.int64(32), T.int64(0):T.int64(8)])
- T.ptx_legacy.ldmatrix(T.bool(False), 4, ".b16", B_warp.data, v1 * T.int64(256) + threadIdx_x * T.int64(8), T.ptr_byte_offset(B_shared_dyn.data, (threadIdx_y * T.int64(2048) + v1 * T.int64(512) + threadIdx_x // T.int64(16) * T.int64(256) + threadIdx_x % T.int64(8) * T.int64(32) + T.bitwise_xor(v0 * T.int64(2) + threadIdx_x % T.int64(16) // T.int64(8), threadIdx_x % T.int64(8) // T.int64(2)) * T.int64(8)) * 2, ty=B_shared_dyn.data.ty), T.int64(0), ty [...]
-
- # fmt: on
- _check_function_transform(before, expected)
-
-
-def test_tensor_c_fp32():
- # fmt: off
- @Ts.function
- def before(O: T.Tensor((T.int64(128), T.int64(128)), 'float16')):
-
- O_shared_dyn = Ts.sblock_alloc_tensor((T.int64(128), T.int64(128)), scope="shared.dyn")
- O_warp = Ts.sblock_alloc_tensor((T.int64(4), T.int64(4), T.int64(32), T.int64(8)), scope="warp")
- for threadIdx_z in T.parallel(T.int64(2), annotations={"thread_binding": "threadIdx.z"}):
- for threadIdx_y in T.parallel(T.int64(2), annotations={"thread_binding": "threadIdx.y"}):
- for threadIdx_x in T.parallel(T.int64(32), annotations={"thread_binding": "threadIdx.x"}):
- for v0, v1 in T.grid(T.int64(4), T.int64(4)):
- with Ts.sblock("O.dyn_warp_o"):
- Ts.sblock_attr({"permuted_layout": 1})
- with Ts.sblock("O.dyn_warp_o"):
- for local_id in range(T.int64(8)):
- O_shared_dyn[threadIdx_z * T.int64(64) + v0 * T.int64(16) + local_id % T.int64(4) // T.int64(2) * T.int64(8) + threadIdx_x // T.int64(4), threadIdx_y * T.int64(64) + v1 * T.int64(16) + local_id // T.int64(4) * T.int64(8) + threadIdx_x % T.int64(4) * T.int64(2) + local_id % T.int64(2)] = O_warp[v0, v1, threadIdx_x, local_id]
- for v0 in range(T.int64(16)):
- for v1 in T.vectorized(T.int64(8)):
- with Ts.sblock("O.dyn"):
- Ts.sblock_attr({"permuted_layout": 1})
- O[v0 * T.int64(8) + threadIdx_z * T.int64(4) + threadIdx_y * T.int64(2) + threadIdx_x // T.int64(16), threadIdx_x % T.int64(16) * T.int64(8) + v1] = T.Cast(O_shared_dyn[v0 * T.int64(8) + threadIdx_z * T.int64(4) + threadIdx_y * T.int64(2) + threadIdx_x // T.int64(16), threadIdx_x % T.int64(16) * T.int64(8) + v1], "float16")
-
- @Ts.function
- def expected(O: T.Tensor((T.int64(128), T.int64(128)), "float16")):
- # with Ts.sblock("root"):
- O_shared_dyn = Ts.sblock_alloc_tensor((T.int64(128), T.int64(128)), scope="shared.dyn")
- O_warp = Ts.sblock_alloc_tensor((T.int64(4), T.int64(4), T.int64(32), T.int64(8)), scope="warp")
- for threadIdx_z in T.parallel(T.int64(2), annotations={"thread_binding": "threadIdx.z"}):
- for threadIdx_y in T.parallel(T.int64(2), annotations={"thread_binding": "threadIdx.y"}):
- for threadIdx_x in T.parallel(T.int64(32), annotations={"thread_binding": "threadIdx.x"}):
- for v0, v1 in T.grid(T.int64(4), T.int64(4)):
- with Ts.sblock("O.dyn_warp_o"):
- Ts.reads(O_warp[v0, v1, threadIdx_x, T.int64(0):T.int64(8)])
- Ts.writes(O_shared_dyn[threadIdx_z * T.int64(64) + v0 * T.int64(16) + threadIdx_x // T.int64(4):threadIdx_z * T.int64(64) + v0 * T.int64(16) + threadIdx_x // T.int64(4) + T.int64(9), threadIdx_y * T.int64(64) + v1 * T.int64(16) + threadIdx_x % T.int64(4) * T.int64(2):threadIdx_y * T.int64(64) + v1 * T.int64(16) + threadIdx_x % T.int64(4) * T.int64(2) + T.int64(10)])
- with Ts.sblock("O.dyn_warp_o"):
- Ts.reads(O_warp[v0, v1, threadIdx_x, T.int64(0):T.int64(8)])
- Ts.writes(O_shared_dyn[threadIdx_z * T.int64(64) + v0 * T.int64(16) + threadIdx_x // T.int64(4):threadIdx_z * T.int64(64) + v0 * T.int64(16) + threadIdx_x // T.int64(4) + T.int64(9), threadIdx_y * T.int64(64) + v1 * T.int64(16) + threadIdx_x % T.int64(4) * T.int64(2):threadIdx_y * T.int64(64) + v1 * T.int64(16) + threadIdx_x % T.int64(4) * T.int64(2) + T.int64(10)])
- for local_id in range(T.int64(8)):
- O_shared_dyn[threadIdx_z * T.int64(64) + v0 * T.int64(16) + local_id % T.int64(4) // T.int64(2) * T.int64(8) + threadIdx_x // T.int64(4), T.bitwise_xor(threadIdx_y * T.int64(8) + v1 * T.int64(2) + local_id // T.int64(4), threadIdx_x // T.int64(4)) * T.int64(8) + threadIdx_x % T.int64(4) * T.int64(2) + local_id % T.int64(2)] = O_warp[v0, v1, threadIdx_x, local_id]
- for v0 in range(T.int64(16)):
- for v1 in T.vectorized(T.int64(8)):
- with Ts.sblock("O.dyn"):
- Ts.reads(O_shared_dyn[v0 * T.int64(8) + threadIdx_z * T.int64(4) + threadIdx_y * T.int64(2) + threadIdx_x // T.int64(16), threadIdx_x % T.int64(16) * T.int64(8) + v1])
- Ts.writes(O[v0 * T.int64(8) + threadIdx_z * T.int64(4) + threadIdx_y * T.int64(2) + threadIdx_x // T.int64(16), threadIdx_x % T.int64(16) * T.int64(8) + v1])
- O[v0 * T.int64(8) + threadIdx_z * T.int64(4) + threadIdx_y * T.int64(2) + threadIdx_x // T.int64(16), threadIdx_x % T.int64(16) * T.int64(8) + v1] = T.Cast(O_shared_dyn[v0 * T.int64(8) + threadIdx_z * T.int64(4) + threadIdx_y * T.int64(2) + threadIdx_x // T.int64(16), T.bitwise_xor(threadIdx_x % T.int64(16), threadIdx_z * T.int64(4) + threadIdx_y * T.int64(2) + threadIdx_x // T.int64(16)) * T.int64(8) + v1], "float16")
-
- # fmt: on
- _check_function_transform(before, expected)
+ def expected(A: T.Tensor((8, columns), "float16"), B: T.Tensor((8, columns), "float16")):
+ shared = Ts.sblock_alloc_tensor((8, columns), "float16", scope=scope)
+ for i, j in T.grid(8, columns):
+ with Ts.sblock("load"):
+ vi, vj = Ts.axis.remap("SS", [i, j])
+ Ts.reads(A[vi, vj])
+ Ts.writes(shared[vi, vj])
+ shared[vi, permuted_column(vi, vj)] = A[vi, vj]
+ for i, j in T.grid(8, columns):
+ with Ts.sblock("store"):
+ vi, vj = Ts.axis.remap("SS", [i, j])
+ Ts.reads(shared[vi, vj])
+ Ts.writes(B[vi, vj])
+ B[vi, vj] = shared[vi, permuted_column(vi, vj)]
+
+ result = tvm.s_tir.transform.InjectPermutedLayout()(tvm.IRModule.from_expr(before))["before"]
+ tvm.ir.assert_structural_equal(
+ result.without_attr("global_symbol"), expected.without_attr("global_symbol")
+ )
if __name__ == "__main__":
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_inject_software_pipeline.py b/tests/python/s_tir/transform/test_s_tir_transform_inject_software_pipeline.py
index ddbfd09f4b..8f751e6f8a 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_inject_software_pipeline.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_inject_software_pipeline.py
@@ -21,20 +21,10 @@ import numpy as np
import pytest
import tvm
-import tvm.s_tir.tensor_intrin.cuda
import tvm.s_tir.testing
import tvm.testing
from tvm import te, tirx
from tvm.s_tir.meta_schedule.testing import te_workload
-from tvm.s_tir.tensor_intrin.cuda import (
- LDMATRIX_f16_A_DYN_INTRIN,
- LDMATRIX_f16_B_DYN_INTRIN,
- MMA_f16f16f32_INTRIN,
- MMA_fill_16x16_f32_INTRIN,
- MMA_store_16x16_f32_global_INTRIN,
- shared_16x16_to_ldmatrix_32x8_layout,
-)
-from tvm.s_tir.testing.tir import mma_schedule
from tvm.script import s_tir as Ts
from tvm.script import tirx as T
from tvm.testing import env
@@ -1477,165 +1467,6 @@ def test_three_stage_compute_two_stage_async():
tvm.ir.assert_structural_equal(mod["main"], ref.with_attr("global_symbol", "main"), True)
-N = K = M = 4096
-
-
-def get_mma_schedule():
- i_factors, j_factors, k_factors = [1, 32, 1, 4, 2], [16, 2, 4, 1, 2], [128, 2, 1]
-
- def index_map(i, j):
- return (
- i // 16,
- j // 16,
- *shared_16x16_to_ldmatrix_32x8_layout(i % 16, j % 16),
- )
-
- workload = te.create_function(
- te_workload.matmul(N, M, K, in_dtype="float16", out_dtype="float32")
- )
-
- return mma_schedule(
- workload,
- 16,
- "float16",
- False,
- i_factors,
- j_factors,
- k_factors,
- index_map,
- index_map,
- index_map,
- LDMATRIX_f16_A_DYN_INTRIN,
- LDMATRIX_f16_B_DYN_INTRIN,
- MMA_f16f16f32_INTRIN,
- MMA_fill_16x16_f32_INTRIN,
- MMA_store_16x16_f32_global_INTRIN,
- "shared.dyn",
- )
-
-
-def build_and_run(sch):
- if tvm.s_tir.testing.is_ampere_or_newer():
- with tvm.ir.transform.PassContext(config={"tirx.use_async_copy": 1}):
- f = tvm.compile(sch.mod["main"], target="cuda")
-
- a_np = np.random.uniform(size=(N, K)).astype("float16")
- b_np = np.random.uniform(size=(K, M)).astype("float16")
- c_np = np.dot(a_np.astype("float32"), b_np.astype("float32"))
-
- def run_and_check():
- dev = tvm.cuda(0)
- a = tvm.runtime.tensor(a_np, dev)
- b = tvm.runtime.tensor(b_np, dev)
- c = tvm.runtime.tensor(np.zeros((N, M), dtype="float32"), dev)
- f(a, b, c)
- tvm.testing.assert_allclose(c.numpy(), c_np, rtol=1e-3)
-
- tvm.testing.run_with_gpu_lock(run_and_check)
-
-
-@pytest.mark.gpu
-@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
-def test_async_pipelined_mma_gemm_simple():
- sch = get_mma_schedule()
-
- k0 = sch.get_loops(sch.get_sblock("C_o_update"))[3]
-
- sch.annotate(k0, ann_key="software_pipeline_stage", ann_val=[0, 0, 3])
- sch.annotate(k0, ann_key="software_pipeline_order", ann_val=[0, 1, 2])
- sch.annotate(k0, ann_key="software_pipeline_async_stages", ann_val=[0])
-
- seq = tvm.ir.transform.Sequential(
- [
- tvm.s_tir.transform.PlanAndUpdateTensorAllocationLocation(),
- tvm.s_tir.transform.ConvertBlocksToOpaque(),
- tvm.s_tir.transform.LowerMatchTensor(),
- tvm.s_tir.transform.InjectSoftwarePipeline(),
- ]
- )
- mod = seq(sch.mod)
-
- pipeline = mod["main"].body[0].block.body[0].body[0].body[0].body[0].block.body[1].block.body
- prologue, body, epilogue = pipeline
-
- commit_group = prologue.block.body[0].body[0].block.body
- assert len(commit_group) == 3
- assert commit_group[-1].value.args[0] == 0
-
- commit_group = body.block.body[0].body[0].block.body
- assert len(commit_group) == 3
- assert commit_group[-1].value.args[0] == 0
-
- assert (
- body.block.body[0].body[1].block.body[0].value.op.same_as(tvm.ir.Op.get("s_tir.async_wait"))
- )
- assert body.block.body[0].body[1].block.body[0].value.args[1] == 3
-
- assert (
- epilogue.block.body[0]
- .body[0]
- .block.body[0]
- .value.op.same_as(tvm.ir.Op.get("s_tir.async_wait"))
- )
- assert str(epilogue.block.body[0].body[0].block.body[0].value.args[1]) == "2 - k_0_0"
-
- build_and_run(sch)
-
-
-@pytest.mark.gpu
-@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
-def test_async_nested_pipeline_mma_gemm_ideal_annotation():
- sch = get_mma_schedule()
-
- k0 = sch.get_loops(sch.get_sblock("C_o_update"))[3]
- k1 = sch.get_loops(sch.get_sblock("C_o_update"))[4]
-
- sch.annotate(k0, ann_key="software_pipeline_stage", ann_val=[0, 0, 2, 3, 3])
- sch.annotate(k0, ann_key="software_pipeline_order", ann_val=[0, 1, 3, 2, 4])
- sch.annotate(k0, ann_key="software_pipeline_async_stages", ann_val=[0])
-
- sch.annotate(k1, ann_key="software_pipeline_stage", ann_val=[0, 0, 1])
- sch.annotate(k1, ann_key="software_pipeline_order", ann_val=[0, 1, 2])
-
- seq = tvm.ir.transform.Sequential(
- [
- tvm.s_tir.transform.PlanAndUpdateTensorAllocationLocation(),
- tvm.s_tir.transform.ConvertBlocksToOpaque(),
- tvm.s_tir.transform.LowerMatchTensor(),
- tvm.s_tir.transform.InjectSoftwarePipeline(),
- ]
- )
- mod = seq(sch.mod)
-
- pipeline = mod["main"].body[0].block.body[0].body[0].body[0].body[0].block.body[1].block.body
- prologue, body, epilogue = pipeline
-
- commit_group = prologue.block.body[0].body[0].block.body
- assert len(commit_group) == 3
- assert commit_group[-1].value.args[0] == 0
-
- assert (
- prologue.block.body[0]
- .body[1]
- .block.body[0]
- .value.op.same_as(tvm.ir.Op.get("s_tir.async_wait"))
- )
- assert prologue.block.body[0].body[1].block.body[0].value.args[1] == 2
-
- commit_group = body.block.body[0].body[0].block.body
- assert len(commit_group) == 3
- assert commit_group[-1].value.args[0] == 0
-
- assert (
- body.block.body[0].body[1].block.body[0].value.op.same_as(tvm.ir.Op.get("s_tir.async_wait"))
- )
- assert body.block.body[0].body[1].block.body[0].value.args[1] == 2
-
- assert str(epilogue.block.body[0].body[0].block.body[0].value.args[1]) == "1 - k_0_0"
-
- build_and_run(sch)
-
-
def test_less_loop_than_num_stage():
@Ts.function
def before(A: T.Tensor((2,), "float32"), E: T.Tensor((2,), "float32")):
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_mma_tensor_layout.py b/tests/python/s_tir/transform/test_s_tir_transform_mma_tensor_layout.py
index a534451680..02007362f2 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_mma_tensor_layout.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_mma_tensor_layout.py
@@ -20,6 +20,8 @@ import pytest
import tvm
import tvm.testing
from tvm import s_tir, tirx
+from tvm.script import s_tir as Ts
+from tvm.script import tirx as T
@pytest.mark.parametrize("scope,shape", [("mma.matrixA", (32, 8)), ("mma.matrixB", (8, 32))])
@@ -37,5 +39,26 @@ def test_explicit_matrix_ab_access_is_rejected(scope, shape, access_kind):
s_tir.transform.TransformMmaTensorLayout()(tvm.IRModule.from_expr(func))
+def test_matrix_c_load_store_lower_without_intrinsic_catalogue():
+ @Ts.function
+ def before(A: T.Tensor((32, 16), "float32"), B: T.Tensor((32, 16), "float32")):
+ fragment = Ts.sblock_alloc_tensor((32, 16), "float32", scope="mma.matrixC")
+ for i, j in T.grid(32, 16):
+ fragment[i, j] = A[i, j]
+ B[i, j] = fragment[i, j]
+
+ @Ts.function
+ def expected(A: T.Tensor((32, 16), "float32"), B: T.Tensor((32, 16), "float32")):
+ fragment = Ts.sblock_alloc_tensor((2, 2, 2, 2), "float32", scope="local")
+ for i, j in T.grid(32, 16):
+ fragment[i // 16, j // 8, i % 16 // 8, j % 2] = A[i, j]
+ B[i, j] = fragment[i // 16, j // 8, i % 16 // 8, j % 2]
+
+ result = s_tir.transform.TransformMmaTensorLayout()(tvm.IRModule.from_expr(before))["before"]
+ tvm.ir.assert_structural_equal(
+ result.without_attr("global_symbol"), expected.without_attr("global_symbol")
+ )
+
+
if __name__ == "__main__":
tvm.testing.main()
diff --git a/tests/python/tirx/base/test_tirx_ptx_ldmatrix.py b/tests/python/tirx/base/test_tirx_ptx_ldmatrix.py
deleted file mode 100644
index 338aeca995..0000000000
--- a/tests/python/tirx/base/test_tirx_ptx_ldmatrix.py
+++ /dev/null
@@ -1,101 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-
-import numpy as np
-import pytest
-
-import tvm
-import tvm.testing
-from tvm.script import tirx as T
-from tvm.testing import env
-
-
-@T.function
-def ptx_ldmatrix(
- A: T.Tensor((16, 16), "float16"), B: T.Tensor((16, 16), "float16"), num: T.int32, trans: T.uint8
-) -> None:
- T.func_attr({"global_symbol": "default_function", "tirx.noalias": True})
- T.launch_thread("blockIdx.x", 1)
- tx = T.launch_thread("threadIdx.x", 32)
- A_shared = T.alloc_tensor(ty_args=[T.Tensor([16, 16], "float16", scope="shared")])
- A_local = T.alloc_tensor(ty_args=[T.Tensor([8], "float16", scope="local")])
- for i in range(8):
- A_shared[i * 2 + tx // 16, tx % 16] = A[i * 2 + tx // 16, tx % 16]
- T.evaluate(
- T.ptx_legacy.ldmatrix(
- trans,
- num,
- ".b16",
- A_local.data,
- 0,
- A_shared.data,
- 16 * (tx % 16) + 8 * (tx // 16),
- ty="float16",
- )
- )
- for k in range(2):
- for j in range(2):
- for i in range(2):
- B[8 * j + tx // 4, 8 * k + tx % 4 * 2 + i] = A_local[4 * k + 2 * j + i]
-
-
-@pytest.mark.gpu
-@pytest.mark.skipif(not env.has_cuda_compute(7, 5), reason="need cuda compute >= 7.5")
-def test_ptx_ldmatrix():
- f = ptx_ldmatrix
- _, _, param_num, param_trans = f.params
-
- for num in [1, 2, 4]:
- for trans in [False, True]:
- mod = tvm.tirx.compile(
- f.specialize({param_num: num, param_trans: trans}), target="cuda"
- )
- A_np = np.random.rand(16, 16).astype("float16")
- A_mask_np = np.zeros_like(A_np)
- if num == 1:
- if trans:
- A_mask_np[:8, :8] = A_np[:8, :8].T
- else:
- A_mask_np[:8, :8] = A_np[:8, :8]
- elif num == 2:
- if trans:
- A_mask_np[:8, :8] = A_np[:8, :8].T
- A_mask_np[8:16, :8] = A_np[8:16, :8].T
- else:
- A_mask_np[:16, :8] = A_np[:16, :8]
- else: # num == 4
- if trans:
- A_mask_np[:8, :8] = A_np[:8, :8].T
- A_mask_np[8:16, :8] = A_np[8:16, :8].T
- A_mask_np[:8, 8:16] = A_np[:8, 8:16].T
- A_mask_np[8:16, 8:16] = A_np[8:16, 8:16].T
- else:
- A_mask_np[:16, :16] = A_np[:16, :16]
- B_np = np.zeros((16, 16)).astype("float16")
-
- def run_and_check():
- dev = tvm.cuda(0)
- A_nd = tvm.runtime.tensor(A_np, device=dev)
- B_nd = tvm.runtime.tensor(B_np, device=dev)
- mod(A_nd, B_nd)
- tvm.testing.assert_allclose(B_nd.numpy(), A_mask_np)
-
- tvm.testing.run_with_gpu_lock(run_and_check)
-
-
-if __name__ == "__main__":
- test_ptx_ldmatrix()
diff --git a/tests/python/tirx/base/test_tirx_ptx_mma.py b/tests/python/tirx/base/test_tirx_ptx_mma.py
deleted file mode 100644
index 103ce7cc7f..0000000000
--- a/tests/python/tirx/base/test_tirx_ptx_mma.py
+++ /dev/null
@@ -1,1264 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-
-import numpy as np
-import pytest
-
-import tvm
-import tvm.testing
-from tvm.script import tirx as T
-from tvm.testing import env
-
-
-@T.function
-def gemm_mma_m8n8k4_row_col_fp64pf64fp64(
- A: T.Tensor([8, 4], dtype="float64"),
- B: T.Tensor([8, 4], dtype="float64"),
- C: T.Tensor([8, 8], dtype="float64"),
-):
- T.func_attr({"global_symbol": "default_function", "tirx.noalias": True})
-
- T.launch_thread("blockIdx.y", 1)
- T.launch_thread("blockIdx.x", 1)
- tx = T.launch_thread("threadIdx.x", 32)
- MultiA = T.alloc_tensor(ty_args=[T.Tensor([1], "float64", scope="local")])
- MultiB = T.alloc_tensor(ty_args=[T.Tensor([1], "float64", scope="local")])
- Accum = T.alloc_tensor(ty_args=[T.Tensor([2], "float64", scope="local")])
- for i in range(2):
- Accum[i] = T.float64(0)
-
- MultiA[0] = A[(tx % 32) // 4, (tx % 32) % 4]
- MultiB[0] = B[(tx % 32) // 4, (tx % 32) % 4]
- T.evaluate(
- T.ptx_legacy.mma(
- "m8n8k4",
- "row",
- "col",
- "float64",
- "float64",
- "float64",
- MultiA.data,
- 0,
- MultiB.data,
- 0,
- Accum.data,
- 0,
- False,
- ty="float64",
- )
- )
- for mma_accum_c_id in range(2):
- C[(tx % 32) // 4, (tx % 32) % 4 * 2 + mma_accum_c_id] = Accum[mma_accum_c_id]
-
-
-@pytest.mark.gpu
-@pytest.mark.skipif(not env.has_cuda_compute(8), reason="need cuda compute >= 8.0")
-def test_gemm_mma_m8n8k4_row_col_fp64pf64fp64():
- mod = tvm.IRModule.from_expr(gemm_mma_m8n8k4_row_col_fp64pf64fp64)
- cuda_mod = tvm.tirx.compile(mod, target="cuda")
-
- A_np = np.random.uniform(-1, 1, [8, 4]).astype("float64")
- B_np = np.random.uniform(-1, 1, [8, 4]).astype("float64")
- C_np = np.zeros([8, 8]).astype("float64")
-
- golden = np.matmul(A_np.astype("float64"), B_np.astype("float64").T)
-
- def run_and_check():
- ctx = tvm.cuda()
- A_tvm = tvm.runtime.tensor(A_np, ctx)
- B_tvm = tvm.runtime.tensor(B_np, ctx)
- C_tvm = tvm.runtime.tensor(C_np, ctx)
- cuda_mod(A_tvm, B_tvm, C_tvm)
- C_numpy = C_tvm.numpy()
- tvm.testing.assert_allclose(golden, C_numpy, atol=1e-3, rtol=1e-3)
-
- tvm.testing.run_with_gpu_lock(run_and_check)
-
-
-@T.function
-def gemm_mma_m8n8k4_row_row_fp16fp16fp16(
- A: T.Tensor([16, 4], dtype="float16"),
- B: T.Tensor([4, 16], dtype="float16"),
- C: T.Tensor([16, 16], dtype="float16"),
-):
- T.func_attr({"global_symbol": "default_function", "tirx.noalias": True})
-
- T.launch_thread("blockIdx.y", 1)
- T.launch_thread("blockIdx.x", 1)
- tx = T.launch_thread("threadIdx.x", 32)
- MultiA = T.alloc_tensor(ty_args=[T.Tensor([4], "float16", scope="local")])
- MultiB = T.alloc_tensor(ty_args=[T.Tensor([4], "float16", scope="local")])
- Accum = T.alloc_tensor(ty_args=[T.Tensor([8], "float16", scope="local")])
- for i in range(8):
- Accum[i] = T.float32(0)
-
- for mma_multi_a_col in T.vectorized(4):
- MultiA[mma_multi_a_col] = A[
- ((tx % 32) % 4) + (4 * (((tx % 32) // 16 + (tx % 32) % 16 // 4 * 2) % 4)),
- mma_multi_a_col,
- ]
- for mma_multi_b_col in T.vectorized(4):
- MultiB[mma_multi_b_col] = B[
- (tx % 32) % 4,
- mma_multi_b_col + (4 * ((tx % 32) // 8)),
- ]
- T.evaluate(
- T.ptx_legacy.mma(
- "m8n8k4",
- "row",
- "row",
- "float16",
- "float16",
- "float16",
- MultiA.data,
- 0,
- MultiB.data,
- 0,
- Accum.data,
- 0,
- False,
- ty="float16",
- )
- )
- for mma_accum_c_id in range(8):
- C[
- ((tx % 32) % 4) + (4 * (((tx % 32) // 16 + (tx % 32) % 16 // 4 * 2) % 4)),
- mma_accum_c_id % 4 + (4 * ((tx % 32) % 16 // 8)) + mma_accum_c_id // 4 * 8,
- ] = Accum[mma_accum_c_id]
-
-
-@pytest.mark.gpu
-@pytest.mark.skipif(not env.has_cuda_compute(7), reason="need cuda compute >= 7.0")
-def test_gemm_mma_m8n8k4_row_row_fp16fp16fp16():
- mod = tvm.IRModule.from_expr(gemm_mma_m8n8k4_row_row_fp16fp16fp16)
- cuda_mod = tvm.tirx.compile(mod, target="cuda")
-
- A_np = np.random.uniform(-1, 1, [16, 4]).astype("float16")
- B_np = np.random.uniform(-1, 1, [4, 16]).astype("float16")
- C_np = np.zeros([16, 16]).astype("float16")
-
- golden = np.matmul(A_np.astype("float16"), B_np.astype("float16"))
-
- def run_and_check():
- ctx = tvm.cuda()
- A_tvm = tvm.runtime.tensor(A_np, ctx)
- B_tvm = tvm.runtime.tensor(B_np, ctx)
- C_tvm = tvm.runtime.tensor(C_np, ctx)
- cuda_mod(A_tvm, B_tvm, C_tvm)
- C_numpy = C_tvm.numpy()
- tvm.testing.assert_allclose(golden, C_numpy, atol=1e-3, rtol=1e-3)
-
- tvm.testing.run_with_gpu_lock(run_and_check)
-
-
-@T.function
-def gemm_mma_m8n8k4_row_row_fp16fp16fp32(
- A: T.Tensor([16, 4], dtype="float16"),
- B: T.Tensor([4, 16], dtype="float16"),
- C: T.Tensor([16, 16], dtype="float32"),
-):
- T.func_attr({"global_symbol": "default_function", "tirx.noalias": True})
-
- T.launch_thread("blockIdx.y", 1)
- T.launch_thread("blockIdx.x", 1)
- tx = T.launch_thread("threadIdx.x", 32)
- MultiA = T.alloc_tensor(ty_args=[T.Tensor([4], "float16", scope="local")])
- MultiB = T.alloc_tensor(ty_args=[T.Tensor([4], "float16", scope="local")])
- Accum = T.alloc_tensor(ty_args=[T.Tensor([8], "float32", scope="local")])
-
- for i in range(8):
- Accum[i] = T.float32(0)
-
- for mma_multi_a_col in T.vectorized(4):
- MultiA[mma_multi_a_col] = A[
- ((tx % 32) % 4) + (4 * (((tx % 32) // 16 + (tx % 32) % 16 // 4 * 2) % 4)),
- mma_multi_a_col,
- ]
- for mma_multi_b_col in T.vectorized(4):
- MultiB[mma_multi_b_col] = B[
- (tx % 32) % 4,
- mma_multi_b_col + (4 * ((tx % 32) // 8)),
- ]
- T.evaluate(
- T.ptx_legacy.mma(
- "m8n8k4",
- "row",
- "row",
- "float16",
- "float16",
- "float32",
- MultiA.data,
- 0,
- MultiB.data,
- 0,
- Accum.data,
- 0,
- False,
- ty="float32",
- )
- )
- for mma_accum_c_id in range(8):
- C[
- ((tx % 32) % 2)
- + ((mma_accum_c_id // 2 % 2) * 2)
- + 4 * ((tx % 32) // 16)
- + ((tx % 32) % 16 // 4) % 2 * 8,
- (tx % 32) % 4 // 2 * 2
- + (tx % 32) % 16 // 8 * 4
- + mma_accum_c_id % 2
- + mma_accum_c_id // 4 * 8,
- ] = Accum[mma_accum_c_id]
-
-
-@pytest.mark.gpu
-@pytest.mark.skipif(not env.has_cuda_compute(7), reason="need cuda compute >= 7.0")
-def test_gemm_mma_m8n8k4_row_row_fp16fp16fp32():
- mod = tvm.IRModule.from_expr(gemm_mma_m8n8k4_row_row_fp16fp16fp32)
- cuda_mod = tvm.tirx.compile(mod, target="cuda")
-
- A_np = np.random.uniform(-1, 1, [16, 4]).astype("float16")
- B_np = np.random.uniform(-1, 1, [4, 16]).astype("float16")
- C_np = np.zeros([16, 16]).astype("float32")
-
- golden = np.matmul(A_np.astype("float32"), B_np.astype("float32"))
-
- def run_and_check():
- ctx = tvm.cuda()
- A_tvm = tvm.runtime.tensor(A_np, ctx)
- B_tvm = tvm.runtime.tensor(B_np, ctx)
- C_tvm = tvm.runtime.tensor(C_np, ctx)
- cuda_mod(A_tvm, B_tvm, C_tvm)
- C_numpy = C_tvm.numpy()
- tvm.testing.assert_allclose(golden, C_numpy, atol=1e-3, rtol=1e-3)
-
- tvm.testing.run_with_gpu_lock(run_and_check)
-
-
-@T.function
-def gemm_mma_m8n8k16_row_col_s8s8s32(
- A: T.Tensor([8, 16], dtype="int8"),
- B: T.Tensor([8, 16], dtype="int8"),
- C: T.Tensor([8, 8], dtype="int32"),
-):
- T.func_attr({"global_symbol": "default_function", "tirx.noalias": True})
-
- T.launch_thread("blockIdx.y", 1)
- T.launch_thread("blockIdx.x", 1)
- tx = T.launch_thread("threadIdx.x", 32)
- MultiA = T.alloc_tensor(ty_args=[T.Tensor([4], "int8", scope="local")])
- MultiB = T.alloc_tensor(ty_args=[T.Tensor([4], "int8", scope="local")])
- Accum = T.alloc_tensor(ty_args=[T.Tensor([2], "int32", scope="local")])
- for i in range(2):
- Accum[i] = T.int32(0)
-
- for mma_multi_a_col in T.vectorized(4):
- MultiA[mma_multi_a_col] = A[(tx % 32) // 4, mma_multi_a_col + (tx % 32) % 4 * 4]
- for mma_multi_b_col in T.vectorized(4):
- MultiB[mma_multi_b_col] = B[(tx % 32) // 4, mma_multi_b_col + (tx % 32) % 4 * 4]
- T.evaluate(
- T.ptx_legacy.mma(
- "m8n8k16",
- "row",
- "col",
- "int8",
- "int8",
- "int32",
- MultiA.data,
- 0,
- MultiB.data,
- 0,
- Accum.data,
- 0,
- False,
- ty="int32",
- )
- )
- for mma_accum_c_id in range(2):
- C[(tx % 32) // 4, (tx % 32) % 4 * 2 + mma_accum_c_id] = Accum[mma_accum_c_id]
-
-
-# This test uses mma instructions that are not available on NVCC 10.1.
-# Failure occurs during the external call to nvcc, when attempting to
-# generate the .fatbin file.
-@pytest.mark.gpu
-@pytest.mark.skipif(not env.has_nvcc_version(11), reason="need nvcc >= 11")
-@pytest.mark.skipif(not env.has_cuda_compute(7, 5), reason="need cuda compute >= 7.5")
-def test_gemm_mma_m8n8k16_row_col_s8s8s32():
- mod = tvm.IRModule.from_expr(gemm_mma_m8n8k16_row_col_s8s8s32)
- cuda_mod = tvm.tirx.compile(mod, target="cuda")
-
- A_np = np.random.uniform(-10, 10, [8, 16]).astype("int8")
- B_np = np.random.uniform(-10, 10, [8, 16]).astype("int8")
- C_np = np.zeros([8, 8]).astype("int32")
-
- golden = np.matmul(A_np.astype("int32"), B_np.astype("int32").T)
-
- def run_and_check():
- ctx = tvm.cuda()
- A_tvm = tvm.runtime.tensor(A_np, ctx)
- B_tvm = tvm.runtime.tensor(B_np, ctx)
- C_tvm = tvm.runtime.tensor(C_np, ctx)
- cuda_mod(A_tvm, B_tvm, C_tvm)
- C_numpy = C_tvm.numpy()
- tvm.testing.assert_allclose(golden, C_numpy, atol=1e-3, rtol=1e-3)
-
- tvm.testing.run_with_gpu_lock(run_and_check)
-
-
-@T.function
-def gemm_mma_m8n8k16_row_col_s8u8s32(
- A: T.Tensor([8, 16], dtype="int8"),
- B: T.Tensor([8, 16], dtype="uint8"),
- C: T.Tensor([8, 8], dtype="int32"),
-):
- T.func_attr({"global_symbol": "default_function", "tirx.noalias": True})
-
- T.launch_thread("blockIdx.y", 1)
- T.launch_thread("blockIdx.x", 1)
- tx = T.launch_thread("threadIdx.x", 32)
- MultiA = T.alloc_tensor(ty_args=[T.Tensor([4], "int8", scope="local")])
- MultiB = T.alloc_tensor(ty_args=[T.Tensor([4], "uint8", scope="local")])
- Accum = T.alloc_tensor(ty_args=[T.Tensor([2], "int32", scope="local")])
- for i in range(2):
- Accum[i] = T.int32(0)
-
- for mma_multi_a_col in T.vectorized(4):
- MultiA[mma_multi_a_col] = A[(tx % 32) // 4, mma_multi_a_col + (tx % 32) % 4 * 4]
- for mma_multi_b_col in T.vectorized(4):
- MultiB[mma_multi_b_col] = B[(tx % 32) // 4, mma_multi_b_col + (tx % 32) % 4 * 4]
- T.evaluate(
- T.ptx_legacy.mma(
- "m8n8k16",
- "row",
- "col",
- "int8",
- "uint8",
- "int32",
- MultiA.data,
- 0,
- MultiB.data,
- 0,
- Accum.data,
- 0,
- False,
- ty="int32",
- )
- )
- for mma_accum_c_id in range(2):
- C[(tx % 32) // 4, (tx % 32) % 4 * 2 + mma_accum_c_id] = Accum[mma_accum_c_id]
-
-
-# This test uses mma instructions that are not available on NVCC 10.1.
-# Failure occurs during the external call to nvcc, when attempting to
-# generate the .fatbin file.
-@pytest.mark.gpu
-@pytest.mark.skipif(not env.has_nvcc_version(11), reason="need nvcc >= 11")
-@pytest.mark.skipif(not env.has_cuda_compute(7, 5), reason="need cuda compute >= 7.5")
-def test_gemm_mma_m8n8k16_row_col_s8u8s32():
- mod = tvm.IRModule.from_expr(gemm_mma_m8n8k16_row_col_s8u8s32)
- cuda_mod = tvm.tirx.compile(mod, target="cuda")
-
- A_np = np.random.uniform(-10, 10, [8, 16]).astype("int8")
- B_np = np.random.uniform(-10, 10, [8, 16]).astype("uint8")
- C_np = np.zeros([8, 8]).astype("int32")
-
- golden = np.matmul(A_np.astype("int32"), B_np.astype("int32").T)
-
- def run_and_check():
- ctx = tvm.cuda()
- A_tvm = tvm.runtime.tensor(A_np, ctx)
- B_tvm = tvm.runtime.tensor(B_np, ctx)
- C_tvm = tvm.runtime.tensor(C_np, ctx)
- cuda_mod(A_tvm, B_tvm, C_tvm)
- C_numpy = C_tvm.numpy()
- tvm.testing.assert_allclose(golden, C_numpy, atol=1e-3, rtol=1e-3)
-
- tvm.testing.run_with_gpu_lock(run_and_check)
-
-
-@T.function
-def gemm_mma_m8n8k32_row_col_s4s4s32(
- A: T.Tensor([8, 32], dtype="int4"),
- B: T.Tensor([8, 32], dtype="int4"),
- C: T.Tensor([8, 8], dtype="int32"),
-):
- T.func_attr({"global_symbol": "default_function", "tirx.noalias": True})
-
- T.launch_thread("blockIdx.y", 1)
- T.launch_thread("blockIdx.x", 1)
- tx = T.launch_thread("threadIdx.x", 32)
- MultiA = T.alloc_tensor(ty_args=[T.Tensor([8], "int4", scope="local")])
- MultiB = T.alloc_tensor(ty_args=[T.Tensor([8], "int4", scope="local")])
- Accum = T.alloc_tensor(ty_args=[T.Tensor([2], "int32", scope="local")])
- for i in range(2):
- Accum[i] = T.int32(0)
-
- for mma_multi_a_col in T.vectorized(8):
- MultiA[mma_multi_a_col] = A[(tx % 32) // 4, mma_multi_a_col + (tx % 32) % 4 * 8]
- for mma_multi_b_col in T.vectorized(8):
- MultiB[mma_multi_b_col] = B[(tx % 32) // 4, mma_multi_b_col + (tx % 32) % 4 * 8]
- T.evaluate(
- T.ptx_legacy.mma(
- "m8n8k32",
- "row",
- "col",
- "int4",
- "int4",
- "int32",
- MultiA.data,
- 0,
- MultiB.data,
- 0,
- Accum.data,
- 0,
- False,
- ty="int32",
- )
- )
- for mma_accum_c_id in range(2):
- C[(tx % 32) // 4, (tx % 32) % 4 * 2 + mma_accum_c_id] = Accum[mma_accum_c_id]
-
-
-# This test uses mma instructions that are not available on NVCC 10.1.
-# Failure occurs during the external call to nvcc, when attempting to
-# generate the .fatbin file.
-@pytest.mark.gpu
-@pytest.mark.skipif(not env.has_nvcc_version(11), reason="need nvcc >= 11")
-@pytest.mark.skipif(not env.has_cuda_compute(7, 5), reason="need cuda compute >= 7.5")
-def test_gemm_mma_m8n8k32_row_col_s4s4s32():
- mod = tvm.IRModule.from_expr(gemm_mma_m8n8k32_row_col_s4s4s32)
- cuda_mod = tvm.tirx.compile(mod, target="cuda")
-
- def run_and_check():
- ctx = tvm.cuda()
- A_tvm = tvm.runtime.empty([8, 32], "int4", ctx)
- B_tvm = tvm.runtime.empty([8, 32], "int4", ctx)
- C_tvm = tvm.runtime.empty([8, 8], "int32", ctx)
- cuda_mod(A_tvm, B_tvm, C_tvm)
-
- tvm.testing.run_with_gpu_lock(run_and_check)
- # Currently the correctness is not checked.
- # TODO: add correctness checking here.
-
-
-@T.function
-def gemm_mma_m8n8k32_row_col_s4u4s32(
- A: T.Tensor([8, 32], dtype="int4"),
- B: T.Tensor([8, 32], dtype="uint4"),
- C: T.Tensor([8, 8], dtype="int32"),
-):
- T.func_attr({"global_symbol": "default_function", "tirx.noalias": True})
-
- T.launch_thread("blockIdx.y", 1)
- T.launch_thread("blockIdx.x", 1)
- tx = T.launch_thread("threadIdx.x", 32)
- MultiA = T.alloc_tensor(ty_args=[T.Tensor([8], "int4", scope="local")])
- MultiB = T.alloc_tensor(ty_args=[T.Tensor([8], "uint4", scope="local")])
- Accum = T.alloc_tensor(ty_args=[T.Tensor([2], "int32", scope="local")])
- for i in range(2):
- Accum[i] = T.int32(0)
-
- for mma_multi_a_col in T.vectorized(8):
- MultiA[mma_multi_a_col] = A[(tx % 32) // 4, mma_multi_a_col + (tx % 32) % 4 * 8]
- for mma_multi_b_col in T.vectorized(8):
- MultiB[mma_multi_b_col] = B[(tx % 32) // 4, mma_multi_b_col + (tx % 32) % 4 * 8]
- T.evaluate(
- T.ptx_legacy.mma(
- "m8n8k32",
- "row",
- "col",
- "int4",
- "uint4",
- "int32",
- MultiA.data,
- 0,
- MultiB.data,
- 0,
- Accum.data,
- 0,
- False,
- ty="int32",
- )
- )
- for mma_accum_c_id in range(2):
- C[(tx % 32) // 4, (tx % 32) % 4 * 2 + mma_accum_c_id] = Accum[mma_accum_c_id]
-
-
-# This test uses mma instructions that are not available on NVCC 10.1.
-# Failure occurs during the external call to nvcc, when attempting to
-# generate the .fatbin file.
-@pytest.mark.gpu
-@pytest.mark.skipif(not env.has_nvcc_version(11), reason="need nvcc >= 11")
-@pytest.mark.skipif(not env.has_cuda_compute(7, 5), reason="need cuda compute >= 7.5")
-def test_gemm_mma_m8n8k32_row_col_s4u4s32():
- mod = tvm.IRModule.from_expr(gemm_mma_m8n8k32_row_col_s4u4s32)
- cuda_mod = tvm.tirx.compile(mod, target="cuda")
-
- def run_and_check():
- ctx = tvm.cuda()
- A_tvm = tvm.runtime.empty([8, 32], "int4", ctx)
- B_tvm = tvm.runtime.empty([8, 32], "uint4", ctx)
- C_tvm = tvm.runtime.empty([8, 8], "int32", ctx)
- cuda_mod(A_tvm, B_tvm, C_tvm)
-
- tvm.testing.run_with_gpu_lock(run_and_check)
- # Currently the correctness is not checked.
- # TODO: add correctness checking here.
-
-
-@T.function
-def gemm_mma_m16n8k8_row_col_fp16fp16fp32(
- A: T.Tensor([16, 8], dtype="float16"),
- B: T.Tensor([8, 8], dtype="float16"),
- C: T.Tensor([16, 8], dtype="float32"),
-):
- T.func_attr({"global_symbol": "default_function", "tirx.noalias": True})
-
- T.launch_thread("blockIdx.y", 1)
- T.launch_thread("blockIdx.x", 1)
- tx = T.launch_thread("threadIdx.x", 32)
- MultiA = T.alloc_tensor(ty_args=[T.Tensor([4], "float16", scope="local")])
- MultiB = T.alloc_tensor(ty_args=[T.Tensor([2], "float16", scope="local")])
- Accum = T.alloc_tensor(ty_args=[T.Tensor([4], "float32", scope="local")])
- for i in range(4):
- Accum[i] = T.float32(0)
-
- for mma_multi_a_col in T.vectorized(4):
- MultiA[mma_multi_a_col] = A[
- (tx % 32) // 4 + mma_multi_a_col // 2 * 8, (tx % 32) % 4 * 2 + mma_multi_a_col % 2
- ]
- for mma_multi_b_col in T.vectorized(4):
- MultiB[mma_multi_b_col] = B[
- (tx % 32) // 4 + mma_multi_b_col // 2 * 8, (tx % 32) % 4 * 2 + mma_multi_b_col % 2
- ]
- T.evaluate(
- T.ptx_legacy.mma(
- "m16n8k8",
- "row",
- "col",
- "float16",
- "float16",
- "float32",
- MultiA.data,
- 0,
- MultiB.data,
- 0,
- Accum.data,
- 0,
- False,
- ty="float32",
- )
- )
- for mma_accum_c_id in range(4):
- C[(tx % 32) // 4 + mma_accum_c_id // 2 * 8, (tx % 32) % 4 * 2 + mma_accum_c_id % 2] = Accum[
- mma_accum_c_id
- ]
-
-
-@pytest.mark.gpu
-@pytest.mark.skipif(not env.has_cuda_compute(8), reason="need cuda compute >= 8.0")
-def test_gemm_mma_m16n8k8_row_col_fp16fp16fp32():
- mod = tvm.IRModule.from_expr(gemm_mma_m16n8k8_row_col_fp16fp16fp32)
- cuda_mod = tvm.tirx.compile(mod, target="cuda")
-
- A_np = np.random.uniform(-1, 1, [16, 8]).astype("float16")
- B_np = np.random.uniform(-1, 1, [8, 8]).astype("float16")
- C_np = np.zeros([16, 8]).astype("float32")
-
- golden = np.matmul(A_np.astype("float32"), B_np.astype("float32").T)
-
- def run_and_check():
- ctx = tvm.cuda()
- A_tvm = tvm.runtime.tensor(A_np, ctx)
- B_tvm = tvm.runtime.tensor(B_np, ctx)
- C_tvm = tvm.runtime.tensor(C_np, ctx)
- cuda_mod(A_tvm, B_tvm, C_tvm)
- C_numpy = C_tvm.numpy()
- tvm.testing.assert_allclose(golden, C_numpy, atol=1e-3, rtol=1e-3)
-
- tvm.testing.run_with_gpu_lock(run_and_check)
-
-
-@T.function
-def gemm_mma_m16n8k16_row_col_fp16fp16fp16(
- A: T.Tensor([16, 16], dtype="float16"),
- B: T.Tensor([8, 16], dtype="float16"),
- C: T.Tensor([16, 8], dtype="float16"),
-):
- T.func_attr({"global_symbol": "default_function", "tirx.noalias": True})
-
- T.launch_thread("blockIdx.y", 1)
- T.launch_thread("blockIdx.x", 1)
- tx = T.launch_thread("threadIdx.x", 32)
- MultiA = T.alloc_tensor(ty_args=[T.Tensor([8], "float16", scope="local")])
- MultiB = T.alloc_tensor(ty_args=[T.Tensor([4], "float16", scope="local")])
- Accum = T.alloc_tensor(ty_args=[T.Tensor([4], "float16", scope="local")])
- for i in range(4):
- Accum[i] = T.float32(0)
-
- for mma_multi_a_col in range(8):
- MultiA[mma_multi_a_col] = A[
- (tx % 32) // 4 + mma_multi_a_col % 4 // 2 * 8,
- (tx % 32) % 4 * 2 + mma_multi_a_col % 2 + mma_multi_a_col // 4 * 8,
- ]
- for mma_multi_b_col in T.vectorized(4):
- MultiB[mma_multi_b_col] = B[
- (tx % 32) // 4,
- (tx % 32) % 4 * 2 + mma_multi_b_col % 2 + mma_multi_b_col // 2 * 8,
- ]
- T.evaluate(
- T.ptx_legacy.mma(
- "m16n8k16",
- "row",
- "col",
- "float16",
- "float16",
- "float16",
- MultiA.data,
- 0,
- MultiB.data,
- 0,
- Accum.data,
- 0,
- False,
- ty="float16",
- )
- )
- for mma_accum_c_id in range(4):
- C[
- (tx % 32) // 4 + mma_accum_c_id // 2 * 8,
- (tx % 32) % 4 * 2 + mma_accum_c_id % 2,
- ] = Accum[mma_accum_c_id]
-
-
-@pytest.mark.gpu
-@pytest.mark.skipif(not env.has_cuda_compute(8), reason="need cuda compute >= 8.0")
-def test_gemm_mma_m16n8k16_row_col_fp16fp16fp16():
- mod = tvm.IRModule.from_expr(gemm_mma_m16n8k16_row_col_fp16fp16fp16)
- cuda_mod = tvm.tirx.compile(mod, target="cuda")
-
- A_np = np.random.uniform(-1, 1, [16, 16]).astype("float16")
- B_np = np.random.uniform(-1, 1, [8, 16]).astype("float16")
- C_np = np.zeros([16, 8]).astype("float16")
-
- golden = np.matmul(A_np.astype("float16"), B_np.astype("float16").T)
-
- def run_and_check():
- ctx = tvm.cuda()
- A_tvm = tvm.runtime.tensor(A_np, ctx)
- B_tvm = tvm.runtime.tensor(B_np, ctx)
- C_tvm = tvm.runtime.tensor(C_np, ctx)
- cuda_mod(A_tvm, B_tvm, C_tvm)
- C_numpy = C_tvm.numpy()
- tvm.testing.assert_allclose(golden, C_numpy, atol=1e-3, rtol=1e-3)
-
- tvm.testing.run_with_gpu_lock(run_and_check)
-
-
-@T.function
-def gemm_mma_m16n8k16_row_col_fp16fp16fp32(
- A: T.Tensor([16, 16], dtype="float16"),
- B: T.Tensor([8, 16], dtype="float16"),
- C: T.Tensor([16, 8], dtype="float32"),
-):
- T.func_attr({"global_symbol": "default_function", "tirx.noalias": True})
-
- T.launch_thread("blockIdx.y", 1)
- T.launch_thread("blockIdx.x", 1)
- tx = T.launch_thread("threadIdx.x", 32)
- MultiA = T.alloc_tensor(ty_args=[T.Tensor([8], "float16", scope="local")])
- MultiB = T.alloc_tensor(ty_args=[T.Tensor([4], "float16", scope="local")])
- Accum = T.alloc_tensor(ty_args=[T.Tensor([4], "float32", scope="local")])
- for i in range(4):
- Accum[i] = T.float32(0)
-
- for mma_multi_a_col in range(8):
- MultiA[mma_multi_a_col] = A[
- (tx % 32) // 4 + mma_multi_a_col % 4 // 2 * 8,
- (tx % 32) % 4 * 2 + mma_multi_a_col % 2 + mma_multi_a_col // 4 * 8,
- ]
- for mma_multi_b_col in T.vectorized(4):
- MultiB[mma_multi_b_col] = B[
- (tx % 32) // 4,
- (tx % 32) % 4 * 2 + mma_multi_b_col % 2 + mma_multi_b_col // 2 * 8,
- ]
- T.evaluate(
- T.ptx_legacy.mma(
- "m16n8k16",
- "row",
- "col",
- "float16",
- "float16",
- "float32",
- MultiA.data,
- 0,
- MultiB.data,
- 0,
- Accum.data,
- 0,
- False,
- ty="float32",
- )
- )
- for mma_accum_c_id in range(4):
- C[
- (tx % 32) // 4 + mma_accum_c_id // 2 * 8,
- (tx % 32) % 4 * 2 + mma_accum_c_id % 2,
- ] = Accum[mma_accum_c_id]
-
-
-@pytest.mark.gpu
-@pytest.mark.skipif(not env.has_cuda_compute(8), reason="need cuda compute >= 8.0")
-def test_gemm_mma_m16n8k16_row_col_fp16fp16fp32():
- mod = tvm.IRModule.from_expr(gemm_mma_m16n8k16_row_col_fp16fp16fp32)
- cuda_mod = tvm.tirx.compile(mod, target="cuda")
-
- A_np = np.random.uniform(-1, 1, [16, 16]).astype("float16")
- B_np = np.random.uniform(-1, 1, [8, 16]).astype("float16")
- C_np = np.zeros([16, 8]).astype("float32")
-
- golden = np.matmul(A_np.astype("float32"), B_np.astype("float32").T)
-
- def run_and_check():
- ctx = tvm.cuda()
- A_tvm = tvm.runtime.tensor(A_np, ctx)
- B_tvm = tvm.runtime.tensor(B_np, ctx)
- C_tvm = tvm.runtime.tensor(C_np, ctx)
- cuda_mod(A_tvm, B_tvm, C_tvm)
- C_numpy = C_tvm.numpy()
- tvm.testing.assert_allclose(golden, C_numpy, atol=1e-3, rtol=1e-3)
-
- tvm.testing.run_with_gpu_lock(run_and_check)
-
-
-@T.function
-def gemm_mma_m16n8k16_row_col_s8s8s32(
- A: T.Tensor([16, 16], dtype="int8"),
- B: T.Tensor([8, 16], dtype="int8"),
- C: T.Tensor([16, 8], dtype="int32"),
-):
- T.func_attr({"global_symbol": "default_function", "tirx.noalias": True})
-
- T.launch_thread("blockIdx.y", 1)
- T.launch_thread("blockIdx.x", 1)
- tx = T.launch_thread("threadIdx.x", 32)
- MultiA = T.alloc_tensor(ty_args=[T.Tensor([8], "int8", scope="local")])
- MultiB = T.alloc_tensor(ty_args=[T.Tensor([4], "int8", scope="local")])
- Accum = T.alloc_tensor(ty_args=[T.Tensor([4], "int32", scope="local")])
- for i in range(4):
- Accum[i] = T.int32(0)
-
- for mma_multi_a_col in range(8):
- MultiA[mma_multi_a_col] = A[
- (tx % 32) // 4 + mma_multi_a_col // 4 * 8,
- (tx % 32) % 4 * 4 + mma_multi_a_col % 4,
- ]
- for mma_multi_b_col in T.vectorized(4):
- MultiB[mma_multi_b_col] = B[
- (tx % 32) // 4,
- (tx % 32) % 4 * 4 + mma_multi_b_col,
- ]
- T.evaluate(
- T.ptx_legacy.mma(
- "m16n8k16",
- "row",
- "col",
- "int8",
- "int8",
- "int32",
- MultiA.data,
- 0,
- MultiB.data,
- 0,
- Accum.data,
- 0,
- False,
- ty="int32",
- )
- )
- for mma_accum_c_id in range(4):
- C[
- (tx % 32) // 4 + mma_accum_c_id // 2 * 8,
- (tx % 32) % 4 * 2 + mma_accum_c_id % 2,
- ] = Accum[mma_accum_c_id]
-
-
-@pytest.mark.gpu
-@pytest.mark.skipif(not env.has_cuda_compute(8), reason="need cuda compute >= 8.0")
-def test_gemm_mma_m16n8k16_row_col_s8s8s32():
- mod = tvm.IRModule.from_expr(gemm_mma_m16n8k16_row_col_s8s8s32)
- cuda_mod = tvm.tirx.compile(mod, target="cuda")
-
- A_np = np.random.uniform(-10, 10, [16, 16]).astype("int8")
- B_np = np.random.uniform(-10, 10, [8, 16]).astype("int8")
- C_np = np.zeros([16, 8]).astype("int32")
-
- golden = np.matmul(A_np.astype("int32"), B_np.astype("int32").T)
-
- def run_and_check():
- ctx = tvm.cuda()
- A_tvm = tvm.runtime.tensor(A_np, ctx)
- B_tvm = tvm.runtime.tensor(B_np, ctx)
- C_tvm = tvm.runtime.tensor(C_np, ctx)
- cuda_mod(A_tvm, B_tvm, C_tvm)
- C_numpy = C_tvm.numpy()
- tvm.testing.assert_allclose(golden, C_numpy, atol=1e-3, rtol=1e-3)
-
- tvm.testing.run_with_gpu_lock(run_and_check)
-
-
-@T.function
-def gemm_mma_m16n8k16_row_col_s8u8s32(
- A: T.Tensor([16, 16], dtype="int8"),
- B: T.Tensor([8, 16], dtype="uint8"),
- C: T.Tensor([16, 8], dtype="int32"),
-):
- T.func_attr({"global_symbol": "default_function", "tirx.noalias": True})
-
- T.launch_thread("blockIdx.y", 1)
- T.launch_thread("blockIdx.x", 1)
- tx = T.launch_thread("threadIdx.x", 32)
- MultiA = T.alloc_tensor(ty_args=[T.Tensor([8], "int8", scope="local")])
- MultiB = T.alloc_tensor(ty_args=[T.Tensor([4], "uint8", scope="local")])
- Accum = T.alloc_tensor(ty_args=[T.Tensor([4], "int32", scope="local")])
- for i in range(4):
- Accum[i] = T.int32(0)
-
- for mma_multi_a_col in range(8):
- MultiA[mma_multi_a_col] = A[
- (tx % 32) // 4 + mma_multi_a_col // 4 * 8,
- (tx % 32) % 4 * 4 + mma_multi_a_col % 4,
- ]
- for mma_multi_b_col in T.vectorized(4):
- MultiB[mma_multi_b_col] = B[
- (tx % 32) // 4,
- (tx % 32) % 4 * 4 + mma_multi_b_col,
- ]
- T.evaluate(
- T.ptx_legacy.mma(
- "m16n8k16",
- "row",
- "col",
- "int8",
- "uint8",
- "int32",
- MultiA.data,
- 0,
- MultiB.data,
- 0,
- Accum.data,
- 0,
- False,
- ty="int32",
- )
- )
- for mma_accum_c_id in range(4):
- C[
- (tx % 32) // 4 + mma_accum_c_id // 2 * 8,
- (tx % 32) % 4 * 2 + mma_accum_c_id % 2,
- ] = Accum[mma_accum_c_id]
-
-
-@pytest.mark.gpu
-@pytest.mark.skipif(not env.has_cuda_compute(8), reason="need cuda compute >= 8.0")
-def test_gemm_mma_m16n8k16_row_col_s8u8s32():
- mod = tvm.IRModule.from_expr(gemm_mma_m16n8k16_row_col_s8u8s32)
- cuda_mod = tvm.tirx.compile(mod, target="cuda")
-
- A_np = np.random.uniform(-10, 10, [16, 16]).astype("int8")
- B_np = np.random.uniform(-10, 10, [8, 16]).astype("uint8")
- C_np = np.zeros([16, 8]).astype("int32")
-
- golden = np.matmul(A_np.astype("int32"), B_np.astype("int32").T)
-
- def run_and_check():
- ctx = tvm.cuda()
- A_tvm = tvm.runtime.tensor(A_np, ctx)
- B_tvm = tvm.runtime.tensor(B_np, ctx)
- C_tvm = tvm.runtime.tensor(C_np, ctx)
- cuda_mod(A_tvm, B_tvm, C_tvm)
- C_numpy = C_tvm.numpy()
- tvm.testing.assert_allclose(golden, C_numpy, atol=1e-3, rtol=1e-3)
-
- tvm.testing.run_with_gpu_lock(run_and_check)
-
-
-@T.function
-def gemm_mma_m16n8k32_row_col_s8s8s32(
- A: T.Tensor([16, 32], dtype="int8"),
- B: T.Tensor([8, 32], dtype="int8"),
- C: T.Tensor([16, 8], dtype="int32"),
-):
- T.func_attr({"global_symbol": "default_function", "tirx.noalias": True})
-
- T.launch_thread("blockIdx.y", 1)
- T.launch_thread("blockIdx.x", 1)
- tx = T.launch_thread("threadIdx.x", 32)
- MultiA = T.alloc_tensor(ty_args=[T.Tensor([16], "int8", scope="local")])
- MultiB = T.alloc_tensor(ty_args=[T.Tensor([8], "int8", scope="local")])
- Accum = T.alloc_tensor(ty_args=[T.Tensor([4], "int32", scope="local")])
- for i in range(4):
- Accum[i] = T.int32(0)
-
- for mma_multi_a_col in range(16):
- MultiA[mma_multi_a_col] = A[
- (tx % 32) // 4 + mma_multi_a_col % 8 // 4 * 8,
- (tx % 32) % 4 * 4 + mma_multi_a_col % 4 + mma_multi_a_col // 8 * 16,
- ]
- for mma_multi_b_col in range(8):
- MultiB[mma_multi_b_col] = B[
- (tx % 32) // 4,
- (tx % 32) % 4 * 4 + mma_multi_b_col % 4 + mma_multi_b_col // 4 * 16,
- ]
- T.evaluate(
- T.ptx_legacy.mma(
- "m16n8k32",
- "row",
- "col",
- "int8",
- "int8",
- "int32",
- MultiA.data,
- 0,
- MultiB.data,
- 0,
- Accum.data,
- 0,
- False,
- ty="int32",
- )
- )
- for mma_accum_c_id in range(4):
- C[
- (tx % 32) // 4 + mma_accum_c_id // 2 * 8,
- (tx % 32) % 4 * 2 + mma_accum_c_id % 2,
- ] = Accum[mma_accum_c_id]
-
-
-@pytest.mark.gpu
-@pytest.mark.skipif(not env.has_cuda_compute(8), reason="need cuda compute >= 8.0")
-def test_gemm_mma_m16n8k32_row_col_s8s8s32():
- mod = tvm.IRModule.from_expr(gemm_mma_m16n8k32_row_col_s8s8s32)
- cuda_mod = tvm.tirx.compile(mod, target="cuda")
-
- A_np = np.random.uniform(-10, 10, [16, 32]).astype("int8")
- B_np = np.random.uniform(-10, 10, [8, 32]).astype("int8")
- C_np = np.zeros([16, 8]).astype("int32")
-
- golden = np.matmul(A_np.astype("int32"), B_np.astype("int32").T)
-
- def run_and_check():
- ctx = tvm.cuda()
- A_tvm = tvm.runtime.tensor(A_np, ctx)
- B_tvm = tvm.runtime.tensor(B_np, ctx)
- C_tvm = tvm.runtime.tensor(C_np, ctx)
- cuda_mod(A_tvm, B_tvm, C_tvm)
- C_numpy = C_tvm.numpy()
- tvm.testing.assert_allclose(golden, C_numpy, atol=1e-3, rtol=1e-3)
-
- tvm.testing.run_with_gpu_lock(run_and_check)
-
-
-@T.function
-def gemm_mma_m16n8k32_row_col_s8u8s32(
- A: T.Tensor([16, 32], dtype="int8"),
- B: T.Tensor([8, 32], dtype="uint8"),
- C: T.Tensor([16, 8], dtype="int32"),
-):
- T.func_attr({"global_symbol": "default_function", "tirx.noalias": True})
-
- T.launch_thread("blockIdx.y", 1)
- T.launch_thread("blockIdx.x", 1)
- tx = T.launch_thread("threadIdx.x", 32)
- MultiA = T.alloc_tensor(ty_args=[T.Tensor([16], "int8", scope="local")])
- MultiB = T.alloc_tensor(ty_args=[T.Tensor([8], "uint8", scope="local")])
- Accum = T.alloc_tensor(ty_args=[T.Tensor([4], "int32", scope="local")])
- for i in range(4):
- Accum[i] = T.int32(0)
-
- for mma_multi_a_col in range(16):
- MultiA[mma_multi_a_col] = A[
- (tx % 32) // 4 + mma_multi_a_col % 8 // 4 * 8,
- (tx % 32) % 4 * 4 + mma_multi_a_col % 4 + mma_multi_a_col // 8 * 16,
- ]
- for mma_multi_b_col in range(8):
- MultiB[mma_multi_b_col] = B[
- (tx % 32) // 4,
- (tx % 32) % 4 * 4 + mma_multi_b_col % 4 + mma_multi_b_col // 4 * 16,
- ]
- T.evaluate(
- T.ptx_legacy.mma(
- "m16n8k32",
- "row",
- "col",
- "int8",
- "uint8",
- "int32",
- MultiA.data,
- 0,
- MultiB.data,
- 0,
- Accum.data,
- 0,
- False,
- ty="int32",
- )
- )
- for mma_accum_c_id in range(4):
- C[
- (tx % 32) // 4 + mma_accum_c_id // 2 * 8,
- (tx % 32) % 4 * 2 + mma_accum_c_id % 2,
- ] = Accum[mma_accum_c_id]
-
-
-@pytest.mark.gpu
-@pytest.mark.skipif(not env.has_cuda_compute(8), reason="need cuda compute >= 8.0")
-def test_gemm_mma_m16n8k32_row_col_s8u8s32():
- mod = tvm.IRModule.from_expr(gemm_mma_m16n8k32_row_col_s8u8s32)
- cuda_mod = tvm.tirx.compile(mod, target="cuda")
-
- A_np = np.random.uniform(-10, 10, [16, 32]).astype("int8")
- B_np = np.random.uniform(-10, 10, [8, 32]).astype("uint8")
- C_np = np.zeros([16, 8]).astype("int32")
-
- golden = np.matmul(A_np.astype("int32"), B_np.astype("int32").T)
-
- def run_and_check():
- ctx = tvm.cuda()
- A_tvm = tvm.runtime.tensor(A_np, ctx)
- B_tvm = tvm.runtime.tensor(B_np, ctx)
- C_tvm = tvm.runtime.tensor(C_np, ctx)
- cuda_mod(A_tvm, B_tvm, C_tvm)
- C_numpy = C_tvm.numpy()
- tvm.testing.assert_allclose(golden, C_numpy, atol=1e-3, rtol=1e-3)
-
- tvm.testing.run_with_gpu_lock(run_and_check)
-
-
-@T.function
-def gemm_mma_m16n8k64_row_col_s4s4s32(
- A: T.Tensor([16, 64], dtype="int4"),
- B: T.Tensor([8, 64], dtype="int4"),
- C: T.Tensor([16, 8], dtype="int32"),
-):
- T.func_attr({"global_symbol": "default_function", "tirx.noalias": True})
-
- T.launch_thread("blockIdx.y", 1)
- T.launch_thread("blockIdx.x", 1)
- tx = T.launch_thread("threadIdx.x", 32)
- MultiA = T.alloc_tensor(ty_args=[T.Tensor([32], "int4", scope="local")])
- MultiB = T.alloc_tensor(ty_args=[T.Tensor([16], "int4", scope="local")])
- Accum = T.alloc_tensor(ty_args=[T.Tensor([4], "int32", scope="local")])
- for i in range(4):
- Accum[i] = T.int32(0)
-
- for mma_multi_a_col in range(32):
- MultiA[mma_multi_a_col] = A[
- (tx % 32) // 4 + mma_multi_a_col % 16 // 8 * 8,
- (tx % 32) % 4 * 8 + mma_multi_a_col % 8 + mma_multi_a_col // 16 * 32,
- ]
- for mma_multi_b_col in range(16):
- MultiB[mma_multi_b_col] = B[
- (tx % 32) // 4,
- (tx % 32) % 4 * 8 + mma_multi_b_col % 8 + mma_multi_b_col // 8 * 32,
- ]
- T.evaluate(
- T.ptx_legacy.mma(
- "m8n8k32",
- "row",
- "col",
- "int4",
- "int4",
- "int32",
- MultiA.data,
- 0,
- MultiB.data,
- 0,
- Accum.data,
- 0,
- False,
- ty="int32",
- )
- )
- for mma_accum_c_id in range(4):
- C[
- (tx % 32) // 4 + mma_accum_c_id // 2 * 8,
- (tx % 32) % 4 * 2 + mma_accum_c_id % 2,
- ] = Accum[mma_accum_c_id]
-
-
-@pytest.mark.gpu
-@pytest.mark.skipif(not env.has_cuda_compute(8), reason="need cuda compute >= 8.0")
-def test_gemm_mma_m16n8k64_row_col_s4s4s32():
- mod = tvm.IRModule.from_expr(gemm_mma_m16n8k64_row_col_s4s4s32)
- cuda_mod = tvm.tirx.compile(mod, target="cuda")
-
- def run_and_check():
- ctx = tvm.cuda()
- A_tvm = tvm.runtime.empty([16, 64], "int4", ctx)
- B_tvm = tvm.runtime.empty([8, 64], "int4", ctx)
- C_tvm = tvm.runtime.empty([16, 8], "int32", ctx)
- cuda_mod(A_tvm, B_tvm, C_tvm)
-
- tvm.testing.run_with_gpu_lock(run_and_check)
- # Currently the correctness is not checked.
- # TODO: add correctness checking here.
-
-
-@T.function
-def gemm_mma_m16n8k64_row_col_s4u4s32(
- A: T.Tensor([16, 64], dtype="int4"),
- B: T.Tensor([8, 64], dtype="uint4"),
- C: T.Tensor([16, 8], dtype="int32"),
-):
- T.func_attr({"global_symbol": "default_function", "tirx.noalias": True})
-
- T.launch_thread("blockIdx.y", 1)
- T.launch_thread("blockIdx.x", 1)
- tx = T.launch_thread("threadIdx.x", 32)
- MultiA = T.alloc_tensor(ty_args=[T.Tensor([32], "int4", scope="local")])
- MultiB = T.alloc_tensor(ty_args=[T.Tensor([16], "uint4", scope="local")])
- Accum = T.alloc_tensor(ty_args=[T.Tensor([4], "int32", scope="local")])
- for i in range(4):
- Accum[i] = T.int32(0)
-
- for mma_multi_a_col in range(32):
- MultiA[mma_multi_a_col] = A[
- (tx % 32) // 4 + mma_multi_a_col % 16 // 8 * 8,
- (tx % 32) % 4 * 8 + mma_multi_a_col % 8 + mma_multi_a_col // 16 * 32,
- ]
- for mma_multi_b_col in range(16):
- MultiB[mma_multi_b_col] = B[
- (tx % 32) // 4,
- (tx % 32) % 4 * 8 + mma_multi_b_col % 8 + mma_multi_b_col // 8 * 32,
- ]
- T.evaluate(
- T.ptx_legacy.mma(
- "m8n8k32",
- "row",
- "col",
- "int4",
- "uint4",
- "int32",
- MultiA.data,
- 0,
- MultiB.data,
- 0,
- Accum.data,
- 0,
- False,
- ty="int32",
- )
- )
- for mma_accum_c_id in range(4):
- C[
- (tx % 32) // 4 + mma_accum_c_id // 2 * 8,
- (tx % 32) % 4 * 2 + mma_accum_c_id % 2,
- ] = Accum[mma_accum_c_id]
-
-
-@pytest.mark.gpu
-@pytest.mark.skipif(not env.has_cuda_compute(8), reason="need cuda compute >= 8.0")
-def test_gemm_mma_m16n8k64_row_col_s4u4s32():
- mod = tvm.IRModule.from_expr(gemm_mma_m16n8k64_row_col_s4u4s32)
- cuda_mod = tvm.tirx.compile(mod, target="cuda")
-
- def run_and_check():
- ctx = tvm.cuda()
- A_tvm = tvm.runtime.empty([16, 64], "int4", ctx)
- B_tvm = tvm.runtime.empty([8, 64], "uint4", ctx)
- C_tvm = tvm.runtime.empty([16, 8], "int32", ctx)
- cuda_mod(A_tvm, B_tvm, C_tvm)
-
- tvm.testing.run_with_gpu_lock(run_and_check)
- # Currently the correctness is not checked.
- # TODO: add correctness checking here.
-
-
-@T.function
-def gemm_mma_m16n8k256_row_col_b1b1s32(
- A: T.Tensor([16, 256], dtype="int1"),
- B: T.Tensor([8, 256], dtype="int1"),
- C: T.Tensor([16, 8], dtype="int32"),
-):
- T.func_attr({"global_symbol": "default_function", "tirx.noalias": True})
-
- T.launch_thread("blockIdx.y", 1)
- T.launch_thread("blockIdx.x", 1)
- tx = T.launch_thread("threadIdx.x", 32)
- MultiA = T.alloc_tensor(ty_args=[T.Tensor([128], "int1", scope="local")])
- MultiB = T.alloc_tensor(ty_args=[T.Tensor([64], "int1", scope="local")])
- Accum = T.alloc_tensor(ty_args=[T.Tensor([4], "int32", scope="local")])
- for i in range(4):
- Accum[i] = T.int32(0)
-
- for mma_multi_a_col in range(128):
- MultiA[mma_multi_a_col] = A[
- (tx % 32) // 4 + mma_multi_a_col % 64 // 32 * 8,
- (tx % 32) % 4 * 32 + mma_multi_a_col % 32 + mma_multi_a_col // 64 * 128,
- ]
- for mma_multi_b_col in range(16):
- MultiB[mma_multi_b_col] = B[
- (tx % 32) // 4,
- (tx % 32) % 4 * 32 + mma_multi_b_col % 32 + mma_multi_b_col // 32 * 128,
- ]
- T.evaluate(
- T.ptx_legacy.mma(
- "m16n8k256",
- "row",
- "col",
- "int1",
- "int1",
- "int32",
- MultiA.data,
- 0,
- MultiB.data,
- 0,
- Accum.data,
- 0,
- False,
- "xor",
- ty="int32",
- )
- )
- for mma_accum_c_id in range(4):
- C[
- (tx % 32) // 4 + mma_accum_c_id // 2 * 8,
- (tx % 32) % 4 * 2 + mma_accum_c_id % 2,
- ] = Accum[mma_accum_c_id]
-
-
-@pytest.mark.gpu
-@pytest.mark.skipif(not env.has_cuda_compute(8), reason="need cuda compute >= 8.0")
-def test_gemm_mma_m16n8k256_row_col_b1b1s32():
- mod = tvm.IRModule.from_expr(gemm_mma_m16n8k256_row_col_b1b1s32)
- cuda_mod = tvm.tirx.compile(mod, target="cuda")
-
- def run_and_check():
- ctx = tvm.cuda()
- A_tvm = tvm.runtime.empty([16, 256], "int1", ctx)
- B_tvm = tvm.runtime.empty([8, 256], "int1", ctx)
- C_tvm = tvm.runtime.empty([16, 8], "int32", ctx)
- cuda_mod(A_tvm, B_tvm, C_tvm)
-
- tvm.testing.run_with_gpu_lock(run_and_check)
- # Currently the correctness is not checked.
- # TODO: add correctness checking here.
-
-
-if __name__ == "__main__":
- tvm.testing.main()
diff --git a/tests/python/tirx/codegen/test_codegen_cuda.py b/tests/python/tirx/codegen/test_codegen_cuda.py
index 4640cbb3c3..9f45ad360b 100644
--- a/tests/python/tirx/codegen/test_codegen_cuda.py
+++ b/tests/python/tirx/codegen/test_codegen_cuda.py
@@ -152,6 +152,60 @@ def test_cuda_module_destructor_preserves_current_device():
torch.cuda.set_device(original_device)
+@pytest.mark.parametrize(
+ "op_name", ["ptx_legacy.mma", "ptx_legacy.ldmatrix", "cuda.mma_store", "cuda.mma_fill"]
+)
+def test_legacy_tensorcore_ops_require_downstream_codegen(op_name):
+ from tvm.backend.cuda import op as cuda_op
+
+ data = tvm.tirx.Var("data", tvm.ir.PtrType("float16", "global"))
+ calls = {
+ "ptx_legacy.mma": lambda: cuda_op.ptx_legacy_mma(
+ "m16n8k16",
+ "row",
+ "col",
+ "float16",
+ "float16",
+ "float32",
+ data,
+ 0,
+ data,
+ 0,
+ data,
+ 0,
+ False,
+ ),
+ "ptx_legacy.ldmatrix": lambda: cuda_op.ptx_legacy_ldmatrix(
+ False, 4, ".b16", data, 0, data, 0, ty="float16"
+ ),
+ "cuda.mma_store": lambda: cuda_op.mma_store("float16", 16, 16, data, data, 0, 16),
+ "cuda.mma_fill": lambda: cuda_op.mma_fill("float16", 8, data, 0),
+ }
+ call = calls[op_name]()
+ assert call.op.name == "tirx." + op_name
+ func = tvm.tirx.Function([data], tvm.ir.Evaluate(call)).with_attr("global_symbol", "main")
+ target = tvm.target.Target({"kind": "cuda", "arch": "sm_80"})
+ with pytest.raises(
+ tvm.error.InternalError, match="Unresolved call.*" + re.escape(call.op.name)
+ ):
+ tvm.get_global_func("target.build.cuda")(tvm.IRModule.from_expr(func), target)
+
+
+@pytest.mark.parametrize("target_kind", ["cuda", "opencl"])
+def test_warp_storage_requires_explicit_local_layout(target_kind):
+ @T.function
+ def main(data: T.handle("float32")):
+ T.func_attr({"global_symbol": "main", "calling_conv": 2})
+ output = T.decl_tensor(data, ty_args=[T.Tensor((1,), "float32")])
+ fragment = T.alloc_tensor(ty_args=[T.Tensor((32,), "float32", scope="warp")])
+ fragment[0] = 1.0
+ output[0] = fragment[0]
+
+ target = tvm.target.Target(target_kind)
+ with pytest.raises(tvm.error.InternalError, match="warp storage is unsupported"):
+ tvm.get_global_func("target.build." + target_kind)(tvm.IRModule.from_expr(main), target)
+
+
def test_typed_buffer_bindings(monkeypatch):
@T.function
def main(data: T.handle("float32"), output: T.handle):