TTNNOpsBackendInterface definition

include/ttmlir/Dialect/TTNN/IR/CMakeLists.txt

@@ -2,6 +2,8 @@ add_mlir_dialect(TTNNOps ttnn)
 add_mlir_doc(TTNNBase TTNNDialect src/autogen/md/Dialect/ -gen-dialect-doc)
 add_mlir_doc(TTNNOps TTNNOp src/autogen/md/Dialect/ -gen-op-doc)
 
+add_mlir_interface(TTNNOpsBackendInterfaces)
+
 set(LLVM_TARGET_DEFINITIONS TTNNOpsEnums.td)
 mlir_tablegen(TTNNOpsEnums.h.inc -gen-enum-decls)
 mlir_tablegen(TTNNOpsEnums.cpp.inc -gen-enum-defs)

include/ttmlir/Dialect/TTNN/IR/TTNNBase.td

@@ -6,6 +6,7 @@
 #define TTMLIR_TTMLIR_DIALECT_TTNN_TTNNDIALECT_TD
 
 include "mlir/IR/OpBase.td"
+include "ttmlir/Dialect/TTNN/IR/TTNNOpsBackendInterfaces.td"
 
 //===----------------------------------------------------------------------===//
 // TTNN dialect definition.
@@ -43,6 +44,6 @@ def TTNN_Dialect : Dialect {
 //===----------------------------------------------------------------------===//
 
 class TTNN_Op<string mnemonic, list<Trait> traits = []> :
-        Op<TTNN_Dialect, mnemonic, traits>;
+        Op<TTNN_Dialect, mnemonic, !listconcat(traits, [TTNNOpBackendInterface])>;
 
 #endif

include/ttmlir/Dialect/TTNN/IR/TTNNOps.h

@@ -15,6 +15,7 @@
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "ttmlir/Dialect/TT/IR/TTOpsTypes.h"
 #include "ttmlir/Dialect/TTNN/IR/TTNNOpsAttrs.h"
+#include "ttmlir/Dialect/TTNN/IR/TTNNOpsBackendInterfaces.h.inc"
 #include "ttmlir/Dialect/TTNN/IR/TTNNOpsTypes.h"
 
 #define GET_OP_CLASSES

include/ttmlir/Dialect/TTNN/IR/TTNNOps.td

@@ -202,7 +202,9 @@ def TTNN_ReciprocalOp : TTNN_ElementwiseUnaryOp<"reciprocal"> {
     }];
 }
 
-def TTNN_ReluOp : TTNN_ElementwiseUnaryOp<"relu"> {
+def TTNN_ReluOp : TTNN_ElementwiseUnaryOp<"relu",
+      [DeclareOpInterfaceMethods<TTNNOpBackendInterface, ["getOpPerfCycles", "getOpL1Usage", "isOpLegal"]>]
+      > {
     let summary = "Eltwise ReLU.";
     let description = [{
       Eltwise ReLU operation.

include/ttmlir/Dialect/TTNN/IR/TTNNOpsBackendInterfaces.td

@@ -0,0 +1,50 @@
+// SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#ifndef TTMLIR_TTMLIR_DIALECT_TTNN_TTNNOPSPERFINTERFACES_TD
+#define TTMLIR_TTMLIR_DIALECT_TTNN_TTNNOPSPERFINTERFACES_TD
+
+include "mlir/IR/OpBase.td"
+
+def TTNNOpBackendInterface : OpInterface<"TTNNOpBackend"> {
+    let description = [{
+        Interface to access a registered method to infer the return types for an
+        operation that can be used during type inference.
+    }];
+    let cppNamespace = "::mlir::tt::ttnn";
+    let methods = [
+        InterfaceMethod<
+            /*desc=*/[{
+                Return the op kernel estimate in clock cycles.
+            }],
+            /*retTy=*/"size_t",
+            /*methodName=*/"getOpPerfCycles",
+            /*args=*/(ins "const tt::LayoutAttr&":$output_layout), // Subject to change
+            /*methodBody=*/"",
+            /*defaultImplementation=*/"return std::numeric_limits<size_t>::max();"
+        >,
+        InterfaceMethod<
+            /*desc=*/[{
+                Return the op kernel estimate in clock cycles.
+            }],
+            /*retTy=*/"size_t",
+            /*methodName=*/"getOpL1Usage",
+            /*args=*/(ins "const tt::LayoutAttr&":$output_layout), // Subject to change
+            /*methodBody=*/"",
+            /*defaultImplementation=*/"return 0;"
+        >,
+        InterfaceMethod<
+            /*desc=*/[{
+                Return the op kernel estimate in clock cycles.
+            }],
+            /*retTy=*/"bool",
+            /*methodName=*/"isOpLegal",
+            /*args=*/(ins "const tt::LayoutAttr&":$output_layout), // Subject to change
+            /*methodBody=*/"",
+            /*defaultImplementation=*/"return true;"
+        >,
+        ];
+}
+
+#endif // TTMLIR_TTMLIR_DIALECT_TTNN_TTNNOPSPERFINTERFACES_TD

lib/Dialect/TTNN/Analysis/LegalGridAnalysis.cpp

@@ -11,7 +11,13 @@ namespace mlir::tt::ttnn {
 
 bool mock_is_output_tensor_legal_for_op(Operation *op, tt::LayoutAttr layout) {
   // Placeholder, needs to be replaced with a call the the TTNN op interface.
-  return true;
+
+  if (TTNNOpBackend backend = dyn_cast<TTNNOpBackend>(op)) {
+    return backend.isOpLegal(layout);
+  }
+
+  assert(false && "Op is not a TTNN op.");
+  return false;
 }
 
 bool tensor_shape_compatible_with_shard(Operation *op, tt::LayoutAttr layout) {

lib/Dialect/TTNN/IR/CMakeLists.txt

@@ -1,6 +1,7 @@
 add_mlir_dialect_library(MLIRTTNNDialect
         TTNNDialect.cpp
         TTNNOps.cpp
+        TTNNOpsBackendInterfaces.cpp
         TTNNOpsTypes.cpp
 
         ADDITIONAL_HEADER_DIRS

lib/Dialect/TTNN/IR/TTNNOpsBackendInterfaces.cpp

@@ -0,0 +1,31 @@
+// SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include "ttmlir/Dialect/TTNN/IR/TTNNOps.h"
+
+#include "ttmlir/Dialect/TTNN/IR/TTNNOpsBackendInterfaces.cpp.inc"
+
+namespace mlir::tt::ttnn {
+
+//===----------------------------------------------------------------------===//
+// ReluOp
+//===----------------------------------------------------------------------===//
+
+// // Relu backend interface
+size_t ReluOp::getOpPerfCycles(const tt::LayoutAttr &output_layout) {
+  // Implement a custom estimate for relu op cycles.
+  return 5;
+}
+
+size_t ReluOp::getOpL1Usage(const tt::LayoutAttr &output_layout) {
+  // Implement a custom estimate for relu op L1 usage.
+  return 10;
+}
+
+bool ReluOp::isOpLegal(const tt::LayoutAttr &output_layout) {
+  // Implement a custom check for relu op legality.
+  return true;
+}
+
+} // namespace mlir::tt::ttnn