From 4320b9cd61f8eb16d64a0b133161a4f37f644754 Mon Sep 17 00:00:00 2001
From: Xavi Zhang <135793234+xazhangAMD@users.noreply.github.com>
Date: Fri, 10 Nov 2023 20:28:25 -0500
Subject: [PATCH] lgc: fix wrong thread id when group size Y/Z is 1 (#2812)

* lgc: fix wrong thread id when group size Y/Z is 1

tidigCompCnt is set to 0 when y and z are 1 and 1 when z is 1. In the
random failed case, the group size is 1x1x1, so the y and z ids are just
random values.
---
 lgc/patch/PatchEntryPointMutate.cpp | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/lgc/patch/PatchEntryPointMutate.cpp b/lgc/patch/PatchEntryPointMutate.cpp
index f5140664ca..96b5899e72 100644
--- a/lgc/patch/PatchEntryPointMutate.cpp
+++ b/lgc/patch/PatchEntryPointMutate.cpp
@@ -378,7 +378,7 @@ void PatchEntryPointMutate::lowerGroupMemcpy(GroupMemcpyOp &groupMemcpyOp) {
         // The local invocation ID is packed to VGPR0 on GFX11+ with the following layout:
         //
         //   +-----------------------+-----------------------+-----------------------+
-        //   | Local Invocation ID Z | Local Invocation ID Y | Local Invocation ID Z |
+        //   | Local Invocation ID Z | Local Invocation ID Y | Local Invocation ID X |
         //   | [29:20]               | [19:10]               | [9:0]                 |
         //   +-----------------------+-----------------------+-----------------------+
         // localInvocationIdZ = localInvocationId[29:20]
@@ -391,9 +391,13 @@ void PatchEntryPointMutate::lowerGroupMemcpy(GroupMemcpyOp &groupMemcpyOp) {
 
       // LocalInvocationIndex is
       // (LocalInvocationId.Z * WorkgroupSize.Y + LocalInvocationId.Y) * WorkGroupSize.X + LocalInvocationId.X
-      threadIndex = builder.CreateMul(threadIdComp[2], builder.getInt32(workgroupSize[1]));
-      threadIndex = builder.CreateAdd(threadIndex, threadIdComp[1]);
-      threadIndex = builder.CreateMul(threadIndex, builder.getInt32(workgroupSize[0]));
+      // tidigCompCnt is not always 3 if groupSizeY and/or groupSizeZ are 1. See RegisterMetadataBuilder.cpp.
+      threadIndex = builder.getInt32(0);
+      if (workgroupSize[2] > 1)
+        threadIndex = builder.CreateMul(threadIdComp[2], builder.getInt32(workgroupSize[1]));
+      if (workgroupSize[1] > 1)
+        threadIndex =
+            builder.CreateMul(builder.CreateAdd(threadIndex, threadIdComp[1]), builder.getInt32(workgroupSize[0]));
       threadIndex = builder.CreateAdd(threadIndex, threadIdComp[0]);
     }
   } else {