[Dawn] Implement DoCompute support in dawn backend and fix sparse descriptor set support (#585)

Also adds Buffer::RecalculateMaxSizeInBytes

8 files changed, 529 insertions, 161 deletions

diff --git a/src/buffer.cc b/src/buffer.cc
index ce33b1a..617ff7a 100644
--- a/src/buffer.cc
+++ b/src/buffer.cc
@@ -115,6 +115,26 @@ Result Buffer::SetData(const std::vector<Value>& data) {
   return SetDataWithOffset(data, 0);
 }
 
+Result Buffer::RecalculateMaxSizeInBytes(const std::vector<Value>& data,
+                                         uint32_t offset) {
+  // Multiply by the input needed because the value count will use the needed
+  // input as the multiplier
+  uint32_t value_count =
+      ((offset / format_->SizeInBytes()) * format_->InputNeededPerElement()) +
+      static_cast<uint32_t>(data.size());
+  uint32_t element_count = value_count;
+  if (format_->GetPackSize() == 0) {
+    // This divides by the needed input values, not the values per element.
+    // The assumption being the values coming in are read from the input,
+    // where components are specified. The needed values maybe less then the
+    // values per element.
+    element_count = value_count / format_->InputNeededPerElement();
+  }
+  if (GetMaxSizeInBytes() < element_count * format_->SizeInBytes())
+    SetMaxSizeInBytes(element_count * format_->SizeInBytes());
+  return {};
+}
+
 Result Buffer::SetDataWithOffset(const std::vector<Value>& data,
                                  uint32_t offset) {
   // Multiply by the input needed because the value count will use the needed
@@ -223,6 +243,23 @@ void Buffer::ResizeTo(uint32_t element_count) {
   bytes_.resize(element_count * format_->SizeInBytes());
 }
 
+void Buffer::SetSizeInBytes(uint32_t size_in_bytes) {
+  assert(size_in_bytes % format_->SizeInBytes() == 0);
+  element_count_ = size_in_bytes / format_->SizeInBytes();
+  bytes_.resize(size_in_bytes);
+}
+
+void Buffer::SetMaxSizeInBytes(uint32_t max_size_in_bytes) {
+  max_size_in_bytes_ = max_size_in_bytes;
+}
+
+uint32_t Buffer::GetMaxSizeInBytes() const {
+  if (max_size_in_bytes_ != 0)
+    return max_size_in_bytes_;
+  else
+    return GetSizeInBytes();
+}
+
 Result Buffer::SetDataFromBuffer(const Buffer* src, uint32_t offset) {
   if (bytes_.size() < offset + src->bytes_.size())
     bytes_.resize(offset + src->bytes_.size());
diff --git a/src/buffer.h b/src/buffer.h
index 5754ddb..906454d 100644
--- a/src/buffer.h
+++ b/src/buffer.h
@@ -150,10 +150,29 @@ class Buffer {
   /// initial count. This requires the format to have been set.
   void ResizeTo(uint32_t element_count);
 
+  /// Resizes the buffer to hold |size_in_bytes| over format_->SizeInBytes()
+  /// number of elements while resizing the buffer to |size_in_bytes| bytes.
+  /// This requires the format to have been set. This is separate from
+  /// ResizeTo() since the given argument here is |size_in_bytes| bytes vs
+  /// |element_count| elements
+  void SetSizeInBytes(uint32_t size_in_bytes);
+
+  /// Sets the max_size_in_bytes_ to |max_size_in_bytes| bytes
+  void SetMaxSizeInBytes(uint32_t max_size_in_bytes);
+  /// Returns max_size_in_bytes_ if it is not zero. Otherwise it means this
+  /// buffer is an amber buffer which has a fix size and returns
+  /// GetSizeInBytes()
+  uint32_t GetMaxSizeInBytes() const;
+
   /// Write |data| into the buffer |offset| bytes from the start. Write
   /// |size_in_bytes| of data.
   Result SetDataWithOffset(const std::vector<Value>& data, uint32_t offset);
 
+  /// At each ubo, ssbo size and ssbo subdata size calls, recalculates
+  /// max_size_in_bytes_ and updates it if underlying buffer got bigger
+  Result RecalculateMaxSizeInBytes(const std::vector<Value>& data,
+                                   uint32_t offset);
+
   /// Writes |src| data into buffer at |offset|.
   Result SetDataFromBuffer(const Buffer* src, uint32_t offset);
 
@@ -181,6 +200,9 @@ class Buffer {
 
   BufferType buffer_type_ = BufferType::kUnknown;
   std::string name_;
+  /// max_size_in_bytes_ is the total size in bytes needed to hold the buffer
+  /// over all ubo, ssbo size and ssbo subdata size calls.
+  uint32_t max_size_in_bytes_ = 0;
   uint32_t element_count_ = 0;
   uint32_t width_ = 0;
   uint32_t height_ = 0;
diff --git a/src/dawn/engine_dawn.cc b/src/dawn/engine_dawn.cc
index 9a3eeef..46699c1 100644
--- a/src/dawn/engine_dawn.cc
+++ b/src/dawn/engine_dawn.cc
@@ -292,6 +292,110 @@ MapResult MapBuffer(const ::dawn::Device& device, const ::dawn::Buffer& buf) {
   return textureCopyView;
 }
 
+Result EngineDawn::MapDeviceTextureToHostBuffer(
+    const RenderPipelineInfo& render_pipeline,
+    const ::dawn::Device& device) {
+  const auto width = render_pipeline.pipeline->GetFramebufferWidth();
+  const auto height = render_pipeline.pipeline->GetFramebufferHeight();
+
+  const auto pixelSize = render_pipeline.pipeline->GetColorAttachments()[0]
+                             .buffer->GetTexelStride();
+  const auto dawn_row_pitch = Align(width * pixelSize, kMinimumImageRowPitch);
+  const auto size = height * dawn_row_pitch;
+  // Create a temporary buffer to hold the color attachment content and can
+  // be mapped
+  ::dawn::BufferDescriptor descriptor;
+  descriptor.size = size;
+  descriptor.usage =
+      ::dawn::BufferUsageBit::CopyDst | ::dawn::BufferUsageBit::MapRead;
+  ::dawn::Buffer copy_buffer = device.CreateBuffer(&descriptor);
+  ::dawn::BufferCopyView copy_buffer_view =
+      CreateBufferCopyView(copy_buffer, 0, dawn_row_pitch, 0);
+  ::dawn::Origin3D origin3D;
+  origin3D.x = 0;
+  origin3D.y = 0;
+  origin3D.z = 0;
+
+  for (uint32_t i = 0;
+       i < render_pipeline.pipeline->GetColorAttachments().size(); i++) {
+    ::dawn::TextureCopyView device_texture_view =
+        CreateTextureCopyView(textures_[i], 0, 0, origin3D);
+    ::dawn::Extent3D copySize = {width, height, 1};
+    auto encoder = device.CreateCommandEncoder();
+    encoder.CopyTextureToBuffer(&device_texture_view, &copy_buffer_view,
+                                &copySize);
+    auto commands = encoder.Finish();
+    auto queue = device.CreateQueue();
+    queue.Submit(1, &commands);
+
+    MapResult mapped_device_texture = MapBuffer(device, copy_buffer);
+    if (!mapped_device_texture.result.IsSuccess())
+      return mapped_device_texture.result;
+
+    auto& host_texture = render_pipeline.pipeline->GetColorAttachments()[i];
+    auto* values = host_texture.buffer->ValuePtr();
+    auto row_stride = pixelSize * width;
+    assert(row_stride * height == host_texture.buffer->GetSizeInBytes());
+    // Each Dawn row has enough data to fill the target row.
+    assert(dawn_row_pitch >= row_stride);
+    values->resize(host_texture.buffer->GetSizeInBytes());
+    // Copy the framebuffer contents back into the host-side
+    // framebuffer-buffer. In the Dawn buffer, the row stride is a multiple of
+    // kMinimumImageRowPitch bytes, so it might have padding therefore memcpy
+    // is done row by row.
+    for (uint h = 0; h < height; h++) {
+      std::memcpy(values->data() + h * row_stride,
+                  static_cast<const uint8_t*>(mapped_device_texture.data) +
+                      h * dawn_row_pitch,
+                  row_stride);
+    }
+    // Always unmap the buffer at the end of the engine's command.
+    copy_buffer.Unmap();
+  }
+  return {};
+}
+
+Result EngineDawn::MapDeviceBufferToHostBuffer(
+    const ComputePipelineInfo& compute_pipeline,
+    const ::dawn::Device& device) {
+  for (uint32_t i = 0; i < compute_pipeline.pipeline->GetBuffers().size();
+       i++) {
+    auto& device_buffer = compute_pipeline.buffers[i];
+    auto& host_buffer = compute_pipeline.pipeline->GetBuffers()[i];
+
+    // Create a copy of device buffer to use it in a map read operation.
+    // It's not possible to simply set this bit on the existing buffers since:
+    // Device error: Only CopyDst is allowed with MapRead
+    ::dawn::BufferDescriptor descriptor;
+    descriptor.size = host_buffer.buffer->GetSizeInBytes();
+    descriptor.usage =
+        ::dawn::BufferUsageBit::CopyDst | ::dawn::BufferUsageBit::MapRead;
+    const auto copy_device_buffer = device.CreateBuffer(&descriptor);
+    const uint64_t source_offset = 0;
+    const uint64_t destination_offset = 0;
+    const uint64_t copy_size =
+        static_cast<uint64_t>(host_buffer.buffer->GetSizeInBytes());
+    auto encoder = device.CreateCommandEncoder();
+    encoder.CopyBufferToBuffer(device_buffer, source_offset, copy_device_buffer,
+                               destination_offset, copy_size);
+    auto commands = encoder.Finish();
+    auto queue = device.CreateQueue();
+    queue.Submit(1, &commands);
+
+    MapResult mapped_device_buffer = MapBuffer(device, copy_device_buffer);
+    auto* values = host_buffer.buffer->ValuePtr();
+    values->resize(host_buffer.buffer->GetSizeInBytes());
+    std::memcpy(values->data(),
+                static_cast<const uint8_t*>(mapped_device_buffer.data),
+                copy_size);
+
+    copy_device_buffer.Unmap();
+    if (!mapped_device_buffer.result.IsSuccess())
+      return mapped_device_buffer.result;
+  }
+  return {};
+}
+
 // Creates a dawn buffer of |size| bytes with TransferDst and the given usage
 // copied from Dawn utils source code
 ::dawn::Buffer CreateBufferFromData(const ::dawn::Device& device,
@@ -568,69 +672,6 @@ Result EngineDawn::Initialize(EngineConfig* config,
   return {};
 }
 
-Result EngineDawn::MapDeviceTextureToHostBuffer(
-    const RenderPipelineInfo& render_pipeline,
-    const ::dawn::Device& device) {
-  const auto width = render_pipeline.pipeline->GetFramebufferWidth();
-  const auto height = render_pipeline.pipeline->GetFramebufferHeight();
-
-  const auto pixelSize = render_pipeline.pipeline->GetColorAttachments()[0]
-                             .buffer->GetTexelStride();
-  const auto dawn_row_pitch = Align(width * pixelSize, kMinimumImageRowPitch);
-  const auto size = height * dawn_row_pitch;
-  // Create a temporary buffer to hold the color attachment content and can
-  // be mapped
-  ::dawn::BufferDescriptor descriptor;
-  descriptor.size = size;
-  descriptor.usage =
-      ::dawn::BufferUsageBit::CopyDst | ::dawn::BufferUsageBit::MapRead;
-  ::dawn::Buffer copy_buffer = device.CreateBuffer(&descriptor);
-  ::dawn::BufferCopyView copy_buffer_view =
-      CreateBufferCopyView(copy_buffer, 0, dawn_row_pitch, 0);
-  ::dawn::Origin3D origin3D;
-  origin3D.x = 0;
-  origin3D.y = 0;
-  origin3D.z = 0;
-
-  for (uint32_t i = 0;
-       i < render_pipeline.pipeline->GetColorAttachments().size(); i++) {
-    ::dawn::TextureCopyView device_texture_view =
-        CreateTextureCopyView(textures_[i], 0, 0, origin3D);
-    ::dawn::Extent3D copySize = {width, height, 1};
-    auto encoder = device.CreateCommandEncoder();
-    encoder.CopyTextureToBuffer(&device_texture_view, &copy_buffer_view,
-                                &copySize);
-    auto commands = encoder.Finish();
-    auto queue = device.CreateQueue();
-    queue.Submit(1, &commands);
-
-    MapResult mapped_device_texture = MapBuffer(device, copy_buffer);
-    if (!mapped_device_texture.result.IsSuccess())
-      return mapped_device_texture.result;
-
-    auto& host_texture = render_pipeline.pipeline->GetColorAttachments()[i];
-    auto* values = host_texture.buffer->ValuePtr();
-    auto row_stride = pixelSize * width;
-    assert(row_stride * height == host_texture.buffer->GetSizeInBytes());
-    // Each Dawn row has enough data to fill the target row.
-    assert(dawn_row_pitch >= row_stride);
-    values->resize(host_texture.buffer->GetSizeInBytes());
-    // Copy the framebuffer contents back into the host-side
-    // framebuffer-buffer. In the Dawn buffer, the row stride is a multiple of
-    // kMinimumImageRowPitch bytes, so it might have padding therefore memcpy
-    // is done row by row.
-    for (uint h = 0; h < height; h++) {
-      std::memcpy(values->data() + h * row_stride,
-                  static_cast<const uint8_t*>(mapped_device_texture.data) +
-                      h * dawn_row_pitch,
-                  row_stride);
-    }
-    // Always unmap the buffer at the end of the engine's command.
-    copy_buffer.Unmap();
-  }
-  return {};
-}
-
 Result EngineDawn::CreatePipeline(::amber::Pipeline* pipeline) {
   if (!device_) {
     return Result("Dawn::CreatePipeline: device is not created");
@@ -661,8 +702,13 @@ Result EngineDawn::CreatePipeline(::amber::Pipeline* pipeline) {
       auto& module = module_for_type[kShaderTypeCompute];
       if (!module)
         return Result("Dawn::CreatePipeline: no compute shader provided");
+
       pipeline_map_[pipeline].compute_pipeline.reset(
           new ComputePipelineInfo(pipeline, module));
+      Result result =
+          AttachBuffers(pipeline_map_[pipeline].compute_pipeline.get());
+      if (!result.IsSuccess())
+        return result;
       break;
     }
 
@@ -1162,10 +1208,44 @@ Result EngineDawn::DoDrawArrays(const DrawArraysCommand* command) {
   result = MapDeviceTextureToHostBuffer(*render_pipeline, *device_);
 
   return result;
-}  // namespace dawn
+}
+
+Result EngineDawn::DoCompute(const ComputeCommand* command) {
+  Result result;
+
+  ComputePipelineInfo* compute_pipeline = GetComputePipeline(command);
+  if (!compute_pipeline)
+    return Result("DoComput: invoked on invalid or missing compute pipeline");
 
-Result EngineDawn::DoCompute(const ComputeCommand*) {
-  return Result("Dawn:DoCompute not implemented");
+  ::dawn::ComputePipelineDescriptor computePipelineDescriptor;
+  computePipelineDescriptor.layout = MakeBasicPipelineLayout(
+      device_->Get(), compute_pipeline->bind_group_layouts);
+
+  ::dawn::PipelineStageDescriptor pipelineStageDescriptor;
+  pipelineStageDescriptor.module = compute_pipeline->compute_shader;
+  pipelineStageDescriptor.entryPoint = "main";
+  computePipelineDescriptor.computeStage = &pipelineStageDescriptor;
+  ::dawn::ComputePipeline pipeline =
+      device_->CreateComputePipeline(&computePipelineDescriptor);
+  ::dawn::CommandEncoder encoder = device_->CreateCommandEncoder();
+  ::dawn::ComputePassEncoder pass = encoder.BeginComputePass();
+  pass.SetPipeline(pipeline);
+  for (uint32_t i = 0; i < compute_pipeline->bind_groups.size(); i++) {
+    if (compute_pipeline->bind_groups[i]) {
+      pass.SetBindGroup(i, compute_pipeline->bind_groups[i], 0, nullptr);
+    }
+  }
+  pass.Dispatch(command->GetX(), command->GetY(), command->GetZ());
+  pass.EndPass();
+  // Finish recording the command buffer.  It only has one command.
+  auto command_buffer = encoder.Finish();
+  // Submit the command.
+  auto queue = device_->CreateQueue();
+  queue.Submit(1, &command_buffer);
+  // Copy result back
+  result = MapDeviceBufferToHostBuffer(*compute_pipeline, *device_);
+
+  return result;
 }
 
 Result EngineDawn::DoEntryPoint(const EntryPointCommand*) {
@@ -1180,30 +1260,43 @@ Result EngineDawn::DoPatchParameterVertices(
 Result EngineDawn::DoBuffer(const BufferCommand* command) {
   Result result;
 
-  // TODO(SarahM0): Make this work for compute pipeline
+  ::dawn::Buffer* dawn_buffer;
+
   RenderPipelineInfo* render_pipeline = GetRenderPipeline(command);
-  if (!render_pipeline)
-    return Result("DoBuffer: invoked on invalid or missing render pipeline");
+  if (render_pipeline) {
+    if (render_pipeline->buffer_map.find(
+            {command->GetDescriptorSet(), command->GetBinding()}) !=
+        render_pipeline->buffer_map.end()) {
+      auto dawn_buffer_index = render_pipeline->buffer_map[{
+          command->GetDescriptorSet(), command->GetBinding()}];
+      dawn_buffer = &render_pipeline->buffers[dawn_buffer_index];
+    }
+  }
+
+  ComputePipelineInfo* compute_pipeline = GetComputePipeline(command);
+  if (compute_pipeline) {
+    if (compute_pipeline->buffer_map.find(
+            {command->GetDescriptorSet(), command->GetBinding()}) !=
+        compute_pipeline->buffer_map.end()) {
+      auto dawn_buffer_index = compute_pipeline->buffer_map[{
+          command->GetDescriptorSet(), command->GetBinding()}];
+      dawn_buffer = &compute_pipeline->buffers[dawn_buffer_index];
+    }
+  }
+
+  if (!render_pipeline && !compute_pipeline)
+    return Result("DoBuffer: invoked on invalid or missing pipeline");
   if (!command->IsSSBO() && !command->IsUniform())
     return Result("DoBuffer: only supports SSBO and uniform buffer type");
 
-  if (render_pipeline->buffer_map.find(
-          {command->GetDescriptorSet(), command->GetBinding()}) !=
-      render_pipeline->buffer_map.end()) {
-    auto dawn_buffer_index =
-        render_pipeline
-            ->buffer_map[{command->GetDescriptorSet(), command->GetBinding()}];
-    ::dawn::Buffer& dawn_buffer = render_pipeline->buffers[dawn_buffer_index];
-
-    Buffer* amber_buffer = command->GetBuffer();
-    if (amber_buffer) {
-      amber_buffer->SetDataWithOffset(command->GetValues(),
-                                      command->GetOffset());
-
-      dawn_buffer.SetSubData(0, amber_buffer->GetSizeInBytes(),
-                             amber_buffer->ValuePtr()->data());
-    }
+  Buffer* amber_buffer = command->GetBuffer();
+  if (amber_buffer) {
+    amber_buffer->SetDataWithOffset(command->GetValues(), command->GetOffset());
+
+    dawn_buffer->SetSubData(0, amber_buffer->GetMaxSizeInBytes(),
+                            amber_buffer->ValuePtr()->data());
   }
+
   return {};
 }
 
@@ -1307,6 +1400,16 @@ Result EngineDawn::AttachBuffersAndTextures(
   uint32_t max_descriptor_set = 0;
 
   // Attach storage/uniform buffers
+  ::dawn::BindGroupLayoutBinding empty_layout_info = {};
+
+  if (!render_pipeline->pipeline->GetBuffers().empty()) {
+    std::vector<uint32_t> max_binding_seen(kMaxDawnBindGroup, -1);
+    for (auto& buf_info : render_pipeline->pipeline->GetBuffers()) {
+      while (layouts_info[buf_info.descriptor_set].size() <= buf_info.binding)
+        layouts_info[buf_info.descriptor_set].push_back(empty_layout_info);
+    }
+  }
+
   for (const auto& buf_info : render_pipeline->pipeline->GetBuffers()) {
     ::dawn::BufferUsageBit bufferUsage;
     ::dawn::BindingType bindingType;
@@ -1330,14 +1433,13 @@ Result EngineDawn::AttachBuffersAndTextures(
     }
 
     if (buf_info.descriptor_set > kMaxDawnBindGroup - 1) {
-      return Result(
-          "AttachBuffersAndTextures: Dawn has maximum of 4 bindGroups "
-          "(descriptor sets)");
+      return Result("AttachBuffers: Dawn has a maximum of " +
+                    std::to_string(kMaxDawnBindGroup) + " (descriptor sets)");
     }
 
     render_pipeline->buffers.emplace_back(
         CreateBufferFromData(*device_, buf_info.buffer->ValuePtr()->data(),
-                             buf_info.buffer->GetSizeInBytes(),
+                             buf_info.buffer->GetMaxSizeInBytes(),
                              bufferUsage | ::dawn::BufferUsageBit::CopySrc |
                                  ::dawn::BufferUsageBit::CopyDst));
 
@@ -1351,22 +1453,14 @@ Result EngineDawn::AttachBuffersAndTextures(
     layout_info.binding = buf_info.binding;
     layout_info.visibility = kAllStages;
     layout_info.type = bindingType;
-    layouts_info[buf_info.descriptor_set].push_back(layout_info);
+    layouts_info[buf_info.descriptor_set][buf_info.binding] = layout_info;
 
     BindingInitializationHelper tempBinding = BindingInitializationHelper(
         buf_info.binding, render_pipeline->buffers.back(), 0,
-        buf_info.buffer->GetSizeInBytes());
+        buf_info.buffer->GetMaxSizeInBytes());
     bindingInitalizerHelper[buf_info.descriptor_set].push_back(tempBinding);
   }
 
-  // TODO(sarahM0): fix issue: Add support for doBuffer with sparse descriptor
-  // sets #573
-  if (render_pipeline->used_descriptor_set.size() != 0 &&
-      render_pipeline->used_descriptor_set.size() != max_descriptor_set + 1) {
-    return Result(
-        "AttachBuffersAndTextures: Sparse descriptor_set is not supported");
-  }
-
   for (uint32_t i = 0; i < kMaxDawnBindGroup; i++) {
     if (layouts_info[i].size() > 0 && bindingInitalizerHelper[i].size() > 0) {
       ::dawn::BindGroupLayout bindGroupLayout =
@@ -1377,6 +1471,115 @@ Result EngineDawn::AttachBuffersAndTextures(
           MakeBindGroup(*device_, render_pipeline->bind_group_layouts[i],
                         bindingInitalizerHelper[i]);
       render_pipeline->bind_groups.push_back(bindGroup);
+    } else if (i < max_descriptor_set) {
+      ::dawn::BindGroupLayout bindGroupLayout =
+          MakeBindGroupLayout(*device_, {});
+      render_pipeline->bind_group_layouts.push_back(bindGroupLayout);
+
+      ::dawn::BindGroup bindGroup =
+          MakeBindGroup(*device_, render_pipeline->bind_group_layouts[i],
+                        bindingInitalizerHelper[i]);
+      render_pipeline->bind_groups.push_back(bindGroup);
+    }
+  }
+  return {};
+}
+
+Result EngineDawn::AttachBuffers(ComputePipelineInfo* compute_pipeline) {
+  Result result;
+
+  // Do not attach pushConstants
+  if (compute_pipeline->pipeline->GetPushConstantBuffer().buffer != nullptr) {
+    return Result("AttachBuffers: Dawn does not support push constants!");
+  }
+
+  std::vector<std::vector<BindingInitializationHelper>> bindingInitalizerHelper(
+      kMaxDawnBindGroup);
+  std::vector<std::vector<::dawn::BindGroupLayoutBinding>> layouts_info(
+      kMaxDawnBindGroup);
+  uint32_t max_descriptor_set = 0;
+
+  // Attach storage/uniform buffers
+  ::dawn::BindGroupLayoutBinding empty_layout_info = {};
+
+  if (!compute_pipeline->pipeline->GetBuffers().empty()) {
+    std::vector<uint32_t> max_binding_seen(kMaxDawnBindGroup, -1);
+    for (auto& buf_info : compute_pipeline->pipeline->GetBuffers()) {
+      while (layouts_info[buf_info.descriptor_set].size() <= buf_info.binding)
+        layouts_info[buf_info.descriptor_set].push_back(empty_layout_info);
+    }
+  }
+
+  for (const auto& buf_info : compute_pipeline->pipeline->GetBuffers()) {
+    ::dawn::BufferUsageBit bufferUsage;
+    ::dawn::BindingType bindingType;
+    switch (buf_info.buffer->GetBufferType()) {
+      case BufferType::kStorage: {
+        bufferUsage = ::dawn::BufferUsageBit::Storage;
+        bindingType = ::dawn::BindingType::StorageBuffer;
+        break;
+      }
+      case BufferType::kUniform: {
+        bufferUsage = ::dawn::BufferUsageBit::Uniform;
+        bindingType = ::dawn::BindingType::UniformBuffer;
+        break;
+      }
+      default: {
+        return Result("AttachBuffers: unknown buffer type: " +
+                      std::to_string(static_cast<uint32_t>(
+                          buf_info.buffer->GetBufferType())));
+        break;
+      }
+    }
+
+    if (buf_info.descriptor_set > kMaxDawnBindGroup - 1) {
+      return Result("AttachBuffers: Dawn has a maximum of " +
+                    std::to_string(kMaxDawnBindGroup) + " (descriptor sets)");
+    }
+
+    compute_pipeline->buffers.emplace_back(
+        CreateBufferFromData(*device_, buf_info.buffer->ValuePtr()->data(),
+                             buf_info.buffer->GetMaxSizeInBytes(),
+                             bufferUsage | ::dawn::BufferUsageBit::CopySrc |
+                                 ::dawn::BufferUsageBit::CopyDst));
+
+    compute_pipeline->buffer_map[{buf_info.descriptor_set, buf_info.binding}] =
+        compute_pipeline->buffers.size() - 1;
+
+    compute_pipeline->used_descriptor_set.insert(buf_info.descriptor_set);
+    max_descriptor_set = std::max(max_descriptor_set, buf_info.descriptor_set);
+
+    ::dawn::BindGroupLayoutBinding layout_info;
+    layout_info.binding = buf_info.binding;
+    layout_info.visibility = ::dawn::ShaderStageBit::Compute;
+    layout_info.type = bindingType;
+    layouts_info[buf_info.descriptor_set][buf_info.binding] = layout_info;
+
+    BindingInitializationHelper tempBinding = BindingInitializationHelper(
+        buf_info.binding, compute_pipeline->buffers.back(), 0,
+        buf_info.buffer->GetMaxSizeInBytes());
+    bindingInitalizerHelper[buf_info.descriptor_set].push_back(tempBinding);
+  }
+
+  for (uint32_t i = 0; i < kMaxDawnBindGroup; i++) {
+    if (layouts_info[i].size() > 0 && bindingInitalizerHelper[i].size() > 0) {
+      ::dawn::BindGroupLayout bindGroupLayout =
+          MakeBindGroupLayout(*device_, layouts_info[i]);
+      compute_pipeline->bind_group_layouts.push_back(bindGroupLayout);
+
+      ::dawn::BindGroup bindGroup =
+          MakeBindGroup(*device_, compute_pipeline->bind_group_layouts[i],
+                        bindingInitalizerHelper[i]);
+      compute_pipeline->bind_groups.push_back(bindGroup);
+    } else if (i < max_descriptor_set) {
+      ::dawn::BindGroupLayout bindGroupLayout =
+          MakeBindGroupLayout(*device_, {});
+      compute_pipeline->bind_group_layouts.push_back(bindGroupLayout);
+
+      ::dawn::BindGroup bindGroup =
+          MakeBindGroup(*device_, compute_pipeline->bind_group_layouts[i],
+                        bindingInitalizerHelper[i]);
+      compute_pipeline->bind_groups.push_back(bindGroup);
     }
   }
 
diff --git a/src/dawn/engine_dawn.h b/src/dawn/engine_dawn.h
index 5fa699f..cd8ad21 100644
--- a/src/dawn/engine_dawn.h
+++ b/src/dawn/engine_dawn.h
@@ -68,15 +68,27 @@ class EngineDawn : public Engine {
       const ::amber::PipelineCommand* command) {
     return pipeline_map_[command->GetPipeline()].render_pipeline.get();
   }
-
+  // Returns the Dawn-specific compute pipeline for the given command,
+  // if it exists.  Returns nullptr otherwise.
+  ComputePipelineInfo* GetComputePipeline(
+      const ::amber::PipelineCommand* command) {
+    return pipeline_map_[command->GetPipeline()].compute_pipeline.get();
+  }
   // Creates and attaches index, vertex, storage, uniform and depth-stencil
   // buffers. Sets up bindings. Also creates textures and texture views if not
-  // created yet.
+  // created yet. Used in the Graphics pipeline creation.
   Result AttachBuffersAndTextures(RenderPipelineInfo* render_pipeline);
+  // Creates and attaches index, vertex, storage, uniform and depth-stencil
+  // buffers.Used in the Compute pipeline creation.
+  Result AttachBuffers(ComputePipelineInfo* compute_pipeline);
   // Creates and submits a command to copy dawn textures back to amber color
-  // attachments
+  // attachments.
   Result MapDeviceTextureToHostBuffer(const RenderPipelineInfo& render_pipeline,
                                       const ::dawn::Device& device);
+  // Creates and submits a command to copy dawn buffers back to amber buffers
+  Result MapDeviceBufferToHostBuffer(
+      const ComputePipelineInfo& compute_pipeline,
+      const ::dawn::Device& device);
 
   // Borrowed from the engine config
   ::dawn::Device* device_ = nullptr;
diff --git a/src/dawn/pipeline_info.h b/src/dawn/pipeline_info.h
index 49b0061..8ea95dd 100644
--- a/src/dawn/pipeline_info.h
+++ b/src/dawn/pipeline_info.h
@@ -82,6 +82,17 @@ struct ComputePipelineInfo {
 
   ::amber::Pipeline* pipeline = nullptr;
   ::dawn::ShaderModule compute_shader;
+
+  // storage and uniform buffers
+  std::vector<::dawn::Buffer> buffers;
+
+  std::vector<::dawn::BindGroup> bind_groups;
+  std::vector<::dawn::BindGroupLayout> bind_group_layouts;
+
+  // Mapping from the <descriptor_set, binding> to dawn buffer index in buffers
+  std::unordered_map<std::pair<uint32_t, uint32_t>, uint32_t, hash_pair>
+      buffer_map;
+  std::set<int> used_descriptor_set;
 };
 
 /// Holds either a render or compute pipeline.
diff --git a/src/vkscript/command_parser.cc b/src/vkscript/command_parser.cc
index 1babe0d..745d320 100644
--- a/src/vkscript/command_parser.cc
+++ b/src/vkscript/command_parser.cc
@@ -621,21 +621,7 @@ Result CommandParser::ProcessSSBO() {
     if (!r.IsSuccess())
       return r;
 
-    // Multiply by the input needed because the value count will use the needed
-    // input as the multiplier
-    uint32_t value_count =
-        ((cmd->GetOffset() / buf->GetFormat()->SizeInBytes()) *
-         buf->GetFormat()->InputNeededPerElement()) +
-        static_cast<uint32_t>(values.size());
-    // The buffer should only be resized to become bigger. This means that if a
-    // command was run to set the buffer size we'll honour that size until a
-    // request happens to make the buffer bigger.
-    if (value_count > buf->ValueCount())
-      buf->SetValueCount(value_count);
-
-    // Even if the value count doesn't change, the buffer is still resized
-    // because this maybe the first time data is set into the buffer.
-    buf->ResizeTo(buf->GetSizeInBytes());
+    buf->RecalculateMaxSizeInBytes(values, cmd->GetOffset());
 
     cmd->SetValues(std::move(values));
 
@@ -789,6 +775,8 @@ Result CommandParser::ProcessUniform() {
   if (!r.IsSuccess())
     return r;
 
+  buf->RecalculateMaxSizeInBytes(values, cmd->GetOffset());
+
   if (cmd->IsPushConstant())
     buf->SetData(values);
   else
diff --git a/tests/cases/multiple_ssbo_with_sparse_descriptor_set_in_compute_pipeline_less_than_4.vkscript b/tests/cases/multiple_ssbo_with_sparse_descriptor_set_in_compute_pipeline_less_than_4.vkscript
new file mode 100644
index 0000000..a6024df
--- /dev/null
+++ b/tests/cases/multiple_ssbo_with_sparse_descriptor_set_in_compute_pipeline_less_than_4.vkscript
@@ -0,0 +1,116 @@
+# Copyright 2018 The Amber 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
+#
+#     https://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.
+
+[compute shader]
+#version 430
+
+layout(set = 1, binding = 0) buffer block1 {
+  float data_set1_binding0[11];
+  float data_set1_binding0_result[11];
+};
+
+layout(set = 1, binding = 2) buffer block2 {
+  float data_set1_binding2[11];
+  float data_set1_binding2_result[11];
+};
+
+layout(set = 3, binding = 1) buffer block3 {
+  float data_set3_binding1[11];
+  float data_set3_binding1_result[11];
+};
+
+layout(set = 0, binding = 3) buffer block4 {
+  float data_set5_binding3[11];
+  float data_set5_binding3_result[11];
+};
+
+void main() {
+  for (int i = 0; i < 11; ++i) {
+    data_set1_binding0_result[i] = data_set1_binding2[i];
+    data_set1_binding2_result[i] = data_set3_binding1[i];
+    data_set3_binding1_result[i] = data_set5_binding3[i];
+    data_set5_binding3_result[i] = data_set1_binding0[i];
+  }
+}
+
+[test]
+ssbo 1:0 subdata float  0 0.1 0.2 0.3 0.4
+ssbo 1:2 subdata float  0 0.1 0.2 0.3 0.4
+ssbo 3:1 subdata float  0 0.1 0.2 0.3 0.4
+ssbo 0:3 subdata float  0 0.1 0.2 0.3 0.4
+
+compute 4 1 1
+
+probe ssbo float 1:0 0  ~= 0.1 0.2  0.3  0.4
+probe ssbo float 1:2 0  ~= 0.1 0.2  0.3  0.4
+probe ssbo float 3:1 0  ~= 0.1 0.2  0.3  0.4
+probe ssbo float 0:3 0  ~= 0.1 0.2  0.3  0.4
+
+ssbo 1:0 subdata float 0 0.1 0.2  0.3  0.4 \
+                         0.5 0.6  0.7  0.8 \
+                         0.9 0.10 0.11     \
+                         0.1 0.2  0.3  0.4 \
+                         0.5 0.6  0.7  0.8 \
+                         0.9 0.10 0.11
+
+ssbo 1:2 subdata float 0 0.57 0.56 0.55 0.54 \
+                         0.53 0.52 0.51 0.50 \
+                         0.49 0.48 0.47      \
+                         0.57 0.56 0.55 0.54 \
+                         0.53 0.52 0.51 0.50 \
+                         0.49 0.48 0.47
+
+ssbo 3:1 subdata float 0 0.21 0.22 0.23 0.24 \
+                         0.25 0.26 0.27 0.28 \
+                         0.29 0.30 0.31      \
+                         0.21 0.22 0.23 0.24 \
+                         0.25 0.26 0.27 0.28 \
+                         0.29 0.30 0.31
+
+ssbo 0:3 subdata float 0 0.23  0.229 0.228 0.227 \
+                         0.226 0.225 0.224 0.223 \
+                         0.222 0.221 0.22        \
+                         0.23  0.229 0.228 0.227 \
+                         0.226 0.225 0.224 0.223 \
+                         0.222 0.221 0.22
+
+compute 4 1 1
+
+probe ssbo float 1:0 0  ~= 0.1 0.2  0.3  0.4 \
+                           0.5 0.6  0.7  0.8 \
+                           0.9 0.10 0.11
+probe ssbo float 1:0 44 ~= 0.57 0.56 0.55 0.54 \
+                           0.53 0.52 0.51 0.50 \
+                           0.49 0.48 0.47
+
+probe ssbo float 1:2 0  ~= 0.57 0.56 0.55 0.54 \
+                           0.53 0.52 0.51 0.50 \
+                           0.49 0.48 0.47
+probe ssbo float 1:2 44 ~= 0.21 0.22 0.23 0.24 \
+                           0.25 0.26 0.27 0.28 \
+                           0.29 0.30 0.31
+
+probe ssbo float 3:1 0  ~= 0.21 0.22 0.23 0.24 \
+                           0.25 0.26 0.27 0.28 \
+                           0.29 0.30 0.31
+probe ssbo float 3:1 44 ~= 0.23  0.229 0.228 0.227 \
+                           0.226 0.225 0.224 0.223 \
+                           0.222 0.221 0.22
+
+probe ssbo float 0:3 0  ~= 0.23  0.229 0.228 0.227 \
+                           0.226 0.225 0.224 0.223 \
+                           0.222 0.221 0.22
+probe ssbo float 0:3 44 ~= 0.1 0.2  0.3  0.4 \
+                           0.5 0.6  0.7  0.8 \
+                           0.9 0.10 0.11
diff --git a/tests/run_tests.py b/tests/run_tests.py
index 7702efa..83abdca 100755
--- a/tests/run_tests.py
+++ b/tests/run_tests.py
@@ -55,37 +55,9 @@ SUPPRESSIONS = {
  }
 
 SUPPRESSIONS_DAWN = [
-  # Dawn does not support push constants, tessellation, geometry shader
+  # Dawn does not support push constants
   "graphics_push_constants.amber",
   "graphics_push_constants.vkscript",
-  "draw_triangle_list_using_geom_shader.vkscript",
-  "draw_triangle_list_using_tessellation.vkscript",
-  # Dawn requires a fragmentStage now and in the medium term
-  "position_to_ssbo.amber",
-  # Draw rect command is not supported in a pipeline with more than one vertex
-  # buffer attached
-  "draw_array_after_draw_rect.vkscript",
-  "draw_rect_after_draw_array.vkscript",
-  "draw_rect_and_draw_array_mixed.vkscript",
-  # DoCompute is not implemented in Dawn backend
-  "compute_accumulated_ubo_definition.amber",
-  "compute_accumulated_ubo_definition.vkscript",
-  "compute_mat2x2.amber",
-  "compute_mat2x2.vkscript",
-  "compute_mat2x2float.vkscript",
-  "compute_mat2x3.vkscript",
-  "compute_mat2x3float.vkscript",
-  "compute_mat3x2.vkscript",
-  "compute_mat3x2float.vkscript",
-  "compute_mat3x3.vkscript",
-  "compute_mat3x3float.vkscript",
-  "compute_mat3x4.vkscript",
-  "compute_mat3x4float.vkscript",
-  "compute_mat4x3.vkscript",
-  "compute_mat4x3float.vkscript",
-  "compute_nothing.vkscript",
-  "compute_nothing_with_ssbo.vkscript",
-  "compute_probe_mat3.vkscript",
   "compute_push_const_mat2x2.vkscript",
   "compute_push_const_mat2x2float.vkscript",
   "compute_push_const_mat2x3.vkscript",
@@ -100,24 +72,25 @@ SUPPRESSIONS_DAWN = [
   "compute_push_const_mat4x3float.vkscript",
   "compute_push_constant_and_ssbo.amber",
   "compute_push_constant_and_ssbo.vkscript",
-  "compute_ssbo.vkscript",
-  "compute_ssbo_with_entrypoint_command.vkscript",
-  "compute_ssbo_with_tolerance.amber",
-  "compute_ssbo_with_tolerance.vkscript",
-  "compute_ssbo_without_probe.vkscript",
-  "compute_ubo_and_ssbo.vkscript",
-  "repeat.amber",
-  "scratch_ssbo.vkscript",
-  "shader_specialization.amber",
-  "ssbo_subdata_size.vkscript",
+  # Dawn does not support tessellation or geometry shader
+  "draw_triangle_list_using_geom_shader.vkscript",
+  "draw_triangle_list_using_tessellation.vkscript",
+  # Dawn does not support sparse descriptor_set
+  "compute_nothing_with_ssbo.vkscript",
+  # Dawn requires a fragmentStage now and in the medium term
+  "position_to_ssbo.amber",
+  # DrawRect command is not supported in a pipeline with more than one vertex
+  # buffer attached
+  "draw_array_after_draw_rect.vkscript",
+  "draw_rect_after_draw_array.vkscript",
+  "draw_rect_and_draw_array_mixed.vkscript",
   # Dawn DoCommands require a pipeline
   "probe_no_compute_with_multiple_ssbo_commands.vkscript",
   "probe_no_compute_with_ssbo.vkscript",
-  # Sparse descriptor sets are not supported in Dawn backend (issue #573)
-  "multiple_ssbo_update_with_graphics_pipeline.vkscript",
+  # Max number of descriptor sets is 4 in Dawn
   "multiple_ssbo_with_sparse_descriptor_set_in_compute_pipeline.vkscript",
-  "multiple_ubo_update_with_graphics_pipeline.vkscript",
-  # DoEntryPoint is not supported in Dawn backend
+  # DoEntryPoint is not supported in Dawn backend, yet
+  "compute_ssbo_with_entrypoint_command.vkscript",
   "entry_point.amber",
   # framebuffer format is not supported according to table "Mandatory format
   # support" in Vulkan spec: VK_FORMAT_FEATURE_COLOR_ATTACHMENT_BIT = 0
@@ -127,9 +100,15 @@ SUPPRESSIONS_DAWN = [
   "draw_triangle_list_in_r8g8b8a8_snorm_color_frame.vkscript",
   "draw_triangle_list_in_r8g8b8a8_srgb_color_frame.vkscript",
   # Currently not working, under investigation
+  "multiple_ubo_update_with_graphics_pipeline.vkscript",
   "draw_triangle_list_with_depth.vkscript",
   "non_default_entry_point.amber",
-  "clear_with_depth.amber"
+  "clear_with_depth.amber",
+  "opencl_bind_buffer.amber",
+  "opencl_c_copy.amber",
+  "opencl_set_arg.amber",
+  "shader_specialization.amber",
+
 ]
 
 class TestCase: