Skip to content

Commit

Permalink
Auto Format
Browse files Browse the repository at this point in the history
  • Loading branch information
@taichi-gardener @squarefk
taichi-gardener authored and squarefk committed Nov 1, 2021
1 parent 01ab4b6 commit 1f7268e
Show file tree
Hide file tree
Showing 4 changed files with 57 additions and 25 deletions.
2 changes: 1 addition & 1 deletion python/taichi/lang/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
from taichi.lang.ops import *
from taichi.lang.quant_impl import quant
from taichi.lang.runtime_ops import async_flush, sync
from taichi.lang.source_builder import SourceBuilder
from taichi.lang.struct import Struct
from taichi.lang.type_factory_impl import type_factory
from taichi.lang.util import (has_clangpp, has_pytorch, is_taichi_class,
Expand All @@ -34,7 +35,6 @@
get_predefined_cupti_metrics)
from taichi.snode.fields_builder import FieldsBuilder
from taichi.type.annotations import any_arr, ext_arr, template
from taichi.lang.source_builder import SourceBuilder

import taichi as ti

Expand Down
76 changes: 54 additions & 22 deletions python/taichi/lang/source_builder.py
View file
Original file line number Diff line number Diff line change
@@ -1,15 +1,16 @@
import atexit
import ctypes
import os
import tempfile
import shutil
import subprocess
import tempfile

from taichi.core.util import ti_core as _ti_core
from taichi.lang import impl
from taichi.lang.exception import TaichiSyntaxError
from taichi.lang.util import get_clangpp, has_clangpp
from taichi.lang.expr import make_expr_group
from taichi.lang import impl
from taichi.core.util import ti_core as _ti_core
from taichi.lang.util import get_clangpp, has_clangpp


class SourceBuilder:
def __init__(self):
Expand All @@ -21,23 +22,35 @@ def __init__(self):
def cleanup():
if self.td is not None:
shutil.rmtree(self.td)

atexit.register(cleanup)

@classmethod
def from_file(cls, filename, compile_fn=None, _temp_dir=None):
self = cls()
self.td = _temp_dir
if filename.endswith((".cpp", ".c", ".cc")):
assert impl.current_cfg().arch in [_ti_core.Arch.x64, _ti_core.Arch.cuda]
assert impl.current_cfg().arch in [
_ti_core.Arch.x64, _ti_core.Arch.cuda
]
if compile_fn is None:
if self.td is None:
self.td = tempfile.mkdtemp()

def compile_fn_impl(filename):
if impl.current_cfg().arch == _ti_core.Arch.x64:
subprocess.call(get_clangpp() + ' -flto -c ' + filename + ' -o ' + os.path.join(self.td, 'source.bc'), shell=True)
subprocess.call(get_clangpp() + ' -flto -c ' +
filename + ' -o ' +
os.path.join(self.td, 'source.bc'),
shell=True)
else:
subprocess.call(get_clangpp() + ' -flto -c ' + filename + ' -o ' + os.path.join(self.td, 'source.bc') + ' -target nvptx64-nvidia-cuda', shell=True)
subprocess.call(get_clangpp() + ' -flto -c ' +
filename + ' -o ' +
os.path.join(self.td, 'source.bc') +
' -target nvptx64-nvidia-cuda',
shell=True)
return os.path.join(self.td, 'source.bc')

compile_fn = compile_fn_impl
self.bc = compile_fn(filename)
self.mode = 'bc'
Expand All @@ -47,11 +60,22 @@ def compile_fn_impl(filename):
if self.td is None:
self.td = tempfile.mkdtemp()
shutil.copy(filename, os.path.join(self.td, 'source.cu'))

def compile_fn_impl(filename):
# Cannot use -o to specify multiple output files
subprocess.call(get_clangpp() + ' ' + os.path.join(self.td, 'source.cu') + ' -S -emit-llvm -std=c++17 --cuda-gpu-arch=sm_50 -nocudalib', cwd=self.td, shell=True)
subprocess.call('llvm-as ' + os.path.join(self.td, 'source-cuda-nvptx64-nvidia-cuda-sm_50.ll'), cwd=self.td, shell=True)
return os.path.join(self.td, 'source-cuda-nvptx64-nvidia-cuda-sm_50.bc')
subprocess.call(
get_clangpp() + ' ' +
os.path.join(self.td, 'source.cu') +
' -S -emit-llvm -std=c++17 --cuda-gpu-arch=sm_50 -nocudalib',
cwd=self.td,
shell=True)
subprocess.call('llvm-as ' + os.path.join(
self.td, 'source-cuda-nvptx64-nvidia-cuda-sm_50.ll'),
cwd=self.td,
shell=True)
return os.path.join(
self.td, 'source-cuda-nvptx64-nvidia-cuda-sm_50.bc')

compile_fn = compile_fn_impl
self.bc = compile_fn(filename)
self.mode = 'bc'
Expand All @@ -60,23 +84,32 @@ def compile_fn_impl(filename):
self.so = ctypes.CDLL(filename)
self.mode = 'so'
elif filename.endswith(".ll"):
assert impl.current_cfg().arch in [_ti_core.Arch.x64, _ti_core.Arch.cuda]
assert impl.current_cfg().arch in [
_ti_core.Arch.x64, _ti_core.Arch.cuda
]
if self.td is None:
self.td = tempfile.mkdtemp()
subprocess.call('llvm-as ' + filename + ' -o ' + os.path.join(self.td, 'source.bc'), shell=True)
subprocess.call('llvm-as ' + filename + ' -o ' +
os.path.join(self.td, 'source.bc'),
shell=True)
self.bc = os.path.join(self.td, 'source.bc')
self.mode = 'bc'
elif filename.endswith(".bc"):
assert impl.current_cfg().arch in [_ti_core.Arch.x64, _ti_core.Arch.cuda]
assert impl.current_cfg().arch in [
_ti_core.Arch.x64, _ti_core.Arch.cuda
]
self.bc = filename
self.mode = 'bc'
else:
raise TaichiSyntaxError('Unsupported file type for external function call.')
raise TaichiSyntaxError(
'Unsupported file type for external function call.')
return self

@classmethod
def from_source(cls, source_code, compile_fn=None):
assert impl.current_cfg().arch in [_ti_core.Arch.x64, _ti_core.Arch.cuda]
assert impl.current_cfg().arch in [
_ti_core.Arch.x64, _ti_core.Arch.cuda
]
_temp_dir = tempfile.mkdtemp()
_temp_source = os.path.join(_temp_dir, '_temp_source.cpp')
with open(_temp_source, 'w') as f:
Expand All @@ -86,19 +119,18 @@ def from_source(cls, source_code, compile_fn=None):
def __getattr__(self, item):
def bitcode_func_call_wrapper(*args):
_ti_core.insert_external_func_call(0, '', self.bc, item,
make_expr_group(args),
make_expr_group([]))

make_expr_group(args),
make_expr_group([]))

if self.mode == 'bc':
return bitcode_func_call_wrapper

def external_func_call_wrapper(args=[], outputs=[]):
func_addr = ctypes.cast(self.so.__getattr__(item), ctypes.c_void_p).value
func_addr = ctypes.cast(self.so.__getattr__(item),
ctypes.c_void_p).value
_ti_core.insert_external_func_call(func_addr, '', '', '',
make_expr_group(args),
make_expr_group(outputs))

make_expr_group(args),
make_expr_group(outputs))

if self.mode == 'so':
return external_func_call_wrapper
Expand Down
3 changes: 2 additions & 1 deletion taichi/codegen/codegen_llvm.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2033,7 +2033,8 @@ void CodeGenLLVM::visit_call_bitcode(ExternalFuncCallStmt *stmt) {
auto *func_ptr = external_module->getFunction(stmt->bc_funcname);
TI_ASSERT_INFO(func_ptr != nullptr, "{} is not found in {}.",
stmt->bc_funcname, stmt->bc_filename);
auto link_error = llvm::Linker::linkModules(*module, std::move(external_module));
auto link_error =
llvm::Linker::linkModules(*module, std::move(external_module));
TI_ASSERT(!link_error);
}
// Retrieve function again. Do it here to detect name conflicting.
Expand Down
1 change: 0 additions & 1 deletion tests/python/test_external_func.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -85,4 +85,3 @@ def func_bc() -> ti.i32:
assert func_bc() == 11**8

shutil.rmtree(td)

0 comments on commit 1f7268e

Please sign in to comment.