Adding sqrt op in ttmlir e2e

docs/src/adding-an-op.md

@@ -124,10 +124,13 @@ graph.
 Inside `runOnOperation` is usually where we define a rewrite pattern set that
 can match much more complicated patterns (nested inside of the `ModuleOp`'s
 [regions](https://mlir.llvm.org/docs/LangRef/#regions))
-than just a single operation.
+than just a single operation. In `runOperation` method you will see the call to
+method `populateTTIRToTTNNPatterns(...)` that actually generates rewrite patterns.
+Method `populateTTIRToTTNNPatterns(...)` is defined
+in `lib/Conversion/TTIRToTTNN/TTIRToTTNN.cpp`.
 
 ```cpp
-{{#include ../../../lib/Conversion/TTIRToTTNN/TTIRToTTNN.cpp:adding_an_op_matmul_rewrite_pattern_set}}
+{{#include ../../../lib/Conversion/TTIRToTTNN/TTIRToTTNN.cpp:op_rewriter_pattern_set}}
 ```
 
 > More information on rewrite patterns and their capabilities can be found in the MLIR documentation [here](https://mlir.llvm.org/docs/PatternRewriter/) and [here](https://mlir.llvm.org/docs/DialectConversion/).

include/ttmlir/Dialect/TTIR/IR/TTIROps.td

@@ -196,6 +196,13 @@ def TTIR_GreaterEqualOp :  TTIR_ElementwiseBinaryOp<"ge"> {
     }];
 }
 
+def TTIR_SqrtOp : TTIR_ElementwiseUnaryOp<"sqrt"> {
+    let summary = "Eltwise square root.";
+    let description = [{
+      Eltwise square root operation.
+    }];
+}
+
 def TTIR_ReluOp : TTIR_ElementwiseUnaryOp<"relu"> {
     let summary = "Eltwise ReLU.";
     let description = [{

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

@@ -63,12 +63,50 @@ class TTNN_ElementwiseOp<string mnemonic, list<Trait> traits = []> :
     let results = (outs Variadic<AnyRankedTensor>:$results);
 }
 
+class TTNN_ElementwiseUnaryOp<string mnemonic, list<Trait> traits = []> :
+    TTNN_ElementwiseOp<mnemonic, traits> {
+    let summary = "Eltwise unary op.";
+    let description = [{
+      Eltwise unary op.
+    }];
+
+    let builders =
+    [
+      OpBuilder<(ins "Value": $in, "Value": $out),
+      [{
+        build($_builder, $_state, {out.getType()}, in, out);
+      }]>
+    ];
+}
+
 class TTNN_ElementwiseBinaryOp<string mnemonic, list<Trait> traits = []> :
     TTNN_ElementwiseOp<mnemonic, traits> {
     let summary = "Eltwise binary op.";
     let description = [{
       Eltwise binary op.
     }];
+
+    let builders =
+    [
+      OpBuilder<(ins "Value": $lhs, "Value": $rhs, "Value": $out),
+      [{
+        build($_builder, $_state, {out.getType()}, {lhs, rhs}, out);
+      }]>
+    ];
+}
+
+def TTNN_SqrtOp : TTNN_ElementwiseUnaryOp<"sqrt"> {
+    let summary = "Eltwise sqrt.";
+    let description = [{
+      Eltwise sqrt operation.
+    }];
+}
+
+def TTNN_ReluOp : TTNN_ElementwiseUnaryOp<"relu"> {
+    let summary = "Eltwise ReLU.";
+    let description = [{
+      Eltwise ReLU operation.
+    }];
 }
 
 def TTNN_AddOp : TTNN_ElementwiseBinaryOp<"add"> {
@@ -131,13 +169,6 @@ def TTNN_MeanOp : TTNN_ReductionOp<"mean"> {
   }];
 }
 
-def TTNN_ReluOp : TTNN_ElementwiseOp<"relu"> {
-    let summary = "Eltwise ReLU.";
-    let description = [{
-      Eltwise ReLU operation.
-    }];
-}
-
 def TTNN_SoftmaxOp : TTNN_NamedDPSOp<"softmax"> {
     let summary = "Softmax op.";
     let description = [{

include/ttmlir/Target/TTNN/program.fbs

@@ -35,6 +35,7 @@ enum EltwiseOpType: uint32 {
   Subtract = 2,
   Relu = 3,
   GreaterEqual = 4,
+  Sqrt = 5,
 }
 
 table EltwiseOp {

lib/Conversion/TTIRToTTNN/TTIRToTTNN.cpp

@@ -149,7 +149,7 @@ namespace mlir::tt {
 void populateTTIRToTTNNPatterns(MLIRContext *ctx, RewritePatternSet &patterns,
                                 TypeConverter &typeConverter) {
   // clang-format off
-  // ANCHOR: adding_an_op_matmul_rewrite_pattern_set
+  // ANCHOR: op_rewriter_pattern_set
   patterns
       .add<TensorEmptyConversionPattern,
            ToLayoutOpConversionPattern,
@@ -158,13 +158,14 @@ void populateTTIRToTTNNPatterns(MLIRContext *ctx, RewritePatternSet &patterns,
            ElementwiseBinaryOpConversionPattern<ttir::MultiplyOp, ttnn::MultiplyOp>,
            ElementwiseBinaryOpConversionPattern<ttir::GreaterEqualOp, ttnn::GreaterEqualOp>,
            ElementwiseBinaryOpConversionPattern<ttir::ReluOp, ttnn::ReluOp>,
+           ElementwiseBinaryOpConversionPattern<ttir::SqrtOp, ttnn::SqrtOp>,
            ReductionOpConversionPattern<ttir::SumOp, ttnn::SumOp>,
            ReductionOpConversionPattern<ttir::MeanOp, ttnn::MeanOp>,
            TransposeOpConversionPattern,
            SoftmaxOpConversionPattern,
            MatmulOpConversionPattern
            >(typeConverter, ctx);
-  // ANCHOR_END: adding_an_op_matmul_rewrite_pattern_set
+  // ANCHOR_END: op_rewriter_pattern_set
   // clang-format on
 }
 

lib/Conversion/TTNNToEmitC/TTNNToEmitC.cpp

@@ -178,6 +178,7 @@ void populateTTNNToEmitCPatterns(mlir::MLIRContext *ctx,
   // Eltwise unary ops
   //
   patterns.add<DefaultOpConversionPattern<ttnn::ReluOp>>(typeConverter, ctx);
+  patterns.add<DefaultOpConversionPattern<ttnn::SqrtOp>>(typeConverter, ctx);
   patterns.add<DefaultOpConversionPattern<ttnn::SoftmaxOp>>(typeConverter, ctx);
 
   // Eltwise binary ops

lib/Target/TTNN/TTNNToFlatbuffer.cpp

@@ -133,6 +133,8 @@ createEltwiseOp(FlatbufferObjectCache &cache, EltwiseOp op) {
     type = ::tt::target::ttnn::EltwiseOpType::GreaterEqual;
   } else if constexpr (std::is_same_v<EltwiseOp, ReluOp>) {
     type = ::tt::target::ttnn::EltwiseOpType::Relu;
+  } else if constexpr (std::is_same_v<EltwiseOp, SqrtOp>) {
+    type = ::tt::target::ttnn::EltwiseOpType::Sqrt;
   } else {
     llvm_unreachable("unhandled EltwiseOp");
   }
@@ -234,6 +236,9 @@ emitTTNNOperation(FlatbufferObjectCache &cache, Operation *op,
   if (auto reluOp = dyn_cast<ReluOp>(op); reluOp) {
     return createOperation(cache, createEltwiseOp(cache, reluOp), debugString);
   }
+  if (auto sqrtOp = dyn_cast<SqrtOp>(op); sqrtOp) {
+    return createOperation(cache, createEltwiseOp(cache, sqrtOp), debugString);
+  }
   if (auto matmulOp = dyn_cast<MatmulOp>(op); matmulOp) {
     return createOperation(cache, createOp(cache, matmulOp), debugString);
   }

runtime/lib/ttnn/program.cpp

@@ -241,6 +241,13 @@ run(::tt::target::ttnn::EltwiseOp const *op, ::ttnn::Device &device,
     liveTensors.insert_or_assign(op->out()->global_id(), &tensorPool.back());
     break;
   }
+  case ::tt::target::ttnn::EltwiseOpType::Sqrt: {
+    assert(op->ins()->size() == 1 && "Unsupported number of inputs");
+    ::ttnn::Tensor &in = *liveTensors.at(op->ins()->Get(0)->global_id());
+    tensorPool.push_back(::ttnn::sqrt(in));
+    liveTensors.insert_or_assign(op->out()->global_id(), &tensorPool.back());
+    break;
+  }
   }
 }
 

test/ttmlir/Dialect/TTNN/eltwise/unary/sqrt/simple_sqrt.mlir

@@ -0,0 +1,13 @@
+// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline %s | FileCheck %s
+#any_device = #tt.operand_constraint<dram|l1|scalar|tile|any_device|any_device_tile>
+module attributes {} {
+  func.func @forward(%arg0: tensor<64x128xf32>) -> tensor<64x128xf32> {
+    // CHECK: %[[C:.*]] = "ttnn.open_device"[[C:.*]]
+    // CHECK: %[[C:.*]] = "ttnn.empty"[[C:.*]]
+    %0 = tensor.empty() : tensor<64x128xf32>
+    // CHECK: %[[C:.*]] = "ttnn.sqrt"[[C:.*]]
+    %1 = "ttir.sqrt"(%arg0, %0) <{operandSegmentSizes = array<i32: 1, 1>, operand_constraints = [#any_device, #any_device]}> : (tensor<64x128xf32>, tensor<64x128xf32>) -> tensor<64x128xf32>
+    // CHECK: "ttnn.close_device"[[C:.*]]
+    return %1 : tensor<64x128xf32>
+  }
+}

test/ttmlir/Silicon/TTNN/eltwise/unary/simple_sqrt.mlir

@@ -0,0 +1,15 @@
+// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline --ttir-load-system-desc="path=%system_desc_path%" %s > %t.mlir
+// RUN: FileCheck %s --input-file=%t.mlir
+// RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn
+#any_device = #tt.operand_constraint<dram|l1|scalar|tile|any_device|any_device_tile>
+module attributes {} {
+  func.func @forward(%arg0: tensor<64x128xf32>) -> tensor<64x128xf32> {
+    // CHECK: %[[C:.*]] = "ttnn.open_device"[[C:.*]]
+    // CHECK: %[[C:.*]] = "ttnn.empty"[[C:.*]]
+    %0 = tensor.empty() : tensor<64x128xf32>
+    // CHECK: %[[C:.*]] = "ttnn.sqrt"[[C:.*]]
+    %1 = "ttir.sqrt"(%arg0, %0) <{operandSegmentSizes = array<i32: 1, 1>, operand_constraints = [#any_device, #any_device]}> : (tensor<64x128xf32>, tensor<64x128xf32>) -> tensor<64x128xf32>
+    // CHECK: "ttnn.close_device"[[C:.*]]
+    return %1 : tensor<64x128xf32>
+  }
+}