generate HW local IDs with forced linear walk order

scottp101 · igcbot · commit cf22c10703a1 · 2022-03-18T17:04:23.000+01:00
diff --git a/IGC/Compiler/CISACodeGen/ComputeShaderBase.cpp b/IGC/Compiler/CISACodeGen/ComputeShaderBase.cpp
@@ -161,40 +161,98 @@ namespace IGC
             m_walkOrder = WO_YXZ;
         }
 
-        uint UNDEF = 999;
+        auto order = selectBestWalkOrder(
+            m_ThreadIDLayout, is_pow2_x, is_pow2_y, is_pow2_z);
+
+        if (order) {
+            m_walkOrder = *order;
+            m_enableHWGenerateLID = true;
+        } else {
+            // Is 2D or 3D dispatch and isnt pow2, so the HW doesn't support it
+            m_enableHWGenerateLID = false;
+            m_ThreadIDLayout = ThreadIDLayout::X;
+            m_walkOrder = WO_XYZ;
+            return;
+        }
+    }
+
+    Optional<CComputeShaderBase::WALK_ORDER>
+    CComputeShaderBase::checkLegalWalkOrder(
+        const std::array<uint32_t, 3>& Dims,
+        const WorkGroupWalkOrderMD& WO)
+    {
+        auto is_pow2 = [](uint32_t dim) {
+            return iSTD::IsPowerOfTwo(dim);
+        };
+
+        const int walkorder_x = WO.dim0;
+        const int walkorder_y = WO.dim1;
+        const int walkorder_z = WO.dim2;
+
+        const uint32_t dim_x = Dims[0];
+        const uint32_t dim_y = Dims[1];
+        const uint32_t dim_z = Dims[2];
+
+        uint order0 = (walkorder_x == 0) ? 0 : (walkorder_y == 0) ? 1 : 2;
+        uint order1 = (walkorder_x == 1) ? 0 : (walkorder_y == 1) ? 1 : 2;
+
+        if (order0 != order1
+            && ((order0 == 0 && is_pow2(dim_x))
+                || (order0 == 1 && is_pow2(dim_y))
+                || (order0 == 2 && is_pow2(dim_z)))
+            && ((order1 == 0 && is_pow2(dim_x))
+                || (order1 == 1 && is_pow2(dim_y))
+                || (order1 == 2 && is_pow2(dim_z)))
+            )
+        {
+            // Legal walk order for HW auto-gen
+            return getWalkOrder(order0, order1);
+        }
+
+        return None;
+    }
+
+    Optional<CComputeShaderBase::WALK_ORDER>
+    CComputeShaderBase::selectBestWalkOrder(
+        ThreadIDLayout Layout,
+        bool is_pow2_x, bool is_pow2_y, bool is_pow2_z)
+    {
+        constexpr uint UNDEF = std::numeric_limits<uint>::max();
         uint order0 = UNDEF;
         uint order1 = UNDEF;
-        if (m_ThreadIDLayout == ThreadIDLayout::TileY) {
+        if (Layout == ThreadIDLayout::TileY)
+        {
             IGC_ASSERT(is_pow2_y);
             order0 = 1;
             order1 = (is_pow2_x ? 0 : (is_pow2_z ? 2 : UNDEF));
-        }else {
+        }
+        else
+        {
             //below is from HAS p-code except tileY
             //try to find walk_order so that HW can generate LID
-            if (is_pow2_x) {
+            if (is_pow2_x)
+            {
                 // (pow2,pow2,z) or (pow2,y,pow2) or illegal
                 order0 = 0;
                 order1 = (is_pow2_y ? 1 : (is_pow2_z ? 2 : UNDEF));
-            }else if (is_pow2_y) {
+            }
+            else if (is_pow2_y)
+            {
                 // (x,pow2,pow2) or illegal
                 order0 = 1;
                 order1 = (is_pow2_z ? 2 : UNDEF);
             }
         }
 
-        if (order1 != UNDEF) {
+        if (order1 != UNDEF)
+        {
             // select walkorder
-            m_walkOrder = getWalkOrder(order0, order1);
-            m_enableHWGenerateLID = true;
-
-        }else {
-            // Is 2D or 3D dispatch and isnt pow2, so the HW doesn't support it
-            m_enableHWGenerateLID = false;
-            m_ThreadIDLayout = ThreadIDLayout::X;
-            m_walkOrder = WO_XYZ;
-            return;
+            return getWalkOrder(order0, order1);
         }
+
+        return None;
     }
+
     //order0: the internal walk dim
     //order1: the intermediate walk dim
     //e.g.: 1, 0 means, YXZ walkorder
@@ -209,8 +267,8 @@ namespace IGC
         case getWalkOrderValue(0, 1): return WALK_ORDER::WO_XYZ; //012
         case getWalkOrderValue(0, 2): return WALK_ORDER::WO_XZY; //021
         case getWalkOrderValue(1, 0): return WALK_ORDER::WO_YXZ; //102
-        case getWalkOrderValue(1, 2): return WALK_ORDER::WO_YZX; //201
-        case getWalkOrderValue(2, 0): return WALK_ORDER::WO_ZXY; //120
+        case getWalkOrderValue(1, 2): return WALK_ORDER::WO_YZX; //120
+        case getWalkOrderValue(2, 0): return WALK_ORDER::WO_ZXY; //201
         case getWalkOrderValue(2, 1): return WALK_ORDER::WO_ZYX; //210
         default:
             IGC_ASSERT_MESSAGE(0, "unhandled case!");
diff --git a/IGC/Compiler/CISACodeGen/ComputeShaderBase.hpp b/IGC/Compiler/CISACodeGen/ComputeShaderBase.hpp
@@ -10,6 +10,9 @@ SPDX-License-Identifier: MIT
 #pragma once
 
 #include "Compiler/CISACodeGen/ShaderCodeGen.hpp"
+#include "common/LLVMWarningsPush.hpp"
+#include "llvm/ADT/Optional.h"
+#include "common/LLVMWarningsPop.hpp"
 
 namespace IGC
 {
@@ -53,7 +56,14 @@ namespace IGC
         bool m_enableHWGenerateLID = false;
 
         void setEmitLocalMask(SGVUsage channelNum);
-    private:
+        static llvm::Optional<WALK_ORDER> selectBestWalkOrder(
+            ThreadIDLayout Layout,
+            bool is_pow2_x, bool is_pow2_y, bool is_pow2_z);
+        // Determines if HW can handle auto generating local IDs with this
+        // order
+        static llvm::Optional<WALK_ORDER> checkLegalWalkOrder(
+            const std::array<uint32_t, 3>& Dims,
+            const WorkGroupWalkOrderMD& WO);
         static WALK_ORDER getWalkOrder(uint order0, uint order1);
     };
 }
diff --git a/IGC/Compiler/CISACodeGen/OpenCLKernelCodeGen.cpp b/IGC/Compiler/CISACodeGen/OpenCLKernelCodeGen.cpp
@@ -159,51 +159,74 @@ namespace IGC
     }
     void COpenCLKernel::tryHWGenerateLocalIDs()
     {
-        if (hasWorkGroupWalkOrder())
-            return;
-
         auto Dims = IGCMetaDataHelper::getThreadGroupDims(
             *m_pMdUtils, entry);
 
         if (!Dims)
             return;
 
+        auto WO = getWorkGroupWalkOrder();
+        bool ForcedWalkOrder = false;
+        if (WO.dim0 != 0 || WO.dim1 != 0 || WO.dim2 != 0)
+        {
+            if (auto Order = checkLegalWalkOrder(*Dims, WO))
+            {
+                ForcedWalkOrder = true;
+                // Don't do TileY if forced in this way.
+                m_ThreadIDLayout = ThreadIDLayout::X;
+                m_walkOrder = *Order;
+            }
+            else
+            {
+                auto WalkOrder = getWalkOrder(WO.dim0, WO.dim1);
+                if (WalkOrder != WO_XYZ)
+                {
+                    IGC_ASSERT_MESSAGE(0, "unhandled walk order!");
+                }
+                return;
+            }
+        }
+
         // OpenCL currently emits all local IDs even if only one dimension
         // is requested. Let's mirror that for now.
         ImplicitArgs implicitArgs(*entry, m_pMdUtils);
         if (implicitArgs.isImplicitArgExist(ImplicitArg::LOCAL_ID_X) ||
             implicitArgs.isImplicitArgExist(ImplicitArg::LOCAL_ID_Y) ||
             implicitArgs.isImplicitArgExist(ImplicitArg::LOCAL_ID_Z))
         {
+            if (ForcedWalkOrder)
+                m_enableHWGenerateLID = true;
             setEmitLocalMask(THREAD_ID_IN_GROUP_Z);
         }
 
-        selectWalkOrder(
-            false,
-            0,
-            0,
-            0, /* dummy 1D accesses */
-            0, /* dummy 2D accesses */
-            0, /* dummy SLM accessed */
-            (*Dims)[0],
-            (*Dims)[1],
-            (*Dims)[2]);
+        if (!ForcedWalkOrder)
+        {
+            selectWalkOrder(
+                false,
+                0,
+                0,
+                0, /* dummy 1D accesses */
+                0, /* dummy 2D accesses */
+                0, /* dummy SLM accessed */
+                (*Dims)[0],
+                (*Dims)[1],
+                (*Dims)[2]);
+        }
         encoder.GetVISABuilder()->SetOption(vISA_autoLoadLocalID, m_enableHWGenerateLID);
     }
 
-    bool COpenCLKernel::hasWorkGroupWalkOrder()
+    WorkGroupWalkOrderMD COpenCLKernel::getWorkGroupWalkOrder()
     {
         const CodeGenContext* pCtx = GetContext();
         const ModuleMetaData* MMD = pCtx->getModuleMetaData();
         if (auto I = MMD->FuncMD.find(entry); I != MMD->FuncMD.end())
         {
             auto& FMD = I->second;
             auto& Order = FMD.workGroupWalkOrder;
-            if (Order.dim0 != 0 || Order.dim1 != 0 || Order.dim2 != 0)
-                return true;
+            return Order;
         }
 
-        return false;
+        return {};
     }
 
     SOpenCLKernelInfo::SResourceInfo COpenCLKernel::getResourceInfo(int argNo)
diff --git a/IGC/Compiler/CISACodeGen/OpenCLKernelCodeGen.hpp b/IGC/Compiler/CISACodeGen/OpenCLKernelCodeGen.hpp
@@ -126,7 +126,7 @@ namespace IGC
 
         void ClearKernelInfo();
     private:
-        bool hasWorkGroupWalkOrder();
+        WorkGroupWalkOrderMD getWorkGroupWalkOrder();
         void tryHWGenerateLocalIDs();
     };
 
