Reduction.cpp

//===----------------------------------------------------------------------===//
//
// 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
// Also available under a BSD-style license. See LICENSE.
//
//===----------------------------------------------------------------------===//

#include "torch-mlir/Conversion/TorchToLinalg/TorchToLinalg.h"

#include "PopulatePatterns.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Complex/IR/Complex.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Math/IR/Math.h"
#include "mlir/IR/Matchers.h"
#include "torch-mlir/Conversion/TorchToLinalg/Utils.h"
#include "torch-mlir/Conversion/Utils/Utils.h"
#include "torch-mlir/Dialect/Torch/IR/TorchOps.h"
#include "torch-mlir/Dialect/Torch/Utils/TorchUpstream.h"
#include "torch-mlir/Dialect/Torch/Utils/Utils.h"
#include "llvm/ADT/APSInt.h"

using namespace mlir;
using namespace mlir::torch;
using namespace mlir::torch::Torch;

namespace {
// Runs an inclusive prefix sum along the middle dimension of a rank-3 tensor.
// The input shape is [outer, scan, inner]; the binary lifting loop keeps the
// implementation rank-independent after aten.cumsum is reshaped into this form.
static Value emitInclusiveScanByPowersOfTwo(Value running,
                                            ConversionPatternRewriter &rewriter,
                                            Location loc) {
  RankedTensorType runningType = cast<RankedTensorType>(running.getType());
  SmallVector<int64_t> runningShape =
      makeShapeTorchCompatible(runningType.getShape());
  int64_t outer = runningShape[0];
  int64_t scanDimSize = runningShape[1];
  int64_t inner = runningShape[2];
  Type elementType = runningType.getElementType();

  Value zero = arith::ConstantOp::create(rewriter, loc,
                                         rewriter.getZeroAttr(elementType));

  SmallVector<OpFoldResult> sliceOffsets(3, rewriter.getIndexAttr(0));
  SmallVector<OpFoldResult> sliceSizes = {rewriter.getIndexAttr(outer),
                                          rewriter.getIndexAttr(scanDimSize),
                                          rewriter.getIndexAttr(inner)};
  SmallVector<OpFoldResult> sliceStrides(3, rewriter.getIndexAttr(1));

  for (int64_t offset = 1; offset < scanDimSize; offset <<= 1) {
    SmallVector<int64_t> lowPad = {0, offset, 0};
    SmallVector<int64_t> highPad = {0, 0, 0};
    Type paddedType =
        tensor::PadOp::inferResultType(runningType, lowPad, highPad);
    SmallVector<OpFoldResult> lowPadValues = {rewriter.getIndexAttr(0),
                                              rewriter.getIndexAttr(offset),
                                              rewriter.getIndexAttr(0)};
    SmallVector<OpFoldResult> highPadValues(3, rewriter.getIndexAttr(0));
    Value padded = tensor::PadOp::create(rewriter, loc, paddedType, running,
                                         lowPadValues, highPadValues, zero);

    Value shifted =
        tensor::ExtractSliceOp::create(rewriter, loc, runningType, padded,
                                       sliceOffsets, sliceSizes, sliceStrides);

    running = torch_to_linalg::createElementwiseLinalgGeneric(
        rewriter, loc, {running, shifted}, elementType,
        [&](OpBuilder &builder, Location loc, ValueRange payloadArgs) {
          Value result;
          if (isa<mlir::FloatType>(elementType))
            result = arith::AddFOp::create(builder, loc, payloadArgs[0],
                                           payloadArgs[1]);
          else if (isa<mlir::IntegerType>(elementType))
            result = arith::AddIOp::create(builder, loc, payloadArgs[0],
                                           payloadArgs[1]);
          else
            llvm_unreachable("unsupported cumsum element type");
          linalg::YieldOp::create(builder, loc, result);
        });
  }

  return running;
}

// Aten max.dim (min.dim) lowering represents the MaxDimOp (MinDimOp) as an
// linalg.indexed_generic op, producing two output buffers.
//
// The first output buffer contains the maximum (minium) value found. It is
// initialized to the minimum (maximum) representable value of the input
// element type.
//
// The second output buffer contains the index of the found maximum (minimum)
// value. It is initialized to 0 and is resulting integer type.
//
// The indexed_generic op updates both the maximum (minimum) value and index
// if the current value exceeds the running max (min).
template <typename OpTy>
class ConvertAtenMinMaxDimOp : public OpConversionPattern<OpTy> {

private:
  bool allowNonFinites;

public:
  using OpConversionPattern<OpTy>::OpConversionPattern;
  using OpConversionPattern<OpTy>::getTypeConverter;

  using OpAdaptor = typename OpTy::Adaptor;

  ConvertAtenMinMaxDimOp(TypeConverter &typeConverter, MLIRContext *context,
                         bool allowNonFinites)
      : OpConversionPattern<OpTy>(typeConverter, context),
        allowNonFinites(allowNonFinites) {}

  LogicalResult
  matchAndRewrite(OpTy op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    static_assert(std::is_same<OpTy, AtenMaxDimOp>() ||
                  std::is_same<OpTy, AtenMinDimOp>());
    constexpr bool isMax = std::is_same<OpTy, AtenMaxDimOp>();
    const llvm::StringRef opName = op->getName().getStringRef();

    Location loc = op.getLoc();
    Value input = adaptor.getSelf();
    auto typec = this->getTypeConverter();
    auto valResultType =
        cast<RankedTensorType>(typec->convertType(op.getResult(0).getType()));
    auto idxResultType =
        cast<RankedTensorType>(typec->convertType(op.getResult(1).getType()));
    RankedTensorType inputType = cast<RankedTensorType>(input.getType());
    Type idxElementType =
        getElementTypeOrSelf(typec->convertType(idxResultType));
    if (!isa<IntegerType>(idxElementType))
      return rewriter.notifyMatchFailure(
          op, opName + " to linalg.* requires integer-like result type");

    bool keepDim = false;
    if (!matchPattern(op.getKeepdim(), m_TorchConstantBool(&keepDim)))
      return rewriter.notifyMatchFailure(
          op, opName + " requires boolean value for keepdim");

    int64_t dim;
    if (!matchPattern(op.getDim(), m_TorchConstantInt(&dim)))
      return rewriter.notifyMatchFailure(
          op, opName + " to linalg.* requires int value for Dim");
    dim = toPositiveDim(dim, inputType.getRank());
    if (!isValidDim(dim, inputType.getRank()))
      return rewriter.notifyMatchFailure(op, "dim is not a valid dim");

    Type inElementType = inputType.getElementType();
    bool isUnsigned = false;
    if (!isa<mlir::FloatType>(inElementType)) {
      if (isa<mlir::IntegerType>(inElementType)) {
        auto integerTy = dyn_cast<mlir::IntegerType>(
            cast<BaseTensorType>(op.getSelf().getType()).getDtype());
        isUnsigned = integerTy.isUnsigned();
      } else {
        return rewriter.notifyMatchFailure(
            op, opName + " to linalg.* requires Float or Integer "
                         "input element type");
      }
    }

    // Constant op to account for the reduction along dim.
    SmallVector<Value> resultShape;
    for (int64_t i = 0; i < inputType.getRank(); i++) {
      if (dim != i) {
        auto currentDimSize = tensor::DimOp::create(rewriter, loc, input, i);
        resultShape.push_back(currentDimSize);
      }
    }
    // First fill the output buffer for the index.
    Value filledTensorIdx =
        createZeroInitTensor(rewriter, loc, resultShape, idxElementType);

    // Second fill the output buffer for the running max or min.
    Value initTensorVal = tensor::EmptyOp::create(
        rewriter, loc, getAsOpFoldResult(resultShape), inElementType);

    Value fillValue;
    if (isa<mlir::FloatType>(inElementType)) {
      fillValue = arith::ConstantOp::create(
          rewriter, loc,
          rewriter.getFloatAttr(
              inElementType,
              getFloatInf(cast<mlir::FloatType>(inElementType),
                          /*Negative=*/isMax, this->allowNonFinites)));
    } else if (!isUnsigned) {
      auto width = cast<mlir::IntegerType>(inElementType).getWidth();
      auto init = isMax ? APSInt::getSignedMinValue(width)
                        : APSInt::getSignedMaxValue(width);
      fillValue = arith::ConstantOp::create(
          rewriter, loc, rewriter.getIntegerAttr(inElementType, init));
    } else if (isUnsigned) {
      auto width = cast<mlir::IntegerType>(inElementType).getWidth();
      auto init = isMax ? APInt::getMinValue(width) : APInt::getMaxValue(width);
      fillValue = arith::ConstantOp::create(
          rewriter, loc, rewriter.getIntegerAttr(inElementType, init));
    }

    Value filledTensorVal =
        linalg::FillOp::create(rewriter, loc, fillValue, initTensorVal)
            .result();

    SmallVector<utils::IteratorType> iteratorTypes(
        inputType.getRank(), utils::IteratorType::parallel);
    iteratorTypes[dim] = utils::IteratorType::reduction;

    // Create the affine expressions that will be used to
    // iterate over the input and output tensors.
    // Here we also set the type of iterator: parallel or reduction.

    SmallVector<AffineExpr> exprs;
    SmallVector<AffineExpr> resultExprs;
    for (auto size :
         llvm::enumerate(makeShapeTorchCompatible(inputType.getShape()))) {
      exprs.push_back(rewriter.getAffineDimExpr(size.index()));
      if (unsigned(dim) != size.index())
        resultExprs.push_back(rewriter.getAffineDimExpr(size.index()));
    }

    auto maps = AffineMap::inferFromExprList({exprs, resultExprs, resultExprs},
                                             rewriter.getContext());
    auto linalgOp = linalg::GenericOp::create(
        rewriter, loc,
        ArrayRef<Type>({filledTensorVal.getType(), filledTensorIdx.getType()}),
        input, ValueRange({filledTensorVal, filledTensorIdx}), maps,
        iteratorTypes,
        [&](OpBuilder &nestedBuilder, Location nestedLoc,
            ValueRange blockArgs) {
          Value newValue = blockArgs[0];
          Value oldValue = blockArgs[1];
          Value oldIndex = blockArgs[2];

          Value newIndex = arith::IndexCastOp::create(
              rewriter, nestedLoc, oldIndex.getType(),
              linalg::IndexOp::create(rewriter, loc, dim));

          Value resultVal, predicate;
          if (isa<mlir::FloatType>(inElementType)) {
            arith::CmpFPredicate predType;
            if (isMax) {
              predType = arith::CmpFPredicate::OGT;
              resultVal = arith::MaximumFOp::create(rewriter, nestedLoc,
                                                    newValue, oldValue);
            } else {
              predType = arith::CmpFPredicate::OLT;
              resultVal = arith::MinimumFOp::create(rewriter, nestedLoc,
                                                    newValue, oldValue);
            }

            predicate = arith::CmpFOp::create(rewriter, nestedLoc, predType,
                                              newValue, oldValue);
          } else {
            arith::CmpIPredicate predType;
            if (isMax) {
              predType = isUnsigned ? arith::CmpIPredicate::ugt
                                    : arith::CmpIPredicate::sgt;
              if (isUnsigned) {
                resultVal = arith::MaxUIOp::create(rewriter, nestedLoc,
                                                   newValue, oldValue);
              } else {
                resultVal = arith::MaxSIOp::create(rewriter, nestedLoc,
                                                   newValue, oldValue);
              }
            } else {
              predType = isUnsigned ? arith::CmpIPredicate::ult
                                    : arith::CmpIPredicate::slt;
              if (isUnsigned) {
                resultVal = arith::MinUIOp::create(rewriter, nestedLoc,
                                                   newValue, oldValue);
              } else {
                resultVal = arith::MinSIOp::create(rewriter, nestedLoc,
                                                   newValue, oldValue);
              }
            }
            predicate = arith::CmpIOp::create(rewriter, nestedLoc, predType,
                                              newValue, oldValue);
          }
          auto resultIndex = arith::SelectOp::create(
              rewriter, nestedLoc, predicate, newIndex, oldIndex);
          linalg::YieldOp::create(nestedBuilder, nestedLoc,
                                  ValueRange({resultVal, resultIndex}));
        });

    if (!keepDim) {
      Value rVal = tensor::CastOp::create(rewriter, loc, valResultType,
                                          linalgOp.getResult(0));
      Value rIdx = tensor::CastOp::create(rewriter, loc, idxResultType,
                                          linalgOp.getResult(1));
      llvm::SmallVector<Value> res{rVal, rIdx};
      rewriter.replaceOp(op, res);
      return success();
    }

    llvm::SmallVector<int64_t> valShape(valResultType.getShape());
    llvm::SmallVector<int64_t> idxShape(idxResultType.getShape());
    for (int i = dim, s = valShape.size() - 1; i < s; ++i) {
      valShape[i] = valShape[i + 1];
      idxShape[i] = idxShape[i + 1];
    }

    valShape.resize(valShape.size() - 1);
    idxShape.resize(idxShape.size() - 1);

    Value rVal = tensor::CastOp::create(
        rewriter, loc, valResultType.clone(valShape), linalgOp.getResult(0));
    Value rIdx = tensor::CastOp::create(
        rewriter, loc, idxResultType.clone(idxShape), linalgOp.getResult(1));

    SmallVector<ReassociationIndices> reassociation(valShape.size());
    if (reassociation.size() > 0) {
      for (int i = 0; i < dim; ++i)
        reassociation[i].push_back(i);
      reassociation[std::max<int64_t>(0, dim - 1)].push_back(dim);
      for (int i = dim, s = reassociation.size(); i < s; ++i)
        reassociation[i].push_back(i + 1);
    }

    valShape.push_back(0);
    idxShape.push_back(0);
    for (int i = dim, s = valShape.size() - 1; i < s; ++i) {
      valShape[i + 1] = valShape[i];
      idxShape[i + 1] = idxShape[i];
    }

    valShape[dim] = 1;
    idxShape[dim] = 1;

    Value unsqueezeVal = tensor::ExpandShapeOp::create(
        rewriter, loc, valResultType, rVal, reassociation);

    Value unsqueezeIdx = tensor::ExpandShapeOp::create(
        rewriter, loc, idxResultType, rIdx, reassociation);

    llvm::SmallVector<Value> unsqueezes = {unsqueezeVal, unsqueezeIdx};
    rewriter.replaceOp(op, unsqueezes);
    return success();
  }
};

} // namespace

static Value createAbsOpForNormOps(OpBuilder &b, Location loc, Value elem,
                                   Type resultElementType) {
  if (isa<mlir::ComplexType>(elem.getType())) {
    return complex::AbsOp::create(b, loc, elem);
  }

  Value self = convertScalarToDtype(b, loc, elem, resultElementType);
  return math::AbsFOp::create(b, loc, self);
}

static Value createInitElementForReduceOp(OpBuilder &b, Location loc,
                                          Operation *op, Type elementType,
                                          bool allowNonFinites) {
  if (isa<AtenSumOp, AtenSumDimIntListOp>(op))
    return arith::ConstantOp::create(b, loc, b.getZeroAttr(elementType));

  if (isa<AtenProdOp, AtenProdDimIntOp, PrimsProdOp>(op)) {
    if (isa<mlir::FloatType>(elementType))
      return arith::ConstantOp::create(b, loc,
                                       b.getFloatAttr(elementType, 1.0));
    else if (isa<mlir::IntegerType>(elementType))
      return arith::ConstantOp::create(b, loc,
                                       b.getIntegerAttr(elementType, 1));
  }

  if (isa<AtenMaxOp>(op)) {
    if (isa<mlir::FloatType>(elementType))
      return arith::ConstantOp::create(
          b, loc,
          b.getFloatAttr(elementType,
                         getFloatInf(cast<mlir::FloatType>(elementType),
                                     /*Negative=*/true, allowNonFinites)));
    else if (isa<mlir::IntegerType>(elementType) &&
             elementType.getIntOrFloatBitWidth() != 8)
      return arith::ConstantOp::create(
          b, loc,
          b.getIntegerAttr(
              elementType,
              APSInt::getSignedMinValue(elementType.getIntOrFloatBitWidth())));
  }

  if (isa<AtenMinOp>(op)) {
    if (isa<mlir::FloatType>(elementType))
      return arith::ConstantOp::create(
          b, loc,
          b.getFloatAttr(elementType,
                         getFloatInf(cast<mlir::FloatType>(elementType),
                                     /*Negative=*/false, allowNonFinites)));
    else if (isa<mlir::IntegerType>(elementType) &&
             elementType.getIntOrFloatBitWidth() != 8)
      return arith::ConstantOp::create(
          b, loc,
          b.getIntegerAttr(
              elementType,
              APSInt::getSignedMaxValue(elementType.getIntOrFloatBitWidth())));
  }

  if (isa<AtenLinalgVectorNormOp>(op) || isa<AtenFrobeniusNormDimOp>(op) ||
      isa<AtenNormScalarOp>(op))
    return arith::ConstantOp::create(b, loc, b.getZeroAttr(elementType));

  if (isa<AtenAllOp, AtenAllDimOp>(op)) {
    return arith::ConstantOp::create(b, loc, b.getBoolAttr(true));
  }

  if (isa<AtenAnyOp, AtenAnyDimsOp, AtenAnyDimOp>(op)) {
    return arith::ConstantOp::create(b, loc, b.getBoolAttr(false));
  }

  op->emitError("unimplemented lowering in createInitElementForReduceOp");
  return nullptr;
}

static Value createLinalgPayloadForReduceOp(OpBuilder &b, Location loc,
                                            ValueRange payloadArgs,
                                            Operation *op,
                                            ArrayRef<Value> operands,
                                            Type resultElementType) {
  if (isa<AtenSumOp, AtenSumDimIntListOp>(op)) {
    Value self =
        convertScalarToDtype(b, loc, payloadArgs[0], resultElementType);
    Value result = payloadArgs[1];
    if (isa<mlir::FloatType>(resultElementType))
      return arith::AddFOp::create(b, loc, self, result);
    else if (isa<mlir::IntegerType>(resultElementType))
      return arith::AddIOp::create(b, loc, self, result);
  } else if (isa<AtenProdOp, AtenProdDimIntOp, PrimsProdOp>(op)) {
    Value self =
        convertScalarToDtype(b, loc, payloadArgs[0], resultElementType);
    Value result = payloadArgs[1];
    if (isa<mlir::FloatType>(resultElementType))
      return arith::MulFOp::create(b, loc, self, result);
    else if (isa<mlir::IntegerType>(resultElementType))
      return arith::MulIOp::create(b, loc, self, result);
  } else if (auto max = dyn_cast<AtenMaxOp>(op)) {
    Value self =
        convertScalarToDtype(b, loc, payloadArgs[0], resultElementType);
    Value result = payloadArgs[1];
    if (isa<mlir::FloatType>(resultElementType))
      return arith::MaximumFOp::create(b, loc, self, result);
    else if (isa<mlir::IntegerType>(resultElementType)) {
      IntegerType intType = dyn_cast<mlir::IntegerType>(
          cast<BaseTensorType>(max.getSelf().getType()).getDtype());
      if (intType.isUnsigned())
        return arith::MaxUIOp::create(b, loc, self, result);
      if (intType.isSigned())
        return arith::MaxSIOp::create(b, loc, self, result);
    }
  } else if (auto min = dyn_cast<AtenMinOp>(op)) {
    Value self =
        convertScalarToDtype(b, loc, payloadArgs[0], resultElementType);
    Value result = payloadArgs[1];
    if (isa<mlir::FloatType>(resultElementType))
      return arith::MinimumFOp::create(b, loc, self, result);
    else if (isa<mlir::IntegerType>(resultElementType)) {
      IntegerType intType = dyn_cast<mlir::IntegerType>(
          cast<BaseTensorType>(min.getSelf().getType()).getDtype());
      if (intType.isUnsigned())
        return arith::MinUIOp::create(b, loc, self, result);
      if (intType.isSigned())
        return arith::MinSIOp::create(b, loc, self, result);
    }
  } else if (isa<AtenNormScalarOp>(op)) {
    // This creates payload for only the first of the two linalg.generic ops.
    // TODO: Short-circuit operations if `p` is zero or one.
    Value elem = payloadArgs[0];
    Value result = payloadArgs[1];

    AtenNormScalarOp::Adaptor adaptor(operands);
    Value p = convertScalarToDtype(b, loc, adaptor.getP(), resultElementType);

    auto abs = createAbsOpForNormOps(b, loc, elem, resultElementType);
    auto pow = math::PowFOp::create(b, loc, abs, p);
    return arith::AddFOp::create(b, loc, pow, result);
  } else if (isa<AtenLinalgVectorNormOp>(op)) {
    // This creates payload for only the first of the two linalg.generic ops.
    // TODO: Short-circuit operations if `ord` is zero or one.
    Value elem = payloadArgs[0];
    Value result = payloadArgs[1];

    AtenLinalgVectorNormOp::Adaptor adaptor(operands);
    Value ord =
        convertScalarToDtype(b, loc, adaptor.getOrd(), resultElementType);

    auto abs = createAbsOpForNormOps(b, loc, elem, resultElementType);
    auto pow = math::PowFOp::create(b, loc, abs, ord);
    return arith::AddFOp::create(b, loc, pow, result);
  } else if (isa<AtenFrobeniusNormDimOp>(op)) {
    Value elem = payloadArgs[0];
    Value result = payloadArgs[1];

    TypedAttr twoAttr = b.getFloatAttr(resultElementType, 2.0);
    auto ord = arith::ConstantOp::create(b, loc, twoAttr);

    auto abs = createAbsOpForNormOps(b, loc, elem, resultElementType);
    auto pow = math::PowFOp::create(b, loc, abs, ord);
    return arith::AddFOp::create(b, loc, pow, result);
  } else if (isa<AtenAllOp, AtenAllDimOp>(op)) {
    Value elem = payloadArgs[0];
    Value result = payloadArgs[1];
    Value self = convertScalarToDtype(b, loc, elem, resultElementType);
    return arith::AndIOp::create(b, loc, self, result);
  } else if (isa<AtenAnyOp, AtenAnyDimsOp, AtenAnyDimOp>(op)) {
    Value elem = payloadArgs[0];
    Value result = payloadArgs[1];
    Value self = convertScalarToDtype(b, loc, elem, resultElementType);
    return arith::OrIOp::create(b, loc, self, result);
  }
  op->emitError("unimplemented lowering in createLinalgPayloadForReduceOp");
  return nullptr;
}

namespace {
class ConvertReductionOp : public ConversionPattern {
private:
  bool allowNonFinites;

  /// Given a reduction operation that has the `keepdim` attribute and the
  /// (optional) `dim` attribute, return the source tensor operand and the
  /// literal values of the attributes or failure otherwise.
  template <typename T>
  FailureOr<torch_to_linalg::ReductionOpInfo>
  computeReductionOpInfoForDimVariantOp(
      T op, ArrayRef<Value> operands,
      ConversionPatternRewriter &rewriter) const {
    auto opInfo = torch_to_linalg::ReductionOpInfo{false, Value{}, {}};
    typename T::Adaptor adaptor(operands);
    opInfo.tensorOperand = adaptor.getSelf();
    auto inputType = cast<RankedTensorType>(opInfo.tensorOperand.getType());

    if (!matchPattern(op.getKeepdim(), m_TorchConstantBool(&opInfo.keepDim)))
      return rewriter.notifyMatchFailure(op,
                                         "`keepdim` must be a constant bool");

    SmallVector<int64_t> dimList;
    int64_t dim;
    bool isNoneOrEmptyDimList = isa<Torch::NoneType>(op.getDim().getType());
    if (matchPattern(op.getDim(), m_TorchListOfConstantInts(dimList))) {
      // Fix negative dimensions, if any, before adding to the list.
      for (int64_t dim : dimList) {
        dim = toPositiveDim(dim, inputType.getRank());
        // Drop invalid dimensions
        if (isValidDim(dim, inputType.getRank()))
          opInfo.dimSet.insert(dim);
      }
      if (dimList.empty())
        isNoneOrEmptyDimList = true;
    } else if (matchPattern(op.getDim(), m_TorchConstantInt(&dim))) {
      dim = toPositiveDim(dim, inputType.getRank());
      if (!isValidDim(dim, inputType.getRank()))
        return rewriter.notifyMatchFailure(
            op, "`dim` argument must be valid, invalid received.");
      opInfo.dimSet.insert(dim);
    } else if (!isNoneOrEmptyDimList) {
      return rewriter.notifyMatchFailure(
          op, "`dim` argument must be a constant int list or None");
    }
    if (isNoneOrEmptyDimList) {
      // If no dimensions were specified, reduce along all dimensions
      for (int64_t i = 0; i < inputType.getRank(); i++)
        opInfo.dimSet.insert(i);
    }

    return opInfo;
  }

  /// Compute reduction info from a tensor operand and an optional Torch list
  /// of constant ints describing the dims to reduce. `keepdim` is inferred
  /// by comparing the rank of the op's result type to the rank of the input
  /// tensor: if they match, the reduced dimensions are kept (size 1). Used
  /// for ops (e.g. `prims.prod`) that lack a `keepdim` operand and use
  /// accessor names that differ from the dim-variant template.
  ///
  /// In typical pytorch-imported programs, `keepdim` is always implicitly
  /// `false` for prims ops; however, we allow intentionally diverging from
  /// this (e.g. in ONNX or IR-emitting situations) for convenience. In
  /// this case, we rely on the result rank to determine `keepdim`.
  FailureOr<torch_to_linalg::ReductionOpInfo>
  computeReductionOpInfoFromDimList(Operation *op, Value tensorOperand,
                                    Value dimsValue,
                                    ConversionPatternRewriter &rewriter) const {
    auto opInfo = torch_to_linalg::ReductionOpInfo{false, tensorOperand, {}};
    auto inputType = cast<RankedTensorType>(tensorOperand.getType());

    SmallVector<int64_t> dimList;
    bool isNoneOrEmptyDimList = isa<Torch::NoneType>(dimsValue.getType());
    if (matchPattern(dimsValue, m_TorchListOfConstantInts(dimList))) {
      for (int64_t dim : dimList) {
        dim = toPositiveDim(dim, inputType.getRank());
        if (isValidDim(dim, inputType.getRank()))
          opInfo.dimSet.insert(dim);
      }
      if (dimList.empty())
        isNoneOrEmptyDimList = true;
    } else if (!isNoneOrEmptyDimList) {
      return rewriter.notifyMatchFailure(
          op, "`dims` argument must be a constant int list or None");
    }
    if (isNoneOrEmptyDimList) {
      for (int64_t i = 0; i < inputType.getRank(); i++)
        opInfo.dimSet.insert(i);
    }

    // Infer `keepdim` from the result type's rank. prims ops do not carry
    // a `keepdim` operand, so the only signal is the shape of the op's
    // declared result. The result must exactly match the expected reduced
    // shape: reduced dimensions are size 1 (keepdim) or dropped (no
    // keepdim), and non-reduced dimensions match the input.
    auto resultTensorType =
        cast<Torch::BaseTensorType>(op->getResult(0).getType());
    if (!resultTensorType.hasSizes())
      return rewriter.notifyMatchFailure(op,
                                         "result tensor must have known sizes");

    int64_t inputRank = inputType.getRank();
    ArrayRef<int64_t> inputSizes = inputType.getShape();
    ArrayRef<int64_t> resultSizes = resultTensorType.getSizes();
    int64_t resultRank = static_cast<int64_t>(resultSizes.size());

    int64_t numReduced = static_cast<int64_t>(opInfo.dimSet.size());
    if (resultRank == inputRank) {
      opInfo.keepDim = true;
      for (int64_t i = 0; i < inputRank; ++i) {
        int64_t expected = opInfo.dimSet.contains(i) ? 1 : inputSizes[i];
        if (resultSizes[i] != expected)
          return rewriter.notifyMatchFailure(
              op, "result shape does not match expected reduced shape");
      }
    } else if (resultRank == inputRank - numReduced) {
      int64_t j = 0;
      for (int64_t i = 0; i < inputRank; ++i) {
        if (opInfo.dimSet.contains(i))
          continue;
        if (resultSizes[j] != inputSizes[i])
          return rewriter.notifyMatchFailure(
              op, "result shape does not match expected reduced shape");
        ++j;
      }
    } else {
      return rewriter.notifyMatchFailure(
          op, "result rank does not match input rank (keepdim) nor "
              "input rank minus the number of reduced dims");
    }

    return opInfo;
  }

  /// Given a reduction operation, return the source tensor operand and the
  /// literal values of the `keepdim` and `dim` attributes, if any, or failure
  /// otherwise.
  FailureOr<torch_to_linalg::ReductionOpInfo>
  computeReductionOpInfo(Operation *op, ArrayRef<Value> operands,
                         ConversionPatternRewriter &rewriter) const {
    auto opInfo = torch_to_linalg::ReductionOpInfo{false, Value{}, {}};

    if (isa<AtenAnyOp, AtenAllOp, AtenMaxOp, AtenMinOp, AtenSumOp, AtenProdOp,
            AtenNormScalarOp>(op)) {
      opInfo.tensorOperand = operands[0];
      auto inputType = cast<RankedTensorType>(opInfo.tensorOperand.getType());

      // `AtenAny`, `AtenAll`, `AtenSumOp`, `AtenProdOp`, `AtenMaxOp`, and
      // `AtenMinOp` each reduce along all the dimensions of the input tensor.
      for (int64_t i = 0; i < inputType.getRank(); i++)
        opInfo.dimSet.insert(i);

      return opInfo;
    }

    if (auto sumOp = dyn_cast<AtenSumDimIntListOp>(op))
      return computeReductionOpInfoForDimVariantOp(sumOp, operands, rewriter);

    if (auto prodOp = dyn_cast<AtenProdDimIntOp>(op))
      return computeReductionOpInfoForDimVariantOp(prodOp, operands, rewriter);

    if (auto normOp = dyn_cast<AtenLinalgVectorNormOp>(op))
      return computeReductionOpInfoForDimVariantOp(normOp, operands, rewriter);

    if (auto normOp = dyn_cast<AtenFrobeniusNormDimOp>(op))
      return computeReductionOpInfoForDimVariantOp(normOp, operands, rewriter);

    if (auto allOp = dyn_cast<AtenAllDimOp>(op))
      return computeReductionOpInfoForDimVariantOp(allOp, operands, rewriter);

    if (auto anyOp = dyn_cast<AtenAnyDimsOp>(op))
      return computeReductionOpInfoForDimVariantOp(anyOp, operands, rewriter);

    if (auto anyDimOp = dyn_cast<AtenAnyDimOp>(op))
      return computeReductionOpInfoForDimVariantOp(anyDimOp, operands,
                                                   rewriter);

    if (auto primsProdOp = dyn_cast<PrimsProdOp>(op)) {
      PrimsProdOp::Adaptor adaptor(operands);
      return computeReductionOpInfoFromDimList(op, adaptor.getInp(),
                                               primsProdOp.getDims(), rewriter);
    }

    return rewriter.notifyMatchFailure(op, "not a supported reduce op");
  }

  /// Generate a linalg.generic operation for pointwise exponentiation of each
  /// element.
  Value createElementwiseExp(Location loc, Type elemType, Value exponent,
                             Value inputTensor,
                             const torch_to_linalg::ReductionOpInfo &opInfo,
                             ConversionPatternRewriter &rewriter) const {
    bool err = false;
    auto powBodyBuilder = [&](OpBuilder &builder, Location loc,
                              ValueRange payloadArgs) {
      Value elem = convertScalarToDtype(builder, loc, payloadArgs[0], elemType);
      auto result = math::PowFOp::create(builder, loc, elem, exponent);
      if (result)
        linalg::YieldOp::create(builder, loc, Value{result});
      err = !result;
    };

    Value powOp = torch_to_linalg::createElementwiseLinalgGeneric(
        rewriter, loc, {inputTensor}, elemType, powBodyBuilder);
    return err ? Value{} : powOp;
  }

  template <typename TOp>
  FailureOr<Value>
  createSecondReductionForNormOp(Location loc, Type elemType, TOp op,
                                 Value ordOp, Value firstReduction,
                                 const torch_to_linalg::ReductionOpInfo &opInfo,
                                 ConversionPatternRewriter &rewriter) const {
    // Cast `ord` to float so that we can readily pass it math.powf.
    Value ordValue = convertScalarToDtype(rewriter, loc, ordOp, elemType);

    // TODO: Add support for ord = {0, +inf, -inf}.
    auto epsilon = 1e-5;
    auto ordLiteral = 0.0;
    if (matchPattern(ordValue, m_TorchConstantFloat(&ordLiteral)) &&
        fabs(ordLiteral) < epsilon)
      return rewriter.notifyMatchFailure(op, "unimplemented: L0 norm");

    if (std::isinf(ordLiteral))
      return rewriter.notifyMatchFailure(op, "unimplemented: ord = +/- inf");

    // Raise each summed value to the inverse of the order of the norm.
    TypedAttr oneAttr = rewriter.getFloatAttr(elemType, 1.0);
    auto oneValue = arith::ConstantOp::create(rewriter, loc, oneAttr);
    auto inverseOrdValue =
        arith::DivFOp::create(rewriter, loc, oneValue, ordValue);

    // Use the results of the first reduction operation from above to generate
    // a second reduction operation.
    Value reduceOp = createElementwiseExp(loc, elemType, inverseOrdValue,
                                          firstReduction, opInfo, rewriter);
    if (!reduceOp)
      return rewriter.notifyMatchFailure(
          op, "failed to create linalg.generic operation for element-wise "
              "exponentiation");

    return reduceOp;
  }

  /// Generate a linalg.generic operation for a reduction.
  Value createReductionOp(Location loc, Type elemType, Operation *op,
                          ArrayRef<Value> operands,
                          const torch_to_linalg::ReductionOpInfo &opInfo,
                          ConversionPatternRewriter &rewriter) const {
    bool err = false;
    auto reductionBodyBuilder = [&](OpBuilder &builder, Location loc,
                                    ValueRange payloadArgs) {
      Value result = createLinalgPayloadForReduceOp(builder, loc, payloadArgs,
                                                    op, operands, elemType);
      if (result)
        linalg::YieldOp::create(builder, loc, result);
      err = !result;
    };

    Value initElem = createInitElementForReduceOp(rewriter, loc, op, elemType,
                                                  this->allowNonFinites);
    Value reduceOp = torch_to_linalg::createReductionLinalgGeneric(
        rewriter, loc, opInfo, initElem, reductionBodyBuilder);
    return err ? Value{} : reduceOp;
  }

  /// Depending on the operation, check validity of the result's element type.
  LogicalResult
  validateReductionElementType(Operation *op, Type elemType,
                               ConversionPatternRewriter &rewriter) const {
    if ((isa<AtenLinalgVectorNormOp>(op) || isa<AtenFrobeniusNormDimOp>(op) ||
         isa<AtenNormScalarOp>(op)) &&
        !isa<mlir::FloatType>(elemType))
      return rewriter.notifyMatchFailure(
          op, "only float types are valid for vector norm ops");
    if (isa<AtenAllDimOp>(op) && isa<mlir::IntegerType>(elemType) &&
        elemType.getIntOrFloatBitWidth() == 8)
      return rewriter.notifyMatchFailure(op, "uint8 is not supported");

    // No checks for all other reduction operations
    return success();
  }

public:
  ConvertReductionOp(TypeConverter &typeConverter, MLIRContext *context,
                     bool allowNonFinites)
      : ConversionPattern(typeConverter, MatchAnyOpTypeTag(), /*benefit=*/1,
                          context),
        allowNonFinites(allowNonFinites) {}

  LogicalResult
  matchAndRewrite(Operation *op, ArrayRef<Value> operands,
                  ConversionPatternRewriter &rewriter) const override {
    if (failed(verifyLinalgCompatibleTypes(op, rewriter)))
      return rewriter.notifyMatchFailure(
          op, "invalid operand or result types to use with linalg on tensors");

    FailureOr<torch_to_linalg::ReductionOpInfo> opInfo =
        computeReductionOpInfo(op, operands, rewriter);
    if (failed(opInfo))
      return opInfo;

    Location loc = op->getLoc();
    auto resultType = cast<RankedTensorType>(
        getTypeConverter()->convertType(op->getResult(0).getType()));
    Type elemType = resultType.getElementType();
    LogicalResult elemTypeCheck =
        validateReductionElementType(op, elemType, rewriter);
    if (failed(elemTypeCheck))
      return elemTypeCheck;

    Value reduceOp =
        createReductionOp(loc, elemType, op, operands, *opInfo, rewriter);
    if (!reduceOp)
      return rewriter.notifyMatchFailure(
          op, "failed to create linalg.generic operation for reduction");

    // If this is aten.norm.Scalar op, then we need to generate another
    // linalg.generic op that references the first linalg.generic op.
    if (isa<AtenNormScalarOp>(op)) {
      AtenNormScalarOp::Adaptor adaptor(operands);
      FailureOr<Value> secondReduceOp = createSecondReductionForNormOp(
          loc, elemType, op, adaptor.getP(), reduceOp, *opInfo, rewriter);
      if (failed(secondReduceOp))
        return secondReduceOp;
      reduceOp = *secondReduceOp;
    }

    // If this is aten.linalg_vector_norm op, then we need to generate another
    // linalg.generic op that references the first linalg.generic op.
    if (auto normOp = dyn_cast<AtenLinalgVectorNormOp>(op)) {
      AtenLinalgVectorNormOp::Adaptor adaptor(operands);
      FailureOr<Value> secondReduceOp = createSecondReductionForNormOp(
          loc, elemType, normOp, adaptor.getOrd(), reduceOp, *opInfo, rewriter);
      if (failed(secondReduceOp))
        return secondReduceOp;
      reduceOp = *secondReduceOp;
    }

    // If it is aten.frobenius_norm.dim op, take the square root of reduceOp as
    // the final result
    if (auto normOp = dyn_cast<AtenFrobeniusNormDimOp>(op)) {
      auto halfAttr = rewriter.getFloatAttr(elemType, 0.5);
      auto exp = arith::ConstantOp::create(rewriter, loc, halfAttr);
      reduceOp =
          createElementwiseExp(loc, elemType, exp, reduceOp, *opInfo, rewriter);
    }

    rewriter.replaceOpWithNewOp<tensor::CastOp>(op, resultType, reduceOp);
    return success();
  }
};
} // namespace

namespace {
class ConvertAtenCumsumOp : public OpConversionPattern<AtenCumsumOp> {
public:
  using OpConversionPattern::OpConversionPattern;

  LogicalResult
  matchAndRewrite(AtenCumsumOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Location loc = op.getLoc();
    Value self = adaptor.getSelf();
    auto selfType = dyn_cast<RankedTensorType>(self.getType());
    if (!selfType || !selfType.hasStaticShape())
      return rewriter.notifyMatchFailure(op,
                                         "only static tensor shapes supported");

    int64_t dim;
    if (!matchPattern(op.getDim(), m_TorchConstantInt(&dim)))
      return rewriter.notifyMatchFailure(op, "dim must be constant");
    dim = toPositiveDim(dim, selfType.getRank());
    if (!isValidDim(dim, selfType.getRank()))
      return rewriter.notifyMatchFailure(op, "dim is not a valid dim");

    auto resultType = dyn_cast<RankedTensorType>(
        getTypeConverter()->convertType(op.getType()));
    if (!resultType || !resultType.hasStaticShape())
      return rewriter.notifyMatchFailure(op, "expected static ranked result");

    Type resultElementType = resultType.getElementType();
    if (!isa<mlir::FloatType, mlir::IntegerType>(resultElementType))
      return rewriter.notifyMatchFailure(
          op, "only floating point and integer element types supported");

    if (selfType.getElementType() != resultElementType)
      self = torch_to_linalg::convertTensorToElementType(rewriter, loc, self,
                                                         resultElementType);

    SmallVector<int64_t> inputShape =
        makeShapeTorchCompatible(selfType.getShape());
    int64_t scanDimSize = inputShape[dim];

    int64_t outer = 1;
    for (int64_t i = 0; i < dim; ++i)
      outer *= inputShape[i];
    int64_t inner = 1;
    for (int64_t i = dim + 1, e = inputShape.size(); i < e; ++i)
      inner *= inputShape[i];

    SmallVector<int64_t> scanShape = {outer, scanDimSize, inner};
    auto scanType = RankedTensorType::get(makeShapeLLVMCompatible(scanShape),
                                          resultElementType);
    auto shapeType = RankedTensorType::get({3}, rewriter.getIntegerType(64));
    SmallVector<Value> scanShapeValues;
    for (int64_t size : scanShape) {
      scanShapeValues.push_back(arith::ConstantOp::create(
          rewriter, loc, rewriter.getI64IntegerAttr(size)));
    }
    Value scanShapeTensor = tensor::FromElementsOp::create(
        rewriter, loc, shapeType, scanShapeValues);
    Value running = tensor::ReshapeOp::create(rewriter, loc, scanType, self,
                                              scanShapeTensor)
                        .getResult();

    running = emitInclusiveScanByPowersOfTwo(running, rewriter, loc);

    auto resultShapeType =
        RankedTensorType::get({resultType.getRank()}, rewriter.getI64Type());
    SmallVector<Value> resultShapeValues;
    for (int64_t size : resultType.getShape()) {
      resultShapeValues.push_back(arith::ConstantOp::create(
          rewriter, loc, rewriter.getI64IntegerAttr(size)));
    }
    Value resultShapeTensor = tensor::FromElementsOp::create(
        rewriter, loc, resultShapeType, resultShapeValues);
    Value result = tensor::ReshapeOp::create(rewriter, loc, resultType, running,
                                             resultShapeTensor)
                       .getResult();

    rewriter.replaceOp(op, result);
    return success();
  }
};
} // namespace

void mlir::torch::torch_to_linalg::populateReductionPatternsAndLegality(
    TypeConverter &typeConverter, RewritePatternSet &patterns,
    ConversionTarget &target, bool allowNonFinites) {
  MLIRContext *context = patterns.getContext();
  target.addIllegalOp<AtenMaxDimOp>();
  patterns.add<ConvertAtenMinMaxDimOp<AtenMaxDimOp>>(typeConverter, context,
                                                     allowNonFinites);
  target.addIllegalOp<AtenMinDimOp>();
  patterns.add<ConvertAtenMinMaxDimOp<AtenMinDimOp>>(typeConverter, context,
                                                     allowNonFinites);
  target.addIllegalOp<AtenSumOp>();
  target.addIllegalOp<AtenAnyOp>();
  target.addIllegalOp<AtenAnyDimOp>();
  target.addIllegalOp<AtenAnyDimsOp>();
  target.addIllegalOp<AtenAllOp>();
  target.addIllegalOp<AtenSumDimIntListOp>();
  target.addIllegalOp<AtenProdOp>();
  target.addIllegalOp<AtenProdDimIntOp>();
  target.addIllegalOp<PrimsProdOp>();
  target.addIllegalOp<AtenMaxOp>();
  target.addIllegalOp<AtenMinOp>();
  target.addIllegalOp<AtenAllDimOp>();
  target.addIllegalOp<AtenNormScalarOp>();
  target.addIllegalOp<AtenLinalgVectorNormOp>();
  target.addIllegalOp<AtenFrobeniusNormDimOp>();
  target.addIllegalOp<AtenCumsumOp>();
  patterns.add<ConvertReductionOp>(typeConverter, context, allowNonFinites);
  patterns.add<ConvertAtenCumsumOp>(typeConverter, context);
}