path: root/source/blender/gpu/metal/mtl_shader_interface.hh

/* SPDX-License-Identifier: GPL-2.0-or-later */

/** \file
 * \ingroup gpu
 */

#pragma once

#include "MEM_guardedalloc.h"

#include "BLI_vector.hh"

#include "gpu_shader_interface.hh"
#include "mtl_capabilities.hh"
#include "mtl_shader_interface_type.hh"

#include "GPU_common.h"
#include "GPU_common_types.h"
#include "GPU_texture.h"
#include "gpu_texture_private.hh"
#include <Metal/Metal.h>
#include <functional>

namespace blender::gpu {

/* MTLShaderInterface describes the layout and properties of a given shader,
 * including input and output bindings, and any special properties or modes
 * that the shader may require.
 *
 * -- Shader input/output bindings --
 *
 * We require custom datastructures for the binding information in Metal.
 * This is because certain bindings contain and require more information to
 * be stored than can be tracked solely within the `ShaderInput` struct.
 * e.g. data sizes and offsets.
 *
 * Upon interface completion, `prepare_common_shader_inputs` is used to
 * populate the global ShaderInput* array to enable correct functionality
 * of shader binding location lookups. These returned locations act as indices
 * into the arrays stored here in the MTLShaderInterace, such that extraction
 * of required information can be performed within the backend.
 *
 * e.g. `int loc = GPU_shader_get_uniform(...)`
 * `loc` will match the index into the MTLShaderUniform uniforms_[] array
 * to fetch the required Metal specific information.
 *
 *
 *
 * -- Argument Buffers and Argument Encoders --
 *
 * We can use ArgumentBuffers (AB's) in Metal to extend the resource bind limitations
 * by providing bindless support.
 *
 * Argument Buffers are used for sampler bindings when the builtin
 * sampler limit of 16 is exceeded, as in all cases for Blender,
 * each individual texture is associated with a given sampler, and this
 * lower limit would otherwise reduce the total availability of textures
 * used in shaders.
 *
 * In future, argument buffers may be extended to support other resource
 * types, if overall bind limits are ever increased within Blender.
 *
 * The ArgumentEncoder cache used to store the generated ArgumentEncoders for a given
 * shader permutation. The ArgumentEncoder is the resource used to write resource binding
 * information to a specified buffer, and is unique to the shader's resource interface.
 */

enum class ShaderStage : uint32_t {
  VERTEX = 1 << 0,
  FRAGMENT = 1 << 1,
  BOTH = (ShaderStage::VERTEX | ShaderStage::FRAGMENT),
};
ENUM_OPERATORS(ShaderStage, ShaderStage::BOTH);

inline uint get_shader_stage_index(ShaderStage stage)
{
  switch (stage) {
    case ShaderStage::VERTEX:
      return 0;
    case ShaderStage::FRAGMENT:
      return 1;
    default:
      BLI_assert_unreachable();
      return 0;
  }
  return 0;
}

/* Shader input/output binding information. */
struct MTLShaderInputAttribute {
  uint32_t name_offset;
  MTLVertexFormat format;
  uint32_t index;
  uint32_t location;
  uint32_t size;
  uint32_t buffer_index;
  uint32_t offset;
  /* For attributes of Matrix/array types, we need to insert "fake" attributes for
   * each element, as matrix types are not natively supported.
   *
   *   > 1 if matrix/arrays are used, specifying number of elements.
   *   = 1 for non-matrix types
   *   = 0 if used as a dummy slot for "fake" matrix attributes. */
  uint32_t matrix_element_count;
};

struct MTLShaderUniformBlock {
  uint32_t name_offset;
  uint32_t size = 0;
  /* Buffer resouce bind index in shader [[buffer(index)]]. */
  uint32_t buffer_index;

  /* Tracking for manual uniform addition. */
  uint32_t current_offset;
  ShaderStage stage_mask;
};

struct MTLShaderUniform {
  uint32_t name_offset;
  /* Index of `MTLShaderUniformBlock` this uniform belongs to. */
  uint32_t size_in_bytes;
  uint32_t byte_offset;
  eMTLDataType type;
  uint32_t array_len;
};

struct MTLShaderTexture {
  bool used;
  uint32_t name_offset;
  /* Texture resource bind slot in shader [[texture(n)]]. */
  int slot_index;
  eGPUTextureType type;
  ShaderStage stage_mask;
};

struct MTLShaderSampler {
  uint32_t name_offset;
  /* Sampler resource bind slot in shader [[sampler(n)]]. */
  uint32_t slot_index = 0;
};

/* Utility Functions. */
MTLVertexFormat mtl_datatype_to_vertex_type(eMTLDataType type);

/**
 * Implementation of Shader interface for Metal Backend.
 **/
class MTLShaderInterface : public ShaderInterface {

 private:
  /* Argument encoders caching.
   * Static size is based on common input permutation variations. */
  static const int ARGUMENT_ENCODERS_CACHE_SIZE = 3;
  struct ArgumentEncoderCacheEntry {
    id<MTLArgumentEncoder> encoder;
    int buffer_index;
  };
  ArgumentEncoderCacheEntry arg_encoders_[ARGUMENT_ENCODERS_CACHE_SIZE] = {};

  /* Vertex input Attribues. */
  uint32_t total_attributes_;
  uint32_t total_vert_stride_;
  MTLShaderInputAttribute attributes_[MTL_MAX_VERTEX_INPUT_ATTRIBUTES];

  /* Uniforms. */
  uint32_t total_uniforms_;
  MTLShaderUniform uniforms_[MTL_MAX_UNIFORMS_PER_BLOCK];

  /* Uniform Blocks. */
  uint32_t total_uniform_blocks_;
  MTLShaderUniformBlock ubos_[MTL_MAX_UNIFORM_BUFFER_BINDINGS];
  MTLShaderUniformBlock push_constant_block_;

  /* Textures. */
  /* Textures support explicit binding indices, so some texture slots
   * remain unused. */
  uint32_t total_textures_;
  int max_texture_index_;
  MTLShaderTexture textures_[MTL_MAX_TEXTURE_SLOTS];

  /* Whether argument buffers are used for sampler bindings. */
  bool sampler_use_argument_buffer_;
  int sampler_argument_buffer_bind_index_vert_;
  int sampler_argument_buffer_bind_index_frag_;

  /* Attribute Mask. */
  uint32_t enabled_attribute_mask_;

  /* Debug. */
  char name[256];

 public:
  MTLShaderInterface(const char *name);
  ~MTLShaderInterface();

  void init();
  void add_input_attribute(uint32_t name_offset,
                           uint32_t attribute_location,
                           MTLVertexFormat format,
                           uint32_t buffer_index,
                           uint32_t size,
                           uint32_t offset,
                           int matrix_element_count = 1);
  uint32_t add_uniform_block(uint32_t name_offset,
                             uint32_t buffer_index,
                             uint32_t size,
                             ShaderStage stage_mask = ShaderStage::BOTH);
  void add_uniform(uint32_t name_offset, eMTLDataType type, int array_len = 1);
  void add_texture(uint32_t name_offset,
                   uint32_t texture_slot,
                   eGPUTextureType tex_binding_type,
                   ShaderStage stage_mask = ShaderStage::FRAGMENT);
  void add_push_constant_block(uint32_t name_offset);

  /* Resolve and cache locations of builtin uniforms and uniform blocks. */
  void map_builtins();
  void set_sampler_properties(bool use_argument_buffer,
                              uint32_t argument_buffer_bind_index_vert,
                              uint32_t argument_buffer_bind_index_frag);

  /* Prepare ShaderInput interface for binding resolution. */
  void prepare_common_shader_inputs();

  /* Fetch Uniforms. */
  const MTLShaderUniform &get_uniform(uint index) const;
  uint32_t get_total_uniforms() const;

  /* Fetch Uniform Blocks. */
  const MTLShaderUniformBlock &get_uniform_block(uint index) const;
  uint32_t get_total_uniform_blocks() const;
  bool has_uniform_block(uint32_t block_index) const;
  uint32_t get_uniform_block_size(uint32_t block_index) const;

  /* Push constant uniform data block should always be available. */
  const MTLShaderUniformBlock &get_push_constant_block() const;

  /* Fetch textures. */
  const MTLShaderTexture &get_texture(uint index) const;
  uint32_t get_total_textures() const;
  uint32_t get_max_texture_index() const;
  bool get_use_argument_buffer_for_samplers(int *vertex_arg_buffer_bind_index,
                                            int *fragment_arg_buffer_bind_index) const;

  /* Fetch Attributes. */
  const MTLShaderInputAttribute &get_attribute(uint index) const;
  uint32_t get_total_attributes() const;
  uint32_t get_total_vertex_stride() const;
  uint32_t get_enabled_attribute_mask() const;

  /* Name buffer fetching. */
  const char *get_name_at_offset(uint32_t offset) const;

  /* Interface name. */
  const char *get_name() const
  {
    return this->name;
  }

  /* Argument buffer encoder management. */
  id<MTLArgumentEncoder> find_argument_encoder(int buffer_index) const;

  void insert_argument_encoder(int buffer_index, id encoder);

  MEM_CXX_CLASS_ALLOC_FUNCS("MTLShaderInterface");
};

}  // namespace blender::gpu