From 576773341de5d7553b4adbaf6b1b597b4a90a794 Mon Sep 17 00:00:00 2001
From: Joseph Schuchart <joseph.schuchart@stonybrook.edu>
Date: Mon, 23 Sep 2024 08:51:49 +0800
Subject: [PATCH] MRA: Fix more mismatched sizes of temporary tensors

Signed-off-by: Joseph Schuchart <joseph.schuchart@stonybrook.edu>
---
 examples/madness/mra-device/kernels.cu | 6 +++---
 examples/madness/mra-device/kernels.h  | 3 +--
 2 files changed, 4 insertions(+), 5 deletions(-)

diff --git a/examples/madness/mra-device/kernels.cu b/examples/madness/mra-device/kernels.cu
index 1e6526ab3..99ab0e516 100644
--- 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
index cd9ab9f36..8447a7f1d 100644
--- 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
           ;
 }