[XLA:FFI] Add a BufferShape alias.

This change introduces a `BufferShape` alias for custom call users' convenience. This new alias eliminates the need for verbose code such as `using Shape = decltype(Buffer<PrimitiveType::F32>::dimensions)`. PiperOrigin-RevId: 632119253
tensorflow · May 9, 2024 · 5d4d94c · 5d4d94c
1 parent d84d042
commit 5d4d94c
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/third_party/xla/xla/ffi/ffi.h b/third_party/xla/xla/ffi/ffi.h
@@ -56,9 +56,11 @@ struct CalledComputation {};  // binds `HloComputation*`
 //===----------------------------------------------------------------------===//
 
 struct BufferBase {
+  using Shape = absl::Span<const int64_t>;
+
   PrimitiveType dtype;
   se::DeviceMemoryBase data;
-  absl::Span<const int64_t> dimensions;
+  Shape dimensions;
 };
 
 namespace internal {
@@ -72,8 +74,10 @@ using NativeType = typename primitive_util::PrimitiveTypeToNative<dtype>::type;
 
 template <PrimitiveType dtype, size_t rank = internal::kDynamicRank>
 struct Buffer {
+  using Shape = BufferBase::Shape;
+
   se::DeviceMemory<internal::NativeType<dtype>> data;
-  absl::Span<const int64_t> dimensions;
+  Shape dimensions;
 };
 
 // clang-format off