d3d12: Fix race condition when getting query results

Before, when an application called into d3d12_query_result, and
the results were not ready to be read, a flush-and-wait would
be attempted via synchronized mapping of the query result resource.

This can end up calling close/execute on the command list while it is
already being executed by the driver thread.

With the current fence value attached to the query, we now wait
for completion if necessary and then map the resource unsynchronized, or
return false if the result is not ready and wait == false.

Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/20217>

src/gallium/drivers/d3d12/d3d12_batch.cpp

@@ -223,8 +223,14 @@ d3d12_end_batch(struct d3d12_context *ctx, struct d3d12_batch *batch)
       count_to_execute--;
    }
    screen->cmdqueue->ExecuteCommandLists(count_to_execute, to_execute);
+
    batch->fence = d3d12_create_fence(screen);
 
+   util_dynarray_foreach(&ctx->ended_queries, struct d3d12_query*, query) {
+      (*query)->fence_value = screen->fence_value;
+   }
+   util_dynarray_clear(&ctx->ended_queries);
+
    mtx_unlock(&screen->submit_mutex);
 }
 

src/gallium/drivers/d3d12/d3d12_context.cpp

@@ -102,6 +102,7 @@ d3d12_context_destroy(struct pipe_context *pctx)
    pipe_resource_reference(&ctx->pstipple.texture, nullptr);
    pipe_sampler_view_reference(&ctx->pstipple.sampler_view, nullptr);
    util_dynarray_fini(&ctx->recently_destroyed_bos);
+   util_dynarray_fini(&ctx->ended_queries);
    FREE(ctx->pstipple.sampler_cso);
 
    u_suballocator_destroy(&ctx->query_allocator);

src/gallium/drivers/d3d12/d3d12_context.h

@@ -248,6 +248,7 @@ struct d3d12_context {
    ID3D12GraphicsCommandList *state_fixup_cmdlist;
 
    struct list_head active_queries;
+   struct util_dynarray ended_queries;
    bool queries_disabled;
 
    struct d3d12_descriptor_pool *sampler_pool;

src/gallium/drivers/d3d12/d3d12_query.cpp

@@ -176,7 +176,7 @@ d3d12_destroy_query(struct pipe_context *pctx,
 static bool
 accumulate_subresult(struct d3d12_context *ctx, struct d3d12_query *q_parent,
                      unsigned sub_query,
-                     union pipe_query_result *result, bool write, bool wait)
+                     union pipe_query_result *result, bool write)
 {
    struct pipe_transfer *transfer = NULL;
    struct d3d12_screen *screen = d3d12_screen(ctx->base.screen);
@@ -186,8 +186,8 @@ accumulate_subresult(struct d3d12_context *ctx, struct d3d12_query *q_parent,
 
    if (write)
       access |= PIPE_MAP_WRITE;
-   if (!wait)
-      access |= PIPE_MAP_DONTBLOCK;
+   access |= PIPE_MAP_UNSYNCHRONIZED;
+
    results = pipe_buffer_map_range(&ctx->base, q->buffer, q->buffer_offset,
                                    q->num_queries * q->query_size,
                                    access, &transfer);
@@ -282,32 +282,32 @@ accumulate_subresult(struct d3d12_context *ctx, struct d3d12_query *q_parent,
 
 static bool
 accumulate_result(struct d3d12_context *ctx, struct d3d12_query *q,
-                  union pipe_query_result *result, bool write, bool wait)
+                  union pipe_query_result *result, bool write)
 {
    union pipe_query_result local_result;
 
    switch (q->type) {
    case PIPE_QUERY_PRIMITIVES_GENERATED:
-      if (!accumulate_subresult(ctx, q, 0, &local_result, write, wait))
+      if (!accumulate_subresult(ctx, q, 0, &local_result, write))
          return false;
       result->u64 = local_result.so_statistics.primitives_storage_needed;
 
-      if (!accumulate_subresult(ctx, q, 1, &local_result, write, wait))
+      if (!accumulate_subresult(ctx, q, 1, &local_result, write))
          return false;
       result->u64 += local_result.pipeline_statistics.gs_primitives;
 
-      if (!accumulate_subresult(ctx, q, 2, &local_result, write, wait))
+      if (!accumulate_subresult(ctx, q, 2, &local_result, write))
          return false;
       result->u64 += local_result.pipeline_statistics.ia_primitives;
       return true;
    case PIPE_QUERY_PRIMITIVES_EMITTED:
-      if (!accumulate_subresult(ctx, q, 0, &local_result, write, wait))
+      if (!accumulate_subresult(ctx, q, 0, &local_result, write))
          return false;
       result->u64 = local_result.so_statistics.num_primitives_written;
       return true;
    default:
       assert(num_sub_queries(q->type) == 1);
-      return accumulate_subresult(ctx, q, 0, result, write, wait);
+      return accumulate_subresult(ctx, q, 0, result, write);
    }
 }
 
@@ -332,6 +332,26 @@ subquery_should_be_active(struct d3d12_context *ctx, struct d3d12_query *q, unsi
    }
 }
 
+static bool 
+query_ensure_ready(struct d3d12_screen* screen, struct d3d12_context* ctx, struct d3d12_query* query, bool wait)
+{
+   // If the query is not flushed, it won't have 
+   // been submitted yet, and won't have a waitable 
+   // fence value
+   if (query->fence_value == UINT64_MAX) {
+      d3d12_flush_cmdlist(ctx);
+   }
+
+   if (screen->fence->GetCompletedValue() < query->fence_value){
+      if (!wait)
+         return false;
+
+      screen->fence->SetEventOnCompletion(query->fence_value, NULL);
+   }
+
+   return true;
+}
+
 static void
 begin_subquery(struct d3d12_context *ctx, struct d3d12_query *q_parent, unsigned sub_query)
 {
@@ -340,7 +360,7 @@ begin_subquery(struct d3d12_context *ctx, struct d3d12_query *q_parent, unsigned
       union pipe_query_result result;
 
       /* Accumulate current results and store in first slot */
-      accumulate_subresult(ctx, q_parent, sub_query, &result, true, true);
+      accumulate_subresult(ctx, q_parent, sub_query, &result, true);
       q->curr_query = 1;
    }
 
@@ -380,7 +400,7 @@ begin_timer_query(struct d3d12_context *ctx, struct d3d12_query *q_parent, bool
 
       /* Accumulate current results and store in first slot */
       d3d12_flush_cmdlist_and_wait(ctx);
-      accumulate_subresult(ctx, q_parent, 0, &result, true, true);
+      accumulate_subresult(ctx, q_parent, 0, &result, true);
       q->curr_query = 2;
    }
 
@@ -463,6 +483,10 @@ d3d12_end_query(struct pipe_context *pctx,
    struct d3d12_context *ctx = d3d12_context(pctx);
    struct d3d12_query *query = (struct d3d12_query *)q;
 
+   // Assign the sentinel and track now that the query is ended
+   query->fence_value = UINT64_MAX;
+   util_dynarray_append(&ctx->ended_queries, d3d12_query*, query);
+
    end_query(ctx, query);
 
    if (query->type != PIPE_QUERY_TIMESTAMP &&
@@ -478,9 +502,13 @@ d3d12_get_query_result(struct pipe_context *pctx,
                       union pipe_query_result *result)
 {
    struct d3d12_context *ctx = d3d12_context(pctx);
+   struct d3d12_screen *screen = d3d12_screen(ctx->base.screen);
    struct d3d12_query *query = (struct d3d12_query *)q;
 
-   return accumulate_result(ctx, query, result, false, wait);
+   if (!query_ensure_ready(screen, ctx, query, wait))
+      return false;
+
+   return accumulate_result(ctx, query, result, false);
 }
 
 void
@@ -551,7 +579,7 @@ d3d12_render_condition(struct pipe_context *pctx,
    if (mode == PIPE_RENDER_COND_WAIT) {
       d3d12_flush_cmdlist_and_wait(ctx);
       union pipe_query_result result;
-      accumulate_result(ctx, (d3d12_query *)pquery, &result, true, true);
+      accumulate_result(ctx, (d3d12_query *)pquery, &result, true);
    }
 
    struct d3d12_resource *res = (struct d3d12_resource *)query->subqueries[0].buffer;
@@ -591,6 +619,7 @@ d3d12_context_query_init(struct pipe_context *pctx)
 {
    struct d3d12_context *ctx = d3d12_context(pctx);
    list_inithead(&ctx->active_queries);
+   util_dynarray_init(&ctx->ended_queries, NULL);
 
    u_suballocator_init(&ctx->query_allocator, &ctx->base, 4096, 0, PIPE_USAGE_STAGING,
                          0, true);

src/gallium/drivers/d3d12/d3d12_query.h

@@ -65,6 +65,13 @@ struct d3d12_query {
    struct list_head active_list;
    struct d3d12_resource* predicate;
 
+   /* 
+   * Used to track if a query's results are ready to be read asynchronously
+   * 
+   * Initialized to 0, it is set to UINT64_MAX when the query is ended but before it is flushed
+   * At flush, it is set to the fence_value associated with the work it was
+   * submitted with
+   */
    uint64_t fence_value;
 };