added tt::LayoutAttr to the backend interface

tenstorrent · Nov 1, 2024 · 3404e19 · 3404e19
1 parent 301383a
commit 3404e19
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 9 deletions.
diff --git a/include/ttmlir/Dialect/TTNN/IR/TTNNOpsBackendInterfaces.td b/include/ttmlir/Dialect/TTNN/IR/TTNNOpsBackendInterfaces.td
@@ -20,7 +20,7 @@ def TTNNOpBackendInterface : OpInterface<"TTNNOpBackend"> {
             }],
             /*retTy=*/"size_t",
             /*methodName=*/"getOpPerfCycles",
-            /*args=*/(ins), // TBD
+            /*args=*/(ins "const tt::LayoutAttr&":$output_layout), // Subject to change
             /*methodBody=*/"",
             /*defaultImplementation=*/"return std::numeric_limits<size_t>::max();"
         >,
@@ -30,7 +30,7 @@ def TTNNOpBackendInterface : OpInterface<"TTNNOpBackend"> {
             }],
             /*retTy=*/"size_t",
             /*methodName=*/"getOpL1Usage",
-            /*args=*/(ins), // TBD
+            /*args=*/(ins "const tt::LayoutAttr&":$output_layout), // Subject to change
             /*methodBody=*/"",
             /*defaultImplementation=*/"return 0;"
         >,
@@ -40,7 +40,7 @@ def TTNNOpBackendInterface : OpInterface<"TTNNOpBackend"> {
             }],
             /*retTy=*/"bool",
             /*methodName=*/"isOpLegal",
-            /*args=*/(ins), // TBD
+            /*args=*/(ins "const tt::LayoutAttr&":$output_layout), // Subject to change
             /*methodBody=*/"",
             /*defaultImplementation=*/"return true;"
         >,

diff --git a/lib/Dialect/TTNN/Analysis/LegalGridAnalysis.cpp b/lib/Dialect/TTNN/Analysis/LegalGridAnalysis.cpp
@@ -16,9 +16,9 @@ bool mock_is_output_tensor_legal_for_op(Operation *op, tt::LayoutAttr layout) {
 
   // if (TTNNOpBackend backend = dyn_cast<TTNNOpBackend>(op)) {
   //   llvm::outs() << op->getName() << "=" << layout << "\n";
-  //   llvm::outs() << "\t[Perf] = " << backend.getOpPerfCycles() << "\n";
-  //   llvm::outs() << "\t[L1] = " << backend.getOpL1Usage() << "\n";
-  //   llvm::outs() << "\t[Legal] = " << backend.isOpLegal() << "\n";
+  //   llvm::outs() << "\t[Perf] = " << backend.getOpPerfCycles(layout) << "\n";
+  //   llvm::outs() << "\t[L1] = " << backend.getOpL1Usage(layout) << "\n";
+  //   llvm::outs() << "\t[Legal] = " << backend.isOpLegal(layout) << "\n";
   // }
 
   return true;

diff --git a/lib/Dialect/TTNN/IR/TTNNOpsBackendInterfaces.cpp b/lib/Dialect/TTNN/IR/TTNNOpsBackendInterfaces.cpp
@@ -13,17 +13,17 @@ namespace mlir::tt::ttnn {
 //===----------------------------------------------------------------------===//
 
 // // Relu backend interface
-size_t ReluOp::getOpPerfCycles() {
+size_t ReluOp::getOpPerfCycles(const tt::LayoutAttr &output_layout) {
   // Implement a custom estimate for relu op cycles.
   return 5;
 }
 
-size_t ReluOp::getOpL1Usage() {
+size_t ReluOp::getOpL1Usage(const tt::LayoutAttr &output_layout) {
   // Implement a custom estimate for relu op L1 usage.
   return 10;
 }
 
-bool ReluOp::isOpLegal() {
+bool ReluOp::isOpLegal(const tt::LayoutAttr &output_layout) {
   // Implement a custom check for relu op legality.
   return true;
 }