chore: import upstream snapshot with attribution

python/tvm/s_tir/tensor_intrin/metal.py

@@ -0,0 +1,351 @@
+# 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
+"""Intrinsics for tensorization on Apple GPU."""
+
+from typing import Literal
+
+from tvm.script import tirx as T
+from tvm.tirx import Buffer, Expr, PrimFunc, TensorIntrin
+
+######## simdgroup matrix intrinsics ########
+
+
+def get_simdgroup_index(buffer: Buffer, stride: Expr, col: int, row: int):
+    """Compute simdgroup index using elem_offset of the buffer"""
+
+    # NOTE: Need further check the usage between `col`` and `row`
+    # Currently, Metal only supports 8x8, which means the values of `col` and `row` are the same
+    frag_index_m = buffer.elem_offset // stride // col
+    frag_index_n = buffer.elem_offset % stride // row
+
+    num_fragments_per_row = stride // row
+    return frag_index_m * num_fragments_per_row + frag_index_n
+
+
+def get_make_filled_simdgroup_matrix_intrin(
+    dtype: str, col: int = 8, row: int = 8
+) -> tuple[PrimFunc, PrimFunc]:
+    @T.prim_func(s_tir=True)
+    def desc(a: T.handle) -> None:
+        A = T.match_buffer(a, (col, row), dtype, scope="metal.simdgroup", offset_factor=1)
+        with T.sblock("root"):
+            T.reads()
+            T.writes(A[0:col, 0:row])
+            for i, j in T.grid(col, row):
+                with T.sblock("init"):
+                    vi, vj = T.axis.remap("SS", [i, j])
+                    A[vi, vj] = T.float32(0)
+
+    @T.prim_func(s_tir=True)
+    def impl(a: T.handle) -> None:
+        d0, d1 = T.int32(), T.int32()
+        A = T.match_buffer(
+            a, (col, row), dtype, scope="metal.simdgroup", strides=[d1, d0], offset_factor=1
+        )
+        with T.sblock("root"):
+            T.reads()
+            T.writes(A[0:col, 0:row])
+            T.metal.make_filled_simdgroup_matrix(
+                A.data,
+                index=get_simdgroup_index(A, d1, col, row),
+                value=T.float32(0),
+                col=col,
+                row=row,
+            )
+
+    return desc, impl
+
+
+def get_simdgroup_load_intrin(
+    dtype: str,
+    scope: Literal["global", "shared"],
+    col: int = 8,
+    row: int = 8,
+    transpose_matrix: bool = False,
+) -> tuple[PrimFunc, PrimFunc]:
+    align = col * row
+
+    @T.prim_func(s_tir=True)
+    def desc(a: T.handle, c: T.handle) -> None:
+        A = T.match_buffer(a, (col, row), dtype, align=align, scope=scope, offset_factor=1)
+        C = T.match_buffer(
+            c, (col, row), dtype, align=align, scope="metal.simdgroup", offset_factor=1
+        )
+        with T.sblock("root"):
+            T.reads(A[0:col, 0:row])
+            T.writes(C[0:col, 0:row])
+            for i, j in T.grid(col, row):
+                with T.sblock("load"):
+                    vii, vjj = T.axis.remap("SS", [i, j])
+                    if transpose_matrix:
+                        # C[vii, vjj] = A[vjj, vii]
+                        C[vjj, vii] = A[vii, vjj]
+                    else:
+                        C[vii, vjj] = A[vii, vjj]
+
+    @T.prim_func(s_tir=True)
+    def impl(a: T.handle, c: T.handle) -> None:
+        s0, s1, d0, d1 = T.int32(), T.int32(), T.int32(), T.int32()
+        A = T.match_buffer(
+            a,
+            (col, row),
+            dtype,
+            align=align,
+            scope=scope,
+            strides=[s1, s0],
+            offset_factor=1,
+        )
+        C = T.match_buffer(
+            c,
+            (col, row),
+            dtype,
+            align=align,
+            scope="metal.simdgroup",
+            strides=[d1, d0],
+            offset_factor=1,
+        )
+        with T.sblock("root"):
+            T.reads(A[0:col, 0:row])
+            T.writes(C[0:col, 0:row])
+            T.metal.simdgroup_load(
+                C.data,
+                index=get_simdgroup_index(C, d1, col, row),
+                ptr=A.access_ptr("r", ptr_type=dtype),
+                stride=s1,
+                col=col,
+                row=row,
+                transpose_matrix=transpose_matrix,
+            )
+
+    return desc, impl
+
+
+def get_simdgroup_store_intrin(
+    dtype: str,
+    scope: Literal["global", "shared"],
+    col: int = 8,
+    row: int = 8,
+    transpose_matrix: bool = False,
+) -> tuple[PrimFunc, PrimFunc]:
+    align = col * row
+
+    @T.prim_func(s_tir=True)
+    def desc(a: T.handle, c: T.handle) -> None:
+        A = T.match_buffer(
+            a, (col, row), dtype, align=align, scope="metal.simdgroup", offset_factor=1
+        )
+        C = T.match_buffer(c, (col, row), dtype, align=align, scope=scope, offset_factor=1)
+        with T.sblock("root"):
+            T.reads(A[0:col, 0:row])
+            T.writes(C[0:col, 0:row])
+            for i, j in T.grid(col, row):
+                with T.sblock("store"):
+                    vii, vjj = T.axis.remap("SS", [i, j])
+                    if transpose_matrix:
+                        C[vjj, vii] = A[vii, vjj]
+                    else:
+                        C[vii, vjj] = A[vii, vjj]
+
+    @T.prim_func(s_tir=True)
+    def impl(a: T.handle, c: T.handle) -> None:
+        s0, s1, d0, d1 = T.int32(), T.int32(), T.int32(), T.int32()
+        A = T.match_buffer(
+            a,
+            (col, row),
+            dtype,
+            align=align,
+            scope="metal.simdgroup",
+            strides=[s1, s0],
+            offset_factor=1,
+        )
+        C = T.match_buffer(
+            c, (col, row), dtype, align=align, scope=scope, strides=[d1, d0], offset_factor=1
+        )
+        with T.sblock("root"):
+            T.reads(A[0:col, 0:row])
+            T.writes(C[0:col, 0:row])
+            T.metal.simdgroup_store(
+                A.data,
+                index=get_simdgroup_index(A, s1, col, row),
+                ptr=C.access_ptr("w", ptr_type=dtype),
+                stride=d1,
+                col=col,
+                row=row,
+                transpose_matrix=transpose_matrix,
+            )
+
+    return desc, impl
+
+
+def get_simdgroup_multiply_accumulate_intrin(
+    m_dim: int, n_dim: int, k_dim: int, dtype: str
+) -> tuple[PrimFunc, PrimFunc]:
+    @T.prim_func(s_tir=True)
+    def desc(a: T.handle, b: T.handle, c: T.handle) -> None:
+        A = T.match_buffer(a, (m_dim, k_dim), dtype, scope="metal.simdgroup", offset_factor=1)
+        B = T.match_buffer(b, (k_dim, n_dim), dtype, scope="metal.simdgroup", offset_factor=1)
+        C = T.match_buffer(c, (m_dim, n_dim), dtype, scope="metal.simdgroup", offset_factor=1)
+        with T.sblock("root"):
+            T.reads(C[0:m_dim, 0:n_dim], A[0:m_dim, 0:k_dim], B[0:k_dim, 0:n_dim])
+            T.writes(C[0:m_dim, 0:n_dim])
+            for i, j, k in T.grid(m_dim, n_dim, k_dim):
+                with T.sblock(""):
+                    vii, vjj, vkk = T.axis.remap("SSR", [i, j, k])
+                    C[vii, vjj] += A[vii, vkk] * B[vkk, vjj]
+
+    @T.prim_func(s_tir=True)
+    def impl(a: T.handle, b: T.handle, c: T.handle) -> None:
+        a0, a1, b0, b1, c0, c1 = T.int32(), T.int32(), T.int32(), T.int32(), T.int32(), T.int32()
+        A = T.match_buffer(
+            a, (m_dim, k_dim), dtype, scope="metal.simdgroup", strides=[a1, a0], offset_factor=1
+        )
+        B = T.match_buffer(
+            b, (k_dim, n_dim), dtype, scope="metal.simdgroup", strides=[b1, b0], offset_factor=1
+        )
+        C = T.match_buffer(
+            c, (m_dim, n_dim), dtype, scope="metal.simdgroup", strides=[c1, c0], offset_factor=1
+        )
+        with T.sblock("root"):
+            T.reads(C[0:m_dim, 0:n_dim], A[0:m_dim, 0:k_dim], B[0:k_dim, 0:n_dim])
+            T.writes(C[0:m_dim, 0:n_dim])
+            T.metal.simdgroup_multiply_accumulate(
+                C.data,
+                get_simdgroup_index(C, c1, m_dim, n_dim),
+                A.data,
+                get_simdgroup_index(A, a1, m_dim, k_dim),
+                B.data,
+                get_simdgroup_index(B, b1, k_dim, n_dim),
+                C.data,
+                get_simdgroup_index(C, c1, m_dim, n_dim),
+            )
+
+    return desc, impl
+
+
+# Make filled simdgroup matrix intrinsics
+
+SIMDGROUP_MAKE_FILLED_8x8x8_f16_INTRIN = "simdgroup_make_filled_8x8x8_f16"
+TensorIntrin.register(
+    SIMDGROUP_MAKE_FILLED_8x8x8_f16_INTRIN,
+    *get_make_filled_simdgroup_matrix_intrin("float16", 8, 8),
+)
+
+SIMDGROUP_FILLED_8x8x8_f32_INTRIN = "simdgroup_fill_8x8x8_f32"
+TensorIntrin.register(
+    SIMDGROUP_FILLED_8x8x8_f32_INTRIN, *get_make_filled_simdgroup_matrix_intrin("float32", 8, 8)
+)
+
+SIMDGROUP_FILLED_8x8x8_bf16_INTRIN = "simdgroup_fill_8x8x8_bf16"
+TensorIntrin.register(
+    SIMDGROUP_FILLED_8x8x8_bf16_INTRIN, *get_make_filled_simdgroup_matrix_intrin("bfloat16", 8, 8)
+)
+
+# Load intrinsics
+
+SIMDGROUP_LOAD_8x8x8_f16_SHARED_INTRIN = "simdgroup_load_8x8x8_f16_shared"
+TensorIntrin.register(
+    SIMDGROUP_LOAD_8x8x8_f16_SHARED_INTRIN,
+    *get_simdgroup_load_intrin("float16", "shared", 8, 8, False),
+)
+
+SIMDGROUP_LOAD_8x8x8_f16_SHARED_TRANS_INTRIN = "simdgroup_load_8x8x8_f16_shared_trans"
+TensorIntrin.register(
+    SIMDGROUP_LOAD_8x8x8_f16_SHARED_TRANS_INTRIN,
+    *get_simdgroup_load_intrin("float16", "shared", 8, 8, True),
+)
+
+# Store intrinsics
+
+SIMDGROUP_STORE_8x8x8_f16_GLOBAL_INTRIN = "simdgroup_store_8x8x8_f16_global"
+TensorIntrin.register(
+    SIMDGROUP_STORE_8x8x8_f16_GLOBAL_INTRIN,
+    *get_simdgroup_store_intrin("float16", "global", 8, 8, False),
+)
+
+SIMDGROUP_STORE_8x8x8_f16_SHARED_INTRIN = "simdgroup_store_8x8x8_f16_shared"
+TensorIntrin.register(
+    SIMDGROUP_STORE_8x8x8_f16_SHARED_INTRIN,
+    *get_simdgroup_store_intrin("float16", "shared", 8, 8, False),
+)
+# Multiply accumulate intrinsics
+
+SIMDGROUP_MULTI_ACC_8x8x8_f16_INTRIN = "simdgroup_multiply_accumulate_8x8x8_f16"
+TensorIntrin.register(
+    SIMDGROUP_MULTI_ACC_8x8x8_f16_INTRIN,
+    *get_simdgroup_multiply_accumulate_intrin(8, 8, 8, "float16"),
+)
+
+
+def get_simdgroup_intrin_group(
+    load_scope: Literal["shared"],
+    store_scope: Literal["global", "shared"],
+    dtype: str,
+    trans_a: bool = False,
+    trans_b: bool = False,
+) -> dict[str, str]:
+    """Get a group of intrinsics for tensorization on Apple GPU.
+
+    Parameters
+    ----------
+    load_scope : Literal["shared"]
+        The memory scope of the input buffer.
+
+    store_scope : Literal["global", "shared"]
+        The memory scope of the result buffer.
+
+    dtype : str
+        The data type of the input and output buffers.
+
+    trans_a : bool
+        Whether the input matrix A is transposed.
+
+    trans_b : bool
+        Whether the input matrix B is transposed.
+
+    Returns
+    -------
+    ret : Dict[str, str]
+        A group of tensor intrinsics.
+    """
+    assert load_scope in ["shared"]
+    assert store_scope in ["global", "shared"]
+    assert dtype in ["float16", "bfloat16", "float32"]
+
+    shape = "8x8x8"
+    dtype = "f16" if dtype == "float16" else "bf16" if dtype == "bfloat16" else "f32"
+    trans_a = "_trans" if trans_a else ""
+    trans_b = "_trans" if trans_b else ""
+
+    # e.g. simdgroup_load_8x8x8_f16_shared
+    load_a_intrin = f"simdgroup_load_{shape}_{dtype}_{load_scope}{trans_a}"
+    # e.g. simdgroup_load_8x8x8_f16_shared_trans
+    load_b_intrin = f"simdgroup_load_{shape}_{dtype}_{load_scope}{trans_b}"
+    # e.g. simdgroup_multiply_accumulate_8x8x8_f16
+    compute_intrin = f"simdgroup_multiply_accumulate_{shape}_{dtype}"
+    # e.g. simdgroup_make_filled_8x8x8_f16
+    init_intrin = f"simdgroup_make_filled_{shape}_{dtype}"
+    # e.g. simdgroup_store_8x8x8_f16_global
+    store_intrin = f"simdgroup_store_{shape}_{dtype}_{store_scope}"
+
+    return {
+        "init": init_intrin,
+        "load_a": load_a_intrin,
+        "load_b": load_b_intrin,
+        "compute": compute_intrin,
+        "store": store_intrin,
+    }