Move internal changes

GitOrigin-RevId: de640057c4adcfee8deb6a6f06405c595d48ac56

Author: Copybara Bot

mlir-tensorrt/tensorrt/lib/Utils/ShapeUtils.cpp

@@ -109,32 +109,33 @@ tensorrt::getBroadcastedShape(ArrayRef<ArrayRef<int64_t>> shapes) {
     // shapes are broadcastable. Don't fail because we can't say for sure it's
     // invalid.
     const bool allEqual = llvm::all_equal(dimSizes);
-    if (allEqual && dimSizes.front() == ShapedType::kDynamic)
-      return ShapedType::kDynamic;
 
-    // Dimensions are all equal to a static size.
+    // Dimensions are all equal to a fixed value or dynamic.
     if (allEqual)
       return dimSizes.front();
 
-    // Some dims are '1', all other dims are equal to another fixed number or
-    // dynamic.
+    // Mixture of fixed or unkown extents.
     std::optional<int64_t> nonUnitSize{};
     for (int64_t dimSize : dimSizes) {
+      // Extent of 1 is always valid.
       if (dimSize == 1)
         continue;
+      // Dynamic extent is always valid.
       if (ShapedType::isDynamic(dimSize))
         continue;
+      // If a extent > 1 is present, check that it matches any previously seen
+      // static >1 extent.
       if (nonUnitSize && dimSize == *nonUnitSize)
         continue;
       if (nonUnitSize && dimSize != *nonUnitSize)
         return failure();
       nonUnitSize = dimSize;
     }
-    if (nonUnitSize)
-      return *nonUnitSize;
 
-    // No other case is valid.
-    return failure();
+    // Return the size >1 is seen, otherwise return dynamic indicator. An
+    // inferred size of 1 is only possible if all extents are 1; this case is
+    // captured by the check before the loop.
+    return nonUnitSize ? *nonUnitSize : ShapedType::kDynamic;
   };
 
   for (auto dim : llvm::seq<unsigned>(0, rank)) {

mlir-tensorrt/tensorrt/test/Dialect/TensorRT/invalid.mlir

@@ -857,6 +857,27 @@ func.func @trt_select(%arg0: tensor<10x10xi1>, %arg1: tensor<1x10xf32>, %arg2: t
 
 // -----
 
+func.func @valid_select_ds_infer(%arg0: tensor<?x?xi1>, %arg1: tensor<?x?xf16>, %arg2: tensor<1x1xf16>) -> tensor<?x?xf16> {
+  %0 = tensorrt.select ins(%arg0, %arg1, %arg2 : tensor<?x?xi1>, tensor<?x?xf16>, tensor<1x1xf16>) -> tensor<?x?xf16>
+  return %0 : tensor<?x?xf16>
+}
+
+// -----
+
+func.func @valid_select_ds_infer2(%arg0: tensor<1x?xi1>, %arg1: tensor<1x?xf16>, %arg2: tensor<1x1xf16>) -> tensor<?x?xf16> {
+  %0 = tensorrt.select ins(%arg0, %arg1, %arg2 : tensor<1x?xi1>, tensor<1x?xf16>, tensor<1x1xf16>) -> tensor<?x?xf16>
+  return %0 : tensor<?x?xf16>
+}
+
+// -----
+
+func.func @valid_select_ds_infer3(%arg0: tensor<1x?xi1>, %arg1: tensor<1x?xf16>, %arg2: tensor<1x1xf16>) -> tensor<1x1xf16> {
+  %0 = tensorrt.select ins(%arg0, %arg1, %arg2 : tensor<1x?xi1>, tensor<1x?xf16>, tensor<1x1xf16>) -> tensor<1x1xf16>
+  return %0 : tensor<1x1xf16>
+}
+
+// -----
+
 func.func @trt_softmax(%arg0: tensor<10x10xf32>) -> tensor<10x10xf32> {
   // expected-error @below {{'tensorrt.softmax' op expected axis to be non-negative and less than 2}}
   %0 = tensorrt.softmax {axis = 2 : i64} %arg0 : tensor<10x10xf32>