nir: Add a ptr_as_array deref type
[mesa.git] / src / compiler / spirv / vtn_cfg.c
index 25ff254bcec5b0daf8aad99c6dfd150598ae8c9a..9cec036d46217f76c5056bd65d80e62f9754f7c3 100644 (file)
 #include "vtn_private.h"
 #include "nir/nir_vla.h"
 
+static struct vtn_pointer *
+vtn_load_param_pointer(struct vtn_builder *b,
+                       struct vtn_type *param_type,
+                       uint32_t param_idx)
+{
+   struct vtn_type *ptr_type = param_type;
+   if (param_type->base_type != vtn_base_type_pointer) {
+      assert(param_type->base_type == vtn_base_type_image ||
+             param_type->base_type == vtn_base_type_sampler);
+      ptr_type = rzalloc(b, struct vtn_type);
+      ptr_type->base_type = vtn_base_type_pointer;
+      ptr_type->deref = param_type;
+      ptr_type->storage_class = SpvStorageClassUniformConstant;
+   }
+
+   return vtn_pointer_from_ssa(b, nir_load_param(&b->nb, param_idx), ptr_type);
+}
+
+static unsigned
+vtn_type_count_function_params(struct vtn_type *type)
+{
+   switch (type->base_type) {
+   case vtn_base_type_array:
+   case vtn_base_type_matrix:
+      return type->length * vtn_type_count_function_params(type->array_element);
+
+   case vtn_base_type_struct: {
+      unsigned count = 0;
+      for (unsigned i = 0; i < type->length; i++)
+         count += vtn_type_count_function_params(type->members[i]);
+      return count;
+   }
+
+   case vtn_base_type_sampled_image:
+      return 2;
+
+   default:
+      return 1;
+   }
+}
+
+static void
+vtn_type_add_to_function_params(struct vtn_type *type,
+                                nir_function *func,
+                                unsigned *param_idx)
+{
+   static const nir_parameter nir_deref_param = {
+      .num_components = 1,
+      .bit_size = 32,
+   };
+
+   switch (type->base_type) {
+   case vtn_base_type_array:
+   case vtn_base_type_matrix:
+      for (unsigned i = 0; i < type->length; i++)
+         vtn_type_add_to_function_params(type->array_element, func, param_idx);
+      break;
+
+   case vtn_base_type_struct:
+      for (unsigned i = 0; i < type->length; i++)
+         vtn_type_add_to_function_params(type->members[i], func, param_idx);
+      break;
+
+   case vtn_base_type_sampled_image:
+      func->params[(*param_idx)++] = nir_deref_param;
+      func->params[(*param_idx)++] = nir_deref_param;
+      break;
+
+   case vtn_base_type_image:
+   case vtn_base_type_sampler:
+      func->params[(*param_idx)++] = nir_deref_param;
+      break;
+
+   case vtn_base_type_pointer:
+      if (type->type) {
+         func->params[(*param_idx)++] = (nir_parameter) {
+            .num_components = glsl_get_vector_elements(type->type),
+            .bit_size = glsl_get_bit_size(type->type),
+         };
+      } else {
+         func->params[(*param_idx)++] = nir_deref_param;
+      }
+      break;
+
+   default:
+      func->params[(*param_idx)++] = (nir_parameter) {
+         .num_components = glsl_get_vector_elements(type->type),
+         .bit_size = glsl_get_bit_size(type->type),
+      };
+   }
+}
+
+static void
+vtn_ssa_value_add_to_call_params(struct vtn_builder *b,
+                                 struct vtn_ssa_value *value,
+                                 struct vtn_type *type,
+                                 nir_call_instr *call,
+                                 unsigned *param_idx)
+{
+   switch (type->base_type) {
+   case vtn_base_type_array:
+   case vtn_base_type_matrix:
+      for (unsigned i = 0; i < type->length; i++) {
+         vtn_ssa_value_add_to_call_params(b, value->elems[i],
+                                          type->array_element,
+                                          call, param_idx);
+      }
+      break;
+
+   case vtn_base_type_struct:
+      for (unsigned i = 0; i < type->length; i++) {
+         vtn_ssa_value_add_to_call_params(b, value->elems[i],
+                                          type->members[i],
+                                          call, param_idx);
+      }
+      break;
+
+   default:
+      call->params[(*param_idx)++] = nir_src_for_ssa(value->def);
+      break;
+   }
+}
+
+static void
+vtn_ssa_value_load_function_param(struct vtn_builder *b,
+                                  struct vtn_ssa_value *value,
+                                  struct vtn_type *type,
+                                  unsigned *param_idx)
+{
+   switch (type->base_type) {
+   case vtn_base_type_array:
+   case vtn_base_type_matrix:
+      for (unsigned i = 0; i < type->length; i++) {
+         vtn_ssa_value_load_function_param(b, value->elems[i],
+                                           type->array_element, param_idx);
+      }
+      break;
+
+   case vtn_base_type_struct:
+      for (unsigned i = 0; i < type->length; i++) {
+         vtn_ssa_value_load_function_param(b, value->elems[i],
+                                           type->members[i], param_idx);
+      }
+      break;
+
+   default:
+      value->def = nir_load_param(&b->nb, (*param_idx)++);
+      break;
+   }
+}
+
+void
+vtn_handle_function_call(struct vtn_builder *b, SpvOp opcode,
+                         const uint32_t *w, unsigned count)
+{
+   struct vtn_type *res_type = vtn_value(b, w[1], vtn_value_type_type)->type;
+   struct vtn_function *vtn_callee =
+      vtn_value(b, w[3], vtn_value_type_function)->func;
+   struct nir_function *callee = vtn_callee->impl->function;
+
+   vtn_callee->referenced = true;
+
+   nir_call_instr *call = nir_call_instr_create(b->nb.shader, callee);
+
+   unsigned param_idx = 0;
+
+   nir_deref_instr *ret_deref = NULL;
+   struct vtn_type *ret_type = vtn_callee->type->return_type;
+   if (ret_type->base_type != vtn_base_type_void) {
+      nir_variable *ret_tmp =
+         nir_local_variable_create(b->nb.impl,
+                                   glsl_get_bare_type(ret_type->type),
+                                   "return_tmp");
+      ret_deref = nir_build_deref_var(&b->nb, ret_tmp);
+      call->params[param_idx++] = nir_src_for_ssa(&ret_deref->dest.ssa);
+   }
+
+   for (unsigned i = 0; i < vtn_callee->type->length; i++) {
+      struct vtn_type *arg_type = vtn_callee->type->params[i];
+      unsigned arg_id = w[4 + i];
+
+      if (arg_type->base_type == vtn_base_type_sampled_image) {
+         struct vtn_sampled_image *sampled_image =
+            vtn_value(b, arg_id, vtn_value_type_sampled_image)->sampled_image;
+
+         call->params[param_idx++] =
+            nir_src_for_ssa(&sampled_image->image->deref->dest.ssa);
+         call->params[param_idx++] =
+            nir_src_for_ssa(&sampled_image->sampler->deref->dest.ssa);
+      } else if (arg_type->base_type == vtn_base_type_pointer ||
+                 arg_type->base_type == vtn_base_type_image ||
+                 arg_type->base_type == vtn_base_type_sampler) {
+         struct vtn_pointer *pointer =
+            vtn_value(b, arg_id, vtn_value_type_pointer)->pointer;
+         call->params[param_idx++] =
+            nir_src_for_ssa(vtn_pointer_to_ssa(b, pointer));
+      } else {
+         vtn_ssa_value_add_to_call_params(b, vtn_ssa_value(b, arg_id),
+                                          arg_type, call, &param_idx);
+      }
+   }
+   assert(param_idx == call->num_params);
+
+   nir_builder_instr_insert(&b->nb, &call->instr);
+
+   if (ret_type->base_type == vtn_base_type_void) {
+      vtn_push_value(b, w[2], vtn_value_type_undef);
+   } else {
+      vtn_push_ssa(b, w[2], res_type, vtn_local_load(b, ret_deref));
+   }
+}
+
 static bool
 vtn_cfg_handle_prepass_instruction(struct vtn_builder *b, SpvOp opcode,
                                    const uint32_t *w, unsigned count)
 {
    switch (opcode) {
    case SpvOpFunction: {
-      assert(b->func == NULL);
+      vtn_assert(b->func == NULL);
       b->func = rzalloc(b, struct vtn_function);
 
       list_inithead(&b->func->body);
@@ -41,34 +253,46 @@ vtn_cfg_handle_prepass_instruction(struct vtn_builder *b, SpvOp opcode,
       struct vtn_value *val = vtn_push_value(b, w[2], vtn_value_type_function);
       val->func = b->func;
 
-      const struct vtn_type *func_type =
-         vtn_value(b, w[4], vtn_value_type_type)->type;
+      b->func->type = vtn_value(b, w[4], vtn_value_type_type)->type;
+      const struct vtn_type *func_type = b->func->type;
 
-      assert(func_type->return_type->type == result_type);
+      vtn_assert(func_type->return_type->type == result_type);
 
       nir_function *func =
          nir_function_create(b->shader, ralloc_strdup(b->shader, val->name));
 
-      func->num_params = func_type->length;
-      func->params = ralloc_array(b->shader, nir_parameter, func->num_params);
-      for (unsigned i = 0; i < func->num_params; i++) {
-         if (func_type->params[i]->base_type == vtn_base_type_pointer &&
-             func_type->params[i]->type == NULL) {
-            func->params[i].type = func_type->params[i]->deref->type;
-         } else {
-            func->params[i].type = func_type->params[i]->type;
-         }
+      unsigned num_params = 0;
+      for (unsigned i = 0; i < func_type->length; i++)
+         num_params += vtn_type_count_function_params(func_type->params[i]);
+
+      /* Add one parameter for the function return value */
+      if (func_type->return_type->base_type != vtn_base_type_void)
+         num_params++;
 
-         /* TODO: We could do something smarter here. */
-         func->params[i].param_type = nir_parameter_inout;
+      func->num_params = num_params;
+      func->params = ralloc_array(b->shader, nir_parameter, num_params);
+
+      unsigned idx = 0;
+      if (func_type->return_type->base_type != vtn_base_type_void) {
+         /* The return value is a regular pointer */
+         func->params[idx++] = (nir_parameter) {
+            .num_components = 1, .bit_size = 32,
+         };
       }
 
-      func->return_type = func_type->return_type->type;
+      for (unsigned i = 0; i < func_type->length; i++)
+         vtn_type_add_to_function_params(func_type->params[i], func, &idx);
+      assert(idx == num_params);
 
       b->func->impl = nir_function_impl_create(func);
+      nir_builder_init(&b->nb, func->impl);
       b->nb.cursor = nir_before_cf_list(&b->func->impl->body);
 
       b->func_param_idx = 0;
+
+      /* The return value is the first parameter */
+      if (func_type->return_type->base_type != vtn_base_type_void)
+         b->func_param_idx++;
       break;
    }
 
@@ -80,51 +304,51 @@ vtn_cfg_handle_prepass_instruction(struct vtn_builder *b, SpvOp opcode,
    case SpvOpFunctionParameter: {
       struct vtn_type *type = vtn_value(b, w[1], vtn_value_type_type)->type;
 
-      assert(b->func_param_idx < b->func->impl->num_params);
-      nir_variable *param = b->func->impl->params[b->func_param_idx++];
-
-      if (type->base_type == vtn_base_type_pointer && type->type == NULL) {
-         struct vtn_variable *vtn_var = rzalloc(b, struct vtn_variable);
-         vtn_var->type = type->deref;
-         vtn_var->var = param;
-
-         assert(vtn_var->type->type == param->type);
-
-         struct vtn_type *without_array = vtn_var->type;
-         while(glsl_type_is_array(without_array->type))
-            without_array = without_array->array_element;
-
-         if (glsl_type_is_image(without_array->type)) {
-            vtn_var->mode = vtn_variable_mode_image;
-            param->interface_type = without_array->type;
-         } else if (glsl_type_is_sampler(without_array->type)) {
-            vtn_var->mode = vtn_variable_mode_sampler;
-            param->interface_type = without_array->type;
-         } else {
-            vtn_var->mode = vtn_variable_mode_param;
-         }
+      vtn_assert(b->func_param_idx < b->func->impl->function->num_params);
 
+      if (type->base_type == vtn_base_type_sampled_image) {
+         /* Sampled images are actually two parameters.  The first is the
+          * image and the second is the sampler.
+          */
+         struct vtn_value *val =
+            vtn_push_value(b, w[2], vtn_value_type_sampled_image);
+
+         val->sampled_image = ralloc(b, struct vtn_sampled_image);
+         val->sampled_image->type = type;
+
+         struct vtn_type *sampler_type = rzalloc(b, struct vtn_type);
+         sampler_type->base_type = vtn_base_type_sampler;
+         sampler_type->type = glsl_bare_sampler_type();
+
+         val->sampled_image->image =
+            vtn_load_param_pointer(b, type, b->func_param_idx++);
+         val->sampled_image->sampler =
+            vtn_load_param_pointer(b, sampler_type, b->func_param_idx++);
+      } else if (type->base_type == vtn_base_type_pointer &&
+                 type->type != NULL) {
+         /* This is a pointer with an actual storage type */
          struct vtn_value *val =
             vtn_push_value(b, w[2], vtn_value_type_pointer);
-
-         /* Name the parameter so it shows up nicely in NIR */
-         param->name = ralloc_strdup(param, val->name);
-
-         val->pointer = vtn_pointer_for_variable(b, vtn_var, type);
+         nir_ssa_def *ssa_ptr = nir_load_param(&b->nb, b->func_param_idx++);
+         val->pointer = vtn_pointer_from_ssa(b, ssa_ptr, type);
+      } else if (type->base_type == vtn_base_type_pointer ||
+                 type->base_type == vtn_base_type_image ||
+                 type->base_type == vtn_base_type_sampler) {
+         struct vtn_value *val =
+            vtn_push_value(b, w[2], vtn_value_type_pointer);
+         val->pointer =
+            vtn_load_param_pointer(b, type, b->func_param_idx++);
       } else {
          /* We're a regular SSA value. */
-         struct vtn_ssa_value *param_ssa =
-            vtn_local_load(b, nir_deref_var_create(b, param));
-         struct vtn_value *val = vtn_push_ssa(b, w[2], type, param_ssa);
-
-         /* Name the parameter so it shows up nicely in NIR */
-         param->name = ralloc_strdup(param, val->name);
+         struct vtn_ssa_value *value = vtn_create_ssa_value(b, type->type);
+         vtn_ssa_value_load_function_param(b, value, type, &b->func_param_idx);
+         vtn_push_ssa(b, w[2], type, value);
       }
       break;
    }
 
    case SpvOpLabel: {
-      assert(b->block == NULL);
+      vtn_assert(b->block == NULL);
       b->block = rzalloc(b, struct vtn_block);
       b->block->node.type = vtn_cf_node_type_block;
       b->block->label = w;
@@ -143,7 +367,7 @@ vtn_cfg_handle_prepass_instruction(struct vtn_builder *b, SpvOp opcode,
 
    case SpvOpSelectionMerge:
    case SpvOpLoopMerge:
-      assert(b->block && b->block->merge == NULL);
+      vtn_assert(b->block && b->block->merge == NULL);
       b->block->merge = w;
       break;
 
@@ -154,7 +378,7 @@ vtn_cfg_handle_prepass_instruction(struct vtn_builder *b, SpvOp opcode,
    case SpvOpReturn:
    case SpvOpReturnValue:
    case SpvOpUnreachable:
-      assert(b->block && b->block->branch == NULL);
+      vtn_assert(b->block && b->block->branch == NULL);
       b->block->branch = w;
       b->block = NULL;
       break;
@@ -170,7 +394,7 @@ vtn_cfg_handle_prepass_instruction(struct vtn_builder *b, SpvOp opcode,
 static void
 vtn_add_case(struct vtn_builder *b, struct vtn_switch *swtch,
              struct vtn_block *break_block,
-             uint32_t block_id, uint32_t val, bool is_default)
+             uint32_t block_id, uint64_t val, bool is_default)
 {
    struct vtn_block *case_block =
       vtn_value(b, block_id, vtn_value_type_block)->block;
@@ -197,7 +421,7 @@ vtn_add_case(struct vtn_builder *b, struct vtn_switch *swtch,
    if (is_default) {
       case_block->switch_case->is_default = true;
    } else {
-      util_dynarray_append(&case_block->switch_case->values, uint32_t, val);
+      util_dynarray_append(&case_block->switch_case->values, uint64_t, val);
    }
 }
 
@@ -231,14 +455,15 @@ vtn_order_case(struct vtn_switch *swtch, struct vtn_case *cse)
 }
 
 static enum vtn_branch_type
-vtn_get_branch_type(struct vtn_block *block,
+vtn_get_branch_type(struct vtn_builder *b,
+                    struct vtn_block *block,
                     struct vtn_case *swcase, struct vtn_block *switch_break,
                     struct vtn_block *loop_break, struct vtn_block *loop_cont)
 {
    if (block->switch_case) {
       /* This branch is actually a fallthrough */
-      assert(swcase->fallthrough == NULL ||
-             swcase->fallthrough == block->switch_case);
+      vtn_assert(swcase->fallthrough == NULL ||
+                 swcase->fallthrough == block->switch_case);
       swcase->fallthrough = block->switch_case;
       return vtn_branch_type_switch_fallthrough;
    } else if (block == loop_break) {
@@ -297,11 +522,24 @@ vtn_cfg_walk_blocks(struct vtn_builder *b, struct list_head *cf_list,
          vtn_cfg_walk_blocks(b, &loop->cont_body, new_loop_cont, NULL, NULL,
                              new_loop_break, NULL, block);
 
+         enum vtn_branch_type branch_type =
+            vtn_get_branch_type(b, new_loop_break, switch_case, switch_break,
+                                loop_break, loop_cont);
+
+         if (branch_type != vtn_branch_type_none) {
+            /* Stop walking through the CFG when this inner loop's break block
+             * ends up as the same block as the outer loop's continue block
+             * because we are already going to visit it.
+             */
+            vtn_assert(branch_type == vtn_branch_type_loop_continue);
+            return;
+         }
+
          block = new_loop_break;
          continue;
       }
 
-      assert(block->node.link.next == NULL);
+      vtn_assert(block->node.link.next == NULL);
       list_addtail(&block->node.link, cf_list);
 
       switch (*block->branch & SpvOpCodeMask) {
@@ -309,7 +547,7 @@ vtn_cfg_walk_blocks(struct vtn_builder *b, struct list_head *cf_list,
          struct vtn_block *branch_block =
             vtn_value(b, block->branch[1], vtn_value_type_block)->block;
 
-         block->branch_type = vtn_get_branch_type(branch_block,
+         block->branch_type = vtn_get_branch_type(b, branch_block,
                                                   switch_case, switch_break,
                                                   loop_break, loop_cont);
 
@@ -349,10 +587,10 @@ vtn_cfg_walk_blocks(struct vtn_builder *b, struct list_head *cf_list,
             if_stmt->control = block->merge[2];
          }
 
-         if_stmt->then_type = vtn_get_branch_type(then_block,
+         if_stmt->then_type = vtn_get_branch_type(b, then_block,
                                                   switch_case, switch_break,
                                                   loop_break, loop_cont);
-         if_stmt->else_type = vtn_get_branch_type(else_block,
+         if_stmt->else_type = vtn_get_branch_type(b, else_block,
                                                   switch_case, switch_break,
                                                   loop_break, loop_cont);
 
@@ -367,7 +605,7 @@ vtn_cfg_walk_blocks(struct vtn_builder *b, struct list_head *cf_list,
          } else if (if_stmt->then_type == vtn_branch_type_none &&
                     if_stmt->else_type == vtn_branch_type_none) {
             /* Neither side of the if is something we can short-circuit. */
-            assert((*block->merge & SpvOpCodeMask) == SpvOpSelectionMerge);
+            vtn_assert((*block->merge & SpvOpCodeMask) == SpvOpSelectionMerge);
             struct vtn_block *merge_block =
                vtn_value(b, block->merge[1], vtn_value_type_block)->block;
 
@@ -379,7 +617,7 @@ vtn_cfg_walk_blocks(struct vtn_builder *b, struct list_head *cf_list,
                                 loop_break, loop_cont, merge_block);
 
             enum vtn_branch_type merge_type =
-               vtn_get_branch_type(merge_block, switch_case, switch_break,
+               vtn_get_branch_type(b, merge_block, switch_case, switch_break,
                                    loop_break, loop_cont);
             if (merge_type == vtn_branch_type_none) {
                block = merge_block;
@@ -404,11 +642,11 @@ vtn_cfg_walk_blocks(struct vtn_builder *b, struct list_head *cf_list,
             }
             continue;
          }
-         unreachable("Should have returned or continued");
+         vtn_fail("Should have returned or continued");
       }
 
       case SpvOpSwitch: {
-         assert((*block->merge & SpvOpCodeMask) == SpvOpSelectionMerge);
+         vtn_assert((*block->merge & SpvOpCodeMask) == SpvOpSelectionMerge);
          struct vtn_block *break_block =
             vtn_value(b, block->merge[1], vtn_value_type_block)->block;
 
@@ -424,16 +662,43 @@ vtn_cfg_walk_blocks(struct vtn_builder *b, struct list_head *cf_list,
          const uint32_t *branch_end =
             block->branch + (block->branch[0] >> SpvWordCountShift);
 
-         vtn_add_case(b, swtch, break_block, block->branch[2], 0, true);
-         for (const uint32_t *w = block->branch + 3; w < branch_end; w += 2)
-            vtn_add_case(b, swtch, break_block, w[1], w[0], false);
+         struct vtn_value *cond_val = vtn_untyped_value(b, block->branch[1]);
+         vtn_fail_if(!cond_val->type ||
+                     cond_val->type->base_type != vtn_base_type_scalar,
+                     "Selector of OpSelect must have a type of OpTypeInt");
+
+         nir_alu_type cond_type =
+            nir_get_nir_type_for_glsl_type(cond_val->type->type);
+         vtn_fail_if(nir_alu_type_get_base_type(cond_type) != nir_type_int &&
+                     nir_alu_type_get_base_type(cond_type) != nir_type_uint,
+                     "Selector of OpSelect must have a type of OpTypeInt");
+
+         bool is_default = true;
+         const unsigned bitsize = nir_alu_type_get_type_size(cond_type);
+         for (const uint32_t *w = block->branch + 2; w < branch_end;) {
+            uint64_t literal = 0;
+            if (!is_default) {
+               if (bitsize <= 32) {
+                  literal = *(w++);
+               } else {
+                  assert(bitsize == 64);
+                  literal = vtn_u64_literal(w);
+                  w += 2;
+               }
+            }
+
+            uint32_t block_id = *(w++);
+
+            vtn_add_case(b, swtch, break_block, block_id, literal, is_default);
+            is_default = false;
+         }
 
          /* Now, we go through and walk the blocks.  While we walk through
           * the blocks, we also gather the much-needed fall-through
           * information.
           */
          list_for_each_entry(struct vtn_case, cse, &swtch->cases, link) {
-            assert(cse->start_block != break_block);
+            vtn_assert(cse->start_block != break_block);
             vtn_cfg_walk_blocks(b, &cse->body, cse->start_block, cse,
                                 break_block, loop_break, loop_cont, NULL);
          }
@@ -441,20 +706,27 @@ vtn_cfg_walk_blocks(struct vtn_builder *b, struct list_head *cf_list,
          /* Finally, we walk over all of the cases one more time and put
           * them in fall-through order.
           */
-         for (const uint32_t *w = block->branch + 2; w < branch_end; w += 2) {
+         for (const uint32_t *w = block->branch + 2; w < branch_end;) {
             struct vtn_block *case_block =
                vtn_value(b, *w, vtn_value_type_block)->block;
 
+            if (bitsize <= 32) {
+               w += 2;
+            } else {
+               assert(bitsize == 64);
+               w += 3;
+            }
+
             if (case_block == break_block)
                continue;
 
-            assert(case_block->switch_case);
+            vtn_assert(case_block->switch_case);
 
             vtn_order_case(swtch, case_block->switch_case);
          }
 
          enum vtn_branch_type branch_type =
-            vtn_get_branch_type(break_block, switch_case, NULL,
+            vtn_get_branch_type(b, break_block, switch_case, NULL,
                                 loop_break, loop_cont);
 
          if (branch_type != vtn_branch_type_none) {
@@ -462,7 +734,7 @@ vtn_cfg_walk_blocks(struct vtn_builder *b, struct list_head *cf_list,
              * for the containing loop.  In this case, we need to bail and let
              * the loop parsing code handle the continue properly.
              */
-            assert(branch_type == vtn_branch_type_loop_continue);
+            vtn_assert(branch_type == vtn_branch_type_loop_continue);
             return;
          }
 
@@ -474,7 +746,7 @@ vtn_cfg_walk_blocks(struct vtn_builder *b, struct list_head *cf_list,
          return;
 
       default:
-         unreachable("Unhandled opcode");
+         vtn_fail("Unhandled opcode");
       }
    }
 }
@@ -519,7 +791,7 @@ vtn_handle_phis_first_pass(struct vtn_builder *b, SpvOp opcode,
    _mesa_hash_table_insert(b->phi_table, w, phi_var);
 
    vtn_push_ssa(b, w[2], type,
-                vtn_local_load(b, nir_deref_var_create(b, phi_var)));
+                vtn_local_load(b, nir_build_deref_var(&b->nb, phi_var)));
 
    return true;
 }
@@ -532,7 +804,7 @@ vtn_handle_phi_second_pass(struct vtn_builder *b, SpvOp opcode,
       return true;
 
    struct hash_entry *phi_entry = _mesa_hash_table_search(b->phi_table, w);
-   assert(phi_entry);
+   vtn_assert(phi_entry);
    nir_variable *phi_var = phi_entry->data;
 
    for (unsigned i = 3; i < count; i += 2) {
@@ -543,7 +815,7 @@ vtn_handle_phi_second_pass(struct vtn_builder *b, SpvOp opcode,
 
       struct vtn_ssa_value *src = vtn_ssa_value(b, w[i]);
 
-      vtn_local_store(b, src, nir_deref_var_create(b, phi_var));
+      vtn_local_store(b, src, nir_build_deref_var(&b->nb, phi_var));
    }
 
    return true;
@@ -555,7 +827,7 @@ vtn_emit_branch(struct vtn_builder *b, enum vtn_branch_type branch_type,
 {
    switch (branch_type) {
    case vtn_branch_type_switch_break:
-      nir_store_var(&b->nb, switch_fall_var, nir_imm_int(&b->nb, NIR_FALSE), 1);
+      nir_store_var(&b->nb, switch_fall_var, nir_imm_false(&b->nb), 1);
       *has_switch_break = true;
       break;
    case vtn_branch_type_switch_fallthrough:
@@ -576,7 +848,7 @@ vtn_emit_branch(struct vtn_builder *b, enum vtn_branch_type branch_type,
       break;
    }
    default:
-      unreachable("Invalid branch type");
+      vtn_fail("Invalid branch type");
    }
 }
 
@@ -604,9 +876,16 @@ vtn_emit_cf_list(struct vtn_builder *b, struct list_head *cf_list,
          nir_builder_instr_insert(&b->nb, &block->end_nop->instr);
 
          if ((*block->branch & SpvOpCodeMask) == SpvOpReturnValue) {
+            vtn_fail_if(b->func->type->return_type->base_type ==
+                        vtn_base_type_void,
+                        "Return with a value from a function returning void");
             struct vtn_ssa_value *src = vtn_ssa_value(b, block->branch[1]);
-            vtn_local_store(b, src,
-                            nir_deref_var_create(b, b->impl->return_var));
+            const struct glsl_type *ret_type =
+               glsl_get_bare_type(b->func->type->return_type->type);
+            nir_deref_instr *ret_deref =
+               nir_build_deref_cast(&b->nb, nir_load_param(&b->nb, 0),
+                                    nir_var_local, ret_type, 0);
+            vtn_local_store(b, src, ret_deref);
          }
 
          if (block->branch_type != vtn_branch_type_none) {
@@ -668,7 +947,7 @@ vtn_emit_cf_list(struct vtn_builder *b, struct list_head *cf_list,
                nir_local_variable_create(b->nb.impl, glsl_bool_type(), "cont");
 
             b->nb.cursor = nir_before_cf_node(&loop->cf_node);
-            nir_store_var(&b->nb, do_cont, nir_imm_int(&b->nb, NIR_FALSE), 1);
+            nir_store_var(&b->nb, do_cont, nir_imm_false(&b->nb), 1);
 
             b->nb.cursor = nir_before_cf_list(&loop->body);
 
@@ -679,7 +958,7 @@ vtn_emit_cf_list(struct vtn_builder *b, struct list_head *cf_list,
 
             nir_pop_if(&b->nb, cont_if);
 
-            nir_store_var(&b->nb, do_cont, nir_imm_int(&b->nb, NIR_TRUE), 1);
+            nir_store_var(&b->nb, do_cont, nir_imm_true(&b->nb), 1);
 
             b->has_loop_continue = true;
          }
@@ -697,7 +976,7 @@ vtn_emit_cf_list(struct vtn_builder *b, struct list_head *cf_list,
           */
          nir_variable *fall_var =
             nir_local_variable_create(b->nb.impl, glsl_bool_type(), "fall");
-         nir_store_var(&b->nb, fall_var, nir_imm_int(&b->nb, NIR_FALSE), 1);
+         nir_store_var(&b->nb, fall_var, nir_imm_false(&b->nb), 1);
 
          /* Next, we gather up all of the conditions.  We have to do this
           * up-front because we also need to build an "any" condition so
@@ -718,9 +997,9 @@ vtn_emit_cf_list(struct vtn_builder *b, struct list_head *cf_list,
             }
 
             nir_ssa_def *cond = NULL;
-            util_dynarray_foreach(&cse->values, uint32_t, val) {
-               nir_ssa_def *is_val =
-                  nir_ieq(&b->nb, sel, nir_imm_int(&b->nb, *val));
+            util_dynarray_foreach(&cse->values, uint64_t, val) {
+               nir_ssa_def *imm = nir_imm_intN_t(&b->nb, *val, sel->bit_size);
+               nir_ssa_def *is_val = nir_ieq(&b->nb, sel, imm);
 
                cond = cond ? nir_ior(&b->nb, cond, is_val) : is_val;
             }
@@ -728,7 +1007,7 @@ vtn_emit_cf_list(struct vtn_builder *b, struct list_head *cf_list,
             any = any ? nir_ior(&b->nb, any, cond) : cond;
             conditions[i++] = cond;
          }
-         assert(i == num_cases);
+         vtn_assert(i == num_cases);
 
          /* Now we can walk the list of cases and actually emit code */
          i = 0;
@@ -736,7 +1015,7 @@ vtn_emit_cf_list(struct vtn_builder *b, struct list_head *cf_list,
             /* Figure out the condition */
             nir_ssa_def *cond = conditions[i++];
             if (cse->is_default) {
-               assert(cond == NULL);
+               vtn_assert(cond == NULL);
                cond = nir_inot(&b->nb, any);
             }
             /* Take fallthrough into account */
@@ -745,19 +1024,19 @@ vtn_emit_cf_list(struct vtn_builder *b, struct list_head *cf_list,
             nir_if *case_if = nir_push_if(&b->nb, cond);
 
             bool has_break = false;
-            nir_store_var(&b->nb, fall_var, nir_imm_int(&b->nb, NIR_TRUE), 1);
+            nir_store_var(&b->nb, fall_var, nir_imm_true(&b->nb), 1);
             vtn_emit_cf_list(b, &cse->body, fall_var, &has_break, handler);
             (void)has_break; /* We don't care */
 
             nir_pop_if(&b->nb, case_if);
          }
-         assert(i == num_cases);
+         vtn_assert(i == num_cases);
 
          break;
       }
 
       default:
-         unreachable("Invalid CF node type");
+         vtn_fail("Invalid CF node type");
       }
    }
 }
@@ -767,6 +1046,7 @@ vtn_function_emit(struct vtn_builder *b, struct vtn_function *func,
                   vtn_instruction_handler instruction_handler)
 {
    nir_builder_init(&b->nb, func->impl);
+   b->func = func;
    b->nb.cursor = nir_after_cf_list(&func->impl->body);
    b->has_loop_continue = false;
    b->phi_table = _mesa_hash_table_create(b, _mesa_hash_pointer,
@@ -777,10 +1057,14 @@ vtn_function_emit(struct vtn_builder *b, struct vtn_function *func,
    vtn_foreach_instruction(b, func->start_block->label, func->end,
                            vtn_handle_phi_second_pass);
 
+   nir_rematerialize_derefs_in_use_blocks_impl(func->impl);
+
    /* Continue blocks for loops get inserted before the body of the loop
     * but instructions in the continue may use SSA defs in the loop body.
     * Therefore, we need to repair SSA to insert the needed phi nodes.
     */
    if (b->has_loop_continue)
       nir_repair_ssa_impl(func->impl);
+
+   func->emitted = true;
 }