<commit_before>#include <metal_stdlib>
using namespace metal;

// -------------- Image Clears -------------- //

typedef struct {
    float4 coords [[attribute(0)]];
} ClearAttributes;

typedef struct {
    float4 position [[position]];
    uint layer [[render_target_array_index]];
} ClearVertexData;

vertex ClearVertexData vs_clear(ClearAttributes in [[stage_in]]) {
    float4 pos = { 0.0, 0.0, 0.0f, 1.0f };
    pos.xy = in.coords.xy * 2.0 - 1.0;
    return ClearVertexData { pos, uint(in.coords.z) };
}

fragment float4 ps_blit_float(
    ClearVertexData in [[stage_in]],
    constant float4 &value [[ buffer(0) ]]
) {
  return value;
}

fragment int4 ps_blit_int(
    ClearVertexData in [[stage_in]],
    constant int4 &value [[ buffer(0) ]]
) {
  return value;
}

fragment uint4 ps_blit_uint(
    ClearVertexData in [[stage_in]],
    constant uint4 &value [[ buffer(0) ]]
) {
  return value;
}

// -------------- Image Blits -------------- //

typedef struct {
    float4 src_coords [[attribute(0)]];
    float4 dst_coords [[attribute(1)]];
} BlitAttributes;

typedef struct {
    float4 position [[position]];
    float4 uv;
    uint layer [[render_target_array_index]];
} BlitVertexData;

vertex BlitVertexData vs_blit(BlitAttributes in [[stage_in]]) {
    float4 pos = { 0.0, 0.0, 0.0f, 1.0f };
    pos.xy = in.dst_coords.xy * 2.0 - 1.0;
    return BlitVertexData { pos, in.src_coords, uint(in.dst_coords.z) };
}

fragment float4 ps_blit_1d_float(
    BlitVertexData in [[stage_in]],
    texture1d<float> tex1D [[ texture(0) ]],
    sampler sampler2D [[ sampler(0) ]]
) {
  return tex1D.sample(sampler2D, in.uv.x);
}

fragment float4 ps_blit_1d_array_float(
    BlitVertexData in [[stage_in]],
    texture1d_array<float> tex1DArray [[ texture(0) ]],
    sampler sampler2D [[ sampler(0) ]]
) {
  return tex1DArray.sample(sampler2D, in.uv.x, uint(in.uv.z));
}

fragment float4 ps_blit_2d_float(
    BlitVertexData in [[stage_in]],
    texture2d<float> tex2D [[ texture(0) ]],
    sampler sampler2D [[ sampler(0) ]]
) {
  return tex2D.sample(sampler2D, in.uv.xy, level(in.uv.w));
}

fragment uint4 ps_blit_2d_uint(
    BlitVertexData in [[stage_in]],
    texture2d<uint> tex2D [[ texture(0) ]],
    sampler sampler2D [[ sampler(0) ]]
) {
  return tex2D.sample(sampler2D, in.uv.xy, level(in.uv.w));
}

fragment int4 ps_blit_2d_int(
    BlitVertexData in [[stage_in]],
    texture2d<int> tex2D [[ texture(0) ]],
    sampler sampler2D [[ sampler(0) ]]
) {
  return tex2D.sample(sampler2D, in.uv.xy, level(in.uv.w));
}

fragment float4 ps_blit_2d_array_float(
    BlitVertexData in [[stage_in]],
    texture2d_array<float> tex2DArray [[ texture(0) ]],
    sampler sampler2D [[ sampler(0) ]]
) {
  return tex2DArray.sample(sampler2D, in.uv.xy, uint(in.uv.z), level(in.uv.w));
}
fragment float4 ps_blit_3d_float(
    BlitVertexData in [[stage_in]],
    texture3d<float> tex3D [[ texture(0) ]],
    sampler sampler2D [[ sampler(0) ]]
) {<commit_msg>[mtl] 2darray versions of int/uint blits<commit_after>#include <metal_stdlib>
using namespace metal;

// -------------- Image Clears -------------- //

typedef struct {
    float4 coords [[attribute(0)]];
} ClearAttributes;

typedef struct {
    float4 position [[position]];
    uint layer [[render_target_array_index]];
} ClearVertexData;

vertex ClearVertexData vs_clear(ClearAttributes in [[stage_in]]) {
    float4 pos = { 0.0, 0.0, 0.0f, 1.0f };
    pos.xy = in.coords.xy * 2.0 - 1.0;
    return ClearVertexData { pos, uint(in.coords.z) };
}

fragment float4 ps_blit_float(
    ClearVertexData in [[stage_in]],
    constant float4 &value [[ buffer(0) ]]
) {
  return value;
}

fragment int4 ps_blit_int(
    ClearVertexData in [[stage_in]],
    constant int4 &value [[ buffer(0) ]]
) {
  return value;
}

fragment uint4 ps_blit_uint(
    ClearVertexData in [[stage_in]],
    constant uint4 &value [[ buffer(0) ]]
) {
  return value;
}

// -------------- Image Blits -------------- //

typedef struct {
    float4 src_coords [[attribute(0)]];
    float4 dst_coords [[attribute(1)]];
} BlitAttributes;

typedef struct {
    float4 position [[position]];
    float4 uv;
    uint layer [[render_target_array_index]];
} BlitVertexData;

vertex BlitVertexData vs_blit(BlitAttributes in [[stage_in]]) {
    float4 pos = { 0.0, 0.0, 0.0f, 1.0f };
    pos.xy = in.dst_coords.xy * 2.0 - 1.0;
    return BlitVertexData { pos, in.src_coords, uint(in.dst_coords.z) };
}

fragment float4 ps_blit_1d_float(
    BlitVertexData in [[stage_in]],
    texture1d<float> tex1D [[ texture(0) ]],
    sampler sampler2D [[ sampler(0) ]]
) {
  return tex1D.sample(sampler2D, in.uv.x);
}

fragment float4 ps_blit_1d_array_float(
    BlitVertexData in [[stage_in]],
    texture1d_array<float> tex1DArray [[ texture(0) ]],
    sampler sampler2D [[ sampler(0) ]]
) {
  return tex1DArray.sample(sampler2D, in.uv.x, uint(in.uv.z));
}

fragment float4 ps_blit_2d_float(
    BlitVertexData in [[stage_in]],
    texture2d<float> tex2D [[ texture(0) ]],
    sampler sampler2D [[ sampler(0) ]]
) {
  return tex2D.sample(sampler2D, in.uv.xy, level(in.uv.w));
}

fragment uint4 ps_blit_2d_uint(
    BlitVertexData in [[stage_in]],
    texture2d<uint> tex2D [[ texture(0) ]],
    sampler sampler2D [[ sampler(0) ]]
) {
  return tex2D.sample(sampler2D, in.uv.xy, level(in.uv.w));
}

fragment int4 ps_blit_2d_int(
    BlitVertexData in [[stage_in]],
    texture2d<int> tex2D [[ texture(0) ]],
    sampler sampler2D [[ sampler(0) ]]
) {
  return tex2D.sample(sampler2D, in.uv.xy, level(in.uv.w));
}

fragment float4 ps_blit_2d_array_float(
    BlitVertexData in [[stage_in]],
    texture2d_array<float> tex2DArray [[ texture(0) ]],
    sampler sampler2D [[ sampler(0) ]]
) {
  return tex2DArray.sample(sampler2D, in.uv.xy, uint(in.uv.z), level(in.uv.w));
}

fragment uint4 ps_blit_2d_array_uint(
    BlitVertexData in [[stage_in]],
    texture2d_array<uint> tex2DArray [[ texture(0) ]],
    sampler sampler2D [[ sampler(0) ]]
) {
  return tex2DArray.sample(sampler2D, in.uv.xy, uint(in.uv.z), level(in.uv.w));
}

fragment int4 ps_blit_2d_array_int(
    BlitVertexData in [[stage_in]],
    texture2d_array<int> tex2DArray [[ texture(0) ]],
    sampler sampler2D [[ sampler(0) ]]
) {
  return tex2DArray.sample(sampler2D, in.uv.xy, uint(in.uv.z), level(in.uv.w));
}

fragment float4 ps_blit_3d_float(
    BlitVertexData in [[stage_in]],
    texture3d<float> tex3D [[ texture(0) ]],
    sampler sampler2D [[ sampler(0) ]]
) {
{'fraction_non_alphanumeric': 0.1651903651903652, 'fraction_numerical': 0.035586635586635586, 'mean_word_length': 3.6435786435786435, 'pattern_counts': {'":': 0, '<': 21, '<?xml version=': 0, '>': 21, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2965782', 'n_tokens_mistral': 2560, 'n_tokens_neox': 2507, 'n_words': 493}