[mlir] Add lowering from math::Log1p to LLVM

[ezhulenev]

Summary

[mlir] Add lowering from math::Log1p to LLVM

[ftynse]

I wonder if this is something that should be instead implemented as an "expansion" pass on math dialect, converting it to std.constant, std.addf and math.log that are already handled by further lowerings.

[ftynse]

In D98662#2629031, @ftynse wrote:

I wonder if this is something that should be instead implemented as an "expansion" pass on math dialect, converting it to std.constant, std.addf and math.log that are already handled by further lowerings.

I haven't seen an answer to this. Is there any reason why we should grow "standard-to-llvm" further instead of having more local conversions?

[mehdi_amini]

I haven't seen an answer to this. Is there any reason why we should grow "standard-to-llvm" further instead of having more local conversions?

+1

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 `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);
+    auto operandType = transformed.operand().getType();
+
+    if (!operandType || !LLVM::isCompatibleType(operandType))
+      return rewriter.notifyMatchFailure(op, "unsupported operand type");
+
+    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 =
+          LLVM::isCompatibleVectorType(operandType)
+              ? rewriter.create<LLVM::ConstantOp>(
+                    loc, operandType,
+                    SplatElementsAttr::get(resultType.cast<ShapedType>(),
+                                           floatOne))
+              : rewriter.create<LLVM::ConstantOp>(loc, operandType, floatOne);
+
+      auto add = rewriter.create<LLVM::FAddOp>(loc, operandType, one,
+                                               transformed.operand());
+      rewriter.replaceOpWithNewOp<LLVM::LogOp>(op, operandType, add);
+      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 add = rewriter.create<LLVM::FAddOp>(loc, llvm1DVectorTy, one,
+                                                   transformed.operand());
+          return rewriter.create<LLVM::LogOp>(loc, llvm1DVectorTy, add);
+        },
+        rewriter);
+  }
+};
+
 // A `rsqrt` is converted into `1 / sqrt`.
 struct RsqrtOpLowering : public ConvertOpToLLVMPattern<math::RsqrtOp> {
   using ConvertOpToLLVMPattern<math::RsqrtOp>::ConvertOpToLLVMPattern;
 
   LogicalResult
   matchAndRewrite(math::RsqrtOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     math::RsqrtOp::Adaptor transformed(operands);
patterns.insert<
       DivFOpLowering,
       ExpOpLowering,
       Exp2OpLowering,
       FloorFOpLowering,
       FmaFOpLowering,
       GenericAtomicRMWOpLowering,
       LogOpLowering,
       Log10OpLowering,
+      Log1pOpLowering,
       Log2OpLowering,
       FPExtLowering,
       FPToSILowering,
       FPToUILowering,
       FPTruncLowering,
       IndexCastOpLowering,
       MulFOpLowering,
       MulIOpLowering,

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

 // RUN: mlir-opt -allow-unregistered-dialect %s -convert-std-to-llvm -split-input-file -verify-diagnostics | FileCheck %s
 
 // CHECK-LABEL: func @address_space(
 // CHECK-SAME:    !llvm.ptr<f32, 7>
 func @address_space(%arg0 : memref<32xf32, affine_map<(d0) -> (d0)>, 7>) {
   %0 = memref.alloc() : memref<32xf32, affine_map<(d0) -> (d0)>, 5>
   %1 = constant 7 : index
   // CHECK: llvm.load %{{.*}} : !llvm.ptr<f32, 5>
   %2 = memref.load %0[%1] : memref<32xf32, affine_map<(d0) -> (d0)>, 5>
   std.return
 }
 
 // -----
 
+// CHECK-LABEL: func @log1p(
+// CHECK-SAME: f32
+func @log1p(%arg0 : f32) {
+  // CHECK: %[[ONE:.*]] = llvm.mlir.constant(1.000000e+00 : f32) : f32
+  // CHECK: %[[ADD:.*]] = llvm.fadd %[[ONE]], %arg0 : f32
+  // CHECK: %[[LOG:.*]] = "llvm.intr.log"(%[[ADD]]) : (f32) -> f32
+  %0 = math.log1p %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