Overhead introduced by constructing buffers for bfloat16 dtypes on read

[dfm-anthropic]

I've been running some benchmarks to test bandwidth on tensorstore reads, and I've found significant overhead for non-standard dtypes (e.g. bfloat16) where tensorstore seems to initialize the array to zero before overwriting it with the data. Here's the benchmark that I've been running:

import time
import tensorstore as ts

dataset = ts.open(
    {
        "driver": "zarr3",
        "kvstore": {
            "driver": "gcs",
            "bucket": "...",
            "path": "...",
        },
        "open": True,
        "create": False,
    }
).result()

start = time.perf_counter()
array = dataset.read().result()
duration = time.perf_counter() - start

gb = array.nbytes / 1024**3
bandwidth = gb / duration
print(f"{gb:.2f} GB loaded in {duration:.2f} s ({bandwidth:.2f} GB/s)")

where the tensor that I'm loading is about 55 GB of bfloat16 data.

If I run this benchmark using the current master branch of tensorstore, I see:

54.93 GB loaded in 58.35 s (0.94 GB/s)

but, if I disable zero initialization on reads (more details below), I get:

54.93 GB loaded in 31.47 s (1.75 GB/s)

which suggests that there's about 30s of overhead introduced by the current implementations construction of the 55 GB array!

I've narrowed the offending line down to:

tensorstore/tensorstore/data_type.cc

Line 86 in 367ef7d

r->construct(n, ptr.get());

I can avoid this overhead by hacking tensorstore to skip that r->construct(...) or replace this line with:

-   constexpr BFloat16() : rep_(0) {}
+   BFloat16() = default;

for bfloat16 specifically.

I wanted to ask here if anyone had advice about what would be the preferred approach for safely avoiding this overhead. I'd be very happy to open a PR if we decide that there's something general to do here. Thanks!!

[laramiel]

It seems to me that we could probably use an = default() constructor on all the custom float types.

[laramiel]

So the challenge here is going to be getting default and value initialization working properly.
The correct behavior should be:

void test_BFloat16_default_value_init() {
  const BFloat16 zero(0);

  BFloat16 d;
  // d is default initialized.
  // EXPECT_NOT_ALWAYS_EQ(d, zero); 

  BFloat16 v = Bfloat16();
  // v is value initialized.
  EXPECT_EQ(v, zero);

  static BFloat16 sd;
  // sd is static, so it must be zero initialized.
  EXPECT_EQ(sd, zero);

  BFloat16 z{};
  // z is zero initialized.
  EXPECT_EQ(z, zero);
}

[dfm-anthropic]

I did some experiments, and this behavior (at least for bf16) does seem like it might already be supported by the switch to the default constructor in #261. But, I must admit that I'm not extremely confident with the theory here!

[laramiel]

After having more thoroughly reviewed the custom float implementation, I agree that it will likely work as expected.
I have a pending internal CL for some related tests, and will get to this pr as soon as I can after that.