def get_module():
from TFUtil import OpCodeCompiler
cc_files = [
"batch_norm_op.cc",
"blocksparse_l2_norm_op.cc", "blocksparse_matmul_op.cc",
"cwise_linear_op.cc", "edge_bias_op.cc", "ew_op.cc", "gpu_types.cc", "layer_norm_op.cc"
]
# not "blocksparse_kernels.cc" (SASS kernels) for now.
# "blocksparse_conv_op.cc" depends on that.
cu_files = [
"batch_norm_op_gpu.cu", "blocksparse_l2_norm_op_gpu.cu", "blocksparse_matmul_op_gpu.cu",
"cwise_linear_op_gpu.cu", "edge_bias_op_gpu.cu", "ew_op_gpu.cu",
"layer_norm_cn_op_gpu.cu", "layer_norm_nc_op_gpu.cu"
]
src_dir = os.path.abspath(os.path.dirname(__file__) + "/../src")
assert os.path.isdir(src_dir)
src_code = ""
for fn in cc_files + cu_files:
src_code += "\n// ----- %s BEGIN -----\n" % fn
src_code += open("%s/%s" % (src_dir, fn)).read()
src_code += "\n// ----- %s END -----\n\n\n" % fn
src_code += """
// custom code
Status GetKernel(std::string& kernel_name, CUfunction* kernel) {
if(*kernel) return Status::OK(); // Only need to get kernel once.
return errors::Internal("Blocksparse GetKernel without SASS: kernel ", kernel_name, " not available");
}
"""
compiler = OpCodeCompiler(base_name="blocksparse", code=src_code, code_version=1, include_paths=[src_dir])
return tf.load_op_library(compiler.get_lib_filename())
def get_tf_mod(verbose=False):
"""
:param bool verbose:
:return: module
"""
global _tf_mod
if _tf_mod:
return _tf_mod
from glob import glob
from TFUtil import OpCodeCompiler
# Collect files.
assert openfst_checked_out(
), "submodule in %r not checked out?" % openfst_dir
files = glob('%s/src/lib/*.cc' % openfst_dir)
assert files, "submodule in %r not checked out?" % openfst_dir
files = sorted(files) # make somewhat deterministic
libs = []
if platform.system() != 'Darwin':
libs.append('rt')
# Put code all together in one big blob.
src_code = ""
for fn in files:
f_code = open(fn).read()
f_code = ''.join([x for x in f_code if ord(x) < 128]) # enforce ASCII
src_code += "\n// ------------ %s : BEGIN { ------------\n" % os.path.basename(
fn)
# https://gcc.gnu.org/onlinedocs/cpp/Line-Control.html#Line-Control
src_code += "#line 1 \"%s\"\n" % os.path.basename(fn)
src_code += f_code
src_code += "\n// ------------ %s : END } --------------\n\n" % os.path.basename(
fn)
src_code += "\n\n// ------------ our code now: ------------\n\n"
src_code += "#line 1 \"TFOpenFst.py:_src_code\"\n"
src_code += _src_code
compiler = OpCodeCompiler(
base_name="OpenFst",
code_version=1,
code=src_code,
include_paths=("%s/src/include" % openfst_dir, ),
c_macro_defines={
"NDEBUG":
1, # https://github.com/tensorflow/tensorflow/issues/17316
},
ld_flags=["-l%s" % lib for lib in libs],
is_cpp=True,
use_cuda_if_available=False,
verbose=verbose)
tf_mod = compiler.load_tf_module()
assert hasattr(
tf_mod, "open_fst_transition"), "content of mod: %r" % (dir(tf_mod), )
_tf_mod = tf_mod
return tf_mod
def get_tf_mod(verbose=False):
"""
:param bool verbose:
:return: module
"""
global _tf_mod
if _tf_mod:
return _tf_mod
import platform
from glob import glob
from TFUtil import OpCodeCompiler
# References:
# https://github.com/kpu/kenlm/blob/master/setup.py
# https://github.com/kpu/kenlm/blob/master/compile_query_only.sh
# Collect files.
files = glob('%s/util/*.cc' % kenlm_dir)
files += glob('%s/lm/*.cc' % kenlm_dir)
files += glob('%s/util/double-conversion/*.cc' % kenlm_dir)
files = [fn for fn in files if not (fn.endswith('main.cc') or fn.endswith('test.cc'))]
assert files, "submodule in %r not checked out?" % kenlm_dir
libs = ["z"]
if platform.system() != 'Darwin':
libs.append('rt')
# Put code all together in one big blob.
src_code = ""
src_code += _kenlm_src_code_workarounds
for fn in files:
f_code = open(fn).read()
f_code = ''.join([x for x in f_code if ord(x) < 128]) # enforce ASCII
# We need to do some replacements to not clash symbol names.
fn_short = os.path.basename(fn).replace(".", "_")
for word in ["kConverter"]:
f_code = f_code.replace(word, "%s_%s" % (fn_short, word))
src_code += "\n// ------------ %s : BEGIN { ------------\n" % os.path.basename(fn)
# https://gcc.gnu.org/onlinedocs/cpp/Line-Control.html#Line-Control
src_code += "#line 1 \"%s\"\n" % os.path.basename(fn)
src_code += f_code
src_code += "\n// ------------ %s : END } --------------\n\n" % os.path.basename(fn)
src_code += "\n\n// ------------ our code now: ------------\n\n"
src_code += _src_code
compiler = OpCodeCompiler(
base_name="KenLM", code_version=1, code=src_code,
include_paths=(kenlm_dir, kenlm_dir + "/util/double-conversion"),
c_macro_defines={"NDEBUG": 1, "KENLM_MAX_ORDER": 6, "HAVE_ZLIB": 1},
ld_flags=["-l%s" % lib for lib in libs],
is_cpp=True, use_cuda_if_available=False,
verbose=verbose)
tf_mod = compiler.load_tf_module()
assert hasattr(tf_mod, "ken_lm_abs_score_strings"), "content of mod: %r" % (dir(tf_mod),)
_tf_mod = tf_mod
return tf_mod
def init_warprnnt(verbose=False):
"""
Initialiazes and compiles the library. Caches the TF module.
:param bool verbose:
"""
global _tf_mod
if _tf_mod:
return
assert is_checked_out(
), "submodule not checked out? Run `git submodule update --init --recursive`"
# References:
# https://github.com/HawkAaron/warp-transducer/blob/master/tensorflow_binding/setup.py
src_files = [
'%s/src/rnnt_entrypoint.cpp' % submodule_dir,
'%s/tensorflow_binding/src/warprnnt_op.cc' % submodule_dir
]
assert all([os.path.isfile(f) for f in src_files
]), "submodule in %r not checked out?" % warprnnt_dir
src_code = ""
for fn in src_files:
f_code = open(fn).read()
src_code += "\n// ------------ %s : BEGIN { ------------\n" % os.path.basename(
fn)
# https://gcc.gnu.org/onlinedocs/cpp/Line-Control.html#Line-Control
src_code += "#line 1 \"%s\"\n" % os.path.basename(fn)
src_code += f_code
src_code += "\n// ------------ %s : END } --------------\n\n" % os.path.basename(
fn)
compiler = OpCodeCompiler(base_name="warprnnt_kernels",
code_version=1,
code=src_code,
include_paths=(submodule_dir + "/include", ),
c_macro_defines={"WITH_OMP": 1},
ld_flags=["-Xcompiler", "-fopenmp"],
is_cpp=True,
use_cuda_if_available=True,
verbose=verbose)
tf_mod = compiler.load_tf_module()
assert hasattr(tf_mod, "WarpRNNT"), "content of mod: %r" % (dir(tf_mod), )
_tf_mod = tf_mod
return tf_mod
def get_tf_mod(verbose=False):
"""
:param bool verbose:
:return: module
"""
global _tf_mod
if _tf_mod:
return _tf_mod
import platform
from glob import glob
from TFUtil import OpCodeCompiler
# References:
# https://github.com/kpu/kenlm/blob/master/setup.py
# https://github.com/kpu/kenlm/blob/master/compile_query_only.sh
# Collect files.
files = glob('%s/util/*.cc' % kenlm_dir)
files += glob('%s/lm/*.cc' % kenlm_dir)
files += glob('%s/util/double-conversion/*.cc' % kenlm_dir)
files = [
fn for fn in files
if not (fn.endswith('main.cc') or fn.endswith('test.cc'))
]
assert files, "submodule in %r not checked out?" % kenlm_dir
libs = ["z"]
if platform.system() != 'Darwin':
libs.append('rt')
# Put code all together in one big blob.
src_code = ""
src_code += _kenlm_src_code_workarounds
for fn in files:
f_code = open(fn).read()
f_code = ''.join([x for x in f_code if ord(x) < 128]) # enforce ASCII
# We need to do some replacements to not clash symbol names.
fn_short = os.path.basename(fn).replace(".", "_")
for word in ["kConverter"]:
f_code = f_code.replace(word, "%s_%s" % (fn_short, word))
src_code += "\n// ------------ %s : BEGIN { ------------\n" % os.path.basename(
fn)
# https://gcc.gnu.org/onlinedocs/cpp/Line-Control.html#Line-Control
src_code += "#line 1 \"%s\"\n" % os.path.basename(fn)
src_code += f_code
src_code += "\n// ------------ %s : END } --------------\n\n" % os.path.basename(
fn)
src_code += "\n\n// ------------ our code now: ------------\n\n"
src_code += _src_code
compiler = OpCodeCompiler(base_name="KenLM",
code_version=1,
code=src_code,
include_paths=(kenlm_dir, kenlm_dir +
"/util/double-conversion"),
c_macro_defines={
"NDEBUG": 1,
"KENLM_MAX_ORDER": 6,
"HAVE_ZLIB": 1
},
ld_flags=["-l%s" % lib for lib in libs],
is_cpp=True,
use_cuda_if_available=False,
verbose=verbose)
tf_mod = compiler.load_tf_module()
assert hasattr(
tf_mod,
"ken_lm_abs_score_strings"), "content of mod: %r" % (dir(tf_mod), )
_tf_mod = tf_mod
return tf_mod