Avoid duplicate evaluations in chaining comparison

python/taichi/lang/impl.py

@@ -36,6 +36,7 @@ def wrap_scalar(x):
 def atomic_add(a, b):
   return a.atomic_add(b)
 
+
 def subscript(value, *indices):
   import numpy as np
   if isinstance(value, np.ndarray):
@@ -66,6 +67,34 @@ def subscript(value, *indices):
     return Expr(taichi_lang_core.subscript(value.ptr, indices_expr_group))
 
 
+def chain_compare(comparators, ops):
+  assert len(comparators) == len(ops) + 1, \
+    f'Chain comparison invoked with {len(comparators)} comparators but {len(ops)} operators'
+  evaluated_comparators = []
+  for i in range(len(comparators)):
+    evaluated_comparators += [expr_init(comparators[i])]
+  ret = expr_init(True)
+  for i in range(len(ops)):
+    lhs = evaluated_comparators[i]
+    rhs = evaluated_comparators[i + 1]
+    if ops[i] == 'Lt':
+      now = lhs < rhs
+    elif ops[i] == 'LtE':
+      now = lhs <= rhs
+    elif ops[i] == 'Gt':
+      now = lhs > rhs
+    elif ops[i] == 'GtE':
+      now = lhs >= rhs
+    elif ops[i] == 'Eq':
+      now = lhs == rhs
+    elif ops[i] == 'NotEq':
+      now = lhs != rhs
+    else:
+      assert False, f'Unknown operator {ops[i]}'
+    ret = ret.logical_and(now)
+  return ret
+
+
 class PyTaichi:
 
   def __init__(self, kernels=[]):

python/taichi/lang/transformer.py

@@ -427,7 +427,12 @@ def visit_Continue(self, node):
     )
 
   def visit_Call(self, node):
-    self.generic_visit(node)
+    if not (isinstance(node.func, ast.Attribute)
+            and isinstance(node.func.value, ast.Name)
+            and node.func.value.id == 'ti'
+            and node.func.attr == 'static'):
+      # Do not apply the generic visitor if the function called is ti.static
+      self.generic_visit(node)
     if isinstance(node.func, ast.Name):
       func_name = node.func.id
       if func_name == 'print':
@@ -524,21 +529,40 @@ def visit_UnaryOp(self, node):
 
   def visit_Compare(self, node):
     self.generic_visit(node)
-    ret = None
     comparators = [node.left] + node.comparators
-    for i in range(len(node.comparators)):
-      new_cmp = ast.Compare(left=comparators[i], ops=[node.ops[i]], comparators=[comparators[i + 1]])
-      ast.copy_location(new_cmp, node)
-      if ret is None:
-        ret = new_cmp
+    ops = []
+    for i in range(len(node.ops)):
+      if isinstance(node.ops[i], ast.Lt):
+        op_str = 'Lt'
+      elif isinstance(node.ops[i], ast.LtE):
+        op_str = 'LtE'
+      elif isinstance(node.ops[i], ast.Gt):
+        op_str = 'Gt'
+      elif isinstance(node.ops[i], ast.GtE):
+        op_str = 'GtE'
+      elif isinstance(node.ops[i], ast.Eq):
+        op_str = 'Eq'
+      elif isinstance(node.ops[i], ast.NotEq):
+        op_str = 'NotEq'
+      elif isinstance(node.ops[i], ast.In):
+        raise TaichiSyntaxError('"in" is not supported in Taichi kernels.')
+      elif isinstance(node.ops[i], ast.NotIn):
+        raise TaichiSyntaxError('"not in" is not supported in Taichi kernels.')
+      elif isinstance(node.ops[i], ast.Is):
+        raise TaichiSyntaxError('"is" is not supported in Taichi kernels.')
+      elif isinstance(node.ops[i], ast.IsNot):
+        raise TaichiSyntaxError('"is not" is not supported in Taichi kernels.')
       else:
-        ret = ast.BoolOp(op=ast.And(), values=[ret, new_cmp])
-        ret = self.visit_BoolOp(ret)
-        ast.copy_location(ret, node)
-
-    self.generic_visit(ret)
-    return ret
-
+        raise Exception(f'Unknown operator {node.ops[i]}')
+      ops += [ast.copy_location(ast.Str(s=op_str), node)]
+
+    call = ast.Call(
+      func=self.parse_expr('ti.chain_compare'),
+      args=[ast.copy_location(ast.List(elts=comparators, ctx=ast.Load()), node),
+            ast.copy_location(ast.List(elts=ops, ctx=ast.Load()), node)],
+      keywords=[])
+    call = ast.copy_location(call, node)
+    return call
 
   def visit_BoolOp(self, node):
     self.generic_visit(node)

tests/python/test_compare.py

@@ -0,0 +1,116 @@
+import taichi as ti
+
+
+@ti.all_archs
+def test_compare_basics():
+  a = ti.var(ti.i32)
+  ti.root.dynamic(ti.i, 256).place(a)
+  b = ti.var(ti.i32, shape=())
+  c = ti.var(ti.i32, shape=())
+
+  @ti.kernel
+  def func():
+    b[None] = 3
+    c[None] = 5
+    a[0] = b < c
+    a[1] = b <= c
+    a[2] = b > c
+    a[3] = b >= c
+    a[4] = b == c
+    a[5] = b != c
+    a[6] = c < b
+    a[7] = c <= b
+    a[8] = c > b
+    a[9] = c >= b
+    a[10] = c == b
+    a[11] = c != b
+    # a[12] = b is not c  # not supported
+
+  func()
+  assert a[0]
+  assert a[1]
+  assert not a[2]
+  assert not a[3]
+  assert not a[4]
+  assert a[5]
+  assert not a[6]
+  assert not a[7]
+  assert a[8]
+  assert a[9]
+  assert not a[10]
+  assert a[11]
+
+
+@ti.all_archs
+def test_compare_equality():
+  a = ti.var(ti.i32)
+  ti.root.dynamic(ti.i, 256).place(a)
+  b = ti.var(ti.i32, shape=())
+  c = ti.var(ti.i32, shape=())
+
+  @ti.kernel
+  def func():
+    b[None] = 3
+    c[None] = 3
+    a[0] = b < c
+    a[1] = b <= c
+    a[2] = b > c
+    a[3] = b >= c
+    a[4] = b == c
+    a[5] = b != c
+    a[6] = c < b
+    a[7] = c <= b
+    a[8] = c > b
+    a[9] = c >= b
+    a[10] = c == b
+    a[11] = c != b
+
+  func()
+  assert not a[0]
+  assert a[1]
+  assert not a[2]
+  assert a[3]
+  assert a[4]
+  assert not a[5]
+  assert not a[6]
+  assert a[7]
+  assert not a[8]
+  assert a[9]
+  assert a[10]
+  assert not a[11]
+
+
+@ti.all_archs
+def test_no_duplicate_eval():
+  a = ti.var(ti.i32)
+  ti.root.dynamic(ti.i, 256).place(a)
+
+  @ti.kernel
+  def func():
+    a[2] = 0 <= ti.append(a.parent(), [], 10) < 1
+
+  func()
+  assert a[0] == 10
+  assert a[1] == 0  # not appended twice
+  assert a[2]  # ti.append returns 0
+
+
+@ti.all_archs
+def test_chain_compare():
+  a = ti.var(ti.i32)
+  ti.root.dynamic(ti.i, 256).place(a)
+  b = ti.var(ti.i32, shape=())
+  c = ti.var(ti.i32, shape=())
+  d = ti.var(ti.i32, shape=())
+
+  @ti.kernel
+  def func():
+    b[None] = 2
+    c[None] = 3
+    d[None] = 3
+    a[0] = c == d != b < d > b >= b <= c
+    a[1] = b <= c != d > b == b
+
+  func()
+  assert a[0]
+  assert not a[1]