ShaderBuilderPostProcessWebGPU.cpp

#include "ShaderBuilderPostProcess.h"
#include "ShaderBuilderHelpersWebGPU.h"
#include "Context.h"
#include "Utilities/Strings.h"
#include "Utilities/Preprocessor.h"
#include "ResourceStore.h"
#include "Renderer/Renderer.h"
#include "Rendering/IGraphicsDevice.h"
#include "Renderer/Tasks/PostProcessTask.h"
#include "Foundation/Logging/Logger.h"
#include "Rendering/IContext.h"
#include "Rendering/IEffects.h"
#include "Renderer/RenderTarget.h"
 
#include <cassert>
#include <sstream>
#include <map>
#include <algorithm>
 
namespace {
  constexpr size_t InitialBufferCapacity = 4096u;
 
  const Cogs::Logging::Log logger = Cogs::Logging::getLogger("ShaderBuilderPostProcessWebGPU");
 
  using namespace Cogs::Core;
 
  void removeSuffix(std::string& dst, const std::string_view& suffix)
  {
    auto pos = dst.find(suffix);
    if (pos == std::string::npos) return;
    dst.erase(pos);
  }
 
  [[nodiscard]] std::string_view parameterType(const ParsedDataType type)
  {
    switch (type) {
    case ParsedDataType::Float:           return "f32"; break;
    case ParsedDataType::Float2:          return "vec2f";  break;
    case ParsedDataType::Float3:          return "vec3f";  break;
    case ParsedDataType::Float4:          return "vec4f";  break;
    case ParsedDataType::Float4x4:        return "mat4x4f";  break;
    case ParsedDataType::Int:             return "i32";   break;
    case ParsedDataType::Int2:            return "vec2i"; break;
    case ParsedDataType::Int3:            return "vec3i"; break;
    case ParsedDataType::Int4:            return "vec4i"; break;
    case ParsedDataType::UInt:            return "u32";   break;
    case ParsedDataType::UInt2:           return "vec2u"; break;
    case ParsedDataType::UInt3:           return "vec3u"; break;
    case ParsedDataType::UInt4:           return "vec4u"; break;
    default:
      LOG_ERROR(logger, "Unsupported attribute type %d", int(type));
      return "<illegal>";
      break;
    }
  }
 
  [[nodiscard]]
  Cogs::BindingTextureSampleType bindingTextureSampleType(const ParsedDataType type, const bool isDepthTexture)
  {
    if (isDepthTexture) {
      return Cogs::BindingTextureSampleType::Depth;
    }
 
    switch (type) {
    case ParsedDataType::Int:
    case ParsedDataType::Int2:
    case ParsedDataType::Int3:
    case ParsedDataType::Int4:
      return Cogs::BindingTextureSampleType::Sint;
    case ParsedDataType::UInt:
    case ParsedDataType::UInt2:
    case ParsedDataType::UInt3:
    case ParsedDataType::UInt4:
      return Cogs::BindingTextureSampleType::Uint;
    default:
      return Cogs::BindingTextureSampleType::Float;
    }
  }
 
  [[nodiscard]]
  std::string convertDefinesToConstExpressions(const std::string& s) {
    std::map<std::string, std::string> addedDefines;
    std::string result;
    result.reserve(s.size());
 
    std::istringstream iss(s);
 
    std::string identifier;
    std::string replacement;
    for (std::string line; std::getline(iss, line); )
    {
      identifier.clear();
      replacement.clear();
 
      std::string_view sv(line);
      auto pos = sv.find_first_not_of(" \t");
      if (pos == std::string_view::npos || sv.substr(pos, 7) != "#define") {
        result += line + "\n";
        continue;
      }
      sv = sv.substr(pos + 7);
 
      pos = sv.find_first_not_of(" \t");
      if (pos == std::string_view::npos) {
        result += line + "\n";
        continue;
      }
      sv = sv.substr(pos);
 
      auto end = sv.find_first_of(" \t");
      if (end == std::string_view::npos) {
        identifier = sv;
        replacement = "1";
      } else {
        identifier = sv.substr(0, end);
        sv = sv.substr(end);
        pos = sv.find_first_not_of(" \t");
        if (pos == std::string_view::npos) {
          replacement = "1";
        } else {
          auto last = sv.find_last_not_of(" \t");
          replacement = sv.substr(pos, last - pos + 1);
        }
      }
 
      auto it = addedDefines.find(identifier);
      if (it != addedDefines.end()) {
        if (replacement != (*it).second) {
          LOG_WARNING(logger, "Shader generation with inconsistent preprosessor define %s set to %s conflicts with previous define %s.", identifier.c_str(), replacement.c_str(), (*it).second.c_str());
        }
        continue;
      }
      result += "const " + identifier + " = " + replacement + ";\n";
      addedDefines.try_emplace(identifier, replacement);
    }
    return result;
  }
 
  void addInterface(std::string& src) {
    src.append(
      R"(
struct VertexIn {
  @builtin(position) position: vec4f,
  @location(0) TexCoords: vec2f,
  @location(1) NormalizedCoords: vec2f,
};
 
)");
  }
 
  void addDefine(std::string& output, const std::string& name, const std::string& value)
  {
    output.append("#define ");
    output.append(name);
    output.append(" ");
    output.append(value);
    output.append("\n");
  }
 
  void addDefines(std::string& output, const std::vector<std::pair<std::string, std::string>>& definitions)
  {
    for (const std::pair<std::string, std::string>& define : definitions) {
      addDefine(output, define.first, define.second);
    }
  }
 
  void addOutputStruct(std::string& src, const PipelineOptions& options, bool writeDepth) {
    bool hasCustomTargets = false;
    src.append("struct FragmentOut {\n");
    if (writeDepth) {
      src.append("  @builtin(frag_depth) fragDepth: f32,\n");
    }
    std::string prefix = "target_";
    for (auto o : options) {
      auto key = o.first;
 
      for (auto& c : key) {
        c = static_cast<decltype(key)::value_type>(std::tolower(c));
      }
      if (key.substr(0, prefix.size()) != prefix) continue;
      hasCustomTargets = true;
 
      int location = std::stoi(key.substr(prefix.size()));
      TokenStream tokens;
      Cogs::Core::split(o.second, " ", tokens);
      if (tokens.size() != 2) {
        LOG_ERROR(logger, "Unable to parse pipeline output option %s", o.second.c_str());
        continue;
      }
      std::string_view datatype = "Undefined";
      switch (tokens[1].hashLowercase()) {
      case Cogs::hash("float"):  datatype = parameterType(ParsedDataType::Float);  break;
      case Cogs::hash("float2"): datatype = parameterType(ParsedDataType::Float2); break;
      case Cogs::hash("float3"): datatype = parameterType(ParsedDataType::Float3); break;
      case Cogs::hash("float4"): datatype = parameterType(ParsedDataType::Float4); break;
      case Cogs::hash("uint"):   datatype = parameterType(ParsedDataType::UInt);   break;
      case Cogs::hash("uint2"):  datatype = parameterType(ParsedDataType::UInt2);  break;
      case Cogs::hash("uint3"):  datatype = parameterType(ParsedDataType::UInt3);  break;
      case Cogs::hash("uint4"):  datatype = parameterType(ParsedDataType::UInt4);  break;
      case Cogs::hash("int"):    datatype = parameterType(ParsedDataType::Int);    break;
      case Cogs::hash("int2"):   datatype = parameterType(ParsedDataType::Int2);   break;
      case Cogs::hash("int3"):   datatype = parameterType(ParsedDataType::Int3);   break;
      case Cogs::hash("int4"):   datatype = parameterType(ParsedDataType::Int4);   break;
      }
      src.append("  @location(" + std::to_string(location) + ") ");
      src.append(tokens[0]);
      src.append(" : ");
      src.append(datatype);
      src.append(",\n");
    }
    if (!hasCustomTargets) {
      src.append("  @location(0) fragColor : vec4f,\n");
    }
    src.append("};\n\n");
  }
 
  struct EffectParameterField
  {
    std::string name;
    ParsedDataType type = ParsedDataType::Unknown;
  };
 
  void emitEffectParametersStruct(std::string& src, const std::vector<EffectParameterField>& fields)
  {
    if (fields.empty()) return;
 
    src.append("struct EffectParameters_t {\n");
    for (const EffectParameterField& field : fields) {
      src.append("  ");
      src.append(field.name);
      src.append(" : ");
      src.append(parameterType(field.type));
      src.append(",\n");
    }
    src.append("};\n");
  }
 
  void addTaskUniforms(std::string& uniforms, Cogs::BindGroupSetDescription& bindings, const std::vector<ProcessTaskProperty>& properties)
  {
    constexpr BindGroup group = BindGroup::Default;
    Cogs::BindGroupDescription& defaultGroup = bindings.groups[static_cast<size_t>(group)];
    bindings.numGroups = std::max(bindings.numGroups, static_cast<uint16_t>(static_cast<size_t>(group) + 1));
 
    std::vector<EffectParameterField> parameterFields;
 
    for (const ProcessTaskProperty& p : properties) {
      switch (p.definition->type) {
      case ParsedDataType::Float:
      case ParsedDataType::Float2:
      case ParsedDataType::Float3:
      case ParsedDataType::Float4:
      case ParsedDataType::Float4x4:
      case ParsedDataType::Int:
      case ParsedDataType::Int2:
      case ParsedDataType::Int3:
      case ParsedDataType::Int4:
      case ParsedDataType::UInt:
      case ParsedDataType::UInt2:
      case ParsedDataType::UInt3:
      case ParsedDataType::UInt4:
        parameterFields.push_back({ p.definition->key, p.definition->type });
        break;
 
      case ParsedDataType::Texture2D:
      {
        assert(defaultGroup.numEntries + 2 <= Cogs::MaxBindGroupEntries);
        const bool isDepthTexture = (p.definition->texture.flags & ParsedValueTextureFlags::DepthTexture) != 0;
        if (p.definition->texture.samples > 1) {
          LOG_ERROR(logger, "Sampling from multisampled textures is not implemented for WebGPU post-process effects (%s)", p.definition->key.c_str());
        }
 
        Cogs::BindGroupEntryDescription& textureEntry = defaultGroup.entries[defaultGroup.numEntries];
        textureEntry.binding = defaultGroup.numEntries++;
        textureEntry.nameHash = Cogs::hash(p.definition->key.c_str());
        textureEntry.resourceType = Cogs::BindingResourceType::Texture;
        textureEntry.visibility = Cogs::BindingVisibilityFragment;
        textureEntry.textureDimension = Cogs::ResourceDimensions::Texture2D;
        textureEntry.textureSampleType = bindingTextureSampleType(p.definition->texture.dataType, isDepthTexture);
        textureEntry.isDepthTexture = isDepthTexture;
        addTextureToSource(uniforms, p.definition->key, group, textureEntry);
 
        const std::string samplerName = p.definition->key + "Sampler";
        Cogs::BindGroupEntryDescription& samplerEntry = defaultGroup.entries[defaultGroup.numEntries];
        samplerEntry.binding = defaultGroup.numEntries++;
        samplerEntry.nameHash = Cogs::hash(samplerName.c_str());
        samplerEntry.resourceType = Cogs::BindingResourceType::Sampler;
        samplerEntry.visibility = Cogs::BindingVisibilityFragment;
        samplerEntry.samplerBindingType = isDepthTexture ? Cogs::BindingSamplerBindingType::NonFiltering : Cogs::BindingSamplerBindingType::Filtering;
        samplerEntry.isDepthTexture = isDepthTexture;
        addSamplerToSource(uniforms, samplerName, group, samplerEntry);
        break;
      }
 
      case ParsedDataType::Buffer:
        LOG_ERROR(logger, "Render buffers are not supported for WebGPU post-process effects (%s)", p.definition->key.c_str());
        break;
 
      case ParsedDataType::ConstantBuffer:
        // Engine-provided constant buffers (e.g. a "ConstantBuffer Cogs.SceneBuffer" property)
        // are picked up automatically from shader usage further down, so nothing to do here.
        break;
 
      default:
        break;
      }
    }
 
    if (!parameterFields.empty()) {
      assert(defaultGroup.numEntries < Cogs::MaxBindGroupEntries);
      Cogs::BindGroupEntryDescription& entry = defaultGroup.entries[defaultGroup.numEntries];
      entry.binding = defaultGroup.numEntries++;
      entry.nameHash = Cogs::hash("EffectParameters");
      entry.resourceType = Cogs::BindingResourceType::UniformBuffer;
      entry.visibility = Cogs::BindingVisibilityFragment;
 
      emitEffectParametersStruct(uniforms, parameterFields);
      uniforms.append("@group(");
      uniforms.append(std::to_string(static_cast<size_t>(group)));
      uniforms.append(") @binding(");
      uniforms.append(std::to_string(entry.binding));
      uniforms.append(") var<uniform> EffectParameters : EffectParameters_t;\n");
    }
  }
}
 
namespace Cogs
{
  namespace Core {
 
    bool buildPostProcessEffectWebGPU(RenderTaskContext* context,
      EffectDescription& desc, PostProcessTask* task) {
      // 1. Fetch the engine's bind group layout so post-process effects share the exact same
      //    groups/bindings as the rest of the engine (SceneBuffer, samplers, textures, ...).
      Cogs::BindGroupSetDescription bindings = getEngineBindGroupDescription();
 
      std::string body;
      addInterface(body);
      addDefine(body, "COGS_MAX_LIGHTS", std::to_string(context->renderer->getMaxLights()));
      addDefines(body, desc.definitions);
      addOutputStruct(body, task->options, task->writeDepth || task->depthTest);
 
      // 2. Add the bindings needed by this specific post-process effect into the Default bind
      //    group and emit their declarations.
      std::string uniforms;
      addTaskUniforms(uniforms, bindings, task->properties);
      body.append(uniforms);
 
      std::string prefix_ps = desc.ps;
      removeSuffix(prefix_ps, ".hlsl");
      body.append("#include \"" + prefix_ps + ".wgsl\"\n");
 
      // 3. Run the partial preprocessor to resolve includes/#ifdefs and learn which identifiers
      //    (engine buffers, samplers, textures) the effect's shader source actually references.
      PartialPreprocessor pp;
      pp.processed.reserve(::InitialBufferCapacity);
      if (!pp.process(context->context, body)) return false;
      body.swap(pp.processed);
      pp.processed.clear();
 
      // 4. Fill a source string with the generated engine bindings that were actually referenced.
      std::string engineUniforms;
      std::vector<bool> bindGroupUsed;
      addEngineUniformsToSourceWebGPU(engineUniforms, pp.identifiersSeen, bindGroupUsed);
      if (!engineUniforms.empty()) {
        if (!pp.process(context->context, engineUniforms)) return false;
        engineUniforms.swap(pp.processed);
      }
 
      bindGroupUsed.resize(static_cast<size_t>(BindGroup::Count), false);
      if (bindings.groups[static_cast<size_t>(BindGroup::Default)].numEntries > 0) {
        bindGroupUsed[static_cast<size_t>(BindGroup::Default)] = true;
      }
 
      Cogs::BindGroupSetDescription usedBindings;
      for (size_t g = 0; g < bindGroupUsed.size(); ++g) {
        if (bindGroupUsed[g]) {
          usedBindings.groups[g] = bindings.groups[g];
          usedBindings.numGroups = static_cast<uint16_t>(g + 1);
        }
      }
      desc.bindGroupLayoutDesc = usedBindings;
 
      std::string all = engineUniforms + body;
      all = convertDefinesToConstExpressions(all);
 
      auto code = hash(all);
      std::string file_prefix = prefix_ps + "_" + std::to_string(code);
      context->context->resourceStore->addResource(file_prefix + ".wgsl", all);
      desc.ps = file_prefix + ".hlsl";
      desc.definitions.clear();
      return true;
    }
  }
}