From df67b47b307732959d674a0e1c26e16ed8a635da Mon Sep 17 00:00:00 2001 From: Dan Lykov Date: Sun, 1 Sep 2024 17:36:45 +0000 Subject: [PATCH] fix writetodisk compressor cupy --- qtensor/compression/Compressor.py | 6 +++--- qtensor/contraction_backends/cupy.py | 8 ++++++-- .../performance_measurement_decorator.py | 2 +- second_try.nc | Bin 1208 -> 0 bytes 4 files changed, 10 insertions(+), 6 deletions(-) delete mode 100644 second_try.nc diff --git a/qtensor/compression/Compressor.py b/qtensor/compression/Compressor.py index 65e28f8d..26d42a94 100644 --- a/qtensor/compression/Compressor.py +++ b/qtensor/compression/Compressor.py @@ -220,7 +220,7 @@ def compress(self, data): isCuPy, num_elements_eff = _get_data_info(data) dtype = data.dtype cmp_bytes, outSize_ptr = self.cuszx_compress(isCuPy, data, num_elements_eff, self.r2r_error, self.r2r_threshold) - # return (cmp_bytes, num_elements_eff, isCuPy, data.shape, dtype, outSize_ptr) + # return (cmpbytes, num_elements_eff, isCuPy, data.shape, dtype, outSize_ptr) return (cmp_bytes, num_elements_eff, isCuPy, data.shape, dtype, outSize_ptr.contents.value) @@ -526,9 +526,9 @@ def _gen_filename(self, info): def compress(self, data): import cupy if isinstance(data, cupy.ndarray): - isCupy=False - else: isCupy=True + else: + isCupy=False fname = self._gen_filename((data.dtype, data.shape, isCupy)) data.tofile(fname) return (fname, data.dtype, data.shape, isCupy) diff --git a/qtensor/contraction_backends/cupy.py b/qtensor/contraction_backends/cupy.py index ee4e4b6c..cb0ead28 100644 --- a/qtensor/contraction_backends/cupy.py +++ b/qtensor/contraction_backends/cupy.py @@ -99,7 +99,11 @@ def get_sliced_buckets(self, buckets, data_dict, slice_dict): # cp.argsort requires input to be cp array #print(tensor.indices) out_indices = list(sorted(tensor.indices, key=int, reverse=True)) - data = data_dict[tensor.data_key] + if tensor.data is None: + data = data_dict[tensor.data_key] + else: + data = tensor.data + #print(f"Data shape: {data.shape}, Tensor indices: {tensor.indices}, out_indices: {out_indices}") data, new_indices = slice_numpy_tensor(data, tensor.indices, out_indices, slice_dict) # transpose indices try: @@ -114,4 +118,4 @@ def get_sliced_buckets(self, buckets, data_dict, slice_dict): return sliced_buckets def get_result_data(self, result): - return cp.transpose(result.data) + return cp.transpose(result.data).get() diff --git a/qtensor/contraction_backends/performance_measurement_decorator.py b/qtensor/contraction_backends/performance_measurement_decorator.py index 39efffa7..e20008cd 100644 --- a/qtensor/contraction_backends/performance_measurement_decorator.py +++ b/qtensor/contraction_backends/performance_measurement_decorator.py @@ -85,7 +85,7 @@ def check_store(self): tensors_sizes=[len(tensor.indices) for tensor in self.object_store.values()] )) # -- - print('MH', self.mem_history[-1]) + #print('MH', self.mem_history[-1]) if cupy_mem>1024**2: self._print("CuPy memory usage", cupy_mem/1024/1024, "MB. Total MB:", mempool.total_bytes()/1024**2) diff --git a/second_try.nc b/second_try.nc deleted file mode 100644 index 8022d972756020c7566b9896a4dab15189ae0875..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 1208 zcmb7^e@IhN6vy9eb7pB^V38CJlNM-@{m`Vi*V0JJlEOfXJ)dpWN@x3F6a>zsCH-My zzlmyz8HLdwN<|Pks0;~OE31DZ)6L=3GQ=7Q);Z_(+^T=-!nvRGJ?GqW&wGz;ecpx` zL0<6yHJJ90WV1+>CX7x4i?cgyX3O4Xnc&X^Qyqj27Uoo2$Q5SDpdEeeqvTZJW393< zmo5uc9x(boWfnX`cLAr1`Ig;PMi@apgf*IK$DWMrh+ZVSG}?SoJS^F4(h<=pIV9PG zM>0$06(&)X3zd1Wm*Qh47UW#HM0EzJJS)zl%!X-v2AmI;x(zmbcW4+(1WN*&2c`$( zVEe)DLc29qhsk2EvWX6JrO8@tt}G`{5UghV0kcCE?BxXm)C&=6pOJpiznSQp_(bf> zdPwX)@rW3(9&>`j2D5XC!{KC#e`|TuiE+RCIK`*fK%bHgWWA=kU|J1v@KO=gg?4C2 z58XD99vzqU#H|Dl2{yU3gXnjy1PT3rG0?D$^hn7ZBwrk~G-5v7UdU0iS?!(3xKo>Y zfo$>#q=SX0iJ|(p#PQ@NV)(p)gTrG!Xo-nEt9ntV zx_n+_Zpi|CA#Bvx<@}8v{5}HvaW2Wb%zkjaj#Q2Z4G!EqPIXZyAlWOlfkxS}^ z@CEXMqIJ7~Y^>D9qxm&((FvNz7Vlh$tnf81hX60Sci95ECJb8?IlC&Y0L@pk)M4ay zYgamudDfl-MCiTqp0vO98j|bwHI*ToOZ95^-+lkQN73x5_c0t3ruZtQn{+6@0LgV% w>NX;qx3*@~|3}~M%#+kdVb59me_$i}1i5cV=pFSfzM10u_r7pM`Ms)t0y