Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fixes to managing descriptor set allocation in a Metal argument buffer. #2392

Merged
merged 1 commit into from
Nov 7, 2024
Merged

Fixes to managing descriptor set allocation in a Metal argument buffer. #2392

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 7 additions & 5 deletions MoltenVK/MoltenVK/GPUObjects/MVKDescriptorSet.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,13 +44,15 @@ typedef struct MVKMetalArgumentBuffer {
void setSamplerState(id<MTLSamplerState> mtlSamp, uint32_t index);
id<MTLBuffer> getMetalArgumentBuffer() { return _mtlArgumentBuffer; }
NSUInteger getMetalArgumentBufferOffset() { return _mtlArgumentBufferOffset; }
void setArgumentBuffer(id<MTLBuffer> mtlArgBuff, NSUInteger mtlArgBuffOfst, id<MTLArgumentEncoder> mtlArgEnc);
NSUInteger getMetalArgumentBufferEncodedSize() { return _mtlArgumentBufferEncodedSize; }
void setArgumentBuffer(id<MTLBuffer> mtlArgBuff, NSUInteger mtlArgBuffOfst, NSUInteger mtlArgBuffEncSize, id<MTLArgumentEncoder> mtlArgEnc);
~MVKMetalArgumentBuffer();
protected:
void* getArgumentPointer(uint32_t index) const;
id<MTLArgumentEncoder> _mtlArgumentEncoder = nil;
id<MTLBuffer> _mtlArgumentBuffer = nil;
NSUInteger _mtlArgumentBufferOffset = 0;
NSUInteger _mtlArgumentBufferEncodedSize = 0;
} MVKMetalArgumentBuffer;


Expand Down Expand Up @@ -135,7 +137,7 @@ class MVKDescriptorSetLayout : public MVKVulkanAPIDeviceObject {
MVKDescriptorSetLayoutBinding* getBinding(uint32_t binding, uint32_t bindingIndexOffset = 0);
uint32_t getBufferSizeBufferArgBuferIndex() { return 0; }
id <MTLArgumentEncoder> getMTLArgumentEncoder(uint32_t variableDescriptorCount);
uint64_t getMetal3ArgumentBufferEncodedLength(uint32_t variableDescriptorCount);
size_t getMetal3ArgumentBufferEncodedLength(uint32_t variableDescriptorCount);
bool checkCanUseArgumentBuffers(const VkDescriptorSetLayoutCreateInfo* pCreateInfo);

MVKSmallVector<MVKDescriptorSetLayoutBinding> _bindings;
Expand Down Expand Up @@ -216,7 +218,8 @@ class MVKDescriptorSet : public MVKVulkanAPIDeviceObject {
MVKDescriptor* getDescriptor(uint32_t binding, uint32_t elementIndex = 0);
VkResult allocate(MVKDescriptorSetLayout* layout,
uint32_t variableDescriptorCount,
NSUInteger mtlArgBufferOffset,
NSUInteger mtlArgBuffOffset,
NSUInteger mtlArgBuffEncSize,
id<MTLArgumentEncoder> mtlArgEnc);
void free(bool isPoolReset);
MVKMTLBufferAllocation* acquireMTLBufferRegion(NSUInteger length);
Expand Down Expand Up @@ -309,7 +312,6 @@ class MVKDescriptorPool : public MVKVulkanAPIDeviceObject {
MVKBitArray _descriptorSetAvailablility;
MVKMTLBufferAllocator _mtlBufferAllocator;
id<MTLBuffer> _metalArgumentBuffer = nil;
NSUInteger _nextMetalArgumentBufferOffset = 0;

MVKDescriptorTypePool<MVKUniformBufferDescriptor> _uniformBufferDescriptors;
MVKDescriptorTypePool<MVKStorageBufferDescriptor> _storageBufferDescriptors;
Expand All @@ -325,7 +327,7 @@ class MVKDescriptorPool : public MVKVulkanAPIDeviceObject {
MVKDescriptorTypePool<MVKStorageTexelBufferDescriptor> _storageTexelBufferDescriptors;

VkDescriptorPoolCreateFlags _flags = 0;
size_t _maxAllocDescSetCount = 0;
size_t _allocatedDescSetCount = 0;
};


Expand Down
59 changes: 31 additions & 28 deletions MoltenVK/MoltenVK/GPUObjects/MVKDescriptorSet.mm
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,9 +34,11 @@

void MVKMetalArgumentBuffer::setArgumentBuffer(id<MTLBuffer> mtlArgBuff,
NSUInteger mtlArgBuffOfst,
NSUInteger mtlArgBuffEncSize,
id<MTLArgumentEncoder> mtlArgEnc) {
_mtlArgumentBuffer = mtlArgBuff;
_mtlArgumentBufferOffset = mtlArgBuffOfst;
_mtlArgumentBufferEncodedSize = mtlArgBuffEncSize;

auto* oldArgEnc = _mtlArgumentEncoder;
_mtlArgumentEncoder = [mtlArgEnc retain]; // retained
Expand Down Expand Up @@ -329,8 +331,8 @@ static void populateAuxBuffer(mvk::SPIRVToMSLConversionConfiguration& shaderConf
}

// Returns the encoded byte length of the resources from a descriptor set in an argument buffer.
uint64_t MVKDescriptorSetLayout::getMetal3ArgumentBufferEncodedLength(uint32_t variableDescriptorCount) {
uint64_t encodedLen = 0;
size_t MVKDescriptorSetLayout::getMetal3ArgumentBufferEncodedLength(uint32_t variableDescriptorCount) {
size_t encodedLen = 0;

// Buffer sizes buffer at front
if (needsBufferSizeAuxBuffer()) {
Expand Down Expand Up @@ -536,12 +538,13 @@ static void populateAuxBuffer(mvk::SPIRVToMSLConversionConfiguration& shaderConf

VkResult MVKDescriptorSet::allocate(MVKDescriptorSetLayout* layout,
uint32_t variableDescriptorCount,
NSUInteger mtlArgBufferOffset,
NSUInteger mtlArgBuffOffset,
NSUInteger mtlArgBuffEncSize,
id<MTLArgumentEncoder> mtlArgEnc) {
_layout = layout;
_layout->retain();
_variableDescriptorCount = variableDescriptorCount;
_argumentBuffer.setArgumentBuffer(_pool->_metalArgumentBuffer, mtlArgBufferOffset, mtlArgEnc);
_argumentBuffer.setArgumentBuffer(_pool->_metalArgumentBuffer, mtlArgBuffOffset, mtlArgBuffEncSize, mtlArgEnc);

uint32_t descCnt = layout->getDescriptorCount(variableDescriptorCount);
_descriptors.reserve(descCnt);
Expand Down Expand Up @@ -581,7 +584,7 @@ static void populateAuxBuffer(mvk::SPIRVToMSLConversionConfiguration& shaderConf
_dynamicOffsetDescriptorCount = 0;
_variableDescriptorCount = 0;

if (isPoolReset) { _argumentBuffer.setArgumentBuffer(_pool->_metalArgumentBuffer, 0, nil); }
if (isPoolReset) { _argumentBuffer.setArgumentBuffer(_pool->_metalArgumentBuffer, 0, 0, nil); }

// If this is a pool reset, and all desciptors are from the pool, we don't need to free them.
if ( !(isPoolReset && _allDescriptorsAreFromPool) ) {
Expand Down Expand Up @@ -734,53 +737,54 @@ static void populateAuxBuffer(mvk::SPIRVToMSLConversionConfiguration& shaderConf
uint32_t variableDescriptorCount,
VkDescriptorSet* pVKDS) {
VkResult rslt = VK_ERROR_FRAGMENTED_POOL;
uint64_t mtlArgBuffEncSize = 0;
size_t mtlArgBuffEncSize = 0;
id<MTLArgumentEncoder> mtlArgEnc = nil;
if (mvkDSL->isUsingMetalArgumentBuffers()) {
bool isUsingMetalArgBuff = mvkDSL->isUsingMetalArgumentBuffers();

if (isUsingMetalArgBuff) {
if (needsMetalArgumentBufferEncoders()) {
mtlArgEnc = mvkDSL->getMTLArgumentEncoder(variableDescriptorCount);
mtlArgBuffEncSize = mtlArgEnc.encodedLength;
} else {
mtlArgBuffEncSize = mvkDSL->getMetal3ArgumentBufferEncodedLength(variableDescriptorCount);
}
}
uint64_t mtlArgBuffEncAlignedSize = mvkAlignByteCount(mtlArgBuffEncSize, getMetalFeatures().mtlBufferAlignment);

size_t dsCnt = _descriptorSetAvailablility.size();
_descriptorSetAvailablility.enumerateEnabledBits([&](size_t dsIdx) {
bool isSpaceAvail = true; // If not using Metal arg buffers, space will always be available.
MVKDescriptorSet* mvkDS = &_descriptorSets[dsIdx];
NSUInteger mtlArgBuffOffset = mvkDS->getMetalArgumentBuffer().getMetalArgumentBufferOffset();

// If the desc set is using a Metal argument buffer, we also need to see if the desc set
// will fit in the slot that might already have been allocated for it in the Metal argument
// buffer from a previous allocation that was returned. If this pool has been reset recently,
// then the desc sets will not have had a Metal argument buffer allocation assigned yet.
if (mtlArgBuffEncSize) {

// If the offset has not been set (and it's not the first desc set except
// on a reset pool), set the offset and update the next available offset value.
if ( !mtlArgBuffOffset && (dsIdx || !_nextMetalArgumentBufferOffset)) {
mtlArgBuffOffset = _nextMetalArgumentBufferOffset;
_nextMetalArgumentBufferOffset += mtlArgBuffEncAlignedSize;
NSUInteger mtlArgBuffOffset = 0;

// If the desc set is using a Metal argument buffer, we must check if the desc set will fit in the slot
// in the Metal argument buffer, if that slot was previously allocated for a returned descriptor set.
if (isUsingMetalArgBuff) {
mtlArgBuffOffset = mvkDS->getMetalArgumentBuffer().getMetalArgumentBufferOffset();

// If the offset has not been set, and this is not the first desc set,
// set the offset to align with the end of the previous desc set.
if ( !mtlArgBuffOffset && dsIdx ) {
auto& prevArgBuff = _descriptorSets[dsIdx - 1].getMetalArgumentBuffer();
mtlArgBuffOffset = (prevArgBuff.getMetalArgumentBufferOffset() +
mvkAlignByteCount(prevArgBuff.getMetalArgumentBufferEncodedSize(),
getMetalFeatures().mtlBufferAlignment));
}

// Get the offset of the next desc set, if one exists and
// its offset has been set, or the end of the arg buffer.
size_t nextDSIdx = dsIdx + 1;
NSUInteger nextOffset = (nextDSIdx < dsCnt ? _descriptorSets[nextDSIdx].getMetalArgumentBuffer().getMetalArgumentBufferOffset() : 0);
NSUInteger nextOffset = (nextDSIdx < _allocatedDescSetCount ? _descriptorSets[nextDSIdx].getMetalArgumentBuffer().getMetalArgumentBufferOffset() : 0);
if ( !nextOffset ) { nextOffset = _metalArgumentBuffer.length; }

isSpaceAvail = (mtlArgBuffOffset + mtlArgBuffEncSize) <= nextOffset;
}

if (isSpaceAvail) {
rslt = mvkDS->allocate(mvkDSL, variableDescriptorCount, mtlArgBuffOffset, mtlArgEnc);
rslt = mvkDS->allocate(mvkDSL, variableDescriptorCount, mtlArgBuffOffset, mtlArgBuffEncSize, mtlArgEnc);
if (rslt) {
freeDescriptorSet(mvkDS, false);
} else {
_descriptorSetAvailablility.disableBit(dsIdx);
_maxAllocDescSetCount = std::max(_maxAllocDescSetCount, dsIdx + 1);
_allocatedDescSetCount = std::max(_allocatedDescSetCount, dsIdx + 1);
*pVKDS = (VkDescriptorSet)mvkDS;
}
return false;
Expand Down Expand Up @@ -818,7 +822,7 @@ static void populateAuxBuffer(mvk::SPIRVToMSLConversionConfiguration& shaderConf
// Free allocated descriptor sets and reset descriptor pools.
// Don't waste time freeing desc sets that were never allocated.
VkResult MVKDescriptorPool::reset(VkDescriptorPoolResetFlags flags) {
for (uint32_t dsIdx = 0; dsIdx < _maxAllocDescSetCount; dsIdx++) {
for (uint32_t dsIdx = 0; dsIdx < _allocatedDescSetCount; dsIdx++) {
freeDescriptorSet(&_descriptorSets[dsIdx], true);
}
_descriptorSetAvailablility.enableAllBits();
Expand All @@ -836,8 +840,7 @@ static void populateAuxBuffer(mvk::SPIRVToMSLConversionConfiguration& shaderConf
_uniformTexelBufferDescriptors.reset();
_storageTexelBufferDescriptors.reset();

_nextMetalArgumentBufferOffset = 0;
_maxAllocDescSetCount = 0;
_allocatedDescSetCount = 0;

return VK_SUCCESS;
}
Expand Down
22 changes: 19 additions & 3 deletions MoltenVK/MoltenVK/Utility/MVKBitArray.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -99,15 +99,31 @@ class MVKBitArray {
_fullyDisabledSectionCount = (uint32_t)secCnt;
}

/** Returns the index of the first bit that is enabled, at or after the specified index. */
/**
* Returns the index of the first enabled bit, at or after the specified index.
* If no bits are enabled, returns the size() of this bit array.
*/
size_t getIndexOfFirstEnabledBit(size_t startIndex = 0) {
size_t secCnt = getSectionCount();
size_t secIdx = getIndexOfSection(startIndex);

// Optimize by skipping all consecutive sections at the beginning that are known to have no enabled bits.
if (secIdx < _fullyDisabledSectionCount) {
secIdx = _fullyDisabledSectionCount;
startIndex = 0;
}
if (secIdx >= getSectionCount()) { return _bitCount; }
return std::min((secIdx * SectionBitCount) + getIndexOfFirstEnabledBitInSection(getSection(secIdx), getBitIndexInSection(startIndex)), _bitCount);

// Search all sections at or after the starting index, and if an enabled bit is found, return the index of it.
while (secIdx < secCnt) {
size_t lclBitIdx = getIndexOfFirstEnabledBitInSection(getSection(secIdx), getBitIndexInSection(startIndex));
if (lclBitIdx < SectionBitCount) {
return (secIdx * SectionBitCount) + lclBitIdx;
}
startIndex = 0;
secIdx++;
}

return _bitCount;
}

/**
Expand Down
Loading