2022-07-04 19:01:19 +08:00
|
|
|
import atexit
|
2022-03-18 15:15:51 +08:00
|
|
|
import functools
|
|
|
|
|
import math
|
2022-08-23 11:59:04 +08:00
|
|
|
import shutil
|
2022-04-01 15:19:02 +08:00
|
|
|
import threading
|
2022-09-06 20:15:56 +08:00
|
|
|
from os import listdir, rmdir, stat
|
2022-07-04 19:01:19 +08:00
|
|
|
from os.path import join
|
|
|
|
|
from tempfile import mkdtemp
|
2022-03-18 15:15:51 +08:00
|
|
|
|
|
|
|
|
import pytest
|
|
|
|
|
|
|
|
|
|
import taichi as ti
|
|
|
|
|
from tests import test_utils
|
|
|
|
|
|
2022-07-04 19:01:19 +08:00
|
|
|
OFFLINE_CACHE_TEMP_DIR = mkdtemp()
|
|
|
|
|
atexit.register(lambda: rmdir(OFFLINE_CACHE_TEMP_DIR))
|
|
|
|
|
|
2023-04-03 18:34:47 +08:00
|
|
|
supported_llvm_archs = {ti.cpu, ti.cuda, ti.amdgpu}
|
|
|
|
|
supported_gfx_archs = {ti.opengl, ti.vulkan, ti.metal, ti.dx11}
|
2023-01-09 20:47:57 +08:00
|
|
|
supported_archs_offline_cache = supported_llvm_archs | supported_gfx_archs
|
2022-10-10 10:12:10 +08:00
|
|
|
supported_archs_offline_cache = {v for v in supported_archs_offline_cache if v in test_utils.expected_archs()}
|
2022-05-17 09:49:45 +08:00
|
|
|
|
|
|
|
|
|
2022-07-04 19:01:19 +08:00
|
|
|
def is_offline_cache_file(filename):
|
2023-03-24 11:22:57 +08:00
|
|
|
suffixes = (".tic",)
|
2022-07-04 19:01:19 +08:00
|
|
|
return filename.endswith(suffixes)
|
|
|
|
|
|
|
|
|
|
|
2022-08-23 11:59:04 +08:00
|
|
|
def cache_files_size(path):
|
2022-07-04 19:01:19 +08:00
|
|
|
files = listdir(path)
|
|
|
|
|
result = 0
|
|
|
|
|
for file in files:
|
|
|
|
|
if is_offline_cache_file(file):
|
2022-08-23 11:59:04 +08:00
|
|
|
result += stat(join(path, file)).st_size
|
2022-07-04 19:01:19 +08:00
|
|
|
return result
|
|
|
|
|
|
|
|
|
|
|
2023-03-28 09:47:33 +08:00
|
|
|
def expected_num_cache_files(num_kernels: int = 0) -> int:
|
|
|
|
|
if num_kernels == 0:
|
|
|
|
|
return 0
|
2023-03-24 11:22:57 +08:00
|
|
|
# code files(*.tic) + metadata files(ticache.tcb)
|
2023-03-28 09:47:33 +08:00
|
|
|
return num_kernels + 1
|
2022-03-18 15:15:51 +08:00
|
|
|
|
|
|
|
|
|
2022-04-01 15:19:02 +08:00
|
|
|
def tmp_offline_cache_file_path():
|
2022-07-04 19:01:19 +08:00
|
|
|
return join(OFFLINE_CACHE_TEMP_DIR, str(threading.currentThread().ident))
|
2022-04-01 15:19:02 +08:00
|
|
|
|
|
|
|
|
|
|
|
|
|
def current_thread_ext_options():
|
|
|
|
|
return {
|
|
|
|
|
"offline_cache": True,
|
|
|
|
|
"offline_cache_file_path": tmp_offline_cache_file_path(),
|
2022-11-16 18:09:13 +08:00
|
|
|
"cuda_stack_limit": 1024,
|
|
|
|
|
"device_memory_GB": 0.2,
|
2022-04-01 15:19:02 +08:00
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
2023-03-28 09:47:33 +08:00
|
|
|
def cache_files_cnt():
|
2022-08-23 11:59:04 +08:00
|
|
|
try:
|
2023-03-28 09:47:33 +08:00
|
|
|
return len(listdir(tmp_offline_cache_file_path()))
|
2022-08-23 11:59:04 +08:00
|
|
|
except FileNotFoundError:
|
|
|
|
|
return 0
|
|
|
|
|
|
|
|
|
|
|
2022-03-18 15:15:51 +08:00
|
|
|
@ti.kernel
|
|
|
|
|
def kernel0() -> ti.i32:
|
|
|
|
|
return 1
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def python_kernel0():
|
|
|
|
|
return 1
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@ti.kernel
|
|
|
|
|
def kernel1(a: ti.i32, b: ti.i32, c: ti.f32) -> ti.f32:
|
|
|
|
|
return a / b + c * b - c + a**2 + ti.log(c)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def python_kernel1(a, b, c):
|
|
|
|
|
return a / b + c * b - c + a**2 + math.log(c)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@ti.kernel
|
|
|
|
|
def kernel2(n: ti.i32) -> ti.i32:
|
|
|
|
|
x = 0
|
|
|
|
|
for i in range(n):
|
|
|
|
|
ti.atomic_add(x, 1)
|
|
|
|
|
return x
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def python_kernel2(n):
|
|
|
|
|
return n
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def kernel3(a, mat):
|
|
|
|
|
mat_type = ti.types.matrix(mat.n, mat.m, ti.i32)
|
|
|
|
|
|
|
|
|
|
@ti.kernel
|
|
|
|
|
def kernel(u: ti.i32, v: mat_type) -> mat_type:
|
|
|
|
|
return u * v
|
|
|
|
|
|
|
|
|
|
return kernel(a, mat)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def python_kernel3(a, mat):
|
|
|
|
|
return a * mat
|
|
|
|
|
|
|
|
|
|
|
2022-07-23 19:16:23 +08:00
|
|
|
@ti.func
|
|
|
|
|
def func_sum(lo: ti.i32, hi: ti.i32) -> ti.i32:
|
|
|
|
|
res = 0
|
|
|
|
|
for i in range(lo, hi):
|
|
|
|
|
res += i
|
|
|
|
|
return res
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@ti.func
|
|
|
|
|
def func_mul(lo: ti.i32, hi: ti.i32) -> ti.i32:
|
|
|
|
|
res = 1
|
|
|
|
|
for i in range(lo, hi):
|
|
|
|
|
res *= i
|
|
|
|
|
return res
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@ti.kernel
|
|
|
|
|
def kernel4(lo: ti.i32, hi: ti.i32, n: ti.i32) -> ti.i32:
|
|
|
|
|
res = 0
|
|
|
|
|
for i in range(n):
|
|
|
|
|
res += func_sum(lo, hi)
|
|
|
|
|
return res
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def python_kernel4(lo: ti.i32, hi: ti.i32, n: ti.i32):
|
|
|
|
|
res = 0
|
|
|
|
|
for i in range(n):
|
|
|
|
|
for j in range(lo, hi):
|
|
|
|
|
res += j
|
|
|
|
|
return res
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@ti.kernel
|
|
|
|
|
def kernel5(lo: ti.i32, hi: ti.i32, n: ti.i32) -> ti.i32:
|
|
|
|
|
res = 1
|
|
|
|
|
for i in range(n):
|
|
|
|
|
res *= func_mul(lo, hi)
|
|
|
|
|
return res
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def python_kernel5(lo: ti.i32, hi: ti.i32, n: ti.i32):
|
|
|
|
|
res = 1
|
|
|
|
|
for i in range(n):
|
|
|
|
|
for j in range(lo, hi):
|
|
|
|
|
res *= j
|
|
|
|
|
return res
|
|
|
|
|
|
|
|
|
|
|
2023-01-18 11:17:59 +08:00
|
|
|
simple_kernels_to_test = [
|
2023-03-28 09:47:33 +08:00
|
|
|
(kernel0, (), python_kernel0),
|
|
|
|
|
(kernel1, (100, 200, 10.2), python_kernel1),
|
|
|
|
|
(kernel2, (1024,), python_kernel2),
|
2023-03-17 14:00:57 +08:00
|
|
|
# FIXME: add this kernel back once we have a better way to compare matrices
|
|
|
|
|
# with test_utils.approx()
|
|
|
|
|
# (kernel3, (10, ti.Matrix([[1, 2], [256, 1024]],
|
2023-03-28 09:47:33 +08:00
|
|
|
# ti.i32)), python_kernel3),
|
2023-01-18 11:17:59 +08:00
|
|
|
# FIXME: add this kernel back once #6221 is fixed
|
2023-03-28 09:47:33 +08:00
|
|
|
# (kernel4, (1, 10, 2), python_kernel4),
|
|
|
|
|
(kernel5, (1, 2, 2), python_kernel5),
|
2023-01-18 11:17:59 +08:00
|
|
|
]
|
2022-03-18 15:15:51 +08:00
|
|
|
|
|
|
|
|
|
|
|
|
|
def _test_offline_cache_dec(func):
|
|
|
|
|
@functools.wraps(func)
|
|
|
|
|
def wrapped(*args, **kwargs):
|
2022-07-04 19:01:19 +08:00
|
|
|
test_utils.mkdir_p(tmp_offline_cache_file_path())
|
2022-03-18 15:15:51 +08:00
|
|
|
ret = None
|
|
|
|
|
try:
|
|
|
|
|
ret = func(*args, **kwargs)
|
|
|
|
|
except Exception as e:
|
|
|
|
|
raise e
|
|
|
|
|
finally:
|
|
|
|
|
ti.reset()
|
2022-08-23 11:59:04 +08:00
|
|
|
shutil.rmtree(tmp_offline_cache_file_path())
|
2022-03-18 15:15:51 +08:00
|
|
|
return ret
|
|
|
|
|
|
|
|
|
|
return wrapped
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@_test_offline_cache_dec
|
2023-03-28 09:47:33 +08:00
|
|
|
def _test_offline_cache_for_a_kernel(curr_arch, kernel, args, result):
|
|
|
|
|
count_of_cache_file = cache_files_cnt()
|
2022-08-23 11:59:04 +08:00
|
|
|
|
2023-03-28 09:47:33 +08:00
|
|
|
def added_files():
|
|
|
|
|
return cache_files_cnt() - count_of_cache_file
|
2022-03-18 15:15:51 +08:00
|
|
|
|
2022-04-01 15:19:02 +08:00
|
|
|
ti.init(arch=curr_arch, enable_fallback=False, **current_thread_ext_options())
|
2022-03-18 15:15:51 +08:00
|
|
|
res1 = kernel(*args)
|
2023-03-28 09:47:33 +08:00
|
|
|
assert added_files() == expected_num_cache_files()
|
2022-03-18 15:15:51 +08:00
|
|
|
|
2022-04-01 15:19:02 +08:00
|
|
|
ti.init(arch=curr_arch, enable_fallback=False, **current_thread_ext_options())
|
2023-03-28 09:47:33 +08:00
|
|
|
assert added_files() == expected_num_cache_files(1)
|
2022-03-18 15:15:51 +08:00
|
|
|
res2 = kernel(*args)
|
|
|
|
|
assert res1 == test_utils.approx(result) and res1 == test_utils.approx(res2)
|
|
|
|
|
|
|
|
|
|
ti.reset()
|
2023-03-28 09:47:33 +08:00
|
|
|
assert added_files() == expected_num_cache_files(1)
|
2022-03-18 15:15:51 +08:00
|
|
|
|
|
|
|
|
|
|
|
|
|
@_test_offline_cache_dec
|
|
|
|
|
def _test_closing_offline_cache_for_a_kernel(curr_arch, kernel, args, result):
|
2023-03-28 09:47:33 +08:00
|
|
|
count_of_cache_file = cache_files_cnt()
|
2022-08-23 11:59:04 +08:00
|
|
|
|
2023-03-28 09:47:33 +08:00
|
|
|
def added_files():
|
|
|
|
|
return cache_files_cnt() - count_of_cache_file
|
2022-03-18 15:15:51 +08:00
|
|
|
|
2022-11-16 18:09:13 +08:00
|
|
|
def my_init():
|
2023-03-28 09:47:33 +08:00
|
|
|
ti.init(
|
|
|
|
|
arch=curr_arch,
|
|
|
|
|
enable_fallback=False,
|
|
|
|
|
offline_cache=False,
|
|
|
|
|
offline_cache_file_path=tmp_offline_cache_file_path(),
|
|
|
|
|
cuda_stack_limit=1024,
|
|
|
|
|
device_memory_GB=0.1,
|
|
|
|
|
)
|
2022-11-16 18:09:13 +08:00
|
|
|
|
|
|
|
|
my_init()
|
2022-03-18 15:15:51 +08:00
|
|
|
res1 = kernel(*args)
|
2023-03-28 09:47:33 +08:00
|
|
|
assert added_files() == expected_num_cache_files()
|
2022-03-18 15:15:51 +08:00
|
|
|
|
2022-11-16 18:09:13 +08:00
|
|
|
my_init()
|
2023-03-28 09:47:33 +08:00
|
|
|
assert added_files() == expected_num_cache_files()
|
2022-03-18 15:15:51 +08:00
|
|
|
res2 = kernel(*args)
|
2022-05-17 09:49:45 +08:00
|
|
|
|
2022-03-18 15:15:51 +08:00
|
|
|
assert res1 == test_utils.approx(result) and res1 == test_utils.approx(res2)
|
|
|
|
|
|
|
|
|
|
ti.reset()
|
2023-03-28 09:47:33 +08:00
|
|
|
assert added_files() == expected_num_cache_files()
|
2022-03-18 15:15:51 +08:00
|
|
|
|
|
|
|
|
|
|
|
|
|
@pytest.mark.parametrize("curr_arch", supported_archs_offline_cache)
|
2022-04-01 15:19:02 +08:00
|
|
|
def test_closing_offline_cache(curr_arch):
|
2023-03-28 09:47:33 +08:00
|
|
|
for kernel, args, get_res in simple_kernels_to_test:
|
2022-03-18 15:15:51 +08:00
|
|
|
_test_closing_offline_cache_for_a_kernel(curr_arch=curr_arch, kernel=kernel, args=args, result=get_res(*args))
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@pytest.mark.parametrize("curr_arch", supported_archs_offline_cache)
|
|
|
|
|
def test_offline_cache_per_kernel(curr_arch):
|
2023-03-28 09:47:33 +08:00
|
|
|
for kernel, args, get_res in simple_kernels_to_test:
|
|
|
|
|
_test_offline_cache_for_a_kernel(curr_arch=curr_arch, kernel=kernel, args=args, result=get_res(*args))
|
2022-03-18 15:15:51 +08:00
|
|
|
|
|
|
|
|
|
|
|
|
|
@pytest.mark.parametrize("curr_arch", supported_archs_offline_cache)
|
|
|
|
|
@_test_offline_cache_dec
|
|
|
|
|
def test_multiple_ib_with_offline_cache(curr_arch):
|
2023-03-28 09:47:33 +08:00
|
|
|
count_of_cache_file = cache_files_cnt()
|
2022-08-23 11:59:04 +08:00
|
|
|
|
2023-03-28 09:47:33 +08:00
|
|
|
def added_files():
|
|
|
|
|
return cache_files_cnt() - count_of_cache_file
|
2022-03-18 15:15:51 +08:00
|
|
|
|
|
|
|
|
def helper():
|
|
|
|
|
x = ti.field(float, (), needs_grad=True)
|
|
|
|
|
y = ti.field(float, (), needs_grad=True)
|
|
|
|
|
|
|
|
|
|
@ti.kernel
|
|
|
|
|
def compute_y():
|
|
|
|
|
for j in range(2):
|
|
|
|
|
for i in range(3):
|
|
|
|
|
y[None] += x[None]
|
|
|
|
|
for i in range(3):
|
|
|
|
|
y[None] += x[None]
|
|
|
|
|
|
|
|
|
|
x[None] = 1.0
|
2022-07-04 19:01:19 +08:00
|
|
|
with ti.ad.Tape(y):
|
2022-03-18 15:15:51 +08:00
|
|
|
compute_y()
|
|
|
|
|
|
|
|
|
|
assert y[None] == 12.0
|
|
|
|
|
assert x.grad[None] == 12.0
|
|
|
|
|
|
2022-04-01 15:19:02 +08:00
|
|
|
ti.init(arch=curr_arch, enable_fallback=False, **current_thread_ext_options())
|
2022-03-18 15:15:51 +08:00
|
|
|
helper()
|
2023-03-28 09:47:33 +08:00
|
|
|
assert added_files() == expected_num_cache_files()
|
2022-03-18 15:15:51 +08:00
|
|
|
|
2022-04-01 15:19:02 +08:00
|
|
|
ti.init(arch=curr_arch, enable_fallback=False, **current_thread_ext_options())
|
2023-03-28 09:47:33 +08:00
|
|
|
assert added_files() == expected_num_cache_files(9)
|
2022-03-18 15:15:51 +08:00
|
|
|
helper()
|
|
|
|
|
|
|
|
|
|
ti.reset()
|
2023-03-28 09:47:33 +08:00
|
|
|
assert added_files() == expected_num_cache_files(9)
|
2022-03-18 15:15:51 +08:00
|
|
|
|
|
|
|
|
|
|
|
|
|
@pytest.mark.parametrize("curr_arch", supported_archs_offline_cache)
|
|
|
|
|
@_test_offline_cache_dec
|
|
|
|
|
def test_calling_a_kernel_with_different_param_list(curr_arch):
|
2023-03-28 09:47:33 +08:00
|
|
|
count_of_cache_file = cache_files_cnt()
|
2022-08-23 11:59:04 +08:00
|
|
|
|
2023-03-28 09:47:33 +08:00
|
|
|
def added_files():
|
|
|
|
|
return cache_files_cnt() - count_of_cache_file
|
2022-08-23 11:59:04 +08:00
|
|
|
|
2022-03-18 15:15:51 +08:00
|
|
|
mat_type = ti.types.matrix(2, 3, ti.i32)
|
|
|
|
|
|
|
|
|
|
@ti.kernel
|
|
|
|
|
def kernel(a: mat_type, b: mat_type) -> mat_type:
|
|
|
|
|
return a + 10 * b
|
|
|
|
|
|
|
|
|
|
def np_kernel(a, b):
|
|
|
|
|
return a + 10 * b
|
|
|
|
|
|
|
|
|
|
mat1 = ti.Matrix([[1, 2, 3], [3, 2, 1]], ti.i32)
|
|
|
|
|
mat2 = ti.Matrix([[1, 2, 3], [3, 2, 1]], ti.i32)
|
|
|
|
|
mat3 = ti.Matrix([[1, 2, 3], [3, 2, 1]], ti.i32)
|
|
|
|
|
np_mat1 = mat1.to_numpy()
|
|
|
|
|
np_mat2 = mat2.to_numpy()
|
|
|
|
|
np_mat3 = mat3.to_numpy()
|
|
|
|
|
|
2023-03-28 09:47:33 +08:00
|
|
|
assert added_files() == expected_num_cache_files()
|
2022-04-01 15:19:02 +08:00
|
|
|
ti.init(arch=curr_arch, enable_fallback=False, **current_thread_ext_options())
|
2022-03-18 15:15:51 +08:00
|
|
|
assert (kernel(mat1, mat1).to_numpy() == np_kernel(np_mat1, np_mat1)).all()
|
|
|
|
|
|
2022-04-01 15:19:02 +08:00
|
|
|
ti.init(arch=curr_arch, enable_fallback=False, **current_thread_ext_options())
|
2023-03-28 09:47:33 +08:00
|
|
|
assert added_files() == expected_num_cache_files(1)
|
2022-03-18 15:15:51 +08:00
|
|
|
|
|
|
|
|
assert (kernel(mat1, mat1).to_numpy() == np_kernel(np_mat1, np_mat1)).all()
|
|
|
|
|
assert (kernel(mat1, mat2).to_numpy() == np_kernel(np_mat1, np_mat2)).all()
|
|
|
|
|
assert (kernel(mat2, mat2).to_numpy() == np_kernel(np_mat2, np_mat2)).all()
|
|
|
|
|
assert (kernel(mat2, mat3).to_numpy() == np_kernel(np_mat2, np_mat3)).all()
|
|
|
|
|
|
|
|
|
|
ti.reset()
|
2023-03-28 09:47:33 +08:00
|
|
|
assert added_files() == expected_num_cache_files(1)
|
2022-03-18 15:15:51 +08:00
|
|
|
|
|
|
|
|
|
|
|
|
|
@pytest.mark.parametrize("curr_arch", supported_archs_offline_cache)
|
|
|
|
|
@_test_offline_cache_dec
|
|
|
|
|
def test_snode_reader_and_writer_with_offline_cache(curr_arch):
|
2023-03-28 09:47:33 +08:00
|
|
|
count_of_cache_file = cache_files_cnt()
|
2022-08-23 11:59:04 +08:00
|
|
|
|
2023-03-28 09:47:33 +08:00
|
|
|
def added_files():
|
|
|
|
|
return cache_files_cnt() - count_of_cache_file
|
2022-03-18 15:15:51 +08:00
|
|
|
|
|
|
|
|
def helper():
|
|
|
|
|
x = ti.field(dtype=ti.f32, shape=())
|
|
|
|
|
y = ti.field(dtype=ti.f32, shape=())
|
|
|
|
|
|
|
|
|
|
x[None] = 3.14
|
|
|
|
|
y[None] = 4.14
|
|
|
|
|
assert x[None] == test_utils.approx(3.14)
|
|
|
|
|
assert y[None] == test_utils.approx(4.14)
|
|
|
|
|
|
|
|
|
|
x[None] = 6.28
|
|
|
|
|
y[None] = 7.28
|
|
|
|
|
assert x[None] == test_utils.approx(6.28)
|
|
|
|
|
assert y[None] == test_utils.approx(7.28)
|
|
|
|
|
|
2023-03-28 09:47:33 +08:00
|
|
|
assert added_files() == expected_num_cache_files()
|
2022-04-01 15:19:02 +08:00
|
|
|
ti.init(arch=curr_arch, enable_fallback=False, **current_thread_ext_options())
|
2022-03-18 15:15:51 +08:00
|
|
|
helper()
|
|
|
|
|
|
2022-04-01 15:19:02 +08:00
|
|
|
ti.init(arch=curr_arch, enable_fallback=False, **current_thread_ext_options())
|
2023-03-28 09:47:33 +08:00
|
|
|
assert added_files() == expected_num_cache_files(4)
|
2022-03-18 15:15:51 +08:00
|
|
|
helper()
|
|
|
|
|
|
|
|
|
|
ti.reset()
|
2023-03-28 09:47:33 +08:00
|
|
|
assert added_files() == expected_num_cache_files(4)
|
2022-03-18 15:15:51 +08:00
|
|
|
|
|
|
|
|
|
|
|
|
|
@pytest.mark.parametrize("curr_arch", supported_archs_offline_cache)
|
|
|
|
|
@_test_offline_cache_dec
|
|
|
|
|
def test_calling_many_kernels(curr_arch):
|
2023-03-28 09:47:33 +08:00
|
|
|
count_of_cache_file = cache_files_cnt()
|
2022-08-23 11:59:04 +08:00
|
|
|
|
2023-03-28 09:47:33 +08:00
|
|
|
def added_files():
|
|
|
|
|
return cache_files_cnt() - count_of_cache_file
|
2022-03-18 15:15:51 +08:00
|
|
|
|
|
|
|
|
def helper():
|
2023-03-28 09:47:33 +08:00
|
|
|
for kernel, args, get_res in simple_kernels_to_test:
|
2022-03-18 15:15:51 +08:00
|
|
|
assert kernel(*args) == test_utils.approx(get_res(*args))
|
|
|
|
|
|
2022-04-01 15:19:02 +08:00
|
|
|
ti.init(arch=curr_arch, enable_fallback=False, **current_thread_ext_options())
|
2022-03-18 15:15:51 +08:00
|
|
|
helper()
|
2023-03-28 09:47:33 +08:00
|
|
|
assert added_files() == expected_num_cache_files()
|
2022-03-18 15:15:51 +08:00
|
|
|
|
2022-04-01 15:19:02 +08:00
|
|
|
ti.init(arch=curr_arch, enable_fallback=False, **current_thread_ext_options())
|
2023-03-28 09:47:33 +08:00
|
|
|
assert added_files() == expected_num_cache_files(len(simple_kernels_to_test))
|
2022-03-18 15:15:51 +08:00
|
|
|
helper()
|
|
|
|
|
ti.reset()
|
2023-03-28 09:47:33 +08:00
|
|
|
assert added_files() == expected_num_cache_files(len(simple_kernels_to_test))
|
2022-04-21 11:16:25 +08:00
|
|
|
|
|
|
|
|
|
|
|
|
|
@pytest.mark.parametrize("curr_arch", supported_archs_offline_cache)
|
|
|
|
|
@_test_offline_cache_dec
|
|
|
|
|
def test_offline_cache_with_changing_compile_config(curr_arch):
|
2023-03-28 09:47:33 +08:00
|
|
|
count_of_cache_file = cache_files_cnt()
|
2022-08-23 11:59:04 +08:00
|
|
|
|
2023-03-28 09:47:33 +08:00
|
|
|
def added_files():
|
|
|
|
|
return cache_files_cnt() - count_of_cache_file
|
2022-04-21 11:16:25 +08:00
|
|
|
|
|
|
|
|
@ti.kernel
|
|
|
|
|
def helper():
|
|
|
|
|
b = 200
|
2023-02-23 08:13:36 +08:00
|
|
|
c = 0
|
2022-04-21 11:16:25 +08:00
|
|
|
for i in range(b):
|
|
|
|
|
c += i
|
|
|
|
|
|
2023-03-28 09:47:33 +08:00
|
|
|
assert added_files() == expected_num_cache_files()
|
2022-04-21 11:16:25 +08:00
|
|
|
ti.init(arch=curr_arch, enable_fallback=False, opt_level=0, **current_thread_ext_options())
|
|
|
|
|
helper()
|
|
|
|
|
|
|
|
|
|
ti.init(arch=curr_arch, enable_fallback=False, opt_level=1, **current_thread_ext_options())
|
2023-03-28 09:47:33 +08:00
|
|
|
assert added_files() == expected_num_cache_files(1)
|
2022-04-21 11:16:25 +08:00
|
|
|
helper()
|
|
|
|
|
|
|
|
|
|
ti.reset()
|
2023-03-28 09:47:33 +08:00
|
|
|
assert added_files() == expected_num_cache_files(2)
|
2023-04-18 18:01:32 +08:00
|
|
|
ti.init(arch=curr_arch, enable_fallback=False, default_fp=ti.f32, **current_thread_ext_options())
|
2022-04-21 11:16:25 +08:00
|
|
|
helper()
|
|
|
|
|
|
|
|
|
|
ti.reset()
|
2023-03-28 09:47:33 +08:00
|
|
|
assert added_files() == expected_num_cache_files(2)
|
2022-07-04 19:01:19 +08:00
|
|
|
|
|
|
|
|
|
2023-03-10 11:44:51 +08:00
|
|
|
@pytest.mark.parametrize("curr_arch", supported_archs_offline_cache)
|
2022-07-04 19:01:19 +08:00
|
|
|
@pytest.mark.parametrize("factor", [0.0, 0.25, 0.85, 1.0])
|
|
|
|
|
@pytest.mark.parametrize("policy", ["never", "version", "lru", "fifo"])
|
|
|
|
|
@_test_offline_cache_dec
|
|
|
|
|
def test_offline_cache_cleaning(curr_arch, factor, policy):
|
|
|
|
|
def only_init(max_size):
|
|
|
|
|
ti.init(
|
|
|
|
|
arch=curr_arch,
|
|
|
|
|
enable_fallback=False,
|
|
|
|
|
offline_cache_cleaning_policy=policy,
|
|
|
|
|
offline_cache_max_size_of_files=max_size, # bytes
|
|
|
|
|
offline_cache_cleaning_factor=factor,
|
|
|
|
|
**current_thread_ext_options()
|
2023-04-16 22:20:35 +08:00
|
|
|
)
|
2022-07-04 19:01:19 +08:00
|
|
|
|
|
|
|
|
def run_simple_kernels(max_size):
|
|
|
|
|
only_init(max_size)
|
2023-03-28 09:47:33 +08:00
|
|
|
for kernel, args, get_res in simple_kernels_to_test:
|
2022-07-04 19:01:19 +08:00
|
|
|
assert kernel(*args) == test_utils.approx(get_res(*args))
|
|
|
|
|
|
|
|
|
|
kernel_count = len(simple_kernels_to_test)
|
2023-03-28 09:47:33 +08:00
|
|
|
count_of_cache_file = cache_files_cnt()
|
2022-08-23 11:59:04 +08:00
|
|
|
|
2023-03-28 09:47:33 +08:00
|
|
|
def added_files():
|
|
|
|
|
return cache_files_cnt() - count_of_cache_file
|
2022-08-23 11:59:04 +08:00
|
|
|
|
2023-03-28 09:47:33 +08:00
|
|
|
assert added_files() == expected_num_cache_files()
|
2022-07-04 19:01:19 +08:00
|
|
|
|
2023-03-10 11:44:51 +08:00
|
|
|
run_simple_kernels(1024**3) # 1GB (>> size_of_cache_files)
|
2022-07-04 19:01:19 +08:00
|
|
|
ti.reset() # Dumping cache data
|
2023-03-28 09:47:33 +08:00
|
|
|
size_of_cache_files = cache_files_size(tmp_offline_cache_file_path())
|
|
|
|
|
assert added_files() == expected_num_cache_files(kernel_count)
|
2022-07-04 19:01:19 +08:00
|
|
|
|
|
|
|
|
only_init(size_of_cache_files * 2)
|
|
|
|
|
ti.reset()
|
2023-03-28 09:47:33 +08:00
|
|
|
assert added_files() == expected_num_cache_files(kernel_count)
|
2022-07-04 19:01:19 +08:00
|
|
|
|
2023-03-10 11:44:51 +08:00
|
|
|
only_init(1) # 1B (<< size_of_cache_files)
|
2022-07-04 19:01:19 +08:00
|
|
|
ti.reset()
|
2022-09-07 00:10:30 +08:00
|
|
|
rem = []
|
2022-07-12 13:18:22 +08:00
|
|
|
if policy in ["never", "version"]:
|
2023-03-28 09:47:33 +08:00
|
|
|
rem = kernel_count
|
2022-07-12 13:18:22 +08:00
|
|
|
else:
|
2022-09-06 20:15:56 +08:00
|
|
|
lo = -min(kernel_count - int(factor * kernel_count), kernel_count)
|
|
|
|
|
lo = kernel_count if lo == 0 else lo
|
2023-03-28 09:47:33 +08:00
|
|
|
rem = len(simple_kernels_to_test[lo:])
|
|
|
|
|
assert added_files() == expected_num_cache_files(rem)
|
2022-07-23 19:16:23 +08:00
|
|
|
|
|
|
|
|
|
2022-08-23 11:59:04 +08:00
|
|
|
# FIXME: Change to `supported_archs_offline_cache` after fixing bugs of real-function on gpu
|
2022-11-16 18:09:13 +08:00
|
|
|
@pytest.mark.run_in_serial
|
2022-11-07 12:35:27 +08:00
|
|
|
@pytest.mark.parametrize("curr_arch", {ti.cpu, ti.cuda} & supported_archs_offline_cache)
|
2022-07-23 19:16:23 +08:00
|
|
|
@_test_offline_cache_dec
|
2023-04-27 16:00:46 +08:00
|
|
|
@test_utils.test(cuda_stack_limit=8192)
|
2022-07-23 19:16:23 +08:00
|
|
|
def test_offline_cache_for_kernels_calling_real_func(curr_arch):
|
2023-03-28 09:47:33 +08:00
|
|
|
count_of_cache_file = cache_files_cnt()
|
2022-08-23 11:59:04 +08:00
|
|
|
|
2023-03-28 09:47:33 +08:00
|
|
|
def added_files():
|
|
|
|
|
return cache_files_cnt() - count_of_cache_file
|
2022-07-23 19:16:23 +08:00
|
|
|
|
|
|
|
|
def helper1():
|
|
|
|
|
@ti.experimental.real_func
|
|
|
|
|
def sum(l: ti.i32, r: ti.i32) -> ti.i32:
|
|
|
|
|
if l == r:
|
|
|
|
|
return l
|
|
|
|
|
else:
|
|
|
|
|
return sum(l, (l + r) // 2) + sum((l + r) // 2 + 1, r)
|
|
|
|
|
|
|
|
|
|
@ti.kernel
|
|
|
|
|
def get_sum() -> ti.i32:
|
|
|
|
|
return sum(0, 99)
|
|
|
|
|
|
|
|
|
|
assert get_sum() == 99 * 50
|
|
|
|
|
|
|
|
|
|
def helper2():
|
|
|
|
|
@ti.experimental.real_func
|
|
|
|
|
def sum(l: ti.i32, r: ti.i32) -> ti.i32:
|
|
|
|
|
if l == r:
|
|
|
|
|
return l
|
|
|
|
|
else:
|
|
|
|
|
return sum((l + r) // 2 + 1, r) + sum(l, (l + r) // 2)
|
|
|
|
|
|
|
|
|
|
@ti.kernel
|
|
|
|
|
def get_sum() -> ti.i32:
|
|
|
|
|
return sum(0, 99)
|
|
|
|
|
|
|
|
|
|
assert get_sum() == 99 * 50
|
|
|
|
|
|
2023-03-28 09:47:33 +08:00
|
|
|
assert added_files() == expected_num_cache_files()
|
2022-11-16 18:09:13 +08:00
|
|
|
|
|
|
|
|
def my_init():
|
|
|
|
|
ti.init(arch=curr_arch, enable_fallback=False, **{**current_thread_ext_options(), "cuda_stack_limit": 4096})
|
|
|
|
|
|
|
|
|
|
my_init()
|
2022-07-23 19:16:23 +08:00
|
|
|
helper1()
|
|
|
|
|
|
2022-11-16 18:09:13 +08:00
|
|
|
my_init()
|
2023-03-28 09:47:33 +08:00
|
|
|
assert added_files() == expected_num_cache_files(1)
|
2022-07-23 19:16:23 +08:00
|
|
|
helper1()
|
|
|
|
|
|
2022-11-16 18:09:13 +08:00
|
|
|
my_init()
|
2023-03-28 09:47:33 +08:00
|
|
|
assert added_files() == expected_num_cache_files(1)
|
2022-07-23 19:16:23 +08:00
|
|
|
helper2()
|
|
|
|
|
|
2022-11-16 18:09:13 +08:00
|
|
|
my_init()
|
2023-03-28 09:47:33 +08:00
|
|
|
assert added_files() == expected_num_cache_files(2)
|
2022-07-23 19:16:23 +08:00
|
|
|
helper2()
|
|
|
|
|
|
|
|
|
|
ti.reset()
|
2023-03-28 09:47:33 +08:00
|
|
|
assert added_files() == expected_num_cache_files(2)
|
