From 7c8b8a5a28c88710506e567b683c3231919e2a09 Mon Sep 17 00:00:00 2001
From: Bezulj Marko <mbezulj@tenstorrent.com>
Date: Mon, 30 Dec 2024 14:05:38 +0000
Subject: [PATCH] fix toCoreRangeSet and virtual/physical grids

---
 include/ttmlir/Dialect/TT/Utils/CoreRangeSet.h | 8 ++++----
 include/ttmlir/Target/Utils/MLIRToFlatbuffer.h | 6 +++++-
 lib/OpModel/TTNN/Conversion.cpp                | 4 +++-
 3 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/include/ttmlir/Dialect/TT/Utils/CoreRangeSet.h b/include/ttmlir/Dialect/TT/Utils/CoreRangeSet.h
index e61e82962..f35b35c54 100644
--- a/include/ttmlir/Dialect/TT/Utils/CoreRangeSet.h
+++ b/include/ttmlir/Dialect/TT/Utils/CoreRangeSet.h
@@ -26,12 +26,12 @@ using locSize2d = std::tuple<std::array<int32_t, 2>,  //  {{locX, locY},
 /// \param virtualGrid The virtual grid attributes.
 /// \returns A vector of core ranges, where each core range is represented by
 ///          a pair of location and size (both 2D).
-inline std::vector<locSize2d> toCoreRangeSet(GridAttr virtualGrid) {
-  llvm::SmallVector<std::int64_t> tensorGridShape(virtualGrid.getShape());
+inline std::vector<locSize2d>
+toCoreRangeSet(const llvm::ArrayRef<int64_t> virtualGridShape,
+               const mlir::AffineMap mapping) {
   std::vector<locSize2d> coreRangeSet;
-  AffineMap mapping = virtualGrid.getMapping();
   ::ttmlir::utils::sample(
-      tensorGridShape, [&](ArrayRef<std::int64_t> virtualCoreCoord) {
+      virtualGridShape, [&](ArrayRef<std::int64_t> virtualCoreCoord) {
         llvm::SmallVector<std::int64_t> coreCoord =
             mapping.compose(virtualCoreCoord);
         assert(coreCoord.size() == PhysGridResultIdx::NumIndices &&
diff --git a/include/ttmlir/Target/Utils/MLIRToFlatbuffer.h b/include/ttmlir/Target/Utils/MLIRToFlatbuffer.h
index c1d4951a1..3b365e748 100644
--- a/include/ttmlir/Target/Utils/MLIRToFlatbuffer.h
+++ b/include/ttmlir/Target/Utils/MLIRToFlatbuffer.h
@@ -370,7 +370,11 @@ toFlatbuffer(FlatbufferObjectCache &cache, GridAttr tensorGrid,
              GridAttr deviceGrid) {
   std::vector<::tt::target::Dim2dRange> coreRangeSet;
 
-  for (const auto &locsize2d : utils::toCoreRangeSet(tensorGrid)) {
+  auto mapping = (tensorGrid.getMapping().isEmpty() == true)
+                     ? deviceGrid.getMapping()
+                     : tensorGrid.getMapping();
+  for (const auto &locsize2d :
+       utils::toCoreRangeSet(tensorGrid.getShape(), mapping)) {
     const auto &[loc, size] = locsize2d;
     coreRangeSet.push_back(
         ::tt::target::Dim2dRange(::tt::target::Dim2d(loc[1], loc[0]),
diff --git a/lib/OpModel/TTNN/Conversion.cpp b/lib/OpModel/TTNN/Conversion.cpp
index 0b4564d67..ce9ab2c9b 100644
--- a/lib/OpModel/TTNN/Conversion.cpp
+++ b/lib/OpModel/TTNN/Conversion.cpp
@@ -71,7 +71,9 @@ getPageLayout(const mlir::tt::ttnn::TTNNLayoutAttr &layout) {
 ::tt::tt_metal::CoreRangeSet
 getCoreRangeSet(const mlir::tt::ttnn::TTNNLayoutAttr &layout) {
   std::set<::tt::tt_metal::CoreRange> coreRangeSet;
-  for (const auto &[loc, size] : utils::toCoreRangeSet(layout.getGrid())) {
+  assert(layout.getGrid().getMapping().isEmpty() == false);
+  for (const auto &[loc, size] : utils::toCoreRangeSet(
+           layout.getGrid().getShape(), layout.getGrid().getMapping())) {
     coreRangeSet.insert(::tt::tt_metal::CoreRange(
         CoreCoord(loc[0], loc[1]),
         CoreCoord(loc[0] + size[0] - 1, loc[1] + size[1] - 1)));