Commit message: Avoid Segfault for scalar shapes.

Calling tensor::FromElementsOp with an empty vector of elements and no type
causes a segfault. We need to let the FromElementsOp know which scalar type it
should have.
Also add back the DynamicBroadcastInDimOp canonicalization patterns, which
previously prevented this bug from happening.
Add a regression test that demonstrates the bug.

PiperOrigin-RevId: 417561444
Change-Id: I6d1d6cfb71aabbad6102422625a00bbe253ac95a
CVE: CVE-2022-23593
CWE: CWE-754
Vulnerability description: Tensorflow is an Open Source Machine Learning Framework. The `simplifyBroadcast` function in the MLIR-TFRT infrastructure in TensorFlow is vulnerable to a segfault (hence, denial of service), if called with scalar shapes. If all shapes are scalar, then `maxRank` is 0, so we build an empty `SmallVector`. The fix will be included in TensorFlow 2.8.0. This is the only affected version.

Project: tensorflow
File: tf_cpurt_symbolic_shape_optimization.cc

  explicit BroadcastOpLowering(MLIRContext* ctx) : OpRewritePattern(ctx) {}

Project: tensorflow
File: tf_cpurt_symbolic_shape_optimization.cc

DynamicBroadcastInDimOpLowering::DynamicBroadcastInDimOpLowering(
    MLIRContext* ctx)
    : Base(ctx) {}

Project: tensorflow
File: tf_cpurt_symbolic_shape_optimization.cc

llvm::Optional<Value> simplifyBroadcast(ShapeComponentAnalysis& analysis,
                                        ValueRange shapes, Location loc,
                                        OpBuilder* builder) {
  // First find the input shape with the largest rank.
  SmallVector<ArrayRef<ShapeComponentAnalysis::SymbolicExpr>> shapes_found;
  size_t maxRank = 0;
  for (const auto &shape : llvm::enumerate(shapes)) {
    auto found_shape = analysis.GetValueInfo(shape.value());
    if (!found_shape) return {};
    shapes_found.push_back(*found_shape);
    maxRank = std::max(maxRank, found_shape->size());
  }
  if (maxRank == 0) {
    return Value(builder->create<tensor::FromElementsOp>(
        loc, shapes[0].getType(), SmallVector<Value>()));
  }

  SmallVector<const ShapeComponentAnalysis::SymbolicExpr*> joined_dimensions(
      maxRank);
  SmallVector<std::pair<Value, int64_t>> shape_and_rank_for_dim(maxRank);
  for (const auto &shape : llvm::enumerate(shapes_found)) {
    for (const auto &dim : llvm::enumerate(llvm::reverse(shape.value()))) {
      // 1 dimensions don't contribute to the final result.
      if (dim.value().isConstant(1)) continue;
      // If it's not a 1 dimension it will be present in the result. Remember
      // where it came from.
      auto index = maxRank - dim.index() - 1;
      if (!joined_dimensions[index]) {
        joined_dimensions[index] = &dim.value();
        shape_and_rank_for_dim[index] =
            std::make_pair(shapes[shape.index()], shape.value().size());
        continue;
      }
      // Bail if the dimensions are neither equal nor 1.
      if (*joined_dimensions[index] != dim.value()) return {};
    }
  }
  // If the output is the same as one of the inputs just return that.
  if (llvm::is_splat(shape_and_rank_for_dim) &&
      shape_and_rank_for_dim[0].first) {
    return shape_and_rank_for_dim[0].first;
  }
  // Otherwise rematerialize the shape from the pieces we have.
  SmallVector<Value> elements;
  for (int i = 0; i != maxRank; ++i) {
    // 1 dimensions are filtered above, recreate the constant.
    if (!shape_and_rank_for_dim[i].first) {
      auto one = builder->getIntegerAttr(
          shapes[0].getType().cast<RankedTensorType>().getElementType(), 1);
      elements.push_back(builder->create<ConstantOp>(loc, one));
      continue;
    }
    // Extract from one of the shapes, accounting for the reverse indexing
    // performed by broadcast.
    Value index = builder->create<ConstantIndexOp>(
        loc, i - maxRank + shape_and_rank_for_dim[i].second);
    elements.push_back(builder->create<tensor::ExtractOp>(
        loc, shape_and_rank_for_dim[i].first, index));
  }
  return Value(builder->create<tensor::FromElementsOp>(loc, elements));
}

Project: tensorflow
File: tf_cpurt_symbolic_shape_optimization.cc

LogicalResult CstrBroadcastableOpLowering::matchAndRewrite(
    shape::CstrBroadcastableOp op, mlir::PatternRewriter& rewriter) const {
  ShapeComponentAnalysis shape_component_analysis;
  if (!isKnownBroadcastable(shape_component_analysis, op.getShapes(),
                            op.getShapes().front()))
    return failure();

  // Replace constraint with a true witness.
  rewriter.replaceOpWithNewOp<shape::ConstWitnessOp>(op, true);

  return success();
}

Project: tensorflow
File: tf_cpurt_symbolic_shape_optimization.cc

  void runOnFunction() override {
    FuncOp func = getFunction();

    MLIRContext* ctx = &getContext();
    mlir::RewritePatternSet patterns(ctx);

    // Rewrite constraints based on the symbolic shapes.
    patterns.insert<CstrBroadcastableOpLowering>(ctx);
    // Rewrite shape.broadcast based on the symbolic shapes.
    patterns.insert<BroadcastOpLowering>(ctx);

    // Rewrite broadcasts based on the symbolic shapes if enabled.
    if (!optimize_only_constraints)
      patterns.insert<DynamicBroadcastInDimOpLowering>(ctx);

    // Add shape dialect canonicalization patterns to fold shape operations
    // after constraints are replaced with constant witness.
    for (auto op : ctx->getRegisteredOperations()) {
      if (llvm::isa<shape::ShapeDialect>(op.getDialect()))
        op.getCanonicalizationPatterns(patterns, ctx);
    }

    (void)mlir::applyPatternsAndFoldGreedily(func, std::move(patterns));
  }

Project: tensorflow
File: tf_cpurt_symbolic_shape_optimization.cc

bool isKnownBroadcastable(ShapeComponentAnalysis& analysis,
                          ValueRange bcasted_shapes, Value output_shape) {
  auto output_shape_dims = analysis.GetValueInfo(output_shape);
  if (!output_shape_dims) return false;
  for (Value shape : bcasted_shapes) {
    auto shape_dims = analysis.GetValueInfo(shape);
    if (!shape_dims) return false;
    // Iterate backwards over the smallest input shape.
    for (auto zip : llvm::zip(llvm::reverse(*output_shape_dims),
                              llvm::reverse(*shape_dims))) {
      const auto& first = std::get<0>(zip);
      const auto& second = std::get<1>(zip);
      // TODO(ezhulenev): What to do with dimensions statically known to be
      // zero?
      // Numpy can only broadcast [0] with [1], however Tensorflow can broadcast
      // [0] with any dimension size, and produces dimension of size [0].
      // Currently we'll conservatively return failure and will not proceed with
      // a rewrite.
      if (first.isConstant(0) || second.isConstant(0)) return false;
      // If either shape has a static one dimension the broadcast will always
      // succeed.
      if (first.isConstant(1) || second.isConstant(1)) continue;
      // Otherwise dims have to be equal.
      if (first != second) return false;
    }
  }
  return true;
}

Project: tensorflow
File: tf_cpurt_symbolic_shape_optimization.cc

llvm::Optional<AffineMap> isNonExpandingBroadcast(
    ShapeComponentAnalysis& analysis, Value from, Value to_shape) {
  auto in_shape = analysis.GetShapeInfo(from);
  auto out_shape = analysis.GetValueInfo(to_shape);
  if (!in_shape || !out_shape) return {};

  SmallVector<AffineExpr> input_map_exprs;
  size_t rank = out_shape->size();
  MLIRContext* ctx = (*out_shape)[0].expr.getContext();
  size_t d = 0;
  auto affine_zero = getAffineConstantExpr(0, ctx);
  for (auto zip :
       llvm::zip(llvm::reverse(*in_shape), llvm::reverse(*out_shape))) {
    const auto& in = std::get<0>(zip);
    const auto& out = std::get<1>(zip);
    bool extend = in.isConstant(1) && !out.isConstant(1);
    input_map_exprs.push_back(extend ? affine_zero
                                     : getAffineDimExpr(rank - d - 1, ctx));
    ++d;

    // Bail if this is neither a known expansion nor a known non-expansion.
    if (!extend && in != out) return {};
  }
  // Any leading dimensions will be expanded.
  input_map_exprs.resize(in_shape->size(), affine_zero);
  std::reverse(input_map_exprs.begin(), input_map_exprs.end());
  return AffineMap::get(/*dimCount=*/rank,
                        /*symbolCount=*/0, input_map_exprs, ctx);
}

Project: tensorflow
File: tf_cpurt_symbolic_shape_optimization.cc

std::unique_ptr<FunctionPass> CreateSymbolicShapeOptimizationPass(
    bool constraints_only) {
  return std::make_unique<SymbolicShapeOptimizationPass>(constraints_only);
}

Project: tensorflow
File: tf_cpurt_symbolic_shape_optimization.cc

LogicalResult BroadcastOpLowering::matchAndRewrite(
    shape::BroadcastOp op, mlir::PatternRewriter& rewriter) const {
  ShapeComponentAnalysis shape_component_analysis;
  auto new_broadcast = simplifyBroadcast(
      shape_component_analysis, op.getShapes(), op.getLoc(), &rewriter);
  if (!new_broadcast) return failure();
  rewriter.replaceOp(op, {*new_broadcast});
  return success();
}

Project: tensorflow
File: tf_cpurt_symbolic_shape_optimization.cc

LogicalResult DynamicBroadcastInDimOpLowering::matchAndRewrite(
    mhlo::DynamicBroadcastInDimOp op, mlir::PatternRewriter& rewriter) const {
  MLIRContext* ctx = getContext();

  auto in_type = op.operand().getType().dyn_cast<RankedTensorType>();
  auto out_type = op.getResult().getType().dyn_cast<RankedTensorType>();
  if (!in_type || !out_type) return failure();

  // Check that broadcast is right-aligned (numpy style), so that operand
  // dimensions broadcasted to match inner-most dimensions of the output.
  auto bcast_dims = op.broadcast_dimensions().getValues<int64_t>();
  auto expected_bcast_dims = llvm::seq<int64_t>(
      out_type.getRank() - in_type.getRank(), out_type.getRank());
  if (!llvm::equal(bcast_dims, expected_bcast_dims)) return failure();

  ShapeComponentAnalysis shape_component_analysis;
  auto input_map = isNonExpandingBroadcast(
      shape_component_analysis, op.operand(), op.output_dimensions());
  if (!input_map) return failure();

  // Resolve dynamic output dimensions for the `linalg.init_tensor` operation.
  SmallVector<Value> output_dyn_dimensions;
  Location loc = op.getLoc();
  int64_t rank = out_type.getRank();
  for (size_t d = 0; d < rank; ++d) {
    int64_t output_dim = out_type.getShape()[d];

    // Skip static output dimensions, they will be resolved from the shape.
    if (output_dim >= 0) continue;

    // Resolve the dynamic size of the output dimension.
    Value output_dyn_dim = rewriter.create<tensor::ExtractOp>(
        loc, op.output_dimensions(),
        ValueRange{rewriter.create<ConstantIndexOp>(loc, d)});

    // Symbolic shape analysis might have given us an i32 or i64. Cast to index.
    if (!output_dyn_dim.getType().isIndex())
      output_dyn_dim = rewriter.create<IndexCastOp>(loc, output_dyn_dim,
                                                    rewriter.getIndexType());

    output_dyn_dimensions.push_back(output_dyn_dim);
  }

  // Create a linalg.tensor_init operation to initialize output.
  Value init = rewriter.create<linalg::InitTensorOp>(loc, output_dyn_dimensions,
                                                     out_type.getShape(),
                                                     out_type.getElementType());

  // Output indexing map is an identity with `rank` number of loops.
  AffineMap output_map = AffineMap::getMultiDimIdentityMap(rank, ctx);

  // All iterators are parallel.
  SmallVector<llvm::StringRef> iterator_types(rank, "parallel");

  rewriter.replaceOpWithNewOp<linalg::GenericOp>(
      op, /*resultTensorTypes=*/TypeRange{init.getType()},
      /*inputs=*/ValueRange{op.operand()},
      /*outputs=*/ValueRange{init},
      /*indexingMaps=*/llvm::makeArrayRef({*input_map, output_map}),
      /*iteratorTypes=*/iterator_types,
      [&](OpBuilder& nested_builder, Location nested_loc, ValueRange args) {
        nested_builder.create<linalg::YieldOp>(nested_loc, args[0]);
      });

  return success();
}

Project: tensorflow
File: tf_cpurt_symbolic_shape_optimization.cc

  explicit SymbolicShapeOptimizationPass(bool constraints_only) {
    this->optimize_only_constraints = constraints_only;
  }

Project: tensorflow
File: tf_cpurt_symbolic_shape_optimization.cc

CstrBroadcastableOpLowering::CstrBroadcastableOpLowering(MLIRContext* ctx)
    : Base(ctx) {}