From 8ba14b9818c12c235266f0d4490d6703242bf8ab Mon Sep 17 00:00:00 2001
From: Jack Lo <36210336+jackl-xilinx@users.noreply.github.com>
Date: Fri, 3 Feb 2023 11:02:46 -0800
Subject: [PATCH] Fixed bug with alloc size in sim (#324)

---
 runtime_lib/test_library.cpp | 12 ++++++++----
 runtime_lib/test_library.h   |  1 +
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/runtime_lib/test_library.cpp b/runtime_lib/test_library.cpp
index f485f6aac5..616c512312 100644
--- a/runtime_lib/test_library.cpp
+++ b/runtime_lib/test_library.cpp
@@ -983,21 +983,25 @@ void mlir_aie_init_mems(aie_libxaie_ctx_t *ctx, int numBufs) {
 #endif
 }
 
+/*
+ * size - size in words (4 bytes)
+ */
 int *mlir_aie_mem_alloc(aie_libxaie_ctx_t *ctx, int bufIdx, int size) {
 #if defined(__AIESIM__)
+  int size_bytes = size * sizeof(int);
   ctx->buffers[bufIdx] = new ext_mem_model_t;
-  (ctx->buffers[bufIdx])->virtualAddr = std::malloc(size * sizeof(int));
+  (ctx->buffers[bufIdx])->virtualAddr = std::malloc(size_bytes);
   if ((ctx->buffers[bufIdx])->virtualAddr) {
-    (ctx->buffers[bufIdx])->size = size;
+    (ctx->buffers[bufIdx])->size = size_bytes;
     // assign physical space in SystemC DDR memory controller
     (ctx->buffers[bufIdx])->physicalAddr = nextAlignedAddr;
     // adjust nextAlignedAddr to the next 128-bit aligned address
-    nextAlignedAddr = nextAlignedAddr + size;
+    nextAlignedAddr = nextAlignedAddr + size_bytes;
     uint64_t gapToAligned = nextAlignedAddr % 16; // 16byte (128bit)
     if (gapToAligned > 0)
       nextAlignedAddr += (16 - gapToAligned);
   } else {
-    printf("ExtMemModel: Failed to allocate %d memory.\n", size);
+    printf("ExtMemModel: Failed to allocate %d memory.\n", size_bytes);
   }
 
   std::cout << "ExtMemModel constructor: virutal address " << std::hex
diff --git a/runtime_lib/test_library.h b/runtime_lib/test_library.h
index ad378a5227..8e65c86a3b 100644
--- a/runtime_lib/test_library.h
+++ b/runtime_lib/test_library.h
@@ -16,6 +16,7 @@
 
 #if defined(__AIESIM__)
 #include "xioutils.h"
+#include <iostream>
 #endif
 
 extern "C" {