Nikolas Feedback

tenstorrent · Aug 7, 2024 · 8d76667 · 8d76667
1 parent 1680064
commit 8d76667
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 5 deletions.
diff --git a/lib/Dialect/TT/IR/TTOpsTypes.cpp b/lib/Dialect/TT/IR/TTOpsTypes.cpp
@@ -503,13 +503,17 @@ uint64_t TileType::getSizeBytes() const {
 }
 
 SystemDescAttr mlir::tt::getCurrentScopeSystemDesc(mlir::Operation *op) {
+  // Walk up scope levels until we find the top level ModuleOp which carries the
+  // system desc
   while (op) {
-    if (auto systemDesc =
-            op->getAttrOfType<SystemDescAttr>(SystemDescAttr::name)) {
+    if (mlir::isa<mlir::ModuleOp>(op)) {
+      auto systemDesc = op->getAttrOfType<SystemDescAttr>(SystemDescAttr::name);
+      assert(systemDesc && "expected system desc to be present on the module");
       return systemDesc;
     }
     op = op->getParentOp();
   }
+  assert(false && "expected system desc to be present in the scope");
   return nullptr;
 }
 

diff --git a/lib/Dialect/TTIR/Transforms/Passes.cpp b/lib/Dialect/TTIR/Transforms/Passes.cpp
@@ -657,7 +657,7 @@ class TTIRAllocate : public impl::TTIRAllocateBase<TTIRAllocate> {
         argumentAllocations.push_back(rewriter.getAttr<ArgumentAllocationAttr>(
             address, sizeBytes, memorySpace));
       }
-      func->setDiscardableAttr("argument_allocations",
+      func->setDiscardableAttr(ArgumentAllocationAttr::name,
                                rewriter.getArrayAttr(argumentAllocations));
 
       func->walk([&](tensor::EmptyOp empty) {

diff --git a/lib/Dialect/TTMetal/Transforms/SerializeToBinary.cpp b/lib/Dialect/TTMetal/Transforms/SerializeToBinary.cpp
@@ -115,8 +115,8 @@ class TTMetalSerializeToBinary
     cqBuilder.name = entry.getSymName().data();
 
     auto argumentAllocations = mlir::cast<ArrayAttr>(
-        entry->getDiscardableAttr("argument_allocations"));
-    assert(argumentAllocations && "expected argument_allocations attribute");
+        entry->getDiscardableAttr(ArgumentAllocationAttr::name));
+    assert(argumentAllocations && "expected tt.argument_allocations attribute");
     for (auto &input : entry.getBody().getArguments()) {
       auto argAlloc = mlir::cast<tt::ArgumentAllocationAttr>(
           argumentAllocations[input.getArgNumber()]);