From 07d8cd0edadce74d7f3c75e0f052d5dbf9fd2d15 Mon Sep 17 00:00:00 2001
From: Tiago Trevisan Jost <tiago.trevisanjost@amd.com>
Date: Thu, 1 Jun 2023 07:59:25 +0000
Subject: [PATCH] Support lowering tosa.custom_op to another dialect operation.

---
 mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
index 3f970befa38dc..a7750a7f7518c 100644
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
@@ -506,6 +506,13 @@ createLinalgBodyCalculationForElementwiseOp(Operation *op, ValueRange args,
     }
   }
 
+  // tosa::CustomOp
+  if (auto customOp = dyn_cast<tosa::CustomOp>(op)) {
+    return llvm::StringSwitch<Value>(customOp.getIdentifierAttr().str())
+        .Case("atan2", rewriter.create<math::Atan2Op>(loc, resultTypes, args))
+        .Default(nullptr);
+  }
+
   (void)rewriter.notifyMatchFailure(
       op, "unhandled op for linalg body calculation for elementwise op");
   return nullptr;
@@ -2067,6 +2074,7 @@ void mlir::tosa::populateTosaToLinalgConversionPatterns(
       PointwiseConverter<tosa::FloorOp>,
       PointwiseConverter<tosa::ClampOp>,
       PointwiseConverter<tosa::SigmoidOp>,
+      PointwiseConverter<tosa::CustomOp>,
       IdentityNConverter<tosa::IdentityOp>,
       ReduceConverter<tosa::ReduceAllOp>,
       ReduceConverter<tosa::ReduceAnyOp>,