[mlir][math] Added math::IPowI conversion to calls of outlined implementations.

[vzakhari]

Summary

Power functions are implemented as linkonce_odr scalar functions
for integer types used by IPowI operations met in a module.
Vector form of IPowI is linearized into a sequence of calls
of the scalar functions.

[vzakhari]

rebase

[vzakhari]

rebase

Please review.

[Mogball]

mlir/include/mlir/Conversion/Passes.td

495

Please change this back to a generic pass.

mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp

284 ↗

(On Diff #449430)

Lowering this to a function call is a little unusual. This lowering makes more sense to be a separate pass, because MathToLLVM is meant for trivial mostly one-to-one conversions

[vzakhari]

mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp

284 ↗

(On Diff #449430)

Okay, this makes sense! I will create a pass like MathToOutlinedImpl - please suggest if you can think of a better name.

[Mogball]

Is it possible to implement the function in arith or other math operations before lowing to LLVM? Probably MathExpandToFuncs

[vzakhari]

In D129810#3716668, @Mogball wrote:

Is it possible to implement the function in arith or other math operations before lowing to LLVM? Probably MathExpandToFuncs

It should be possible. Let me try it.

[Mogball]

mlir/include/mlir/Conversion/Passes.td

522

mlir/lib/Conversion/MathToFuncs/MathToFuncs.cpp

62–63

The type should only be omitted if it's explicit on the right hand side on an assignment due to a cast, or if it's "too long" (e.g. vector<int32_t>::iterator)

65–69

Can you add a failure reason with rewriter.notifyMatchFailure? It can notoriously difficult to debug pattern matching without those messages.

70

Spell out the auto.

71

Should this be unsigned?

78

Please don't use auto for a loop index.

89

129

Please spell out all the autos as appropriate.

139

Would generating this function be made any easier by using scf instead of cf?

327–329

Since this is a partial dialect conversion, can you not just add math::IPowI as an illegal op?

334–338

Nit: typically this comes before the pass declaration/definition

[vzakhari]

mlir/lib/Conversion/MathToFuncs/MathToFuncs.cpp

62–63

I will fix auto declarations. I will keep auto for results of rewriter.create<T> - it looks like it is a common practice to write it with auto.

65–69

Sure! Thanks for the suggestion!

71

int64_t should be correct, since ShapedType::getNumElements is defined as returning int64_t in IR/BuiltinAttributes.h.

139

It seems to me unstructured control flow fits better for the code with early returns from the function.

327–329

Will do.

[vzakhari]

Friendly ping. Please review.

[Mogball]

The high level idea makes sense to me, but I think the pass should be structured as:

1. Collect all element types for which functions need to be instantiated and instantiate those functions
2. Run the vector to scalar conversions and scalar to function call rewrites

mlir/include/mlir/Conversion/MathToFuncs/MathToFuncs.h

13

You can remove this include

24

These functions normally don't take a PatternBenefit.

27

These functions normally return just a Pass

mlir/include/mlir/Conversion/Passes.td

523

Why is this type of linkage needed? It's unfortunate that this pass depends on the LLVM dialect just for this.

mlir/lib/Conversion/MathToFuncs/MathToFuncs.cpp

125

It's generally not valid for patterns to access and modify parent operations, especially since these "patterns" are exposed publicly. The generation of these software implementations should not be done inside patterns. Patterns should be local rewrites of operations.

[vzakhari]

In D129810#3736207, @Mogball wrote:

The high level idea makes sense to me, but I think the pass should be structured as:

1. Collect all element types for which functions need to be instantiated and instantiate those functions
2. Run the vector to scalar conversions and scalar to function call rewrites

Can you please explain the benefit of such a split? There are 4 existing conversion passes that access their getNearestSymbolTable() "on the fly" (ComplexToLibm, MathToLibm, MemRefToSPIRV and SCFToOpenMP), and I do not really want to invent the wheel here :)

Thank you for the review! I will upload updated files shortly.

mlir/include/mlir/Conversion/MathToFuncs/MathToFuncs.h

24

There are some (e.g. ComplexToLLVM, MathToLibm), but it is currently not needed so I will remove it. It may be used in future in case there is target/project specific handling for such operations, e.g. in some cases there is target/project specific conversion for some flavors of the operation that is more preferable than MathToFuncs conversion - as I understand, in this case the patterns may be selected based on their benefit.

27

I see that many-many passes that are defined as ModuleOp passes in mlir/lib/Conversion/Passes.td (i.e. they are derived from Pass<"...", "ModuleOp">) all return std::unique_ptr<OperationPass<ModuleOp>> in the CPP files in mlir/lib/Conversion/*/*.cpp. I am not comfortable with stepping away from the common pattern. I guess it should be written this way to make it clear that the pass is not a generic pass but rather a specific operation pass.

mlir/include/mlir/Conversion/Passes.td

523

In FuncToLLVM conversion all func.func operations have external linkage by default. We cannot use external linkage for the generated functions, because there may be multiple definitions in different translation modules and the linking will fail. So we have to specify either internal or linkonce_odr linkage - the latter is better, because it allows the linker to keep only one copy of the generated function across multiple modules being linked.

I do not know of any other way to control linkage in MLIR.

mlir/lib/Conversion/MathToFuncs/MathToFuncs.cpp

125

I understand the concern. I followed the same approach that MathToLibm and ComplexToLibm use.

I guess I can add a comment in MathToFuncs.h to warn about valid usage of populateMathToFuncsConversionPatterns - does it make sense?

[vzakhari]

Addressed review remarks.

[Mogball]

The benefit is that the patterns won't 1. Keep looking up functions in the symbol table, which is painfully slow and 2. Won't violate the API contract for patterns.

mlir/include/mlir/Conversion/MathToFuncs/MathToFuncs.h

24

In that case, they wouldn't use these patterns or this pass. The validity of composing patterns is inconsistent.

27

It is inconsistent throughout the MLIR codebase, but please return a Pass. There is no benefit to returning OperationPass<ModuleOp> because the pass manager API only accepts the former anyways.

mlir/include/mlir/Conversion/Passes.td

523

Adding linkage to FuncDialect appears to be a TODO. It's fine for now then.

mlir/lib/Conversion/MathToFuncs/MathToFuncs.cpp

125

No. Those other patterns are also invalid. Patterns should not be modifying parent operations. It's something that "just happens to work" because no one has needed to compose these patterns sets in a parallelized pass (which would result in race conditions). This is especially true in conversion patterns, because there is extra bookkeeping in the dialect conversion pass that is violated if this happens.

Please don't propagate what has been done in the older passes and separate these.

[vzakhari]

mlir/include/mlir/Conversion/MathToFuncs/MathToFuncs.h

27

Will do.

mlir/lib/Conversion/MathToFuncs/MathToFuncs.cpp

125

Thank you for the explanation!

Then it looks like exposing the patterns is unsafe, in general, and I should rewrite the pass as you suggested without externalizing the match patterns (populateMathToFuncsConversionPatterns). The module pass will first visit all IPowI operations for collecting information about the kinds of implementation functions that need to be created in the module, then it will create the new functions, and then it will apply populateMathToFuncsConversionPatterns patterns to rewrite the operations into the calls.

[vzakhari]

Addressed review comments.

[Mogball]

Good evolution! The direction of the patch LGTM. I just have some minor complaints about the implementation details.

mlir/include/mlir/Conversion/MathToFuncs/MathToFuncs.h

15–17

mlir/include/mlir/Conversion/Passes.td

519

522–523

mlir/lib/Conversion/MathToFuncs/MathToFuncs.cpp

10–26

Can you trim these includes? I don't think all of these are needed

31–38

This isn't necessary. FunctionType can be hashed directly in a DenseMap.

154

I would prefer this to be a standalone function, not a static function on the pattern.

341–344

These two checks aren't necessary since they are invariants enforced by the verifier. You just need to check that this is a scalar integer op (done by the first condition).

377

384–385

386–388

If all the types are the same, then you only need to hash based on the element type.

[Mogball]

Do you also mind marking resolved comments as "done"?

[vzakhari]

Thank you for the prompt review!

mlir/lib/Conversion/MathToFuncs/MathToFuncs.cpp

384–385

Since I am going to add FPowI support later, does it make sense to keep the switch now?

386–388

This is true for IPowI but not for FPowI. I prefer to keep it as a function type so that the keys are consistent in the map, otherwise, for IPowI we will have IntegerType keys and for FPowI we will have FunctionType keys.

[vzakhari]

In D129810#3747847, @Mogball wrote:

Do you also mind marking resolved comments as "done"?

Done :)

[Mogball]

mlir/lib/Conversion/MathToFuncs/MathToFuncs.cpp

384–385

Yes it does. Thanks.

386–388

Then you can just use Type as the key.

386–388

In any case, you don't need std::map and a custom comparator. DenseMap<Type, T> works just fine.

[vzakhari]

mlir/lib/Conversion/MathToFuncs/MathToFuncs.cpp

386–388

Should be done now. Thanks!

[Mogball]

LGTM. Just one last complaint about the stringify/hashing

mlir/lib/Conversion/MathToFuncs/MathToFuncs.cpp

40

We generally don't use std::function

62

66

I think the "function type stringify" is too generic. For FPowI, just the floating point and integer type will suffice. This doesn't need to be done generally. I think you can just remove this function.

152–153

362–376

Can you remove the extra logic for now? I know you want to add support for other ops, but I think it'd be less automagical to just add another Case for FPowI

[vzakhari]

mlir/lib/Conversion/MathToFuncs/MathToFuncs.cpp

66

Okay, I will remove it.

152–153

Will do. flush should not be necessary according to this: https://llvm.org/doxygen/classllvm_1_1raw__string__ostream.html#details

[vzakhari]

Please let me know if everything looks good now. Thank you for all the feedback!

[Mogball]

LGTM.

mlir/include/mlir/Conversion/MathToFuncs/MathToFuncs.h

+//===- MathToFuncs.h - Math to outlined impl conversion ---------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_CONVERSION_MATHTOFUNCS_MATHTOFUNCS_H
+#define MLIR_CONVERSION_MATHTOFUNCS_MATHTOFUNCS_H
+
+#include "mlir/IR/PatternMatch.h"
+#include <memory>
+
+namespace mlir {
+
+class ModuleOp;
+template <typename T>
+class OperationPass;
+class PatternBenefit;
+class RewritePatternSet;
+
+void populateMathToFuncsConversionPatterns(RewritePatternSet &patterns,
+                                           PatternBenefit benefit);
+
+std::unique_ptr<OperationPass<ModuleOp>> createConvertMathToFuncsPass();
+} // namespace mlir
+
+#endif // MLIR_CONVERSION_MATHTOFUNCS_MATHTOFUNCS_H

mlir/include/mlir/Conversion/Passes.h

 #include "mlir/Conversion/GPUCommon/GPUCommonPass.h"
 #include "mlir/Conversion/GPUToNVVM/GPUToNVVMPass.h"
 #include "mlir/Conversion/GPUToROCDL/GPUToROCDLPass.h"
 #include "mlir/Conversion/GPUToSPIRV/GPUToSPIRVPass.h"
 #include "mlir/Conversion/GPUToVulkan/ConvertGPUToVulkanPass.h"
 #include "mlir/Conversion/LinalgToLLVM/LinalgToLLVM.h"
 #include "mlir/Conversion/LinalgToSPIRV/LinalgToSPIRVPass.h"
 #include "mlir/Conversion/LinalgToStandard/LinalgToStandard.h"
+#include "mlir/Conversion/MathToFuncs/MathToFuncs.h"
 #include "mlir/Conversion/MathToLLVM/MathToLLVM.h"
 #include "mlir/Conversion/MathToLibm/MathToLibm.h"
 #include "mlir/Conversion/MathToSPIRV/MathToSPIRVPass.h"
 #include "mlir/Conversion/MemRefToLLVM/MemRefToLLVM.h"
 #include "mlir/Conversion/MemRefToSPIRV/MemRefToSPIRVPass.h"
 #include "mlir/Conversion/NVGPUToNVVM/NVGPUToNVVM.h"
 #include "mlir/Conversion/OpenACCToLLVM/ConvertOpenACCToLLVM.h"
 #include "mlir/Conversion/OpenACCToSCF/ConvertOpenACCToSCF.h"

mlir/include/mlir/Conversion/Passes.td

let dependentDialects = [
     "vector::VectorDialect",
   ];
 }
 
 //===----------------------------------------------------------------------===//
 // MathToLLVM
 //===----------------------------------------------------------------------===//
 
 def ConvertMathToLLVM : Pass<"convert-math-to-llvm"> {
   let summary = "Convert Math dialect to LLVM dialect";
   let description = [{
     This pass converts supported Math ops to LLVM dialect intrinsics.
   }];
   let constructor = "mlir::createConvertMathToLLVMPass()";
   let dependentDialects = ["LLVM::LLVMDialect"];
 }
 
 //===----------------------------------------------------------------------===//
 // MathToSPIRV
 //===----------------------------------------------------------------------===//
 
 def ConvertMathToSPIRV : Pass<"convert-math-to-spirv"> {
   let summary = "Convert Math dialect to SPIR-V dialect";
   let constructor = "mlir::createConvertMathToSPIRVPass()";
   let dependentDialects = ["spirv::SPIRVDialect"];
 }
 
 //===----------------------------------------------------------------------===//
+// MathToFuncs
+//===----------------------------------------------------------------------===//
+
+def ConvertMathToFuncs : Pass<"convert-math-to-funcs", "ModuleOp"> {
+  let summary = "Convert Math operations to calls of outlined implementations";
+  let description = [{
+    This pass converts supported Math ops to calls of compiler generated
+    functions implementing these operations.
+    LLVMDialect is used for LinkonceODR linkage of the generated functions.
+  }];
+  let constructor = "mlir::createConvertMathToFuncsPass()";
+  let dependentDialects = [
+    "arith::ArithmeticDialect",
+    "cf::ControlFlowDialect",
+    "func::FuncDialect",
+    "vector::VectorDialect",
+    "LLVM::LLVMDialect",
+  ];
+}
+
+//===----------------------------------------------------------------------===//
 // MemRefToLLVM
 //===----------------------------------------------------------------------===//
 
 def ConvertMemRefToLLVM : Pass<"convert-memref-to-llvm", "ModuleOp"> {
   let summary = "Convert operations from the MemRef dialect to the LLVM "
                 "dialect";
   let constructor = "mlir::createMemRefToLLVMPass()";
   let dependentDialects = ["LLVM::LLVMDialect"];

mlir/lib/Conversion/CMakeLists.txt

 add_subdirectory(GPUToNVVM)
 add_subdirectory(GPUToROCDL)
 add_subdirectory(GPUToSPIRV)
 add_subdirectory(GPUToVulkan)
 add_subdirectory(LinalgToLLVM)
 add_subdirectory(LinalgToSPIRV)
 add_subdirectory(LinalgToStandard)
 add_subdirectory(LLVMCommon)
+add_subdirectory(MathToFuncs)
 add_subdirectory(MathToLibm)
 add_subdirectory(MathToLLVM)
 add_subdirectory(MathToSPIRV)
 add_subdirectory(MemRefToLLVM)
 add_subdirectory(MemRefToSPIRV)
 add_subdirectory(NVGPUToNVVM)
 add_subdirectory(OpenACCToLLVM)
 add_subdirectory(OpenACCToSCF)

mlir/lib/Conversion/MathToFuncs/CMakeLists.txt

+add_mlir_conversion_library(MLIRMathToFuncs
+  MathToFuncs.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Conversion/MathToFuncs
+
+  DEPENDS
+  MLIRConversionPassIncGen
+
+  LINK_COMPONENTS
+  Core
+
+  LINK_LIBS PUBLIC
+  MLIRArithmeticDialect
+  MLIRControlFlowDialect
+  MLIRFuncDialect
+  MLIRLLVMDialect
+  MLIRMathDialect
+  MLIRPass
+  MLIRTransforms
+  MLIRVectorDialect
+  MLIRVectorUtils
+  )

mlir/lib/Conversion/MathToFuncs/MathToFuncs.cpp

+//===- MathToFuncs.cpp - Math to outlined implementation conversion -------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Conversion/MathToFuncs/MathToFuncs.h"
+#include "../PassDetail.h"
+#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Dialect/Math/IR/Math.h"
+#include "mlir/Dialect/Utils/IndexingUtils.h"
+#include "mlir/Dialect/Vector/IR/VectorOps.h"
+#include "mlir/Dialect/Vector/Utils/VectorUtils.h"
+#include "mlir/IR/BuiltinDialect.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+using namespace mlir;
+
+namespace {
+struct ConvertMathToFuncsPass
+    : public ConvertMathToFuncsBase<ConvertMathToFuncsPass> {
+  ConvertMathToFuncsPass() = default;
+
+  void runOnOperation() override;
+};
+
+// Pattern to convert vector operations to scalar operations.
+template <typename Op>
+struct VecOpToScalarOp : public OpRewritePattern<Op> {
+public:
+  using OpRewritePattern<Op>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(Op op, PatternRewriter &rewriter) const final;
+};
+
+struct IPowIOpLowering : public OpRewritePattern<math::IPowIOp> {
+private:
+  func::FuncOp getElementFunc(math::IPowIOp op, IntegerType elementType,
+                              PatternRewriter &rewriter) const;
+
+public:
+  IPowIOpLowering(MLIRContext *context, PatternBenefit benefit)
+      : OpRewritePattern<math::IPowIOp>(context, benefit) {}
+
+  /// Convert IPowI into a call to a local function implementing
+  /// the power operation. The local function computes a scalar result,
+  /// so vector forms of IPowI are linearized.
+  LogicalResult matchAndRewrite(math::IPowIOp op,
+                                PatternRewriter &rewriter) const final;
+};
+
+} // namespace
+
+template <typename Op>
+LogicalResult
+VecOpToScalarOp<Op>::matchAndRewrite(Op op, PatternRewriter &rewriter) const {
+  auto opType = op.getType();
+  auto loc = op.getLoc();
+  auto vecType = opType.template dyn_cast<VectorType>();
+
+  if (!vecType)
+    return failure();
+  if (!vecType.hasRank())
+    return failure();
+  auto shape = vecType.getShape();
+  int64_t numElements = vecType.getNumElements();
+
+  Value result = rewriter.create<arith::ConstantOp>(
+      loc, DenseElementsAttr::get(
+               vecType, IntegerAttr::get(vecType.getElementType(), 0)));
+  SmallVector<int64_t> ones(shape.size(), 1);
+  SmallVector<int64_t> strides = computeStrides(shape, ones);
+  for (auto linearIndex = 0; linearIndex < numElements; ++linearIndex) {
+    SmallVector<int64_t> positions = delinearize(strides, linearIndex);
+    SmallVector<Value> operands;
+    for (auto input : op->getOperands())
+      operands.push_back(
+          rewriter.create<vector::ExtractOp>(loc, input, positions));
+    Value scalarOp =
+        rewriter.create<Op>(loc, vecType.getElementType(), operands);
+    result =
+        rewriter.create<vector::InsertOp>(loc, scalarOp, result, positions);
+  }
+  rewriter.replaceOp(op, {result});
+  return success();
+}
+
+/// Create linkonce_odr function to implement the power function with
+/// the given \p elementType scalar type.
+///
+/// template <typename T>
+/// T __mlir_math_ipowi_*(T b, T p) {
+///   if (p == T(0))
+///     return T(1);
+///   if (p < T(0)) {
+///     if (b == T(0))
+///       return T(1) / T(0); // trigger div-by-zero
+///     if (b == T(1))
+///       return T(1);
+///     if (b == T(-1)) {
+///       if (p & T(1))
+///         return T(-1);
+///       return T(1);
+///     }
+///     return T(0);
+///   }
+///   T result = T(1);
+///   while (true) {
+///     if (p & T(1))
+///       result *= b;
+///     p >>= T(1);
+///     if (p == T(0))
+///       return result;
+///     b *= b;
+///   }
+/// }
+func::FuncOp IPowIOpLowering::getElementFunc(math::IPowIOp op,
+                                             IntegerType elementType,
+                                             PatternRewriter &rewriter) const {
+  std::string funcName("__mlir_math_ipowi_");
+  llvm::raw_string_ostream nameOS(funcName);
+  elementType.print(nameOS);
+
+  auto module = SymbolTable::getNearestSymbolTable(op);
+  auto funcType = FunctionType::get(rewriter.getContext(),
+                                    {elementType, elementType}, {elementType});
+  if (auto funcOp = dyn_cast_or_null<func::FuncOp>(
+          SymbolTable::lookupSymbolIn(module, funcName))) {
+    assert(funcOp.getFunctionTypeAttr().getValue() == funcType &&
+           "ipowi function type mismatch");
+    return funcOp;
+  }
+
+  OpBuilder::InsertionGuard guard(rewriter);
+  auto loc = rewriter.getUnknownLoc();
+  rewriter.setInsertionPointToEnd(&module->getRegion(0).front());
+  auto funcOp = rewriter.create<func::FuncOp>(loc, funcName, funcType);
+  auto inlineLinkage = LLVM::linkage::Linkage::LinkonceODR;
+  auto linkage = LLVM::LinkageAttr::get(rewriter.getContext(), inlineLinkage);
+  funcOp->setAttr("llvm.linkage", linkage);
+  funcOp.setPrivate();
+
+  auto *entryBlock = funcOp.addEntryBlock();
+  auto *funcBody = entryBlock->getParent();
+
+  auto bArg = funcOp.getArgument(0);
+  auto pArg = funcOp.getArgument(1);
+  rewriter.setInsertionPointToEnd(entryBlock);
+  Value zeroValue = rewriter.create<arith::ConstantOp>(
+      loc, elementType, rewriter.getIntegerAttr(elementType, 0));
+  Value oneValue = rewriter.create<arith::ConstantOp>(
+      loc, elementType, rewriter.getIntegerAttr(elementType, 1));
+  Value minusOneValue = rewriter.create<arith::ConstantOp>(
+      loc, elementType,
+      rewriter.getIntegerAttr(elementType,
+                              APInt(elementType.getIntOrFloatBitWidth(), -1ULL,
+                                    /*isSigned=*/true)));
+
+  // if (p == T(0))
+  //   return T(1);
+  auto pIsZero = rewriter.create<arith::CmpIOp>(loc, arith::CmpIPredicate::eq,
+                                                pArg, zeroValue);
+  auto *thenBlock = rewriter.createBlock(funcBody);
+  rewriter.create<func::ReturnOp>(loc, oneValue);
+  auto *fallthroughBlock = rewriter.createBlock(funcBody);
+  // Set up conditional branch for (p == T(0)).
+  rewriter.setInsertionPointToEnd(pIsZero->getBlock());
+  rewriter.create<cf::CondBranchOp>(loc, pIsZero, thenBlock, fallthroughBlock);
+
+  // if (p < T(0)) {
+  rewriter.setInsertionPointToEnd(fallthroughBlock);
+  auto pIsNeg = rewriter.create<arith::CmpIOp>(loc, arith::CmpIPredicate::sle,
+                                               pArg, zeroValue);
+  //   if (b == T(0))
+  rewriter.createBlock(funcBody);
+  auto bIsZero = rewriter.create<arith::CmpIOp>(loc, arith::CmpIPredicate::eq,
+                                                bArg, zeroValue);
+  //     return T(1) / T(0);
+  thenBlock = rewriter.createBlock(funcBody);
+  rewriter.create<func::ReturnOp>(
+      loc,
+      rewriter.create<arith::DivSIOp>(loc, oneValue, zeroValue).getResult());
+  fallthroughBlock = rewriter.createBlock(funcBody);
+  // Set up conditional branch for (b == T(0)).
+  rewriter.setInsertionPointToEnd(bIsZero->getBlock());
+  rewriter.create<cf::CondBranchOp>(loc, bIsZero, thenBlock, fallthroughBlock);
+
+  //   if (b == T(1))
+  rewriter.setInsertionPointToEnd(fallthroughBlock);
+  auto bIsOne = rewriter.create<arith::CmpIOp>(loc, arith::CmpIPredicate::eq,
+                                               bArg, oneValue);
+  //    return T(1);
+  thenBlock = rewriter.createBlock(funcBody);
+  rewriter.create<func::ReturnOp>(loc, oneValue);
+  fallthroughBlock = rewriter.createBlock(funcBody);
+  // Set up conditional branch for (b == T(1)).
+  rewriter.setInsertionPointToEnd(bIsOne->getBlock());
+  rewriter.create<cf::CondBranchOp>(loc, bIsOne, thenBlock, fallthroughBlock);
+
+  //   if (b == T(-1)) {
+  rewriter.setInsertionPointToEnd(fallthroughBlock);
+  auto bIsMinusOne = rewriter.create<arith::CmpIOp>(
+      loc, arith::CmpIPredicate::eq, bArg, minusOneValue);
+  //     if (p & T(1))
+  rewriter.createBlock(funcBody);
+  auto pIsOdd = rewriter.create<arith::CmpIOp>(
+      loc, arith::CmpIPredicate::ne,
+      rewriter.create<arith::AndIOp>(loc, pArg, oneValue), zeroValue);
+  //       return T(-1);
+  thenBlock = rewriter.createBlock(funcBody);
+  rewriter.create<func::ReturnOp>(loc, minusOneValue);
+  fallthroughBlock = rewriter.createBlock(funcBody);
+  // Set up conditional branch for (p & T(1)).
+  rewriter.setInsertionPointToEnd(pIsOdd->getBlock());
+  rewriter.create<cf::CondBranchOp>(loc, pIsOdd, thenBlock, fallthroughBlock);
+
+  //     return T(1);
+  //   } // b == T(-1)
+  rewriter.setInsertionPointToEnd(fallthroughBlock);
+  rewriter.create<func::ReturnOp>(loc, oneValue);
+  fallthroughBlock = rewriter.createBlock(funcBody);
+  // Set up conditional branch for (b == T(-1)).
+  rewriter.setInsertionPointToEnd(bIsMinusOne->getBlock());
+  rewriter.create<cf::CondBranchOp>(loc, bIsMinusOne, pIsOdd->getBlock(),
+                                    fallthroughBlock);
+
+  //   return T(0);
+  // } // (p < T(0))
+  rewriter.setInsertionPointToEnd(fallthroughBlock);
+  rewriter.create<func::ReturnOp>(loc, zeroValue);
+  auto *loopHeader = rewriter.createBlock(
+      funcBody, funcBody->end(), {elementType, elementType, elementType},
+      {loc, loc, loc});
+  // Set up conditional branch for (p < T(0)).
+  rewriter.setInsertionPointToEnd(pIsNeg->getBlock());
+  // Set initial values of 'result', 'b' and 'p' for the loop.
+  rewriter.create<cf::CondBranchOp>(loc, pIsNeg, bIsZero->getBlock(),
+                                    loopHeader,
+                                    ValueRange{oneValue, bArg, pArg});
+
+  // T result = T(1);
+  // while (true) {
+  //   if (p & T(1))
+  //     result *= b;
+  //   p >>= T(1);
+  //   if (p == T(0))
+  //     return result;
+  //   b *= b;
+  // }
+  Value resultTmp = loopHeader->getArgument(0);
+  Value baseTmp = loopHeader->getArgument(1);
+  Value powerTmp = loopHeader->getArgument(2);
+  rewriter.setInsertionPointToEnd(loopHeader);
+
+  //   if (p & T(1))
+  auto powerTmpIsOdd = rewriter.create<arith::CmpIOp>(
+      loc, arith::CmpIPredicate::ne,
+      rewriter.create<arith::AndIOp>(loc, powerTmp, oneValue), zeroValue);
+  thenBlock = rewriter.createBlock(funcBody);
+  //     result *= b;
+  Value newResultTmp = rewriter.create<arith::MulIOp>(loc, resultTmp, baseTmp);
+  fallthroughBlock =
+      rewriter.createBlock(funcBody, funcBody->end(), elementType, loc);
+  rewriter.setInsertionPointToEnd(thenBlock);
+  rewriter.create<cf::BranchOp>(loc, newResultTmp, fallthroughBlock);
+  // Set up conditional branch for (p & T(1)).
+  rewriter.setInsertionPointToEnd(powerTmpIsOdd->getBlock());
+  rewriter.create<cf::CondBranchOp>(loc, powerTmpIsOdd, thenBlock,
+                                    fallthroughBlock, resultTmp);
+  // Merged 'result'.
+  newResultTmp = fallthroughBlock->getArgument(0);
+
+  //   p >>= T(1);
+  rewriter.setInsertionPointToEnd(fallthroughBlock);
+  Value newPowerTmp = rewriter.create<arith::ShRUIOp>(loc, powerTmp, oneValue);
+
+  //   if (p == T(0))
+  auto newPowerIsZero = rewriter.create<arith::CmpIOp>(
+      loc, arith::CmpIPredicate::eq, newPowerTmp, zeroValue);
+  //     return result;
+  thenBlock = rewriter.createBlock(funcBody);
+  rewriter.create<func::ReturnOp>(loc, newResultTmp);
+  fallthroughBlock = rewriter.createBlock(funcBody);
+  // Set up conditional branch for (p == T(0)).
+  rewriter.setInsertionPointToEnd(newPowerIsZero->getBlock());
+  rewriter.create<cf::CondBranchOp>(loc, newPowerIsZero, thenBlock,
+                                    fallthroughBlock);
+
+  //   b *= b;
+  // }
+  rewriter.setInsertionPointToEnd(fallthroughBlock);
+  Value newBaseTmp = rewriter.create<arith::MulIOp>(loc, baseTmp, baseTmp);
+  // Pass new values for 'result', 'b' and 'p' to the loop header.
+  rewriter.create<cf::BranchOp>(
+      loc, ValueRange{newResultTmp, newBaseTmp, newPowerTmp}, loopHeader);
+  return funcOp;
+}
+
+/// Convert IPowI into a call to a local function implementing
+/// the power operation. The local function computes a scalar result,
+/// so vector forms of IPowI are linearized.
+LogicalResult
+IPowIOpLowering::matchAndRewrite(math::IPowIOp op,
+                                 PatternRewriter &rewriter) const {
+  auto baseType = op.getOperands()[0].getType().dyn_cast<IntegerType>();
+  auto exponentType = op.getOperands()[1].getType().dyn_cast<IntegerType>();
+
+  if (!baseType || baseType != exponentType)
+    return failure();
+
+  auto elementFunc = getElementFunc(op, baseType, rewriter);
+  rewriter.replaceOpWithNewOp<func::CallOp>(op, elementFunc, op.getOperands());
+  return success();
+}
+
+void ConvertMathToFuncsPass::runOnOperation() {
+  auto module = getOperation();
+
+  RewritePatternSet patterns(&getContext());
+  populateMathToFuncsConversionPatterns(patterns, /*benefit=*/1);
+  ConversionTarget target(getContext());
+  target.addLegalDialect<arith::ArithmeticDialect, BuiltinDialect,
+                         cf::ControlFlowDialect, func::FuncDialect,
+                         vector::VectorDialect>();
+  if (failed(applyPartialConversion(module, target, std::move(patterns))))
+    signalPassFailure();
+}
+
+void mlir::populateMathToFuncsConversionPatterns(RewritePatternSet &patterns,
+                                                 PatternBenefit benefit) {
+  patterns.add<VecOpToScalarOp<math::IPowIOp>>(patterns.getContext(), benefit);
+  patterns.add<IPowIOpLowering>(patterns.getContext(), benefit);
+}
+
+std::unique_ptr<OperationPass<ModuleOp>> mlir::createConvertMathToFuncsPass() {
+  return std::make_unique<ConvertMathToFuncsPass>();
+}

mlir/test/Conversion/MathToFuncs/math-to-funcs.mlir

+// RUN: mlir-opt %s -split-input-file -pass-pipeline="convert-math-to-funcs" | FileCheck %s
+
+// -----
+
+// CHECK-LABEL: func @ipowi(
+// CHECK-SAME: %[[ARG0:.+]]: i64,
+// CHECK-SAME: %[[ARG1:.+]]: i64)
+func.func @ipowi(%arg0: i64, %arg1: i64) {
+  // CHECK: call @__mlir_math_ipowi_i64(%[[ARG0]], %[[ARG1]]) : (i64, i64) -> i64
+  %0 = math.ipowi %arg0, %arg1 : i64
+  func.return
+}
+
+// CHECK-LABEL:   func.func private @__mlir_math_ipowi_i64(
+// CHECK-SAME:      %[[VAL_0:.*]]: i64,
+// CHECK-SAME:      %[[VAL_1:.*]]: i64) -> i64
+// CHECK-SAME:        attributes {llvm.linkage = #llvm.linkage<linkonce_odr>} {
+// CHECK:           %[[VAL_2:.*]] = arith.constant 0 : i64
+// CHECK:           %[[VAL_3:.*]] = arith.constant 1 : i64
+// CHECK:           %[[VAL_4:.*]] = arith.constant -1 : i64
+// CHECK:           %[[VAL_5:.*]] = arith.cmpi eq, %[[VAL_1]], %[[VAL_2]] : i64
+// CHECK:           cf.cond_br %[[VAL_5]], ^bb1, ^bb2
+// CHECK:         ^bb1:
+// CHECK:           return %[[VAL_3]] : i64
+// CHECK:         ^bb2:
+// CHECK:           %[[VAL_6:.*]] = arith.cmpi sle, %[[VAL_1]], %[[VAL_2]] : i64
+// CHECK:           cf.cond_br %[[VAL_6]], ^bb3, ^bb12(%[[VAL_3]], %[[VAL_0]], %[[VAL_1]] : i64, i64, i64)
+// CHECK:         ^bb3:
+// CHECK:           %[[VAL_7:.*]] = arith.cmpi eq, %[[VAL_0]], %[[VAL_2]] : i64
+// CHECK:           cf.cond_br %[[VAL_7]], ^bb4, ^bb5
+// CHECK:         ^bb4:
+// CHECK:           %[[VAL_8:.*]] = arith.divsi %[[VAL_3]], %[[VAL_2]]  : i64
+// CHECK:           return %[[VAL_8]] : i64
+// CHECK:         ^bb5:
+// CHECK:           %[[VAL_9:.*]] = arith.cmpi eq, %[[VAL_0]], %[[VAL_3]] : i64
+// CHECK:           cf.cond_br %[[VAL_9]], ^bb6, ^bb7
+// CHECK:         ^bb6:
+// CHECK:           return %[[VAL_3]] : i64
+// CHECK:         ^bb7:
+// CHECK:           %[[VAL_10:.*]] = arith.cmpi eq, %[[VAL_0]], %[[VAL_4]] : i64
+// CHECK:           cf.cond_br %[[VAL_10]], ^bb8, ^bb11
+// CHECK:         ^bb8:
+// CHECK:           %[[VAL_11:.*]] = arith.andi %[[VAL_1]], %[[VAL_3]]  : i64
+// CHECK:           %[[VAL_12:.*]] = arith.cmpi ne, %[[VAL_11]], %[[VAL_2]] : i64
+// CHECK:           cf.cond_br %[[VAL_12]], ^bb9, ^bb10
+// CHECK:         ^bb9:
+// CHECK:           return %[[VAL_4]] : i64
+// CHECK:         ^bb10:
+// CHECK:           return %[[VAL_3]] : i64
+// CHECK:         ^bb11:
+// CHECK:           return %[[VAL_2]] : i64
+// CHECK:         ^bb12(%[[VAL_13:.*]]: i64, %[[VAL_14:.*]]: i64, %[[VAL_15:.*]]: i64):
+// CHECK:           %[[VAL_16:.*]] = arith.andi %[[VAL_15]], %[[VAL_3]]  : i64
+// CHECK:           %[[VAL_17:.*]] = arith.cmpi ne, %[[VAL_16]], %[[VAL_2]] : i64
+// CHECK:           cf.cond_br %[[VAL_17]], ^bb13, ^bb14(%[[VAL_13]] : i64)
+// CHECK:         ^bb13:
+// CHECK:           %[[VAL_18:.*]] = arith.muli %[[VAL_13]], %[[VAL_14]]  : i64
+// CHECK:           cf.br ^bb14(%[[VAL_18]] : i64)
+// CHECK:         ^bb14(%[[VAL_19:.*]]: i64):
+// CHECK:           %[[VAL_20:.*]] = arith.shrui %[[VAL_15]], %[[VAL_3]]  : i64
+// CHECK:           %[[VAL_21:.*]] = arith.cmpi eq, %[[VAL_20]], %[[VAL_2]] : i64
+// CHECK:           cf.cond_br %[[VAL_21]], ^bb15, ^bb16
+// CHECK:         ^bb15:
+// CHECK:           return %[[VAL_19]] : i64
+// CHECK:         ^bb16:
+// CHECK:           %[[VAL_22:.*]] = arith.muli %[[VAL_14]], %[[VAL_14]]  : i64
+// CHECK:           cf.br ^bb12(%[[VAL_19]], %[[VAL_22]], %[[VAL_20]] : i64, i64, i64)
+// CHECK:         }
+
+// -----
+
+// CHECK-LABEL: func @ipowi(
+// CHECK-SAME: %[[ARG0:.+]]: i8,
+// CHECK-SAME: %[[ARG1:.+]]: i8)
+  // CHECK: call @__mlir_math_ipowi_i8(%[[ARG0]], %[[ARG1]]) : (i8, i8) -> i8
+func.func @ipowi(%arg0: i8, %arg1: i8) {
+  %0 = math.ipowi %arg0, %arg1 : i8
+  func.return
+}
+
+// CHECK-LABEL:   func.func private @__mlir_math_ipowi_i8(
+// CHECK-SAME:      %[[VAL_0:.*]]: i8,
+// CHECK-SAME:      %[[VAL_1:.*]]: i8) -> i8
+// CHECK-SAME:        attributes {llvm.linkage = #llvm.linkage<linkonce_odr>} {
+// CHECK:           %[[VAL_2:.*]] = arith.constant 0 : i8
+// CHECK:           %[[VAL_3:.*]] = arith.constant 1 : i8
+// CHECK:           %[[VAL_4:.*]] = arith.constant -1 : i8
+// CHECK:           %[[VAL_5:.*]] = arith.cmpi eq, %[[VAL_1]], %[[VAL_2]] : i8
+// CHECK:           cf.cond_br %[[VAL_5]], ^bb1, ^bb2
+// CHECK:         ^bb1:
+// CHECK:           return %[[VAL_3]] : i8
+// CHECK:         ^bb2:
+// CHECK:           %[[VAL_6:.*]] = arith.cmpi sle, %[[VAL_1]], %[[VAL_2]] : i8
+// CHECK:           cf.cond_br %[[VAL_6]], ^bb3, ^bb12(%[[VAL_3]], %[[VAL_0]], %[[VAL_1]] : i8, i8, i8)
+// CHECK:         ^bb3:
+// CHECK:           %[[VAL_7:.*]] = arith.cmpi eq, %[[VAL_0]], %[[VAL_2]] : i8
+// CHECK:           cf.cond_br %[[VAL_7]], ^bb4, ^bb5
+// CHECK:         ^bb4:
+// CHECK:           %[[VAL_8:.*]] = arith.divsi %[[VAL_3]], %[[VAL_2]]  : i8
+// CHECK:           return %[[VAL_8]] : i8
+// CHECK:         ^bb5:
+// CHECK:           %[[VAL_9:.*]] = arith.cmpi eq, %[[VAL_0]], %[[VAL_3]] : i8
+// CHECK:           cf.cond_br %[[VAL_9]], ^bb6, ^bb7
+// CHECK:         ^bb6:
+// CHECK:           return %[[VAL_3]] : i8
+// CHECK:         ^bb7:
+// CHECK:           %[[VAL_10:.*]] = arith.cmpi eq, %[[VAL_0]], %[[VAL_4]] : i8
+// CHECK:           cf.cond_br %[[VAL_10]], ^bb8, ^bb11
+// CHECK:         ^bb8:
+// CHECK:           %[[VAL_11:.*]] = arith.andi %[[VAL_1]], %[[VAL_3]]  : i8
+// CHECK:           %[[VAL_12:.*]] = arith.cmpi ne, %[[VAL_11]], %[[VAL_2]] : i8
+// CHECK:           cf.cond_br %[[VAL_12]], ^bb9, ^bb10
+// CHECK:         ^bb9:
+// CHECK:           return %[[VAL_4]] : i8
+// CHECK:         ^bb10:
+// CHECK:           return %[[VAL_3]] : i8
+// CHECK:         ^bb11:
+// CHECK:           return %[[VAL_2]] : i8
+// CHECK:         ^bb12(%[[VAL_13:.*]]: i8, %[[VAL_14:.*]]: i8, %[[VAL_15:.*]]: i8):
+// CHECK:           %[[VAL_16:.*]] = arith.andi %[[VAL_15]], %[[VAL_3]]  : i8
+// CHECK:           %[[VAL_17:.*]] = arith.cmpi ne, %[[VAL_16]], %[[VAL_2]] : i8
+// CHECK:           cf.cond_br %[[VAL_17]], ^bb13, ^bb14(%[[VAL_13]] : i8)
+// CHECK:         ^bb13:
+// CHECK:           %[[VAL_18:.*]] = arith.muli %[[VAL_13]], %[[VAL_14]]  : i8
+// CHECK:           cf.br ^bb14(%[[VAL_18]] : i8)
+// CHECK:         ^bb14(%[[VAL_19:.*]]: i8):
+// CHECK:           %[[VAL_20:.*]] = arith.shrui %[[VAL_15]], %[[VAL_3]]  : i8
+// CHECK:           %[[VAL_21:.*]] = arith.cmpi eq, %[[VAL_20]], %[[VAL_2]] : i8
+// CHECK:           cf.cond_br %[[VAL_21]], ^bb15, ^bb16
+// CHECK:         ^bb15:
+// CHECK:           return %[[VAL_19]] : i8
+// CHECK:         ^bb16:
+// CHECK:           %[[VAL_22:.*]] = arith.muli %[[VAL_14]], %[[VAL_14]]  : i8
+// CHECK:           cf.br ^bb12(%[[VAL_19]], %[[VAL_22]], %[[VAL_20]] : i8, i8, i8)
+// CHECK:         }
+
+// -----
+
+// CHECK-LABEL:   func.func @ipowi_vec(
+// CHECK-SAME:                          %[[VAL_0:.*]]: vector<2x3xi64>,
+// CHECK-SAME:                          %[[VAL_1:.*]]: vector<2x3xi64>) {
+func.func @ipowi_vec(%arg0: vector<2x3xi64>, %arg1: vector<2x3xi64>) {
+// CHECK:   %[[CST:.*]] = arith.constant dense<0> : vector<2x3xi64>
+// CHECK:   %[[B00:.*]] = vector.extract %[[VAL_0]][0, 0] : vector<2x3xi64>
+// CHECK:   %[[E00:.*]] = vector.extract %[[VAL_1]][0, 0] : vector<2x3xi64>
+// CHECK:   %[[R00:.*]] = call @__mlir_math_ipowi_i64(%[[B00]], %[[E00]]) : (i64, i64) -> i64
+// CHECK:   %[[TMP00:.*]] = vector.insert %[[R00]], %[[CST]] [0, 0] : i64 into vector<2x3xi64>
+// CHECK:   %[[B01:.*]] = vector.extract %[[VAL_0]][0, 1] : vector<2x3xi64>
+// CHECK:   %[[E01:.*]] = vector.extract %[[VAL_1]][0, 1] : vector<2x3xi64>
+// CHECK:   %[[R01:.*]] = call @__mlir_math_ipowi_i64(%[[B01]], %[[E01]]) : (i64, i64) -> i64
+// CHECK:   %[[TMP01:.*]] = vector.insert %[[R01]], %[[TMP00]] [0, 1] : i64 into vector<2x3xi64>
+// CHECK:   %[[B02:.*]] = vector.extract %[[VAL_0]][0, 2] : vector<2x3xi64>
+// CHECK:   %[[E02:.*]] = vector.extract %[[VAL_1]][0, 2] : vector<2x3xi64>
+// CHECK:   %[[R02:.*]] = call @__mlir_math_ipowi_i64(%[[B02]], %[[E02]]) : (i64, i64) -> i64
+// CHECK:   %[[TMP02:.*]] = vector.insert %[[R02]], %[[TMP01]] [0, 2] : i64 into vector<2x3xi64>
+// CHECK:   %[[B10:.*]] = vector.extract %[[VAL_0]][1, 0] : vector<2x3xi64>
+// CHECK:   %[[E10:.*]] = vector.extract %[[VAL_1]][1, 0] : vector<2x3xi64>
+// CHECK:   %[[R10:.*]] = call @__mlir_math_ipowi_i64(%[[B10]], %[[E10]]) : (i64, i64) -> i64
+// CHECK:   %[[TMP10:.*]] = vector.insert %[[R10]], %[[TMP02]] [1, 0] : i64 into vector<2x3xi64>
+// CHECK:   %[[B11:.*]] = vector.extract %[[VAL_0]][1, 1] : vector<2x3xi64>
+// CHECK:   %[[E11:.*]] = vector.extract %[[VAL_1]][1, 1] : vector<2x3xi64>
+// CHECK:   %[[R11:.*]] = call @__mlir_math_ipowi_i64(%[[B11]], %[[E11]]) : (i64, i64) -> i64
+// CHECK:   %[[TMP11:.*]] = vector.insert %[[R11]], %[[TMP10]] [1, 1] : i64 into vector<2x3xi64>
+// CHECK:   %[[B12:.*]] = vector.extract %[[VAL_0]][1, 2] : vector<2x3xi64>
+// CHECK:   %[[E12:.*]] = vector.extract %[[VAL_1]][1, 2] : vector<2x3xi64>
+// CHECK:   %[[R12:.*]] = call @__mlir_math_ipowi_i64(%[[B12]], %[[E12]]) : (i64, i64) -> i64
+// CHECK:   %[[TMP12:.*]] = vector.insert %[[R12]], %[[TMP11]] [1, 2] : i64 into vector<2x3xi64>
+// CHECK:   return
+// CHECK: }
+  %0 = math.ipowi %arg0, %arg1 : vector<2x3xi64>
+  func.return
+}