message: [NFC] Do the code review comments.
diff: @@ -1358,6 +1358,9 @@ Status VerifyBatchNormForThunkEmission(
  // fusion with only point-wise operations, scalar broadcasting and row
  // broadcasting, we can trigger a kernel that vectorize the row loads.
  // This speed up the kernel, in particular on A100.
  +// Returns a pair<bool, int>. The bool mean should we try to enable
  +// row vectorization. The int is the number of inputs with the higher
  +// rank.
  std::pair<bool, int> RowVectorizationEnabled(mlir::lmhlo::FusionOp fusion) {
  const auto is_row_major = [](mlir::Value value) {
  // Only tested when the inputs are row-major. So only
  @@ -1380,13 +1383,12 @@ std::pair<bool, int> RowVectorizationEnabled(mlir::lmhlo::FusionOp fusion) {
  // We also detect at the same time if there is a row broadcasting
  // operation.
  bool some_row_broadcasting = false;
  - auto out_rank = TypeToShape(fusion.getFusionResults()[0].getType()).rank();
  - int nb_big_input = 0;
  + auto out_rank = fusion.getFusionResults()[0].getType().cast<mlir::ShapedType>().getRank();
  + int num_big_inputs = 0;
  for (mlir::Operation& op : fusion.region().front()) {
  - if (mlir::isa<mlir::memref::TensorLoadOp>(op)) {
  - auto load = mlir::dyn_cast<mlir::memref::TensorLoadOp>(op);
  - auto rank = TypeToShape(load.getResult().getType()).rank();
  - nb_big_input += rank == out_rank;
  + if (auto load = mlir::dyn_cast<mlir::memref::TensorLoadOp>(op)) {
  + auto rank = load.getResult().getType().cast<mlir::ShapedType>().getRank();
  + num_big_inputs += rank == out_rank;
  continue;
  } else if (mlir::isa<mlir::memref::TensorStoreOp,
  mlir::lmhlo::TerminatorOp, mlir::mhlo::ReturnOp,
  @@ -1420,7 +1422,7 @@ std::pair<bool, int> RowVectorizationEnabled(mlir::lmhlo::FusionOp fusion) {
  }
  // Trigger only when there is a row broadcasting.
  - return std::make_pair(row_vectorized && some_row_broadcasting, nb_big_input);
  + return std::make_pair(row_vectorized && some_row_broadcasting, num_big_inputs);
  }
  } // namespace
  @@ -1926,9 +1928,10 @@ Status IrEmitterUnnested::EmitLoopFusionFromMlir(
  }
  auto row_vec_enabled = RowVectorizationEnabled(fusion);
  - bool row_vectorized = std::get<0>(row_vec_enabled);
  - int nb_big_inputs = std::get<1>(row_vec_enabled);
  - bool few_waves = [fusion, row_vectorized, nb_big_inputs]() mutable {
  + bool row_vectorized;
  + int num_big_inputs;
  + std::tie(row_vectorized, num_big_inputs) = RowVectorizationEnabled(fusion);
  + bool few_waves = [fusion, row_vectorized, num_big_inputs]() mutable {
  for (mlir::Operation& op : fusion.region().front()) {
  if (mlir::isa<mlir::memref::TensorLoadOp, mlir::memref::TensorStoreOp,
  mlir::lmhlo::TerminatorOp, mlir::mhlo::ReturnOp,
  @@ -1942,7 +1945,7 @@ Status IrEmitterUnnested::EmitLoopFusionFromMlir(
  if (auto broadcast = mlir::dyn_cast<mlir::mhlo::BroadcastInDimOp>(op)) {
  if (broadcast.broadcast_dimensions().size() == 0 ||
  // More then 2 bit inputs cause one speed regression.
  - (row_vectorized && nb_big_inputs <= 2)) {
  + (row_vectorized && num_big_inputs <= 2)) {
  continue;
  }
  }
  