wehub-resource-sync
352 lines
12 KiB
Python
352 lines
12 KiB
Python
# 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,
|
|
}
|