[mlir][tosa] Constant optimizations for reduce operations

Replace the different reduce operations which is getting a constant tensor as an input argument with a constant tensor. As the arguement of the reduce operation is constant tensor and has only a single user we could calculate the resulted constant tensor in compilation time and replace it with reduced memory tensor This optimization has been implemented for: tosa.reduce_sum tosa.reduce_prod tosa.reduce_any tosa.reduce_all tosa.reduce_max tosa.reduce_min Reviewed By: rsuderman Differential Revision: https://reviews.llvm.org/D154832

Author: amirBish

mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td

@@ -1273,6 +1273,11 @@ def Tosa_ReduceAllOp : Tosa_InferTensorTypeOp<"reduce_all"> {
  /// Returns true when two result types are compatible for this op;
  /// Method used by InferTypeOpInterface.
  static bool isCompatibleReturnTypes(TypeRange l, TypeRange r);
+
+ /// Return the AND result between two integer operands
+ static inline APInt calcOneElement(APInt leftOperand, APInt rightOperand) {
+ return leftOperand & rightOperand;
+ }
  }];
 }
 
@@ -1301,6 +1306,11 @@ def Tosa_ReduceAnyOp : Tosa_InferTensorTypeOp<"reduce_any"> {
  /// Returns true when two result types are compatible for this op;
  /// Method used by InferTypeOpInterface.
  static bool isCompatibleReturnTypes(TypeRange l, TypeRange r);
+
+ /// Return the OR result between two integer operands
+ static inline APInt calcOneElement(APInt leftOperand, APInt rightOperand) {
+ return leftOperand | rightOperand;
+ }
  }];
 }
 
@@ -1329,6 +1339,12 @@ def Tosa_ReduceMaxOp : Tosa_InferTensorTypeOp<"reduce_max"> {
  /// Returns true when two result types are compatible for this op;
  /// Method used by InferTypeOpInterface.
  static bool isCompatibleReturnTypes(TypeRange l, TypeRange r);
+
+ /// Return the max of the two integer operands
+ static inline APInt calcOneElement(APInt leftOperand, APInt rightOperand) {
+ const llvm::APInt subtractRes = leftOperand - rightOperand;
+ return (!subtractRes.isNegative()) ? leftOperand : rightOperand;
+ }
  }];
 }
 
@@ -1357,6 +1373,12 @@ def Tosa_ReduceMinOp : Tosa_InferTensorTypeOp<"reduce_min"> {
  /// Returns true when two result types are compatible for this op;
  /// Method used by InferTypeOpInterface.
  static bool isCompatibleReturnTypes(TypeRange l, TypeRange r);
+
+ /// Return the min of the two integer operands
+ static inline APInt calcOneElement(APInt leftOperand, APInt rightOperand) {
+ const llvm::APInt subtractRes = leftOperand - rightOperand;
+ return (!subtractRes.isNegative()) ? rightOperand : leftOperand;
+ }
  }];
 }
 
@@ -1385,6 +1407,11 @@ def Tosa_ReduceProdOp : Tosa_InferTensorTypeOp<"reduce_prod"> {
  /// Returns true when two result types are compatible for this op;
  /// Method used by InferTypeOpInterface.
  static bool isCompatibleReturnTypes(TypeRange l, TypeRange r);
+
+ /// Return the prod of the two integer operands
+ static inline APInt calcOneElement(APInt leftOperand, APInt rightOperand) {
+ return leftOperand * rightOperand;
+ }
  }];
 }
 
@@ -1406,13 +1433,17 @@ def Tosa_ReduceSumOp : Tosa_InferTensorTypeOp<"reduce_sum"> {
  let results = (outs
  Tosa_Tensor:$output
  );
-
  let hasFolder = 1;
 
  let extraClassDeclaration = [{
  /// Returns true when two result types are compatible for this op;
  /// Method used by InferTypeOpInterface.
  static bool isCompatibleReturnTypes(TypeRange l, TypeRange r);
+
+ /// Return the sum of the two integer operands
+ static inline APInt calcOneElement(APInt leftOperand, APInt rightOperand) {
+ return leftOperand + rightOperand;
+ }
  }];
 }
 

mlir/include/mlir/Dialect/Tosa/Transforms/Passes.h

@@ -34,6 +34,8 @@ void populateTosaFoldConstantReciprocalPatterns(MLIRContext *ctx,
  RewritePatternSet &patterns);
 void populateTosaFoldConstantTransposePatterns(MLIRContext *ctx,
  RewritePatternSet &patterns);
+void populateTosaConstantReduction(MLIRContext *ctx,
+ RewritePatternSet &patterns);
 
 std::unique_ptr<Pass> createTosaLayerwiseConstantFoldPass();
 std::unique_ptr<Pass> createTosaInferShapesPass();

mlir/lib/Dialect/Tosa/Transforms/TosaFolders.cpp

@@ -11,6 +11,7 @@
 //===----------------------------------------------------------------------===//
 
 #include <functional>
+#include <numeric>
 
 #include "mlir/Dialect/Tosa/IR/TosaOps.h"
 #include "mlir/Dialect/Tosa/Transforms/Passes.h"
@@ -289,8 +290,130 @@ struct TosaFoldConstantReciprocal : public OpRewritePattern<ReciprocalOp> {
  }
 };
 
+/// Getting the axes position of the element which is located
+/// in the tensor at the counter index
+
+llvm::SmallVector<int64_t>
+getPositionFromIndex(int64_t index, llvm::ArrayRef<int64_t> tensorShape) {
+ int64_t remaining = index;
+ llvm::SmallVector<int64_t> position(tensorShape.size(), 0);
+ for (int64_t i = tensorShape.size() - 1; i >= 0; --i) {
+ position[i] = remaining % tensorShape[i];
+ remaining /= tensorShape[i];
+ }
+ return position;
+}
+
+/// Getting the index of the element which is located at the
+/// axes position in the tensor
+
+int64_t getIndexFromPosition(llvm::ArrayRef<int64_t> position,
+ llvm::ArrayRef<int64_t> tensorShape) {
+ int64_t index = 0;
+ int64_t multiplierTmp = 1;
+ for (int64_t i = position.size() - 1; i >= 0; --i) {
+ index += position[i] * multiplierTmp;
+ multiplierTmp *= tensorShape[i];
+ }
+ return index;
+}
+
+template <typename OperationType>
+llvm::APInt calculateReducedValue(const mlir::ElementsAttr &oldTensorAttr,
+ llvm::ArrayRef<int64_t> oldShape,
+ int64_t reductionAxis,
+ int64_t reductionIndex) {
+
+ llvm::SmallVector<int64_t> newShape(oldShape);
+ newShape[reductionAxis] = 1;
+ /// Let's calculate the position of the index
+ llvm::SmallVector<int64_t> position =
+ getPositionFromIndex(reductionIndex, newShape);
+ auto oldTensor = oldTensorAttr.getValues<llvm::APInt>();
+ /// Starting from the first positon along the reduction axis
+ position[reductionAxis] = 0;
+ int64_t indexAtOldTensor = getIndexFromPosition(position, oldShape);
+ llvm::APInt reducedValue = oldTensor[indexAtOldTensor];
+
+ for (int64_t reductionAxisVal = 1; reductionAxisVal < oldShape[reductionAxis];
+ ++reductionAxisVal) {
+
+ int64_t stride = std::accumulate(oldShape.begin() + reductionAxis + 1,
+ oldShape.end(), 1, std::multiplies<int>());
+ int64_t index = indexAtOldTensor + stride * reductionAxisVal;
+ reducedValue =
+ OperationType::calcOneElement(reducedValue, oldTensor[index]);
+ }
+ return reducedValue;
+}
+
+template <typename OperationType>
+struct ReduceConstantOptimization : public OpRewritePattern<OperationType> {
+
+ using OpRewritePattern<OperationType>::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(OperationType op,
+ PatternRewriter &rewriter) const override {
+ Value inputOp = op.getInput();
+ auto constOp = inputOp.getDefiningOp<tosa::ConstOp>();
+
+ if (!constOp)
+ return rewriter.notifyMatchFailure(
+ op, "reduce input must be const operation");
+
+ if (!inputOp.hasOneUse())
+ return rewriter.notifyMatchFailure(
+ op, "input operation has more than one user");
+
+ auto resultType = cast<ShapedType>(op.getOutput().getType());
+
+ if (!resultType.hasStaticShape())
+ return rewriter.notifyMatchFailure(op, "result type shape is not static");
+
+ auto reductionAxis = op.getAxis();
+ const auto denseElementsAttr = constOp.getValue();
+ const auto shapedOldElementsValues =
+ denseElementsAttr.getType().cast<ShapedType>();
+
+ if (!llvm::isa<IntegerType>(shapedOldElementsValues.getElementType()))
+ return rewriter.notifyMatchFailure(
+ op, "reduce input currently supported with integer type");
+
+ auto oldShape = shapedOldElementsValues.getShape();
+ auto newShape = resultType.getShape();
+
+ auto newNumOfElements = std::accumulate(newShape.begin(), newShape.end(), 1,
+ std::multiplies<int>());
+ llvm::SmallVector<APInt> newReducedTensor(newNumOfElements);
+
+ for (int64_t reductionIndex = 0; reductionIndex < newNumOfElements;
+ ++reductionIndex) {
+
+ /// Let's reduce all the elements along this reduction axis
+ newReducedTensor[reductionIndex] = calculateReducedValue<OperationType>(
+ denseElementsAttr, oldShape, reductionAxis, reductionIndex);
+ }
+
+ auto rankedTensorType = cast<RankedTensorType>(resultType);
+ auto denseAttr =
+ mlir::DenseElementsAttr::get(rankedTensorType, newReducedTensor);
+ rewriter.replaceOpWithNewOp<tosa::ConstOp>(op, rankedTensorType, denseAttr);
+ return success();
+ }
+};
+
 } // namespace
 
+void mlir::tosa::populateTosaConstantReduction(MLIRContext *ctx,
+ RewritePatternSet &patterns) {
+ patterns.add<ReduceConstantOptimization<ReduceAllOp>>(ctx);
+ patterns.add<ReduceConstantOptimization<ReduceAnyOp>>(ctx);
+ patterns.add<ReduceConstantOptimization<ReduceMaxOp>>(ctx);
+ patterns.add<ReduceConstantOptimization<ReduceMinOp>>(ctx);
+ patterns.add<ReduceConstantOptimization<ReduceProdOp>>(ctx);
+ patterns.add<ReduceConstantOptimization<ReduceSumOp>>(ctx);
+}
+
 void mlir::tosa::populateTosaFoldConstantTransposePatterns(
  MLIRContext *ctx, RewritePatternSet &patterns) {
  patterns.add<TosaFoldConstantTranspose>(ctx);

mlir/lib/Dialect/Tosa/Transforms/TosaLayerwiseConstantFoldPass.cpp

@@ -52,6 +52,7 @@ struct TosaLayerwiseConstantFoldPass
 
  mlir::tosa::populateTosaFoldConstantReciprocalPatterns(ctx, patterns);
  mlir::tosa::populateTosaFoldConstantTransposePatterns(ctx, patterns);
+ mlir::tosa::populateTosaConstantReduction(ctx, patterns);
  populateTosaOpsCanonicalizationPatterns(ctx, patterns);
 
  if (applyPatternsAndFoldGreedily(func, std::move(patterns)).failed())