src/amd/vulkan/radv_nir_lower_ycbcr_textures.c

   1 /*
   2  * Copyright © 2017 Intel Corporation
   3  *
   4  * Permission is hereby granted, free of charge, to any person obtaining a
   5  * copy of this software and associated documentation files (the "Software"),
   6  * to deal in the Software without restriction, including without limitation
   7  * the rights to use, copy, modify, merge, publish, distribute, sublicense,
   8  * and/or sell copies of the Software, and to permit persons to whom the
   9  * Software is furnished to do so, subject to the following conditions:
  10  *
  11  * The above copyright notice and this permission notice (including the next
  12  * paragraph) shall be included in all copies or substantial portions of the
  13  * Software.
  14  *
  15  * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
  16  * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
  17  * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
  18  * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
  19  * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
  20  * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
  21  * IN THE SOFTWARE.
  22  */
  23
  24 #include "radv_private.h"
  25 #include "radv_shader.h"
  26 #include "vk_format.h"
  27 #include "nir/nir.h"
  28 #include "nir/nir_builder.h"
  29 #include "nir/nir_vulkan.h"
  30
  31 struct ycbcr_state {
  32         nir_builder *builder;
  33         nir_ssa_def *image_size;
  34         nir_tex_instr *origin_tex;
  35         nir_deref_instr *tex_deref;
  36         const struct radv_sampler_ycbcr_conversion *conversion;
  37 };
  38
  39 static nir_ssa_def *
  40 get_texture_size(struct ycbcr_state *state, nir_deref_instr *texture)
  41 {
  42         nir_builder *b = state->builder;
  43         const struct glsl_type *type = texture->type;
  44         nir_tex_instr *tex = nir_tex_instr_create(b->shader, 1);
  45
  46         tex->op = nir_texop_txs;
  47         tex->sampler_dim = glsl_get_sampler_dim(type);
  48         tex->is_array = glsl_sampler_type_is_array(type);
  49         tex->is_shadow = glsl_sampler_type_is_shadow(type);
  50         tex->dest_type = nir_type_int;
  51
  52         tex->src[0].src_type = nir_tex_src_texture_deref;
  53         tex->src[0].src = nir_src_for_ssa(&texture->dest.ssa);
  54
  55         nir_ssa_dest_init(&tex->instr, &tex->dest,
  56                           nir_tex_instr_dest_size(tex), 32, NULL);
  57         nir_builder_instr_insert(b, &tex->instr);
  58
  59         return nir_i2f32(b, &tex->dest.ssa);
  60 }
  61
  62 static nir_ssa_def *
  63 implicit_downsampled_coord(nir_builder *b,
  64                            nir_ssa_def *value,
  65                            nir_ssa_def *max_value,
  66                            int div_scale)
  67 {
  68         return nir_fadd(b,
  69                         value,
  70                         nir_fdiv(b,
  71                                  nir_imm_float(b, 1.0f),
  72                                  nir_fmul(b,
  73                                           nir_imm_float(b, div_scale),
  74                                           max_value)));
  75 }
  76
  77 static nir_ssa_def *
  78 implicit_downsampled_coords(struct ycbcr_state *state,
  79                             nir_ssa_def *old_coords)
  80 {
  81         nir_builder *b = state->builder;
  82         const struct radv_sampler_ycbcr_conversion *conversion = state->conversion;
  83         nir_ssa_def *image_size = NULL;
  84         nir_ssa_def *comp[4] = { NULL, };
  85         const struct vk_format_description *fmt_desc = vk_format_description(state->conversion->format);
  86         const unsigned divisors[2] = {fmt_desc->width_divisor, fmt_desc->height_divisor};
  87
  88         for (int c = 0; c < old_coords->num_components; c++) {
  89                 if (c < ARRAY_SIZE(divisors) && divisors[c] > 1 &&
  90                     conversion->chroma_offsets[c] == VK_CHROMA_LOCATION_COSITED_EVEN) {
  91                         if (!image_size)
  92                                 image_size = get_texture_size(state, state->tex_deref);
  93
  94                         comp[c] = implicit_downsampled_coord(b,
  95                                                              nir_channel(b, old_coords, c),
  96                                                              nir_channel(b, image_size, c),
  97                                                              divisors[c]);
  98                 } else {
  99                         comp[c] = nir_channel(b, old_coords, c);
 100                 }
 101         }
 102
 103         return nir_vec(b, comp, old_coords->num_components);
 104 }
 105
 106 static nir_ssa_def *
 107 create_plane_tex_instr_implicit(struct ycbcr_state *state,
 108                                 uint32_t plane)
 109 {
 110         nir_builder *b = state->builder;
 111         nir_tex_instr *old_tex = state->origin_tex;
 112         nir_tex_instr *tex = nir_tex_instr_create(b->shader, old_tex->num_srcs+ 1);
 113         for (uint32_t i = 0; i < old_tex->num_srcs; i++) {
 114                 tex->src[i].src_type = old_tex->src[i].src_type;
 115
 116                 switch (old_tex->src[i].src_type) {
 117                 case nir_tex_src_coord:
 118                         if (plane && true/*state->conversion->chroma_reconstruction*/) {
 119                                 assert(old_tex->src[i].src.is_ssa);
 120                                 tex->src[i].src =
 121                                         nir_src_for_ssa(implicit_downsampled_coords(state,
 122                                                                                     old_tex->src[i].src.ssa));
 123                                 break;
 124                         }
 125                 /* fall through */
 126                 default:
 127                         nir_src_copy(&tex->src[i].src, &old_tex->src[i].src, tex);
 128                         break;
 129                 }
 130         }
 131
 132         tex->src[tex->num_srcs - 1].src = nir_src_for_ssa(nir_imm_int(b, plane));
 133         tex->src[tex->num_srcs - 1].src_type = nir_tex_src_plane;
 134
 135         tex->sampler_dim = old_tex->sampler_dim;
 136         tex->dest_type = old_tex->dest_type;
 137         tex->is_array = old_tex->is_array;
 138
 139         tex->op = old_tex->op;
 140         tex->coord_components = old_tex->coord_components;
 141         tex->is_new_style_shadow = old_tex->is_new_style_shadow;
 142         tex->component = old_tex->component;
 143
 144         tex->texture_index = old_tex->texture_index;
 145         tex->sampler_index = old_tex->sampler_index;
 146
 147         nir_ssa_dest_init(&tex->instr, &tex->dest,
 148                           old_tex->dest.ssa.num_components,
 149                           nir_dest_bit_size(old_tex->dest), NULL);
 150         nir_builder_instr_insert(b, &tex->instr);
 151
 152         return &tex->dest.ssa;
 153 }
 154
 155 struct swizzle_info {
 156         unsigned plane[4];
 157         unsigned swizzle[4];
 158 };
 159
 160 static struct swizzle_info
 161 get_plane_swizzles(VkFormat format)
 162 {
 163         int planes = vk_format_get_plane_count(format);
 164         switch (planes) {
 165         case 3:
 166                 return (struct swizzle_info) {
 167                         {2, 0, 1, 0},
 168                         {0, 0, 0, 3}
 169                 };
 170         case 2:
 171                 return (struct swizzle_info) {
 172                         {1, 0, 1, 0},
 173                         {1, 0, 0, 3}
 174                 };
 175         case 1:
 176                 return (struct swizzle_info) {
 177                         {0, 0, 0, 0},
 178                         {0, 1, 2, 3}
 179                 };
 180         default:
 181                 unreachable("unhandled plane count for ycbcr swizzling");
 182         }
 183 }
 184
 185
 186 static nir_ssa_def *
 187 build_swizzled_components(nir_builder *builder,
 188                           VkFormat format,
 189                           VkComponentMapping mapping,
 190                           nir_ssa_def **plane_values)
 191 {
 192         struct swizzle_info plane_swizzle = get_plane_swizzles(format);
 193         enum vk_swizzle swizzles[4];
 194         nir_ssa_def *values[4];
 195
 196         vk_format_compose_swizzles(&mapping, (const unsigned char[4]){0,1,2,3}, swizzles);
 197
 198         nir_ssa_def *zero = nir_imm_float(builder, 0.0f);
 199         nir_ssa_def *one = nir_imm_float(builder, 1.0f);
 200
 201         for (unsigned i = 0; i < 4; ++i) {
 202                 switch(swizzles[i]) {
 203                 case VK_SWIZZLE_X:
 204                 case VK_SWIZZLE_Y:
 205                 case VK_SWIZZLE_Z:
 206                 case VK_SWIZZLE_W: {
 207                         unsigned channel = swizzles[i] - VK_SWIZZLE_X;
 208                         values[i] = nir_channel(builder,
 209                                                 plane_values[plane_swizzle.plane[channel]],
 210                                                 plane_swizzle.swizzle[channel]);
 211                         break;
 212                 }
 213                 case VK_SWIZZLE_0:
 214                         values[i] = zero;
 215                         break;
 216                 case VK_SWIZZLE_1:
 217                         values[i] = one;
 218                         break;
 219                 default:
 220                         unreachable("unhandled swizzle");
 221                 }
 222         }
 223         return nir_vec(builder, values, 4);
 224 }
 225
 226 static bool
 227 try_lower_tex_ycbcr(const struct radv_pipeline_layout *layout,
 228                     nir_builder *builder,
 229                     nir_tex_instr *tex)
 230 {
 231         int deref_src_idx = nir_tex_instr_src_index(tex, nir_tex_src_texture_deref);
 232         assert(deref_src_idx >= 0);
 233         nir_deref_instr *deref = nir_src_as_deref(tex->src[deref_src_idx].src);
 234
 235         nir_variable *var = nir_deref_instr_get_variable(deref);
 236         const struct radv_descriptor_set_layout *set_layout =
 237                 layout->set[var->data.descriptor_set].layout;
 238         const struct radv_descriptor_set_binding_layout *binding =
 239                 &set_layout->binding[var->data.binding];
 240         const struct radv_sampler_ycbcr_conversion *ycbcr_samplers =
 241                 radv_immutable_ycbcr_samplers(set_layout, var->data.binding);
 242
 243         if (!ycbcr_samplers)
 244                 return false;
 245
 246         /* For the following instructions, we don't apply any change and let the
 247          * instruction apply to the first plane.
 248          */
 249         if (tex->op == nir_texop_txs ||
 250             tex->op == nir_texop_query_levels ||
 251             tex->op == nir_texop_lod)
 252                 return false;
 253
 254         assert(tex->texture_index == 0);
 255         unsigned array_index = 0;
 256         if (deref->deref_type != nir_deref_type_var) {
 257                 assert(deref->deref_type == nir_deref_type_array);
 258                 if (!nir_src_is_const(deref->arr.index))
 259                         return false;
 260                 array_index = nir_src_as_uint(deref->arr.index);
 261                 array_index = MIN2(array_index, binding->array_size - 1);
 262         }
 263         const struct radv_sampler_ycbcr_conversion *ycbcr_sampler = ycbcr_samplers + array_index;
 264
 265         if (ycbcr_sampler->format == VK_FORMAT_UNDEFINED)
 266                 return false;
 267
 268         struct ycbcr_state state = {
 269                 .builder = builder,
 270                 .origin_tex = tex,
 271                 .tex_deref = deref,
 272                 .conversion = ycbcr_sampler,
 273         };
 274
 275         builder->cursor = nir_before_instr(&tex->instr);
 276
 277         VkFormat format = state.conversion->format;
 278         const int plane_count = vk_format_get_plane_count(format);
 279         nir_ssa_def *plane_values[3];
 280
 281         for (int p = 0; p < plane_count; ++p) {
 282                 plane_values[p] = create_plane_tex_instr_implicit(&state, p);
 283         }
 284
 285         nir_ssa_def *result = build_swizzled_components(builder, format, ycbcr_sampler->components, plane_values);
 286         if (state.conversion->ycbcr_model != VK_SAMPLER_YCBCR_MODEL_CONVERSION_RGB_IDENTITY) {
 287                 VkFormat first_format = vk_format_get_plane_format(format, 0);
 288                 uint32_t bits = vk_format_get_component_bits(first_format, VK_FORMAT_COLORSPACE_RGB, VK_SWIZZLE_X);
 289                 /* TODO: swizzle and bpcs */
 290                 uint32_t bpcs[3] = {bits, bits, bits};
 291                 result = nir_convert_ycbcr_to_rgb(builder,
 292                                                   state.conversion->ycbcr_model,
 293                                                   state.conversion->ycbcr_range,
 294                                                   result,
 295                                                   bpcs);
 296         }
 297
 298         nir_ssa_def_rewrite_uses(&tex->dest.ssa, nir_src_for_ssa(result));
 299         nir_instr_remove(&tex->instr);
 300
 301         return true;
 302 }
 303
 304 bool
 305 radv_nir_lower_ycbcr_textures(nir_shader *shader,
 306                              const struct radv_pipeline_layout *layout)
 307 {
 308         bool progress = false;
 309
 310         nir_foreach_function(function, shader) {
 311                 if (!function->impl)
 312                         continue;
 313
 314                 bool function_progress = false;
 315                 nir_builder builder;
 316                 nir_builder_init(&builder, function->impl);
 317
 318                 nir_foreach_block(block, function->impl) {
 319                         nir_foreach_instr_safe(instr, block) {
 320                                 if (instr->type != nir_instr_type_tex)
 321                                         continue;
 322
 323                                 nir_tex_instr *tex = nir_instr_as_tex(instr);
 324                                 function_progress |= try_lower_tex_ycbcr(layout, &builder, tex);
 325                         }
 326                 }
 327
 328                 if (function_progress) {
 329                         nir_metadata_preserve(function->impl,
 330                                               nir_metadata_block_index |
 331                                               nir_metadata_dominance);
 332                 }
 333
 334                 progress |= function_progress;
 335         }
 336
 337         return progress;
 338 }