gallivm: add subgroup read invocation support

Reviewed-by: Roland Scheidegger <sroland@vmware.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/9645>

3 files changed, 51 insertions, 0 deletions

diff --git a/src/gallium/auxiliary/gallivm/lp_bld_nir.c b/src/gallium/auxiliary/gallivm/lp_bld_nir.c
index fb9b670e440..c194290b2bc 100644
--- a/src/gallium/auxiliary/gallivm/lp_bld_nir.c
+++ b/src/gallium/auxiliary/gallivm/lp_bld_nir.c
@@ -1797,6 +1797,15 @@ static void visit_intrinsic(struct lp_build_nir_context *bld_base,
    case nir_intrinsic_ballot:
       bld_base->ballot(bld_base, cast_type(bld_base, get_src(bld_base, instr->src[0]), nir_type_int, 32), instr, result);
       break;
+   case nir_intrinsic_read_invocation:
+   case nir_intrinsic_read_first_invocation: {
+      LLVMValueRef src1 = NULL;
+
+      if (instr->intrinsic == nir_intrinsic_read_invocation)
+         src1 = cast_type(bld_base, get_src(bld_base, instr->src[1]), nir_type_int, 32);
+      bld_base->read_invocation(bld_base, get_src(bld_base, instr->src[0]), nir_src_bit_size(instr->src[0]), src1, result);
+      break;
+   }
    case nir_intrinsic_interp_deref_at_offset:
    case nir_intrinsic_interp_deref_at_centroid:
    case nir_intrinsic_interp_deref_at_sample:
diff --git a/src/gallium/auxiliary/gallivm/lp_bld_nir.h b/src/gallium/auxiliary/gallivm/lp_bld_nir.h
index a4cef097dfd..1a92bbc03c5 100644
--- a/src/gallium/auxiliary/gallivm/lp_bld_nir.h
+++ b/src/gallium/auxiliary/gallivm/lp_bld_nir.h
@@ -188,6 +188,9 @@ struct lp_build_nir_context
    void (*elect)(struct lp_build_nir_context *bld_base, LLVMValueRef dst[4]);
    void (*reduce)(struct lp_build_nir_context *bld_base, LLVMValueRef src, nir_intrinsic_instr *instr, LLVMValueRef dst[4]);
    void (*ballot)(struct lp_build_nir_context *bld_base, LLVMValueRef src, nir_intrinsic_instr *instr, LLVMValueRef dst[4]);
+   void (*read_invocation)(struct lp_build_nir_context *bld_base,
+                           LLVMValueRef src, unsigned bit_size, LLVMValueRef invoc,
+                           LLVMValueRef dst[4]);
    void (*helper_invocation)(struct lp_build_nir_context *bld_base, LLVMValueRef *dst);
 
    void (*interp_at)(struct lp_build_nir_context *bld_base,
diff --git a/src/gallium/auxiliary/gallivm/lp_bld_nir_soa.c b/src/gallium/auxiliary/gallivm/lp_bld_nir_soa.c
index 0302876e39b..6a390ecbc55 100644
--- a/src/gallium/auxiliary/gallivm/lp_bld_nir_soa.c
+++ b/src/gallium/auxiliary/gallivm/lp_bld_nir_soa.c
@@ -2093,6 +2093,44 @@ static void emit_reduce(struct lp_build_nir_context *bld_base, LLVMValueRef src,
       result[0] = LLVMBuildLoad(builder, res_store, "");
 }
 
+static void emit_read_invocation(struct lp_build_nir_context *bld_base,
+                                 LLVMValueRef src,
+                                 unsigned bit_size,
+                                 LLVMValueRef invoc,
+                                 LLVMValueRef result[4])
+{
+   struct gallivm_state *gallivm = bld_base->base.gallivm;
+   LLVMBuilderRef builder = gallivm->builder;
+   LLVMValueRef idx;
+   struct lp_build_context *uint_bld = get_int_bld(bld_base, true, bit_size);
+   if (invoc) {
+      idx = invoc;
+      idx = LLVMBuildExtractElement(gallivm->builder, idx, lp_build_const_int32(gallivm, 0), "");
+   } else {
+      /* have to find the first active invocation */
+      LLVMValueRef exec_mask = mask_vec(bld_base);
+      struct lp_build_loop_state loop_state;
+      LLVMValueRef res_store = lp_build_alloca(gallivm, bld_base->int_bld.elem_type, "");
+      LLVMValueRef outer_cond = LLVMBuildICmp(builder, LLVMIntNE, exec_mask, bld_base->uint_bld.zero, "");
+      lp_build_loop_begin(&loop_state, gallivm, lp_build_const_int32(gallivm, bld_base->uint_bld.type.length));
+
+      LLVMValueRef if_cond = LLVMBuildExtractElement(gallivm->builder, outer_cond, loop_state.counter, "");
+      struct lp_build_if_state ifthen;
+
+      lp_build_if(&ifthen, gallivm, if_cond);
+      LLVMBuildStore(builder, loop_state.counter, res_store);
+      lp_build_endif(&ifthen);
+
+      lp_build_loop_end_cond(&loop_state, lp_build_const_int32(gallivm, -1),
+                             lp_build_const_int32(gallivm, -1), LLVMIntEQ);
+      idx = LLVMBuildLoad(builder, res_store, "");
+   }
+
+   LLVMValueRef value = LLVMBuildExtractElement(gallivm->builder,
+                                                src, idx, "");
+   result[0] = lp_build_broadcast_scalar(uint_bld, value);
+}
+
 static void
 emit_interp_at(struct lp_build_nir_context *bld_base,
                unsigned num_components,
@@ -2341,6 +2379,7 @@ void lp_build_nir_soa(struct gallivm_state *gallivm,
    bld.bld_base.elect = emit_elect;
    bld.bld_base.reduce = emit_reduce;
    bld.bld_base.ballot = emit_ballot;
+   bld.bld_base.read_invocation = emit_read_invocation;
    bld.bld_base.helper_invocation = emit_helper_invocation;
    bld.bld_base.interp_at = emit_interp_at;
    bld.bld_base.load_scratch = emit_load_scratch;