chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,306 @@
|
||||
# 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 re
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
import tvm.tirx as tirx
|
||||
from tvm.ir.module import IRModule
|
||||
from tvm.runtime.executable import Executable
|
||||
from tvm.script import tirx as T
|
||||
from tvm.support.nvcc import have_fp16
|
||||
from tvm.testing import env
|
||||
|
||||
VECTOR_N_INPUTS = 8
|
||||
|
||||
|
||||
def make_prim_func(
|
||||
name: str,
|
||||
dtype: str,
|
||||
num_inputs: int,
|
||||
op: Callable[[tirx.Expr, ...], tirx.Expr],
|
||||
) -> tirx.PrimFunc:
|
||||
"""Make a primitive function that applies the given operation to the input buffer."""
|
||||
if num_inputs == 1:
|
||||
|
||||
@T.prim_func
|
||||
def kernel(
|
||||
A: T.Buffer((VECTOR_N_INPUTS,), dtype),
|
||||
B: T.Buffer((VECTOR_N_INPUTS,), dtype),
|
||||
):
|
||||
T.func_attr({"global_symbol": name + "_kernel", "tirx.noalias": True})
|
||||
for i in T.thread_binding(VECTOR_N_INPUTS, thread="threadIdx.x"):
|
||||
B[i] = op(A[i])
|
||||
|
||||
return kernel
|
||||
elif num_inputs == 2:
|
||||
|
||||
@T.prim_func
|
||||
def kernel(
|
||||
A: T.Buffer((VECTOR_N_INPUTS,), dtype),
|
||||
E: T.Buffer((VECTOR_N_INPUTS,), dtype),
|
||||
B: T.Buffer((VECTOR_N_INPUTS,), dtype),
|
||||
):
|
||||
T.func_attr({"global_symbol": name + "_kernel", "tirx.noalias": True})
|
||||
for i in T.thread_binding(VECTOR_N_INPUTS, thread="threadIdx.x"):
|
||||
B[i] = op(A[i], E[i])
|
||||
|
||||
return kernel
|
||||
else:
|
||||
raise ValueError(f"Unsupported number of inputs: {num_inputs}")
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MathCase:
|
||||
name: str
|
||||
op: Callable[[tirx.Expr, ...], tirx.Expr]
|
||||
num_inputs: int
|
||||
default_intrinsic_f16: str
|
||||
default_intrinsic_bf16: str
|
||||
default_intrinsic_f32: str
|
||||
default_intrinsic_f64: str
|
||||
fast_math_intrinsic_f32: str
|
||||
np_ref: object
|
||||
rtol: float = 1e-5
|
||||
atol: float = 1e-6
|
||||
|
||||
|
||||
MATH_CASES = [
|
||||
MathCase(
|
||||
"exp_case",
|
||||
T.exp,
|
||||
1,
|
||||
"hexp",
|
||||
"hexp",
|
||||
"expf",
|
||||
"exp",
|
||||
"__expf",
|
||||
lambda x: np.exp(x),
|
||||
),
|
||||
MathCase(
|
||||
"exp10_case",
|
||||
T.exp10,
|
||||
1,
|
||||
"hexp10",
|
||||
"hexp10",
|
||||
"exp10f",
|
||||
"exp10",
|
||||
"__exp10f",
|
||||
lambda x: np.power(10.0, x),
|
||||
),
|
||||
MathCase(
|
||||
"log_case",
|
||||
T.log,
|
||||
1,
|
||||
"hlog",
|
||||
"hlog",
|
||||
"logf",
|
||||
"log",
|
||||
"__logf",
|
||||
lambda x: np.log(x),
|
||||
),
|
||||
MathCase(
|
||||
"log2_case",
|
||||
T.log2,
|
||||
1,
|
||||
"hlog2",
|
||||
"hlog2",
|
||||
"log2f",
|
||||
"log2",
|
||||
"__log2f",
|
||||
lambda x: np.log2(x),
|
||||
),
|
||||
MathCase(
|
||||
"log10_case",
|
||||
T.log10,
|
||||
1,
|
||||
"hlog10",
|
||||
"hlog10",
|
||||
"log10f",
|
||||
"log10",
|
||||
"__log10f",
|
||||
lambda x: np.log10(x),
|
||||
),
|
||||
MathCase(
|
||||
"tan_case",
|
||||
T.tan,
|
||||
1,
|
||||
"htan",
|
||||
"htan",
|
||||
"tanf",
|
||||
"tan",
|
||||
"tanf",
|
||||
lambda x: np.tan(x),
|
||||
),
|
||||
MathCase(
|
||||
"cos_case",
|
||||
T.cos,
|
||||
1,
|
||||
"hcos",
|
||||
"hcos",
|
||||
"cosf",
|
||||
"cos",
|
||||
"__cosf",
|
||||
lambda x: np.cos(x),
|
||||
),
|
||||
MathCase(
|
||||
"sin_case",
|
||||
T.sin,
|
||||
1,
|
||||
"hsin",
|
||||
"hsin",
|
||||
"sinf",
|
||||
"sin",
|
||||
"__sinf",
|
||||
lambda x: np.sin(x),
|
||||
),
|
||||
MathCase(
|
||||
"tanh_case",
|
||||
T.tanh,
|
||||
1,
|
||||
"htanh",
|
||||
"htanh",
|
||||
"tanhf",
|
||||
"tanh",
|
||||
"__tanhf",
|
||||
lambda x: np.tanh(x),
|
||||
),
|
||||
MathCase(
|
||||
"pow_case",
|
||||
T.pow,
|
||||
2,
|
||||
"hpow",
|
||||
"hpow",
|
||||
"powf",
|
||||
"pow",
|
||||
"__powf",
|
||||
lambda x, y: np.power(x, y),
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def make_mod(
|
||||
dtype: str, case: MathCase, enable_fast_math: bool
|
||||
) -> tuple[tvm.target.Target, tvm.IRModule]:
|
||||
"""Make a module for the given dtype and case."""
|
||||
target = tvm.target.Target("cuda")
|
||||
prim_func = make_prim_func(case.name, dtype, case.num_inputs, case.op)
|
||||
return target, tvm.IRModule.from_expr(prim_func.with_attr("target", target))
|
||||
|
||||
|
||||
def expected_intrinsic(dtype: str, case: MathCase, enable_fast_math: bool) -> str:
|
||||
"""Get the expected intrinsic for the given dtype and case."""
|
||||
if dtype == "float16":
|
||||
return case.default_intrinsic_f16
|
||||
elif dtype == "bfloat16":
|
||||
return case.default_intrinsic_bf16
|
||||
elif dtype == "float32":
|
||||
return case.fast_math_intrinsic_f32 if enable_fast_math else case.default_intrinsic_f32
|
||||
elif dtype == "float64":
|
||||
return case.default_intrinsic_f64
|
||||
else:
|
||||
raise ValueError(f"Unsupported dtype: {dtype}")
|
||||
|
||||
|
||||
def check_lowered_ir(
|
||||
dtype: str, case: MathCase, enable_fast_math: bool
|
||||
) -> tuple[tvm.target.Target, IRModule]:
|
||||
"""Check the lowered IR for the given dtype and case."""
|
||||
target, mod = make_mod(dtype, case, enable_fast_math)
|
||||
with tvm.transform.PassContext(config={"tirx.enable_fast_math": enable_fast_math}):
|
||||
lowered_mod = tvm.tirx.transform.LowerIntrin()(mod)
|
||||
script = lowered_mod.script(show_meta=False)
|
||||
expected = expected_intrinsic(dtype, case, enable_fast_math)
|
||||
assert re.search(rf"""["']{re.escape(expected)}["']""", script)
|
||||
return target, lowered_mod
|
||||
|
||||
|
||||
def check_cuda_source(
|
||||
target: tvm.target.Target,
|
||||
mod: IRModule,
|
||||
dtype: str,
|
||||
case: MathCase,
|
||||
enable_fast_math: bool,
|
||||
) -> Executable:
|
||||
"""Check the CUDA source for the given dtype and case."""
|
||||
with tvm.transform.PassContext(config={"tirx.enable_fast_math": enable_fast_math}):
|
||||
executable = tvm.compile(mod, target=target)
|
||||
source = executable.mod.imports[0].inspect_source()
|
||||
expected = expected_intrinsic(dtype, case, enable_fast_math)
|
||||
assert re.search(rf"(?<!_)\b{re.escape(expected)}\s*\(", source)
|
||||
return executable
|
||||
|
||||
|
||||
def make_numpy_inputs(dtype: str, case: MathCase):
|
||||
"""Make the numpy inputs for the given dtype and case."""
|
||||
lhs = np.array([0.25, 0.5, 1.0, 2.0, 4.0, 9.0, 16.0, 10.0], dtype=dtype)
|
||||
if case.num_inputs == 1:
|
||||
return [lhs]
|
||||
elif case.num_inputs == 2:
|
||||
rhs = np.array([2.0, 3.0, 0.5, 1.5, 0.25, 0.5, 2.0, 1.0], dtype=dtype)
|
||||
return [lhs, rhs]
|
||||
else:
|
||||
raise ValueError(f"Unsupported number of inputs: {case.num_inputs}")
|
||||
|
||||
|
||||
def check_runtime(dtype: str, case: MathCase, executable: Executable):
|
||||
"""Check the runtime for the given dtype and case."""
|
||||
np_inputs = make_numpy_inputs(dtype, case)
|
||||
expected = case.np_ref(*[arr.astype(dtype) for arr in np_inputs]).astype(dtype)
|
||||
|
||||
def run_and_check():
|
||||
dev = tvm.cuda(0)
|
||||
tvm_inputs = [tvm.runtime.tensor(arr, device=dev) for arr in np_inputs]
|
||||
output = tvm.runtime.empty((VECTOR_N_INPUTS,), dtype, dev)
|
||||
|
||||
executable(*tvm_inputs, output)
|
||||
|
||||
actual = output.numpy()
|
||||
np.testing.assert_allclose(actual, expected, rtol=case.rtol, atol=case.atol)
|
||||
|
||||
tvm.testing.run_with_gpu_lock(run_and_check)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("enable_fast_math", [False, True], ids=["default", "fast_math"])
|
||||
def test_cuda_math_intrinsic_lowering_pass_context(enable_fast_math):
|
||||
check_lowered_ir("float32", MATH_CASES[0], enable_fast_math)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
||||
@pytest.mark.parametrize(
|
||||
"dtype",
|
||||
["float16", "bfloat16", "float32", "float64"],
|
||||
)
|
||||
@pytest.mark.parametrize("case", MATH_CASES, ids=lambda case: f"{case.name}")
|
||||
@pytest.mark.parametrize("enable_fast_math", [False, True], ids=["default", "fast_math"])
|
||||
def test_cuda_math_intrinsic_lowering_source_and_runtime(dtype, case, enable_fast_math):
|
||||
if dtype == "float16" and not have_fp16(tvm.cuda(0).compute_version):
|
||||
pytest.skip("GPU does not support float16")
|
||||
if dtype == "bfloat16" and case.name.startswith("pow_"):
|
||||
pytest.skip("pow_argnames=case is only supported for float")
|
||||
|
||||
target, lowered_mod = check_lowered_ir(dtype, case, enable_fast_math)
|
||||
executable = check_cuda_source(target, lowered_mod, dtype, case, enable_fast_math)
|
||||
check_runtime(dtype, case, executable)
|
||||
Reference in New Issue
Block a user