Fix automatic usage of D3D12_HEAP_FLAG_ALLOW_SHADER_ATOMICS; allow additional heapFlags in Allocator::AllocateMemory
diff --git a/src/D3D12MemAlloc.cpp b/src/D3D12MemAlloc.cpp
index aa4e974..4eef684 100644
--- a/src/D3D12MemAlloc.cpp
+++ b/src/D3D12MemAlloc.cpp
@@ -86,13 +86,15 @@
    #define D3D12MA_DEFAULT_BLOCK_SIZE (256ull * 1024 * 1024)

 #endif

 

-#ifndef D3D12MA_EXTRA_DEFAULT_TYPE_HEAP_FLAGS

+#ifndef D3D12MA_ALLOW_SHADER_ATOMICS

     /*

-    Here you can control additional heap flags added to heaps/resources created in DEFAULT heap type.

-    It's mostly for automatic usage of the cryptic, undocumented flag D3D12_HEAP_FLAG_ALLOW_SHADER_ATOMICS.

-    Its absence doesn't seem to change anything but better to use it always, just in case.

+    Set this to 1 to make the library automatically adding flag

+    D3D12_HEAP_FLAG_ALLOW_SHADER_ATOMICS to any allocated heap that might contain a

+    buffer or a texture used as a UAV. You can change it to 0 if you are sure your

+    program doesn't use atomic functions in its shaders. In practice this flag

+    doesn't seem to change anything.

     */

-    #define D3D12MA_EXTRA_DEFAULT_TYPE_HEAP_FLAGS (D3D12_HEAP_FLAG_ALLOW_SHADER_ATOMICS)

+    #define D3D12MA_ALLOW_SHADER_ATOMICS (1)

 #endif

 

 ////////////////////////////////////////////////////////////////////////////////

@@ -3686,6 +3688,7 @@
     D3D12MA_ASSERT(resAllocInfo.SizeInBytes > 0);

 

     const UINT defaultPoolIndex = CalcDefaultPoolIndex(*pAllocDesc, resourceDesc2);

+    D3D12MA_ASSERT(defaultPoolIndex != UINT32_MAX);

     BlockVector* blockVector = m_BlockVectors[defaultPoolIndex];

     D3D12MA_ASSERT(blockVector);

 

@@ -3777,15 +3780,17 @@
     ALLOCATION_DESC finalAllocDesc = *pAllocDesc;

 

     const UINT defaultPoolIndex = CalcDefaultPoolIndex(*pAllocDesc, heapFlags);

-    if(defaultPoolIndex == UINT32_MAX)

+    bool requireCommittedMemory = (defaultPoolIndex == UINT32_MAX);

+    if(requireCommittedMemory)

     {

-        return E_INVALIDARG;

+        return AllocateHeap(&finalAllocDesc, heapFlags, *pAllocInfo, ppAllocation);

     }

+

     BlockVector* blockVector = m_BlockVectors[defaultPoolIndex];

     D3D12MA_ASSERT(blockVector);

 

     const UINT64 preferredBlockSize = blockVector->GetPreferredBlockSize();

-    bool preferCommittedMemory =

+    const bool preferCommittedMemory =

         m_AlwaysCommitted ||

         // Heuristics: Allocate committed memory if requested size if greater than half of preferred block size.

         pAllocInfo->SizeInBytes > preferredBlockSize / 2;

@@ -3850,8 +3855,14 @@
     D3D12_HEAP_PROPERTIES heapProps = {};

     heapProps.Type = pAllocDesc->HeapType;

 

-    const D3D12_HEAP_FLAGS heapFlags = pAllocDesc->HeapType == D3D12_HEAP_TYPE_DEFAULT ?

-        D3D12MA_EXTRA_DEFAULT_TYPE_HEAP_FLAGS : D3D12_HEAP_FLAG_NONE;

+    D3D12_HEAP_FLAGS heapFlags = D3D12_HEAP_FLAG_NONE;

+#if D3D12MA_ALLOW_SHADER_ATOMICS

+    if((pResourceDesc->Flags & D3D12_RESOURCE_FLAG_ALLOW_UNORDERED_ACCESS) != 0)

+    {

+        D3D12MA_ASSERT(pAllocDesc->HeapType == D3D12_HEAP_TYPE_DEFAULT);

+        heapFlags = D3D12_HEAP_FLAG_ALLOW_SHADER_ATOMICS;

+    }

+#endif

 

     ID3D12Resource* res = NULL;

     HRESULT hr = m_Device->CreateCommittedResource(

@@ -3902,10 +3913,12 @@
         }

     }

 

+#if D3D12MA_ALLOW_SHADER_ATOMICS

     if(pAllocDesc->HeapType == D3D12_HEAP_TYPE_DEFAULT)

     {

-        heapFlags |= D3D12MA_EXTRA_DEFAULT_TYPE_HEAP_FLAGS;

+        heapFlags |= D3D12_HEAP_FLAG_ALLOW_SHADER_ATOMICS;

     }

+#endif

 

     D3D12_HEAP_DESC heapDesc = {};

     heapDesc.SizeInBytes = allocInfo.SizeInBytes;

@@ -3971,6 +3984,17 @@
 

 UINT AllocatorPimpl::CalcDefaultPoolIndex(const ALLOCATION_DESC& allocDesc, const D3D12_HEAP_FLAGS heapFlags) const

 {

+    D3D12_HEAP_FLAGS heapFlagsToIgnore =

+        D3D12_HEAP_FLAG_DENY_BUFFERS | D3D12_HEAP_FLAG_DENY_RT_DS_TEXTURES | D3D12_HEAP_FLAG_DENY_NON_RT_DS_TEXTURES;

+#if D3D12MA_ALLOW_SHADER_ATOMICS

+    heapFlagsToIgnore |= D3D12_HEAP_FLAG_ALLOW_SHADER_ATOMICS;

+#endif

+    const D3D12_HEAP_FLAGS additionalHeapFlags = heapFlags & ~heapFlagsToIgnore;

+    if(additionalHeapFlags != 0)

+    {

+        return UINT32_MAX;

+    }

+

     UINT poolIndex = UINT_MAX;

     switch(allocDesc.HeapType)

     {

@@ -4045,10 +4069,12 @@
         D3D12MA_ASSERT(0);

     }

 

+#if D3D12MA_ALLOW_SHADER_ATOMICS

     if(outHeapType == D3D12_HEAP_TYPE_DEFAULT)

     {

-        outHeapFlags |= D3D12MA_EXTRA_DEFAULT_TYPE_HEAP_FLAGS;

+        outHeapFlags |= D3D12_HEAP_FLAG_ALLOW_SHADER_ATOMICS;

     }

+#endif

 }

 

 void AllocatorPimpl::RegisterCommittedAllocation(Allocation* alloc, D3D12_HEAP_TYPE heapType)

@@ -4750,8 +4776,7 @@
             pAllocInfo->Alignment == D3D12_DEFAULT_RESOURCE_PLACEMENT_ALIGNMENT ||

             pAllocInfo->Alignment == D3D12_DEFAULT_MSAA_RESOURCE_PLACEMENT_ALIGNMENT) ||

         pAllocInfo->SizeInBytes == 0 ||

-        pAllocInfo->SizeInBytes % (64ull * 1024) != 0 ||

-        (heapFlags & ~(D3D12_HEAP_FLAG_DENY_BUFFERS | D3D12_HEAP_FLAG_DENY_RT_DS_TEXTURES | D3D12_HEAP_FLAG_DENY_NON_RT_DS_TEXTURES)) != 0)

+        pAllocInfo->SizeInBytes % (64ull * 1024) != 0)

     {

         D3D12MA_ASSERT(0 && "Invalid arguments passed to Allocator::AllocateMemory.");

         return E_INVALIDARG;

diff --git a/src/D3D12MemAlloc.h b/src/D3D12MemAlloc.h
index 55932b2..df1c0d1 100644
--- a/src/D3D12MemAlloc.h
+++ b/src/D3D12MemAlloc.h
@@ -777,12 +777,13 @@
     This function is similar to `ID3D12Device::CreateHeap`, but it may really assign

     part of a larger, existing heap to the allocation.

 

-    If ResourceHeapTier = 1, `heapFlags` must be one of these values, depending on type

+    If ResourceHeapTier = 1, `heapFlags` must contain one of these values, depending on type

     of resources you are going to create in this memory:

     `D3D12_HEAP_FLAG_ALLOW_ONLY_BUFFERS`,

     `D3D12_HEAP_FLAG_ALLOW_ONLY_NON_RT_DS_TEXTURES`,

     `D3D12_HEAP_FLAG_ALLOW_ONLY_RT_DS_TEXTURES`.

-    If ResourceHeapTier = 2, `heapFlags` may be `D3D12_HEAP_FLAG_ALLOW_ALL_BUFFERS_AND_TEXTURES`.

+    If ResourceHeapTier = 2, `heapFlags` may be `D3D12_HEAP_FLAG_ALLOW_ALL_BUFFERS_AND_TEXTURES` = 0.

+    Additional flags in `heapFlags` are allowed as well.

 

     `pAllocInfo->SizeInBytes` must be multiply of 64KB.

     `pAllocInfo->Alignment` must be one of the legal values as described in documentation of `D3D12_HEAP_DESC`.

diff --git a/src/Tests.cpp b/src/Tests.cpp
index b758b4c..817aea0 100644
--- a/src/Tests.cpp
+++ b/src/Tests.cpp
@@ -155,7 +155,7 @@
     ResourceWithAllocation resources[count];

 

     D3D12MA::ALLOCATION_DESC allocDesc = {};

-    allocDesc.HeapType = D3D12_HEAP_TYPE_UPLOAD;

+    allocDesc.HeapType = D3D12_HEAP_TYPE_DEFAULT;

     allocDesc.Flags = D3D12MA::ALLOCATION_FLAG_COMMITTED;

 

     D3D12_RESOURCE_DESC resourceDesc;

@@ -169,7 +169,7 @@
         CHECK_HR( ctx.allocator->CreateResource(

             &allocDesc,

             &resourceDesc,

-            D3D12_RESOURCE_STATE_GENERIC_READ,

+            D3D12_RESOURCE_STATE_COPY_DEST,

             NULL,

             &alloc,

             __uuidof(ID3D12Resource),

@@ -213,6 +213,29 @@
     ctx.allocator->FreeStatsString(jsonString);

 }

 

+static void TestCustomHeapFlags(const TestContext& ctx)

+{

+    wprintf(L"Test custom heap flags\n");

+

+    D3D12MA::ALLOCATION_DESC allocDesc = {};

+    allocDesc.HeapType = D3D12_HEAP_TYPE_DEFAULT;

+

+    const D3D12_HEAP_FLAGS heapFlags = D3D12_HEAP_FLAG_ALLOW_ONLY_NON_RT_DS_TEXTURES |

+        D3D12_HEAP_FLAG_SHARED; // Extra flag.

+

+    D3D12_RESOURCE_ALLOCATION_INFO resAllocInfo = {};

+    resAllocInfo.SizeInBytes = D3D12_DEFAULT_RESOURCE_PLACEMENT_ALIGNMENT;

+    resAllocInfo.Alignment = D3D12_DEFAULT_RESOURCE_PLACEMENT_ALIGNMENT;

+

+    D3D12MA::Allocation* alloc = nullptr;

+    CHECK_HR( ctx.allocator->AllocateMemory(&allocDesc, heapFlags, &resAllocInfo, &alloc) );

+    ResourceWithAllocation res;

+    res.allocation.reset(alloc);

+

+    // Must be created as separate allocation.

+    CHECK_BOOL( res.allocation->GetOffset() == 0 );

+}

+

 static void TestPlacedResources(const TestContext& ctx)

 {

     wprintf(L"Test placed resources\n");

@@ -775,6 +798,7 @@
 {

     TestFrameIndexAndJson(ctx);

     TestCommittedResourcesAndJson(ctx);

+    TestCustomHeapFlags(ctx);

     TestPlacedResources(ctx);

     TestAliasingMemory(ctx);

     TestMapping(ctx);