From cc8a1229070898997fd165e4e60a3a2c8bd7f40a Mon Sep 17 00:00:00 2001
From: Enrico Galli <enrico.galli@intel.com>
Date: Thu, 5 Aug 2021 13:42:05 -0700
Subject: [PATCH] microsoft/compiler: Add support for local_invocation_index

Reviewed-by: Jesse Natalie <jenatali@microsoft.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/12228>
---
 src/microsoft/compiler/dxil_function.c |  1 +
 src/microsoft/compiler/nir_to_dxil.c   | 38 ++++++++++++++++++++++++++
 2 files changed, 39 insertions(+)

diff --git a/src/microsoft/compiler/dxil_function.c b/src/microsoft/compiler/dxil_function.c
index 4f55992c89b..88af18e84c6 100644
--- a/src/microsoft/compiler/dxil_function.c
+++ b/src/microsoft/compiler/dxil_function.c
@@ -43,6 +43,7 @@ static struct  predefined_func_descr predefined_funcs[] = {
 {"dx.op.tertiary", "O", "iOOO", DXIL_ATTR_KIND_READ_NONE},
 {"dx.op.threadId", "i", "ii", DXIL_ATTR_KIND_READ_NONE},
 {"dx.op.threadIdInGroup", "i", "ii", DXIL_ATTR_KIND_READ_NONE},
+{"dx.op.flattenedThreadIdInGroup", "i", "i", DXIL_ATTR_KIND_READ_NONE},
 {"dx.op.groupId", "i", "ii", DXIL_ATTR_KIND_READ_NONE},
 {"dx.op.unary", "O", "iO", DXIL_ATTR_KIND_READ_NONE},
 {"dx.op.unaryBits", "i", "iO", DXIL_ATTR_KIND_READ_NONE},
diff --git a/src/microsoft/compiler/nir_to_dxil.c b/src/microsoft/compiler/nir_to_dxil.c
index 300efbaabf7..55f3cbd66c5 100644
--- a/src/microsoft/compiler/nir_to_dxil.c
+++ b/src/microsoft/compiler/nir_to_dxil.c
@@ -249,6 +249,7 @@ enum dxil_intr {
    DXIL_INTR_THREAD_ID = 93,
    DXIL_INTR_GROUP_ID = 94,
    DXIL_INTR_THREAD_ID_IN_GROUP = 95,
+   DXIL_INTR_FLATTENED_THREAD_ID_IN_GROUP = 96,
 
    DXIL_INTR_EMIT_STREAM = 97,
    DXIL_INTR_CUT_STREAM = 98,
@@ -557,6 +558,26 @@ emit_threadidingroup_call(struct ntd_context *ctx,
    return dxil_emit_call(&ctx->mod, func, args, ARRAY_SIZE(args));
 }
 
+static const struct dxil_value *
+emit_flattenedthreadidingroup_call(struct ntd_context *ctx)
+{
+   const struct dxil_func *func = dxil_get_function(&ctx->mod, "dx.op.flattenedThreadIdInGroup", DXIL_I32);
+
+   if (!func)
+      return NULL;
+
+   const struct dxil_value *opcode = dxil_module_get_int32_const(&ctx->mod,
+      DXIL_INTR_FLATTENED_THREAD_ID_IN_GROUP);
+   if (!opcode)
+      return NULL;
+
+   const struct dxil_value *args[] = {
+     opcode
+   };
+
+   return dxil_emit_call(&ctx->mod, func, args, ARRAY_SIZE(args));
+}
+
 static const struct dxil_value *
 emit_groupid_call(struct ntd_context *ctx, const struct dxil_value *comp)
 {
@@ -2287,6 +2308,21 @@ emit_load_local_invocation_id(struct ntd_context *ctx,
    return true;
 }
 
+static bool
+emit_load_local_invocation_index(struct ntd_context *ctx,
+                                 nir_intrinsic_instr *intr)
+{
+   assert(intr->dest.is_ssa);
+
+   const struct dxil_value
+      *flattenedthreadidingroup = emit_flattenedthreadidingroup_call(ctx);
+   if (!flattenedthreadidingroup)
+      return false;
+   store_dest_value(ctx, &intr->dest, 0, flattenedthreadidingroup);
+   
+   return true;
+}
+
 static bool
 emit_load_local_workgroup_id(struct ntd_context *ctx,
                               nir_intrinsic_instr *intr)
@@ -3480,6 +3516,8 @@ emit_intrinsic(struct ntd_context *ctx, nir_intrinsic_instr *intr)
       return emit_load_global_invocation_id(ctx, intr);
    case nir_intrinsic_load_local_invocation_id:
       return emit_load_local_invocation_id(ctx, intr);
+   case nir_intrinsic_load_local_invocation_index:
+      return emit_load_local_invocation_index(ctx, intr);
    case nir_intrinsic_load_workgroup_id:
    case nir_intrinsic_load_workgroup_id_zero_base:
       return emit_load_local_workgroup_id(ctx, intr);