spirv/nir: Add support for OpAtomicLoad/Store
[mesa.git] / src / compiler / spirv / spirv_to_nir.c
index 0d6a70e4660f17e091734990872a81d0e7593228..4d88fdbb1fae84a6b6d9e135d293cfb2ca1e4b6e 100644 (file)
@@ -1589,6 +1589,47 @@ vtn_handle_texture(struct vtn_builder *b, SpvOp opcode,
    nir_builder_instr_insert(&b->nb, &instr->instr);
 }
 
+static void
+fill_common_atomic_sources(struct vtn_builder *b, SpvOp opcode,
+                           const uint32_t *w, nir_src *src)
+{
+   switch (opcode) {
+   case SpvOpAtomicIIncrement:
+      src[0] = nir_src_for_ssa(nir_imm_int(&b->nb, 1));
+      break;
+
+   case SpvOpAtomicIDecrement:
+      src[0] = nir_src_for_ssa(nir_imm_int(&b->nb, -1));
+      break;
+
+   case SpvOpAtomicISub:
+      src[0] =
+         nir_src_for_ssa(nir_ineg(&b->nb, vtn_ssa_value(b, w[6])->def));
+      break;
+
+   case SpvOpAtomicCompareExchange:
+      src[0] = nir_src_for_ssa(vtn_ssa_value(b, w[8])->def);
+      src[1] = nir_src_for_ssa(vtn_ssa_value(b, w[7])->def);
+      break;
+      /* Fall through */
+
+   case SpvOpAtomicExchange:
+   case SpvOpAtomicIAdd:
+   case SpvOpAtomicSMin:
+   case SpvOpAtomicUMin:
+   case SpvOpAtomicSMax:
+   case SpvOpAtomicUMax:
+   case SpvOpAtomicAnd:
+   case SpvOpAtomicOr:
+   case SpvOpAtomicXor:
+      src[0] = nir_src_for_ssa(vtn_ssa_value(b, w[6])->def);
+      break;
+
+   default:
+      unreachable("Invalid SPIR-V atomic");
+   }
+}
+
 static nir_ssa_def *
 get_image_coord(struct vtn_builder *b, uint32_t value)
 {
@@ -1640,6 +1681,18 @@ vtn_handle_image(struct vtn_builder *b, SpvOp opcode,
       image = *vtn_value(b, w[3], vtn_value_type_image_pointer)->image;
       break;
 
+   case SpvOpAtomicLoad: {
+      image.image =
+         vtn_value(b, w[3], vtn_value_type_access_chain)->access_chain;
+      break;
+   }
+
+   case SpvOpAtomicStore: {
+      image.image =
+         vtn_value(b, w[1], vtn_value_type_access_chain)->access_chain;
+      break;
+   }
+
    case SpvOpImageQuerySize:
       image.image =
          vtn_value(b, w[3], vtn_value_type_access_chain)->access_chain;
@@ -1685,6 +1738,8 @@ vtn_handle_image(struct vtn_builder *b, SpvOp opcode,
    OP(ImageQuerySize,         size)
    OP(ImageRead,              load)
    OP(ImageWrite,             store)
+   OP(AtomicLoad,             load)
+   OP(AtomicStore,            store)
    OP(AtomicExchange,         atomic_exchange)
    OP(AtomicCompareExchange,  atomic_comp_swap)
    OP(AtomicIIncrement,       atomic_add)
@@ -1723,19 +1778,19 @@ vtn_handle_image(struct vtn_builder *b, SpvOp opcode,
    }
 
    switch (opcode) {
+   case SpvOpAtomicLoad:
    case SpvOpImageQuerySize:
    case SpvOpImageRead:
       break;
+   case SpvOpAtomicStore:
+      intrin->src[2] = nir_src_for_ssa(vtn_ssa_value(b, w[4])->def);
+      break;
    case SpvOpImageWrite:
       intrin->src[2] = nir_src_for_ssa(vtn_ssa_value(b, w[3])->def);
       break;
+
    case SpvOpAtomicIIncrement:
-      intrin->src[2] = nir_src_for_ssa(nir_imm_int(&b->nb, 1));
-      break;
    case SpvOpAtomicIDecrement:
-      intrin->src[2] = nir_src_for_ssa(nir_imm_int(&b->nb, -1));
-      break;
-
    case SpvOpAtomicExchange:
    case SpvOpAtomicIAdd:
    case SpvOpAtomicSMin:
@@ -1745,16 +1800,7 @@ vtn_handle_image(struct vtn_builder *b, SpvOp opcode,
    case SpvOpAtomicAnd:
    case SpvOpAtomicOr:
    case SpvOpAtomicXor:
-      intrin->src[2] = nir_src_for_ssa(vtn_ssa_value(b, w[6])->def);
-      break;
-
-   case SpvOpAtomicCompareExchange:
-      intrin->src[2] = nir_src_for_ssa(vtn_ssa_value(b, w[8])->def);
-      intrin->src[3] = nir_src_for_ssa(vtn_ssa_value(b, w[7])->def);
-      break;
-
-   case SpvOpAtomicISub:
-      intrin->src[2] = nir_src_for_ssa(nir_ineg(&b->nb, vtn_ssa_value(b, w[6])->def));
+      fill_common_atomic_sources(b, opcode, w, &intrin->src[2]);
       break;
 
    default:
@@ -1784,6 +1830,8 @@ static nir_intrinsic_op
 get_ssbo_nir_atomic_op(SpvOp opcode)
 {
    switch (opcode) {
+   case SpvOpAtomicLoad:      return nir_intrinsic_load_ssbo;
+   case SpvOpAtomicStore:     return nir_intrinsic_store_ssbo;
 #define OP(S, N) case SpvOp##S: return nir_intrinsic_ssbo_##N;
    OP(AtomicExchange,         atomic_exchange)
    OP(AtomicCompareExchange,  atomic_comp_swap)
@@ -1808,6 +1856,8 @@ static nir_intrinsic_op
 get_shared_nir_atomic_op(SpvOp opcode)
 {
    switch (opcode) {
+   case SpvOpAtomicLoad:      return nir_intrinsic_load_var;
+   case SpvOpAtomicStore:     return nir_intrinsic_store_var;
 #define OP(S, N) case SpvOp##S: return nir_intrinsic_var_##N;
    OP(AtomicExchange,         atomic_exchange)
    OP(AtomicCompareExchange,  atomic_comp_swap)
@@ -1829,31 +1879,21 @@ get_shared_nir_atomic_op(SpvOp opcode)
 }
 
 static void
-fill_common_atomic_sources(struct vtn_builder *b, SpvOp opcode,
-                           const uint32_t *w, nir_src *src)
+vtn_handle_ssbo_or_shared_atomic(struct vtn_builder *b, SpvOp opcode,
+                                 const uint32_t *w, unsigned count)
 {
+   struct vtn_access_chain *chain;
+   nir_intrinsic_instr *atomic;
+
    switch (opcode) {
+   case SpvOpAtomicLoad:
+   case SpvOpAtomicExchange:
+   case SpvOpAtomicCompareExchange:
+   case SpvOpAtomicCompareExchangeWeak:
    case SpvOpAtomicIIncrement:
-      src[0] = nir_src_for_ssa(nir_imm_int(&b->nb, 1));
-      break;
-
    case SpvOpAtomicIDecrement:
-      src[0] = nir_src_for_ssa(nir_imm_int(&b->nb, -1));
-      break;
-
-   case SpvOpAtomicISub:
-      src[0] =
-         nir_src_for_ssa(nir_ineg(&b->nb, vtn_ssa_value(b, w[6])->def));
-      break;
-
-   case SpvOpAtomicCompareExchange:
-      src[0] = nir_src_for_ssa(vtn_ssa_value(b, w[8])->def);
-      src[1] = nir_src_for_ssa(vtn_ssa_value(b, w[7])->def);
-      break;
-      /* Fall through */
-
-   case SpvOpAtomicExchange:
    case SpvOpAtomicIAdd:
+   case SpvOpAtomicISub:
    case SpvOpAtomicSMin:
    case SpvOpAtomicUMin:
    case SpvOpAtomicSMax:
@@ -1861,21 +1901,18 @@ fill_common_atomic_sources(struct vtn_builder *b, SpvOp opcode,
    case SpvOpAtomicAnd:
    case SpvOpAtomicOr:
    case SpvOpAtomicXor:
-      src[0] = nir_src_for_ssa(vtn_ssa_value(b, w[6])->def);
+      chain =
+         vtn_value(b, w[3], vtn_value_type_access_chain)->access_chain;
+      break;
+
+   case SpvOpAtomicStore:
+      chain =
+         vtn_value(b, w[1], vtn_value_type_access_chain)->access_chain;
       break;
 
    default:
       unreachable("Invalid SPIR-V atomic");
    }
-}
-
-static void
-vtn_handle_ssbo_or_shared_atomic(struct vtn_builder *b, SpvOp opcode,
-                                 const uint32_t *w, unsigned count)
-{
-   struct vtn_access_chain *chain =
-      vtn_value(b, w[3], vtn_value_type_access_chain)->access_chain;
-   nir_intrinsic_instr *atomic;
 
    /*
    SpvScope scope = w[4];
@@ -1897,18 +1934,58 @@ vtn_handle_ssbo_or_shared_atomic(struct vtn_builder *b, SpvOp opcode,
       nir_intrinsic_op op = get_ssbo_nir_atomic_op(opcode);
 
       atomic = nir_intrinsic_instr_create(b->nb.shader, op);
-      atomic->src[0] = nir_src_for_ssa(index);
-      atomic->src[1] = nir_src_for_ssa(offset);
-      fill_common_atomic_sources(b, opcode, w, &atomic->src[2]);
+
+      switch (opcode) {
+      case SpvOpAtomicLoad:
+         atomic->num_components = glsl_get_vector_elements(type->type);
+         atomic->src[0] = nir_src_for_ssa(index);
+         atomic->src[1] = nir_src_for_ssa(offset);
+         break;
+
+      case SpvOpAtomicStore:
+         atomic->num_components = glsl_get_vector_elements(type->type);
+         nir_intrinsic_set_write_mask(atomic, (1 << atomic->num_components) - 1);
+         atomic->src[0] = nir_src_for_ssa(vtn_ssa_value(b, w[4])->def);
+         atomic->src[1] = nir_src_for_ssa(index);
+         atomic->src[2] = nir_src_for_ssa(offset);
+         break;
+
+      case SpvOpAtomicExchange:
+      case SpvOpAtomicCompareExchange:
+      case SpvOpAtomicCompareExchangeWeak:
+      case SpvOpAtomicIIncrement:
+      case SpvOpAtomicIDecrement:
+      case SpvOpAtomicIAdd:
+      case SpvOpAtomicISub:
+      case SpvOpAtomicSMin:
+      case SpvOpAtomicUMin:
+      case SpvOpAtomicSMax:
+      case SpvOpAtomicUMax:
+      case SpvOpAtomicAnd:
+      case SpvOpAtomicOr:
+      case SpvOpAtomicXor:
+         atomic->src[0] = nir_src_for_ssa(index);
+         atomic->src[1] = nir_src_for_ssa(offset);
+         fill_common_atomic_sources(b, opcode, w, &atomic->src[2]);
+         break;
+
+      default:
+         unreachable("Invalid SPIR-V atomic");
+      }
    }
 
-   nir_ssa_dest_init(&atomic->instr, &atomic->dest, 1, 32, NULL);
+   if (opcode != SpvOpAtomicStore) {
+      struct vtn_type *type = vtn_value(b, w[1], vtn_value_type_type)->type;
+
+      nir_ssa_dest_init(&atomic->instr, &atomic->dest,
+                        glsl_get_vector_elements(type->type),
+                        glsl_get_bit_size(type->type), NULL);
 
-   struct vtn_type *type = vtn_value(b, w[1], vtn_value_type_type)->type;
-   struct vtn_value *val = vtn_push_value(b, w[2], vtn_value_type_ssa);
-   val->ssa = rzalloc(b, struct vtn_ssa_value);
-   val->ssa->def = &atomic->dest.ssa;
-   val->ssa->type = type->type;
+      struct vtn_value *val = vtn_push_value(b, w[2], vtn_value_type_ssa);
+      val->ssa = rzalloc(b, struct vtn_ssa_value);
+      val->ssa->def = &atomic->dest.ssa;
+      val->ssa->type = type->type;
+   }
 
    nir_builder_instr_insert(&b->nb, &atomic->instr);
 }
@@ -2690,6 +2767,7 @@ vtn_handle_body_instruction(struct vtn_builder *b, SpvOp opcode,
       break;
    }
 
+   case SpvOpAtomicLoad:
    case SpvOpAtomicExchange:
    case SpvOpAtomicCompareExchange:
    case SpvOpAtomicCompareExchangeWeak:
@@ -2714,6 +2792,17 @@ vtn_handle_body_instruction(struct vtn_builder *b, SpvOp opcode,
       break;
    }
 
+   case SpvOpAtomicStore: {
+      struct vtn_value *pointer = vtn_untyped_value(b, w[1]);
+      if (pointer->value_type == vtn_value_type_image_pointer) {
+         vtn_handle_image(b, opcode, w, count);
+      } else {
+         assert(pointer->value_type == vtn_value_type_access_chain);
+         vtn_handle_ssbo_or_shared_atomic(b, opcode, w, count);
+      }
+      break;
+   }
+
    case SpvOpSNegate:
    case SpvOpFNegate:
    case SpvOpNot: