[xla] Optimize ShapeTree construction time for arrays

ezhulenev · tensorflower-gardener · commit 2ddb5fdc4faf · 2025-07-23T10:07:39.000-07:00
1. Add inplace constructor that avoids InlinedVector copy constructor
2. Do not compute index table for non-tuple types

PiperOrigin-RevId: 786322120
diff --git a/third_party/xla/xla/BUILD b/third_party/xla/xla/BUILD
@@ -1037,6 +1037,7 @@ cc_library(
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/types:span",
+        "@com_google_absl//absl/utility",
     ],
 )
 
diff --git a/third_party/xla/xla/shape_tree.cc b/third_party/xla/xla/shape_tree.cc
@@ -16,32 +16,22 @@ limitations under the License.
 #include "xla/shape_tree.h"
 
 #include <cstddef>
-#include <cstdint>
 
 #include "absl/base/optimization.h"
 #include "absl/log/check.h"
 #include "absl/types/span.h"
 #include "xla/shape.h"
-#include "xla/shape_util.h"
-#include "xla/tsl/platform/logging.h"  // IWYU pragma: keep
 
-namespace xla {
-namespace internal {
+namespace xla::internal {
 
-// Computes the total size of all nested tuples in the given shape.
-//
-// If `is_known_tuple` is true, then the shape is known to be a tuple, and we
-// can skip the run time check for `IsTuple()`.
-template <bool is_known_tuple = false>
+// Computes the total size of all nested tuples in the given tuple shape.
 static size_t IndexTableTuplesSize(const Shape& shape) {
-  if (!is_known_tuple && ABSL_PREDICT_TRUE(!shape.IsTuple())) {
-    return 0;
-  }
+  DCHECK(shape.IsTuple()) << "Shape must be a tuple";
 
   size_t size = shape.tuple_shapes().size();
   for (const Shape& subshape : shape.tuple_shapes()) {
     if (ABSL_PREDICT_FALSE(subshape.IsTuple())) {
-      size += IndexTableTuplesSize</*is_known_tuple=*/true>(subshape);
+      size += IndexTableTuplesSize(subshape);
     }
   }
 
@@ -75,22 +65,18 @@ static void InitializeIndexTable(const Shape& shape,
   }
 }
 
-IndexTable::IndexTable(const Shape& shape)
-    : entries_(1 + IndexTableTuplesSize(shape)) {
+IndexTable::IndexTable(const Shape& shape) {
+  if (!shape.IsTuple()) {
+    return;
+  }
+
+  // Allocate storage for the index table.
+  entries_.emplace(1 + IndexTableTuplesSize(shape));
+
   size_t next_node_id = 0;
   size_t next_children_start_index = 1;
-  InitializeIndexTable(shape, absl::MakeSpan(entries_), 0, next_node_id,
+  InitializeIndexTable(shape, absl::MakeSpan(*entries_), 0, next_node_id,
                        next_children_start_index);
 }
 
-const IndexTable::Entry& IndexTable::operator[](ShapeIndexView index) const {
-  const Entry* result = &entries_.front();
-  for (int64_t i : index) {
-    CHECK_GE(result->children_start_id, 0);
-    result = &entries_[result->children_start_id + i];
-  }
-  return *result;
-}
-
-}  // namespace internal
-}  // namespace xla
+}  // namespace xla::internal
diff --git a/third_party/xla/xla/shape_tree.h b/third_party/xla/xla/shape_tree.h
@@ -21,6 +21,7 @@ limitations under the License.
 #include <cstdint>
 #include <iterator>
 #include <memory>
+#include <optional>
 #include <type_traits>
 #include <utility>
 
@@ -30,6 +31,7 @@ limitations under the License.
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "absl/types/span.h"
+#include "absl/utility/utility.h"
 #include "xla/shape.h"
 #include "xla/shape_util.h"
 #include "xla/tsl/lib/gtl/iterator_range.h"
@@ -55,12 +57,25 @@ class IndexTable {
   IndexTable() = default;
   explicit IndexTable(const Shape& shape);
 
-  bool empty() const { return entries_.empty(); }
+  const Entry& operator[](ShapeIndexView index) const {
+    static constexpr Entry kRootEntry = {0, -1};
 
-  const Entry& operator[](ShapeIndexView index) const;
+    if (!entries_.has_value()) {
+      DCHECK(index.empty());
+      return kRootEntry;
+    }
+
+    const Entry* result = &entries_->front();
+    for (int64_t i : index) {
+      DCHECK_GE(result->children_start_id, 0);
+      result = &(*entries_)[result->children_start_id + i];
+    }
+    return *result;
+  }
 
  private:
-  absl::InlinedVector<Entry, 1> entries_;
+  // Entries are computed only if the shape is a tuple.
+  std::optional<absl::InlinedVector<Entry, 1>> entries_;
 };
 
 }  // namespace internal
@@ -111,14 +126,14 @@ class ShapeTree {
       : ShapeTree(std::make_shared<Shape>(std::move(shape))) {}
 
   explicit ShapeTree(const Shape* shape)
-      : ShapeTree(shape, CreateNodes(*shape)) {}
+      : ShapeTree(absl::in_place_t{}, shape) {}
 
   // Create ShapeTree with the given shape, and init_value for all nodes.
   ShapeTree(Shape shape, const T& init_value)
       : ShapeTree(std::make_shared<Shape>(std::move(shape)), init_value) {}
 
   ShapeTree(const Shape* shape, const T& init_value)
-      : ShapeTree(shape, CreateNodes(*shape, init_value)) {}
+      : ShapeTree(absl::in_place_t{}, shape, init_value) {}
 
   // Returns the data element associated with the array in the shape at the
   // given index (see ShapeUtil::GetSubshape for how indexes are defined).
@@ -375,21 +390,20 @@ class ShapeTree {
   }
 
   template <typename... Ts>
-  static Nodes CreateNodes(const Shape& shape, Ts&&... args) {
-    Nodes nodes;
-    nodes.reserve(ShapeUtil::SubshapeCount(shape));
+  ShapeTree(absl::in_place_t, const Shape* shape, Ts&&... args)
+      : index_table_(*shape), shape_(shape) {
+    nodes_.reserve(ShapeUtil::SubshapeCount(*shape));
     ShapeUtil::ForEachSubshape(
-        shape, [&](const Shape&, const ShapeIndex& index) {
-          nodes.emplace_back(index, T(std::forward<Ts>(args)...));
+        *shape, [&](const Shape&, const ShapeIndex& index) {
+          nodes_.emplace_back(index, T(std::forward<Ts>(args)...));
         });
-    return nodes;
   }
 
   // The nodes in this shape tree.
   Nodes nodes_;
 
   // Index table for node lookups. Each entry contains the index of the first
-  // child of the node at that index, or -1 for leaf nodes. Evaluated lazily.
+  // child of the node at that index, or -1 for leaf nodes.
   IndexTable index_table_;
 
   // If we own our Shape, this field contains it, and shape_ is a pointer into
diff --git a/third_party/xla/xla/shape_tree_test.cc b/third_party/xla/xla/shape_tree_test.cc
@@ -654,7 +654,7 @@ void BM_Iterate(::testing::benchmark::State& state) {
 }
 
 #define BENCHMARK_WITH_ARGS(name) \
-  BENCHMARK(name)->ArgPair(2, 8)->ArgPair(1, 1000)
+  BENCHMARK(name)->ArgPair(0, 0)->ArgPair(2, 8)->ArgPair(1, 1000)
 
 BENCHMARK_WITH_ARGS(BM_Construct);
 BENCHMARK_WITH_ARGS(BM_ConstructUnowned);

Original file line number	Diff line number	Diff line change
`@@ -1037,6 +1037,7 @@ cc_library(`
`1037`	`1037`	`"@com_google_absl//absl/status",`
`1038`	`1038`	`"@com_google_absl//absl/status:statusor",`
`1039`	`1039`	`"@com_google_absl//absl/types:span",`
	`1040`	`+ "@com_google_absl//absl/utility",`
`1040`	`1041`	`],`
`1041`	`1042`	`)`
`1042`	`1043`
Original file line number	Diff line number	Diff line change
`@@ -654,7 +654,7 @@ void BM_Iterate(::testing::benchmark::State& state) {`
`654`	`654`	`}`
`655`	`655`
`656`	`656`	`#define BENCHMARK_WITH_ARGS(name) \`
`657`		`- BENCHMARK(name)->ArgPair(2, 8)->ArgPair(1, 1000)`
	`657`	`+ BENCHMARK(name)->ArgPair(0, 0)->ArgPair(2, 8)->ArgPair(1, 1000)`
`658`	`658`
`659`	`659`	`BENCHMARK_WITH_ARGS(BM_Construct);`
`660`	`660`	`BENCHMARK_WITH_ARGS(BM_ConstructUnowned);`