[buddy-mlir] add new constructor for Memref class

Signed-off-by: Avimitin <[email protected]>
chipsalliance · Aug 24, 2024 · b36543e · b36543e
1 parent ca18bb5
commit b36543e
Showing 1 changed file with 17 additions and 2 deletions.
diff --git a/tests/pytorch/memref.hpp b/tests/pytorch/memref.hpp
@@ -7,6 +7,7 @@
 template <typename T, size_t N> class MemRef {
 public:
   constexpr MemRef(T *data, const int32_t sizes[N]);
+  constexpr MemRef(T *data, T init, const int32_t sizes[N]);
 
 protected:
   inline void setStrides();
@@ -19,8 +20,8 @@ template <typename T, size_t N> class MemRef {
   int32_t strides[N];
 };
 
-template <typename T, std::size_t N> constexpr
-MemRef<T, N>::MemRef(T *data, const int32_t sizes[N]) {
+template <typename T, std::size_t N>
+constexpr MemRef<T, N>::MemRef(T *data, const int32_t sizes[N]) {
   for (size_t i = 0; i < N; i++) {
     this->sizes[i] = sizes[i];
   }
@@ -31,6 +32,20 @@ MemRef<T, N>::MemRef(T *data, const int32_t sizes[N]) {
   aligned = data;
 }
 
+template <typename T, std::size_t N>
+constexpr MemRef<T, N>::MemRef(T *data, T init, const int32_t sizes[N])
+    : MemRef(data, sizes) {
+
+  int32_t total_size = 0;
+  for (size_t i = 0; i < N; i++) {
+    total_size += sizes[i];
+  }
+
+  for (int32_t i = 0; i < total_size; i++) {
+    aligned[i] = init;
+  }
+}
+
 template <typename T, std::size_t N> inline void MemRef<T, N>::setStrides() {
   strides[N - 1] = 1;
   if (N < 2)