Multi-threaded unit testing by pytest

python/taichi/main.py

@@ -20,13 +20,27 @@ def test_python(test_files=None, verbose=False):
   if test_files:
     # run individual tests
     for f in test_files:
+      # auto-compelete
+      if not f.startswith('test_'):
+        f = 'test_' + f
+      if not f.endswith('.py'):
+        f = f + '.py'
       args.append(os.path.join(test_dir, f))
   else:
     # run all the tests
     args = [test_dir]
   if verbose:
     args += ['-s']
+  if not test_files or len(test_files) > 4:
+    try:
+      from multiprocessing import cpu_count
+      cpu_count = cpu_count()
+    except:
+      cpu_count = 4
+    print(f'Starting {cpu_count} testing thread(s)...')
+    args += ['-n', str(cpu_count)]
 
+  print('XXX')
   return int(pytest.main(args))
 
 
@@ -104,11 +118,11 @@ def main(debug=False):
   elif mode == "test_cpp":
     return test_cpp()
   elif mode == "test":
-    if test_python() != 0:
+    if test_python(test_files=sys.argv[2:]) != 0 or len(sys.argv) > 2:
       return -1
     return test_cpp()
   elif mode == "test_verbose":
-    if test_python(verbose=True) != 0:
+    if test_python(test_files=sys.argv[2:], verbose=True) != 0:
       return -1
     return test_cpp()
   elif mode == "build":

taichi/transforms/make_adjoint.cpp

@@ -63,6 +63,14 @@ class MakeAdjoint : public IRVisitor {
     return insert<UnaryOpStmt>(UnaryOpType::sin, load(op1));
   }
 
+  Stmt *log(Stmt *op1) {
+    return insert<UnaryOpStmt>(UnaryOpType::log, load(op1));
+  }
+
+  Stmt *pow(Stmt *op1, Stmt *op2) {
+    return insert<BinaryOpStmt>(BinaryOpType::pow, load(op1), load(op2));
+  }
+
  public:
   Block *current_block;
   int for_depth;
@@ -189,6 +197,7 @@ class MakeAdjoint : public IRVisitor {
       accumulate(bin->lhs, adjoint(bin));
       accumulate(bin->rhs, negate(adjoint(bin)));
     } else if (bin->op_type == BinaryOpType::mul) {
+      // d (x * y) = y * dx + x * dy
       accumulate(bin->lhs, mul(adjoint(bin), bin->rhs));
       accumulate(bin->rhs, mul(adjoint(bin), bin->lhs));
     } else if (bin->op_type == BinaryOpType::mod) {
@@ -201,6 +210,11 @@ class MakeAdjoint : public IRVisitor {
       auto numerator = add(sqr(bin->lhs), sqr(bin->rhs));
       accumulate(bin->lhs, div(mul(adjoint(bin), bin->rhs), numerator));
       accumulate(bin->rhs, negate(div(mul(adjoint(bin), bin->lhs), numerator)));
+    } else if (bin->op_type == BinaryOpType::pow) {
+      // d (x ^ y) = y * x ^ (y-1)  dx  +  log(x) * x ^ y  dy
+      accumulate(bin->lhs, mul(adjoint(bin),
+            mul(pow(bin->lhs, sub(bin->rhs, constant(1))), bin->rhs)));
+      accumulate(bin->rhs, mul(adjoint(bin), mul(log(bin->lhs), pow(bin->lhs, bin->rhs))));
     } else if (bin->op_type == BinaryOpType::min ||
                bin->op_type == BinaryOpType::max) {
       auto cmp = bin->op_type == BinaryOpType::min ? cmp_lt(bin->lhs, bin->rhs)