wehub-resource-sync
346 lines
11 KiB
Python
346 lines
11 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.
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
import tvm
|
|
import tvm.testing
|
|
from tvm.script import tirx as T
|
|
from tvm.testing import env
|
|
|
|
|
|
def gen_2in4_mask(m: int, n: int):
|
|
assert n % 4 == 0
|
|
return np.array(
|
|
[[np.sort(np.random.choice(4, 2, replace=False)) for _ in range(n // 4)] for _ in range(m)]
|
|
).astype("uint8")
|
|
|
|
|
|
def get_dense_mat_by_mask(val, mask):
|
|
m, n_chunks, _ = mask.shape
|
|
val = val.reshape(m, n_chunks, 2)
|
|
ret = np.zeros((m, n_chunks, 4)).astype(val.dtype)
|
|
for i in range(m):
|
|
for j in range(n_chunks):
|
|
for k in range(2):
|
|
ret[i, j, mask[i, j, k]] = val[i, j, k]
|
|
return ret.reshape(m, n_chunks * 4)
|
|
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def mma_sp_m16n8k16_f16f16f16(a: T.handle, b: T.handle, c: T.handle, _metadata: T.handle):
|
|
T.func_attr({"global_symbol": "default_function", "tirx.noalias": True})
|
|
A = T.match_buffer(a, [16, 8], dtype="float16")
|
|
B = T.match_buffer(b, [16, 8], dtype="float16")
|
|
C = T.match_buffer(c, [16, 8], dtype="float16")
|
|
metadata = T.match_buffer(_metadata, [8], dtype="uint32")
|
|
brow = T.env_thread("blockIdx.y")
|
|
bcol = T.env_thread("blockIdx.x")
|
|
tx = T.env_thread("threadIdx.x")
|
|
T.launch_thread(brow, 1)
|
|
T.launch_thread(bcol, 1)
|
|
T.launch_thread(tx, 32)
|
|
multi_a = T.decl_buffer([4], "float16", scope="local")
|
|
multi_b = T.decl_buffer([4], "float16", scope="local")
|
|
accum = T.decl_buffer([4], "float16", scope="local")
|
|
meta_local = T.decl_buffer([1], "uint32", scope="local")
|
|
for i in range(4):
|
|
accum[i] = T.float16(0)
|
|
|
|
for i in range(4):
|
|
multi_a[i] = A[tx // 4 + i // 2 * 8, tx % 4 * 2 + i % 2]
|
|
|
|
for i in range(4):
|
|
multi_b[i] = B[tx % 4 * 2 + i % 2 + i // 2 * 8, tx // 4]
|
|
|
|
meta_local[0] = metadata[tx // 4]
|
|
|
|
T.evaluate(
|
|
T.ptx.mma.sp(
|
|
"m16n8k16",
|
|
"row",
|
|
"col",
|
|
"fp16",
|
|
"fp16",
|
|
"fp16",
|
|
multi_a.data,
|
|
0,
|
|
multi_b.data,
|
|
0,
|
|
accum.data,
|
|
0,
|
|
meta_local.data,
|
|
0,
|
|
0,
|
|
False,
|
|
dtype="float16",
|
|
)
|
|
)
|
|
|
|
for i in range(4):
|
|
C[i // 2 * 8 + tx // 4, tx % 4 * 2 + i % 2] = accum[i]
|
|
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def mma_sp_m16n8k16_f16f16f32(a: T.handle, b: T.handle, c: T.handle, _metadata: T.handle):
|
|
T.func_attr({"global_symbol": "default_function", "tirx.noalias": True})
|
|
A = T.match_buffer(a, [16, 8], dtype="float16")
|
|
B = T.match_buffer(b, [16, 8], dtype="float16")
|
|
C = T.match_buffer(c, [16, 8], dtype="float32")
|
|
metadata = T.match_buffer(_metadata, [8], dtype="uint32")
|
|
brow = T.env_thread("blockIdx.y")
|
|
bcol = T.env_thread("blockIdx.x")
|
|
tx = T.env_thread("threadIdx.x")
|
|
T.launch_thread(brow, 1)
|
|
T.launch_thread(bcol, 1)
|
|
T.launch_thread(tx, 32)
|
|
multi_a = T.decl_buffer([4], "float16", scope="local")
|
|
multi_b = T.decl_buffer([4], "float16", scope="local")
|
|
accum = T.decl_buffer([4], "float32", scope="local")
|
|
meta_local = T.decl_buffer([1], "uint32", scope="local")
|
|
for i in range(4):
|
|
accum[i] = T.float16(0)
|
|
|
|
for i in range(4):
|
|
multi_a[i] = A[tx // 4 + i // 2 * 8, tx % 4 * 2 + i % 2]
|
|
|
|
for i in range(4):
|
|
multi_b[i] = B[tx % 4 * 2 + i % 2 + i // 2 * 8, tx // 4]
|
|
|
|
meta_local[0] = metadata[tx // 4]
|
|
|
|
T.evaluate(
|
|
T.ptx.mma.sp(
|
|
"m16n8k16",
|
|
"row",
|
|
"col",
|
|
"fp16",
|
|
"fp16",
|
|
"fp32",
|
|
multi_a.data,
|
|
0,
|
|
multi_b.data,
|
|
0,
|
|
accum.data,
|
|
0,
|
|
meta_local.data,
|
|
0,
|
|
0,
|
|
False,
|
|
dtype="float32",
|
|
)
|
|
)
|
|
|
|
for i in range(4):
|
|
C[i // 2 * 8 + tx // 4, tx % 4 * 2 + i % 2] = accum[i]
|
|
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def mma_sp_m16n8k32_f16f16f16(a: T.handle, b: T.handle, c: T.handle, _metadata: T.handle):
|
|
T.func_attr({"global_symbol": "default_function", "tirx.noalias": True})
|
|
A = T.match_buffer(a, [16, 16], dtype="float16")
|
|
B = T.match_buffer(b, [32, 8], dtype="float16")
|
|
C = T.match_buffer(c, [16, 8], dtype="float16")
|
|
metadata = T.match_buffer(_metadata, [16], dtype="uint32")
|
|
brow = T.env_thread("blockIdx.y")
|
|
bcol = T.env_thread("blockIdx.x")
|
|
tx = T.env_thread("threadIdx.x")
|
|
T.launch_thread(brow, 1)
|
|
T.launch_thread(bcol, 1)
|
|
T.launch_thread(tx, 32)
|
|
multi_a = T.decl_buffer([8], "float16", scope="local")
|
|
multi_b = T.decl_buffer([8], "float16", scope="local")
|
|
accum = T.decl_buffer([4], "float16", scope="local")
|
|
meta_local = T.decl_buffer([1], "uint32", scope="local")
|
|
for i in range(4):
|
|
accum[i] = T.float16(0)
|
|
|
|
for i in range(8):
|
|
multi_a[i] = A[(i % 4) // 2 * 8 + tx // 4, i // 4 * 8 + tx % 4 * 2 + i % 2]
|
|
|
|
for i in range(8):
|
|
multi_b[i] = B[i // 2 * 8 + tx % 4 * 2 + i % 2, tx // 4]
|
|
|
|
meta_local[0] = metadata[tx // 4 * 2 + tx % 2]
|
|
|
|
T.evaluate(
|
|
T.ptx.mma.sp(
|
|
"m16n8k32",
|
|
"row",
|
|
"col",
|
|
"fp16",
|
|
"fp16",
|
|
"fp16",
|
|
multi_a.data,
|
|
0,
|
|
multi_b.data,
|
|
0,
|
|
accum.data,
|
|
0,
|
|
meta_local.data,
|
|
0,
|
|
0,
|
|
False,
|
|
dtype="float16",
|
|
)
|
|
)
|
|
|
|
for i in range(4):
|
|
C[i // 2 * 8 + tx // 4, tx % 4 * 2 + i % 2] = accum[i]
|
|
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def mma_sp_m16n8k32_f16f16f32(a: T.handle, b: T.handle, c: T.handle, _metadata: T.handle):
|
|
T.func_attr({"global_symbol": "default_function", "tirx.noalias": True})
|
|
A = T.match_buffer(a, [16, 16], dtype="float16")
|
|
B = T.match_buffer(b, [32, 8], dtype="float16")
|
|
C = T.match_buffer(c, [16, 8], dtype="float32")
|
|
metadata = T.match_buffer(_metadata, [16], dtype="uint32")
|
|
brow = T.env_thread("blockIdx.y")
|
|
bcol = T.env_thread("blockIdx.x")
|
|
tx = T.env_thread("threadIdx.x")
|
|
T.launch_thread(brow, 1)
|
|
T.launch_thread(bcol, 1)
|
|
T.launch_thread(tx, 32)
|
|
multi_a = T.decl_buffer([8], "float16", scope="local")
|
|
multi_b = T.decl_buffer([8], "float16", scope="local")
|
|
accum = T.decl_buffer([4], "float32", scope="local")
|
|
meta_local = T.decl_buffer([1], "uint32", scope="local")
|
|
for i in range(4):
|
|
accum[i] = T.float16(0)
|
|
|
|
for i in range(8):
|
|
multi_a[i] = A[(i % 4) // 2 * 8 + tx // 4, i // 4 * 8 + tx % 4 * 2 + i % 2]
|
|
|
|
for i in range(8):
|
|
multi_b[i] = B[i // 2 * 8 + tx % 4 * 2 + i % 2, tx // 4]
|
|
|
|
meta_local[0] = metadata[tx // 4 * 2 + tx % 2]
|
|
|
|
T.evaluate(
|
|
T.ptx.mma.sp(
|
|
"m16n8k32",
|
|
"row",
|
|
"col",
|
|
"fp16",
|
|
"fp16",
|
|
"fp32",
|
|
multi_a.data,
|
|
0,
|
|
multi_b.data,
|
|
0,
|
|
accum.data,
|
|
0,
|
|
meta_local.data,
|
|
0,
|
|
0,
|
|
False,
|
|
dtype="float32",
|
|
)
|
|
)
|
|
|
|
for i in range(4):
|
|
C[i // 2 * 8 + tx // 4, tx % 4 * 2 + i % 2] = accum[i]
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_cuda_compute(8), reason="need cuda compute >= 8.0")
|
|
def test_mma_sp_m16n8k16_f16():
|
|
def get_meta_m16n8k16_half(mask):
|
|
assert mask.shape == (16, 4, 2)
|
|
mask = mask.reshape(16, 8)
|
|
ret = np.zeros((8,)).astype("uint32")
|
|
|
|
for i in range(8):
|
|
base = 1
|
|
for blk in range(2):
|
|
for j in range(8):
|
|
ret[i] |= int(mask[blk * 8 + i, j]) * base
|
|
base = base << 2
|
|
return ret
|
|
|
|
for out_dtype in ["float16", "float32"]:
|
|
func = mma_sp_m16n8k16_f16f16f16 if out_dtype == "float16" else mma_sp_m16n8k16_f16f16f32
|
|
sch = tvm.s_tir.Schedule(func)
|
|
cuda_mod = tvm.compile(sch.mod, target="cuda")
|
|
|
|
A_np = np.random.uniform(-1, 1, [16, 8]).astype("float16")
|
|
B_np = np.random.uniform(-1, 1, [16, 8]).astype("float16")
|
|
mask = gen_2in4_mask(16, 16)
|
|
A_dense_np = get_dense_mat_by_mask(A_np, mask)
|
|
C_np = np.matmul(A_dense_np, B_np).astype(out_dtype)
|
|
meta = get_meta_m16n8k16_half(mask)
|
|
|
|
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(np.zeros_like(C_np), ctx)
|
|
meta_tvm = tvm.runtime.tensor(meta, ctx)
|
|
cuda_mod(A_tvm, B_tvm, C_tvm, meta_tvm)
|
|
tvm.testing.assert_allclose(C_tvm.numpy(), C_np, atol=1e-3, rtol=1e-3)
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
|
|
|
|
@pytest.mark.gpu
|
|
@pytest.mark.skipif(not env.has_cuda_compute(8), reason="need cuda compute >= 8.0")
|
|
def test_mma_sp_m16n8k32_f16():
|
|
def get_meta_m16n8k32_half(mask):
|
|
assert mask.shape == (16, 8, 2)
|
|
mask = mask.reshape(16, 2, 8)
|
|
ret = np.zeros((8, 2)).astype("uint32")
|
|
|
|
for i in range(8):
|
|
for k in range(2):
|
|
base = 1
|
|
for blk in range(2):
|
|
for j in range(8):
|
|
ret[i, k] |= int(mask[blk * 8 + i, k, j]) * base
|
|
base = base << 2
|
|
|
|
return ret.reshape(16)
|
|
|
|
for out_dtype in ["float16", "float32"]:
|
|
func = mma_sp_m16n8k32_f16f16f16 if out_dtype == "float16" else mma_sp_m16n8k32_f16f16f32
|
|
sch = tvm.s_tir.Schedule(func)
|
|
cuda_mod = tvm.compile(sch.mod, target="cuda")
|
|
|
|
A_np = np.random.uniform(-1, 1, [16, 16]).astype("float16")
|
|
B_np = np.random.uniform(-1, 1, [32, 8]).astype("float16")
|
|
mask = gen_2in4_mask(16, 32)
|
|
A_dense_np = get_dense_mat_by_mask(A_np, mask)
|
|
C_np = np.matmul(A_dense_np, B_np).astype(out_dtype)
|
|
meta = get_meta_m16n8k32_half(mask)
|
|
|
|
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(np.zeros_like(C_np), ctx)
|
|
meta_tvm = tvm.runtime.tensor(meta, ctx)
|
|
cuda_mod(A_tvm, B_tvm, C_tvm, meta_tvm)
|
|
tvm.testing.assert_allclose(C_tvm.numpy(), C_np, atol=1e-3, rtol=1e-3)
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_mma_sp_m16n8k16_f16()
|
|
test_mma_sp_m16n8k32_f16()
|