Update test_caffe2 to latest python bindings

Tested internally

Author: Nicolas Vasilache

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(