[mlir][tosa] Add tosa.reverse lowering to linalg.generic

[rsuderman]

Summary

Reverse lowers to a linalg.generic op by reversing the read order
in the index map.

[rsuderman]

Fixed unused variable.

[rsuderman]

Sync to head, pass names were made consistent.

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

public:
   using OpRewritePattern<SrcOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(SrcOp op,
                                 PatternRewriter &rewriter) const final {
     return elementwiseMatchAndRewriteHelper(op, rewriter);
   }
 };
 
-class ReshapeOpConverter : public OpConversionPattern<tosa::ReshapeOp> {
+class ReshapeConverter : public OpConversionPattern<tosa::ReshapeOp> {
 public:
   using OpConversionPattern<tosa::ReshapeOp>::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(tosa::ReshapeOp reshape, ArrayRef<Value> args,
                   ConversionPatternRewriter &rewriter) const final {
     typename tosa::ReshapeOp::Adaptor operands(args);
 
rewriter.replaceOpWithNewOp<linalg::GenericOp>(
         getNParallelLoopsAttrs(resultTy.getRank()),
         [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange args) {
           nestedBuilder.create<linalg::YieldOp>(op.getLoc(), *args.begin());
         });
     return success();
   }
 };
 
-class RescaleOpConverter : public OpRewritePattern<tosa::RescaleOp> {
+class RescaleConverter : public OpRewritePattern<tosa::RescaleOp> {
 public:
   using OpRewritePattern<tosa::RescaleOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(tosa::RescaleOp op,
                                 PatternRewriter &rewriter) const final {
     auto loc = op.getLoc();
     auto input = op.input();
     auto inputTy = op.input().getType().cast<ShapedType>();
public:
   using OpRewritePattern<SrcOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(SrcOp reduceOp,
                                 PatternRewriter &rewriter) const final {
     return reduceMatchAndRewriteHelper(reduceOp, reduceOp.axis(), rewriter);
   }
 };
 
-struct ConcatOpConversion : public OpConversionPattern<tosa::ConcatOp> {
+struct ConcatConverter : public OpConversionPattern<tosa::ConcatOp> {
   using OpConversionPattern<tosa::ConcatOp>::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(tosa::ConcatOp op, ArrayRef<Value> args,
                   ConversionPatternRewriter &rewriter) const override {
     auto resultType = op.getType().dyn_cast<RankedTensorType>();
     if (!resultType || !resultType.hasStaticShape()) {
       return rewriter.notifyMatchFailure(op,
for (auto arg : args) {
                                                   sizes, strides);
       offsets[axis] = rewriter.create<AddIOp>(loc, offsets[axis], sizes[axis]);
     }
     rewriter.replaceOp(op, result);
     return success();
   }
 };
 
+class ReverseConverter : public OpRewritePattern<tosa::ReverseOp> {
+public:
+  using OpRewritePattern<tosa::ReverseOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(tosa::ReverseOp op,
+                                PatternRewriter &rewriter) const final {
+    auto loc = op.getLoc();
+    Value input = op.input();
+    auto inputTy = input.getType().template cast<ShapedType>();
+    auto resultTy = op.getType().template cast<ShapedType>();
+    auto rank = resultTy.getRank();
+    auto axis = op.axis();
+
+    if (!inputTy.hasStaticShape())
+      return rewriter.notifyMatchFailure(
+          op, "No initial value found for reduction operation");
+
+    // First fill the output buffer with the init value.
+    auto initTensor = rewriter
+                          .create<linalg::InitTensorOp>(
+                              loc, ArrayRef<Value>({}), inputTy.getShape(),
+                              inputTy.getElementType())
+                          .result();
+
+    SmallVector<AffineExpr, 2> inputExprs;
+    inputExprs.resize(resultTy.getRank());
+
+    for (int i = 0; i < rank; i++)
+      inputExprs[i] = rewriter.getAffineDimExpr(i);
+
+    inputExprs[axis] =
+        rewriter.getAffineConstantExpr(inputTy.getDimSize(axis) - 1) -
+        inputExprs[axis];
+
+    SmallVector<AffineMap, 2> affineMaps = {
+        AffineMap::get(resultTy.getRank(), /*symbolCount=*/0, inputExprs,
+                       rewriter.getContext()),
+        rewriter.getMultiDimIdentityMap(resultTy.getRank())};
+
+    rewriter.replaceOpWithNewOp<linalg::GenericOp>(
+        op, resultTy, op.input(), ValueRange{initTensor}, affineMaps,
+        getNParallelLoopsAttrs(resultTy.getRank()),
+        [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange args) {
+          nestedBuilder.create<linalg::YieldOp>(op.getLoc(), *args.begin());
+        });
+
+    return success();
+  }
+};
+
 } // namespace
 
 void mlir::tosa::populateTosaToLinalgOnTensorsConversionPatterns(
     MLIRContext *context, OwningRewritePatternList *patterns) {
   patterns->insert<
       PointwiseConverter<tosa::AddOp>, PointwiseConverter<tosa::SubOp>,
       PointwiseConverter<tosa::MulOp>, PointwiseConverter<tosa::NegateOp>,
       PointwiseConverter<tosa::PowOp>, PointwiseConverter<tosa::RsqrtOp>,
patterns->insert<
       PointwiseConverter<tosa::SelectOp>, PointwiseConverter<tosa::GreaterOp>,
       PointwiseConverter<tosa::GreaterEqualOp>,
       PointwiseConverter<tosa::MaximumOp>, PointwiseConverter<tosa::MinimumOp>,
       PointwiseConverter<tosa::CeilOp>, PointwiseConverter<tosa::FloorOp>,
       PointwiseConverter<tosa::ClampOp>, PointwiseConverter<tosa::ReluNOp>,
       IdentityNConverter<tosa::IdentityOp>,
       IdentityNConverter<tosa::IdentityNOp>, ReduceConverter<tosa::ReduceMinOp>,
       ReduceConverter<tosa::ReduceMaxOp>, ReduceConverter<tosa::ReduceSumOp>,
-      ReduceConverter<tosa::ReduceProdOp>, ConcatOpConversion,
-      ReshapeOpConverter, TransposeConverter, RescaleOpConverter>(context);
+      ReduceConverter<tosa::ReduceProdOp>, ConcatConverter, ReshapeConverter,
+      RescaleConverter, ReverseConverter, TransposeConverter>(context);
 }

mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir

 // CHECK-LABEL: @rescaleUnnecessaryDoubleRound
 func @rescaleUnnecessaryDoubleRound(%arg0 : tensor<1xi8>) -> (tensor<1xi8>) {
   // CHECK: linalg.generic
   // CHECK: "tosa.apply_scale"
   // CHECK-SAME:  {double_round = false}
   %0 = "tosa.rescale"(%arg0) {input_zp = 243 : i32, output_zp = 252 : i32, multiplier = [19689 : i32], shift = [15 : i32], scale32 = true, double_round = true, per_channel = false} : (tensor<1xi8>)  -> (tensor<1xi8>)
   return %0 : tensor<1xi8>
 }
+
+// -----
+
+// CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (-d0 + 4, d1)>
+// CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d0, d1)>
+// CHECK: #[[$MAP2:.*]] = affine_map<(d0, d1) -> (d0, -d1 + 3)>
+
+// CHECK-LABEL: @reverse
+func @reverse(%arg0: tensor<5x4xi32>) -> () {
+  // CHECK: [[INIT:%.+]] = linalg.init_tensor [5, 4]
+  // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["parallel", "parallel"]} ins(%arg0 : tensor<5x4xi32>) outs([[INIT]] : tensor<5x4xi32>) {
+  // CHECK: ^bb0(%arg1: i32, %arg2: i32):
+  // CHECK:   linalg.yield %arg1 : i32
+  %0 = "tosa.reverse"(%arg0) {axis = 0 : i64} : (tensor<5x4xi32>) -> tensor<5x4xi32>
+
+  // CHECK: [[INIT:%.+]] = linalg.init_tensor [5, 4]
+  // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP2]], #[[$MAP1]]], iterator_types = ["parallel", "parallel"]} ins(%arg0 : tensor<5x4xi32>) outs([[INIT]] : tensor<5x4xi32>) {
+  // CHECK: ^bb0(%arg1: i32, %arg2: i32):
+  // CHECK:   linalg.yield %arg1 : i32
+  %1 = "tosa.reverse"(%arg0) {axis = 1 : i64} : (tensor<5x4xi32>) -> tensor<5x4xi32>
+
+  return
+}