support lowering or aten.log1p

Author: zezhang

lib/Conversion/TorchToTcp/Elementwise.cpp

@@ -478,6 +478,37 @@ class ConvertAtenSqrtOp : public OpConversionPattern<AtenSqrtOp> {
   }
 };
 
+class ConvertAtenLog1pOp : public OpConversionPattern<AtenLog1pOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(AtenLog1pOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    Value input = adaptor.getSelf();
+    RankedTensorType inputType = input.getType().dyn_cast<RankedTensorType>();
+
+    if (!inputType)
+      return rewriter.notifyMatchFailure(
+          op, "Only Ranked Tensor types are supported in TCP");
+
+    auto elementType = inputType.getElementType();
+    if (!isa<mlir::FloatType>(elementType))
+      return rewriter.notifyMatchFailure(
+          op, "Only floating-point datatype is supported");
+
+    auto constOp = torch_to_tcp::getConstTensor<float>(
+                       rewriter, op, llvm::ArrayRef((float)1.0), {})
+                       .value();
+    constOp = torch_to_tcp::broadcast0DOr1DToNDAndMatchShape(
+        rewriter, constOp, input, elementType);
+    auto addOp =
+        rewriter.create<tcp::AddOp>(op.getLoc(), inputType, input, constOp);
+    rewriter.replaceOpWithNewOp<tcp::LogOp>(op, inputType, addOp);
+    return success();
+  }
+};
+
 template <typename AtenOpT, typename TcpOpT>
 class ConvertAtenUnaryIntOrFpOp : public OpConversionPattern<AtenOpT> {
 public:
@@ -694,6 +725,7 @@ void torch_to_tcp::populateElementwisePatternsAndLegality(
   INSERT_ATEN_ELEMENTWISE_OP_PATTERN(AtenBatchNormOp);
   INSERT_ATEN_ELEMENTWISE_OP_PATTERN(AtenAtan2Op);
   INSERT_ATEN_ELEMENTWISE_OP_PATTERN(AtenSqrtOp);
+  INSERT_ATEN_ELEMENTWISE_OP_PATTERN(AtenLog1pOp);
 #undef INSERT_ATEN_ELEMENTWISE_OP_PATTERN
 
 #define INSERT_ATEN_ELEMENTWISE_ADD_SUB_PATTERN(AtenOp, TcpOp)                 \

test/Conversion/TorchToTcp/elementwise.mlir

@@ -762,3 +762,22 @@ func.func @torch.aten.to.dtype(%arg0: !torch.vtensor<[?,?],i1>) -> !torch.vtenso
   %0 = torch.aten.to.dtype %arg0, %int11, %false, %false, %none : !torch.vtensor<[?,?],i1>, !torch.int, !torch.bool, !torch.bool, !torch.none -> !torch.vtensor<[?,?],ui8>
   return %0 : !torch.vtensor<[?,?],ui8>
 }
+
+// -----
+
+// CHECK-LABEL:  func.func @torch.aten.log1p(
+// CHECK-SAME:         %[[ARG0:.*]]: !torch.vtensor<[?,4,19,2],f32>) -> !torch.vtensor<[?,4,19,2],f32> {
+// CHECK-DAG:     %[[TO_BUILTIN0:.*]] = torch_c.to_builtin_tensor %[[ARG0]] : !torch.vtensor<[?,4,19,2],f32> -> tensor<?x4x19x2xf32>
+// CHECK:         %[[CONST:.*]] =  tcp.const {value = dense<1.000000e+00> : tensor<f32>} : tensor<f32>
+// CHECK:         %[[EXPAND_SHAPE:.*]] = tensor.expand_shape %[[CONST]] [] output_shape [1, 1, 1, 1] : tensor<f32> into tensor<1x1x1x1xf32>
+// CHECK:         %[[CONST0:.*]] = arith.constant 0 : index
+// CHECK:         %[[DIM0:.*]] = tensor.dim %[[TO_BUILTIN0]], %[[CONST0]] : tensor<?x4x19x2xf32>
+// CHECK:         %[[BROADCAST:.*]] = tcp.broadcast %[[EXPAND_SHAPE]], %[[DIM0]]
+// CHECK:         %[[ADD:.*]] = tcp.add %[[TO_BUILTIN0]], %[[BROADCAST]] : tensor<?x4x19x2xf32>, tensor<?x4x19x2xf32> -> tensor<?x4x19x2xf32>
+// CHECK:         %[[LOG:.*]] = tcp.log %[[ADD]] : tensor<?x4x19x2xf32> -> tensor<?x4x19x2xf32>
+// CHECK:         %[[FROM_BUILTIN:.*]] = torch_c.from_builtin_tensor %[[LOG]] : tensor<?x4x19x2xf32> -> !torch.vtensor<[?,4,19,2],f32>
+// CHECK:         return %[[FROM_BUILTIN]] : !torch.vtensor<[?,4,19,2],f32>
+func.func @torch.aten.log1p(%arg0: !torch.vtensor<[?,4,19,2],f32>) -> !torch.vtensor<[?,4,19,2],f32> {
+  %1 = torch.aten.log1p %arg0 : !torch.vtensor<[?,4,19,2],f32> -> !torch.vtensor<[?,4,19,2],f32>
+  return %1 : !torch.vtensor<[?,4,19,2],f32>
+}