[mlir][tosa] Make tosa MakeBroadcastable pass handle unreanked tensors.

[rsuderman]

Summary

If this pass executes without shape inference its possible for unranked tensors
to appear in the IR. This pass should gracefully handle unranked tensors.

mlir/lib/Dialect/Tosa/Transforms/TosaMakeBroadcastable.cpp

if (lowerRightIndex != 0 && higherRightIndex != 0) {
     }
   }
 }
 
 /// Common code to create the reshape op where necessary to make the rank of the
 /// operations equal. Returns the updated input1 and input2 for the original
 /// input. The caller is expected to use these to rewrite the original operator
 /// with the RESHAPE now in the graph.
-static int reshapeLowerToHigher(PatternRewriter &rewriter, Location loc,
-                                RankedTensorType outputType, Value input1,
-                                Value input2, Value &outInput1,
-                                Value &outInput2) {
+static LogicalResult reshapeLowerToHigher(PatternRewriter &rewriter,
+                                          Location loc,
+                                          RankedTensorType outputType,
+                                          Value input1, Value input2,
+                                          Value &outInput1, Value &outInput2) {
+  auto input1Ty = input1.getType().dyn_cast<RankedTensorType>();
+  auto input2Ty = input2.getType().dyn_cast<RankedTensorType>();
 
-  int64_t input1Rank = input1.getType().cast<RankedTensorType>().getRank();
-  int64_t input2Rank = input2.getType().cast<RankedTensorType>().getRank();
+  if (!input1Ty || !input2Ty)
+    return failure();
+
+  int64_t input1Rank = input1Ty.getRank();
+  int64_t input2Rank = input2Ty.getRank();
 
   Value higherTensorValue, lowerTensorValue;
-  // return if rank already match
+  // Cannot rewrite as its already correct.
   if (input1Rank == input2Rank)
-    return 1;
+    return failure();
 
   if (input1Rank > input2Rank) {
     higherTensorValue = input1;
     lowerTensorValue = input2;
   } else {
     higherTensorValue = input2;
     lowerTensorValue = input1;
   }
 
-  ArrayRef<int64_t> outputRankShape = outputType.getShape();
   ArrayRef<int64_t> higherRankShape =
       higherTensorValue.getType().cast<RankedTensorType>().getShape();
   (void)higherRankShape;
   ArrayRef<int64_t> lowerRankShape =
       lowerTensorValue.getType().cast<RankedTensorType>().getShape();
 
-  // outputRank == higherRank == max(input1Rank, input2Rank)
-  assert(higherRankShape.size() == outputRankShape.size());
-
   SmallVector<int64_t, 4> reshapeOutputShape;
 
-  computeReshapeOutput(outputRankShape, lowerRankShape, reshapeOutputShape);
+  computeReshapeOutput(higherRankShape, lowerRankShape, reshapeOutputShape);
 
   auto reshapeInputType = lowerTensorValue.getType().cast<RankedTensorType>();
   auto reshapeOutputType = RankedTensorType::get(
       ArrayRef<int64_t>(reshapeOutputShape), reshapeInputType.getElementType());
 
+  // Verify the rank agrees with the output type if the output type is ranked.
+  if (outputType) {
+    if (outputType.getShape().size() != reshapeOutputShape.size() ||
+        outputType.getShape().size() != higherRankShape.size())
+      return failure();
+  }
+
   auto reshapeLower = rewriter.create<tosa::ReshapeOp>(
       loc, reshapeOutputType, lowerTensorValue,
       rewriter.getI64ArrayAttr(reshapeOutputShape));
 
   if (input1Rank > input2Rank) {
     outInput1 = higherTensorValue;
     outInput2 = reshapeLower.getResult();
   } else {
     outInput1 = reshapeLower.getResult();
     outInput2 = higherTensorValue;
   }
 
-  return 0;
+  return success();
 }
 
 namespace {
 template <typename OpTy>
 struct ConvertTosaOp : public OpRewritePattern<OpTy> {
   using OpRewritePattern<OpTy>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(OpTy tosaBinaryOp,
                                 PatternRewriter &rewriter) const override {
 
     Value input1 = tosaBinaryOp.input1();
     Value input2 = tosaBinaryOp.input2();
     Value output = tosaBinaryOp.getResult();
-    auto outputType = output.getType().cast<RankedTensorType>();
+
+    auto outputType = output.getType().dyn_cast<RankedTensorType>();
 
     Value outInput1, outInput2;
     if (reshapeLowerToHigher(rewriter, tosaBinaryOp.getLoc(), outputType,
-                             input1, input2, outInput1, outInput2))
+                             input1, input2, outInput1, outInput2)
+            .failed())
       return failure();
 
     rewriter.replaceOpWithNewOp<OpTy>(tosaBinaryOp, outputType, outInput1,
                                       outInput2);
 
     return success();
   }
 };
 
 // The MulOp has an extra parameter 'shift' not present in other elementwise
 // binary ops, that necessitates special handling of its builder.
 template <>
 struct ConvertTosaOp<tosa::MulOp> : public OpRewritePattern<tosa::MulOp> {
   using OpRewritePattern<tosa::MulOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(tosa::MulOp tosaBinaryOp,
                                 PatternRewriter &rewriter) const override {
 
     Value input1 = tosaBinaryOp.input1();
     Value input2 = tosaBinaryOp.input2();
     int32_t shift = tosaBinaryOp.shift();
     Value output = tosaBinaryOp.getResult();
-    auto outputType = output.getType().cast<RankedTensorType>();
+    auto outputType = output.getType().dyn_cast<RankedTensorType>();
 
     Value outInput1, outInput2;
     if (reshapeLowerToHigher(rewriter, tosaBinaryOp.getLoc(), outputType,
-                             input1, input2, outInput1, outInput2))
+                             input1, input2, outInput1, outInput2)
+            .failed())
       return failure();
 
     rewriter.replaceOpWithNewOp<tosa::MulOp>(tosaBinaryOp, outputType,
                                              outInput1, outInput2, shift);
 
     return success();
   }
 };
LogicalResult matchAndRewrite(tosa::ArithmeticRightShiftOp tosaBinaryOp,
     Value input1 = tosaBinaryOp.input1();
     Value input2 = tosaBinaryOp.input2();
     int32_t round = tosaBinaryOp.round();
     Value output = tosaBinaryOp.getResult();
     auto outputType = output.getType().dyn_cast<RankedTensorType>();
 
     Value outInput1, outInput2;
     if (reshapeLowerToHigher(rewriter, tosaBinaryOp.getLoc(), outputType,
-                             input1, input2, outInput1, outInput2))
+                             input1, input2, outInput1, outInput2)
+            .failed())
       return failure();
 
     rewriter.replaceOpWithNewOp<tosa::ArithmeticRightShiftOp>(
         tosaBinaryOp, outputType, outInput1, outInput2, round);
 
     return success();
   }
 };