ASAN fix of failing tests (#1569)

Both bugs are a result of dangling references of `ArrayRef` objects, because they receive temporary `SmallVector` objects. Fixes #1565
tenstorrent · Dec 13, 2024 · 6744952 · 6744952
1 parent 9d34762
commit 6744952
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git 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
@@ -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};
   }