chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,353 @@
|
||||
# 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_ffi
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm.script import ir as I
|
||||
from tvm.script import tirx as T
|
||||
from tvm.testing import env
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_metal(), reason="need metal")
|
||||
def test_metal_inf_nan():
|
||||
target = "metal"
|
||||
|
||||
def check_inf_nan(n, value, dtype):
|
||||
@I.ir_module(s_tir=True)
|
||||
class Module:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(
|
||||
A: T.Buffer((1,), dtype),
|
||||
C: T.Buffer((1,), dtype),
|
||||
):
|
||||
T.func_attr({"tirx.noalias": True})
|
||||
for i in T.thread_binding(1, thread="threadIdx.x"):
|
||||
with T.sblock("C"):
|
||||
v_i = T.axis.spatial(1, i)
|
||||
T.reads()
|
||||
T.writes(C[v_i])
|
||||
C[v_i] = T.Cast(dtype, value)
|
||||
|
||||
fun = tvm.compile(Module, target=target)
|
||||
|
||||
def run_and_check():
|
||||
dev = tvm.device(target, 0)
|
||||
a = tvm.runtime.empty((n,), dtype, dev)
|
||||
c = tvm.runtime.empty((n,), dtype, dev)
|
||||
fun(a, c)
|
||||
|
||||
tvm.testing.run_with_gpu_lock(run_and_check)
|
||||
|
||||
check_inf_nan(1, -float("inf"), "float32")
|
||||
check_inf_nan(1, -float("inf"), "float16")
|
||||
check_inf_nan(1, float("inf"), "float32")
|
||||
check_inf_nan(1, float("inf"), "float16")
|
||||
check_inf_nan(1, float("nan"), "float32")
|
||||
check_inf_nan(1, float("nan"), "float16")
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_metal(), reason="need metal")
|
||||
def test_unaligned_vectorize():
|
||||
@tvm.script.ir_module
|
||||
class IRModule:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(A: T.Buffer((2, 3), "float32"), B: T.Buffer((6,), "float32")):
|
||||
T.func_attr({"global_symbol": "main"})
|
||||
for i0_1 in T.thread_binding(3, thread="threadIdx.x"):
|
||||
for i0_0 in T.vectorized(2):
|
||||
with T.sblock("block"):
|
||||
vi0 = T.axis.spatial(6, i0_0 * 3 + i0_1)
|
||||
B[vi0] = A[vi0 // 3, vi0 % 3]
|
||||
|
||||
target = "metal"
|
||||
a = (np.arange(6).reshape(2, 3)).astype("float32")
|
||||
f = tvm.compile(IRModule, target=target)
|
||||
|
||||
def run_and_check():
|
||||
dev = tvm.metal()
|
||||
a_nd = tvm.runtime.tensor(a, dev)
|
||||
b_nd = tvm.runtime.empty((6,), "float32", dev)
|
||||
f(a_nd, b_nd)
|
||||
tvm.testing.assert_allclose(b_nd.numpy(), a.reshape(6), atol=1e-5, rtol=1e-5)
|
||||
|
||||
tvm.testing.run_with_gpu_lock(run_and_check)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_metal(), reason="need metal")
|
||||
def test_metal_erf():
|
||||
target = "metal"
|
||||
|
||||
def check_erf(n, dtype):
|
||||
@I.ir_module(s_tir=True)
|
||||
class Module:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(
|
||||
A: T.Buffer((1,), dtype),
|
||||
C: T.Buffer((1,), dtype),
|
||||
):
|
||||
T.func_attr({"tirx.noalias": True})
|
||||
for i0 in T.thread_binding(1, thread="threadIdx.x"):
|
||||
with T.sblock("C"):
|
||||
v_i0 = T.axis.spatial(1, i0)
|
||||
T.reads(A[v_i0])
|
||||
T.writes(C[v_i0])
|
||||
C[v_i0] = T.erf(A[v_i0])
|
||||
|
||||
fun = tvm.compile(Module, target=target)
|
||||
|
||||
def run_and_check():
|
||||
dev = tvm.device(target, 0)
|
||||
a = tvm.runtime.empty((n,), dtype, dev)
|
||||
c = tvm.runtime.empty((n,), dtype, dev)
|
||||
fun(a, c)
|
||||
|
||||
tvm.testing.run_with_gpu_lock(run_and_check)
|
||||
|
||||
check_erf(1, "float32")
|
||||
check_erf(1, "float16")
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_metal(), reason="need metal")
|
||||
def test_ramp():
|
||||
target = "metal"
|
||||
|
||||
@tvm.script.ir_module
|
||||
class IRModule:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(A: T.Buffer((1, 2), "int32")):
|
||||
T.func_attr({"global_symbol": "main"})
|
||||
for i in T.thread_binding(1, thread="threadIdx.x"):
|
||||
with T.sblock("block"):
|
||||
tx = T.axis.spatial(1, i)
|
||||
r = T.ramp(tx, 3, 2)
|
||||
A[0, T.ramp(0, 1, 2)] = r
|
||||
|
||||
f = tvm.compile(IRModule, target=target)
|
||||
|
||||
def run_and_check():
|
||||
dev = tvm.metal()
|
||||
a_nd = tvm.runtime.empty((1, 2), "int32", dev)
|
||||
f(a_nd)
|
||||
assert tuple(a_nd.numpy()[0, :]) == (0, 3)
|
||||
|
||||
tvm.testing.run_with_gpu_lock(run_and_check)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_metal(), reason="need metal")
|
||||
def test_select_vectorize():
|
||||
@tvm.script.ir_module
|
||||
class IRModule:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(A: T.Buffer((6), "float32"), B: T.Buffer((6,), "float32")):
|
||||
T.func_attr({"global_symbol": "main"})
|
||||
for i0_1 in T.thread_binding(3, thread="threadIdx.x"):
|
||||
for i0_0 in T.vectorized(2):
|
||||
with T.sblock("block"):
|
||||
vi0 = T.axis.spatial(6, i0_0 * 3 + i0_1)
|
||||
B[vi0] = T.Select((vi0 % 2) == 0, A[vi0], T.float32(0))
|
||||
|
||||
target = "metal"
|
||||
a = np.arange(6).astype("float32")
|
||||
f = tvm.compile(IRModule, target=target)
|
||||
a.reshape(3, 2)[:, 1] = 0
|
||||
|
||||
def run_and_check():
|
||||
dev = tvm.metal()
|
||||
a_nd = tvm.runtime.tensor(a, dev)
|
||||
b_nd = tvm.runtime.empty((6,), "float32", dev)
|
||||
f(a_nd, b_nd)
|
||||
tvm.testing.assert_allclose(b_nd.numpy(), a, atol=1e-5, rtol=1e-5)
|
||||
|
||||
tvm.testing.run_with_gpu_lock(run_and_check)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_metal(), reason="need metal")
|
||||
def test_vectorized_uint8():
|
||||
@T.prim_func(s_tir=True)
|
||||
def func(A: T.Buffer((16), "uint8"), B: T.Buffer((16), "float32")):
|
||||
for i in T.thread_binding(4, thread="threadIdx.x"):
|
||||
for j in T.vectorized(4):
|
||||
with T.sblock("block"):
|
||||
vi = T.axis.spatial(16, i * 4 + j)
|
||||
B[vi] = T.Cast("float32", A[vi])
|
||||
|
||||
a = np.arange(16).astype("uint8")
|
||||
f = tvm.compile(func, target="metal")
|
||||
|
||||
def run_and_check():
|
||||
dev = tvm.metal()
|
||||
a_nd = tvm.runtime.tensor(a, dev)
|
||||
b_nd = tvm.runtime.empty((16,), "float32", dev)
|
||||
f(a_nd, b_nd)
|
||||
tvm.testing.assert_allclose(b_nd.numpy(), a.astype("float32"), atol=1e-5, rtol=1e-5)
|
||||
|
||||
tvm.testing.run_with_gpu_lock(run_and_check)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_metal(), reason="need metal")
|
||||
def test_func_with_trailing_pod_params():
|
||||
from tvm.support import xcode # pylint: disable=import-outside-toplevel
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def func(A: T.Buffer((16), "float32"), B: T.Buffer((16), "float32"), x: T.float32):
|
||||
for i in T.thread_binding(16, thread="threadIdx.x"):
|
||||
with T.sblock("block"):
|
||||
vi = T.axis.spatial(16, i)
|
||||
B[vi] = A[vi] + x
|
||||
|
||||
@tvm.register_global_func("tvm_callback_metal_compile")
|
||||
def compile_metal(src, target):
|
||||
return xcode.compile_metal(src)
|
||||
|
||||
mod = tvm.IRModule({"main": func})
|
||||
|
||||
f = tvm.tirx.build(mod, target="metal")
|
||||
src: str = f.imports[0].inspect_source()
|
||||
occurrences = src.count("struct func_kernel_args_t")
|
||||
assert occurrences == 1, occurrences
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_metal(), reason="need metal")
|
||||
def test_metal_compile_callback_source_passthrough():
|
||||
n = 1024
|
||||
|
||||
@I.ir_module(s_tir=True)
|
||||
class Module:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(A: T.Buffer((n,), "float32"), B: T.Buffer((n,), "float32")):
|
||||
T.func_attr({"tirx.noalias": True})
|
||||
for i_0 in T.thread_binding(n // 32, thread="blockIdx.x"):
|
||||
for i_1 in T.thread_binding(32, thread="threadIdx.x"):
|
||||
with T.sblock("B"):
|
||||
v_i = T.axis.spatial(n, i_0 * 32 + i_1)
|
||||
T.reads(A[v_i])
|
||||
T.writes(B[v_i])
|
||||
B[v_i] = A[v_i] + 1.0
|
||||
|
||||
seen = {}
|
||||
|
||||
def inspect_callback(src, target):
|
||||
# Pure inspection callback: capture the source, return it untouched and
|
||||
# declare it is still textual MSL so it is compiled at load time.
|
||||
seen["src"] = src
|
||||
return (src, "metal")
|
||||
|
||||
tvm.register_global_func("tvm_callback_metal_compile", inspect_callback, override=True)
|
||||
try:
|
||||
f = tvm.compile(Module, target="metal")
|
||||
dev = tvm.metal()
|
||||
a = np.random.rand(n).astype("float32")
|
||||
a_nd = tvm.runtime.tensor(a, dev)
|
||||
b_nd = tvm.runtime.empty((n,), "float32", dev)
|
||||
f(a_nd, b_nd)
|
||||
dev.sync()
|
||||
finally:
|
||||
tvm_ffi.registry.remove_global_func("tvm_callback_metal_compile")
|
||||
|
||||
assert "src" in seen and len(seen["src"]) > 0
|
||||
tvm.testing.assert_allclose(b_nd.numpy(), a + 1.0, atol=1e-5, rtol=1e-5)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_metal(), reason="need metal")
|
||||
def test_metal_compile_callback_mixed_formats_rejected():
|
||||
n = 1024
|
||||
|
||||
@I.ir_module(s_tir=True)
|
||||
class Module:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(
|
||||
A: T.Buffer((n,), "float32"),
|
||||
B: T.Buffer((n,), "float32"),
|
||||
C: T.Buffer((n,), "float32"),
|
||||
):
|
||||
T.func_attr({"tirx.noalias": True})
|
||||
# Two independent thread-bound regions -> two device kernels, so the
|
||||
# compile callback is invoked twice within one module.
|
||||
for i_0 in T.thread_binding(n // 32, thread="blockIdx.x"):
|
||||
for i_1 in T.thread_binding(32, thread="threadIdx.x"):
|
||||
with T.sblock("B"):
|
||||
v_i = T.axis.spatial(n, i_0 * 32 + i_1)
|
||||
T.reads(A[v_i])
|
||||
T.writes(B[v_i])
|
||||
B[v_i] = A[v_i] + 1.0
|
||||
for j_0 in T.thread_binding(n // 32, thread="blockIdx.x"):
|
||||
for j_1 in T.thread_binding(32, thread="threadIdx.x"):
|
||||
with T.sblock("C"):
|
||||
v_j = T.axis.spatial(n, j_0 * 32 + j_1)
|
||||
T.reads(A[v_j])
|
||||
T.writes(C[v_j])
|
||||
C[v_j] = A[v_j] + 2.0
|
||||
|
||||
calls = {"n": 0}
|
||||
|
||||
def mixed_callback(src, target):
|
||||
calls["n"] += 1
|
||||
if calls["n"] == 1:
|
||||
# Treated as a compiled metallib payload.
|
||||
return src
|
||||
# Second kernel declares textual MSL, contradicting the metallib above.
|
||||
return (src, "metal")
|
||||
|
||||
tvm.register_global_func("tvm_callback_metal_compile", mixed_callback, override=True)
|
||||
try:
|
||||
with pytest.raises(Exception, match="inconsistent formats"):
|
||||
tvm.compile(Module, target="metal")
|
||||
finally:
|
||||
tvm_ffi.registry.remove_global_func("tvm_callback_metal_compile")
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_metal(), reason="need metal")
|
||||
def test_export_load_with_fallback(monkeypatch, tmp_path):
|
||||
"""Force the codegen wrapper into the fallback branch, then export."""
|
||||
n = 1024
|
||||
|
||||
@I.ir_module(s_tir=True)
|
||||
class Module:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(A: T.Buffer((n,), "float32"), B: T.Buffer((n,), "float32")):
|
||||
T.func_attr({"tirx.noalias": True})
|
||||
for i_0 in T.thread_binding(n // 32, thread="blockIdx.x"):
|
||||
for i_1 in T.thread_binding(32, thread="threadIdx.x"):
|
||||
with T.sblock("B"):
|
||||
v_i = T.axis.spatial(n, i_0 * 32 + i_1)
|
||||
T.reads(A[v_i])
|
||||
T.writes(B[v_i])
|
||||
B[v_i] = A[v_i] + 1.0
|
||||
|
||||
monkeypatch.setenv("TVM_COMPILE_FORCE_FALLBACK", "1")
|
||||
host_lib = tvm.compile(Module, target="metal")
|
||||
monkeypatch.delenv("TVM_COMPILE_FORCE_FALLBACK")
|
||||
|
||||
lib_path = str(tmp_path / "lib.so")
|
||||
host_lib.export_library(lib_path)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
Reference in New Issue
Block a user