[dawn][native] Don't handle availability in timestamp quantization shader

The timestamp quantization shader was used both to quantize, normalize
(to ns values) and zero-out timestamps in the query resolve buffer. The
zeroing-out was redundant because all backends either have guarantees
from the target API that the queries not available are 0, or have special
handling in the backend to zero-out unavailable queries already.

Remove that logic to simplify things a bit. This also allows removing
TimestampParams::first. The dispatch for the quantization is also
updated to be sized for exactly the number of queries that are resolved
instead of running on the whole buffer.

Bug: 499140183
Change-Id: I9b53b595c406d7e3fe05f1405010d64428d133f5
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/305055
Reviewed-by: Brandon Jones <bajones@chromium.org>
Reviewed-by: Kai Ninomiya <kainino@chromium.org>
Commit-Queue: Corentin Wallez <cwallez@chromium.org>
diff --git a/src/dawn/native/CommandEncoder.cpp b/src/dawn/native/CommandEncoder.cpp
index d744e9b..4b944e7 100644
--- a/src/dawn/native/CommandEncoder.cpp
+++ b/src/dawn/native/CommandEncoder.cpp
@@ -1065,29 +1065,13 @@
                                                    uint64_t destinationOffset) {
     DeviceBase* device = encoder->GetDevice();
 
-    // The availability got from query set is a reference to vector<bool>, need to covert
-    // bool to uint32_t due to a user input in pipeline must not contain a bool type in
-    // WGSL.
-    std::vector<uint32_t> availability{querySet->GetQueryAvailability().begin(),
-                                       querySet->GetQueryAvailability().end()};
-
-    // Timestamp availability storage buffer
-    BufferDescriptor availabilityDesc = {};
-    availabilityDesc.usage = wgpu::BufferUsage::Storage | wgpu::BufferUsage::CopyDst;
-    availabilityDesc.size = querySet->GetQueryCount() * sizeof(uint32_t);
-    Ref<BufferBase> availabilityBuffer;
-    DAWN_TRY_ASSIGN(availabilityBuffer, device->CreateBuffer(&availabilityDesc));
-
-    DAWN_TRY(device->GetQueue()->WriteBuffer(availabilityBuffer.Get(), 0, availability.data(),
-                                             availability.size() * sizeof(uint32_t)));
-
     const uint32_t quantization_mask = (device->IsToggleEnabled(Toggle::TimestampQuantization))
                                            ? kTimestampQuantizationMask
                                            : 0xFFFFFFFF;
 
     // Timestamp params uniform buffer
-    TimestampParams params(firstQuery, queryCount, static_cast<uint32_t>(destinationOffset),
-                           quantization_mask, device->GetTimestampPeriodInNS());
+    TimestampParams params(queryCount, static_cast<uint32_t>(destinationOffset), quantization_mask,
+                           device->GetTimestampPeriodInNS());
 
     BufferDescriptor parmsDesc = {};
     parmsDesc.usage = wgpu::BufferUsage::Uniform | wgpu::BufferUsage::CopyDst;
@@ -1097,7 +1081,7 @@
 
     DAWN_TRY(device->GetQueue()->WriteBuffer(paramsBuffer.Get(), 0, &params, sizeof(params)));
 
-    return EncodeConvertTimestampsToNanoseconds(encoder, destination, availabilityBuffer.Get(),
+    return EncodeConvertTimestampsToNanoseconds(encoder, queryCount, destination,
                                                 paramsBuffer.Get());
 }
 
diff --git a/src/dawn/native/QueryHelper.cpp b/src/dawn/native/QueryHelper.cpp
index a4688c6..6555ec4 100644
--- a/src/dawn/native/QueryHelper.cpp
+++ b/src/dawn/native/QueryHelper.cpp
@@ -46,13 +46,15 @@
 namespace {
 
 // Assert the offsets in dawn::native::TimestampParams are same with the ones in the shader
-static_assert(offsetof(dawn::native::TimestampParams, first) == 0);
-static_assert(offsetof(dawn::native::TimestampParams, count) == 4);
-static_assert(offsetof(dawn::native::TimestampParams, offset) == 8);
-static_assert(offsetof(dawn::native::TimestampParams, quantizationMask) == 12);
-static_assert(offsetof(dawn::native::TimestampParams, multiplier) == 16);
-static_assert(offsetof(dawn::native::TimestampParams, rightShift) == 20);
+static_assert(offsetof(dawn::native::TimestampParams, count) == 0);
+static_assert(offsetof(dawn::native::TimestampParams, offset) == 4);
+static_assert(offsetof(dawn::native::TimestampParams, quantizationMask) == 8);
+static_assert(offsetof(dawn::native::TimestampParams, multiplier) == 12);
+static_assert(offsetof(dawn::native::TimestampParams, rightShift) == 16);
 
+// TODO(https://crbug.com/465714204): Use immediates instead of a params buffer, which will make
+// EncodeConvertTimestampsToNanoseconds take a TimestampParams directly, and removes the need for
+// the duplicate queryCount parameter.
 static const char sConvertTimestampsToNanoseconds[] = DAWN_MULTILINE(
     struct Timestamp {
         low  : u32,
@@ -63,12 +65,7 @@
         t : array<Timestamp>
     }
 
-    struct AvailabilityArr {
-        v : array<u32>
-    }
-
     struct TimestampParams {
-        first  : u32,
         count  : u32,
         offset : u32,
         quantization_mask : u32,
@@ -77,8 +74,7 @@
     }
 
     @group(0) @binding(0) var<storage, read_write> timestamps : TimestampArr;
-    @group(0) @binding(1) var<storage, read> availability : AvailabilityArr;
-    @group(0) @binding(2) var<uniform> params : TimestampParams;
+    @group(0) @binding(1) var<uniform> params : TimestampParams;
 
     const sizeofTimestamp : u32 = 8u;
 
@@ -87,14 +83,6 @@
         if (GlobalInvocationID.x >= params.count) { return; }
 
         var index = GlobalInvocationID.x + params.offset / sizeofTimestamp;
-
-        // Return 0 for the unavailable value.
-        if (availability.v[GlobalInvocationID.x + params.first] == 0u) {
-            timestamps.t[index].low = 0u;
-            timestamps.t[index].high = 0u;
-            return;
-        }
-
         var timestamp = timestamps.t[index];
 
         // TODO(dawn:1250): Consider using the umulExtended and uaddCarry intrinsics once
@@ -146,15 +134,14 @@
 
         // Create binding group layout
         Ref<BindGroupLayoutBase> bgl;
-        DAWN_TRY_ASSIGN(
-            bgl, utils::MakeBindGroupLayout(
-                     device,
-                     {
-                         {0, wgpu::ShaderStage::Compute, kInternalStorageBufferBinding},
-                         {1, wgpu::ShaderStage::Compute, wgpu::BufferBindingType::ReadOnlyStorage},
-                         {2, wgpu::ShaderStage::Compute, wgpu::BufferBindingType::Uniform},
-                     },
-                     /* allowInternalBinding */ true));
+        DAWN_TRY_ASSIGN(bgl,
+                        utils::MakeBindGroupLayout(
+                            device,
+                            {
+                                {0, wgpu::ShaderStage::Compute, kInternalStorageBufferBinding},
+                                {1, wgpu::ShaderStage::Compute, wgpu::BufferBindingType::Uniform},
+                            },
+                            /* allowInternalBinding */ true));
 
         // Create pipeline layout
         Ref<PipelineLayoutBase> layout;
@@ -176,12 +163,11 @@
 
 }  // anonymous namespace
 
-TimestampParams::TimestampParams(uint32_t first,
-                                 uint32_t count,
+TimestampParams::TimestampParams(uint32_t count,
                                  uint32_t offset,
                                  uint32_t quantizationMask,
                                  float period)
-    : first(first), count(count), offset(offset), quantizationMask(quantizationMask) {
+    : count(count), offset(offset), quantizationMask(quantizationMask) {
     // The overall conversion happening, if p is the period, m the multiplier, s the shift, is::
     //
     //   m = round(p * 2^s)
@@ -205,8 +191,8 @@
 }
 
 MaybeError EncodeConvertTimestampsToNanoseconds(CommandEncoder* encoder,
+                                                uint32_t queryCount,
                                                 BufferBase* timestamps,
-                                                BufferBase* availability,
                                                 BufferBase* params) {
     DeviceBase* device = encoder->GetDevice();
     DAWN_ASSERT(device->IsLockedByCurrentThreadIfNeeded());
@@ -220,17 +206,14 @@
 
     // Create bind group after all binding entries are set.
     Ref<BindGroupBase> bindGroup;
-    DAWN_TRY_ASSIGN(
-        bindGroup,
-        utils::MakeBindGroup(device, layout, {{0, timestamps}, {1, availability}, {2, params}},
-                             UsageValidationMode::Internal));
+    DAWN_TRY_ASSIGN(bindGroup, utils::MakeBindGroup(device, layout, {{0, timestamps}, {1, params}},
+                                                    UsageValidationMode::Internal));
 
     // Create compute encoder and issue dispatch.
     Ref<ComputePassEncoder> pass = encoder->BeginComputePass();
     pass->APISetPipeline(pipeline);
     pass->APISetBindGroup(0, bindGroup.Get());
-    pass->APIDispatchWorkgroups(
-        static_cast<uint32_t>((timestamps->GetSize() / sizeof(uint64_t) + 7) / 8));
+    pass->APIDispatchWorkgroups((queryCount + 7) / 8);
     pass->APIEnd();
 
     return {};
diff --git a/src/dawn/native/QueryHelper.h b/src/dawn/native/QueryHelper.h
index 7086272..29eb37b 100644
--- a/src/dawn/native/QueryHelper.h
+++ b/src/dawn/native/QueryHelper.h
@@ -37,13 +37,8 @@
 class CommandEncoder;
 
 struct TimestampParams {
-    TimestampParams(uint32_t first,
-                    uint32_t count,
-                    uint32_t offset,
-                    uint32_t quantizationMask,
-                    float period);
+    TimestampParams(uint32_t count, uint32_t offset, uint32_t quantizationMask, float period);
 
-    uint32_t first;
     uint32_t count;
     uint32_t offset;
     uint32_t quantizationMask;
@@ -52,8 +47,8 @@
 };
 
 MaybeError EncodeConvertTimestampsToNanoseconds(CommandEncoder* encoder,
+                                                uint32_t queryCount,
                                                 BufferBase* timestamps,
-                                                BufferBase* availability,
                                                 BufferBase* params);
 
 }  // namespace dawn::native
diff --git a/src/dawn/tests/white_box/QueryInternalShaderTests.cpp b/src/dawn/tests/white_box/QueryInternalShaderTests.cpp
index f6888d2..f014f47 100644
--- a/src/dawn/tests/white_box/QueryInternalShaderTests.cpp
+++ b/src/dawn/tests/white_box/QueryInternalShaderTests.cpp
@@ -27,8 +27,6 @@
 
 #include <vector>
 
-#include "dawn/native/Buffer.h"
-#include "dawn/native/CommandEncoder.h"
 #include "dawn/native/QueryHelper.h"
 #include "dawn/tests/DawnTest.h"
 #include "dawn/utils/WGPUHelpers.h"
@@ -37,12 +35,12 @@
 namespace {
 
 void EncodeConvertTimestampsToNanoseconds(wgpu::CommandEncoder encoder,
+                                          uint32_t queryCount,
                                           wgpu::Buffer timestamps,
-                                          wgpu::Buffer availability,
                                           wgpu::Buffer params) {
     ASSERT_TRUE(native::EncodeConvertTimestampsToNanoseconds(
-                    native::FromAPI(encoder.Get()), native::FromAPI(timestamps.Get()),
-                    native::FromAPI(availability.Get()), native::FromAPI(params.Get()))
+                    native::FromAPI(encoder.Get()), queryCount, native::FromAPI(timestamps.Get()),
+                    native::FromAPI(params.Get()))
                     .IsSuccess());
 }
 
@@ -113,8 +111,6 @@
     uint32_t mQuantizationMask;
 };
 
-constexpr static uint64_t kSentinelValue = ~uint64_t(0u);
-
 class QueryInternalShaderTests : public DawnTest {
   protected:
     void SetUp() override {
@@ -135,23 +131,20 @@
 
     // Original timestamp values in query set for testing
     const std::vector<uint64_t> querySetValues = {
-        kSentinelValue,  // garbage data which is not written at beginning
-        10079569507,     // t0
-        10394415012,     // t1
-        kSentinelValue,  // garbage data which is not written between timestamps
-        11713454943,     // t2
-        38912556941,     // t3 (big value)
-        10080295766,     // t4 (reset)
-        12159966783,     // t5 (after reset)
-        12651224612,     // t6
-        39872473956,     // t7
+        0,            // A zero (written prior for an unavailable query) stays zero.
+        10079569507,  // t0
+        10394415012,  // t1
+        0,            // Same for a zero in the middle.
+        11713454943,  // t2
+        38912556941,  // t3 (big value)
+        10080295766,  // t4 (reset)
+        12159966783,  // t5 (after reset)
+        12651224612,  // t6
+        39872473956,  // t7
     };
 
     const uint32_t kQueryCount = querySetValues.size();
 
-    // Timestamps available state
-    const std::vector<uint32_t> availabilities = {0, 1, 1, 0, 1, 1, 1, 1, 1, 1};
-
     const std::vector<uint64_t> GetExpectedResults(const std::vector<uint64_t>& origin,
                                                    uint32_t start,
                                                    uint32_t firstQuery,
@@ -159,15 +152,10 @@
                                                    float period) {
         std::vector<uint64_t> expected(origin.begin(), origin.end());
         for (size_t i = 0; i < queryCount; i++) {
-            if (availabilities[firstQuery + i] == 0) {
-                // Not a available timestamp, write 0
-                expected[start + i] = 0u;
-            } else {
-                // Maybe the timestamp * period is larger than the maximum of uint64, so cast the
-                // delta value to double (higher precision than float)
-                expected[start + i] =
-                    static_cast<uint64_t>(static_cast<double>(origin[start + i]) * period);
-            }
+            // Maybe the timestamp * period is larger than the maximum of uint64, so cast the
+            // delta value to double (higher precision than float)
+            expected[start + i] =
+                static_cast<uint64_t>(static_cast<double>(origin[start + i]) * period);
         }
         return expected;
     }
@@ -195,23 +183,16 @@
         for (uint32_t i = 0; i < queryCount; i++) {
             timestampValues[start + i] = querySetValues[firstQuery + i];
         }
-        // Write sentinel values and orignal timestamps to timestamps buffer
+        // Write sentinel values and original timestamps to timestamps buffer
         queue.WriteBuffer(timestampsBuffer, 0, timestampValues.data(), size);
 
-        // The buffer indicating which values are available timestamps
-        wgpu::Buffer availabilityBuffer =
-            utils::CreateBufferFromData(device, availabilities.data(),
-                                        kQueryCount * sizeof(uint32_t), wgpu::BufferUsage::Storage);
-
         // The params uniform buffer
-        native::TimestampParams params(firstQuery, queryCount, destinationOffset, quantizationMask,
-                                       period);
+        native::TimestampParams params(queryCount, destinationOffset, quantizationMask, period);
         wgpu::Buffer paramsBuffer = utils::CreateBufferFromData(device, &params, sizeof(params),
                                                                 wgpu::BufferUsage::Uniform);
 
         wgpu::CommandEncoder encoder = device.CreateCommandEncoder();
-        EncodeConvertTimestampsToNanoseconds(encoder, timestampsBuffer, availabilityBuffer,
-                                             paramsBuffer);
+        EncodeConvertTimestampsToNanoseconds(encoder, queryCount, timestampsBuffer, paramsBuffer);
         wgpu::CommandBuffer commands = encoder.Finish();
         queue.Submit(1, &commands);