From 67449521f8b6ec323ba304ca5ac6a1ea2cf5eb0b Mon Sep 17 00:00:00 2001
From: Aleksandar Zecevic <azecevic@tenstorrent.com>
Date: Fri, 13 Dec 2024 18:44:08 +0100
Subject: [PATCH] ASAN fix of failing tests (#1569)

Both bugs are a result of
dangling references of `ArrayRef` objects, because they receive
temporary `SmallVector` objects.

Fixes https://github.com/tenstorrent/tt-mlir/issues/1565
---
 lib/Dialect/TTNN/Transforms/Passes.cpp    | 5 +++--
 lib/Dialect/TTNN/Utils/TransformUtils.cpp | 2 +-
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/lib/Dialect/TTNN/Transforms/Passes.cpp b/lib/Dialect/TTNN/Transforms/Passes.cpp
index c842c4075..20172f4fd 100644
--- a/lib/Dialect/TTNN/Transforms/Passes.cpp
+++ b/lib/Dialect/TTNN/Transforms/Passes.cpp
@@ -143,7 +143,7 @@ class TTNNDecomposeLayouts
     ttnn::Layout layoutEnum;
     DataType dataType;
     ttnn::TensorMemoryLayoutAttr tensorMemoryLayout;
-    llvm::ArrayRef<int64_t> shardShape;
+    llvm::SmallVector<int64_t> shardShape;
 
     ttnn::MemoryConfigAttr createMemoryConfigAttr(MLIRContext *context) const {
       return ttnn::MemoryConfigAttr::get(
@@ -233,7 +233,8 @@ class TTNNDecomposeLayouts
     output.tensorMemoryLayout = outputMemoryConfig.getTensorMemoryLayout();
 
     input.shardShape = inputLayoutAttr.getShardShape();
-    output.shardShape = outputMemoryConfig.getShardShapeArray();
+    output.shardShape =
+        llvm::SmallVector<int64_t>{outputMemoryConfig.getShardShapeArray()};
     return {input, output};
   }
 
diff --git a/lib/Dialect/TTNN/Utils/TransformUtils.cpp b/lib/Dialect/TTNN/Utils/TransformUtils.cpp
index ee5c1fd85..bef022fda 100644
--- a/lib/Dialect/TTNN/Utils/TransformUtils.cpp
+++ b/lib/Dialect/TTNN/Utils/TransformUtils.cpp
@@ -21,7 +21,7 @@ Value getOrInsertDevice(PatternRewriter &rewriter, Operation *op) {
   DeviceAttr deviceAttr = getCurrentScopeDevice(op);
   auto currentInsertionPoint = rewriter.saveInsertionPoint();
   rewriter.setInsertionPoint(block, block->begin());
-  auto meshShape = deviceAttr.getMeshShape();
+  llvm::SmallVector<int64_t> meshShape{deviceAttr.getMeshShape()};
   if (meshShape.empty()) {
     meshShape = llvm::SmallVector<int64_t, 2>{1, 1};
   }