Seele/src/Engine/Graphics/Metal/PipelineCache.mm

#include "PipelineCache.h"
#include "Descriptor.h"
#include "Enums.h"
#include "Foundation/NSError.hpp"
#include "Foundation/NSString.hpp"
#include "Graphics.h"
#include "Graphics/Descriptor.h"
#include "Graphics/Enums.h"
#include "Graphics/Metal/Pipeline.h"
#include "Metal/MTLComputePipeline.hpp"
#include "Metal/MTLDevice.hpp"
#include "Metal/MTLRenderCommandEncoder.hpp"
#include "Metal/MTLRenderPipeline.hpp"
#include "Metal/MTLVertexDescriptor.hpp"
#include "Shader.h"
#include "Texture.h"
#include <iostream>

using namespace Seele;
using namespace Seele::Metal;

PipelineCache::PipelineCache(PGraphics graphics, const std::string &name)
    : graphics(graphics), cacheFile(name) {}

PipelineCache::~PipelineCache() {}

PGraphicsPipeline
PipelineCache::createPipeline(Gfx::LegacyPipelineCreateInfo createInfo) {
  PPipelineLayout layout =
      Gfx::PPipelineLayout(createInfo.pipelineLayout).cast<PipelineLayout>();

  MTL::RenderPipelineDescriptor *pipelineDescriptor =
      MTL::RenderPipelineDescriptor::alloc()->init();

  MTL::VertexDescriptor *vertexDescriptor =
      pipelineDescriptor->vertexDescriptor()->init();
  MTL::VertexAttributeDescriptorArray *attributes =
      vertexDescriptor->attributes()->init();
  if (createInfo.vertexInput != nullptr) {
    const auto &vertexInfo = createInfo.vertexInput->getInfo();
    for (size_t attr = 0; attr < vertexInfo.attributes.size(); ++attr) {
      MTL::VertexAttributeDescriptor *attribute =
          attributes->object(attr + METAL_VERTEXATTRIBUTE_OFFSET)->init();
      attribute->setBufferIndex(vertexInfo.attributes[attr].binding +
                                METAL_VERTEXBUFFER_OFFSET);
      switch (vertexInfo.attributes[attr].format) {
      case Gfx::SE_FORMAT_R32G32B32_SFLOAT:
        attribute->setFormat(MTL::VertexFormatFloat3);
        break;
      default:
        throw std::logic_error("TODO");
      }
      attribute->setOffset(vertexInfo.attributes[attr].offset);
    }

    MTL::VertexBufferLayoutDescriptorArray *bufferLayout =
        vertexDescriptor->layouts()->init();
    for (size_t binding = 0; binding < vertexInfo.bindings.size(); ++binding) {
      MTL::VertexBufferLayoutDescriptor *buffer =
          bufferLayout->object(binding + METAL_VERTEXBUFFER_OFFSET)->init();
      buffer->setStride(vertexInfo.bindings[binding].stride);
      buffer->setStepRate(1);
      switch (vertexInfo.bindings[binding].inputRate) {
      case Gfx::SE_VERTEX_INPUT_RATE_VERTEX:
        buffer->setStepFunction(MTL::VertexStepFunctionPerVertex);
        break;
      case Gfx::SE_VERTEX_INPUT_RATE_INSTANCE:
        buffer->setStepFunction(MTL::VertexStepFunctionPerInstance);
        break;
      }
    }
  }
  pipelineDescriptor->setVertexDescriptor(vertexDescriptor);

  pipelineDescriptor->setVertexFunction(
      createInfo.vertexShader.cast<VertexShader>()->getFunction());
  if (createInfo.fragmentShader != nullptr) {
    pipelineDescriptor->setFragmentFunction(
        createInfo.fragmentShader.cast<FragmentShader>()->getFunction());
  }
  pipelineDescriptor->setInputPrimitiveTopology(cast(createInfo.topology));
  if (createInfo.renderPass->getLayout().depthAttachment.getTexture() !=
      nullptr) {
    pipelineDescriptor->setDepthAttachmentPixelFormat(
        cast(createInfo.renderPass->getLayout()
                 .depthAttachment.getTexture()
                 .cast<Texture2D>()
                 ->getFormat()));
  }
  pipelineDescriptor->setAlphaToCoverageEnabled(
      createInfo.multisampleState.alphaCoverageEnable);
  pipelineDescriptor->setAlphaToOneEnabled(
      createInfo.multisampleState.alphaToOneEnable);
  pipelineDescriptor->setRasterSampleCount(createInfo.multisampleState.samples);
  pipelineDescriptor->setRasterizationEnabled(
      !createInfo.rasterizationState.rasterizerDiscardEnable);

  uint32 hash = pipelineDescriptor->hash();

  if (graphicsPipelines.contains(hash)) {
    return graphicsPipelines[hash];
  }
  MTL::PrimitiveType type = MTL::PrimitiveTypeTriangle;
  switch (createInfo.topology) {
  case Gfx::SE_PRIMITIVE_TOPOLOGY_POINT_LIST:
    type = MTL::PrimitiveTypePoint;
    break;
  case Gfx::SE_PRIMITIVE_TOPOLOGY_LINE_LIST:
    type = MTL::PrimitiveTypeLine;
    break;
  case Gfx::SE_PRIMITIVE_TOPOLOGY_LINE_STRIP:
    type = MTL::PrimitiveTypeLineStrip;
    break;
  case Gfx::SE_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST:
    type = MTL::PrimitiveTypeTriangle;
    break;
  case Gfx::SE_PRIMITIVE_TOPOLOGY_TRIANGLE_STRIP:
    type = MTL::PrimitiveTypeTriangleStrip;
    break;
  case Gfx::SE_PRIMITIVE_TOPOLOGY_TRIANGLE_FAN:
    type = MTL::PrimitiveTypeTriangle;
    break;
  case Gfx::SE_PRIMITIVE_TOPOLOGY_LINE_LIST_WITH_ADJACENCY:
    type = MTL::PrimitiveTypeLine;
    break;
  case Gfx::SE_PRIMITIVE_TOPOLOGY_LINE_STRIP_WITH_ADJACENCY:
    type = MTL::PrimitiveTypeLineStrip;
    break;
  case Gfx::SE_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST_WITH_ADJACENCY:
    type = MTL::PrimitiveTypeTriangle;
    break;
  case Gfx::SE_PRIMITIVE_TOPOLOGY_TRIANGLE_STRIP_WITH_ADJACENCY:
    type = MTL::PrimitiveTypeTriangleStrip;
    break;
  case Gfx::SE_PRIMITIVE_TOPOLOGY_PATCH_LIST:
    type = MTL::PrimitiveTypeTriangle;
    break;
  }
  NS::Error *error;
  graphicsPipelines[hash] = new GraphicsPipeline(
      graphics, type,
      graphics->getDevice()->newRenderPipelineState(pipelineDescriptor, &error),
      std::move(createInfo.pipelineLayout));
  if (error) {
    std::cout << error->localizedDescription()->cString(NS::ASCIIStringEncoding)
              << std::endl;
    assert(false);
  }

  pipelineDescriptor->release();
  return graphicsPipelines[hash];
}

PGraphicsPipeline
PipelineCache::createPipeline(Gfx::MeshPipelineCreateInfo createInfo) {
  MTL::MeshRenderPipelineDescriptor *pipelineDescriptor =
      MTL::MeshRenderPipelineDescriptor::alloc()->init();

  pipelineDescriptor->setMeshFunction(
      createInfo.meshShader.cast<MeshShader>()->getFunction());
  if (createInfo.taskShader != nullptr) {
    pipelineDescriptor->setObjectFunction(
        createInfo.taskShader.cast<TaskShader>()->getFunction());
  }
  if (createInfo.fragmentShader != nullptr) {
    pipelineDescriptor->setFragmentFunction(
        createInfo.fragmentShader.cast<FragmentShader>()->getFunction());
  }
  if (createInfo.renderPass->getLayout().depthAttachment.getTexture() !=
      nullptr) {
    pipelineDescriptor->setDepthAttachmentPixelFormat(
        cast(createInfo.renderPass->getLayout()
                 .depthAttachment.getTexture()
                 .cast<Texture2D>()
                 ->getFormat()));
  }
  pipelineDescriptor->setAlphaToCoverageEnabled(
      createInfo.multisampleState.alphaCoverageEnable);
  pipelineDescriptor->setAlphaToOneEnabled(
      createInfo.multisampleState.alphaToOneEnable);
  pipelineDescriptor->setRasterSampleCount(createInfo.multisampleState.samples);
  pipelineDescriptor->setRasterizationEnabled(
      !createInfo.rasterizationState.rasterizerDiscardEnable);

  uint32 hash = pipelineDescriptor->hash();

  if (graphicsPipelines.contains(hash)) {
    return graphicsPipelines[hash];
  }
  NS::Error *error = nullptr;
  MTL::AutoreleasedRenderPipelineReflection reflection;
  graphicsPipelines[hash] = new GraphicsPipeline(
      graphics, MTL::PrimitiveTypeTriangle,
      graphics->getDevice()->newRenderPipelineState(
          pipelineDescriptor, MTL::PipelineOptionNone, &reflection, &error),
      std::move(createInfo.pipelineLayout));

  pipelineDescriptor->release();
  return graphicsPipelines[hash];
}

PComputePipeline
PipelineCache::createPipeline(Gfx::ComputePipelineCreateInfo createInfo) {
  PComputeShader shader = createInfo.computeShader.cast<ComputeShader>();
  uint32 hash = shader->getShaderHash();
  if (computePipelines.contains(hash)) {
    return computePipelines[hash];
  }

  NS::Error *error;
  computePipelines[hash] =
      new ComputePipeline(graphics,
                          graphics->getDevice()->newComputePipelineState(
                              shader->getFunction(), &error),
                          std::move(createInfo.pipelineLayout));
  assert(!error);
  return computePipelines[hash];
}