From 1ebf7ce950bb72599055d2f2789fd604a02b6d15 Mon Sep 17 00:00:00 2001
From: Tres Popp <tpopp@google.com>
Date: Tue, 1 Jun 2021 12:17:33 +0200
Subject: [PATCH] [mlir] Use interfaces in MathToLibm

Previously, this assumed use of ModuleOp and FuncOp. There is no need to
restrict this, and using interfaces allows these patterns to be used
during dialect conversion to LLVM.

Some assertions were removed due to inconsistent implementation of
FunctionLikeOps.

Differential Revision: https://reviews.llvm.org/D103447
---
 mlir/lib/Conversion/MathToLibm/MathToLibm.cpp | 16 ++++++++--------
 1 file changed, 8 insertions(+), 8 deletions(-)

diff --git a/mlir/lib/Conversion/MathToLibm/MathToLibm.cpp b/mlir/lib/Conversion/MathToLibm/MathToLibm.cpp
index 8512432681c2..7b71386a5294 100644
--- a/mlir/lib/Conversion/MathToLibm/MathToLibm.cpp
+++ b/mlir/lib/Conversion/MathToLibm/MathToLibm.cpp
@@ -81,30 +81,30 @@ template <typename Op>
 LogicalResult
 ScalarOpToLibmCall<Op>::matchAndRewrite(Op op,
                                         PatternRewriter &rewriter) const {
-  auto module = op->template getParentOfType<ModuleOp>();
+  auto module = SymbolTable::getNearestSymbolTable(op);
   auto type = op.getType();
   // TODO: Support Float16 by upcasting to Float32
   if (!type.template isa<Float32Type, Float64Type>())
     return failure();
 
   auto name = type.getIntOrFloatBitWidth() == 64 ? doubleFunc : floatFunc;
-  auto opFunc = module.template lookupSymbol<FuncOp>(name);
+  auto opFunc = dyn_cast_or_null<SymbolOpInterface>(
+      SymbolTable::lookupSymbolIn(module, name));
   // Forward declare function if it hasn't already been
   if (!opFunc) {
     OpBuilder::InsertionGuard guard(rewriter);
-    rewriter.setInsertionPointToStart(module.getBody());
+    rewriter.setInsertionPointToStart(&module->getRegion(0).front());
     auto opFunctionTy = FunctionType::get(
         rewriter.getContext(), op->getOperandTypes(), op->getResultTypes());
     opFunc =
         rewriter.create<FuncOp>(rewriter.getUnknownLoc(), name, opFunctionTy);
     opFunc.setPrivate();
   }
-  assert(opFunc.getType().template cast<FunctionType>().getResults() ==
-         op->getResultTypes());
-  assert(opFunc.getType().template cast<FunctionType>().getInputs() ==
-         op->getOperandTypes());
+  assert(SymbolTable::lookupSymbolIn(module, name)
+             ->template hasTrait<mlir::OpTrait::FunctionLike>());
 
-  rewriter.replaceOpWithNewOp<CallOp>(op, opFunc, op->getOperands());
+  rewriter.replaceOpWithNewOp<CallOp>(op, name, op.getType(),
+                                      op->getOperands());
 
   return success();
 }