* Patch from Anna to handle the data capture and storage issue.

* Removed the skipping of Python tests

-----------------------

Signed-off-by: Pradnya Khalate <[email protected]>
Co-authored-by: Anna Gringauze <[email protected]>

Author: khalatepradnya

python/cudaq/kernel/kernel_decorator.py

@@ -211,6 +211,12 @@ def compile(self):
         if self.module != None:
             return
 
+        # FIX: Cleanup up the captured data if the module needs recompilation.
+        if self.capturedDataStorage != None:
+            self.capturedDataStorage.__del__()
+        self.capturedDataStorage = self.createStorage()
+
+        # Caches the module and stores captured data into self.capturedDataStorage.
         self.module, self.argTypes, extraMetadata = compile_to_mlir(
             self.astModule,
             self.capturedDataStorage,
@@ -420,7 +426,8 @@ def __call__(self, *args):
             return
 
         # Prepare captured state storage for the run
-        self.capturedDataStorage = self.createStorage()
+        # FIX: moved to compile()
+        # self.capturedDataStorage = self.createStorage()
 
         # Compile, no-op if the module is not None
         self.compile()
@@ -501,8 +508,9 @@ def __call__(self, *args):
                                             self.module,
                                             *processedArgs,
                                             callable_names=callableNames)
-            self.capturedDataStorage.__del__()
-            self.capturedDataStorage = None
+            # FIX: do not delete the captured storage
+            # self.capturedDataStorage.__del__()
+            # self.capturedDataStorage = None
         else:
             result = cudaq_runtime.pyAltLaunchKernelR(
                 self.name,
@@ -511,8 +519,9 @@ def __call__(self, *args):
                 *processedArgs,
                 callable_names=callableNames)
 
-            self.capturedDataStorage.__del__()
-            self.capturedDataStorage = None
+            # FIX: do not delete the captured storage
+            # self.capturedDataStorage.__del__()
+            # self.capturedDataStorage = None
             return result
 
 

python/tests/builder/test_qalloc_init_state.py

@@ -154,7 +154,6 @@ def test_kernel_complex_params_f32():
 
 
 @skipIfNvidiaFP64NotInstalled
-@pytest.mark.skip(reason="kernel.compile() causes a crash")
 def test_kernel_complex_capture_f64():
     cudaq.reset_target()
     cudaq.set_target('nvidia-fp64')
@@ -207,7 +206,6 @@ def test_kernel_complex128_capture_f64():
 
 
 @skipIfNvidiaNotInstalled
-@pytest.mark.skip(reason="kernel.compile() causes a crash")
 def test_kernel_complex64_capture_f32():
     cudaq.reset_target()
     cudaq.set_target('nvidia')

python/tests/kernel/test_kernel_qvector_state_init.py

@@ -175,7 +175,6 @@ def test_kernel_complex_params_f32():
 
 
 @skipIfNvidiaFP64NotInstalled
-@pytest.mark.skip(reason="kernel.compile() causes a crash")
 def test_kernel_complex_capture_f64():
     cudaq.reset_target()
     cudaq.set_target('nvidia', option='fp64')
@@ -228,7 +227,6 @@ def test_kernel_complex128_capture_f64():
 
 
 @skipIfNvidiaNotInstalled
-@pytest.mark.skip(reason="kernel.compile() causes a crash")
 def test_kernel_complex64_capture_f32():
     cudaq.reset_target()
     cudaq.set_target('nvidia')
@@ -316,7 +314,6 @@ def kernel(vec: cudaq.State):
 
 
 @skipIfNvidiaFP64NotInstalled
-@pytest.mark.skip(reason="kernel.compile() causes a crash")
 def test_kernel_simulation_dtype_capture_f64():
     cudaq.reset_target()
     cudaq.set_target('nvidia', option='fp64')
@@ -337,7 +334,6 @@ def kernel():
 
 
 @skipIfNvidiaNotInstalled
-@pytest.mark.skip(reason="kernel.compile() causes a crash")
 def test_kernel_simulation_dtype_capture_f32():
     cudaq.reset_target()
     cudaq.set_target('nvidia')
@@ -418,7 +414,6 @@ def kernel(initialState: cudaq.State):
     assert not '01' in counts
 
 
-@pytest.mark.skip(reason="kernel0.compile() causes a crash")
 def test_inner_kernels_state():
     c = np.array([1. / np.sqrt(2.) + 0j, 0., 0., 1. / np.sqrt(2.)],
                  dtype=cudaq.complex())