storage image usage fix

cocos/rendering/custom/web-pipeline.ts

@@ -238,7 +238,11 @@ export class WebSetter {
         this._copyToBuffer(v, offset, Type.FLOAT);
     }
     public setBuffer (name: string, buffer: Buffer): void {
-        // TODO
+        if (this._getCurrDescriptorBlock(name) === -1) {
+            return;
+        }
+        const num = this._lg.attributeIndex.get(name)!;
+        this._data.buffers.set(num, buffer);
     }
     public setTexture (name: string, texture: Texture): void {
         if (this._getCurrDescriptorBlock(name) === -1) {
@@ -1344,7 +1348,7 @@ export class WebComputePassBuilder extends WebSetter implements ComputePassBuild
         this._addComputeResource(name, accessType, slotName);
     }
     addStorageImage (name: string, accessType: AccessType, slotName: string): void {
-        throw new Error('Method not implemented.');
+        this._addComputeResource(name, accessType, slotName);
     }
     addMaterialTexture (resourceName: string, flags?: ShaderStageFlagBit | undefined): void {
         throw new Error('Method not implemented.');
@@ -1493,6 +1497,9 @@ export class WebPipeline implements BasicPipeline {
     }
     updateRenderWindow (name: string, renderWindow: RenderWindow): void {
         const resId = this.resourceGraph.vertex(name);
+        const desc = this.resourceGraph.getDesc(resId);
+        desc.width = renderWindow.width;
+        desc.height = renderWindow.height;
         const currFbo = this.resourceGraph._vertices[resId]._object;
         if (currFbo !== renderWindow.framebuffer) {
             this.resourceGraph._vertices[resId]._object = renderWindow.framebuffer;

cocos/rendering/custom/web-program-library.ts

@@ -1,3 +1,4 @@
+/* eslint-disable no-console */
 /****************************************************************************
  Copyright (c) 2021-2023 Xiamen Yaji Software Co., Ltd.
 
@@ -91,7 +92,7 @@ function findBinding (shaderInfo: ShaderInfo, name: string): { set: number, bind
 
 function overwriteShaderSourceBinding (shaderInfo: ShaderInfo, source: string): string {
     let code = source;
-    const samplerExp = /layout\s*\(([^\)])+\)\s+uniform\s+(\b\w+\b\s+)?sampler(\w+)\s+(\b\w+\b)/g;
+    const samplerExp = /layout\s*\(([^)])+\)\s+uniform\s+(\b\w+\b\s+)?sampler(\w+)\s+(\b\w+\b)/g;
     let samplerIter = samplerExp.exec(code);
     while (samplerIter) {
         const name = samplerIter[4];
@@ -101,13 +102,20 @@ function overwriteShaderSourceBinding (shaderInfo: ShaderInfo, source: string):
         code = code.replace(samplerIter[0], replaceStr);
         samplerIter = samplerExp.exec(code);
     }
-    const blockExp = /layout\s*\(([^\)])+\)\s*(readonly)?\s*\b(uniform|buffer)\b\s+(\b\w+\b)\s*[{;]/g;
+    const blockExp = /layout\s*\(([^)]+)\)\s*(readonly|writeonly)?\s*\b((uniform\s*|buffer\s*|image2D\s*){1,2})\b\s*(\b\w+\b)\s*[{;]/g;
     let blockIter = blockExp.exec(code);
     while (blockIter) {
-        const name = blockIter[4];
+        const name = blockIter[5];
         const { set, binding } = findBinding(shaderInfo, name);
         const accessStr = blockIter[2] ? blockIter[2] : '';
-        const replaceStr = `layout(set = ${set}, binding = ${binding}) ${accessStr} ${blockIter[3]} ${blockIter[4]} {`;
+        let endStr = ' {';
+        if (blockIter[3].includes('image')) {
+            endStr = `;`;
+        }
+        let desc = blockIter[1];
+        desc = desc.replace(/set\s*=\s*\d+/g, `set = ${set}`);
+        desc = desc.replace(/binding\s*=\s*\d+/g, `binding = ${binding}`);
+        const replaceStr = `layout(${desc}) ${accessStr} ${blockIter[3]} ${blockIter[5]}${endStr}`;
         code = code.replace(blockIter[0], replaceStr);
         blockIter = blockExp.exec(code);
     }
@@ -474,7 +482,7 @@ function calculateFlattenedBinding (
     shaderInfo: ShaderInfo,
 ): void {
     // Descriptors of UniformBlock starts from 0, and Descriptors of SamplerTexture starts from the end of UniformBlock.
-    const uniformBlockCapacities = new Array(4);
+    const uniformBlockCapacities = new Array<number>(4);
     {
         const passCapacity = descriptorSets[UpdateFrequency.PER_PASS]?.uniformBlockCapacity || 0;
         const phaseCapacity = descriptorSets[UpdateFrequency.PER_PHASE]?.uniformBlockCapacity || 0;
@@ -496,7 +504,7 @@ function calculateFlattenedBinding (
         const batchOffset = instanceOffset + instanceCapacity;
 
         // save uniform block offsets by set index
-        const uniformBlockOffsets = new Array(4);
+        const uniformBlockOffsets = new Array<number>(4);
         uniformBlockOffsets[_setIndex[UpdateFrequency.PER_PASS]] = passOffset;
         uniformBlockOffsets[_setIndex[UpdateFrequency.PER_PHASE]] = phaseOffset;
         uniformBlockOffsets[_setIndex[UpdateFrequency.PER_BATCH]] = batchOffset;
@@ -521,7 +529,7 @@ function calculateFlattenedBinding (
         const batchOffset = instanceOffset + instanceCapacity;
 
         // save sampler texture offsets by set index
-        const samplerTextureOffsets = new Array(4);
+        const samplerTextureOffsets = new Array<number>(4);
         samplerTextureOffsets[_setIndex[UpdateFrequency.PER_PASS]] = passOffset;
         samplerTextureOffsets[_setIndex[UpdateFrequency.PER_PHASE]] = phaseOffset;
         samplerTextureOffsets[_setIndex[UpdateFrequency.PER_BATCH]] = batchOffset;

native/cocos/renderer/gfx-base/GFXBarrier.cpp

@@ -241,7 +241,7 @@ constexpr AccessElem ACCESS_MAP[] = {
      AccessFlags::VERTEX_SHADER_READ_TEXTURE},
 
     {IGNORE_MEMUSAGE & IGNORE_RESTYPE,
-     ACCESS_READ | SHADERSTAGE_VERT | CMN_STORAGE,
+     ACCESS_READ | ACCESS_WRITE | SHADERSTAGE_VERT | CMN_STORAGE,
      AccessFlags::VERTEX_SHADER_READ_OTHER},
 
     {IGNORE_MEMUSAGE,
@@ -251,7 +251,7 @@ constexpr AccessElem ACCESS_MAP[] = {
     {CARE_MEMACCESS | CARE_RESTYPE | CARE_SHADERSTAGE | CARE_CMNUSAGE,
      ACCESS_READ | RES_TEXTURE | SHADERSTAGE_FRAG | CMN_ROM,
      AccessFlags::FRAGMENT_SHADER_READ_TEXTURE,
-     CMN_STORAGE | CMN_VB_OR_DS},
+     CMN_VB_OR_DS},
 
     {IGNORE_MEMUSAGE,
      ACCESS_READ | RES_TEXTURE | SHADERSTAGE_FRAG | CMN_IB_OR_CA | CMN_INDIRECT_OR_INPUT,
@@ -262,7 +262,7 @@ constexpr AccessElem ACCESS_MAP[] = {
      AccessFlags::FRAGMENT_SHADER_READ_DEPTH_STENCIL_INPUT_ATTACHMENT},
 
     {IGNORE_MEMUSAGE & IGNORE_RESTYPE,
-     ACCESS_READ | SHADERSTAGE_FRAG | CMN_STORAGE,
+     ACCESS_READ | ACCESS_WRITE | SHADERSTAGE_FRAG | CMN_STORAGE,
      AccessFlags::FRAGMENT_SHADER_READ_OTHER,
      CMN_SHADING_RATE},
 
@@ -285,9 +285,9 @@ constexpr AccessElem ACCESS_MAP[] = {
 
     // shading rate has its own flag
     {CARE_MEMACCESS | CARE_SHADERSTAGE | CARE_CMNUSAGE,
-     ACCESS_READ | SHADERSTAGE_COMP | CMN_STORAGE,
+     ACCESS_READ | ACCESS_WRITE | SHADERSTAGE_COMP | CMN_STORAGE,
      AccessFlags::COMPUTE_SHADER_READ_OTHER,
-     RES_TEXTURE | CMN_ROM},
+     CMN_ROM},
 
     {CARE_MEMACCESS | CARE_CMNUSAGE,
      ACCESS_READ | CMN_COPY_SRC,

native/cocos/renderer/gfx-vulkan/VKGPUObjects.h

@@ -1062,8 +1062,11 @@ class CCVKGPUDescriptorHub final {
         if (hasFlag(texture->gpuTexture->flags, TextureFlagBit::GENERAL_LAYOUT)) {
             descriptor->imageLayout = VK_IMAGE_LAYOUT_GENERAL;
         } else {
-            if (hasAllFlags(flags, AccessFlagBit::FRAGMENT_SHADER_READ_COLOR_INPUT_ATTACHMENT | AccessFlagBit::COLOR_ATTACHMENT_WRITE) ||
-                hasAllFlags(flags, AccessFlagBit::FRAGMENT_SHADER_READ_DEPTH_STENCIL_INPUT_ATTACHMENT | AccessFlagBit::DEPTH_STENCIL_ATTACHMENT_WRITE)) {
+            bool inoutAttachment = hasAllFlags(flags, AccessFlagBit::FRAGMENT_SHADER_READ_COLOR_INPUT_ATTACHMENT | AccessFlagBit::COLOR_ATTACHMENT_WRITE) ||
+                                   hasAllFlags(flags, AccessFlagBit::FRAGMENT_SHADER_READ_DEPTH_STENCIL_INPUT_ATTACHMENT | AccessFlagBit::DEPTH_STENCIL_ATTACHMENT_WRITE);
+            bool storageWrite = hasAnyFlags(flags, AccessFlagBit::VERTEX_SHADER_WRITE | AccessFlagBit::FRAGMENT_SHADER_WRITE | AccessFlagBit::COMPUTE_SHADER_WRITE);
+
+            if (inoutAttachment || storageWrite) {
                 descriptor->imageLayout = VK_IMAGE_LAYOUT_GENERAL;
             } else if (hasFlag(texture->gpuTexture->usage, TextureUsage::DEPTH_STENCIL_ATTACHMENT)) {
                 descriptor->imageLayout = VK_IMAGE_LAYOUT_DEPTH_STENCIL_READ_ONLY_OPTIMAL;

native/cocos/renderer/gfx-wgpu/WGPUDescriptorSet.cpp

@@ -193,7 +193,7 @@ void CCWGPUDescriptorSet::update() {
             auto &bindGroupEntry = _gpuBindGroupObj->bindGroupEntries.emplace_back();
             bindGroupEntry.binding = binding.binding;
             bindGroupEntry.textureView = static_cast<WGPUTextureView>(ccTexture->getPlaneView(0));
-            dsLayout->updateTextureLayout(i, ccTexture);
+            dsLayout->updateSampledTextureLayout(i, ccTexture);
             _gpuDescriptorObj->gpuDescriptors[i].texture = ccTexture;
 
             auto &sBindGroupEntry = _gpuBindGroupObj->bindGroupEntries.emplace_back();
@@ -204,14 +204,18 @@ void CCWGPUDescriptorSet::update() {
             _gpuDescriptorObj->gpuDescriptors[i].sampler = ccSampler;
         } else if (DescriptorType::STORAGE_IMAGE == bindings[i].descriptorType || DescriptorType::TEXTURE == bindings[i].descriptorType) {
             auto *ccTexture = _textures[resourceIndex].ptr ? static_cast<CCWGPUTexture *>(_textures[resourceIndex].ptr) : CCWGPUTexture::defaultCommonTexture();
-            auto &bindGroupEntry = _gpuBindGroupObj->bindGroupEntries[resourceIndex];
+            auto &bindGroupEntry = _gpuBindGroupObj->bindGroupEntries.emplace_back();
             bindGroupEntry.binding = binding.binding;
             bindGroupEntry.textureView = static_cast<WGPUTextureView>(ccTexture->getPlaneView(0));
-            dsLayout->updateTextureLayout(i, ccTexture);
+            if (DescriptorType::STORAGE_IMAGE == bindings[i].descriptorType) {
+                dsLayout->updateStorageTextureLayout(i, ccTexture);
+            } else {
+                dsLayout->updateSampledTextureLayout(i, ccTexture);
+            }
             _gpuDescriptorObj->gpuDescriptors[i].texture = ccTexture;
         } else if (DescriptorType::SAMPLER == bindings[i].descriptorType) {
             auto *ccSampler = _samplers[resourceIndex].ptr ? static_cast<CCWGPUSampler *>(_samplers[resourceIndex].ptr) : CCWGPUSampler::defaultFilterableSampler();
-            auto &bindGroupEntry = _gpuBindGroupObj->bindGroupEntries[resourceIndex];
+            auto &bindGroupEntry = _gpuBindGroupObj->bindGroupEntries.emplace_back();
             bindGroupEntry.binding = binding.binding;
             bindGroupEntry.sampler = ccSampler->gpuSampler();
             dsLayout->updateSamplerLayout(i, ccSampler);

native/cocos/renderer/gfx-wgpu/WGPUDescriptorSetLayout.cpp

@@ -131,29 +131,33 @@ WGPUTextureSampleType sampletypeTraits(const CCWGPUTexture *texture, uint32_t pl
 }
 } // namespace
 
-void CCWGPUDescriptorSetLayout::updateTextureLayout(uint8_t index, const CCWGPUTexture *texture, uint32_t plane) {
+void CCWGPUDescriptorSetLayout::updateSampledTextureLayout(uint8_t index, const CCWGPUTexture *texture, uint32_t plane) {
     WGPUBindGroupLayoutEntry textureEntry{};
     textureEntry.binding = _bindings[index].binding;
     textureEntry.visibility = toWGPUShaderStageFlag(_bindings[index].stageFlags);
 
     CC_ASSERT(texture);
 
-    if (texture->getInfo().usage == TextureUsageBit::STORAGE) {
-        WGPUStorageTextureBindingLayout storageTextureLayout{};
-        storageTextureLayout.access = WGPUStorageTextureAccess::WGPUStorageTextureAccess_WriteOnly;
-        storageTextureLayout.format = formatTraits(texture, plane);
-        TextureType type = texture->isTextureView() ? texture->getViewInfo().type : texture->getInfo().type;
-        storageTextureLayout.viewDimension = toWGPUTextureViewDimension(type);
-        textureEntry.storageTexture = storageTextureLayout;
-    } else {
         WGPUTextureBindingLayout textureLayout{};
         textureLayout.sampleType = sampletypeTraits(texture, plane); // textureSampleTypeTrait(texture->getFormat());
         const CCWGPUTexture *ccTex = static_cast<const CCWGPUTexture *>(texture->isTextureView() ? texture->getViewInfo().texture : texture);
         TextureType type = ccTex->getViewInfo().type;
         textureLayout.viewDimension = toWGPUTextureViewDimension(type);
         textureLayout.multisampled = ccTex->getInfo().samples != SampleCount::X1;
         textureEntry.texture = textureLayout;
+    _gpuLayoutEntryObj->bindGroupLayoutEntries[textureEntry.binding] = textureEntry;
     }
+void CCWGPUDescriptorSetLayout::updateStorageTextureLayout(uint8_t index, const CCWGPUTexture *texture, uint32_t plane) {
+    WGPUBindGroupLayoutEntry textureEntry{};
+    textureEntry.binding = _bindings[index].binding;
+    textureEntry.visibility = toWGPUShaderStageFlag(_bindings[index].stageFlags);
+    CC_ASSERT(texture);
+    WGPUStorageTextureBindingLayout storageTextureLayout{};
+    storageTextureLayout.access = WGPUStorageTextureAccess::WGPUStorageTextureAccess_WriteOnly;
+    storageTextureLayout.format = formatTraits(texture, plane);
+    TextureType type = texture->isTextureView() ? texture->getViewInfo().type : texture->getInfo().type;
+    storageTextureLayout.viewDimension = toWGPUTextureViewDimension(type);
+    textureEntry.storageTexture = storageTextureLayout;
     _gpuLayoutEntryObj->bindGroupLayoutEntries[textureEntry.binding] = textureEntry;
 }
 

native/cocos/renderer/gfx-wgpu/WGPUDescriptorSetLayout.h

@@ -45,7 +45,8 @@ class CCWGPUDescriptorSetLayout final : public DescriptorSetLayout {
     inline CCWGPUBindGroupLayoutObject *gpuLayoutEntryObject() { return _gpuLayoutEntryObj; }
 
     void updateBufferLayout(uint8_t index, const CCWGPUBuffer *buffer, AccessFlags flags);
-    void updateTextureLayout(uint8_t index, const CCWGPUTexture *texture, uint32_t plane = 0);
+    void updateSampledTextureLayout(uint8_t index, const CCWGPUTexture *texture, uint32_t plane = 0);
+    void updateStorageTextureLayout(uint8_t index, const CCWGPUTexture *texture, uint32_t plane = 0);
     void updateSamplerLayout(uint8_t index, const CCWGPUSampler *sampler);
 
     inline void setBindings(const DescriptorSetLayoutBindingList &list) { _bindings.assign(list.begin(), list.end()); }

native/cocos/renderer/pipeline/custom/NativeExecutor.cpp

@@ -537,10 +537,20 @@ gfx::DescriptorSet* initDescriptorSet(
 
                     auto iter = resourceIndex.find(d.descriptorID);
                     if (iter != resourceIndex.end()) {
+                        gfx::AccessFlags access = gfx::AccessFlagBit::NONE;
+                        if (accessNode != nullptr) {
+                            const auto& resID = iter->second;
+                            // whole access only now.
+                            auto parentID = parent(resID, resg);
+                            parentID = parentID == ResourceGraph::null_vertex() ? resID : parentID;
+                            const auto& resName = get(ResourceGraph::NameTag{}, resg, parentID);
+                            access = accessNode->resourceStatus.at(resName).accessFlag;
+                        }
+
                         // render graph textures
                         auto* texture = resg.getTexture(iter->second);
                         CC_ENSURES(texture);
-                        newSet->bindTexture(bindID, texture);
+                        newSet->bindTexture(bindID, texture, 0, access);
                     }
                     bindID += d.count;
                 }
@@ -940,11 +950,12 @@ struct RenderGraphUploadVisitor : boost::dfs_visitor<> {
             auto& set = iter->second;
             const auto& user = get(RenderGraph::DataTag{}, ctx.g, vertID);
             auto& node = ctx.context.layoutGraphResources.at(layoutID);
+            const auto& accessNode = ctx.fgd.getAccessNode(vertID);
             auto* perPassSet = initDescriptorSet(
                 ctx.resourceGraph,
                 ctx.device, ctx.cmdBuff,
                 *ctx.context.defaultResource, ctx.lg,
-                resourceIndex, set, user, node);
+                resourceIndex, set, user, node, &accessNode);
             CC_ENSURES(perPassSet);
             ctx.renderGraphDescriptorSet[vertID] = perPassSet;
         } else if (holds<QueueTag>(vertID, ctx.g)) {