ShaderBuilderHelpersWebGPU.cpp

#include "ShaderBuilderHelpersWebGPU.h"
#include "Rendering/IBindGroups.h"
#include "Context.h"
#include "Utilities/Strings.h"
#include "Utilities/Preprocessor.h"
#include "ResourceStore.h"
#include "Renderer/EnginePermutations.h"
#include "Rendering/IGraphicsDevice.h"
 
#include "Foundation/Logging/Logger.h"
 
#include <array>
#include <sstream>
#include <map>
 
namespace
{
  using Cogs::Core::BindGroup;
  const Cogs::Logging::Log logger = Cogs::Logging::getLogger("ShaderBuilderHelpersWebGPU");
 
  std::vector<std::string> sceneGetters = {
  "getClipFromViewMatrix",
  "getClipFromWorldMatrix",
  "getViewFromWorldMatrix",
  "getViewFromClipMatrix",
  "getWorldFromViewMatrix",
  "getInverseViewMatrix",
  "getViewFromViewportMatrix",
  "getViewportFromViewMatrix",
  "getPeriodicWorldPosAndCell",
  "getPeriodicWorldPos"
  };
 
  std::vector<std::string> objectGetters = {
  "getWorldFromLocalMatrix",
  "getObjectId"
  };
 
  struct EngineBufferDesc {
    std::string name;
    BindGroup group;
    std::vector<std::string> getters;
  };
 
 
  const std::array engineBuffers = {
    EngineBufferDesc{"SceneBuffer", BindGroup::Scene, sceneGetters},
    EngineBufferDesc{"ObjectBuffer", BindGroup::Object, objectGetters},
  };
 
  struct EngineTexture {
    std::string name;
    Cogs::ResourceDimensions dimensions;
    bool isDepthTexture = false;
    BindGroup group = BindGroup::Default;
  };
 
  const std::array engineTextures = {
   EngineTexture{.name = "environmentSky", .dimensions = Cogs::ResourceDimensions::TextureCube, .group = BindGroup::LightTextures},
   EngineTexture{.name = "environmentRadiance", .dimensions = Cogs::ResourceDimensions::TextureCube, .group = BindGroup::LightTextures},
   EngineTexture{.name = "environmentIrradiance", .dimensions = Cogs::ResourceDimensions::TextureCube, .group = BindGroup::LightTextures},
   EngineTexture{.name = "ambientIrradiance", .dimensions = Cogs::ResourceDimensions::TextureCube, .group = BindGroup::LightTextures},
   EngineTexture{.name = "brdfLUT", .dimensions = Cogs::ResourceDimensions::Texture2D, .group = BindGroup::LightTextures},
   EngineTexture{.name = "cascadedShadowMap", .dimensions = Cogs::ResourceDimensions::Texture2DArray, .isDepthTexture = true, .group = BindGroup::LightTextures},
   EngineTexture{.name = "cubeShadowMap", .dimensions = Cogs::ResourceDimensions::TextureCube, .isDepthTexture = true, .group = BindGroup::LightTextures},
  };
 
  struct EngineSampler {
    std::string name;
    BindGroup group = BindGroup::Default;
    bool isDepthTexture = false;
  };
 
  const std::array engineSamplers = {
    EngineSampler{.name = "linearSampler", .group = BindGroup::LightTextures},
    EngineSampler{.name = "linearClampSampler", .group = BindGroup::LightTextures},
    EngineSampler{.name = "pointSampler", .group = BindGroup::LightTextures},
    EngineSampler{.name = "pointClampSampler", .group = BindGroup::LightTextures},
    EngineSampler{.name = "shadowSampler", .group = BindGroup::LightTextures, .isDepthTexture = true},
  };
 
 
  const uint32_t SceneGlobalVisibility = (Cogs::BindingVisibilityVertex | Cogs::BindingVisibilityFragment | Cogs::BindingVisibilityCompute);
 
  Cogs::BindGroupEntryDescription makeBufferEntry(const uint32_t binding, const char* name)
  {
    Cogs::BindGroupEntryDescription entry;
    entry.binding = binding;
    entry.nameHash = Cogs::hash(name);
    entry.resourceType = Cogs::BindingResourceType::UniformBuffer;
    entry.visibility = SceneGlobalVisibility;
    return entry;
  }
 
  Cogs::BindGroupEntryDescription makeTextureEntry(
    const uint32_t binding,
    const char* name,
    const Cogs::ResourceDimensions dimension,
    const bool isDepthTexture)
  {
    Cogs::BindGroupEntryDescription entry;
    entry.binding = binding;
    entry.nameHash = Cogs::hash(name);
    entry.resourceType = Cogs::BindingResourceType::Texture;
    entry.visibility = SceneGlobalVisibility;
    entry.textureDimension = dimension;
    entry.isDepthTexture = isDepthTexture;
    return entry;
  }
 
  Cogs::BindGroupEntryDescription makeSamplerEntry(const uint32_t binding, const char* name, bool isDepthTexture)
  {
    Cogs::BindGroupEntryDescription entry;
    entry.binding = binding;
    entry.nameHash = Cogs::hash(name);
    entry.resourceType = Cogs::BindingResourceType::Sampler;
    entry.samplerBindingType = isDepthTexture ? Cogs::BindingSamplerBindingType::Comparison : Cogs::BindingSamplerBindingType::Filtering;
    entry.visibility = SceneGlobalVisibility;
    return entry;
  }
 
  [[nodiscard]]
  std::string textureDataType(Cogs::BindingTextureSampleType  type)
  {
    switch (type) {
    case Cogs::BindingTextureSampleType::Unknown: // Defaults to f32
    case Cogs::BindingTextureSampleType::Float:
    case Cogs::BindingTextureSampleType::UnfilterableFloat:
    case Cogs::BindingTextureSampleType::Depth:
      return "f32"; break;
    case Cogs::BindingTextureSampleType::Sint:
      return "i32"; break;
    case Cogs::BindingTextureSampleType::Uint:
      return "u32"; break;
    default:
      LOG_ERROR(logger, "Unsupported texture datatype %d", int(type));
      return "<illegal>";
      break;
    }
  }
 
}
 
namespace Cogs::Core
{
  const Cogs::BindGroupEntryDescription* findEntryByName(const Cogs::BindGroupDescription& group, std::string_view name)
  {
    size_t nameHash = Cogs::hash(name);
    for (uint16_t i = 0; i < group.numEntries; ++i) {
      const Cogs::BindGroupEntryDescription& entry = group.entries[i];
      if (entry.nameHash == nameHash) {
        return &entry;
      }
    }
    return nullptr;
  }
 
  void addBufferToSource(std::string& source, const std::string& name, BindGroup group, const Cogs::BindGroupEntryDescription& entry)
  {
    source.append("#include \"Engine/");
    source.append(name);
    source.append(".wgsl\"\n");
    source.append("@group(");
    source.append(std::to_string(static_cast<size_t>(group)));
    source.append(") @binding(");
    source.append(std::to_string(entry.binding));
    source.append(") var<uniform> ");
    source.append(name);
    source.append(" : ");
    source.append(name);
    source.append("_t;\n");
  }
 
  void addSamplerToSource(std::string& source, const std::string& name, BindGroup group, const Cogs::BindGroupEntryDescription& entry)
  {
    source.append("@group(");
    source.append(std::to_string(static_cast<size_t>(group)));
    source.append(") @binding(");
    source.append(std::to_string(entry.binding));
    source.append(") var ");
    source.append(name);
    source.append(" : ");
    source.append(entry.samplerBindingType == Cogs::BindingSamplerBindingType::Comparison ? "sampler_comparison" : "sampler");
    source.append(";\n");
  }
 
  void addTextureToSource(std::string& source, const std::string& name, BindGroup group, const Cogs::BindGroupEntryDescription& entry)
  {
    std::string type;
    type.reserve(50);
    type += "texture";
 
    if (entry.isDepthTexture) {
      type += "_depth";
    }
 
    switch (entry.textureDimension) {
    case Cogs::ResourceDimensions::Texture2D:
      type += "_2d";
      break;
    case Cogs::ResourceDimensions::TextureCube:
      type += "_cube";
      break;
    case Cogs::ResourceDimensions::Texture2DArray:
      type += "_2d_array";
      break;
    case Cogs::ResourceDimensions::Texture3D:
      type += "_3d";
      break;
    default:
      LOG_ERROR(logger, "Unsupported texture dimension %d", int(entry.textureDimension));
      break;
    }
 
    if (!entry.isDepthTexture) {
      std::string dataType = textureDataType(entry.textureSampleType);
      type += "<" + dataType + ">";
    }
 
    source.append("@group(");
    source.append(std::to_string(static_cast<size_t>(group)));
    source.append(") @binding(");
    source.append(std::to_string(entry.binding));
    source.append(") var");
    source.append(" ");
    source.append(name);
    source.append(" : ");
    source.append(type);
    source.append(";\n");
  }
 
 
  const Cogs::BindGroupSetDescription& getEngineBindGroupDescription()
  {
    static Cogs::BindGroupSetDescription staticDesc = {};
 
    if (staticDesc.numGroups == 0) {
      Cogs::BindGroupSetDescription groups = {};
      groups.numGroups = static_cast<uint16_t>(BindGroup::Count);
      for (const auto& buffer : engineBuffers) {
        Cogs::BindGroupDescription& g = groups.groups[static_cast<size_t>(buffer.group)];
        uint16_t binding = g.numEntries;
        g.entries[g.numEntries++] = makeBufferEntry(binding, buffer.name.c_str());
        ;
      }
 
      for (const auto& texture : engineTextures) {
        Cogs::BindGroupDescription& g = groups.groups[static_cast<size_t>(texture.group)];
        uint16_t binding = g.numEntries;
        g.entries[g.numEntries++] = makeTextureEntry(binding, texture.name.c_str(), texture.dimensions, texture.isDepthTexture);
      }
 
      for (const auto& sampler : engineSamplers) {
        Cogs::BindGroupDescription& g = groups.groups[static_cast<size_t>(sampler.group)];
        uint16_t binding = g.numEntries;
        g.entries[g.numEntries++] = makeSamplerEntry(binding, sampler.name.c_str(), sampler.isDepthTexture);
      }
      staticDesc = groups;
    }
    return staticDesc;
  }
 
  void addEngineUniformsToSourceWebGPU(std::string& source, const std::unordered_set<Cogs::Core::StringRef>& identifiersSeen, std::vector<bool>& bindGroupUsed) {
    bindGroupUsed.resize(static_cast<size_t>(BindGroup::Count), false);
    const Cogs::BindGroupSetDescription& desc = getEngineBindGroupDescription();
 
    for (const auto& buffer : engineBuffers) {
      bool add = false;
      for (auto& get : buffer.getters) {
        if (identifiersSeen.contains(Strings::add(get))) {
          add = true;
          break;
        }
      }
      if (identifiersSeen.contains(Cogs::Core::Strings::add(buffer.name))) {
        add = true;
      }
      if (!add) continue;
      bindGroupUsed[static_cast<size_t>(buffer.group)] = true;
      const Cogs::BindGroupEntryDescription* entry = findEntryByName(desc.groups[static_cast<size_t>(buffer.group)], buffer.name);
      if (entry) {
        addBufferToSource(source, buffer.name, buffer.group, *entry);
      }
    }
 
    for (const auto& texture : engineTextures) {
      if (!identifiersSeen.contains(Cogs::Core::Strings::add(texture.name))) {
        continue;
      }
 
      const Cogs::BindGroupEntryDescription* entry = findEntryByName(desc.groups[static_cast<size_t>(texture.group)], texture.name);
      if (entry) {
        addTextureToSource(source, texture.name, texture.group, *entry);
        bindGroupUsed[static_cast<size_t>(texture.group)] = true;
      }
    }
 
    for (const auto& sampler : engineSamplers) {
      if (!identifiersSeen.contains(Cogs::Core::Strings::add(sampler.name))) {
        continue;
      }
      const Cogs::BindGroupEntryDescription* entry = findEntryByName(desc.groups[static_cast<size_t>(sampler.group)], sampler.name);
      if (entry) {
        bindGroupUsed[static_cast<size_t>(sampler.group)] = true;
        addSamplerToSource(source, sampler.name, sampler.group, *entry);
      }
    }
  }
 
}