[MLIR][Shape] Concretize broadcast result type if possible

[frgossen]

Summary

As a canonicalization, infer the resulting shape rank if possible.

Depends On D101376

[frgossen]

Rebase

mlir/include/mlir/Dialect/Shape/IR/Shape.h

 #include "mlir/Interfaces/SideEffectInterfaces.h"
 
 namespace mlir {
 class PatternRewriter;
 
 namespace shape {
 
 /// Alias type for extent tensors.
-RankedTensorType getExtentTensorType(MLIRContext *ctx);
+RankedTensorType getExtentTensorType(MLIRContext *ctx,
+                                     int64_t rank = ShapedType::kDynamicSize);
 
 // Given an input shape Value, try to obtain the shape's values.
 LogicalResult getShapeVec(Value input, SmallVectorImpl<int64_t> &shapeValues);
 
 /// The shape descriptor type represents rank and dimension sizes.
 class ShapeType : public Type::TypeBase<ShapeType, Type, TypeStorage> {
 public:
   using Base::Base;

mlir/lib/Dialect/Shape/IR/Shape.cpp

 
 using namespace mlir;
 using namespace mlir::shape;
 
 namespace {
 #include "ShapeCanonicalization.inc"
 }
 
-RankedTensorType shape::getExtentTensorType(MLIRContext *ctx) {
-  return RankedTensorType::get({ShapedType::kDynamicSize}, IndexType::get(ctx));
+RankedTensorType shape::getExtentTensorType(MLIRContext *ctx, int64_t rank) {
+  return RankedTensorType::get({rank}, IndexType::get(ctx));
 }
 
 LogicalResult shape::getShapeVec(Value input,
                                  SmallVectorImpl<int64_t> &shapeValues) {
   if (auto inputOp = input.getDefiningOp<ShapeOfOp>()) {
     auto type = inputOp.arg().getType().dyn_cast<ShapedType>();
     if (!type.hasRank())
       return failure();
LogicalResult matchAndRewrite(OpTy op,
 
     // Rewrite op if any change required.
     if (!anyChange)
       return failure();
     rewriter.replaceOpWithNewOp<OpTy>(op, op->getResultTypes(), newOperands);
     return success();
   }
 };
+
+struct BroadcastConcretizeResultTypePattern
+    : public OpRewritePattern<BroadcastOp> {
+  using OpRewritePattern<BroadcastOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(BroadcastOp op,
+                                PatternRewriter &rewriter) const override {
+    // Only concretize dynamic extent tensor result types.
+    auto resultTy = op.getType().dyn_cast<RankedTensorType>();
+    if (!resultTy || !resultTy.isDynamicDim(0))
+      return failure();
+
+    // Infer resulting shape rank if possible.
+    int64_t maxRank = 0;
+    for (Value shape : op.shapes()) {
+      if (auto extentTensorTy = shape.getType().dyn_cast<RankedTensorType>()) {
+        // Cannot infer resulting shape rank if any operand is dynamically
+        // ranked.
+        if (extentTensorTy.isDynamicDim(0))
+          return failure();
+        maxRank = std::max(maxRank, extentTensorTy.getDimSize(0));
+      }
+    }
+
+    auto newOp = rewriter.create<BroadcastOp>(
+        op.getLoc(), getExtentTensorType(getContext(), maxRank), op.shapes());
+    rewriter.replaceOpWithNewOp<tensor::CastOp>(op, op.getType(), newOp);
+    return success();
+  }
+};
 } // namespace
 
 void BroadcastOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
                                               MLIRContext *context) {
-  patterns.add<BroadcastFoldConstantOperandsPattern,
+  patterns.add<BroadcastConcretizeResultTypePattern,
+               BroadcastFoldConstantOperandsPattern,
                BroadcastForwardSingleOperandPattern,
                CanonicalizeCastExtentTensorOperandsPattern<BroadcastOp>,
                RemoveDuplicateOperandsPattern<BroadcastOp>,
                RemoveEmptyShapeOperandsPattern<BroadcastOp>>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // ConcatOp

mlir/test/Dialect/Shape/canonicalize.mlir

 // -----
 
 // CHECK-LABEL: @cast_extent_tensor_operands
 // CHECK-SAME: (%[[ARG0:.*]]: tensor<?xindex>, %[[ARG1:.*]]: tensor<3xindex>)
 func @cast_extent_tensor_operands(%arg0 : tensor<?xindex>,
     %arg1 : tensor<3xindex>) -> (!shape.witness, tensor<?xindex>) {
   // CHECK: %[[CAST_ARG0:.*]] = tensor.cast %[[ARG0]] : tensor<?xindex> to tensor<3xindex>
   // CHECK: %[[WIT:.*]] = shape.cstr_broadcastable %[[CAST_ARG0]], %[[ARG1]] : tensor<3xindex>, tensor<3xindex>
-  // CHECK: %[[RES:.*]] = shape.broadcast %[[CAST_ARG0]], %[[ARG1]] : tensor<3xindex>, tensor<3xindex>
+  // CHECK: %[[UNCAST_RES:.*]] = shape.broadcast %[[CAST_ARG0]], %[[ARG1]] : tensor<3xindex>, tensor<3xindex> -> tensor<3xindex>
+  // CHECK: %[[RES:.*]] = tensor.cast %[[UNCAST_RES]] : tensor<3xindex> to tensor<?xindex>
   // CHECK: return %[[WIT]], %[[RES]]
   %0 = tensor.cast %arg0 : tensor<?xindex> to tensor<3xindex>
   %1 = tensor.cast %arg1 : tensor<3xindex> to tensor<?xindex>
   %2 = shape.cstr_broadcastable %0, %1 : tensor<3xindex>, tensor<?xindex>
   %3 = shape.broadcast %0, %1 :tensor<3xindex>, tensor<?xindex>
       -> tensor<?xindex>
   return %2, %3 : !shape.witness, tensor<?xindex>
 }
+
+// -----
+
+// CHECK-LABEL: @concretize_broadcast_result_type
+// CHECK-SAME:  (%[[ARG0:.*]]: tensor<2xindex>, %[[ARG1:.*]]: tensor<3xindex>)
+func @concretize_broadcast_result_type(%arg0 : tensor<2xindex>,
+    %arg1 : tensor<3xindex>) -> tensor<?xindex> {
+  // CHECK: %[[CONCR:.*]] = shape.broadcast %[[ARG0]], %[[ARG1]] : tensor<2xindex>, tensor<3xindex> -> tensor<3xindex>
+  // CHECK: %[[RES:.*]] = tensor.cast %[[CONCR]] : tensor<3xindex> to tensor<?xindex>
+  // CHECK: return %[[RES]]
+  %0 = shape.broadcast %arg0, %arg1 : tensor<2xindex>, tensor<3xindex>
+      -> tensor<?xindex>
+  return %0 : tensor<?xindex>
+}