From 8cc7c463c69a7b8ef10197e30b65789bc9f80b19 Mon Sep 17 00:00:00 2001
From: dgolubovic <dgolubovic@tenstorrent.com>
Date: Wed, 7 Aug 2024 15:17:41 +0000
Subject: [PATCH] [lower_to_mlir] Use named_attrs from tt-forge ops instead of
 hardcoding attributes for each op, while lowering to mlir.

---
 pybuda/csrc/passes/lower_to_mlir.cpp | 38 ++++++++++++++++++++--------
 pybuda/pybuda/op/nn.py               |  4 +--
 2 files changed, 30 insertions(+), 12 deletions(-)

diff --git a/pybuda/csrc/passes/lower_to_mlir.cpp b/pybuda/csrc/passes/lower_to_mlir.cpp
index 10bb809ec..72bc4171d 100644
--- a/pybuda/csrc/passes/lower_to_mlir.cpp
+++ b/pybuda/csrc/passes/lower_to_mlir.cpp
@@ -45,6 +45,7 @@ using namespace tt;
 /**
  * @brief Implementation of TT-MLIR emission from the TTForge graph.
  */
+
 class MLIRGenerator
 {
     public:
@@ -109,6 +110,25 @@ class MLIRGenerator
             symbolTable_[node->name()] = {value, node};
         }
 
+        // Convert a TTForge attribute to an MLIR attribute.
+        mlir::Attribute convert_to_mlir_attribute(const tt::BudaOpAttr& value) {
+            return std::visit([this](auto&& arg) -> mlir::Attribute {
+                using T = std::decay_t<decltype(arg)>;
+                if constexpr (std::is_same_v<T, std::string>) {
+                    return builder_.getStringAttr(arg);
+                } else if constexpr (std::is_same_v<T, bool>) {
+                    return builder_.getBoolAttr(arg);
+                } else if constexpr (std::is_same_v<T, int>) {
+                    return builder_.getI32IntegerAttr(arg);
+                } else if constexpr (std::is_same_v<T, float>) {
+                    return builder_.getF32FloatAttr(arg);
+                } else {
+                    // If type not handled, throw an exception or handle it appropriately
+                    throw std::runtime_error("Unhandled attribute type");
+                }
+            }, value);
+        }
+
         /// Emit a new function in MLIR.
         /// A function represents a set of TTForge operations that are executed to produce output results.
         /// This function will generate the MLIR code for each TTForge operation in the graph and emit the return operation for the function.
@@ -204,7 +224,7 @@ class MLIRGenerator
             ::llvm::ArrayRef<::llvm::StringRef> operation_attributes = TTIROp::getAttributeNames();
             for(auto attribute_name: operation_attributes)
             {
-                if(attribute_name.equals("operand_constraints"))
+                if(attribute_name == "operand_constraints")
                 {
                     // Create operation constraint attributes
                     mlir::NamedAttribute operand_constraints_attribute = builder_.getNamedAttr(
@@ -212,7 +232,7 @@ class MLIRGenerator
                     builder_.getArrayAttr(get_mlir_operand_constraint_attributes(graph, op_node)));
                     attributes.push_back(operand_constraints_attribute);
                 }
-                else if(attribute_name.equals(mlir::OpTrait::AttrSizedOperandSegments<void>::getOperandSegmentSizeAttr()))
+                else if(attribute_name == mlir::OpTrait::AttrSizedOperandSegments<void>::getOperandSegmentSizeAttr())
                 {
                     // Create operation segment sizes attributes
                     mlir::NamedAttribute operand_segment_sizes_attribute = builder_.getNamedAttr(
@@ -225,15 +245,13 @@ class MLIRGenerator
                 }
             }
 
-            // Workaround for now, need to figure out how to handle this properly
-            if(op_node->op_name() == "softmax")
+            for(const auto & attribute: op_node->op_type().named_attrs)
             {
-                log_info("Softmax");
-                int32_t dimension = std::get<int>(op_node->op_attrs()[0]);
-                mlir::NamedAttribute dimension_attribute = builder_.getNamedAttr(
-                        "dimension",
-                    builder_.getSI32IntegerAttr(dimension));
-                attributes.push_back(dimension_attribute);
+                // convert atribute to mlir atribute
+                auto mlir_atribute = convert_to_mlir_attribute(attribute.second);
+                mlir::NamedAttribute named_attribute = builder_.getNamedAttr(
+                                        attribute.first, mlir_atribute);
+                attributes.push_back(named_attribute);
             }
 
             auto op = builder_.create<TTIROp>(
diff --git a/pybuda/pybuda/op/nn.py b/pybuda/pybuda/op/nn.py
index 4dfe9282c..54228ed8a 100644
--- a/pybuda/pybuda/op/nn.py
+++ b/pybuda/pybuda/op/nn.py
@@ -51,7 +51,7 @@ def Softmax(
     Tensor
         Buda tensor
     """
-    return op("softmax", name, operandA, attrs=(dim, stable)).get_tensor()
+    return op("softmax", name, operandA, attrs=(dim, stable), dimension=dim).get_tensor()
 
 
 def LogSoftmax(
@@ -82,7 +82,7 @@ def LogSoftmax(
     Tensor
         Buda tensor
     """
-    return op("log_softmax", name, operandA, attrs=(dim, stable)).get_tensor()
+    return op("log_softmax", name, operandA, attrs=(dim, stable), dimension=dim).get_tensor()
 
 def Layernorm(
         name: str,