gallivm: add subgroup shuffle support

Signed-off-by: Autumn on Tape <autumn@cyfox.net>
Reviewed-by: Dave Airlie <airlied@redhat.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/13671>

3 files changed, 64 insertions, 0 deletions

diff --git a/src/gallium/auxiliary/gallivm/lp_bld_nir.c b/src/gallium/auxiliary/gallivm/lp_bld_nir.c
index d7741abe7c1..499942ab10c 100644
--- a/src/gallium/auxiliary/gallivm/lp_bld_nir.c
+++ b/src/gallium/auxiliary/gallivm/lp_bld_nir.c
@@ -1731,6 +1731,20 @@ static void visit_global_atomic(struct lp_build_nir_context *bld_base,
                            val_bitsize, addr, val, val2, &result[0]);
 }
 
+#if LLVM_VERSION_MAJOR >= 10
+static void visit_shuffle(struct lp_build_nir_context *bld_base,
+                          nir_intrinsic_instr *instr,
+                          LLVMValueRef dst[4])
+{
+   LLVMValueRef src = get_src(bld_base, instr->src[0]);
+   src = cast_type(bld_base, src, nir_type_int, nir_src_bit_size(instr->src[0]));
+   LLVMValueRef index = get_src(bld_base, instr->src[1]);
+   index = cast_type(bld_base, index, nir_type_uint, nir_src_bit_size(instr->src[1]));
+
+   bld_base->shuffle(bld_base, src, index, instr, dst);
+}
+#endif
+
 static void visit_interp(struct lp_build_nir_context *bld_base,
                          nir_intrinsic_instr *instr,
                          LLVMValueRef result[NIR_MAX_VEC_COMPONENTS])
@@ -1959,6 +1973,11 @@ 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;
+#if LLVM_VERSION_MAJOR >= 10
+   case nir_intrinsic_shuffle:
+      visit_shuffle(bld_base, instr, result);
+      break;
+#endif
    case nir_intrinsic_read_invocation:
    case nir_intrinsic_read_first_invocation: {
       LLVMValueRef src1 = NULL;
@@ -2518,6 +2537,7 @@ void lp_build_opt_nir(struct nir_shader *nir)
         .ballot_components = 1,
 	.lower_to_scalar = true,
 	.lower_subgroup_masks = true,
+        .lower_relative_shuffle = true,
       };
       NIR_PASS_V(nir, nir_lower_subgroups, &subgroups_options);
 
diff --git a/src/gallium/auxiliary/gallivm/lp_bld_nir.h b/src/gallium/auxiliary/gallivm/lp_bld_nir.h
index a71181f9ddb..1ca26caf9a0 100644
--- a/src/gallium/auxiliary/gallivm/lp_bld_nir.h
+++ b/src/gallium/auxiliary/gallivm/lp_bld_nir.h
@@ -209,6 +209,13 @@ 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]);
+#if LLVM_VERSION_MAJOR >= 10
+   void (*shuffle)(struct lp_build_nir_context *bld_base,
+                   LLVMValueRef src,
+                   LLVMValueRef index,
+                   nir_intrinsic_instr *instr,
+                   LLVMValueRef dst[4]);
+#endif
    void (*read_invocation)(struct lp_build_nir_context *bld_base,
                            LLVMValueRef src, unsigned bit_size, LLVMValueRef invoc,
                            LLVMValueRef dst[4]);
diff --git a/src/gallium/auxiliary/gallivm/lp_bld_nir_soa.c b/src/gallium/auxiliary/gallivm/lp_bld_nir_soa.c
index 558245ca641..cb8a2c31e1f 100644
--- a/src/gallium/auxiliary/gallivm/lp_bld_nir_soa.c
+++ b/src/gallium/auxiliary/gallivm/lp_bld_nir_soa.c
@@ -2110,6 +2110,40 @@ static void emit_elect(struct lp_build_nir_context *bld_base, LLVMValueRef resul
                                       "");
 }
 
+#if LLVM_VERSION_MAJOR >= 10
+static void emit_shuffle(struct lp_build_nir_context *bld_base, LLVMValueRef src, LLVMValueRef index,
+                        nir_intrinsic_instr *instr, LLVMValueRef result[4])
+{
+   assert(instr->intrinsic == nir_intrinsic_shuffle);
+
+   struct gallivm_state *gallivm = bld_base->base.gallivm;
+   LLVMBuilderRef builder = gallivm->builder;
+   uint32_t bit_size = nir_src_bit_size(instr->src[0]);
+   struct lp_build_context *int_bld = get_int_bld(bld_base, true, bit_size);
+
+   LLVMValueRef res_store = lp_build_alloca(gallivm, int_bld->vec_type, "");
+   struct lp_build_loop_state loop_state;
+   lp_build_loop_begin(&loop_state, gallivm, lp_build_const_int32(gallivm, 0));
+
+   LLVMValueRef index_value = LLVMBuildExtractElement(builder, index, loop_state.counter, "");
+
+   LLVMValueRef src_value = LLVMBuildExtractElement(builder, src, index_value, "");
+   /* freeze `src_value` in case an out-of-bounds index or an index into an
+    * inactive invocation results in poison
+    */
+   src_value = LLVMBuildFreeze(builder, src_value, "");
+
+   LLVMValueRef res = LLVMBuildLoad(builder, res_store, "");
+   res = LLVMBuildInsertElement(builder, res, src_value, loop_state.counter, "");
+   LLVMBuildStore(builder, res, res_store);
+
+   lp_build_loop_end_cond(&loop_state, lp_build_const_int32(gallivm, bld_base->uint_bld.type.length),
+                          NULL, LLVMIntUGE);
+
+   result[0] = LLVMBuildLoad(builder, res_store, "");
+}
+#endif
+
 static void emit_reduce(struct lp_build_nir_context *bld_base, LLVMValueRef src,
                         nir_intrinsic_instr *instr, LLVMValueRef result[4])
 {
@@ -2613,6 +2647,9 @@ 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;
+#if LLVM_VERSION_MAJOR >= 10
+   bld.bld_base.shuffle = emit_shuffle;
+#endif
    bld.bld_base.read_invocation = emit_read_invocation;
    bld.bld_base.helper_invocation = emit_helper_invocation;
    bld.bld_base.interp_at = emit_interp_at;