Lowering input and weight names into the MLIR argument attributes (#92)

tenstorrent · Aug 16, 2024 · 82db124 · 82db124
1 parent 8d1b35c
commit 82db124
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/pybuda/csrc/passes/lower_to_mlir.cpp b/pybuda/csrc/passes/lower_to_mlir.cpp
@@ -174,6 +174,15 @@ class MLIRGenerator
             auto funcType = builder_.getType<mlir::FunctionType>(mlir::TypeRange(argument_types), mlir::TypeRange(returns));
             auto func = builder_.create<mlir::func::FuncOp>(graphModule_.getLoc(), "main", funcType);
 
+            // Set the function argument names
+            for(size_t i = 0; i < argument_nodes.size(); i++)
+            {
+                graphlib::Node* argument_node = argument_nodes[i];
+                llvm::SmallVector<mlir::NamedAttribute, 1> named_attributes;
+                named_attributes.push_back(builder_.getNamedAttr("ttir.name", builder_.getStringAttr(argument_node->name())));
+                func.setArgAttrs(i, named_attributes);
+            }
+
             // Start the body of the function by creating an entry block.
             mlir::Block *entryBlock = func.addEntryBlock();