File size: 1,163 Bytes
27813b0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 | """Load resample-poly natively on hosts without a published variant.
import load_local
rp = load_local.load()
"""
import os
import sys
import types
_cached = None
def load(verbose=False):
"""Return the resample_poly module, JIT-built from the local source."""
global _cached
if _cached is not None:
return _cached
from torch.utils.cpp_extension import load as _jit
root = os.path.dirname(os.path.abspath(__file__))
cflags = ["/O2"] if os.name == "nt" else ["-O3"]
ext = _jit(name="resample_poly_ext",
sources=[os.path.join(root, "local_bind.cpp"),
os.path.join(root, "resample_csrc", "resample.cpp")],
extra_include_paths=[os.path.join(root, "torch-ext")],
extra_cflags=cflags,
verbose=verbose)
ops_mod = types.ModuleType("resample_poly._ops")
ops_mod.ops = type("_Ops", (), {
"rp_resample": staticmethod(ext.rp_resample),
})()
sys.modules["resample_poly._ops"] = ops_mod
sys.path.insert(0, os.path.join(root, "torch-ext"))
import resample_poly
_cached = resample_poly
return resample_poly
|