[MLIR][TOSA] Lower tosa.transpose to linalg.generic

[rsuderman]

Summary

Lowers the transpose operation to a generic linalg op when permutations
is a constant value.

[nicolasvasilache]

This one is probably better rewritten as a linalg.copy with the proper permutations (at least until we have the generic -> named ops inference).

[mravishankar]

In D97508#2589742, @nicolasvasilache wrote:

This one is probably better rewritten as a linalg.copy with the proper permutations (at least until we have the generic -> named ops inference).

This would require making CopyOp work on tensors. @rsuderman do you think you can make that change? Not a requirement though, this would be fine as is....

Few minor comments below.

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

368

Does TOSA op only support static shapes. If not then you need to have a way to get the dynamic values for the non-static dimensions.

378

I dont think you need to do ValueRange{...} . Just op.input1() should work,

[rsuderman]

Updated for ValueRange comment.

[rsuderman]

Made the minor changes possible. Deferring on the CopyOp changes as lack of tensor support is outside the scope of TOSA work. If I have some time I will look into updating linalg.copy

[rsuderman]

Rebased.

[rsuderman]

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

368

I'm only supporting static shapes for now to get things bootstrapped. Once we have more dynamic cases we can swing back around to dynamic support.

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Conversion/TosaToLinalg/TosaToLinalg.h"
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
 #include "mlir/Dialect/Math/IR/Math.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Dialect/Tosa/IR/TosaOps.h"
+#include "mlir/IR/Matchers.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 #include <numeric>
 
 using namespace mlir;
 
matchAndRewrite(tosa::ReshapeOp reshape, ArrayRef<Value> args,
 
     // Compute the reassociation maps for the linalg operation.
     ArrayRef<int64_t> expandedShape =
         (operandTy.getRank() > resultTy.getRank() ? operandTy.getShape()
                                                   : resultTy.getShape());
     ArrayRef<int64_t> collapsedShape =
         (operandTy.getRank() > resultTy.getRank() ? resultTy.getShape()
                                                   : operandTy.getShape());
     unsigned currSrcDim = 0, currDstDim = 0;
     SmallVector<linalg::ReassociationExprs, 4> reassociationMap(
         collapsedShape.size());
 
     // First scan all dimensions in the source shapes to see whether we have a
     // perfect case where consecutive dimensions in source are collapsed. For
     // such case we can just generate one single linalg.reshape.
     bool isCollapsingSource = true;
     while (currSrcDim < expandedShape.size() &&
            currDstDim < collapsedShape.size()) {
       int64_t dstSize = collapsedShape[currDstDim];
       int64_t srcSize = expandedShape[currSrcDim];
       while (srcSize < dstSize && currSrcDim < expandedShape.size()) {
         reassociationMap[currDstDim].push_back(
             rewriter.getAffineDimExpr(currSrcDim++));
         srcSize *= expandedShape[currSrcDim];
       }
       if (srcSize == dstSize) {
         reassociationMap[currDstDim].push_back(
if (!isCollapsingSource) {
       rewriter.replaceOpWithNewOp<linalg::TensorReshapeOp>(
           reshape, resultTy, collapsedOp, expandingMap);
 
       return success();
     }
 
     rewriter.replaceOpWithNewOp<linalg::TensorReshapeOp>(
         reshape, resultTy, args[0], reassociationMap);
+
+    return success();
+  }
+};
+
+class TransposeConverter : public OpRewritePattern<tosa::TransposeOp> {
+public:
+  using OpRewritePattern<tosa::TransposeOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(tosa::TransposeOp op,
+                                PatternRewriter &rewriter) const final {
+    DenseIntElementsAttr perms;
+    if (!matchPattern(op.perms(), m_Constant(&perms))) {
+      return failure();
+    }
+
+    auto resultTy = op.getType().cast<ShapedType>();
+    if (!resultTy.hasStaticShape())
+      return failure();
+
+    SmallVector<AffineExpr, 2> inputExprs;
+    inputExprs.resize(resultTy.getRank());
+    for (auto permutation : llvm::enumerate(perms.getIntValues())) {
+      inputExprs[permutation.value().getZExtValue()] =
+          rewriter.getAffineDimExpr(permutation.index());
+    }
+
+    auto initTensor = rewriter.create<linalg::InitTensorOp>(
+        op.getLoc(), ArrayRef<Value>({}), resultTy.getShape(),
+        resultTy.getElementType());
+
+    SmallVector<AffineMap, 2> affineMaps = {
+        AffineMap::get(resultTy.getRank(), /*symbolCount=*/0, inputExprs,
+                       rewriter.getContext()),
+        rewriter.getMultiDimIdentityMap(resultTy.getRank())};
+
+    rewriter.replaceOpWithNewOp<linalg::GenericOp>(
+        op, resultTy, op.input1(), ValueRange{initTensor}, affineMaps,
+        getNParallelLoopsAttrs(resultTy.getRank()),
+        [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange args) {
+          nestedBuilder.create<linalg::YieldOp>(op.getLoc(), *args.begin());
+        });
     return success();
   }
 };
 
 // At the codegen level any identity operations should be removed. Any cases
 // where identity is load-bearing (e.g. cross device computation) should be
 // handled before lowering to codegen.
 template <typename SrcOp>
patterns->insert<
       PointwiseConverter<tosa::LogicalLeftShiftOp>,
       PointwiseConverter<tosa::LogicalRightShiftOp>,
       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>, ReshapeOpConverter>(context);
+      IdentityNConverter<tosa::IdentityNOp>,
+      ReshapeOpConverter, TransposeConverter>(context);
 }

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

func @test_identity(%arg0: tensor<1xf32>, %arg1: tensor<1xi32>) -> (tensor<1xf32>, tensor<1xi32>) {
   %0 = "tosa.identity"(%arg0) : (tensor<1xf32>) -> tensor<1xf32>
   %1 = "tosa.identity"(%arg1) : (tensor<1xi32>) -> tensor<1xi32>
 
   %2:2 = "tosa.identityn"(%0, %1) : (tensor<1xf32>, tensor<1xi32>) -> (tensor<1xf32>, tensor<1xi32>)
 
   // CHECK: return %arg0, %arg1
   return %2#0, %2#1 : tensor<1xf32>, tensor<1xi32>
 }
+
+// -----
+
+// CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1, d2) -> (d2, d0, d1)>
+// CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+
+// CHECK-LABEL: @test_transpose
+// CHECK-SAME: ([[ARG0:%.+]]: tensor<1x2x3xi32>)
+func @test_transpose(%arg0: tensor<1x2x3xi32>) -> () {
+  %0 = constant dense<[1, 2, 0]> : tensor<3xi32>
+  // CHECK: [[INIT:%.+]] = linalg.init_tensor [2, 3, 1]
+  // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["parallel", "parallel", "parallel"]} ins([[ARG0]] : tensor<1x2x3xi32>) outs([[OUT:%.+]] : tensor<2x3x1xi32>)
+  // CHECK: ^bb0([[ARG1:%.+]]: i32, [[ARG2:%.+]]: i32)
+  // CHECK:   linalg.yield [[ARG1]]
+  // CHECK: }
+  %1 = "tosa.transpose"(%arg0, %0) : (tensor<1x2x3xi32>, tensor<3xi32>) -> (tensor<2x3x1xi32>)
+  return
+}