[Torch] support AtenExp2Op (#3832)

- support AtenExp2Op by decomposing it to aten.pow.scalar
- refine stablehlo pow.scalar pow.Tensor_Scalar pow.Tensor_Tensor
lowering according to https://github.com/llvm/torch-mlir/pull/2983
- Close https://github.com/llvm/torch-mlir/pull/2983

include/torch-mlir/Dialect/Torch/IR/GeneratedTorchOps.td

@@ -941,6 +941,51 @@ def Torch_AtenExp_Op : Torch_Op<"aten.exp_", [
   }];
 }
 
+def Torch_AtenExp2Op : Torch_Op<"aten.exp2", [
+    AllowsTypeRefinement,
+    HasValueSemantics,
+    ReadOnly
+  ]> {
+  let summary = "Generated op for `aten::exp2 : (Tensor) -> (Tensor)`";
+  let arguments = (ins
+    AnyTorchTensorType:$self
+  );
+  let results = (outs
+    AnyTorchOptionalTensorType:$result
+  );
+  let hasCustomAssemblyFormat = 1;
+  let extraClassDefinition = [{
+    ParseResult AtenExp2Op::parse(OpAsmParser &parser, OperationState &result) {
+      return parseDefaultTorchOp(parser, result, 1, 1);
+    }
+    void AtenExp2Op::print(OpAsmPrinter &printer) {
+      printDefaultTorchOp(printer, *this, 1, 1);
+    }
+  }];
+}
+
+def Torch_AtenExp2_Op : Torch_Op<"aten.exp2_", [
+    IsTrailingUnderscoreInplaceVariant,
+    AllowsTypeRefinement
+  ]> {
+  let summary = "Generated op for `aten::exp2_ : (Tensor) -> (Tensor)`";
+  let arguments = (ins
+    Torch_NonValueTensorType:$self
+  );
+  let results = (outs
+    AnyTorchOptionalNonValueTensorType:$result
+  );
+  let hasCustomAssemblyFormat = 1;
+  let extraClassDefinition = [{
+    ParseResult AtenExp2_Op::parse(OpAsmParser &parser, OperationState &result) {
+      return parseDefaultTorchOp(parser, result, 1, 1);
+    }
+    void AtenExp2_Op::print(OpAsmPrinter &printer) {
+      printDefaultTorchOp(printer, *this, 1, 1);
+    }
+  }];
+}
+
 def Torch_AtenExpm1Op : Torch_Op<"aten.expm1", [
     AllowsTypeRefinement,
     HasValueSemantics,

lib/Conversion/TorchToStablehlo/Basic.cpp

@@ -931,79 +931,49 @@ LogicalResult ConvertAtenOp<AtenReciprocalOp>::matchAndRewrite(
   return success();
 }
 
-// AtenPowTensorScalarOp
+namespace {
-template <>
+template <typename AtenOpT>
-LogicalResult ConvertAtenOp<AtenPowTensorScalarOp>::matchAndRewrite(
+class ConvertAtenPowOp : public OpConversionPattern<AtenOpT> {
-    AtenPowTensorScalarOp op, OpAdaptor adaptor,
+public:
-    ConversionPatternRewriter &rewriter) const {
+  using OpConversionPattern<AtenOpT>::OpConversionPattern;
-  Value lhs = adaptor.getSelf();
+  using OpAdaptor = typename AtenOpT::Adaptor;
-  auto lhsType = dyn_cast<TensorType>(lhs.getType());
+  LogicalResult
-  Value rhs = adaptor.getExponent();
+  matchAndRewrite(AtenOpT op, OpAdaptor adaptor,
-  TensorType rhsType = dyn_cast<TensorType>(rhs.getType());
+                  ConversionPatternRewriter &rewriter) const override {
+    auto outType = cast<TensorType>(
+        OpConversionPattern<AtenPowScalarOp>::getTypeConverter()->convertType(
+            op.getType()));
 
-  if (!lhsType)
+    Type outElemTy = outType.getElementType();
-    return op.emitError("only Tensor types supported in StableHLO");
+    if (!outElemTy.isIntOrFloat()) {
+      return op.emitError(
+          "only floating-point or integer datatype legalization supported");
+    }
 
-  auto outType = cast<TensorType>(
+    Value lhs = adaptor.getSelf();
-      OpConversionPattern<AtenPowTensorScalarOp>::getTypeConverter()
+    auto lhsType = dyn_cast<TensorType>(lhs.getType());
-          ->convertType(op.getType()));
+    Value rhs = adaptor.getExponent();
+    auto rhsType = dyn_cast<TensorType>(rhs.getType());
 
-  Type outElemTy = outType.getElementType();
+    if (!lhsType && !rhsType) {
-  if (!outElemTy.isIntOrFloat()) {
+      return op.emitError("only Tensor types supported in StableHLO");
-    return op.emitError(
+    }
-        "only floating-point or integer datatype legalization supported");
+    if (!lhsType) {
+      lhs = hlo::scalarToStablehloTensor(rewriter, op, lhs, outElemTy);
+    }
+    if (!rhsType) {
+      rhs = hlo::scalarToStablehloTensor(rewriter, op, rhs, outElemTy);
+    }
+
+    lhs = hlo::promoteType(rewriter, op.getLoc(), lhs, outElemTy);
+    rhs = hlo::promoteType(rewriter, op.getLoc(), rhs, outElemTy);
+    DenseI64ArrayAttr bcastDimensions;
+    rewriter.replaceOpWithNewOp<chlo::BroadcastPowOp>(op, outType, lhs, rhs,
+                                                      bcastDimensions);
+    return success();
   }
-
+};
-  if (!rhsType) {
+} // namespace
-    rhs = hlo::scalarToStablehloTensor(rewriter, op, rhs, outElemTy);
-  }
-  DenseI64ArrayAttr bcastDimensions;
-  lhs = hlo::promoteType(rewriter, op.getLoc(), lhs, outElemTy);
-  rhs = hlo::promoteType(rewriter, op.getLoc(), rhs, outElemTy);
-  auto loc = op.getLoc();
-  Value result = rewriter.create<chlo::BroadcastPowOp>(loc, outType, lhs, rhs,
-                                                       bcastDimensions);
-
-  rewriter.replaceOp(op, result);
-  return success();
-}
-
-// AtenPowScalarOp
-template <>
-LogicalResult ConvertAtenOp<AtenPowScalarOp>::matchAndRewrite(
-    AtenPowScalarOp op, OpAdaptor adaptor,
-    ConversionPatternRewriter &rewriter) const {
-  Value lhs = adaptor.getSelf();
-  auto lhsType = dyn_cast<TensorType>(lhs.getType());
-  Value rhs = adaptor.getExponent();
-  auto rhsType = dyn_cast<TensorType>(rhs.getType());
-
-  if (!rhsType)
-    return op.emitError("only Tensor types supported in StableHLO");
-
-  auto outType = cast<TensorType>(
-      OpConversionPattern<AtenPowScalarOp>::getTypeConverter()->convertType(
-          op.getType()));
-
-  Type outElemTy = outType.getElementType();
-  if (!outElemTy.isIntOrFloat()) {
-    return op.emitError(
-        "only floating-point or integer datatype legalization supported");
-  }
-
-  if (!lhsType) {
-    lhs = hlo::scalarToStablehloTensor(rewriter, op, lhs, outElemTy);
-  }
-  DenseI64ArrayAttr bcastDimensions;
-  lhs = hlo::promoteType(rewriter, op.getLoc(), lhs, outElemTy);
-  rhs = hlo::promoteType(rewriter, op.getLoc(), rhs, outElemTy);
-  auto loc = op.getLoc();
-  Value result = rewriter.create<chlo::BroadcastPowOp>(loc, outType, lhs, rhs,
-                                                       bcastDimensions);
-
-  rewriter.replaceOp(op, result);
-  return success();
-}
 
 // PrimNumToTensorScalarOp
 template <>
@@ -1797,29 +1767,6 @@ LogicalResult ConvertAtenOp<AtenGeluBackwardOp>::matchAndRewrite(
   return success();
 }
 
-template <>
-LogicalResult ConvertAtenOp<AtenPowTensorTensorOp>::matchAndRewrite(
-    AtenPowTensorTensorOp op, OpAdaptor adaptor,
-    ConversionPatternRewriter &rewriter) const {
-  Value lhs = adaptor.getSelf();
-  auto lhsTy = cast<TensorType>(lhs.getType());
-  Value rhs = adaptor.getExponent();
-  auto rhsTy = cast<TensorType>(rhs.getType());
-
-  if (!lhsTy || !rhsTy)
-    return op.emitError("only Tensor types supported");
-
-  auto outTy =
-      cast<TensorType>(this->getTypeConverter()->convertType(op.getType()));
-
-  lhs = hlo::promoteType(rewriter, op.getLoc(), lhs, outTy.getElementType());
-  rhs = hlo::promoteType(rewriter, op.getLoc(), rhs, outTy.getElementType());
-
-  rewriter.replaceOpWithNewOp<chlo::BroadcastPowOp>(op, outTy, lhs, rhs,
-                                                    /*broadcast_attr*/ nullptr);
-  return success();
-}
-
 // Converts `aten.empty.memory_format` to `tensor.empty` op.
 template <>
 LogicalResult ConvertAtenOp<AtenEmptyMemoryFormatOp>::matchAndRewrite(
@@ -2250,6 +2197,14 @@ void mlir::torch::torch_to_stablehlo::populateBasicOpPatternsAndLegality(
 
 #undef INSERT_BINARY_LOGICAL_PATTERN
 
+#define INSERT_BINARY_POW_PATTERN(AtenOp)                                      \
+  target.addIllegalOp<AtenOp>();                                               \
+  patterns.add<ConvertAtenPowOp<AtenOp>>(typeConverter, context)
+  INSERT_BINARY_POW_PATTERN(AtenPowTensorScalarOp);
+  INSERT_BINARY_POW_PATTERN(AtenPowTensorTensorOp);
+  INSERT_BINARY_POW_PATTERN(AtenPowScalarOp);
+#undef INSERT_BINARY_ADDSUB_PATTERN
+
 #define INSERT_ATENOP_PATTERN(AtenOp)                                          \
   target.addIllegalOp<AtenOp>();                                               \
   patterns.add<ConvertAtenOp<AtenOp>>(typeConverter, context, options)
@@ -2260,8 +2215,6 @@ void mlir::torch::torch_to_stablehlo::populateBasicOpPatternsAndLegality(
   INSERT_ATENOP_PATTERN(ValueTensorLiteralOp);
   INSERT_ATENOP_PATTERN(AtenTensorIntOp);
   INSERT_ATENOP_PATTERN(AtenReciprocalOp);
-  INSERT_ATENOP_PATTERN(AtenPowTensorScalarOp);
-  INSERT_ATENOP_PATTERN(AtenPowScalarOp);
   INSERT_ATENOP_PATTERN(PrimNumToTensorScalarOp);
   INSERT_ATENOP_PATTERN(AtenScalarImplicitOp);
   INSERT_ATENOP_PATTERN(AtenContiguousOp);
@@ -2285,7 +2238,6 @@ void mlir::torch::torch_to_stablehlo::populateBasicOpPatternsAndLegality(
   INSERT_ATENOP_PATTERN(AtenSizeIntOp);
   INSERT_ATENOP_PATTERN(AtenToDtypeOp);
   INSERT_ATENOP_PATTERN(AtenWhereSelfOp);
-  INSERT_ATENOP_PATTERN(AtenPowTensorTensorOp);
 
   INSERT_ATENOP_PATTERN(AtenEmptyMemoryFormatOp);
   INSERT_ATENOP_PATTERN(AtenFillScalarOp);

lib/Dialect/Torch/Transforms/AbstractInterpLibrary.cpp

@@ -6487,6 +6487,10 @@ StringRef mlir::torch::Torch::getAbstractInterpLibrary() {
 "    %0 = call @__torch__.torch.jit._shape_functions.unary(%arg0) : (!torch.list<int>) -> !torch.list<int>\n"
 "    return %0 : !torch.list<int>\n"
 "  }\n"
+"  func.func @\"__torch_mlir_shape_fn.aten.exp2\"(%arg0: !torch.list<int>) -> !torch.list<int> {\n"
+"    %0 = call @__torch__.torch.jit._shape_functions.unary(%arg0) : (!torch.list<int>) -> !torch.list<int>\n"
+"    return %0 : !torch.list<int>\n"
+"  }\n"
 "  func.func @\"__torch_mlir_shape_fn.aten.expm1\"(%arg0: !torch.list<int>) -> !torch.list<int> {\n"
 "    %0 = call @__torch__.torch.jit._shape_functions.unary(%arg0) : (!torch.list<int>) -> !torch.list<int>\n"
 "    return %0 : !torch.list<int>\n"
@@ -11142,6 +11146,11 @@ StringRef mlir::torch::Torch::getAbstractInterpLibrary() {
 "    %1 = call @__torch__._get_dtype_of_floating_point_op(%0#1) : (!torch.int) -> !torch.int\n"
 "    return %1 : !torch.int\n"
 "  }\n"
+"  func.func @\"__torch_mlir_dtype_fn.aten.exp2\"(%arg0: !torch.tuple<int, int>) -> !torch.int {\n"
+"    %0:2 = torch.prim.TupleUnpack %arg0 : !torch.tuple<int, int> -> !torch.int, !torch.int\n"
+"    %1 = call @__torch__._get_dtype_of_floating_point_op(%0#1) : (!torch.int) -> !torch.int\n"
+"    return %1 : !torch.int\n"
+"  }\n"
 "  func.func @\"__torch_mlir_dtype_fn.aten.expm1\"(%arg0: !torch.tuple<int, int>) -> !torch.int {\n"
 "    %0:2 = torch.prim.TupleUnpack %arg0 : !torch.tuple<int, int> -> !torch.int, !torch.int\n"
 "    %1 = call @__torch__._get_dtype_of_floating_point_op(%0#1) : (!torch.int) -> !torch.int\n"

lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp

@@ -8590,6 +8590,24 @@ class DecomposeAtenBinaryCrossEntropyWithLogitsOp
 };
 } // namespace
 
+namespace {
+class DecomposeAtenExp2Op : public OpRewritePattern<AtenExp2Op> {
+  using OpRewritePattern<AtenExp2Op>::OpRewritePattern;
+  LogicalResult matchAndRewrite(AtenExp2Op op,
+                                PatternRewriter &rewriter) const override {
+    Location loc = op.getLoc();
+    Value self = op.getSelf();
+
+    auto two =
+        rewriter.create<ConstantIntOp>(loc, rewriter.getI64IntegerAttr(2));
+    rewriter.replaceOpWithNewOp<AtenPowScalarOp>(op, op.getType(), two, self);
+
+    return success();
+  }
+};
+
+} // namespace
+
 namespace {
 class DecomposeAtenOneHotOp : public OpRewritePattern<AtenOneHotOp> {
   using OpRewritePattern<AtenOneHotOp>::OpRewritePattern;
@@ -9721,6 +9739,7 @@ public:
     addPatternIfTargetOpIsIllegal<DecomposePrimTolistOp>(patterns);
     addPatternIfTargetOpIsIllegal<DecomposePrimsSqueezeOp>(patterns);
     addPatternIfTargetOpIsIllegal<DecomposeAtenMovedimIntOp>(patterns);
+    addPatternIfTargetOpIsIllegal<DecomposeAtenExp2Op>(patterns);
     addPatternIfTargetOpIsIllegal<DecomposeAtenOneHotOp>(patterns);
     addPatternIfTargetOpIsIllegal<DecomposeAtenCrossEntropyLossOp>(patterns);
     addPatternIfTargetOpIsIllegal<DecomposeAtenBinaryCrossEntropyWithLogitsOp>(

projects/pt1/e2e_testing/xfail_sets.py

@@ -2311,6 +2311,9 @@ ONNX_XFAIL_SET = {
     "ElementwiseLog2IntModule_basic",
     "ElementwiseFminModule_basic",
     "ElementwiseFmaxModule_basic",
+    "Exp2StaticModule_basic",
+    "MultinomialModule2D_basic",
+    "MultinomialModule2D_F32",
     "PixelShuffleModuleStaticRank4Float32_basic",
     "ReflectionPad1dModule2dInput_Right",
     "ReflectionPad1dModule2dInput_basic",

projects/pt1/python/torch_mlir/jit_ir_importer/build_tools/abstract_interp_lib_gen.py

@@ -216,6 +216,9 @@ def aten〇silu〡shape(self: List[int]) -> List[int]:
 def aten〇exp〡shape(self: List[int]) -> List[int]:
     return upstream_shape_functions.unary(self)
 
+def aten〇exp2〡shape(self: List[int]) -> List[int]:
+    return upstream_shape_functions.unary(self)
+
 def aten〇expm1〡shape(self: List[int]) -> List[int]:
     return upstream_shape_functions.unary(self)
 
@@ -2513,6 +2516,11 @@ def aten〇exp〡dtype(self_rank_dtype: Tuple[int, int]) -> int:
     self_rank, self_dtype = self_rank_dtype
     return _get_dtype_of_floating_point_op(self_dtype)
 
+@check_dtype_function(_check_tensors_with_the_same_dtype(num_of_tensors=1))
+def aten〇exp2〡dtype(self_rank_dtype: Tuple[int, int]) -> int:
+    self_rank, self_dtype = self_rank_dtype
+    return _get_dtype_of_floating_point_op(self_dtype)
+
 @check_dtype_function(_check_tensors_with_the_same_dtype(num_of_tensors=1))
 def aten〇expm1〡dtype(self_rank_dtype: Tuple[int, int]) -> int:
     self_rank, self_dtype = self_rank_dtype

projects/pt1/python/torch_mlir/jit_ir_importer/build_tools/torch_ods_gen.py

@@ -316,6 +316,7 @@ def emit_ops(emitter_td: TextEmitter, registry: Registry):
         "aten::asin : (Tensor) -> (Tensor)",
         "aten::asinh : (Tensor) -> (Tensor)",
         "aten::exp : (Tensor) -> (Tensor)",
+        "aten::exp2 : (Tensor) -> (Tensor)",
         "aten::expm1 : (Tensor) -> (Tensor)",
         "aten::cos : (Tensor) -> (Tensor)",
         "aten::cosh : (Tensor) -> (Tensor)",

projects/pt1/python/torch_mlir_e2e_test/test_suite/elementwise.py

@@ -2772,6 +2772,29 @@ def ElementwiseSgnModule_basic(module, tu: TestUtils):
 # ==============================================================================
 
 
+class Exp2StaticModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args(
+        [
+            None,
+            ([3, 2], torch.float32, True),
+        ]
+    )
+    def forward(self, x):
+        return torch.ops.aten.exp2(x)
+
+
+@register_test_case(module_factory=lambda: Exp2StaticModule())
+def Exp2StaticModule_basic(module, tu: TestUtils):
+    module.forward(tu.rand(3, 2))
+
+
+# ==============================================================================
+
+
 class ElementwisePowModule(torch.nn.Module):
     def __init__(self):
         super().__init__()