fix toCoreRangeSet and virtual/physical grids

tenstorrent · Dec 30, 2024 · 7c8b8a5 · 7c8b8a5
1 parent 9b501fa
commit 7c8b8a5
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 6 deletions.
diff --git 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
@@ -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
@@ -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)));