MRA: Fix more mismatched sizes of temporary tensors

Signed-off-by: Joseph Schuchart <[email protected]>
devreal · Sep 23, 2024 · 5767733 · 5767733
1 parent 71790c7
commit 5767733
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 5 deletions.
diff --git a/examples/madness/mra-device/kernels.cu b/examples/madness/mra-device/kernels.cu
@@ -428,10 +428,10 @@ GLOBALSCOPE void compress_kernel(
     SHARED TensorView<T,2> hgT;
     if (is_t0) {
       s = TensorView<T,NDIM>(&tmp[0], 2*K);
+      workspace = TensorView<T, NDIM>(&tmp[TWOK2NDIM], 2*K);
       d = TensorView<T,NDIM>(result_ptr, 2*K);
       p = TensorView<T,NDIM>(p_ptr, K);
-      hgT = TensorView<T,2>(hgT_ptr, K);
-      workspace = TensorView<T, NDIM>(&tmp[TWOK2NDIM], K);
+      hgT = TensorView<T,2>(hgT_ptr, 2*K);
     }
     SYNCTHREADS();
     d = 0.0;
@@ -512,7 +512,7 @@ GLOBALSCOPE void reconstruct_kernel(
     node        = TensorView<T, NDIM>(node_ptr, 2*K);
     s           = TensorView<T, NDIM>(&tmp_ptr[0], 2*K);
     workspace   = TensorView<T, NDIM>(&tmp_ptr[TWOK2NDIM], 2*K);
-    hg          = TensorView<T, 2>(hg_ptr, K);
+    hg          = TensorView<T, 2>(hg_ptr, 2*K);
     from_parent = TensorView<T, NDIM>(from_parent_ptr, K);
   }
   SYNCTHREADS();

diff --git a/examples/madness/mra-device/kernels.h b/examples/madness/mra-device/kernels.h
@@ -40,8 +40,7 @@ template<mra::Dimension NDIM>
 std::size_t compress_tmp_size(std::size_t K) {
   const size_t TWOK2NDIM = std::pow(2*K,NDIM);
   const size_t K2NDIM = std::pow(K,NDIM);
-  return (TWOK2NDIM) // s
-          + K2NDIM // workspace
+  return (2*TWOK2NDIM) // s & workspace
           + mra::Key<NDIM>::num_children() // sumsq for each child and result
           ;
 }