Loop sinking tuning to sink load chains

admitric · igcbot · commit 02e2e6b6c6dd · 2024-02-16T11:38:20.000+01:00
Prepopulate all load chains for a loop and sink the instructions with
awareness of the loads that are already in loop
diff --git a/IGC/Compiler/CISACodeGen/CodeSinking.cpp b/IGC/Compiler/CISACodeGen/CodeSinking.cpp
@@ -1440,8 +1440,11 @@ namespace IGC {
         MovedInsts.clear();
         UndoLocas.clear();
 
+        if (IGC_IS_FLAG_ENABLED(PrepopulateLoadChainLoopSink))
+            prepopulateLoadChains(L, LoadChains);
+
         bool AllowLoadSinking = IGC_IS_FLAG_ENABLED(ForceLoadsLoopSink);
-        bool AllowOnlyLoadChainSinking = false;
+        bool AllowOnlySingleUseLoadChainSinking = false;
         bool IterChanged = false;
 
         uint MaxLoopPressure = 0;
@@ -1476,7 +1479,7 @@ namespace IGC {
 
                 if (isLoopSinkCandidate(I, L, AllowLoadSinking))
                 {
-                    if (!AllowOnlyLoadChainSinking || isLoadChain(I, LoadChains))
+                    if (!AllowOnlySingleUseLoadChainSinking || isLoadChain(I, LoadChains, true))
                     {
                         SinkCandidates.push_back(I);
                     }
@@ -1510,7 +1513,7 @@ namespace IGC {
                     if (IGC_IS_FLAG_ENABLED(EnableLoadChainLoopSink) && !LoadChains.empty())
                     {
                         PrintDump("Allowing only chain sinking...\n");
-                        AllowOnlyLoadChainSinking = true;
+                        AllowOnlySingleUseLoadChainSinking = true;
                     }
                     else
                     {
@@ -1580,15 +1583,54 @@ namespace IGC {
 
     // Check that this instruction is a part of address calc
     // chain of an already sinked load
-    bool CodeSinking::isLoadChain(Instruction *I, SmallPtrSet<Instruction *, 32> &LoadChains)
+    bool CodeSinking::isLoadChain(Instruction *I, SmallPtrSet<Instruction *, 32> &LoadChains, bool EnsureSingleUser)
     {
         if (!isa<BinaryOperator>(I) && !isa<CastInst>(I))
             return false;
         User *InstrUser = IGCLLVM::getUniqueUndroppableUser(I);
-        if (!InstrUser)
+        if (EnsureSingleUser && !InstrUser)
             return false;
-        Instruction *UI = dyn_cast<Instruction>(InstrUser);
-        return UI && LoadChains.count(UI);
+
+        return std::all_of(I->user_begin(), I->user_end(),
+            [&](User *U)
+            {
+                Instruction *UI = dyn_cast<Instruction>(U);
+                return UI && LoadChains.count(UI);
+            });
+    }
+
+    // Prepopulate load chain with the loads that are already in the loop
+    void CodeSinking::prepopulateLoadChains(Loop *L, SmallPtrSet<Instruction *, 32> &LoadChains)
+    {
+        std::function<void(Value *)> addInstructionIfLoadChain = [&](Value *V)-> void
+        {
+            Instruction *I = dyn_cast<Instruction>(V);
+            if (!I)
+                return;
+
+            if (!L->contains(I))
+                return;
+
+            if (!isLoadChain(I, LoadChains))
+                return;
+
+            LoadChains.insert(I);
+            for (auto &U : I->operands()) {
+                addInstructionIfLoadChain(U);
+            }
+        };
+
+        for (BasicBlock *BB: L->blocks())
+        {
+            for (Instruction &I : *BB)
+            {
+                if (LoadInst *LI = dyn_cast<LoadInst>(&I))
+                {
+                    LoadChains.insert(&I);
+                    addInstructionIfLoadChain(LI->getPointerOperand());
+                }
+            }
+        }
     }
 
     bool CodeSinking::isLoopSinkCandidate(Instruction *I, Loop *L, bool AllowLoadSinking)
diff --git a/IGC/Compiler/CISACodeGen/CodeSinking.hpp b/IGC/Compiler/CISACodeGen/CodeSinking.hpp
@@ -77,7 +77,9 @@ namespace IGC {
             llvm::SmallPtrSetImpl<llvm::Instruction*>& Stores);
         bool isSafeToLoopSinkLoad(llvm::Instruction* I, llvm::Loop* Loop, llvm::AliasAnalysis* AA);
         bool isAlwaysSinkInstruction(llvm::Instruction* I);
-        bool isLoadChain(llvm::Instruction* I, SmallPtrSet<Instruction *, 32> &LoadChains);
+        bool isLoadChain(llvm::Instruction* I, SmallPtrSet<Instruction*, 32>& LoadChains, bool EnsureSingleUser=false);
+        void prepopulateLoadChains(llvm::Loop* I, SmallPtrSet<Instruction*, 32>& LoadChains);
+
 
         /// rollback sinking. Uses MovedInsts and UndoLocas members implicitly
         void rollbackSinking(bool ReverseOrder, llvm::BasicBlock* BB);
diff --git a/IGC/Compiler/tests/CodeSinking/LoopSinking/load-chain-sinking.ll b/IGC/Compiler/tests/CodeSinking/LoopSinking/load-chain-sinking.ll
@@ -0,0 +1,72 @@
+;=========================== begin_copyright_notice ============================
+;
+; Copyright (C) 2024 Intel Corporation
+;
+; SPDX-License-Identifier: MIT
+;
+;============================ end_copyright_notice =============================
+; REQUIRES: regkeys
+; RUN: igc_opt --regkey ForceLoopSink=1 --regkey LoopSinkMinSave=1 --regkey LoopSinkThresholdDelta=10 --regkey CodeSinkingMinSize=10 --regkey PrepopulateLoadChainLoopSink=1 --igc-wi-analysis %enable-basic-aa% --igc-code-sinking -S %s | FileCheck %s
+
+; Check that address computations are considered beneficial and are sinked
+; even if the loads are already in the loop (when PrepopulateLoadChainLoopSink=1 is passed)
+
+define spir_kernel void @dsmm23x23x23(i16 %localIdX) {
+; CHECK-LABEL: @dsmm23x23x23(
+; CHECK:       for.body:
+; CHECK:         [[SINK_IDXPROM36:%.*]] = zext i16 [[LOCALIDX:%.*]] to i64
+; CHECK:         [[SINK_317:%.*]] = shl nuw nsw i64 [[SINK_IDXPROM36]], 3
+; CHECK:         [[TMP40:%.*]] = add i64 [[TMP39:%.*]], [[SINK_317]]
+; CHECK:         [[TMP41:%.*]] = inttoptr i64 [[TMP40]] to double addrspace(2)*
+; CHECK:         [[TMP42:%.*]] = load double, double addrspace(2)* [[TMP41]], align 8
+; CHECK:         [[SINK_CONV2:%.*]] = zext i16 [[LOCALIDX]] to i32
+; CHECK:         [[SINK_ADD33_1:%.*]] = add nuw nsw i32 [[SINK_CONV2]], 23
+; CHECK:         [[SINK_IDXPROM34_1:%.*]] = zext i32 [[SINK_ADD33_1]] to i64
+; CHECK:         [[SINK_316:%.*]] = shl nuw nsw i64 [[SINK_IDXPROM34_1]], 3
+; CHECK:         [[TMP43:%.*]] = add i64 [[TMP39]], [[SINK_316]]
+; CHECK:         [[TMP44:%.*]] = inttoptr i64 [[TMP43]] to double addrspace(2)*
+; CHECK:         [[TMP45:%.*]] = load double, double addrspace(2)* [[TMP44]], align 8
+; CHECK:         [[SINK_ADD33_2:%.*]] = add nuw nsw i32 [[SINK_CONV2]], 46
+; CHECK:         [[SINK_IDXPROM34_2:%.*]] = zext i32 [[SINK_ADD33_2]] to i64
+; CHECK:         [[SINK_315:%.*]] = shl nuw nsw i64 [[SINK_IDXPROM34_2]], 3
+; CHECK:         [[TMP48:%.*]] = add i64 [[TMP39]], [[SINK_315]]
+; CHECK:         [[TMP49:%.*]] = inttoptr i64 [[TMP48]] to double addrspace(2)*
+; CHECK:         [[TMP50:%.*]] = load double, double addrspace(2)* [[TMP49]], align 8
+; CHECK:         [[SINK_ADD33_3:%.*]] = add nuw nsw i32 [[SINK_CONV2]], 69
+; CHECK:         [[SINK_IDXPROM34_3:%.*]] = zext i32 [[SINK_ADD33_3]] to i64
+; CHECK:         [[SINK_314:%.*]] = shl nuw nsw i64 [[SINK_IDXPROM34_3]], 3
+; CHECK:         [[TMP51:%.*]] = add i64 [[TMP39]], [[SINK_314]]
+; CHECK:         [[TMP52:%.*]] = inttoptr i64 [[TMP51]] to double addrspace(2)*
+; CHECK:         [[TMP53:%.*]] = load double, double addrspace(2)* [[TMP52]], align 8
+
+entry:
+  %conv2 = zext i16 %localIdX to i32
+  %idxprom36 = zext i16 %localIdX to i64
+  %add33.1 = add nuw nsw i32 %conv2, 23
+  %idxprom34.1 = zext i32 %add33.1 to i64
+  %add33.2 = add nuw nsw i32 %conv2, 46
+  %idxprom34.2 = zext i32 %add33.2 to i64
+  %add33.3 = add nuw nsw i32 %conv2, 69
+  %idxprom34.3 = zext i32 %add33.3 to i64
+  %0 = shl nuw nsw i64 %idxprom36, 3
+  %1 = shl nuw nsw i64 %idxprom34.1, 3
+  %2 = shl nuw nsw i64 %idxprom34.2, 3
+  %3 = shl nuw nsw i64 %idxprom34.3, 3
+  br label %for.body
+
+for.body:                                         ; preds = %for.body, %entry
+  %4 = add i64 0, 0
+  %5 = add i64 %4, %0
+  %6 = inttoptr i64 %5 to double addrspace(2)*
+  %7 = load double, double addrspace(2)* %6, align 8
+  %8 = add i64 %4, %1
+  %9 = inttoptr i64 %8 to double addrspace(2)*
+  %10 = load double, double addrspace(2)* %9, align 8
+  %11 = add i64 %4, %2
+  %12 = inttoptr i64 %11 to double addrspace(2)*
+  %13 = load double, double addrspace(2)* %12, align 8
+  %14 = add i64 %4, %3
+  %15 = inttoptr i64 %14 to double addrspace(2)*
+  %16 = load double, double addrspace(2)* %15, align 8
+  br label %for.body
+}
diff --git a/IGC/common/igc_flags.h b/IGC/common/igc_flags.h
@@ -134,7 +134,8 @@ DECLARE_IGC_REGKEY(bool, DisableLoopSink,               false, "Disable sinking
 DECLARE_IGC_REGKEY(bool, ForceLoopSink,                 false, "Force sinking in all loops", false)
 DECLARE_IGC_REGKEY(bool, EnableLoadsLoopSink,           true, "Allow sinking of loads in the loop", false)
 DECLARE_IGC_REGKEY(bool, ForceLoadsLoopSink,            false, "Force sinking of loads in the loop from the beginning", false)
-DECLARE_IGC_REGKEY(bool, EnableLoadChainLoopSink,       true, "Allow sinking of load address calculation when the load was sinked to the loop, even if the needed regpressure is achieved", false)
+DECLARE_IGC_REGKEY(bool, PrepopulateLoadChainLoopSink,  true, "Check the loop for loop chains before sinking to use the existing chains in a heuristic", false)
+DECLARE_IGC_REGKEY(bool, EnableLoadChainLoopSink,       true, "Allow sinking of load address calculation when the load was sinked to the loop, even if the needed regpressure is achieved (only single use instructions)", false)
 DECLARE_IGC_REGKEY(DWORD, LoopSinkRegpressureMargin,    10,  "Sink into the loop until the pressure becomes less than #grf-margin", false)
 DECLARE_IGC_REGKEY(DWORD, CodeSinkingMinSize,           32,  "Don't sink if the number of instructions in the kernel is less", false)
 DECLARE_IGC_REGKEY(DWORD, CodeSinkingLoadSchedulingInstr, 20,  "Instructions number to step to schedule loads in advance before the load use to cover latency. 1 to insert it immediately before use", false)
