Class KernelCode

gpu.hpp:291–403

* @brief KernelCode is the representation of WGSL GPU code with template * substitutions applied. It is a type around the code string with additional * metadata for workgroup size and precision since they are specified in the * WGSL code. Additionally, label and entryPoint are used by `createKernel()` * to specify the label and entry point of the kernel. */

 * to specify the label and entry point of the kernel.
 */
struct KernelCode {
  /**
   * @brief Constructor to create a code object from a template
   * string and optional workgroup size and precision.
   *
   * @param[in] pData Shader template string with placeholders
   * @param[in] workgroupSize Shape of the workgroup. Unlike tensor shapes which
   * can be of arbitrary rank, workgroup size is always of rank 3 corresponding
   * to x y and z. workgroupSize is stored as a field in the KernelCode instance
   * that is returned by createShader().
   * @param[in] precision Data type precision to be substituted for
   * {{precision}} in the WGSL code. As with workgroupSize, precision is stored
   * as a field in the KernelCode instance that is returned by createShader().
   * @code
   * KernelCode code = {kShaderText, {256, 1, 1}, kf32};
   * @endcode
   */
  inline KernelCode(const std::string &pData = "", size_t workgroupSize = 256,
                    NumType precision = kf32)
      : data(pData), workgroupSize({workgroupSize, 1, 1}),
        precision(precision) {
    if (precision == kf16) {
      data = "enable f16;\n" + data;
    }
    replaceAll(data, "{{workgroupSize}}", toString({workgroupSize, 1, 1}));
    replaceAll(data, "{{precision}}", toString(precision));
    LOG(kDefLog, kTrace, "Shader code:\n%s", data.c_str());
  }

  /**
   * @brief Overload of the constructor to create a code object from a template
   * string and workgroup size. This overload takes a single size_t
   * workgroupSize parameter instead of a 3D shape for the workgroup size and
   * instantiates a 3D shape with the workgroupSize in the x dimension and 1 in
   * the y and z dimensions.
   *
   * @param[in] pData Shader template string with placeholders @param[in]
   * workgroupSize 3D Workgroup size 
   * @param[in] precision Data type precision for the shader
   *
   * @code KernelCode code = {kPuzzle1, 256, kf32}; @endcode
   */
  inline KernelCode(const std::string &pData, const Shape &workgroupSize =
      {256, 1, 1}, NumType precision = kf32) : data(pData),
  workgroupSize(workgroupSize), precision(precision) { if (precision == kf16) {
    data = "enable f16;\n" + data; } replaceAll(data, "{{workgroupSize}}",
        toString(workgroupSize)); replaceAll(data, "{{precision}}",
        toString(precision)); LOG(kDefLog, kInfo, "Shader code:\n%s",
        data.c_str()); }


  /**
   * @brief Overload of the constructor, adding totalWorkgroups parameter to
   * perform a string replacement for the total number of workgroups in the
   * kernel code.
   *
   * @param[in] pData Shader template string with placeholders
   * @param[in] workgroupSize 3D Workgroup size