Track IndirectDrawIndex for all indirect draws

Adds an IndirectDrawIndex that tracks the order of indirect draw
calls in a render pass so that they can be associated with
validated args later. Includes updating the indices of indirect
draws merged in from a RenderBundle after the validation batches
have been merged.

Bug: 495489174
Change-Id: I1639f4f1f827fe8935b6c366b0ab913612fd95e7
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/305976
Commit-Queue: Brandon Jones <bajones@chromium.org>
Reviewed-by: Corentin Wallez <cwallez@chromium.org>

src/dawn/native/IndirectDrawMetadata.cpp

diff --git a/src/dawn/native/IndirectDrawMetadata.cpp b/src/dawn/native/IndirectDrawMetadata.cpp
index 72fa797..4ab3ab2 100644
--- a/src/dawn/native/IndirectDrawMetadata.cpp
+++ b/src/dawn/native/IndirectDrawMetadata.cpp
@@ -103,12 +103,23 @@
     mBatches.insert(it, std::move(newBatch));
 }
 
+void AdjustValidatedDrawIndex(std::vector<IndirectDrawMetadata::IndirectDraw>& draws,
+                              std::vector<IndirectDrawMetadata::IndirectDraw>::iterator begin,
+                              IndirectDrawIndex indirectDrawIndexOffset) {
+    // Ensure that the validatedDrawIndex is properly offset for every newly inserted draw.
+    for (auto it = begin; it != draws.end(); ++it) {
+        it->validatedDrawIndex += indirectDrawIndexOffset;
+    }
+}
+
 void IndirectDrawMetadata::IndexedIndirectBufferValidationInfo::AddBatch(
     uint32_t maxDrawCallsPerIndirectValidationBatch,
     uint64_t maxBatchOffsetRange,
-    const IndirectValidationBatch& newBatch) {
+    const IndirectValidationBatch& newBatch,
+    IndirectDrawIndex indirectDrawIndexOffset) {
     auto it = mBatches.begin();
     while (it != mBatches.end()) {
+        // TODO(crbug.com/495489174): Investigate simplifying this.
         IndirectValidationBatch& batch = *it;
         uint64_t min = std::min(newBatch.minOffset, batch.minOffset);
         uint64_t max = std::max(newBatch.maxOffset, batch.maxOffset);
@@ -117,7 +128,10 @@
             // This batch fits within the limits of an existing batch. Merge it.
             batch.minOffset = min;
             batch.maxOffset = max;
-            batch.draws.insert(batch.draws.end(), newBatch.draws.begin(), newBatch.draws.end());
+
+            auto insertedDraws =
+                batch.draws.insert(batch.draws.end(), newBatch.draws.begin(), newBatch.draws.end());
+            AdjustValidatedDrawIndex(batch.draws, insertedDraws, indirectDrawIndexOffset);
             return;
         }
 
@@ -127,7 +141,11 @@
 
         ++it;
     }
-    mBatches.push_back(newBatch);
+    {
+        mBatches.push_back(newBatch);
+        IndirectValidationBatch& batch = mBatches.back();
+        AdjustValidatedDrawIndex(batch.draws, batch.draws.begin(), indirectDrawIndexOffset);
+    }
 }
 
 const std::vector<IndirectDrawMetadata::IndirectValidationBatch>&
@@ -168,13 +186,17 @@
     for (const auto& [config, validationInfo] :
          bundle->GetIndirectDrawMetadata().mIndexedIndirectBufferValidationInfo) {
         auto it = mIndexedIndirectBufferValidationInfo.lower_bound(config);
-        if (it != mIndexedIndirectBufferValidationInfo.end() && it->first == config) {
-            // We already have batches for the same config. Merge the new ones in.
-            for (const IndirectValidationBatch& batch : validationInfo.GetBatches()) {
-                it->second.AddBatch(mMaxDrawCallsPerBatch, mMaxBatchOffsetRange, batch);
-            }
-        } else {
-            mIndexedIndirectBufferValidationInfo.emplace_hint(it, config, validationInfo);
+        if (it == mIndexedIndirectBufferValidationInfo.end() || it->first != config) {
+            it = mIndexedIndirectBufferValidationInfo.emplace_hint(
+                it, config,
+                IndexedIndirectBufferValidationInfo(validationInfo.GetIndirectBuffer()));
+        }
+
+        // Merge the bundle's batches in.
+        for (const IndirectValidationBatch& batch : validationInfo.GetBatches()) {
+            it->second.AddBatch(mMaxDrawCallsPerBatch, mMaxBatchOffsetRange, batch,
+                                mNextIndirectDrawIndex);
+            mNextIndirectDrawIndex += IndirectDrawIndex(batch.draws.size());
         }
     }
 }
@@ -211,6 +233,7 @@
     }
 
     IndirectDraw draw{};
+    draw.validatedDrawIndex = mNextIndirectDrawIndex++;
     draw.inputBufferOffset = indirectOffset;
     draw.numIndexBufferElements = numIndexBufferElements;
     draw.indexBufferOffsetInElements = indexBufferOffsetInElements;
@@ -232,6 +255,7 @@
     }
 
     IndirectDraw draw{};
+    draw.validatedDrawIndex = mNextIndirectDrawIndex++;
     draw.inputBufferOffset = indirectOffset;
     draw.numIndexBufferElements = 0;
     draw.cmd = cmd;

src/dawn/native/IndirectDrawMetadata.h

diff --git a/src/dawn/native/IndirectDrawMetadata.h b/src/dawn/native/IndirectDrawMetadata.h
index e58813b..f34297b 100644
--- a/src/dawn/native/IndirectDrawMetadata.h
+++ b/src/dawn/native/IndirectDrawMetadata.h
@@ -46,6 +46,11 @@
 class RenderBundleBase;
 struct CombinedLimits;
 
+// The IndirectDrawIndex indicates the order in which an indirect draw is executed in a render
+// pass. It is tracked to enable association of validated buffer/offset data with the original
+// draw after draws are batched for validation.
+using IndirectDrawIndex = TypedInteger<struct IndirectDrawIndexT, uint64_t>;
+
 // In the unlikely scenario that indirect offsets used over a single buffer span more than
 // this length of the buffer, we split the validation work into multiple batches.
 uint64_t ComputeMaxIndirectValidationBatchOffsetRange(const CombinedLimits& limits);
@@ -62,6 +67,7 @@
     };
 
     struct IndirectDraw {
+        IndirectDrawIndex validatedDrawIndex;
         uint64_t inputBufferOffset;
         uint64_t numIndexBufferElements;
         uint64_t indexBufferOffsetInElements;
@@ -115,7 +121,8 @@
         // it's added to mBatch.
         void AddBatch(uint32_t maxDrawCallsPerIndirectValidationBatch,
                       uint64_t maxBatchOffsetRange,
-                      const IndirectValidationBatch& batch);
+                      const IndirectValidationBatch& batch,
+                      IndirectDrawIndex indirectDrawIndexOffset);
 
         const std::vector<IndirectValidationBatch>& GetBatches() const;
 
@@ -192,6 +199,8 @@
 
     std::vector<IndirectMultiDraw> mMultiDraws;
 
+    IndirectDrawIndex mNextIndirectDrawIndex{0};
+
     uint64_t mMaxBatchOffsetRange;
     uint32_t mMaxDrawCallsPerBatch;
 };