Rename some XLA internal API from rank() to num_dimensions() to avoid confusion.

PiperOrigin-RevId: 738831506

Author: Google-ML-Automation

xla/backends/gpu/codegen/triton/tma_utils_test.cc

@@ -70,7 +70,7 @@ TEST(Create2DTmaDescriptorTest, ValidInputReturnCorrectDescriptor) {
       TmaDescriptor tma_desc,
       Create2DTmaDescriptor(global_shape, block_shape, element_type));
   EXPECT_EQ(tma_desc.element_size(), 4);
-  EXPECT_EQ(tma_desc.rank(), 2);
+  EXPECT_EQ(tma_desc.num_dimensions(), 2);
   EXPECT_THAT(tma_desc.global_dims(), ElementsAre(128, 256));
   EXPECT_THAT(tma_desc.global_strides(), ElementsAre(128 * 4));
   EXPECT_THAT(tma_desc.box_dims(), ElementsAre(32, 64));

xla/mlir/tools/mlir_interpreter/dialects/linalg.cc

@@ -149,7 +149,7 @@ llvm::SmallVector<InterpreterValue> Map(InterpreterState& state,
       isa<TensorType>(op.getInit().getType()) ? init.Clone() : init;
 
   InterpreterScope scope(state);
-  SmallVector<int64_t> ivs(output.View().Rank());
+  SmallVector<int64_t> ivs(output.View().num_dimensions());
   scope.SetSideChannel(std::make_shared<IterationIndexSideChannel>(ivs));
   for (const auto& indices : output.View().Indices()) {
     std::copy(indices.begin(), indices.end(), ivs.begin());

xla/mlir/tools/mlir_interpreter/dialects/memref.cc

@@ -129,15 +129,15 @@ InterpreterValue Subview(InterpreterState& state, memref::SubViewOp subview,
     return {};
   }
 
-  if (subview.getResult().getType().getRank() == out_view.Rank()) {
+  if (subview.getResult().getType().getRank() == out_view.num_dimensions()) {
     return out;
   }
 
   auto shape = subview.getResult().getType().getShape();
   // TODO(jreiffers): Check why subview.getDroppedDims() yields the wrong shape
   // here for 1x2x2x3 (-> 1x2x1x3) -> 1x2x3 (claiming 0 is dropped).
   int64_t dim = 0;
-  while (dim < out_view.Rank() && dim < shape.size()) {
+  while (dim < out_view.num_dimensions() && dim < shape.size()) {
     if (shape[dim] != 1 && out_view.sizes[dim] == 1) {
       out_view.sizes.erase(out_view.sizes.begin() + dim);
       out_view.strides.erase(out_view.strides.begin() + dim);
@@ -147,7 +147,7 @@ InterpreterValue Subview(InterpreterState& state, memref::SubViewOp subview,
       ++dim;
     }
   }
-  while (dim < out_view.Rank()) {
+  while (dim < out_view.num_dimensions()) {
     assert(out_view.sizes.back() == 1 && "expected remaining dims to be 1");
     out_view.sizes.pop_back();
     out_view.strides.pop_back();

xla/mlir/tools/mlir_interpreter/dialects/mhlo.cc

@@ -321,13 +321,13 @@ llvm::SmallVector<InterpreterValue> Gather(
   auto& operand = args[0];
   auto& start_indices = args[1];
   const auto& operand_view = operand.View();
-  int64_t operand_rank = operand_view.Rank();
+  int64_t operand_rank = operand_view.num_dimensions();
 
   // Make a fake BufferView for the start indices.
   BufferView start_indices_view = start_indices.View();
-  auto output_rank =
-      static_cast<int64_t>(start_indices_view.Rank() + offset_dims.size());
-  if (index_vector_dim < start_indices_view.Rank()) {
+  auto output_rank = static_cast<int64_t>(start_indices_view.num_dimensions() +
+                                          offset_dims.size());
+  if (index_vector_dim < start_indices_view.num_dimensions()) {
     --output_rank;
     start_indices_view.sizes[index_vector_dim] = 1;
   }
@@ -365,7 +365,7 @@ llvm::SmallVector<InterpreterValue> Gather(
   for (auto start_indices_index : start_indices_view.Indices()) {
     SmallVector<int64_t> operand_base_indices(operand_rank);
     for (auto [i, dim] : llvm::enumerate(start_index_map)) {
-      if (index_vector_dim < start_indices_view.Rank()) {
+      if (index_vector_dim < start_indices_view.num_dimensions()) {
         start_indices_index[index_vector_dim] = static_cast<int64_t>(i);
       }
       operand_base_indices[dim] = std::max<int64_t>(
@@ -410,9 +410,9 @@ llvm::SmallVector<InterpreterValue> Scatter(
   auto update_window_dims = dims.getUpdateWindowDims();
 
   auto input_view = n_inputs.front().View();
-  int64_t operand_rank = input_view.Rank();
-  int64_t updates_rank = n_updates.front().View().Rank();
-  int64_t indices_rank = scatter_indices.View().Rank();
+  int64_t operand_rank = input_view.num_dimensions();
+  int64_t updates_rank = n_updates.front().View().num_dimensions();
+  int64_t indices_rank = scatter_indices.View().num_dimensions();
 
   llvm::SmallVector<int64_t> batch_dims;
   for (int64_t dim = 0; dim < operand_rank; ++dim) {
@@ -689,23 +689,23 @@ InterpreterValue DotGeneralImpl(InterpreterValue& lhs, InterpreterValue& rhs,
   for (int64_t lhs_dim : lhs_batch) {
     dimensions.push_back(lhsv.sizes[lhs_dim]);
   }
-  for (int64_t i = 0; i < lhsv.Rank(); i++) {
+  for (int64_t i = 0; i < lhsv.num_dimensions(); i++) {
     if (!llvm::is_contained(lhs_contracting, i) &&
         !llvm::is_contained(lhs_batch, i)) {
       dimensions.push_back(lhsv.sizes[i]);
       lhs_non_batch.push_back(i);
     }
   }
-  for (int64_t i = 0; i < rhs.View().Rank(); i++) {
+  for (int64_t i = 0; i < rhs.View().num_dimensions(); i++) {
     if (!llvm::is_contained(rhs_contracting, i) &&
         !llvm::is_contained(rhs_batch, i)) {
       dimensions.push_back(rhsv.sizes[i]);
       rhs_non_batch.push_back(i);
     }
   }
 
-  SmallVector<int64_t> lhs_index(lhsv.Rank());
-  SmallVector<int64_t> rhs_index(rhsv.Rank());
+  SmallVector<int64_t> lhs_index(lhsv.num_dimensions());
+  SmallVector<int64_t> rhs_index(rhsv.num_dimensions());
   SmallVector<int64_t> output_index(dimensions.size());
   auto output = lhs.TypedAlike(dimensions);
 
@@ -778,7 +778,7 @@ InterpreterValue Dot(InterpreterState& state, mhlo::DotOp op,
   auto ty = cast<ShapedType>(op->getResultTypes()[0]);
   auto result = lhs.TypedAlike(ty.getShape());
 
-  if (lhs.View().Rank() == 1 && rhs.View().Rank() == 1) {
+  if (lhs.View().num_dimensions() == 1 && rhs.View().num_dimensions() == 1) {
     DispatchScalarType(ty, [&](auto dummy) {
       using T = decltype(dummy);
       using TT = TensorOrMemref<T>;
@@ -791,7 +791,8 @@ InterpreterValue Dot(InterpreterState& state, mhlo::DotOp op,
       }
       std::get<TT>(result.storage).at({}) += product;
     });
-  } else if (lhs.View().Rank() == 2 && rhs.View().Rank() == 1) {
+  } else if (lhs.View().num_dimensions() == 2 &&
+             rhs.View().num_dimensions() == 1) {
     DispatchScalarType(ty, [&](auto dummy) {
       using TT = TensorOrMemref<decltype(dummy)>;
       auto lhs_tensor = std::get<TT>(lhs.storage);
@@ -803,7 +804,8 @@ InterpreterValue Dot(InterpreterState& state, mhlo::DotOp op,
         }
       }
     });
-  } else if (lhs.View().Rank() == 2 && rhs.View().Rank() == 2) {
+  } else if (lhs.View().num_dimensions() == 2 &&
+             rhs.View().num_dimensions() == 2) {
     DispatchScalarType(ty, [&](auto dummy) {
       using TT = TensorOrMemref<decltype(dummy)>;
       auto lhs_tensor = std::get<TT>(lhs.storage);

xla/mlir/tools/mlir_interpreter/dialects/tensor.cc

@@ -126,7 +126,7 @@ llvm::SmallVector<InterpreterValue> ExtractSlice(
   int64_t dim = 0;
   const auto& result_sizes = extract.getResultType().getShape();
   const auto& static_sizes = extract.getStaticSizes();
-  while (dim < out_view.Rank()) {
+  while (dim < out_view.num_dimensions()) {
     if (static_sizes[num_dropped + dim] == 1 &&
         (dim >= result_sizes.size() || result_sizes[dim] != 1)) {
       out_view.sizes.erase(out_view.sizes.begin() + dim);

xla/mlir/tools/mlir_interpreter/dialects/util.cc

@@ -121,7 +121,7 @@ InterpreterValue TransposeImpl(const InterpreterValue& in,
 
 int64_t DimImpl(const InterpreterValue& in, int64_t index,
                 InterpreterState& state) {
-  if (index < 0 || index >= in.View().Rank()) {
+  if (index < 0 || index >= in.View().num_dimensions()) {
     state.AddFailure("dimension index out of bounds");
     return 0;
   }

xla/mlir/tools/mlir_interpreter/dialects/vector.cc

@@ -299,7 +299,7 @@ InterpreterValue Extract(InterpreterState& state, vector::ExtractOp extract,
   for (int64_t offset : extract.getStaticPosition()) {
     state.CheckSuccess(result_view.Slice(0, offset), "index out of bounds");
   }
-  return result_view.Rank() == 0 ? result.ExtractElement({}) : result;
+  return result_view.num_dimensions() == 0 ? result.ExtractElement({}) : result;
 }
 
 InterpreterValue ExtractElement(InterpreterState& state,
@@ -635,7 +635,7 @@ InterpreterValue Shuffle(InterpreterState& state, vector::ShuffleOp shuffle,
   result_view.is_vector = true;
 
   auto mask = shuffle.getMask();
-  bool is_zero_dim = v0.View().Rank() == 0;
+  bool is_zero_dim = v0.View().num_dimensions() == 0;
   int64_t size0 = is_zero_dim ? 1 : v0.View().sizes[0];
   for (auto [dst_index, src_index] : llvm::enumerate(mask)) {
     auto src = src_index < size0 ? v0 : v1;
@@ -698,9 +698,9 @@ std::optional<InterpreterValue> ExtractMemorySlice(
   auto mem_slice = memory;
   auto& mem_slice_view = mem_slice.View();
   auto& vector_view = vector.View();
-  for (int64_t i = 0; i < mem_slice_view.Rank(); ++i) {
+  for (int64_t i = 0; i < mem_slice_view.num_dimensions(); ++i) {
     bool found = false;
-    for (int64_t j = 0; !found && j < vector_view.Rank(); ++j) {
+    for (int64_t j = 0; !found && j < vector_view.num_dimensions(); ++j) {
       if (map.getResult(j).isFunctionOfDim(i)) {
         int64_t size = mem_slice_view.sizes[i] - offsets[i];
         bool is_in_bounds = size >= vector_view.sizes[j];
@@ -801,7 +801,8 @@ llvm::SmallVector<InterpreterValue> TransferWrite(
   }
 
   const auto& src_view = src.View();
-  assert(transfer.getPermutationMap().getNumResults() == src_view.Rank() &&
+  assert(transfer.getPermutationMap().getNumResults() ==
+             src_view.num_dimensions() &&
          "expected matching number of results");
 
   dst = transfer.getSource().getType().isa<TensorType>() ? dst.Clone() : dst;
@@ -832,7 +833,7 @@ InterpreterValue Transpose(InterpreterState&, vector::TransposeOp transpose,
 
 InterpreterValue TypeCast(InterpreterState&, vector::TypeCastOp,
                           InterpreterValue vector) {
-  vector.View().num_vector_dims = vector.View().Rank();
+  vector.View().num_vector_dims = vector.View().num_dimensions();
   return vector;
 }
 

xla/mlir/tools/mlir_interpreter/framework/interpreter_value.cc

@@ -83,7 +83,7 @@ struct InterpreterValuePrinter {
     } else {
       os << TypeStr::Get(T{}) << ">: ";
     }
-    SmallVector<int64_t> indices(t.view.Rank() +
+    SmallVector<int64_t> indices(t.view.num_dimensions() +
                                  t.view.num_vector_dims.value_or(0));
     std::function<void(int64_t)> print;
     print = [&](int64_t dim) {

xla/mlir/tools/mlir_interpreter/framework/tensor_or_memref.cc

@@ -77,22 +77,22 @@ SmallVector<int64_t> BufferView::GetStridesForLayout(ArrayRef<int64_t> sizes,
 }
 
 LogicalResult BufferView::Slice(int64_t dim_index, int64_t dim_offset) {
-  llvm::SmallVector<int64_t> offsets(Rank(), 0);
+  llvm::SmallVector<int64_t> offsets(num_dimensions(), 0);
   offsets[dim_index] = dim_offset;
   if (auto new_offset = GetPhysicalIndex(offsets)) {
     offset = *new_offset;
   } else {
     return failure();
   }
-  if (dim_index >= Rank()) --*num_vector_dims;
+  if (dim_index >= num_dimensions()) --*num_vector_dims;
   strides.erase(strides.begin() + dim_index);
   sizes.erase(sizes.begin() + dim_index);
   return success();
 }
 
 LogicalResult BufferView::Slice(int64_t dim_index, int64_t dim_offset,
                                 int64_t dim_size, int64_t dim_stride) {
-  llvm::SmallVector<int64_t> offsets(Rank(), 0);
+  llvm::SmallVector<int64_t> offsets(num_dimensions(), 0);
   offsets[dim_index] = dim_offset;
   if (dim_size == 0) {
     offset = 0;

xla/mlir/tools/mlir_interpreter/framework/tensor_or_memref.h

@@ -73,7 +73,9 @@ struct BufferView {
   std::optional<int64_t> num_vector_dims = std::nullopt;
   bool is_vector = false;
 
-  int64_t Rank() const { return sizes.size() - num_vector_dims.value_or(0); }
+  int64_t num_dimensions() const {
+    return sizes.size() - num_vector_dims.value_or(0);
+  }
 
   // Removes the dimension from the view. If you need to keep it, use the
   // overload below with dim_size = 1.
@@ -154,7 +156,7 @@ struct BufferView {
       return {
           view_,
           llvm::SmallVector<int64_t>(
-              view_->Rank() +
+              view_->num_dimensions() +
               (include_vector_dims_ ? view_->num_vector_dims.value_or(0) : 0)),
           include_vector_dims_};
     }
@@ -314,8 +316,10 @@ struct TensorOrMemref {
   TensorOrMemref VectorAt(ArrayRef<int64_t> indices) const {
     auto offset = view.GetPhysicalIndex(indices);
     BufferView subview;
-    subview.strides = {view.strides.begin() + view.Rank(), view.strides.end()};
-    subview.sizes = {view.sizes.begin() + view.Rank(), view.sizes.end()};
+    subview.strides = {view.strides.begin() + view.num_dimensions(),
+                       view.strides.end()};
+    subview.sizes = {view.sizes.begin() + view.num_dimensions(),
+                     view.sizes.end()};
     if (offset) {
       subview.offset = *offset;
     } else {

xla/mlir/tools/mlir_interpreter/framework/tests/tensor_or_memref_test.cc

@@ -56,7 +56,7 @@ std::optional<int64_t> GetCollapsedStrideNaive(llvm::ArrayRef<int64_t> dims,
   // Find all physical indices for the dimensions.
   llvm::SmallBitVector v(view.GetNumElements());
   for (const auto& indices : f.Indices()) {
-    SmallVector<int64_t> view_indices(view.Rank());
+    SmallVector<int64_t> view_indices(view.num_dimensions());
     for (auto [dim, index] : llvm::zip(dims, indices)) {
       view_indices[dim] = index;
     }
@@ -83,9 +83,9 @@ TEST(TensorOrMemrefTest, CollapsedStride) {
                   .strides = BufferView::GetDefaultStrides({1, 2, 3, 1, 5})};
 
   auto check_all = [&]() {
-    for (int64_t i = 0; i < (1 << view.Rank()); ++i) {
+    for (int64_t i = 0; i < (1 << view.num_dimensions()); ++i) {
       SmallVector<int64_t> dims;
-      for (int64_t dim = 0; dim < view.Rank(); ++dim) {
+      for (int64_t dim = 0; dim < view.num_dimensions(); ++dim) {
         if (i & (1 << dim)) dims.push_back(dim);
       }