Skip to content

Commit

Permalink
Merge pull request #2392 from billhollings/fix-ds-alloc-in-mtlargbuff
Browse files Browse the repository at this point in the history 
Fixes to managing descriptor set allocation in a Metal argument buffer.
  • Loading branch information
@billhollings
billhollings authored Nov 7, 2024
2 parents 758c4f2 + bdc646c commit 26c7b02
Show file tree
Hide file tree
Showing 3 changed files with 57 additions and 36 deletions.
12 changes: 7 additions & 5 deletions MoltenVK/MoltenVK/GPUObjects/MVKDescriptorSet.h
View file
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
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
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

0 comments on commit 26c7b02

Please sign in to comment.