Cupybackend, mixed backends

qtensor/contraction_backends/__init__.py

@@ -4,38 +4,134 @@
 from .torch import TorchBackend
 from .cupy import CuPyBackend
 from .mkl import CMKLExtendedBackend
+from .cupy import CuPyBackend
+# from .cutensor import CuTensorBackend
 from .transposed import TransposedBackend
 from .opt_einsum import OptEinusmBackend
-from .performance_measurement_decorator import PerfNumpyBackend, PerfBackend
+# from .transpose_backend import NumpyTranspoedBackend, TorchTransposedBackend, CupyTransposedBackend, CutensorTransposedBackend
+from .torch_mix import TorchMixBackend
+from .performance_measurement_decorator import PerfNumpyBackend, PerfBackend, GPUPerfBackend
+from .mix_decorator import MixBackend
+from .torch_profiled import TorchEmbeddedBackend
 
 def get_backend(name):
     backend_dict = {
+        'mkl': CMKLExtendedBackend,
         'einsum':NumpyBackend,
-        'torch': TorchBackend,
-        'mkl':CMKLExtendedBackend,
-        'tr_einsum': TransposedBackend,
         'opt_einsum': OptEinusmBackend,
-        'cupy': CuPyBackend
+        'torch_cpu': TorchBackend,
+        'torch_gpu': TorchBackend,
+        'cupy': CuPyBackend,
+        # 'cutensor': CuTensorBackend,
+        # 'tr_einsum': NumpyTranspoedBackend,
+        # 'tr_torch': TorchTransposedBackend,
+        # 'tr_cupy': CupyTransposedBackend,
+        # 'tr_cutensor': CutensorTransposedBackend,
+        'torch_mix': TorchMixBackend
+    }
+    if name == "torch_cpu":
+        return backend_dict['torch_cpu'](device = "cpu")
+    else:
+        return backend_dict[name]()
+
+
+def get_embedded_backend(name):
+    backend_dict = {
+        'torch_gpu': TorchEmbeddedBackend
     }
-    if name in ["torch_gpu", "torch_cpu"]:
-        return backend_dict['torch'](device = name[-3:])
+    if name == "torch_cpu":
+        return backend_dict['torch_cpu'](device = "cpu")
     else:
         return backend_dict[name]()
 
-def get_perf_backend(name):
+def get_cpu_perf_backend(name):
     class MyPerfBackend(PerfBackend):
         Backend = {
+            'mkl': CMKLExtendedBackend,
             'einsum':NumpyBackend,
-            'torch': TorchBackend,
-            'torch_cpu': TorchBackend,
-            'torch_gpu': TorchBackend,
-            'mkl':CMKLExtendedBackend,
-            'tr_einsum': TransposedBackend,
             'opt_einsum': OptEinusmBackend,
+            'torch_cpu': TorchBackend,     
+            # 'tr_einsum': NumpyTranspoedBackend,
+        }[name]
+
+    if name == "torch_cpu":
+        return MyPerfBackend(device="cpu")
+    else:
+        return MyPerfBackend()
+
+def get_gpu_perf_backend(name):
+    class MyPerfBackend(GPUPerfBackend):
+        Backend = {
+            'torch_gpu': TorchBackend,
             'cupy': CuPyBackend,
+            # 'cutensor': CuTensorBackend,
+            # 'tr_torch': TorchTransposedBackend,
+            # 'tr_cupy': CupyTransposedBackend,
+            # 'tr_cutensor': CutensorTransposedBackend
         }[name]
 
-    if name in ["torch_gpu", "torch_cpu"]:
-        return MyPerfBackend(device = name[-3:])
+    if name == "torch_gpu":
+        return MyPerfBackend(device="gpu")
     else:
         return MyPerfBackend()
+
+
+
+# def get_mixed_backend(cpu_name, gpu_name):
+#     class MyMixedBackend(MixBackend):
+#         CBE = {
+#             'einsum':NumpyBackend,
+#             'torch_cpu': TorchBackend,
+#             'mkl':CMKLExtendedBackend,
+#             'opt_einsum': OptEinusmBackend,
+#             'tr_einsum': NumpyTranspoedBackend,
+#             'opt_einsum': OptEinusmBackend,
+#         }[cpu_name]
+#         GBE = {
+#             'torch_gpu': TorchBackend,
+#             'cupy': CuPyBackend,
+#             'tr_torch': TorchTransposedBackend,
+#             'tr_cupy': CupyTransposedBackend,
+#             'tr_cutensor': CutensorTransposedBackend
+#         }[gpu_name]
+
+#     if cpu_name == "torch_cpu":
+#         return MyMixedBackend(device = "cpu")
+#     else:
+#         return MyMixedBackend()
+
+def get_mixed_backend(cpu_name, gpu_name, threshold = 11):
+
+    cpu_be = get_backend(cpu_name)
+    gpu_be = get_backend(gpu_name)
+    # mix_name = cpu_name +"-"+ gpu_name
+
+    # threshold_dict = {
+    #     "einsum-cupy":16,
+    #     "torch_cpu-torch_gpu":11,
+    #     "einsum-torch_gpu":13
+    # }
+
+    # if mix_name in threshold_dict:
+    #     threshold = threshold_dict[mix_name]
+
+    return MixBackend(cpu_be, gpu_be, threshold)
+
+
+def get_mixed_perf_backend(cpu_name, gpu_name,threshold = 11):
+
+    cpu_be = get_cpu_perf_backend(cpu_name)
+    gpu_be = get_gpu_perf_backend(gpu_name)
+    mix_name = cpu_name +"-"+ gpu_name
+
+    threshold_dict = {
+        "einsum-cupy":16,
+        "torch_cpu-torch_gpu":11,
+        "einsum-torch_gpu":13
+    }
+
+    if mix_name in threshold_dict:
+        threshold = threshold_dict[mix_name]
+
+    return MixBackend(cpu_be, gpu_be, threshold)
+

qtensor/contraction_backends/cupy.py

@@ -1,6 +1,8 @@
 import qtree
 from qtensor.tools.lazy_import import cupy as cp
+import numpy as np
 from qtensor.contraction_backends import ContractionBackend
+from qtensor.contraction_backends.numpy import get_einsum_expr
 
 mempool = mempool = cp.get_default_memory_pool()
 
@@ -20,8 +22,8 @@ def process_bucket(self, bucket, no_sum=False):
             '''
             Change 1: Using cp.einsum not torch.einsum
             ''' 
-            result_data = cp.einsum(expr, result_data, tensor.data)
-
+            result_data = cp.einsum(expr, cp.asarray(result_data), cp.asarray(tensor.data))
+            #print("result data: ",type(result_data))
             # Merge and sort indices and shapes
             result_indices = tuple(sorted(
                 set(result_indices + tensor.indices),
@@ -47,6 +49,36 @@ def process_bucket(self, bucket, no_sum=False):
                                 data=cp.sum(result_data, axis=0))
         return result
 
+    def process_bucket_merged(self, ixs, bucket, no_sum=False):
+        result_indices = bucket[0].indices
+        result_data = bucket[0].data
+        all_indices = set(sum((list(t.indices) for t in bucket), []))
+        result_indices = all_indices - set(ixs)
+        all_indices_list = list(all_indices)
+        to_small_int = lambda x: all_indices_list.index(x)
+
+        expr = get_einsum_expr(bucket, all_indices_list, result_indices)
+
+        params = []
+        for tensor in bucket:
+            params.append(tensor.data)
+            params.append(list(map(to_small_int, tensor.indices)))
+        params.append(list(map(to_small_int, result_indices)))
+
+        expect = len(result_indices)
+
+        result_data = cp.einsum(*params)
+
+        if len(result_indices) > 0:
+            first_index, *_ = result_indices
+            tag = str(first_index)
+        else:
+            tag = 'f'
+
+        result = qtree.optimizer.Tensor(f'E{tag}', result_indices,
+                            data=result_data)
+        return result
+
     def get_sliced_buckets(self, buckets, data_dict, slice_dict):
         mempool.free_all_blocks()
         sliced_buckets = []
@@ -105,4 +137,8 @@ def get_sliced_buckets(self, buckets, data_dict, slice_dict):
         return sliced_buckets
 
     def get_result_data(self, result):
-        return result.data
+        #print(type(result.data))
+        try:
+            return result.data.get()
+        except :
+            return result.data

qtensor/contraction_backends/cutensor.py

@@ -0,0 +1,85 @@
+import qtree
+from qtensor.tools.lazy_import import cupy as cp
+from qtensor.contraction_backends import ContractionBackend
+from qtensor.contraction_backends import CuPyBackend
+from cupy import cutensor as cupy_cutensor
+
+mempool = mempool = cp.get_default_memory_pool()
+
+class CuTensorBackend(CuPyBackend):
+
+    # Replace all torch methods with cupy's analog
+
+    def process_bucket(self, bucket, no_sum=False):
+        result_indices = bucket[0].indices
+        result_data = bucket[0].data
+        for tensor in bucket[1:]:
+
+            expr = qtree.utils.get_einsum_expr(
+                list(map(int, result_indices)), list(map(int, tensor.indices))
+            )
+
+            a, b, c, mode_a, mode_b, mode_c, desc_a, desc_b, desc_c = self.get_ready(expr, result_data, tensor.data)
+            result_data = cupy_cutensor.contraction(1.0, a, desc_a, mode_a, 
+                        b, desc_b, mode_b, 0, 
+                        c, desc_c, mode_c)
+
+            # Merge and sort indices and shapes
+            result_indices = tuple(sorted(
+                set(result_indices + tensor.indices),
+                key=int)
+            )
+
+        if len(result_indices) > 0:
+            if not no_sum:  # trim first index
+                first_index, *result_indices = result_indices
+            else:
+                first_index, *_ = result_indices
+            tag = first_index.identity
+        else:
+            tag = 'f'
+            result_indices = []
+
+        # reduce
+        if no_sum:
+            result = qtree.optimizer.Tensor(f'E{tag}', result_indices,
+                                data=result_data)
+        else:
+            result = qtree.optimizer.Tensor(f'E{tag}', result_indices,
+                                data=cp.sum(result_data, axis=0))
+        return result
+
+
+    @staticmethod
+    def get_ready(contraction, a, b):
+        # get ready
+        inp, out = contraction.split('->')
+        size = inp.split(',')
+        out_size =cp.full(len(out), 2).tolist() # fill_number = 2
+
+        # TO DO
+        mode_a = tuple(size[0])
+        mode_b = tuple(size[1])
+        mode_c = tuple(out)
+
+        # generate tensor c
+        shape_a, shape_b = a.shape, b.shape
+        shape_c = tuple(out_size)
+
+        c = cp.random.rand(*shape_c).astype(a.dtype)
+
+        # manually cast b to a's type
+        b = b.astype(a.dtype)
+
+        # generate tensor descriptor
+        desc_a = cupy_cutensor.create_tensor_descriptor(a)
+        desc_b = cupy_cutensor.create_tensor_descriptor(b)
+        desc_c = cupy_cutensor.create_tensor_descriptor(c)
+
+
+        if not a.flags['C_CONTIGUOUS']:
+            a = cp.ascontiguousarray(a)
+        if not b.flags['C_CONTIGUOUS']:
+            b = cp.ascontiguousarray(b)
+
+        return a, b, c, mode_a, mode_b, mode_c, desc_a, desc_b, desc_c

qtensor/contraction_backends/mix_decorator.py

@@ -0,0 +1,92 @@
+from qtensor.contraction_backends import ContractionBackend
+"""
+class MixedBe(ConBE):
+    be1: cpu_be
+    be2: gpu_be
+
+    def get_sliced_bucket():
+        normal slicing
+            either use be1_get_sliuced or naive implementation
+            np.array for convinience
+
+    def process_bucket():
+        - Check input bucket width
+        - If larger than 8, use gpu_be.process(bucket)
+        - Else, use less than 8, use cpu_be.process(bucket)
+
+    def get_result_data():
+        - always use gpu_be's get_result
+        - so that the gpu_be shall handle the gpu-cpu transfer all the time
+
+"""
+
+
+'''
+Input:  Array of either (indices) or a tensor
+        If a tensor, use t.indices for conversion
+'''
+def bucketWidth(bucket):
+    bucket_width = 0
+    for tensor in bucket:
+        tensor_len = 0
+        if type(tensor) is tuple:
+            tensor_len = len(tensor)
+            if tensor_len > bucket_width:
+                bucket_width = tensor_len
+        else:
+            tensor_len = len(tensor.indices)
+            if tensor_len > bucket_width:
+                bucket_width = tensor_len
+    return bucket_width
+
+
+def AccurateBucketWidth(bucket):
+    indices_set = set()
+    for tensor in bucket:
+        if type(tensor) is tuple:
+            for index in tensor:
+                indices_set.add(index)
+        else:
+            for index in tensor.indices:
+                indices_set.add(index)
+
+    return len(indices_set)
+
+
+
+
+'''
+I/O: Actual BE Objects -> Wrapped Class
+TODO: process_bucket_merged() for its own, when when we test the merged simulators, we create the merged simulator
+'''
+
+class MixBackend(ContractionBackend):
+
+    def __init__(self, cpu_be, gpu_be, watershed):
+        self.cpu_be = cpu_be
+        self.gpu_be = gpu_be
+        self.watershed = watershed
+
+    def process_bucket(self, bucket, no_sum = False):
+        # bucket_width = bucketWidth(bucket)
+        accc_width = AccurateBucketWidth(bucket)
+        # if accc_width != bucket_width:
+        #      print(accc_width, bucket_width)
+        if accc_width >= self.watershed:
+            #print(f"In GPU, width {accc_width}")
+            return self.gpu_be.process_bucket(bucket, no_sum)
+        else:
+            return self.cpu_be.process_bucket(bucket, no_sum)
+
+    def process_bucket_merged(self,ixs, bucket, no_sum=False):
+        accu_width = AccurateBucketWidth(bucket)
+        if accu_width >= self.watershed:
+            return self.gpu_be.process_bucket_merged(ixs, bucket, no_sum)
+        else:
+            return self.cpu_be.process_bucket_merged(ixs, bucket, no_sum)
+
+    def get_sliced_buckets(self, buckets, data_dict, slice_dict):
+        return self.cpu_be.get_sliced_buckets(buckets, data_dict, slice_dict)
+
+    def get_result_data(self, result):
+        return self.gpu_be.get_result_data(result)