Removed create_golden_Tensor, created "safe" custom Constructor (#2136)

- PR closes #2027 
- Changes container for `GoldenTensor` in Python to be a shared_ptr,
ensuring data is maintained and not copied.
- Uses `GoldenTensor` custom constructor instead of
`create_golden_tensor`.
- Creates copy of supplied data_ptr (size of data must also be provided,
example in `ttir_builder.py`) and stores within GoldenTensor object.

include/ttmlir/Target/Utils/MLIRToFlatbuffer.h

@@ -33,23 +33,13 @@ struct GoldenTensor {
   std::vector<int64_t> shape;
   std::vector<int64_t> strides;
   ::tt::target::DataType dtype;
-  std::uint8_t *data;
+  std::vector<std::uint8_t> data;
 
   GoldenTensor(std::string name, std::vector<int64_t> shape,
                std::vector<int64_t> strides, ::tt::target::DataType dtype,
-               std::uint8_t *data)
-      : name(name), shape(shape), strides(strides), dtype(dtype), data(data) {}
-
-  std::vector<std::uint8_t> convertDataToVector() {
-    int totalDataSize = std::accumulate(this->shape.begin(), this->shape.end(),
-                                        1, std::multiplies<int64_t>()) *
-                        sizeof(float);
-
-    std::vector<std::uint8_t> dataVec(totalDataSize);
-    std::memcpy(dataVec.data(), this->data, totalDataSize);
-
-    return dataVec;
-  }
+               std::vector<std::uint8_t> &&_data)
+      : name(name), shape(shape), strides(strides), dtype(dtype),
+        data(std::move(_data)) {}
 };
 
 inline ::tt::target::OOBVal toFlatbuffer(FlatbufferObjectCache &,

lib/Target/TTNN/TTNNToFlatbuffer.cpp

@@ -1540,13 +1540,12 @@ std::shared_ptr<void> ttnnToFlatbuffer(
   std::vector<::flatbuffers::Offset<::tt::target::GoldenKV>> goldenKVList;
   goldenKVList.reserve(goldenMap.size());
 
-  for (auto element : goldenMap) {
-    std::vector<std::uint8_t> dataTensor = element.second.convertDataToVector();
+  for (const auto &[key, value] : goldenMap) {
     auto goldenTensor = ::tt::target::CreateGoldenTensorDirect(
-        fbb, element.second.name.c_str(), &element.second.shape,
-        &element.second.strides, element.second.dtype, &dataTensor);
-    auto goldenKV = ::tt::target::CreateGoldenKVDirect(
-        fbb, element.first.c_str(), goldenTensor);
+        fbb, value.name.c_str(), &value.shape, &value.strides, value.dtype,
+        &value.data);
+    auto goldenKV =
+        ::tt::target::CreateGoldenKVDirect(fbb, key.c_str(), goldenTensor);
     goldenKVList.push_back(goldenKV);
   }
 

python/Passes.cpp

@@ -12,6 +12,7 @@
 #include <cstdint>
 #include <pybind11/stl_bind.h>
 
+// Make Opaque so Casts & Copies don't occur
 PYBIND11_MAKE_OPAQUE(std::shared_ptr<void>);
 PYBIND11_MAKE_OPAQUE(std::vector<std::pair<std::string, std::string>>);
 
@@ -254,23 +255,26 @@ void populatePassesModule(py::module &m) {
     return ::tt::target::DataType::MIN;
   });
 
-  py::class_<mlir::tt::GoldenTensor>(m, "GoldenTensor")
-      .def(py::init<std::string, std::vector<int64_t>, std::vector<int64_t>,
-                    ::tt::target::DataType, std::uint8_t *>())
+  // Preserve the Data by holding it in a SharedPtr.
+  py::class_<mlir::tt::GoldenTensor, std::shared_ptr<mlir::tt::GoldenTensor>>(
+      m, "GoldenTensor")
+      .def(py::init([](std::string name, std::vector<int64_t> shape,
+                       std::vector<int64_t> strides,
+                       ::tt::target::DataType dtype, std::uintptr_t ptr,
+                       std::size_t dataSize) {
+        // Create Golden Tensor and move ownership to GoldenTensor
+        auto *dataPtr = reinterpret_cast<std::uint8_t *>(ptr);
+
+        return std::make_shared<mlir::tt::GoldenTensor>(
+            name, shape, strides, dtype,
+            std::vector<std::uint8_t>(dataPtr, dataPtr + dataSize));
+      }))
       .def_readwrite("name", &mlir::tt::GoldenTensor::name)
       .def_readwrite("shape", &mlir::tt::GoldenTensor::shape)
       .def_readwrite("strides", &mlir::tt::GoldenTensor::strides)
       .def_readwrite("dtype", &mlir::tt::GoldenTensor::dtype)
       .def_readwrite("data", &mlir::tt::GoldenTensor::data);
 
-  m.def("create_golden_tensor",
-        [](std::string name, std::vector<int64_t> shape,
-           std::vector<int64_t> strides, ::tt::target::DataType dtype,
-           std::uintptr_t ptr) {
-          return mlir::tt::GoldenTensor(name, shape, strides, dtype,
-                                        reinterpret_cast<std::uint8_t *>(ptr));
-        });
-
   py::class_<mlir::tt::MLIRModuleLogger,
              std::shared_ptr<mlir::tt::MLIRModuleLogger>>(m, "MLIRModuleLogger")
       .def(py::init<>())

python/test_infra/ttir_builder.py

@@ -9,7 +9,7 @@
 from typing import List, Optional, Union, Tuple, Callable, Dict, Any
 from ttmlir.ir import *
 from ttmlir.dialects import ttir, tt, tensor
-from ttmlir.passes import create_golden_tensor, DataType
+from ttmlir.passes import GoldenTensor, DataType
 import torch
 
 # Alias for operands of ops which can be either BlockArguments, Values, or other
@@ -172,12 +172,13 @@ def get_golden_map(self) -> Dict:
         golden_info = {}
         for name, golden_tensor in self.id_golden_map.items():
             golden_tensor = golden_tensor.contiguous()
-            golden_info[name] = create_golden_tensor(
+            golden_info[name] = GoldenTensor(
                 name,
                 list(golden_tensor.tensor.shape),
                 list(golden_tensor.tensor.stride()),
                 DataType.Float32,
                 golden_tensor.tensor.data_ptr(),
+                golden_tensor.tensor.numel() * golden_tensor.tensor.dtype.itemsize,
             )
         return golden_info
 

tools/explorer/tt_adapter/src/tt_adapter/runner.py

@@ -241,21 +241,20 @@ def compile_and_run(self, model_path, overrides_string):
                 data = entry["value"]
                 # Turn this into a Torch Tensor to easily format it for the GoldenMap
                 # data is a uint8_t buffer type that contains the data in the format of dtype
-                # We will need to render this data as a buffer reference for the create_golden_tensor function
+                # We will need to render this data as a buffer reference for the GoldenTensor constructor
                 import array
 
                 # B is unsigned char in the array library
-                # This will parse the data as a 1D Buffer of uint8_t, exactly the pointer type expected by create_golden_tensor
+                # This will parse the data as a 1D Buffer of uint8_t, exactly the pointer type expected
                 data_arr = array.array("B", data["data"])
-                kept_alive_data_arrs.append(data_arr)
-                # Weird keepalive measure for the GoldenData...?
 
-                rendered_golden_map[entry["key"]] = passes.create_golden_tensor(
+                rendered_golden_map[entry["key"]] = passes.GoldenTensor(
                     data["name"],
                     data["shape"],
                     data["stride"],
                     passes.lookup_dtype(data["dtype"]),
                     data_arr.buffer_info()[0],
+                    data_arr.buffer_info()[1],
                 )
 
             # Get module from file