From bc9ea32f44c49413a31c4a88846a1d7437219910 Mon Sep 17 00:00:00 2001
From: Yuanming Hu <yuanmhu@gmail.com>
Date: Wed, 9 Oct 2019 00:17:18 -0400
Subject: [PATCH] torch_tensor ad within taichi

---
 examples/{ => difftaichi}/volume_renderer.py |  0
 examples/torch_tensor_ad.py                  | 77 ++++++++++++++++++++
 examples/torch_tensor_io.py                  | 23 +++---
 3 files changed, 87 insertions(+), 13 deletions(-)
 rename examples/{ => difftaichi}/volume_renderer.py (100%)
 create mode 100644 examples/torch_tensor_ad.py

diff --git a/examples/volume_renderer.py b/examples/difftaichi/volume_renderer.py
similarity index 100%
rename from examples/volume_renderer.py
rename to examples/difftaichi/volume_renderer.py
diff --git a/examples/torch_tensor_ad.py b/examples/torch_tensor_ad.py
new file mode 100644
index 0000000000000..b44d4d6a97657
--- /dev/null
+++ b/examples/torch_tensor_ad.py
@@ -0,0 +1,77 @@
+import taichi as ti
+import numpy as np
+import torch
+
+# ti.set_gdb_trigger(True)
+ti.cfg.arch = ti.cuda
+
+# n = 1024 * 1024
+n = 32
+
+x = ti.var(ti.f32)
+y = ti.var(ti.f32)
+
+# https://pytorch.org/tutorials/beginner/examples_autograd/two_layer_net_custom_function.html
+
+@ti.layout
+def values():
+  # actually useless in thie example
+  ti.root.dense(ti.i, n).place(x)
+  ti.root.dense(ti.i, n).place(y)
+  ti.root.lazy_grad()
+
+@ti.kernel
+def torch_kernel():
+  for i in range(n):
+    y[i] = x[i] * x[i]
+    
+    
+def copy_from(taichi_tensor):
+  @ti.kernel
+  def ker(torch_tensor: np.ndarray):
+    for i in taichi_tensor:
+      taichi_tensor[i] = torch_tensor[i]
+      
+  ker.materialize()
+  return lambda x: ker(x.contiguous())
+
+def copy_to(taichi_tensor):
+  @ti.kernel
+  def ker(torch_tensor: np.ndarray):
+    for i in taichi_tensor:
+      torch_tensor[i] = taichi_tensor[i]
+  
+  ker.materialize()
+  return lambda x: ker(x.contiguous())
+  
+x_copy_from = copy_from(x)
+y_copy_to = copy_to(y)
+
+y_grad_copy_from = copy_from(y.grad)
+x_grad_copy_to = copy_to(x.grad)
+
+class Sqr(torch.autograd.Function):
+  @staticmethod
+  def forward(ctx, inp):
+    outp = torch.zeros_like(inp)
+    x_copy_from(inp)
+    torch_kernel()
+    y_copy_to(outp)
+    return outp
+  
+  @staticmethod
+  def backward(ctx, outp_grad):
+    ti.clear_all_gradients()
+    inp_grad = torch.zeros_like(outp_grad)
+    
+    y_grad_copy_from(outp_grad)
+    torch_kernel.grad()
+    x_grad_copy_to(inp_grad)
+    
+    return inp_grad
+
+sqr = Sqr.apply
+X = torch.tensor(2 * np.ones((n, ), dtype=np.float32), device=torch.device('cuda:0'), requires_grad=True)
+sqr(X).sum().backward()
+print(X.grad.cpu())
+
diff --git a/examples/torch_tensor_io.py b/examples/torch_tensor_io.py
index c57e1fd044895..2621176e34ced 100644
--- a/examples/torch_tensor_io.py
+++ b/examples/torch_tensor_io.py
@@ -8,29 +8,28 @@
 # n = 1024 * 1024
 n = 32
 
-y = ti.var(ti.i32)
+y = ti.var(ti.f32)
 
 # https://pytorch.org/tutorials/beginner/examples_autograd/two_layer_net_custom_function.html
 
-z = np.array((n,), dtype=np.float32)
-
 @ti.layout
 def values():
-  ti.root.place(y)
-
+  # actually useless in thie example
+  ti.root.dense(ti.i, n).place(y)
+  ti.root.lazy_grad()
 
 @ti.kernel
 def torch_kernel(t: np.ndarray, o: np.ndarray):
   for i in range(n):
     o[i] = t[i] * t[i]
-    
+
 @ti.kernel
 def torch_kernel_2(t_grad: np.ndarray, t:np.ndarray, o_grad: np.ndarray):
   for i in range(n):
     ti.print(o_grad[i])
     t_grad[i] = 2 * t[i] * o_grad[i]
-  
-  
+
+
 class Sqr(torch.autograd.Function):
   @staticmethod
   def forward(ctx, inp):
@@ -41,7 +40,6 @@ def forward(ctx, inp):
   
   @staticmethod
   def backward(ctx, outp_grad):
-    print(outp_grad.cpu())
     outp_grad = outp_grad.contiguous()
     inp_grad = torch.zeros_like(outp_grad)
     inp, = ctx.saved_tensors
@@ -49,8 +47,7 @@ def backward(ctx, outp_grad):
     return inp_grad
 
 sqr = Sqr.apply
-x = torch.tensor(2 * np.ones((n, ), dtype=np.float32), device=torch.device('cuda:0'), requires_grad=True)
-sqr(x).sum().backward()
-# print(sqr(x).sum())#.backward()
-print(x.grad.cpu())
+X = torch.tensor(2 * np.ones((n, ), dtype=np.float32), device=torch.device('cuda:0'), requires_grad=True)
+sqr(X).sum().backward()
+print(X.grad.cpu())