[mlir][TensorOps][Canonicalize] FoldDimOfExpandShape better handling for dynamic dim expansion

[hockyy]

https://github.com/llvm/llvm-project/blame/476b208e0115e766605e9f850982996a1d51c287/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp#L1988

@matthias-springer

Why the assert?

    // `dim` is the only dynamic dimension in `group`. (Otherwise, the
    // ExpandShapeOp would be ambiguous.)
    int64_t product = 1;
    ReassociationIndices grp = expandShapeOp.getReassociationIndices()[srcDim];
    bool multiDynamicDim = false;
    for (int64_t d : grp) {
      if (d != dim) {
        if (resultType.isDynamicDim(d)) {
          multiDynamicDim = true;
          break;
        }
        //        assert(!resultType.isDynamicDim(d) && "expected static dim");
        product *= resultType.getDimSize(d);
      }
    }
    LLVM_DEBUG(llvm::dbgs() << expandShapeOp << "\n";);
    int dynIndex = 0;
    for (int i = 0; i < *dim; i++) {
      if (ShapedType::isDynamic(expandShapeOp.getStaticOutputShape()[i]))
        dynIndex++;
    }
    for (auto tmp : expandShapeOp.getOutputShape()) {
      LLVM_DEBUG(llvm::dbgs() << tmp << "\n";);
    }
    for (auto tmp : expandShapeOp.getStaticOutputShape()) {
      LLVM_DEBUG(llvm::dbgs() << tmp << "\n";);
    }

    if (multiDynamicDim) {
      rewriter.replaceOp(dimOp, expandShapeOp.getOutputShape()[dynIndex]);
      return success();
    }

can we do this for better handling?

[hockyy]

module {
  func.func @broadcast_test(%arg0: tensor<4x8x?x3x?xf32>) -> tensor<1x4x3x8x1x1x?x3x1x1x?x1x14xf32> {
    %c4 = arith.constant 4 : index
    %dim = tensor.dim %arg0, %c4 : tensor<4x8x?x3x?xf32>
    %c2 = arith.constant 2 : index
    %dim_0 = tensor.dim %arg0, %c2 : tensor<4x8x?x3x?xf32>
    %collapsed = tensor.collapse_shape %arg0 [[0], [1, 2, 3, 4]] : tensor<4x8x?x3x?xf32> into tensor<4x?xf32>
    %c2_1 = arith.constant 2 : index
    %c4_2 = arith.constant 4 : index
    %dim_3 = tensor.dim %arg0, %c2_1 : tensor<4x8x?x3x?xf32>
    %dim_4 = tensor.dim %arg0, %c4_2 : tensor<4x8x?x3x?xf32>
    %0 = tensor.empty(%dim_3, %dim_4) : tensor<4x3x8x?x3x?x14xf32>
    %collapsed_5 = tensor.collapse_shape %0 [[0], [1], [2, 3, 4, 5], [6]] : tensor<4x3x8x?x3x?x14xf32> into tensor<4x3x?x14xf32>
    %broadcasted = linalg.broadcast ins(%collapsed : tensor<4x?xf32>) outs(%collapsed_5 : tensor<4x3x?x14xf32>) dimensions = [1, 3]
    %expanded = tensor.expand_shape %broadcasted [[0], [1], [2, 3, 4, 5], [6]] output_shape [4, 3, 8, %dim_0, 3, %dim, 14] : tensor<4x3x?x14xf32> into tensor<4x3x8x?x3x?x14xf32>
    %expanded_6 = tensor.expand_shape %expanded [[0, 1], [2], [3, 4, 5], [6], [7, 8, 9], [10, 11], [12]] output_shape [1, 4, 3, 8, 1, 1, %dim_3, 3, 1, 1, %dim_4, 1, 14] : tensor<4x3x8x?x3x?x14xf32> into tensor<1x4x3x8x1x1x?x3x1x1x?x1x14xf32>
    %c10 = arith.constant 10 : index
    %dim_7 = tensor.dim %expanded_6, %c10 : tensor<1x4x3x8x1x1x?x3x1x1x?x1x14xf32>
    %c6 = arith.constant 6 : index
    %dim_8 = tensor.dim %expanded_6, %c6 : tensor<1x4x3x8x1x1x?x3x1x1x?x1x14xf32>
    %collapsed_9 = tensor.collapse_shape %expanded_6 [[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]] : tensor<1x4x3x8x1x1x?x3x1x1x?x1x14xf32> into tensor<?xf32>
    %expanded_10 = tensor.expand_shape %collapsed_9 [[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]] output_shape [1, 4, 3, 8, 1, 1, %dim_8, 3, 1, 1, %dim_7, 1, 14] : tensor<?xf32> into tensor<1x4x3x8x1x1x?x3x1x1x?x1x14xf32>
    return %expanded_10 : tensor<1x4x3x8x1x1x?x3x1x1x?x1x14xf32>
  }
}

This crashes canonicalizer