Skip to content

Commit c822105

Browse files
MarijnS95claude
MarijnS95
and
claude
committed
Fold the AS pointer into ResourceSet / ResourceRef
The AS-binding commit immediately above had each backend cross-reference IS.AccelStructs by index arithmetic (R.TLASPtr - &P.AccelStructs.TLAS[0] plus the BLAS-count offset) every time the binding loop hit an AS resource. Pull that lookup into createBuffers / createResource (where the per-resource record is built) and cache the resulting TLAS pointer on the per-array-element struct so the binding loop just reads Bundle[0].AS. DX and Metal add an `AS` field to `ResourceSet`; Vulkan adds an `AS` field to `ResourceRef`. The bundle stays a plain `SmallVector<ResourceSet>` typedef on DX and Metal and the existing struct on Vulkan — the AS pointer co-locates with its sibling upload/buffer/image fields at the right granularity (per array element), which is what the existing tagged-union shape of these structs already does for buffer-vs-texture-vs-sampler. `IS.AccelStructs` still owns the AS objects (BLASes need to live somewhere too, and lifetime stays simple); the per-element record only carries a non-owning pointer. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
1 parent a1a661f commit c822105

3 files changed

Lines changed: 58 additions & 51 deletions

File tree

‎lib/API/DX/Device.cpp‎

Lines changed: 19 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -976,27 +976,31 @@ class DXDevice : public offloadtest::Device {
976976
ComPtr<ID3D12Resource> Buffer;
977977
std::unique_ptr<offloadtest::Buffer> Readback;
978978
ComPtr<ID3D12Heap> Heap;
979+
// Populated for AS resources; mutually exclusive with the
980+
// Upload/Buffer/Readback/Heap fields above.
981+
DXAccelerationStructure *AS = nullptr;
979982
ResourceSet(ComPtr<ID3D12Resource> Upload, ComPtr<ID3D12Resource> Buffer,
980983
std::unique_ptr<offloadtest::Buffer> Readback,
981984
ComPtr<ID3D12Heap> Heap = nullptr)
982985
: Upload(Upload), Buffer(Buffer), Readback(std::move(Readback)),
983986
Heap(Heap) {}
987+
explicit ResourceSet(DXAccelerationStructure *AS) : AS(AS) {}
984988
ResourceSet(const ResourceSet &) = delete;
985989
ResourceSet(ResourceSet &&A)
986990
: Upload(A.Upload), Buffer(A.Buffer), Readback(std::move(A.Readback)),
987-
Heap(A.Heap) {}
991+
Heap(A.Heap), AS(A.AS) {}
988992
ResourceSet &operator=(const ResourceSet &) = delete;
989993
ResourceSet &operator=(ResourceSet &&A) {
990994
Upload = A.Upload;
991995
Buffer = A.Buffer;
992996
Readback = std::move(A.Readback);
993997
Heap = A.Heap;
998+
AS = A.AS;
994999
return *this;
9951000
}
9961001
};
9971002

998-
// ResourceBundle will contain one ResourceSet for a singular resource
999-
// or multiple ResourceSets for resource array.
1003+
// ResourceBundle holds one ResourceSet per array element (singular = 1).
10001004
using ResourceBundle = llvm::SmallVector<ResourceSet>;
10011005
using ResourcePair = std::pair<offloadtest::Resource *, ResourceBundle>;
10021006

@@ -2166,13 +2170,20 @@ class DXDevice : public offloadtest::Device {
21662170

21672171
llvm::Error createBuffers(Pipeline &P, InvocationState &IS) {
21682172
auto CreateBuffer =
2169-
[&IS,
2173+
[&P, &IS,
21702174
this](Resource &R,
21712175
llvm::SmallVectorImpl<ResourcePair> &Resources) -> llvm::Error {
2172-
// Acceleration structures are created separately; descriptor binding for
2173-
// AS resources is wired up in the per-table binding loop below.
2176+
// OutAS layout from buildPipelineAccelerationStructures: BLASes
2177+
// first, then TLASes — both in P.AccelStructs declaration order.
21742178
if (R.isAccelerationStructure()) {
2175-
Resources.push_back(std::make_pair(&R, ResourceBundle{}));
2179+
assert(R.TLASPtr && "AS resource must be resolved to a TLAS");
2180+
assert(R.getArraySize() == 1 && "AS arrays not yet supported");
2181+
const size_t TLASIdx = R.TLASPtr - &P.AccelStructs.TLAS[0];
2182+
const size_t ASIdx = P.AccelStructs.BLAS.size() + TLASIdx;
2183+
ResourceBundle Bundle;
2184+
Bundle.emplace_back(
2185+
llvm::cast<DXAccelerationStructure>(IS.AccelStructs[ASIdx].get()));
2186+
Resources.push_back(std::make_pair(&R, std::move(Bundle)));
21762187
return llvm::Error::success();
21772188
}
21782189
switch (getDescriptorKind(R.Kind)) {
@@ -2219,15 +2230,7 @@ class DXDevice : public offloadtest::Device {
22192230
D3D12_DESCRIPTOR_HEAP_TYPE_CBV_SRV_UAV);
22202231
for (auto &T : IS.DescTables) {
22212232
for (auto &R : T.Resources) {
2222-
if (R.first->isAccelerationStructure()) {
2223-
assert(R.first->TLASPtr && "AS resource must be resolved to a TLAS");
2224-
assert(R.first->getArraySize() == 1 && "AS arrays not yet supported");
2225-
// OutAS layout from buildPipelineAccelerationStructures: BLASes
2226-
// first, then TLASes — both in P.AccelStructs declaration order.
2227-
const size_t TLASIdx = R.first->TLASPtr - &P.AccelStructs.TLAS[0];
2228-
const size_t ASIdx = P.AccelStructs.BLAS.size() + TLASIdx;
2229-
auto *DXAS =
2230-
llvm::cast<DXAccelerationStructure>(IS.AccelStructs[ASIdx].get());
2233+
if (DXAccelerationStructure *DXAS = R.second[0].AS) {
22312234
D3D12_SHADER_RESOURCE_VIEW_DESC SRVDesc = {};
22322235
SRVDesc.Format = DXGI_FORMAT_UNKNOWN;
22332236
SRVDesc.ViewDimension =

‎lib/API/MTL/MTLDevice.cpp‎

Lines changed: 16 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -931,11 +931,13 @@ class MTLDevice : public offloadtest::Device {
931931

932932
struct ResourceSet {
933933
MTLPtr<MTL::Resource> Resource;
934+
// Populated for AS resources; mutually exclusive with Resource above.
935+
MetalAccelerationStructure *AS = nullptr;
934936
ResourceSet(MTL::Resource *Resource) : Resource(Resource) {}
937+
explicit ResourceSet(MetalAccelerationStructure *AS) : AS(AS) {}
935938
};
936939

937-
// ResourceBundle will contain one ResourceSet for a singular resource
938-
// or multiple ResourceSets for resource array.
940+
// ResourceBundle holds one ResourceSet per array element (singular = 1).
939941
using ResourceBundle = llvm::SmallVector<ResourceSet>;
940942
using ResourcePair = std::pair<offloadtest::Resource *, ResourceBundle>;
941943

@@ -1305,13 +1307,20 @@ class MTLDevice : public offloadtest::Device {
13051307

13061308
llvm::Error createBuffers(Pipeline &P, InvocationState &IS) {
13071309
auto CreateBuffer =
1308-
[&IS,
1310+
[&P, &IS,
13091311
this](Resource &R,
13101312
llvm::SmallVectorImpl<ResourcePair> &Resources) -> llvm::Error {
1311-
// Acceleration structures are created separately; descriptor binding for
1312-
// AS resources is wired up in the per-table binding loop below.
1313+
// OutAS layout from buildPipelineAccelerationStructures: BLASes
1314+
// first, then TLASes — both in P.AccelStructs declaration order.
13131315
if (R.isAccelerationStructure()) {
1314-
Resources.emplace_back(&R, ResourceBundle{});
1316+
assert(R.TLASPtr && "AS resource must be resolved to a TLAS");
1317+
assert(R.getArraySize() == 1 && "AS arrays not yet supported");
1318+
const size_t TLASIdx = R.TLASPtr - &P.AccelStructs.TLAS[0];
1319+
const size_t ASIdx = P.AccelStructs.BLAS.size() + TLASIdx;
1320+
ResourceBundle Bundle;
1321+
Bundle.emplace_back(llvm::cast<MetalAccelerationStructure>(
1322+
IS.AccelStructs[ASIdx].get()));
1323+
Resources.emplace_back(&R, std::move(Bundle));
13151324
return llvm::Error::success();
13161325
}
13171326
switch (getDescriptorKind(R.Kind)) {
@@ -1356,16 +1365,7 @@ class MTLDevice : public offloadtest::Device {
13561365
uint32_t HeapIndex = 0;
13571366
for (auto &T : IS.DescTables) {
13581367
for (auto &R : T.Resources) {
1359-
if (R.first->isAccelerationStructure()) {
1360-
assert(R.first->TLASPtr && "AS resource must be resolved to a TLAS");
1361-
assert(R.first->getArraySize() == 1 && "AS arrays not yet supported");
1362-
// OutAS layout from buildPipelineAccelerationStructures:
1363-
// BLASes first, then TLASes — both in declaration order.
1364-
const size_t TLASIdx = R.first->TLASPtr - &P.AccelStructs.TLAS[0];
1365-
const size_t ASIdx = P.AccelStructs.BLAS.size() + TLASIdx;
1366-
auto *MTLAS = llvm::cast<MetalAccelerationStructure>(
1367-
IS.AccelStructs[ASIdx].get());
1368-
1368+
if (MetalAccelerationStructure *MTLAS = R.second[0].AS) {
13691369
// Metal shader converter doesn't bind the AS directly; the shader
13701370
// reads a buffer containing an IRRaytracingAccelerationStructureGPU-
13711371
// Header that holds the AS's gpuResourceID plus a pointer to an

‎lib/API/VK/Device.cpp‎

Lines changed: 23 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -1203,10 +1203,14 @@ class VulkanDevice : public offloadtest::Device {
12031203
struct ResourceRef {
12041204
ResourceRef(BufferRef H, BufferRef D) : Host(H), Device(D) {}
12051205
ResourceRef(BufferRef H, ImageRef I) : Host(H), Image(I) {}
1206-
1207-
BufferRef Host;
1208-
BufferRef Device;
1209-
ImageRef Image;
1206+
explicit ResourceRef(VulkanAccelerationStructure *AS) : AS(AS) {}
1207+
1208+
BufferRef Host{};
1209+
BufferRef Device{};
1210+
ImageRef Image{};
1211+
// Populated for AS resources; mutually exclusive with the buffer/image
1212+
// fields above.
1213+
VulkanAccelerationStructure *AS = nullptr;
12101214
};
12111215

12121216
struct ResourceBundle {
@@ -3029,12 +3033,19 @@ class VulkanDevice : public offloadtest::Device {
30293033
return ResourceRef(Host, ImageRef{0, Sampler, 0});
30303034
}
30313035

3032-
llvm::Error createResource(Resource &R, InvocationState &IS) {
3033-
// Acceleration structures are created separately; push a placeholder so
3034-
// the resource index stays in sync with the descriptor set layout.
3036+
llvm::Error createResource(Pipeline &P, Resource &R, InvocationState &IS) {
3037+
// OutAS layout from buildPipelineAccelerationStructures: BLASes first,
3038+
// then TLASes — both in P.AccelStructs declaration order.
30353039
if (R.isAccelerationStructure()) {
3036-
const ResourceBundle Bundle{getDescriptorType(R.Kind), 0, nullptr};
3037-
IS.Resources.push_back(Bundle);
3040+
assert(R.TLASPtr && "AS resource must be resolved to a TLAS");
3041+
assert(R.getArraySize() == 1 && "AS arrays not yet supported");
3042+
const size_t TLASIdx = R.TLASPtr - &P.AccelStructs.TLAS[0];
3043+
const size_t ASIdx = P.AccelStructs.BLAS.size() + TLASIdx;
3044+
auto *VkAS =
3045+
llvm::cast<VulkanAccelerationStructure>(IS.AccelStructs[ASIdx].get());
3046+
ResourceBundle Bundle{getDescriptorType(R.Kind), 0, nullptr};
3047+
Bundle.ResourceRefs.push_back(ResourceRef{VkAS});
3048+
IS.Resources.push_back(std::move(Bundle));
30383049
return llvm::Error::success();
30393050
}
30403051

@@ -3153,7 +3164,7 @@ class VulkanDevice : public offloadtest::Device {
31533164
llvm::Error createResources(Pipeline &P, InvocationState &IS) {
31543165
for (auto &D : P.Sets) {
31553166
for (auto &R : D.Resources) {
3156-
if (auto Err = createResource(R, IS))
3167+
if (auto Err = createResource(P, R, IS))
31573168
return Err;
31583169
}
31593170
}
@@ -3322,15 +3333,8 @@ class VulkanDevice : public offloadtest::Device {
33223333
for (uint32_t RIdx = 0; RIdx < P.Sets[SetIdx].Resources.size();
33233334
++RIdx, ++OverallResIdx) {
33243335
const Resource &R = P.Sets[SetIdx].Resources[RIdx];
3325-
if (R.isAccelerationStructure()) {
3326-
assert(R.TLASPtr && "AS resource must be resolved to a TLAS");
3327-
assert(R.getArraySize() == 1 && "AS arrays not yet supported");
3328-
// OutAS layout from buildPipelineAccelerationStructures: BLASes
3329-
// first, then TLASes — both in P.AccelStructs declaration order.
3330-
const size_t TLASIdx = R.TLASPtr - &P.AccelStructs.TLAS[0];
3331-
const size_t ASIdx = P.AccelStructs.BLAS.size() + TLASIdx;
3332-
auto *VkAS = llvm::cast<VulkanAccelerationStructure>(
3333-
IS.AccelStructs[ASIdx].get());
3336+
if (VulkanAccelerationStructure *VkAS =
3337+
IS.Resources[OverallResIdx].ResourceRefs[0].AS) {
33343338
const size_t HandleStart = ASHandles.size();
33353339
ASHandles.push_back(VkAS->AccelStruct);
33363340
VkWriteDescriptorSetAccelerationStructureKHR ASWrite = {};

0 commit comments

Comments
 (0)