Introduce BufferLocation

This is a simple RefCounted holder of a Buffer ref and offset. Encoded
commands can store a ref to this object instead of directly storing an
inline Buffer ref and offset, allowing other commands to dynamically
patch in a different buffer+offset as needed, in a memory-safe way;
as opposed to e.g. retaining an unmanaged pointer to the encoded
command itself and modifying it in-place.

Validation commands will use this to rewrite buffer references in
encoded commands, so that they execute over validated inputs rather
than over their original client-provided inputs.

No net functional changes in this CL, just some groundwork for indirect
draw/dispatch validation.

Bug: dawn:809
Change-Id: I0570521a610fe3ea08190a525b4904749d7b7f24
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/64420
Reviewed-by: Austin Eng <enga@chromium.org>
Reviewed-by: Corentin Wallez <cwallez@chromium.org>
Commit-Queue: Ken Rockot <rockot@google.com>

src/dawn_native/BUILD.gn

@@ -186,6 +186,8 @@ source_set("dawn_native_sources") {
     "BuddyMemoryAllocator.h",
     "Buffer.cpp",
     "Buffer.h",
+    "BufferLocation.cpp",
+    "BufferLocation.h",
     "CachedObject.cpp",
     "CachedObject.h",
     "CallbackTaskManager.cpp",

src/dawn_native/Buffer.cpp

@@ -589,4 +589,5 @@ namespace dawn_native {
     bool BufferBase::IsFullBufferRange(uint64_t offset, uint64_t size) const {
         return offset == 0 && size == GetSize();
     }
+
 }  // namespace dawn_native

src/dawn_native/BufferLocation.cpp

@@ -0,0 +1,54 @@
+// Copyright 2021 The Dawn Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "dawn_native/BufferLocation.h"
+
+namespace dawn_native {
+
+    BufferLocation::BufferLocation() = default;
+
+    BufferLocation::BufferLocation(BufferBase* buffer, uint64_t offset)
+        : mBuffer(buffer), mOffset(offset) {
+    }
+
+    BufferLocation::~BufferLocation() = default;
+
+    // static
+    Ref<BufferLocation> BufferLocation::New() {
+        return AcquireRef(new BufferLocation());
+    }
+
+    // static
+    Ref<BufferLocation> BufferLocation::New(BufferBase* buffer, uint64_t offset) {
+        return AcquireRef(new BufferLocation(buffer, offset));
+    }
+
+    bool BufferLocation::IsNull() const {
+        return mBuffer.Get() == nullptr;
+    }
+
+    BufferBase* BufferLocation::GetBuffer() const {
+        return mBuffer.Get();
+    }
+
+    uint64_t BufferLocation::GetOffset() const {
+        return mOffset;
+    }
+
+    void BufferLocation::Set(BufferBase* buffer, uint64_t offset) {
+        mBuffer = buffer;
+        mOffset = offset;
+    }
+
+}  // namespace dawn_native

src/dawn_native/BufferLocation.h

@@ -0,0 +1,49 @@
+// Copyright 2021 The Dawn Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef DAWNNATIVE_BUFFERLOCATION_H_
+#define DAWNNATIVE_BUFFERLOCATION_H_
+
+#include "common/RefCounted.h"
+#include "dawn_native/Buffer.h"
+
+#include <cstdint>
+
+namespace dawn_native {
+
+    // A ref-counted wrapper around a Buffer ref and an offset into the buffer.
+    class BufferLocation : public RefCounted {
+      public:
+        BufferLocation();
+        BufferLocation(BufferBase* buffer, uint64_t offset = 0);
+        ~BufferLocation();
+
+        static Ref<BufferLocation> New();
+        static Ref<BufferLocation> New(BufferBase* buffer, uint64_t offset = 0);
+
+        bool IsNull() const;
+
+        BufferBase* GetBuffer() const;
+        uint64_t GetOffset() const;
+
+        void Set(BufferBase* buffer, uint64_t offset);
+
+      private:
+        Ref<BufferBase> mBuffer;
+        uint64_t mOffset = 0;
+    };
+
+}  // namespace dawn_native
+
+#endif  // DAWNNATIVE_BUFFERLOCATION_H_

src/dawn_native/Commands.h

@@ -19,6 +19,7 @@
 
 #include "dawn_native/AttachmentState.h"
 #include "dawn_native/BindingInfo.h"
+#include "dawn_native/BufferLocation.h"
 #include "dawn_native/Texture.h"
 
 #include "dawn_native/dawn_platform.h"
@@ -177,8 +178,7 @@ namespace dawn_native {
     };
 
     struct DrawIndexedIndirectCmd {
-        Ref<BufferBase> indirectBuffer;
+        Ref<BufferLocation> indirectBufferLocation;
-        uint64_t indirectOffset;
     };
 
     struct EndComputePassCmd {};

src/dawn_native/RenderEncoderBase.cpp

@@ -17,6 +17,7 @@
 #include "common/Constants.h"
 #include "common/Log.h"
 #include "dawn_native/Buffer.h"
+#include "dawn_native/BufferLocation.h"
 #include "dawn_native/CommandEncoder.h"
 #include "dawn_native/CommandValidation.h"
 #include "dawn_native/Commands.h"
@@ -178,8 +179,7 @@ namespace dawn_native {
 
             DrawIndexedIndirectCmd* cmd =
                 allocator->Allocate<DrawIndexedIndirectCmd>(Command::DrawIndexedIndirect);
-            cmd->indirectBuffer = indirectBuffer;
+            cmd->indirectBufferLocation = BufferLocation::New(indirectBuffer, indirectOffset);
-            cmd->indirectOffset = indirectOffset;
 
             mUsageTracker.BufferUsedAs(indirectBuffer, wgpu::BufferUsage::Indirect);
 

src/dawn_native/d3d12/CommandBufferD3D12.cpp

@@ -1383,14 +1383,16 @@ namespace dawn_native { namespace d3d12 {
 
                 case Command::DrawIndexedIndirect: {
                     DrawIndexedIndirectCmd* draw = iter->NextCommand<DrawIndexedIndirectCmd>();
+                    ASSERT(!draw->indirectBufferLocation->IsNull());
 
                     DAWN_TRY(bindingTracker->Apply(commandContext));
                     vertexBufferTracker.Apply(commandList, lastPipeline);
-                    Buffer* buffer = ToBackend(draw->indirectBuffer.Get());
+                    Buffer* buffer = ToBackend(draw->indirectBufferLocation->GetBuffer());
                     ComPtr<ID3D12CommandSignature> signature =
                         ToBackend(GetDevice())->GetDrawIndexedIndirectSignature();
                     commandList->ExecuteIndirect(signature.Get(), 1, buffer->GetD3D12Resource(),
-                                                 draw->indirectOffset, nullptr, 0);
+                                                 draw->indirectBufferLocation->GetOffset(), nullptr,
+                                                 0);
                     break;
                 }
 

src/dawn_native/metal/CommandBufferMTL.mm

@@ -1332,20 +1332,21 @@ namespace dawn_native { namespace metal {
                 }
 
                 case Command::DrawIndexedIndirect: {
-                    DrawIndirectCmd* draw = iter->NextCommand<DrawIndirectCmd>();
+                    DrawIndexedIndirectCmd* draw = iter->NextCommand<DrawIndexedIndirectCmd>();
 
                     vertexBuffers.Apply(encoder, lastPipeline, enableVertexPulling);
                     bindGroups.Apply(encoder);
                     storageBufferLengths.Apply(encoder, lastPipeline, enableVertexPulling);
 
-                    Buffer* buffer = ToBackend(draw->indirectBuffer.Get());
+                    ASSERT(!draw->indirectBufferLocation->IsNull());
+                    Buffer* buffer = ToBackend(draw->indirectBufferLocation->GetBuffer());
                     id<MTLBuffer> indirectBuffer = buffer->GetMTLBuffer();
                     [encoder drawIndexedPrimitives:lastPipeline->GetMTLPrimitiveTopology()
                                          indexType:indexBufferType
                                        indexBuffer:indexBuffer
                                  indexBufferOffset:indexBufferBaseOffset
                                     indirectBuffer:indirectBuffer
-                              indirectBufferOffset:draw->indirectOffset];
+                              indirectBufferOffset:draw->indirectBufferLocation->GetOffset()];
                     break;
                 }
 

src/dawn_native/opengl/CommandBufferGL.cpp

@@ -1182,16 +1182,17 @@ namespace dawn_native { namespace opengl {
 
                 case Command::DrawIndexedIndirect: {
                     DrawIndexedIndirectCmd* draw = iter->NextCommand<DrawIndexedIndirectCmd>();
+                    ASSERT(!draw->indirectBufferLocation->IsNull());
+
                     vertexStateBufferBindingTracker.Apply(gl);
                     bindGroupTracker.Apply(gl);
 
-                    uint64_t indirectBufferOffset = draw->indirectOffset;
+                    Buffer* indirectBuffer = ToBackend(draw->indirectBufferLocation->GetBuffer());
-                    Buffer* indirectBuffer = ToBackend(draw->indirectBuffer.Get());
-
                     gl.BindBuffer(GL_DRAW_INDIRECT_BUFFER, indirectBuffer->GetHandle());
                     gl.DrawElementsIndirect(
                         lastPipeline->GetGLPrimitiveTopology(), indexBufferFormat,
-                        reinterpret_cast<void*>(static_cast<intptr_t>(indirectBufferOffset)));
+                        reinterpret_cast<void*>(
+                            static_cast<intptr_t>(draw->indirectBufferLocation->GetOffset())));
                     break;
                 }
 

src/dawn_native/vulkan/CommandBufferVk.cpp

@@ -1079,13 +1079,15 @@ namespace dawn_native { namespace vulkan {
                 }
 
                 case Command::DrawIndexedIndirect: {
-                    DrawIndirectCmd* draw = iter->NextCommand<DrawIndirectCmd>();
+                    DrawIndexedIndirectCmd* draw = iter->NextCommand<DrawIndexedIndirectCmd>();
-                    VkBuffer indirectBuffer = ToBackend(draw->indirectBuffer)->GetHandle();
+                    ASSERT(!draw->indirectBufferLocation->IsNull());
+                    VkBuffer indirectBuffer =
+                        ToBackend(draw->indirectBufferLocation->GetBuffer())->GetHandle();
 
                     descriptorSets.Apply(device, recordingContext, VK_PIPELINE_BIND_POINT_GRAPHICS);
                     device->fn.CmdDrawIndexedIndirect(
-                        commands, indirectBuffer, static_cast<VkDeviceSize>(draw->indirectOffset),
+                        commands, indirectBuffer,
-                        1, 0);
+                        static_cast<VkDeviceSize>(draw->indirectBufferLocation->GetOffset()), 1, 0);
                     break;
                 }