arguments = parser.parse_args()
func_list_llvm = []
func_list_cuda = []
// TODO: attach instruction features to the library, e.g., avx-512, etc.
for operator_def in __OP_DEF__:
for sch, args, name in operator_def.invoke_all():
After Change
arguments = parser.parse_args()
mod_llvm = tvm.IRModule({})
mod_cuda = tvm.IRModule({})
has_cuda = False
// TODO: attach instruction features to the library, e.g., avx-512, etc.
for operator_def in __OP_DEF__:
for sch, args, name in operator_def.invoke_all():
name = operator_def.get_op_name(name, args)
if tvm.runtime.module.enabled(get_target(operator_def.target)):
func_lower = tvm.lower(sch, args,
name=name,
binds=operator_def.get_binds(args))
if operator_def.target == "cpu":
mod = mod_llvm.update(func_lower)
else:
has_cuda = True
mod_cuda.update(func_lower)
lowered_funcs = {get_target("cpu"): mod_llvm}
if has_cuda > 0:
lowered_funcs[get_target("cuda")] = mod_cuda
