Currently unittests are slow when using flashinfer jit because we only compile kernels the first time we run it, it's blocking and didn't compile multiple ops in parallel. This PR add a warmup pre-hook to kernel unittests, so that we compile all necessary kernels before running the unittests in JIT mode, which greatly accelerate the unittests.
This PR also fixes the several issues with #628 :
using thread-safe make_dirs(..., exist_ok=True) instead of relying on os.path.exists
change the signature of parallel_load_modules to lists of (jit_module_creation_func, args) instead of lambda function, because lambda function captures variable by ref instead of value, which may cause some unexpected errors.
Currently unittests are slow when using flashinfer jit because we only compile kernels the first time we run it, it's blocking and didn't compile multiple ops in parallel. This PR add a warmup pre-hook to kernel unittests, so that we compile all necessary kernels before running the unittests in JIT mode, which greatly accelerate the unittests.
This PR also fixes the several issues with #628 :
make_dirs(..., exist_ok=True)instead of relying onos.path.existsparallel_load_modulesto lists of(jit_module_creation_func, args)instead of lambda function, because lambda function captures variable by ref instead of value, which may cause some unexpected errors.