Add pass to Split Indirect EE to sel to avoid VxH mov.

It's a separate pass, as it wants to capture result of LowerGEPIntoPrivMem.

Heuristic currently handles calculating cost cases:
- less than <4 x type> vec size
- when EEI index is a multiply(+ add) of some ConstantInt as in below example.

From:
  %268 = mul nuw i32 %res.i2.i, 3
  %269 = extractelement <12 x float> %234, i32 %268
  %274 = extractelement <12 x float> %198, i32 %268
  %270 = add i32 %268, 1
  %271 = extractelement <12 x float> %234, i32 %270
  %275 = extractelement <12 x float> %198, i32 %270
  %272 = add i32 %268, 2
  %273 = extractelement <12 x float> %234, i32 %272
  %276 = extractelement <12 x float> %198, i32 %272

To:
  %250 = icmp eq i32 %res.i2.i, i16 1
  %251 = select i1 %250, float %206, float %200
  %252 = select i1 %250, float %208, float %202
  %253 = select i1 %250, float %210, float %204
  %254 = select i1 %250, float %48, float %32
  %255 = select i1 %250, float %49, float %33
  %256 = select i1 %250, float %50, float %34
  %257 = icmp eq i32 %res.i2.i, i16 2
  %258 = select i1 %257, float %214, float %251
  %259 = select i1 %257, float %215, float %252
  %260 = select i1 %257, float %216, float %253
  %261 = select i1 %257, float %64, float %254
  %262 = select i1 %257, float %65, float %255
  %263 = select i1 %257, float %66, float %256

Change-Id: Ibd4ec1da3b81f18cef6a405d5f599d86eb36c900

Author: bgajdaINTC

IGC/Compiler/CISACodeGen/ShaderCodeGen.cpp

@@ -551,6 +551,12 @@ static void AddLegalizationPasses(CodeGenContext& ctx, IGCPassManager& mpm, PSSi
     // Since we don't support switch statements, switch lowering is needed after the last CFG simplication
     mpm.add(llvm::createLowerSwitchPass());
 
+    // There's no particular reason for this exact place, but it should be after LowerGEPForPrivMem
+    if (IGC_IS_FLAG_ENABLED(EnableSplitIndirectEEtoSel))
+    {
+        mpm.add(createSplitIndirectEEtoSelPass());
+    }
+
     // Split big vector & 3-element load/store, etc.
     mpm.add(createVectorPreProcessPass());
 

IGC/Compiler/CustomSafeOptPass.cpp

@@ -45,6 +45,7 @@ The passes are
     VectorBitCastOpt
     GenStrengthReduction
     FlattenSmallSwitch
+    SplitIndirectEEtoSel
 
 CustomSafeOptPass does peephole optimizations
 For example, reduce the alloca size so there is a chance to promote indexed temp.
@@ -69,6 +70,9 @@ GenStrengthReduction performs a fdiv optimization.
 FlattenSmallSwitch flatten the if/else or switch structure and use cmp+sel
 instead if the structure is small.
 
+SplitIndirectEEtoSel splits extractelements with very small vec to a series of
+cmp+sel to avoid expensive VxH mov.
+
 =============================================================================*/
 
 #include "Compiler/CustomSafeOptPass.hpp"
@@ -4460,6 +4464,210 @@ void FCmpPaternMatch::visitSelectInst(SelectInst& I)
 IGC_INITIALIZE_PASS_BEGIN(FlattenSmallSwitch, "flattenSmallSwitch", "flattenSmallSwitch", false, false)
 IGC_INITIALIZE_PASS_END(FlattenSmallSwitch, "flattenSmallSwitch", "flattenSmallSwitch", false, false)
 
+
+
+/*======================== SplitIndirectEEtoSel =============================
+
+This class changes extract element for small vectors to series of cmp+sel to avoid VxH mov.
+before:
+  %268 = mul nuw i32 %res.i2.i, 3
+  %269 = extractelement <12 x float> %234, i32 %268
+  %270 = add i32 %268, 1
+  %271 = extractelement <12 x float> %234, i32 %270
+  %272 = add i32 %268, 2
+  %273 = extractelement <12 x float> %234, i32 %272
+  %274 = extractelement <12 x float> %198, i32 %268
+  %275 = extractelement <12 x float> %198, i32 %270
+  %276 = extractelement <12 x float> %198, i32 %272
+
+after:
+  %250 = icmp eq i32 %res.i2.i, i16 1
+  %251 = select i1 %250, float %206, float %200
+  %252 = select i1 %250, float %208, float %202
+  %253 = select i1 %250, float %210, float %204
+  %254 = select i1 %250, float %48, float %32
+  %255 = select i1 %250, float %49, float %33
+  %256 = select i1 %250, float %50, float %34
+  %257 = icmp eq i32 %res.i2.i, i16 2
+  %258 = select i1 %257, float %214, float %251
+  %259 = select i1 %257, float %215, float %252
+  %260 = select i1 %257, float %216, float %253
+  %261 = select i1 %257, float %64, float %254
+  %262 = select i1 %257, float %65, float %255
+  %263 = select i1 %257, float %66, float %256
+
+  It is a bit similar to SimplifyConstant::isCmpSelProfitable for OCL, but not restricted to api.
+  And to GenSimplification::visitExtractElement() but not restricted to vec of 2, and later.
+  TODO: for known vectors check how many unique items there are.
+===========================================================================*/
+namespace {
+    class SplitIndirectEEtoSel : public FunctionPass, public llvm::InstVisitor<SplitIndirectEEtoSel>
+    {
+    public:
+        static char ID;
+        SplitIndirectEEtoSel() : FunctionPass(ID)
+        {
+            initializeSplitIndirectEEtoSelPass(*PassRegistry::getPassRegistry());
+        }
+        virtual llvm::StringRef getPassName() const { return "Split Indirect EE to ICmp Plus Sel"; }
+        virtual bool runOnFunction(Function& F);
+        void visitExtractElementInst(llvm::ExtractElementInst& I);
+    private:
+        bool isProfitableToSplit(uint64_t num, int64_t mul, int64_t add);
+        bool didSomething;
+    };
+
+} // namespace
+
+
+char SplitIndirectEEtoSel::ID = 0;
+FunctionPass* IGC::createSplitIndirectEEtoSelPass() { return new SplitIndirectEEtoSel(); }
+
+bool SplitIndirectEEtoSel::runOnFunction(Function& F)
+{
+    didSomething = false;
+    visit(F);
+    return didSomething;
+}
+
+bool SplitIndirectEEtoSel::isProfitableToSplit(uint64_t num, int64_t mul, int64_t add)
+{
+    /* Assumption:
+       Pass is profitable when: (X * cmp + Y * sel) < (ExecSize * mov VxH).
+    */
+
+    const int64_t assumedVXHCost = IGC_GET_FLAG_VALUE(SplitIndirectEEtoSelThreshold);
+    int64_t possibleCost = 0;
+
+    /* for: extractelement <4 x float> , %index
+       cost is (4 - 1)  * (icmp + sel) = 6;
+    */
+    possibleCost = ((int64_t)num -1) * 2;
+    if (possibleCost < assumedVXHCost)
+        return true;
+
+    /* for: extractelement <12 x float> , (mul %real_index, 3)
+       cost is ((12/3) - 1) * (icmp + sel) = 6;
+    */
+
+    if (mul > 0) // not tested negative options
+    {
+        int64_t differentOptions = 1 + ((int64_t)num - 1) / mul; // ceil(num/mul)
+        possibleCost = (differentOptions - 1) * 2;
+
+        if (possibleCost < assumedVXHCost)
+            return true;
+    }
+
+    return false;
+}
+
+void SplitIndirectEEtoSel::visitExtractElementInst(llvm::ExtractElementInst& I)
+{
+    using namespace llvm::PatternMatch;
+
+    VectorType* vecTy = I.getVectorOperandType();
+    uint64_t num = vecTy->getNumElements();
+    Type* eleType = vecTy->getElementType();
+
+    Value* vec = I.getVectorOperand();
+    Value* index = I.getIndexOperand();
+
+    // ignore constant index
+    if (dyn_cast<ConstantInt>(index))
+    {
+        return;
+    }
+
+    // ignore others for now (did not yet evaluate perf. impact)
+    if (!(eleType->isIntegerTy(32) || eleType->isFloatTy()))
+    {
+        return;
+    }
+
+    // used to calculate offsets
+    int64_t add = 0;
+    int64_t mul = 1;
+
+
+    /* strip mul/add from index calculation and remember it for later:
+       %268 = mul nuw i32 %res.i2.i, 3
+       %270 = add i32 %268, 1
+       %271 = extractelement <12 x float> %234, i32 %270
+    */
+    Value* Val1 = nullptr;
+    ConstantInt* ci_add = nullptr;
+    ConstantInt* ci_mul = nullptr;
+
+    auto pat1 = m_Add(m_Mul(m_Value(Val1), m_ConstantInt(ci_mul)), m_ConstantInt(ci_add));
+    auto pat2 = m_Mul(m_Value(Val1), m_ConstantInt(ci_mul));
+    // Some code shows `shl+or` instead of mul+add.
+    auto pat21 = m_Or(m_Shl(m_Value(Val1), m_ConstantInt(ci_mul)), m_ConstantInt(ci_add));
+    auto pat22 = m_Shl(m_Value(Val1), m_ConstantInt(ci_mul));
+
+    if (match(index, pat1) || match(index, pat2))
+    {
+        add = ci_add ? ci_add->getSExtValue() : 0;
+        mul = ci_mul ? ci_mul->getSExtValue() : 1;
+        index = Val1;
+    }
+    else if (match(index, pat21) || match(index, pat22))
+    {
+        add = ci_add ? ci_add->getSExtValue() : 0;
+        mul = ci_mul ? (1LL << ci_mul->getSExtValue()) : 1LL;
+        index = Val1;
+    }
+
+    if (!isProfitableToSplit(num, mul, add))
+        return;
+
+    Value* vTemp = llvm::UndefValue::get(eleType);
+    IRBuilder<> builder(I.getNextNode());
+
+    // returns true if we can skip this icmp, such as:
+    // icmp eq (add (mul %index, 3), 2), 1
+    // icmp eq (mul %index, 3), 1
+    auto canSafelySkipThis = [&](int64_t add, int64_t mul, int64_t & newIndex) {
+        if (mul)
+        {
+            newIndex -= add;
+            if ((newIndex % mul) != 0)
+                return true;
+            newIndex = newIndex / mul;
+        }
+        return false;
+    };
+
+    // Generate combinations
+    for (uint64_t elemIndex = 0; elemIndex < num; elemIndex++)
+    {
+        int64_t cmpIndex = elemIndex;
+
+        if (canSafelySkipThis(add, mul, cmpIndex))
+            continue;
+
+        // Those 2 might be different, when cmp will get altered by it's operands, but EE index stays the same
+        ConstantInt* cmpIndexCI = llvm::ConstantInt::get(builder.getInt32Ty(), (uint64_t)cmpIndex);
+        ConstantInt* eeiIndexCI = llvm::ConstantInt::get(builder.getInt32Ty(), (uint64_t)elemIndex);
+
+        Value* cmp = builder.CreateICmp(CmpInst::Predicate::ICMP_EQ, index, cmpIndexCI);
+        Value* subcaseEE = builder.CreateExtractElement(vec, eeiIndexCI);
+        Value* sel = builder.CreateSelect(cmp, subcaseEE, vTemp);
+        vTemp = sel;
+        didSomething = true;
+    }
+
+    // In theory there's no situation where we don't do something up to this point.
+    if (didSomething)
+    {
+        I.replaceAllUsesWith(vTemp);
+    }
+}
+
+
+IGC_INITIALIZE_PASS_BEGIN(SplitIndirectEEtoSel, "SplitIndirectEEtoSel", "SplitIndirectEEtoSel", false, false)
+IGC_INITIALIZE_PASS_END(SplitIndirectEEtoSel, "SplitIndirectEEtoSel", "SplitIndirectEEtoSel", false, false)
+
 ////////////////////////////////////////////////////////////////////////
 // LogicalAndToBranch trying to find logical AND like below:
 //    res = simpleCond0 && complexCond1

IGC/Compiler/CustomSafeOptPass.hpp

@@ -250,6 +250,7 @@ namespace IGC
     llvm::FunctionPass* createGenStrengthReductionPass();
     llvm::FunctionPass* createNanHandlingPass();
     llvm::FunctionPass* createFlattenSmallSwitchPass();
+    llvm::FunctionPass* createSplitIndirectEEtoSelPass();
     llvm::FunctionPass* createIGCIndirectICBPropagaionPass();
     llvm::FunctionPass* createBlendToDiscardPass();
     llvm::FunctionPass* createMarkReadOnlyLoadPass();

IGC/Compiler/InitializePasses.h

@@ -68,6 +68,7 @@ void initializeGenUpdateCBPass(llvm::PassRegistry&);
 void initializeGenStrengthReductionPass(llvm::PassRegistry&);
 void initializeNanHandlingPass(llvm::PassRegistry&);
 void initializeFlattenSmallSwitchPass(llvm::PassRegistry&);
+void initializeSplitIndirectEEtoSelPass(llvm::PassRegistry&);
 void initializeGenXFunctionGroupAnalysisPass(llvm::PassRegistry&);
 void initializeGenXCodeGenModulePass(llvm::PassRegistry&);
 void initializeEstimateFunctionSizePass(llvm::PassRegistry&);

IGC/common/igc_flags.def

@@ -135,6 +135,8 @@ DECLARE_IGC_REGKEY(bool, EnableSimplifyGEP,             true,  "Enable IGC to si
 DECLARE_IGC_REGKEY(bool, DisableCustomUnsafeOpt,        false, "Disable IGC to run custom unsafe optimizations", false)
 DECLARE_IGC_REGKEY(bool, EnableFastMath,                false, "Enable fast math optimizations in IGC", false)
 DECLARE_IGC_REGKEY(bool, DisableFlattenSmallSwitch,     false, "Disable the flatten small switch pass", false)
+DECLARE_IGC_REGKEY(bool, EnableSplitIndirectEEtoSel,    true,  "Enable the split indirect extractelement to icmp+sel pass", false)
+DECLARE_IGC_REGKEY(bool, SplitIndirectEEtoSelThreshold, 8,     "Split indirect extractelement cost threshold", false)
 DECLARE_IGC_REGKEY(bool, DisableImmConstantOpt,         false, "Disable IGC IndirectICBPropagaion optimization", false)
 DECLARE_IGC_REGKEY(DWORD,MaxImmConstantSizePushed,      256,   "Set the max size of immediate constant buffer pushed", false)
 DECLARE_IGC_REGKEY(bool, EnableCustomLoopVersioning,    true,  "Enable IGC to do custom loop versioning", false)