[mlir] Add lowering from math.expm1 to LLVM.

[akuegel]

Summary

This is the straightforward lowering to exp(x) - 1.

[akuegel]

Fix comment that was still copy/pasted from rsqrt.

[ftynse]

I won't ask to split out the math conversion from "standard" conversion, yet, but have you considered implementing this (as well as rsqrt from which most of the code is duplicated) as a sort of expand-math conversion instead?

[herhut]

In D96776#2565842, @ftynse wrote:

I won't ask to split out the math conversion from "standard" conversion, yet, but have you considered implementing this (as well as rsqrt from which most of the code is duplicated) as a sort of expand-math conversion instead?

Splitting this out is on my todo list. I though about having a general lowering to libm / llvm intrinsics that converts everything into a form that can be mapped that way. Additionally, we will have a lowering to GPU intrinsics and also the explicit approximations that are currently being worked on. Then the user of these lowerings can choose between those three sets of patterns, depending on their use-case. WDYT?

[mehdi_amini]

Then the user of these lowerings can choose between those three sets of patterns, depending on their use-case. WDYT?

As long as the patterns are reusable so that the user can combine them and decide what is legal/illegal to chose a lowering strategy "per operation" :)

[akuegel]

Sorry for not following up on this patch. I was very busy this week to work on a new op that we need for Kernel Generator broadcast specialization, and that is blocking and thus has priority over this patch.
I hope I can do the necessary adjustments to this patch soon.

[cota]

Ping.
I just added log1p's approximation (https://reviews.llvm.org/D101765). I'd like this patch to go in so that I can also add expm1's approximation.

[akuegel]

Rebase

[akuegel]

In D96776#2734777, @cota wrote:

Ping.
I just added log1p's approximation (https://reviews.llvm.org/D101765). I'd like this patch to go in so that I can also add expm1's approximation.

Essentially this patch is blocked by same thing that was also commented on https://reviews.llvm.org/D98662, that we should not add additional copy/paste conversions here.
Personally I would be ok with landing this patch as-is to unblock you, if that is acceptable. But I am still working on other things which have a higher priority to my team than working on a cleanup here by adding a new math conversion pass.
@herhut Please correct me in case I got the prioritization wrong.

[ftynse]

Like I said, I won't block this (and I haven't blocked https://reviews.llvm.org/D98662 either), but let's try and find time across teams to proceed with the refactoring -- it's not a priority for me either, but maintaining an ever-growing list of patterns may become hard.

[akuegel]

In D96776#2735967, @ftynse wrote:

Like I said, I won't block this (and I haven't blocked https://reviews.llvm.org/D98662 either), but let's try and find time across teams to proceed with the refactoring -- it's not a priority for me either, but maintaining an ever-growing list of patterns may become hard.

Maybe we should have a fixit week for such cleanups :)
And I also agree that a separate pass would be better (so that we would not have to deal with the "is it a vector type" separately in each such op, and instead can rely on that handling inside the lowering of the ops into which it is composed).

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

auto deadBeefPtr =
         rewriter.create<LLVM::IntToPtrOp>(loc, elementPtrType, deadBeefConst);
 
     // Both allocated and aligned pointers are same. We could potentially stash
     // a nullptr for the allocated pointer since we do not expect any dealloc.
     return std::make_tuple(deadBeefPtr, gep);
   }
 };
 
+// A `expm1` is converted into `exp - 1`.
+struct ExpM1OpLowering : public ConvertOpToLLVMPattern<math::ExpM1Op> {
+  using ConvertOpToLLVMPattern<math::ExpM1Op>::ConvertOpToLLVMPattern;
+
+  LogicalResult
+  matchAndRewrite(math::ExpM1Op op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    math::ExpM1Op::Adaptor transformed(operands);
+    auto operandType = transformed.operand().getType();
+
+    if (!operandType || !LLVM::isCompatibleType(operandType))
+      return failure();
+
+    auto loc = op.getLoc();
+    auto resultType = op.getResult().getType();
+    auto floatType = getElementTypeOrSelf(resultType).cast<FloatType>();
+    auto floatOne = rewriter.getFloatAttr(floatType, 1.0);
+
+    if (!operandType.isa<LLVM::LLVMArrayType>()) {
+      LLVM::ConstantOp one;
+      if (LLVM::isCompatibleVectorType(operandType)) {
+        one = rewriter.create<LLVM::ConstantOp>(
+            loc, operandType,
+            SplatElementsAttr::get(resultType.cast<ShapedType>(), floatOne));
+      } else {
+        one = rewriter.create<LLVM::ConstantOp>(loc, operandType, floatOne);
+      }
+      auto exp = rewriter.create<LLVM::ExpOp>(loc, transformed.operand());
+      rewriter.replaceOpWithNewOp<LLVM::FSubOp>(op, operandType, exp, one);
+      return success();
+    }
+
+    auto vectorType = resultType.dyn_cast<VectorType>();
+    if (!vectorType)
+      return rewriter.notifyMatchFailure(op, "expected vector result type");
+
+    return handleMultidimensionalVectors(
+        op.getOperation(), operands, *getTypeConverter(),
+        [&](Type llvm1DVectorTy, ValueRange operands) {
+          auto splatAttr = SplatElementsAttr::get(
+              mlir::VectorType::get(
+                  {LLVM::getVectorNumElements(llvm1DVectorTy).getFixedValue()},
+                  floatType),
+              floatOne);
+          auto one =
+              rewriter.create<LLVM::ConstantOp>(loc, llvm1DVectorTy, splatAttr);
+          auto exp =
+              rewriter.create<LLVM::ExpOp>(loc, llvm1DVectorTy, operands[0]);
+          return rewriter.create<LLVM::FSubOp>(loc, llvm1DVectorTy, exp, one);
+        },
+        rewriter);
+  }
+};
+
 // A `log1p` is converted into `log(1 + ...)`.
 struct Log1pOpLowering : public ConvertOpToLLVMPattern<math::Log1pOp> {
   using ConvertOpToLLVMPattern<math::Log1pOp>::ConvertOpToLLVMPattern;
 
   LogicalResult
   matchAndRewrite(math::Log1pOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     math::Log1pOp::Adaptor transformed(operands);
patterns.add<
       CondBranchOpLowering,
       CopySignOpLowering,
       CosOpLowering,
       ConstantOpLowering,
       DialectCastOpLowering,
       DivFOpLowering,
       ExpOpLowering,
       Exp2OpLowering,
+      ExpM1OpLowering,
       FloorFOpLowering,
       FmaFOpLowering,
       GenericAtomicRMWOpLowering,
       LogOpLowering,
       Log10OpLowering,
       Log1pOpLowering,
       Log2OpLowering,
       FPExtLowering,

mlir/test/Conversion/StandardToLLVM/standard-to-llvm.mlir

func @log1p_2dvector(%arg0 : vector<4x3xf32>) {
   // CHECK: %[[LOG:.*]] = "llvm.intr.log"(%[[ADD]]) : (vector<3xf32>) -> vector<3xf32>
   // CHECK: %[[INSERT:.*]] = llvm.insertvalue %[[LOG]], %0[0] : !llvm.array<4 x vector<3xf32>>
   %0 = math.log1p %arg0 : vector<4x3xf32>
   std.return
 }
 
 // -----
 
+// CHECK-LABEL: func @expm1(
+// CHECK-SAME: f32
+func @expm1(%arg0 : f32) {
+  // CHECK: %[[ONE:.*]] = llvm.mlir.constant(1.000000e+00 : f32) : f32
+  // CHECK: %[[EXP:.*]] = "llvm.intr.exp"(%arg0) : (f32) -> f32
+  // CHECK: %[[SUB:.*]] = llvm.fsub %[[EXP]], %[[ONE]] : f32
+  %0 = math.expm1 %arg0 : f32
+  std.return
+}
+
+// -----
+
 // CHECK-LABEL: func @rsqrt(
 // CHECK-SAME: f32
 func @rsqrt(%arg0 : f32) {
   // CHECK: %[[ONE:.*]] = llvm.mlir.constant(1.000000e+00 : f32) : f32
   // CHECK: %[[SQRT:.*]] = "llvm.intr.sqrt"(%arg0) : (f32) -> f32
   // CHECK: %[[DIV:.*]] = llvm.fdiv %[[ONE]], %[[SQRT]] : f32
   %0 = math.rsqrt %arg0 : f32
   std.return