anv: Add support for binding acceleration structures

Reviewed-by: Lionel Landwerlin <lionel.g.landwerlin@intel.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/8637>

src/intel/vulkan/anv_cmd_buffer.c

@@ -1433,6 +1433,22 @@ void anv_CmdPushDescriptorSetKHR(
          }
          break;
 
+      case VK_DESCRIPTOR_TYPE_ACCELERATION_STRUCTURE_KHR: {
+         const VkWriteDescriptorSetAccelerationStructureKHR *accel_write =
+            vk_find_struct_const(write, WRITE_DESCRIPTOR_SET_ACCELERATION_STRUCTURE_KHR);
+         assert(accel_write->accelerationStructureCount ==
+                write->descriptorCount);
+         for (uint32_t j = 0; j < write->descriptorCount; j++) {
+            ANV_FROM_HANDLE(anv_acceleration_structure, accel,
+                            accel_write->pAccelerationStructures[j]);
+            anv_descriptor_set_write_acceleration_structure(cmd_buffer->device,
+                                                            set, accel,
+                                                            write->dstBinding,
+                                                            write->dstArrayElement + j);
+         }
+         break;
+      }
+
       default:
          break;
       }

src/intel/vulkan/anv_descriptor_set.c

@@ -91,6 +91,10 @@ anv_descriptor_data_for_type(const struct anv_physical_device *device,
       data = ANV_DESCRIPTOR_INLINE_UNIFORM;
       break;
 
+   case VK_DESCRIPTOR_TYPE_ACCELERATION_STRUCTURE_KHR:
+      data = ANV_DESCRIPTOR_ADDRESS_RANGE;
+      break;
+
    default:
       unreachable("Unsupported descriptor type");
    }
@@ -1548,6 +1552,35 @@ anv_descriptor_set_write_inline_uniform_data(struct anv_device *device,
    memcpy(desc_map + offset, data, size);
 }
 
+void
+anv_descriptor_set_write_acceleration_structure(struct anv_device *device,
+                                                struct anv_descriptor_set *set,
+                                                struct anv_acceleration_structure *accel,
+                                                uint32_t binding,
+                                                uint32_t element)
+{
+   const struct anv_descriptor_set_binding_layout *bind_layout =
+      &set->layout->binding[binding];
+   struct anv_descriptor *desc =
+      &set->descriptors[bind_layout->descriptor_index + element];
+
+   assert(bind_layout->type == VK_DESCRIPTOR_TYPE_ACCELERATION_STRUCTURE_KHR);
+   *desc = (struct anv_descriptor) {
+      .type = VK_DESCRIPTOR_TYPE_ACCELERATION_STRUCTURE_KHR,
+   };
+
+   struct anv_address_range_descriptor desc_data = { };
+   if (accel != NULL) {
+      desc_data.address = anv_address_physical(accel->address);
+      desc_data.range = accel->size;
+   }
+   assert(anv_descriptor_size(bind_layout) == sizeof(desc_data));
+
+   void *desc_map = set->desc_mem.map + bind_layout->descriptor_offset +
+                    element * sizeof(desc_data);
+   memcpy(desc_map, &desc_data, sizeof(desc_data));
+}
+
 void anv_UpdateDescriptorSets(
     VkDevice                                    _device,
     uint32_t                                    descriptorWriteCount,
@@ -1621,6 +1654,21 @@ void anv_UpdateDescriptorSets(
          break;
       }
 
+      case VK_DESCRIPTOR_TYPE_ACCELERATION_STRUCTURE_KHR: {
+         const VkWriteDescriptorSetAccelerationStructureKHR *accel_write =
+            vk_find_struct_const(write, WRITE_DESCRIPTOR_SET_ACCELERATION_STRUCTURE_KHR);
+         assert(accel_write->accelerationStructureCount ==
+                write->descriptorCount);
+         for (uint32_t j = 0; j < write->descriptorCount; j++) {
+            ANV_FROM_HANDLE(anv_acceleration_structure, accel,
+                            accel_write->pAccelerationStructures[j]);
+            anv_descriptor_set_write_acceleration_structure(device, set, accel,
+                                                            write->dstBinding,
+                                                            write->dstArrayElement + j);
+         }
+         break;
+      }
+
       default:
          break;
       }

src/intel/vulkan/anv_device.c

@@ -1363,7 +1363,7 @@ void anv_GetPhysicalDeviceFeatures2(
          features->accelerationStructureCaptureReplay = false;
          features->accelerationStructureIndirectBuild = false;
          features->accelerationStructureHostCommands = false;
-         features->descriptorBindingAccelerationStructureUpdateAfterBind = false;
+         features->descriptorBindingAccelerationStructureUpdateAfterBind = true;
          break;
       }
 
@@ -2161,10 +2161,10 @@ void anv_GetPhysicalDeviceProperties2(
          props->maxGeometryCount = (1u << 24) - 1;
          props->maxInstanceCount = (1u << 24) - 1;
          props->maxPrimitiveCount = (1u << 29) - 1;
-         props->maxPerStageDescriptorAccelerationStructures = 0;
-         props->maxPerStageDescriptorUpdateAfterBindAccelerationStructures = 0;
-         props->maxDescriptorSetAccelerationStructures = 0;
-         props->maxDescriptorSetUpdateAfterBindAccelerationStructures = 0;
+         props->maxPerStageDescriptorAccelerationStructures = UINT16_MAX;
+         props->maxPerStageDescriptorUpdateAfterBindAccelerationStructures = UINT16_MAX;
+         props->maxDescriptorSetAccelerationStructures = UINT16_MAX;
+         props->maxDescriptorSetUpdateAfterBindAccelerationStructures = UINT16_MAX;
          props->minAccelerationStructureScratchOffsetAlignment = 64;
          break;
       }

src/intel/vulkan/anv_nir_apply_pipeline_layout.c

@@ -753,6 +753,48 @@ try_lower_direct_buffer_intrinsic(nir_builder *b,
    return true;
 }
 
+static bool
+lower_load_accel_struct_desc(nir_builder *b,
+                             nir_intrinsic_instr *load_desc,
+                             struct apply_pipeline_layout_state *state)
+{
+   assert(load_desc->intrinsic == nir_intrinsic_load_vulkan_descriptor);
+
+   nir_intrinsic_instr *idx_intrin = nir_src_as_intrinsic(load_desc->src[0]);
+
+   /* It doesn't really matter what address format we choose as
+    * everything will constant-fold nicely.  Choose one that uses the
+    * actual descriptor buffer.
+    */
+   const nir_address_format addr_format =
+      nir_address_format_64bit_bounded_global;
+
+   uint32_t set = UINT32_MAX, binding = UINT32_MAX;
+   nir_ssa_def *res_index =
+      build_res_index_for_chain(b, idx_intrin, addr_format,
+                                &set, &binding, state);
+
+   const struct anv_descriptor_set_binding_layout *bind_layout =
+      &state->layout->set[set].layout->binding[binding];
+
+   b->cursor = nir_before_instr(&load_desc->instr);
+
+   nir_ssa_def *desc_addr =
+      build_desc_addr(b, bind_layout, bind_layout->type,
+                      res_index, addr_format, state);
+
+   /* Acceleration structure descriptors are always uint64_t */
+   nir_ssa_def *desc = build_load_descriptor_mem(b, desc_addr, 0, 1, 64, state);
+
+   assert(load_desc->dest.is_ssa);
+   assert(load_desc->dest.ssa.bit_size == 64);
+   assert(load_desc->dest.ssa.num_components == 1);
+   nir_ssa_def_rewrite_uses(&load_desc->dest.ssa, desc);
+   nir_instr_remove(&load_desc->instr);
+
+   return true;
+}
+
 static bool
 lower_direct_buffer_instr(nir_builder *b, nir_instr *instr, void *_state)
 {
@@ -809,6 +851,12 @@ lower_direct_buffer_instr(nir_builder *b, nir_instr *instr, void *_state)
       return true;
    }
 
+   case nir_intrinsic_load_vulkan_descriptor:
+      if (nir_intrinsic_desc_type(intrin) ==
+          VK_DESCRIPTOR_TYPE_ACCELERATION_STRUCTURE_KHR)
+         return lower_load_accel_struct_desc(b, intrin, state);
+      return false;
+
    default:
       return false;
    }

src/intel/vulkan/anv_private.h

@@ -81,6 +81,7 @@ struct anv_batch;
 struct anv_buffer;
 struct anv_buffer_view;
 struct anv_image_view;
+struct anv_acceleration_structure;
 struct anv_instance;
 
 struct intel_aux_map_context;
@@ -2145,6 +2146,14 @@ anv_descriptor_set_write_buffer(struct anv_device *device,
                                 uint32_t element,
                                 VkDeviceSize offset,
                                 VkDeviceSize range);
+
+void
+anv_descriptor_set_write_acceleration_structure(struct anv_device *device,
+                                                struct anv_descriptor_set *set,
+                                                struct anv_acceleration_structure *accel,
+                                                uint32_t binding,
+                                                uint32_t element);
+
 void
 anv_descriptor_set_write_inline_uniform_data(struct anv_device *device,
                                              struct anv_descriptor_set *set,