Skip to content

Commit

Permalink
kernel templates working
Browse files Browse the repository at this point in the history
  • Loading branch information
yuanming-hu committed Oct 22, 2019
1 parent 0fb4477 commit ac7585b
Show file tree
Hide file tree
Showing 2 changed files with 14 additions and 12 deletions.
18 changes: 12 additions & 6 deletions python/taichi/lang/impl.py
Original file line number Diff line number Diff line change
Expand Up @@ -284,8 +284,8 @@ class Kernel:
def __init__(self, func, is_grad):
self.func = func
self.is_grad = is_grad
self.materialized = False
self.arguments = []
self.argument_names = []
self.extract_arguments()
self.template_slot_locations = []
for i in range(len(self.arguments)):
Expand Down Expand Up @@ -320,16 +320,15 @@ def extract_arguments(self):
if param.annotation is inspect.Parameter.empty:
raise KernelDefError('Taichi kernels parameters must be type annotated')
self.arguments.append(param.annotation)
self.argument_names.append(param.name)

def materialize(self, key=None, args=None, extra_frame_backtrace=-1):
if key is None:
key = (self.func, 0)
if not self.materialized:
self.materialized = True
else:
return
if not pytaichi.materialized:
pytaichi.materialize()
if key in self.compiled_functions:
return
grad_suffix = ""
if self.is_grad:
grad_suffix = "_grad"
Expand Down Expand Up @@ -357,8 +356,15 @@ def materialize(self, key=None, args=None, extra_frame_backtrace=-1):
frame = inspect.currentframe()
for t in range(extra_frame_backtrace + 2):
frame = frame.f_back
globals = dict(frame.f_globals, **frame.f_locals)
# inject template parameters into globals

for i in self.template_slot_locations:
template_var_name = self.argument_names[i]
globals[template_var_name] = args[i]

exec(compile(tree, filename=inspect.getsourcefile(self.func), mode='exec'),
dict(frame.f_globals, **frame.f_locals), locals())
globals, locals())
pytaichi.inside_kernel = False
compiled = locals()[self.func.__name__]

Expand Down
8 changes: 2 additions & 6 deletions tests/python/test_kernel_templates.py
Original file line number Diff line number Diff line change
@@ -1,9 +1,7 @@
import taichi as ti

def test_kernel_template_mapper():
return
ti.reset()
ti.get_runtime().print_preprocessed = True
x = ti.var(ti.i32)
y = ti.var(ti.f32)

Expand All @@ -19,10 +17,8 @@ def inc(a: ti.template(), b: ti.template()):
a[i] += b

inc(x, 1)
# inc(y, 2)
inc(y, 2)

for i in range(n):
assert x[i] == 1
assert y[i] == 1

test_kernel_template_mapper()
assert y[i] == 2

0 comments on commit ac7585b

Please sign in to comment.