Query API: QuerySet on Metal

- Add query set creation on Metal backend
- Enable end2end tests for query set creation

Bug: dawn:434
Change-Id: I7fe013192ae215b6b97cfdb646a8dd6f2596d4af
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/28800
Commit-Queue: Corentin Wallez <cwallez@chromium.org>
Reviewed-by: Corentin Wallez <cwallez@chromium.org>
diff --git a/src/dawn_native/BUILD.gn b/src/dawn_native/BUILD.gn
index 23bc55a..b9dfd14 100644
--- a/src/dawn_native/BUILD.gn
+++ b/src/dawn_native/BUILD.gn
@@ -388,6 +388,8 @@
       "metal/Forward.h",
       "metal/PipelineLayoutMTL.h",
       "metal/PipelineLayoutMTL.mm",
+      "metal/QuerySetMTL.h",
+      "metal/QuerySetMTL.mm",
       "metal/QueueMTL.h",
       "metal/QueueMTL.mm",
       "metal/RenderPipelineMTL.h",
diff --git a/src/dawn_native/Toggles.cpp b/src/dawn_native/Toggles.cpp
index 6bd0625..e98fd80 100644
--- a/src/dawn_native/Toggles.cpp
+++ b/src/dawn_native/Toggles.cpp
@@ -104,6 +104,13 @@
                "Disables the use of sampler compare on Metal. This is unsupported before A9 "
                "processors.",
                "https://crbug.com/dawn/342"}},
+             {Toggle::MetalUseSharedModeForCounterSampleBuffer,
+              {"metal_use_shared_mode_for_counter_sample_buffer",
+               "The query set on Metal need to create MTLCounterSampleBuffer which storage mode "
+               "must be either MTLStorageModeShared or MTLStorageModePrivate. But the private mode "
+               "does not work properly on Intel platforms. The workaround is use shared mode "
+               "instead.",
+               "https://crbug.com/dawn/434"}},
              {Toggle::DisableBaseVertex,
               {"disable_base_vertex",
                "Disables the use of non-zero base vertex which is unsupported on some platforms.",
diff --git a/src/dawn_native/Toggles.h b/src/dawn_native/Toggles.h
index e3821bd..a84edc2 100644
--- a/src/dawn_native/Toggles.h
+++ b/src/dawn_native/Toggles.h
@@ -36,6 +36,7 @@
         SkipValidation,
         VulkanUseD32S8,
         MetalDisableSamplerCompare,
+        MetalUseSharedModeForCounterSampleBuffer,
         DisableBaseVertex,
         DisableBaseInstance,
         UseD3D12SmallShaderVisibleHeapForTesting,
diff --git a/src/dawn_native/metal/BackendMTL.mm b/src/dawn_native/metal/BackendMTL.mm
index a5e7c47..6916c6d 100644
--- a/src/dawn_native/metal/BackendMTL.mm
+++ b/src/dawn_native/metal/BackendMTL.mm
@@ -214,13 +214,16 @@
             if ([mDevice supportsFeatureSet:MTLFeatureSet_macOS_GPUFamily1_v1]) {
                 mSupportedExtensions.EnableExtension(Extension::TextureCompressionBC);
             }
-
-            if (@available(macOS 10.15, *)) {
-                mSupportedExtensions.EnableExtension(Extension::PipelineStatisticsQuery);
-                mSupportedExtensions.EnableExtension(Extension::TimestampQuery);
-            }
 #endif
 
+            if (@available(macOS 10.15, iOS 14.0, *)) {
+                if ([mDevice supportsFamily:MTLGPUFamilyMac2] ||
+                    [mDevice supportsFamily:MTLGPUFamilyApple5]) {
+                    mSupportedExtensions.EnableExtension(Extension::PipelineStatisticsQuery);
+                    mSupportedExtensions.EnableExtension(Extension::TimestampQuery);
+                }
+            }
+
             mSupportedExtensions.EnableExtension(Extension::ShaderFloat16);
         }
 
diff --git a/src/dawn_native/metal/DeviceMTL.mm b/src/dawn_native/metal/DeviceMTL.mm
index e6b072d..8ef2b6e 100644
--- a/src/dawn_native/metal/DeviceMTL.mm
+++ b/src/dawn_native/metal/DeviceMTL.mm
@@ -14,6 +14,7 @@
 
 #include "dawn_native/metal/DeviceMTL.h"
 
+#include "common/GPUInfo.h"
 #include "common/Platform.h"
 #include "dawn_native/BackendConnection.h"
 #include "dawn_native/BindGroupLayout.h"
@@ -25,6 +26,7 @@
 #include "dawn_native/metal/CommandBufferMTL.h"
 #include "dawn_native/metal/ComputePipelineMTL.h"
 #include "dawn_native/metal/PipelineLayoutMTL.h"
+#include "dawn_native/metal/QuerySetMTL.h"
 #include "dawn_native/metal/QueueMTL.h"
 #include "dawn_native/metal/RenderPipelineMTL.h"
 #include "dawn_native/metal/SamplerMTL.h"
@@ -106,6 +108,14 @@
 
         // TODO(jiawei.shao@intel.com): tighten this workaround when the driver bug is fixed.
         SetToggle(Toggle::AlwaysResolveIntoZeroLevelAndLayer, true);
+
+        // TODO(hao.x.li@intel.com): Use MTLStorageModeShared instead of MTLStorageModePrivate when
+        // creating MTLCounterSampleBuffer in QuerySet on Intel platforms, otherwise it fails to
+        // create the buffer. Change to use MTLStorageModePrivate when the bug is fixed.
+        if (@available(macOS 10.15, iOS 14.0, *)) {
+            bool useSharedMode = gpu_info::IsIntel(this->GetAdapter()->GetPCIInfo().vendorId);
+            SetToggle(Toggle::MetalUseSharedModeForCounterSampleBuffer, useSharedMode);
+        }
     }
 
     ResultOrError<BindGroupBase*> Device::CreateBindGroupImpl(
@@ -132,7 +142,7 @@
         return new PipelineLayout(this, descriptor);
     }
     ResultOrError<QuerySetBase*> Device::CreateQuerySetImpl(const QuerySetDescriptor* descriptor) {
-        return DAWN_UNIMPLEMENTED_ERROR("Waiting for implementation");
+        return QuerySet::Create(this, descriptor);
     }
     ResultOrError<RenderPipelineBase*> Device::CreateRenderPipelineImpl(
         const RenderPipelineDescriptor* descriptor) {
diff --git a/src/dawn_native/metal/QuerySetMTL.h b/src/dawn_native/metal/QuerySetMTL.h
new file mode 100644
index 0000000..d1ac5df
--- /dev/null
+++ b/src/dawn_native/metal/QuerySetMTL.h
@@ -0,0 +1,50 @@
+// Copyright 2020 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_METAL_QUERYSETMTL_H_
+#define DAWNNATIVE_METAL_QUERYSETMTL_H_
+
+#include "dawn_native/QuerySet.h"
+
+#import <Metal/Metal.h>
+
+namespace dawn_native { namespace metal {
+
+    class Device;
+
+    class QuerySet final : public QuerySetBase {
+      public:
+        static ResultOrError<QuerySet*> Create(Device* device,
+                                               const QuerySetDescriptor* descriptor);
+
+        id<MTLBuffer> GetVisibilityBuffer() const;
+        id<MTLCounterSampleBuffer> GetCounterSampleBuffer() const
+            API_AVAILABLE(macos(10.15), ios(14.0));
+
+      private:
+        ~QuerySet() override;
+        using QuerySetBase::QuerySetBase;
+        MaybeError Initialize();
+
+        // Dawn API
+        void DestroyImpl() override;
+
+        id<MTLBuffer> mVisibilityBuffer = nil;
+        id<MTLCounterSampleBuffer> mCounterSampleBuffer API_AVAILABLE(macos(10.15),
+                                                                      ios(14.0)) = nil;
+    };
+
+}}  // namespace dawn_native::metal
+
+#endif  // DAWNNATIVE_METAL_QUERYSETMTL_H_
diff --git a/src/dawn_native/metal/QuerySetMTL.mm b/src/dawn_native/metal/QuerySetMTL.mm
new file mode 100644
index 0000000..4e395e3
--- /dev/null
+++ b/src/dawn_native/metal/QuerySetMTL.mm
@@ -0,0 +1,134 @@
+// Copyright 2020 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/metal/QuerySetMTL.h"
+
+#include "common/Math.h"
+#include "common/Platform.h"
+#include "dawn_native/metal/DeviceMTL.h"
+
+namespace dawn_native { namespace metal {
+
+    namespace {
+
+        ResultOrError<id<MTLCounterSampleBuffer>> CreateCounterSampleBuffer(
+            Device* device,
+            MTLCommonCounterSet counterSet,
+            uint32_t count) API_AVAILABLE(macos(10.15), ios(14.0)) {
+            MTLCounterSampleBufferDescriptor* descriptor = [MTLCounterSampleBufferDescriptor new];
+
+            // To determine which counters are available from a device, we need to iterate through
+            // the counterSets property of a MTLDevice. Then configure which counters will be
+            // sampled by creating a MTLCounterSampleBufferDescriptor and setting its counterSet
+            // property to the matched one of the available set.
+            for (id<MTLCounterSet> set in device->GetMTLDevice().counterSets) {
+                if ([set.name isEqualToString:counterSet]) {
+                    descriptor.counterSet = set;
+                    break;
+                }
+            }
+            ASSERT(descriptor.counterSet != nil);
+            descriptor.sampleCount = count;
+            descriptor.storageMode = MTLStorageModePrivate;
+            if (device->IsToggleEnabled(Toggle::MetalUseSharedModeForCounterSampleBuffer)) {
+                descriptor.storageMode = MTLStorageModeShared;
+            }
+
+            NSError* error = nil;
+            id<MTLCounterSampleBuffer> counterSampleBuffer =
+                [device->GetMTLDevice() newCounterSampleBufferWithDescriptor:descriptor
+                                                                       error:&error];
+            if (error != nil) {
+                const char* errorString = [error.localizedDescription UTF8String];
+                return DAWN_INTERNAL_ERROR(std::string("Error creating query set: ") + errorString);
+            }
+
+            return counterSampleBuffer;
+        }
+    }
+
+    // static
+    ResultOrError<QuerySet*> QuerySet::Create(Device* device,
+                                              const QuerySetDescriptor* descriptor) {
+        Ref<QuerySet> queryset = AcquireRef(new QuerySet(device, descriptor));
+        DAWN_TRY(queryset->Initialize());
+        return queryset.Detach();
+    }
+
+    MaybeError QuerySet::Initialize() {
+        Device* device = ToBackend(GetDevice());
+
+        switch (GetQueryType()) {
+            case wgpu::QueryType::Occlusion: {
+                // Create buffer for writing 64-bit results.
+                NSUInteger bufferSize = static_cast<NSUInteger>(GetQueryCount() * sizeof(uint64_t));
+                mVisibilityBuffer =
+                    [device->GetMTLDevice() newBufferWithLength:bufferSize
+                                                        options:MTLResourceStorageModePrivate];
+                break;
+            }
+            case wgpu::QueryType::PipelineStatistics:
+                if (@available(macOS 10.15, iOS 14.0, *)) {
+                    DAWN_TRY_ASSIGN(mCounterSampleBuffer,
+                                    CreateCounterSampleBuffer(device, MTLCommonCounterSetStatistic,
+                                                              GetQueryCount()));
+                } else {
+                    UNREACHABLE();
+                }
+                break;
+            case wgpu::QueryType::Timestamp:
+                if (@available(macOS 10.15, iOS 14.0, *)) {
+                    DAWN_TRY_ASSIGN(mCounterSampleBuffer,
+                                    CreateCounterSampleBuffer(device, MTLCommonCounterSetTimestamp,
+                                                              GetQueryCount()));
+                } else {
+                    UNREACHABLE();
+                }
+                break;
+            default:
+                UNREACHABLE();
+                break;
+        }
+
+        return {};
+    }
+
+    id<MTLBuffer> QuerySet::GetVisibilityBuffer() const {
+        return mVisibilityBuffer;
+    }
+
+    id<MTLCounterSampleBuffer> QuerySet::GetCounterSampleBuffer() const
+        API_AVAILABLE(macos(10.15), ios(14.0)) {
+        return mCounterSampleBuffer;
+    }
+
+    QuerySet::~QuerySet() {
+        DestroyInternal();
+    }
+
+    void QuerySet::DestroyImpl() {
+        if (mVisibilityBuffer != nil) {
+            [mVisibilityBuffer release];
+            mVisibilityBuffer = nil;
+        }
+
+        if (@available(macOS 10.15, iOS 14.0, *)) {
+            if (mCounterSampleBuffer != nil) {
+                [mCounterSampleBuffer release];
+                mCounterSampleBuffer = nil;
+            }
+        }
+    }
+
+}}  // namespace dawn_native::metal
diff --git a/src/tests/end2end/QueryTests.cpp b/src/tests/end2end/QueryTests.cpp
index 943ead2..dfcebc1 100644
--- a/src/tests/end2end/QueryTests.cpp
+++ b/src/tests/end2end/QueryTests.cpp
@@ -51,7 +51,7 @@
     querySet.Destroy();
 }
 
-DAWN_INSTANTIATE_TEST(OcclusionQueryTests, D3D12Backend(), VulkanBackend());
+DAWN_INSTANTIATE_TEST(OcclusionQueryTests, D3D12Backend(), MetalBackend(), VulkanBackend());
 
 class PipelineStatisticsQueryTests : public QueryTests {
   protected:
@@ -85,7 +85,10 @@
     device.CreateQuerySet(&descriptor);
 }
 
-DAWN_INSTANTIATE_TEST(PipelineStatisticsQueryTests, D3D12Backend(), VulkanBackend());
+DAWN_INSTANTIATE_TEST(PipelineStatisticsQueryTests,
+                      D3D12Backend(),
+                      MetalBackend(),
+                      VulkanBackend());
 
 class TimestampExpectation : public detail::Expectation {
   public:
@@ -138,6 +141,9 @@
 
 // Test calling timestamp query from command encoder
 TEST_P(TimestampQueryTests, TimestampOnCommandEncoder) {
+    // TODO(hao.x.li@intel.com): Waiting for timestamp query implementation on Metal
+    DAWN_SKIP_TEST_IF(IsMetal());
+
     constexpr uint32_t kQueryCount = 2;
 
     wgpu::QuerySet querySet = CreateQuerySetForTimestamp(kQueryCount);
@@ -155,6 +161,9 @@
 
 // Test calling timestamp query from render pass encoder
 TEST_P(TimestampQueryTests, TimestampOnRenderPass) {
+    // TODO(hao.x.li@intel.com): Waiting for timestamp query implementation on Metal
+    DAWN_SKIP_TEST_IF(IsMetal());
+
     constexpr uint32_t kQueryCount = 2;
 
     wgpu::QuerySet querySet = CreateQuerySetForTimestamp(kQueryCount);
@@ -175,6 +184,9 @@
 
 // Test calling timestamp query from compute pass encoder
 TEST_P(TimestampQueryTests, TimestampOnComputePass) {
+    // TODO(hao.x.li@intel.com): Waiting for timestamp query implementation on Metal
+    DAWN_SKIP_TEST_IF(IsMetal());
+
     constexpr uint32_t kQueryCount = 2;
 
     wgpu::QuerySet querySet = CreateQuerySetForTimestamp(kQueryCount);
@@ -197,6 +209,9 @@
     // TODO(hao.x.li@intel.com): Failed on old Intel Vulkan driver on Windows, need investigation.
     DAWN_SKIP_TEST_IF(IsWindows() && IsIntel() && IsVulkan());
 
+    // TODO(hao.x.li@intel.com): Waiting for timestamp query implementation on Metal
+    DAWN_SKIP_TEST_IF(IsMetal());
+
     constexpr uint32_t kQueryCount = 2;
     constexpr uint64_t kZero = 0;
 
@@ -232,4 +247,4 @@
     }
 }
 
-DAWN_INSTANTIATE_TEST(TimestampQueryTests, D3D12Backend(), VulkanBackend());
+DAWN_INSTANTIATE_TEST(TimestampQueryTests, D3D12Backend(), MetalBackend(), VulkanBackend());