Merge pull request #545 from nicolasvasilache/pr/caffe2-update

Update to support upcoming Caffe2 API

Author: nicolasvasilache

python/benchmarks/caffe2_benchmark.py

@@ -19,6 +19,7 @@
 import numpy as np
 import os
 import tensor_comprehensions as tc
+import torch
 
 from caffe2.python import core, dyndep, workspace, utils
 
@@ -85,18 +86,22 @@ def main():
 
 @utils.debug
 def tune(args):
-    fc = tc.define(FC_LANG, name="func_fc")
-    options = fc.autotune(
-        (args.batch_size, args.input_dim),
-        (args.output_dim, args.input_dim),
-        (args.output_dim,),
-        cache = args.tuner_cache_file,
-        threads = args.tuner_threads,
-        generations = args.tuner_gen_generations,
-        pop_size = args.tuner_gen_pop_size,
-    )
-    print(options.toString())
-    return options
+    tuner_config = (
+        tc.TunerConfig()
+        .generations(args.tuner_gen_generations)
+        .devices(args.tuner_devices)
+        .threads(args.tuner_threads)
+        .pop_size(args.tuner_gen_pop_size))
+    return tc.autotune(
+        FC_LANG,
+        'func_fc',
+        torch.randn(args.batch_size, args.input_dim, device='cuda'),
+        torch.randn(args.output_dim, args.input_dim, device='cuda'),
+        torch.randn(args.output_dim, device='cuda'),
+        starting_options = tc.MappingOptions('naive'),
+        tuner_config = tuner_config,
+        cache_filename = args.tuner_cache_file,
+        store_to_cache = True)
 
 
 @utils.debug

python/examples/min_distance.py

@@ -103,9 +103,9 @@ def argmin_1d(float(N) S, int32(D) MinIdx) -> (min_idx) {
 
 T = tc.define(lang, tc.make_naive_options_factory())
 S = T.reduce_codes(luts_t, codes_t)
-V = T.min_2d(S.view((D, N / D)))
+V = T.min_2d(S.view((D, N // D)))
 v = T.min_1d(V)
-MinIdx = T.argmin_2d(S.view((D, N / D)), v)
+MinIdx = T.argmin_2d(S.view((D, N // D)), v)
 min_idx = T.argmin_1d(S, MinIdx)
 print("minval: {} minidx: {}".format(v, min_idx))
 
@@ -159,8 +159,8 @@ def argmin_1d(float(K, N) S, int32(K, D) MinIdx2) -> (MinIdx) {
 
 T = tc.define(lang, tc.make_naive_options_factory())
 S = T.reduce_codes(luts_t, codes_t)
-V2 = T.min_2d(S.view((K, D, N / D)))
+V2 = T.min_2d(S.view((K, D, N // D)))
 V = T.min_1d(V2)
-MinIdx2 = T.argmin_2d(S.view((K, D, N / D)), V)
+MinIdx2 = T.argmin_2d(S.view((K, D, N // D)), V)
 MinIdx = T.argmin_1d(S, MinIdx2)
 print("minval: {} minidx: {}".format(V, MinIdx))

python/tests/test_caffe2.py

@@ -18,6 +18,7 @@
 import hypothesis.strategies as st
 import caffe2.python.hypothesis_test_util as hu
 import tensor_comprehensions as tc
+import torch
 
 from hypothesis import given, settings
 from caffe2.python import core, dyndep
@@ -33,9 +34,6 @@
 def matmul(float(M,N) A, float(N,K) B) -> (output) {
     output(m, k) +=! A(m, r_n) * B(r_n, k)
 }
-"""
-
-MATMUL_GRAD_LANG = """
 def matmul_grad(float(M, N) A, float(N, K) B, float(M, K) d_O) -> (d_A, d_B) {
     d_A(m, n) +=! d_O(m, r_k) * B(n, r_k)
     d_B(n, k) +=! d_O(r_m, k) * A(r_m, n)
@@ -61,7 +59,7 @@ def ref(X, W):
             "TcOp", ["X", "Y"], "out",
             tc_def=MATMUL_LANG,
             tc_name="matmul",
-            tc_grad_def=MATMUL_GRAD_LANG,
+            tc_grad_def=MATMUL_LANG,
             tc_grad_name="matmul_grad",
             inputs_used_by_gradient=[0, 1],
             output_gradients_used_by_gradient=[0],
@@ -91,24 +89,23 @@ def ref(X, W):
            **hu.gcs_gpu_only)
     @settings(max_examples=2)
     def test_matmul_tune_and_run(self, n, m, k, seed, gc, dc):
-        matmul = tc.define(MATMUL_LANG, name="matmul")
-        matmul_grad = tc.define(MATMUL_GRAD_LANG, name="matmul_grad")
-
-        mapping_options = matmul.autotune(
-            (n, k), (k, m),
-            generations=3,
-            threads=32,
-            pop_size=2,
-            tuner_min_launch_total_threads=1,
-        )
-
-        grad_mapping_options = matmul_grad.autotune(
-            (n, k), (k, m), (n, m),
-            generations=1,
-            threads=32,
-            pop_size=2,
-            tuner_min_launch_total_threads=1,
-        )
+        tuner = tc.Tuner(MATMUL_LANG)
+        tuner_config = (
+            tc.TunerConfig().generations(3).threads(32).pop_size(2)
+            .tuner_min_launch_total_threads(1))
+        matmul_top1 = tuner.tune(
+            'matmul',
+            (torch.randn(n, k, device='cuda'),
+             torch.randn(k, m, device='cuda')),
+            tc.MappingOptions('naive'),
+            tuner_config)
+        matmul_grad_top1 = tuner.tune(
+            'matmul_grad',
+            (torch.randn(n, k, device='cuda'),
+             torch.randn(k, m, device='cuda'),
+             torch.randn(n, m, device='cuda')),
+            tc.MappingOptions('naive'),
+            tuner_config)
 
         X = np.random.rand(m, k).astype(np.float32)
         W = np.random.rand(k, n).astype(np.float32)
@@ -120,13 +117,13 @@ def ref(X, W):
             "TcOp", ["X", "Y"], "out",
             tc_def=MATMUL_LANG,
             tc_name="matmul",
-            tc_grad_def=MATMUL_GRAD_LANG,
+            tc_grad_def=MATMUL_LANG,
             tc_grad_name="matmul_grad",
             inputs_used_by_gradient=[0, 1],
             output_gradients_used_by_gradient=[0],
             inputs_to_compute_gradients_of=[0, 1],
-            mapping_options=mapping_options.serialize(),
-            grad_mapping_options=grad_mapping_options.serialize(),
+            mapping_options=matmul_top1.serialize(),
+            grad_mapping_options=matmul_grad_top1.serialize(),
         )
 
         self.assertReferenceChecks(

test/caffe2/test_harness-inl.h

@@ -69,8 +69,11 @@ caffe2::Tensor<typename Caffe2Backend::Context> GetNamedTensor(
     caffe2::Workspace& ws,
     const std::string& name) {
   // Resolved dynamically
-  return caffe2::Tensor<typename Caffe2Backend::Context>(
-      ws.GetBlob(name)->Get<typename Caffe2Backend::Tensor>());
+  auto& t = ws.GetBlob(name)->Get<typename Caffe2Backend::Tensor>();
+  caffe2::Tensor<typename Caffe2Backend::Context> res;
+  res.ResizeLike(t);
+  res.ShareData(t);
+  return res;
 }
 
 // helper functions to construct an ATen tensor from a caffe2 tensor