C++ _jit_tensorwise - first prototype of C++ implementation of tensorwise (#22)

Author: cpuhrsch

benchmarks/jit_tensorwise.py

@@ -0,0 +1,68 @@
+import torch
+import nestedtensor
+import utils
+import time
+
+
+@nestedtensor._C.jit_tensorwise()
+@torch.jit.script
+def f(i, w):
+    return torch.conv2d(i, w)
+
+def loop_f(inp1, w):
+    for inp in inp1:
+        torch.conv2d(inp, w)
+
+
+if __name__ == "__main__":
+    w = torch.randn(64, 3, 9, 9).cuda()
+    inp1 = list(torch.randn(128, 1, 3, 16, 16).cuda().unbind())
+    inp3 = nestedtensor.as_nested_tensor(inp1)._impl
+    # print(sum(inp.numel() for inp in inp1))
+    # print(inp3.numel())
+
+    fc = nestedtensor._C.jit_tensorwise()(torch.conv2d)
+
+    t0 = time.time()
+    count = 0
+    while(time.time() - t0 < 5.0):
+        r2 = fc(inp3, w)
+        torch.cuda.synchronize()
+        count += 1
+    print("jit: " + str(count))
+
+    t0 = time.time()
+    count = 0
+    while(time.time() - t0 < 5.0):
+        loop_f(inp1, w)
+        torch.cuda.synchronize()
+        count += 1
+    print("for loop: " + str(count))
+
+    
+    # print(r.nested_size())
+
+    # na = nestedtensor._C.jit_tensorwise()(torch.mul)
+
+    # print("111")
+    # out = nestedtensor.as_nested_tensor([torch.randn(1, 2)])
+    # print(na(
+    #     nestedtensor.as_nested_tensor([torch.randn(1, 2)])._impl,
+    #     4.0,
+    #     ))
+    # print("222")
+    # print('out')
+    # print(out)
+
+    # nv = nestedtensor._C.jit_tensorwise()(torch.mv)
+    # print(nv(
+    #     nestedtensor._C._ListNestedTensor([torch.randn(1, 2)]),
+    #     nestedtensor._C._ListNestedTensor([torch.randn(2)]),
+    #     ))
+
+    # print("333")
+    # print(na(
+    #     torch.randn(1, 2),
+    #     torch.randn(1, 2),
+    #     ))
+    # print("444")

benchmarks/nearest_neighbors.py

@@ -1,11 +1,11 @@
-from nestedtensor import torch
 import nestedtensor
+import torch
 import argparse
 import time
 import random
 import pprint
 
-EMBED_DIM = 1024
+EMBED_DIM = 128
 
 SEED = 0
 
@@ -60,8 +60,8 @@ def gen_algorithm_nested_mv(keys, sub_clusters):
     for sub_cluster in sub_clusters:
         new_sub_cluster = [torch.tensor(list(map(list, cluster))) for cluster in sub_cluster]
         new_sub_clusters.append(new_sub_cluster)
-    nested_sub_clusters = torch.nested_tensor(sub_clusters).to_tensor(2)
-    nested_keys = torch.nested_tensor(keys)
+    nested_sub_clusters = nestedtensor.nested_tensor(sub_clusters).to_tensor(2)
+    nested_keys = nestedtensor.nested_tensor(keys)
     def _nested_mv():
         return torch.mv(nested_sub_clusters, nested_keys)
     return _nested_mv
@@ -74,18 +74,16 @@ def gen_algorithm_nested_jit_mv(keys, sub_clusters):
         for cluster in sub_cluster:
             new_sub_cluster.append(torch.stack(cluster))
         new_sub_clusters.append(new_sub_cluster)
-    nested_sub_clusters = nestedtensor._ListNestedTensor(new_sub_clusters)
-    print("HERE")
-    print(nested_sub_clusters.nested_size())
-    nested_keys = nestedtensor._ListNestedTensor(keys)
-    print(nested_keys.nested_size())
+    nested_sub_clusters = nestedtensor.as_nested_tensor(new_sub_clusters)
+    nested_keys = nestedtensor.as_nested_tensor(keys)
 
+    @nestedtensor._C.jit_tensorwise()
     @torch.jit.script
     def my_fun(x, y):
         return torch.mv(x, y)
 
     def _nested_jit_mv():
-        return nestedtensor._C.jit_apply_function((nested_sub_clusters, nested_keys), my_fun)
+        return my_fun(nested_sub_clusters, nested_keys)
     return _nested_jit_mv
 
 
@@ -139,12 +137,12 @@ def benchmark_fn(fn, run_time = 15.0):
     gen_results_naive = gen_algorithm_naive(keys, sub_clusters)
     gen_results_mv = gen_algorithm_mv(keys, sub_clusters)
     gen_results_nested_mv = gen_algorithm_nested_mv(keys, sub_clusters)
-    gen_results_nested_jit_mv = gen_algorithm_nested_jit_mv(keys, sub_clusters)
+    # gen_results_nested_jit_mv = gen_algorithm_nested_jit_mv(keys, sub_clusters)
 
-    # print(benchmark_fn(gen_results_naive))
-    # print(benchmark_fn(gen_results_mv))
-    # print(benchmark_fn(gen_results_nested_mv))
-    print(benchmark_fn(gen_results_nested_jit_mv))
+    print(benchmark_fn(gen_results_nested_mv))
+    print(benchmark_fn(gen_results_naive))
+    print(benchmark_fn(gen_results_mv))
+    # print(benchmark_fn(gen_results_nested_jit_mv))
     # import cProfile, pstats, io
     # pr = cProfile.Profile()
     # pr.enable()

nestedtensor/csrc/buffer_nested_tensor.cpp

@@ -71,8 +71,8 @@ std::pair<int64_t, TensorNode> _build_structure(
     for (size_t i = 0; i < nested_size.degree(); i++) {
       std::pair<int64_t, TensorNode> result_i = _build_structure(
           index, buffers, nested_size.children(i), nested_stride.children(i));
+      index = std::get<0>(result_i);
       result.push_back(std::get<1>(result_i));
-      index++;
     }
     return std::pair<int64_t, TensorNode>(index, TensorNode(result));
   }

nestedtensor/csrc/buffer_nested_tensor.h

@@ -125,8 +125,11 @@ struct TORCH_API _BufferNestedTensor {
       new_size.push_back(start->degree());
       start = start->children_data(0);
     }
-    for (size_t i = 0; i < start->payload(0).size(); i++) {
-      new_size.push_back(start->payload(0)[i]);
+    new_size.push_back(start->size());
+    if (start->size() > 0) {
+      for (size_t i = 0; i < start->payload(0).size(); i++) {
+        new_size.push_back(start->payload(0)[i]);
+      }
     }
     return _buffer.reshape(at::IntArrayRef(new_size));
   }