mmanzoorTT · mmanzoorTT · Jul 17, 2024
diff --git a/include/ttmlir/Dialect/TTIR/IR/TTIROps.td b/include/ttmlir/Dialect/TTIR/IR/TTIROps.td
@@ -153,6 +153,13 @@ def TTIR_MultiplyOp : TTIR_ElementwiseBinaryOp<"multiply"> {
     }];
 }
 
+def TTIR_GreaterEqualOp :  TTIR_ElementwiseBinaryOp<"ge"> {
+    let summary = "Eltwise greater than or equal to.";
+    let description = [{
+      Eltwise greater than or equal to operation.
+    }];
+}
+
 class TTIR_ReductionOp<string mnemonic, list<Trait> traits = []> : TTIR_DPSOp<mnemonic, traits> {
     let summary = "Reduction op.";
     let description = [{

diff --git a/include/ttmlir/Dialect/TTNN/IR/TTNNOps.td b/include/ttmlir/Dialect/TTNN/IR/TTNNOps.td
@@ -92,6 +92,13 @@ def TTNN_MultiplyOp : TTNN_ElementwiseBinaryOp<"multiply"> {
     }];
 }
 
+def TTNN_GreaterEqualOp : TTNN_ElementwiseBinaryOp<"ge"> {
+    let summary = "Eltwise greater than or equal to.";
+    let description = [{
+      Eltwise greater than or equal to operation.
+    }];
+}
+
 class TTNN_ReductionOp<string mnemonic, list<Trait> traits = []> : TTNN_NamedDPSOp<mnemonic, traits> {
     let summary = "Reduction op.";
     let description = [{

diff --git a/include/ttmlir/Target/TTNN/program.fbs b/include/ttmlir/Target/TTNN/program.fbs
@@ -29,6 +29,7 @@ enum EltwiseOpType: uint32 {
   Multiply = 1,
   Subtract = 2,
   Relu = 3,
+  GreaterEqual = 4,
 }
 
 table EltwiseOp {

diff --git a/lib/Conversion/TTNNToEmitC/TTNNToEmitC.cpp b/lib/Conversion/TTNNToEmitC/TTNNToEmitC.cpp
@@ -52,6 +52,8 @@ void populateTTNNToEmitCPatterns(mlir::MLIRContext *ctx,
   patterns.add<DefaultOpConversionPattern<ttnn::AddOp>>(typeConverter, ctx);
   patterns.add<DefaultOpConversionPattern<ttnn::SubtractOp>>(typeConverter,
                                                              ctx);
+  patterns.add<DefaultOpConversionPattern<ttnn::GreaterEqualOp>>(typeConverter,
+                                                                 ctx);
   patterns.add<DefaultOpConversionPattern<ttnn::SumOp>>(typeConverter, ctx);
   patterns.add<DefaultOpConversionPattern<ttnn::SoftmaxOp>>(typeConverter, ctx);
   patterns.add<DefaultOpConversionPattern<ttnn::MultiplyOp>>(typeConverter,

diff --git a/lib/Dialect/TTIR/Transforms/Passes.cpp b/lib/Dialect/TTIR/Transforms/Passes.cpp
@@ -75,6 +75,9 @@ class ConvertTosaToTTIR
                  TosaToTTIREltwiseBinaryRewriter<tosa::MulOp, ttir::MultiplyOp,
                                                  OperandConstraint::AnyDevice>,
                  TosaToTTIREltwiseBinaryRewriter<tosa::SubOp, ttir::SubtractOp,
+                                                 OperandConstraint::AnyDevice>,
+                 TosaToTTIREltwiseBinaryRewriter<tosa::GreaterEqualOp,
+                                                 ttir::GreaterEqualOp,
                                                  OperandConstraint::AnyDevice>>(
         &getContext());
     FrozenRewritePatternSet patternSet(std::move(patterns));
@@ -115,6 +118,9 @@ class TTIRNamedToKernelRewriter : public OpRewritePattern<TTIROp> {
     } else if constexpr (std::is_same<TTIROp, ttir::SubtractOp>::value) {
       kernelName = "subtract";
       kernelKind = "eltwise";
+    } else if constexpr (std::is_same<TTIROp, ttir::GreaterEqualOp>::value) {
+      kernelName = "ge";
+      kernelKind = "eltwise";
     } else if constexpr (std::is_same<TTIROp, ttir::ReluOp>::value) {
       kernelName = "relu";
       kernelKind = "eltwise";
@@ -272,6 +278,7 @@ class TTIRGeneric : public impl::TTIRGenericBase<TTIRGeneric> {
                  TTIRNamedToKernelRewriter<AddOp>,
                  TTIRNamedToKernelRewriter<MultiplyOp>,
                  TTIRNamedToKernelRewriter<SubtractOp>,
+                 TTIRNamedToKernelRewriter<GreaterEqualOp>,
                  TTIRNamedToKernelRewriter<ReluOp>>(&getContext());
     FrozenRewritePatternSet patternSet(std::move(patterns));
     if (failed(applyPatternsAndFoldGreedily(getOperation(), patternSet))) {
@@ -567,6 +574,7 @@ class TTIRLayout : public impl::TTIRLayoutBase<TTIRLayout> {
           TTIRLayoutOperandsRewriter<AddOp>,
           TTIRLayoutOperandsRewriter<MultiplyOp>,
           TTIRLayoutOperandsRewriter<SubtractOp>,
+          TTIRLayoutOperandsRewriter<GreaterEqualOp>,
           TTIRLayoutOperandsRewriter<ReluOp>, TTIRLayoutOperandsRewriter<SumOp>,
           TTIRLayoutOperandsRewriter<SoftmaxOp>,
           TTIRLayoutOperandsRewriter<MatmulOp>, TTIRLayoutFuncReturnRewriter>(

diff --git a/lib/Dialect/TTNN/Transforms/Passes.cpp b/lib/Dialect/TTNN/Transforms/Passes.cpp
@@ -159,6 +159,7 @@ class ConvertTTIRToTTNN
         .add<TTIRToTTNNLayoutRewriter, TTIRToTTNNOpRewriter<ttir::AddOp, AddOp>,
              TTIRToTTNNOpRewriter<ttir::MultiplyOp, MultiplyOp>,
              TTIRToTTNNOpRewriter<ttir::SubtractOp, SubtractOp>,
+             TTIRToTTNNOpRewriter<ttir::GreaterEqualOp, GreaterEqualOp>,
              TTIRToTTNNOpRewriter<ttir::ReluOp, ReluOp>,
              TTIRToTTNNBinaryOpRewriter<ttir::MatmulOp, MatmulOp>,
              TTIRToTTNNReductionOpRewriter<ttir::SumOp, SumOp>,

diff --git a/lib/Dialect/TTNN/Transforms/TTNNToSerializedBinary.cpp b/lib/Dialect/TTNN/Transforms/TTNNToSerializedBinary.cpp
@@ -113,6 +113,8 @@ createEltwiseOp(FlatbufferObjectCache &cache, EltwiseOp op) {
     type = ::tt::target::ttnn::EltwiseOpType::Multiply;
   } else if constexpr (std::is_same_v<EltwiseOp, SubtractOp>) {
     type = ::tt::target::ttnn::EltwiseOpType::Subtract;
+  } else if constexpr (std::is_same_v<EltwiseOp, GreaterEqualOp>) {
+    type = ::tt::target::ttnn::EltwiseOpType::GreaterEqual;
   } else if constexpr (std::is_same_v<EltwiseOp, ReluOp>) {
     type = ::tt::target::ttnn::EltwiseOpType::Relu;
   } else {
@@ -191,6 +193,9 @@ emitTTNNOperation(FlatbufferObjectCache &cache, Operation *op,
     return createOperation(cache, createEltwiseOp(cache, subtractOp),
                            debugString);
   }
+  if (auto geOp = dyn_cast<GreaterEqualOp>(op); geOp) {
+    return createOperation(cache, createEltwiseOp(cache, geOp), debugString);
+  }
   if (auto reluOp = dyn_cast<ReluOp>(op); reluOp) {
     return createOperation(cache, createEltwiseOp(cache, reluOp), debugString);
   }

diff --git a/runtime/lib/ttnn/program.cpp b/runtime/lib/ttnn/program.cpp
@@ -92,6 +92,14 @@ run(::tt::target::ttnn::EltwiseOp const *op, ::ttnn::Device &device,
     liveTensors.try_emplace(op->out()->global_id(), &tensorPool.back());
     break;
   }
+  case ::tt::target::ttnn::EltwiseOpType::GreaterEqual: {
+    assert(op->ins()->size() == 2 && "Unsupported number of inputs");
+    ::ttnn::Tensor &lhs = *liveTensors.at(op->ins()->Get(0)->global_id());
+    ::ttnn::Tensor &rhs = *liveTensors.at(op->ins()->Get(1)->global_id());
+    tensorPool.push_back(::ttnn::ge(lhs, rhs));
+    liveTensors.try_emplace(op->out()->global_id(), &tensorPool.back());
+    break;
+  }
   default:
     throw std::runtime_error("Unsupported elementwise operation type");
   }

diff --git a/test/ttmlir/Dialect/TTNN/simple_ge.mlir b/test/ttmlir/Dialect/TTNN/simple_ge.mlir
@@ -0,0 +1,15 @@
+// RUN: ttmlir-opt --ttir-layout --ttnn-open-device --convert-ttir-to-ttnn %s | FileCheck %s
+#any_device = #tt.operand_constraint<dram|l1|scalar|tile|any_device|any_device_tile>
+module attributes {tt.system_desc = #tt.system_desc<[{arch = <wormhole_b0>, grid = <8x8>, l1_size = 1048576, num_dram_channels = 12, dram_channel_size = 1048576, noc_l1_address_align_bytes = 16, pcie_address_align_bytes = 32, noc_dram_address_align_bytes = 32}], [0], [<pcie|host_mmio>], [<0, 0, 0, 0>]>} {
+  func.func @forward(%arg0: tensor<64x128xf32>, %arg1: tensor<64x128xf32>) -> tensor<64x128xf32> {
+    // CHECK: %[[C:.*]] = "ttnn.open_device"[[C:.*]]
+    // CHECK: %[[C:.*]] = "ttnn.full"[[C:.*]]
+    %0 = tensor.empty() : tensor<64x128xf32>
+    // CHECK: %[[C:.*]] = "ttnn.to_memory_config"[[C:.*]]
+    // CHECK: %[[C:.*]] = "ttnn.ge"[[C:.*]]
+    %1 = "ttir.ge"(%arg0, %arg1, %0) <{operandSegmentSizes = array<i32: 2, 1>, operand_constraints = [#any_device, #any_device, #any_device]}> : (tensor<64x128xf32>, tensor<64x128xf32>, tensor<64x128xf32>) -> tensor<64x128xf32>
+    // CHECK: %[[C:.*]] = "ttnn.to_memory_config"[[C:.*]]
+    // CHECK: "ttnn.close_device"[[C:.*]]
+    return %1 : tensor<64x128xf32>
+  }
+}