Pick getTosaConstShape helper from 571a987

Author: jorickert

mlir/include/mlir/Dialect/Tosa/Utils/ConversionUtils.h

@@ -84,6 +84,12 @@ LogicalResult EqualizeRanks(PatternRewriter &rewriter, Location loc,
 LogicalResult EqualizeRanks(ImplicitLocOpBuilder &builder, Value &input1,
                             Value &input2);
 
+Value getTosaConstShape(ImplicitLocOpBuilder &builder,
+                        llvm::ArrayRef<int64_t> shape);
+
+Value getTosaConstShape(PatternRewriter &rewriter, Location loc,
+                        llvm::ArrayRef<int64_t> shape);
+
 namespace {
 
 // Creates a TOSA operation and performs shape inference on the individual
@@ -217,7 +223,8 @@ TosaOp CreateOpAndInferShape(PatternRewriter &rewriter, Location loc,
 }
 
 // Apply an int32_t permutation to some input, that should be of the same
-// size as perms. Perms should contain some permutation of 0 - perms.size() - 1.
+// size as perms. Perms should contain some permutation of 0 - perms.size()
+// - 1.
 template <typename T>
 SmallVector<T> applyTOSAPermutation(ArrayRef<T> input,
                                     ArrayRef<int32_t> perms) {

mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp

@@ -93,15 +93,11 @@ struct SelfConcatToTile : public OpRewritePattern<tosa::ConcatOp> {
     }
     SmallVector<int64_t> multiplies(concatType.getRank(), 1);
     multiplies[concatOp.getAxis()] = concatOp->getNumOperands();
-    const int64_t rank = multiplies.size();
-    auto constantShapeOp = rewriter.create<ConstShapeOp>(
-        concatOp->getLoc(), shapeType::get(concatOp->getContext(), rank),
-        DenseIntElementsAttr::get(
-            RankedTensorType::get({rank}, rewriter.getIndexType()),
-            multiplies));
+    auto constantShapeValue =
+        getTosaConstShape(rewriter, concatOp->getLoc(), multiplies);
     auto tileOp = rewriter.createOrFold<tosa::TileOp>(
         concatOp->getLoc(), concatOp.getType(), concatOp->getOperand(0),
-        constantShapeOp);
+        constantShapeValue);
     rewriter.replaceOp(concatOp, {tileOp});
     return success();
   }
@@ -140,19 +136,15 @@ struct FuseChainedTile : public OpRewritePattern<tosa::TileOp> {
     for (auto [idx, multiplier] : llvm::enumerate(inputTileMultiples)) {
       multiplies[idx] *= multiplier;
     }
-    auto constantShapeOp = rewriter.create<ConstShapeOp>(
+    auto constantShapeValue = getTosaConstShape(
+        rewriter,
         rewriter.getFusedLoc(
             {op.getMultiples().getLoc(), inputTile.getMultiples().getLoc()}),
-        op.getMultiples().getType(),
-        DenseIntElementsAttr::get(
-            RankedTensorType::get(
-                {cast<shapeType>(op.getMultiples().getType()).getRank()},
-                rewriter.getIndexType()),
-            multiplies));
+        multiplies);
 
     rewriter.modifyOpInPlace(op, [&]() {
       op.setOperand(0, inputTile->getOperand(0));
-      op.setOperand(1, constantShapeOp);
+      op.setOperand(1, constantShapeValue);
       op.getOperation()->setLoc(
           FusedLoc::get(getContext(), {inputTile->getLoc(), op.getLoc()}));
     });
@@ -828,16 +820,11 @@ struct TileSliceOptimization : public OpRewritePattern<tosa::SliceOp> {
          llvm::zip_equal(newTileShape, requiredMultipliers)) {
       newShape *= multiplier;
     }
-    auto constantShapeOp = rewriter.create<ConstShapeOp>(
-        tileOp.getMultiples().getLoc(), tileOp.getMultiples().getType(),
-        DenseIntElementsAttr::get(
-            RankedTensorType::get(
-                {cast<shapeType>(tileOp.getMultiples().getType()).getRank()},
-                rewriter.getIndexType()),
-            requiredMultipliers));
+    auto constantShapeValue = getTosaConstShape(
+        rewriter, tileOp.getMultiples().getLoc(), requiredMultipliers);
     auto newTile = rewriter.create<tosa::TileOp>(
         tileOp->getLoc(), tileOpInputType.clone(newTileShape),
-        tileOp->getOperand(0), constantShapeOp);
+        tileOp->getOperand(0), constantShapeValue);
     auto newSlice = rewriter.create<tosa::SliceOp>(
         sliceOp->getLoc(), sliceOp.getType(), newTile,
         rewriter.getDenseI64ArrayAttr(newTileStarts), sliceOp.getSizeAttr());

mlir/lib/Dialect/Tosa/Utils/ConversionUtils.cpp

@@ -160,3 +160,25 @@ LogicalResult mlir::tosa::EqualizeRanks(ImplicitLocOpBuilder &builder,
 
   return success();
 }
+
+namespace {
+SmallVector<int64_t> convertFromMlirShape(ArrayRef<int64_t> shape) {
+  return to_vector(llvm::map_range(shape, [](int64_t dim) {
+    return ShapedType::isDynamic(dim) ? -1 : dim;
+  }));
+}
+} // namespace
+
+Value mlir::tosa::getTosaConstShape(ImplicitLocOpBuilder &builder,
+                                    llvm::ArrayRef<int64_t> shape) {
+  auto attr = builder.getIndexTensorAttr(convertFromMlirShape(shape));
+  auto type = mlir::tosa::shapeType::get(builder.getContext(), shape.size());
+  mlir::Operation *mlir_op = builder.create<tosa::ConstShapeOp>(type, attr);
+  return mlir_op->getResult(0);
+}
+
+Value mlir::tosa::getTosaConstShape(PatternRewriter &rewriter, Location loc,
+                                    llvm::ArrayRef<int64_t> shape) {
+  ImplicitLocOpBuilder builder(loc, rewriter);
+  return getTosaConstShape(builder, shape);
+}