Fix tfq_simulate_ops_cuda op test.

Author: jaeyoo

tensorflow_quantum/core/ops/tfq_simulate_ops_cuda_test.py

@@ -24,18 +24,23 @@
 from tensorflow_quantum.core.ops import tfq_simulate_ops_cuda
 from tensorflow_quantum.python import util
 
+def measure_average_runtime(fn, tag, num_samples=10):
+    avg_time = []
+    for _ in range(num_samples):
+        begin_time = time.time()
+        result = fn()
+        duration = time.time() - begin_time
+        avg_time.append(duration)
+    avg_time = sum(avg_time) / float(num_samples)
+    print(f"\n\t{tag} time: {avg_time}\n")
+    return avg_time, result
+
 
 class SimulateExpectationTest(tf.test.TestCase):
     """Tests tfq_simulate_expectation."""
 
-    def test_simulate_expectation_diff(self):
-        """Make sure that cpu & gpu ops have the same results."""
-        # TF 2
-        # gpus = tf.config.list_physical_devices('GPU')
-        # self.assertLess(len(gpus), 1)
-        # if len(gpus) < 1:
-        #     self.skipTest("Expected at least 1 GPU but found {} GPUs".format(
-        #         len(gpus)))
+    def test_simulate_expectation_cpu_vs_cuda(self):
+        """Make sure that cpu & gpu(cuda) ops have the same results."""
         n_qubits = 20
         batch_size = 5
         symbol_names = ['alpha']
@@ -54,77 +59,27 @@ def test_simulate_expectation_diff(self):
         pauli_sums = util.random_pauli_sums(qubits, 3, batch_size)
         pauli_sums_tensor = util.convert_to_tensor([[x] for x in pauli_sums])
 
-        cpu_avg_time = []
-        for _ in range(10):
-            cpu_time = time.time()
-            res_cpu = tfq_simulate_ops.tfq_simulate_expectation(
+        cpu_avg_time, res_cpu = measure_average_runtime(
+            lambda: tfq_simulate_ops.tfq_simulate_expectation(
                 circuit_batch_tensor,
                 symbol_names, symbol_values_array.astype(np.float64),
-                pauli_sums_tensor)
-            cpu_time = time.time() - cpu_time
-            cpu_avg_time.append(cpu_time)
-        cpu_avg_time = sum(cpu_avg_time) / 10.0
-        print("\n\tCPU time: ", cpu_avg_time, "\n")
-
-        avg_cpu_with_gpu_time = []
-        for _ in range(10):
-            cpu_with_gpu_time = time.time()
-            with tf.device("/device:GPU:0"):
-                res_cpu_with_gpu = tfq_simulate_ops_cuda.tfq_simulate_expectation(
-                    circuit_batch_tensor,
-                    symbol_names, symbol_values_array.astype(np.float64),
-                    pauli_sums_tensor)
-            cpu_with_gpu_time = time.time() - cpu_with_gpu_time
-            avg_cpu_with_gpu_time.append(cpu_with_gpu_time)
-        avg_cpu_with_gpu_time = sum(avg_cpu_with_gpu_time) / 10.0
-
-        # Both are CPU devices.
-        self.assertEqual(res_cpu.device, res_cpu_with_gpu.device)
-        np.testing.assert_allclose(res_cpu, res_cpu_with_gpu)
-        print("\n\tCPU with GPU device time: ", avg_cpu_with_gpu_time, "\n")
-
-        @tf.function
-        def cpu_with_gpu_fn():
-            with tf.device("/device:GPU:0"):
-                return tfq_simulate_ops_cuda.tfq_simulate_expectation(
-                    circuit_batch_tensor,
-                    symbol_names, symbol_values_array.astype(np.float64),
-                    pauli_sums_tensor)
+                pauli_sums_tensor),
+            "CPU"
+        )
 
-        avg_fn_cpu_with_gpu_time = []
-        for _ in range(10):
-            fn_cpu_with_gpu_time = time.time()
-            res_fn_cpu_with_gpu = cpu_with_gpu_fn()
-            fn_cpu_with_gpu_time = time.time() - fn_cpu_with_gpu_time
-            avg_fn_cpu_with_gpu_time.append(fn_cpu_with_gpu_time)
-        avg_fn_cpu_with_gpu_time = sum(avg_fn_cpu_with_gpu_time) / 10.0
-
-        # CPU & GPU devices.
-        self.assertNotEqual(res_cpu.device, res_fn_cpu_with_gpu.device)
-        np.testing.assert_allclose(res_cpu, res_fn_cpu_with_gpu)
-        print("\n\ttf.function, CPU with GPU device time: ",
-              avg_fn_cpu_with_gpu_time, "\n")
-
-        # avg_gpu_time = []
-        # for _ in range(10):
-        #     gpu_time = time.time()
-        #     res_gpu = tfq_simulate_ops_gpu_cpu.tfq_simulate_expectation(
-        #         circuit_batch_tensor,
-        #         symbol_names, symbol_values_array.astype(np.float64),
-        #         pauli_sums_tensor)
-        #     gpu_time = time.time() - gpu_time
-        #     avg_gpu_time.append(gpu_time)
-        # avg_gpu_time = sum(avg_gpu_time) / 10.0
-        # print("\n\tGPU version time: ", avg_gpu_time, "\n")
+        cuda_avg_time, res_cuda = measure_average_runtime(
+            lambda: tfq_simulate_ops_cuda.tfq_simulate_expectation(
+                circuit_batch_tensor,
+                symbol_names, symbol_values_array.astype(np.float64),
+                pauli_sums_tensor),
+            "CUDA"
+        )
 
+        # The result should be the similar within a tolerance.
+        np.testing.assert_allclose(res_cpu, res_cuda, atol=1e-5)
 
-        # # This guarantees that both tensors are not in the same devices
-        # # (e.g. CPU vs GPU)
-        # # self.assertNotEqual(res.device, res_gpu.device)
-        # # -> this doesn't work anymore because TFQ op itself is in CPU.
-        # # only qsim::SimulatorCUDA is in GPU
-        # np.testing.assert_allclose(res_cpu, res_gpu)
-        # self.assertGreater(cpu_avg_time, avg_gpu_time)
+        # CUDA op should be faster than CPU op.
+        self.assertGreater(cpu_avg_time, cuda_avg_time)
 
 
 if __name__ == "__main__":