[VPlan] Replace disjoint or with add instead of dropping disjoint. #83821

fhahn · 2024-03-04T11:06:51Z

Dropping disjoint from an OR may yield incorrect results, as some analysis may have converted it to an Add implicitly (e.g. SCEV used for dependence analysis). Instead, replace it with an equivalent Add.

This is possible as all users of the disjoint OR only access lanes where the operands are disjoint or poison otherwise.

Note that replacing all disjoint ORs with ADDs instead of dropping the flags is not strictly necessary. It is only needed for disjoint ORs that SCEV treated as ADDs, but those are not tracked.

There are other places that may drop poison-generating flags; those likely need similar treatment.

Fixes #81872

Dropping disjoint from an OR may yield incorrect results, as some analysis may have converted it to an Add implicitly (e.g. SCEV used for dependence analysis). Instead, replace it with an equivalent Add. This is possible as all users of the disjoint OR only access lanes where the operands are disjoint or poison otherwise. Note that replacing all disjoint ORs with ADDs instead of dropping the flags is not strictly necessary. It is only needed for disjoint ORs that SCEV treated as ADDs, but those are not tracked.

llvmbot · 2024-03-04T11:07:39Z

@llvm/pr-subscribers-llvm-transforms

Author: Florian Hahn (fhahn)

Changes

Dropping disjoint from an OR may yield incorrect results, as some analysis may have converted it to an Add implicitly (e.g. SCEV used for dependence analysis). Instead, replace it with an equivalent Add.

This is possible as all users of the disjoint OR only access lanes where the operands are disjoint or poison otherwise.

Note that replacing all disjoint ORs with ADDs instead of dropping the flags is not strictly necessary. It is only needed for disjoint ORs that SCEV treated as ADDs, but those are not tracked.

There are other places that may drop poison-generating flags; those likely need similar treatment.

Full diff: https://github.com/llvm/llvm-project/pull/83821.diff

4 Files Affected:

(modified) llvm/lib/Transforms/Vectorize/LoopVectorizationPlanner.h (+3)
(modified) llvm/lib/Transforms/Vectorize/VPlan.h (+8)
(modified) llvm/lib/Transforms/Vectorize/VPlanPatternMatch.h (+23-4)
(modified) llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp (+18)

diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorizationPlanner.h b/llvm/lib/Transforms/Vectorize/LoopVectorizationPlanner.h
index a7ebf78e54ceb6..b94859864fff3c 100644
--- a/llvm/lib/Transforms/Vectorize/LoopVectorizationPlanner.h
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorizationPlanner.h
@@ -68,6 +68,9 @@ class VPBuilder {
 public:
   VPBuilder() = default;
   VPBuilder(VPBasicBlock *InsertBB) { setInsertPoint(InsertBB); }
+  VPBuilder(VPRecipeBase *InsertPt) {
+    setInsertPoint(InsertPt->getParent(), InsertPt->getIterator());
+  }
 
   /// Clear the insertion point: created instructions will not be inserted into
   /// a block.
diff --git a/llvm/lib/Transforms/Vectorize/VPlan.h b/llvm/lib/Transforms/Vectorize/VPlan.h
index 16c09a83e777dd..b565b4351e16d4 100644
--- a/llvm/lib/Transforms/Vectorize/VPlan.h
+++ b/llvm/lib/Transforms/Vectorize/VPlan.h
@@ -1127,6 +1127,12 @@ class VPRecipeWithIRFlags : public VPSingleDefRecipe {
     return WrapFlags.HasNSW;
   }
 
+  bool isDisjoint() const {
+    assert(OpType == OperationType::DisjointOp &&
+           "recipe cannot have a disjoing flag");
+    return DisjointFlags.IsDisjoint;
+  }
+
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
   void printFlags(raw_ostream &O) const;
 #endif
@@ -2136,6 +2142,8 @@ class VPReplicateRecipe : public VPRecipeWithIRFlags {
     assert(isPredicated() && "Trying to get the mask of a unpredicated recipe");
     return getOperand(getNumOperands() - 1);
   }
+
+  unsigned getOpcode() const { return getUnderlyingInstr()->getOpcode(); }
 };
 
 /// A recipe for generating conditional branches on the bits of a mask.
diff --git a/llvm/lib/Transforms/Vectorize/VPlanPatternMatch.h b/llvm/lib/Transforms/Vectorize/VPlanPatternMatch.h
index b90c588b607564..4b5b6b8cc3dbcf 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanPatternMatch.h
+++ b/llvm/lib/Transforms/Vectorize/VPlanPatternMatch.h
@@ -73,12 +73,12 @@ template <typename Op0_t, unsigned Opcode> struct UnaryVPInstruction_match {
   }
 };
 
-template <typename Op0_t, typename Op1_t, unsigned Opcode>
-struct BinaryVPInstruction_match {
+template <typename RecipeTy, typename Op0_t, typename Op1_t, unsigned Opcode>
+struct BinaryRecipe_match {
   Op0_t Op0;
   Op1_t Op1;
 
-  BinaryVPInstruction_match(Op0_t Op0, Op1_t Op1) : Op0(Op0), Op1(Op1) {}
+  BinaryRecipe_match(Op0_t Op0, Op1_t Op1) : Op0(Op0), Op1(Op1) {}
 
   bool match(const VPValue *V) {
     auto *DefR = V->getDefiningRecipe();
@@ -86,15 +86,27 @@ struct BinaryVPInstruction_match {
   }
 
   bool match(const VPRecipeBase *R) {
-    auto *DefR = dyn_cast<VPInstruction>(R);
+    auto *DefR = dyn_cast<RecipeTy>(R);
     if (!DefR || DefR->getOpcode() != Opcode)
       return false;
     assert(DefR->getNumOperands() == 2 &&
            "recipe with matched opcode does not have 2 operands");
     return Op0.match(DefR->getOperand(0)) && Op1.match(DefR->getOperand(1));
   }
+
+  bool match(const VPSingleDefRecipe *R) {
+    return match(static_cast<const VPRecipeBase *>(R));
+  }
 };
 
+template <typename Op0_t, typename Op1_t, unsigned Opcode>
+using BinaryVPInstruction_match =
+    BinaryRecipe_match<VPInstruction, Op0_t, Op1_t, Opcode>;
+
+template <typename Op0_t, typename Op1_t, unsigned Opcode>
+using BinaryVPReplicate_match =
+    BinaryRecipe_match<VPReplicateRecipe, Op0_t, Op1_t, Opcode>;
+
 template <unsigned Opcode, typename Op0_t>
 inline UnaryVPInstruction_match<Op0_t, Opcode>
 m_VPInstruction(const Op0_t &Op0) {
@@ -130,6 +142,13 @@ inline BinaryVPInstruction_match<Op0_t, Op1_t, VPInstruction::BranchOnCount>
 m_BranchOnCount(const Op0_t &Op0, const Op1_t &Op1) {
   return m_VPInstruction<VPInstruction::BranchOnCount>(Op0, Op1);
 }
+
+template <unsigned Opcode, typename Op0_t, typename Op1_t>
+inline BinaryVPReplicate_match<Op0_t, Op1_t, Opcode>
+m_VPReplicate(const Op0_t &Op0, const Op1_t &Op1) {
+  return BinaryVPReplicate_match<Op0_t, Op1_t, Opcode>(Op0, Op1);
+}
+
 } // namespace VPlanPatternMatch
 } // namespace llvm
 
diff --git a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
index 9d6deb802e2090..818647d5ea1bac 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp
@@ -1249,6 +1249,24 @@ void VPlanTransforms::dropPoisonGeneratingRecipes(
       // load/store. If the underlying instruction has poison-generating flags,
       // drop them directly.
       if (auto *RecWithFlags = dyn_cast<VPRecipeWithIRFlags>(CurRec)) {
+        VPValue *A, *B;
+        using namespace llvm::VPlanPatternMatch;
+        // Dropping disjoint from an OR may yield incorrect results, as some
+        // analysis may have converted it to an Add implicitly (e.g. SCEV used
+        // for dependence analysis). Instead, replace it with an equivalent Add.
+        // This is possible as all users of the disjoint OR only access lanes
+        // where the operands are disjoint or poison otherwise.
+        if (match(RecWithFlags,
+                  m_VPReplicate<Instruction::Or>(m_VPValue(A), m_VPValue(B))) &&
+            RecWithFlags->isDisjoint()) {
+          VPBuilder Builder(RecWithFlags);
+          VPInstruction *New = Builder.createOverflowingOp(
+              Instruction::Add, {A, B}, {false, false},
+              RecWithFlags->getDebugLoc());
+          RecWithFlags->replaceAllUsesWith(New);
+          RecWithFlags->eraseFromParent();
+          CurRec = New;
+        }
         RecWithFlags->dropPoisonGeneratingFlags();
       } else {
         Instruction *Instr = dyn_cast_or_null<Instruction>(

nikic

Looks reasonable conceptually, but I'm not the right person to review VPlan changes.

Test case seems to be missing?

annamthomas · 2024-03-04T20:50:38Z

Thank you for the fix! I'll land my test case from the other MR and close it off in favour of this one.

annamthomas · 2024-03-11T17:31:14Z

I've landed the test and the post commit comments. Pls rebase. Thanks.

…t-or-with-add

fhahn · 2024-03-18T15:13:04Z

@annamthomas thanks for landing the fix! Rebased, also the majority of the pattern-match changes have been split up and submitted & used separately.

github-actions · 2024-03-18T15:13:29Z

✅ With the latest revision this PR passed the C/C++ code formatter.

annamthomas · 2024-03-18T17:25:04Z

Change LGTM.

Could you pls link the commit to #81872 ? Thanks!

…oint. (#83821)" This reverts commit c2c1e6e. It creates a use after free. ==8342==ERROR: AddressSanitizer: heap-use-after-free on address 0x50f000001760 at pc 0x55b9fb84a8fb bp 0x7ffc18468a10 sp 0x7ffc18468a08 READ of size 1 at 0x50f000001760 thread T0 #0 0x55b9fb84a8fa in dropPoisonGeneratingFlags llvm/lib/Transforms/Vectorize/VPlan.h:1040:13 #1 0x55b9fb84a8fa in llvm::VPlanTransforms::dropPoisonGeneratingRecipes(llvm::VPlan&, llvm::function_ref<bool (llvm::BasicBlock*)>)::$_0::operator()(llvm::VPRecipeBase*) const llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp:1236:23 #2 0x55b9fb84a196 in llvm::VPlanTransforms::dropPoisonGeneratingRecipes(llvm::VPlan&, llvm::function_ref<bool (llvm::BasicBlock*)>) llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp Can be reproduced with asan on Transforms/LoopVectorize/AArch64/sve-interleaved-masked-accesses.ll Transforms/LoopVectorize/X86/pr81872.ll Transforms/LoopVectorize/X86/x86-interleaved-accesses-masked-group.ll

…lvm#83821) Dropping disjoint from an OR may yield incorrect results, as some analysis may have converted it to an Add implicitly (e.g. SCEV used for dependence analysis). Instead, replace it with an equivalent Add. This is possible as all users of the disjoint OR only access lanes where the operands are disjoint or poison otherwise. Note that replacing all disjoint ORs with ADDs instead of dropping the flags is not strictly necessary. It is only needed for disjoint ORs that SCEV treated as ADDs, but those are not tracked. There are other places that may drop poison-generating flags; those likely need similar treatment. Fixes llvm#81872 PR: llvm#83821

…oint. (llvm#83821)" This reverts commit c2c1e6e. It creates a use after free. ==8342==ERROR: AddressSanitizer: heap-use-after-free on address 0x50f000001760 at pc 0x55b9fb84a8fb bp 0x7ffc18468a10 sp 0x7ffc18468a08 READ of size 1 at 0x50f000001760 thread T0 #0 0x55b9fb84a8fa in dropPoisonGeneratingFlags llvm/lib/Transforms/Vectorize/VPlan.h:1040:13 #1 0x55b9fb84a8fa in llvm::VPlanTransforms::dropPoisonGeneratingRecipes(llvm::VPlan&, llvm::function_ref<bool (llvm::BasicBlock*)>)::$_0::operator()(llvm::VPRecipeBase*) const llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp:1236:23 #2 0x55b9fb84a196 in llvm::VPlanTransforms::dropPoisonGeneratingRecipes(llvm::VPlan&, llvm::function_ref<bool (llvm::BasicBlock*)>) llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp Can be reproduced with asan on Transforms/LoopVectorize/AArch64/sve-interleaved-masked-accesses.ll Transforms/LoopVectorize/X86/pr81872.ll Transforms/LoopVectorize/X86/x86-interleaved-accesses-masked-group.ll

danilaml · 2024-03-27T13:35:59Z

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

+          RecWithFlags->replaceAllUsesWith(New);
+          RecWithFlags->eraseFromParent();
+          CurRec = New;
+        }
        RecWithFlags->dropPoisonGeneratingFlags();


Should this be in the else branch?

It should be one of the two:

in else branch (to get rid of the use-after-free error since RecWithFlags is removed within the if branch above), OR

CurRec->dropPoisonGeneratingFlags

Makes sense to just have it in the else branch.

@fhahn the change was reverted because of the use-after-free above.

Missed that this was reverted, recommitted with a fix.

ayalz

Adding some post-commit (and post-revert) thoughts.

ayalz · 2024-03-27T18:33:56Z

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

+        // analysis may have converted it to an Add implicitly (e.g. SCEV used
+        // for dependence analysis). Instead, replace it with an equivalent Add.
+        // This is possible as all users of the disjoint OR only access lanes
+        // where the operands are disjoint or poison otherwise.


Should this be taken care of by VPRecipeWithIRFlags::dropPoisonGeneratingFlags(), i.e., have it replace the opcode of a disjoint Or with Add?

nit (independent): collectPoisonGeneratingInstrsInBackwardSlice() does more than collect Instrs, it also updates them.

Note that replacing all disjoint ORs with ADDs instead of dropping the flags is not strictly necessary. It is only needed for disjoint ORs that SCEV treated as ADDs, but those are not tracked.

Should code based on such SCEV expressions be generated (expanded(?)) directly?

(This case is reminiscent of dropping assumptions based on conditions subject to predication.)

Should this be taken care of by VPRecipeWithIRFlags::dropPoisonGeneratingFlags(), i.e., have it replace the opcode of a disjoint Or with Add?

Hmmm, it may work if there's no need to directly delete the old recipe, i.e. it will get cleaned up. Let me check.

Should code based on such SCEV expressions be generated (expanded(?)) directly?

(This case is reminiscent of dropping assumptions based on conditions subject to predication.)

At the moment, SCEV is only used for analysis of existing IR instructions in this case; switching use SCEV expansion for pointer expressions would probably introduce some other potential issues.

…sjoint. (#83821)" Recommit with a fix for the use-after-free causing the revert. This reverts the revert commit f872043. Original commit message: Dropping disjoint from an OR may yield incorrect results, as some analysis may have converted it to an Add implicitly (e.g. SCEV used for dependence analysis). Instead, replace it with an equivalent Add. This is possible as all users of the disjoint OR only access lanes where the operands are disjoint or poison otherwise. Note that replacing all disjoint ORs with ADDs instead of dropping the flags is not strictly necessary. It is only needed for disjoint ORs that SCEV treated as ADDs, but those are not tracked. There are other places that may drop poison-generating flags; those likely need similar treatment. Fixes #81872 PR: #83821

fhahn requested review from nikic, preames, annamthomas and ayalz March 4, 2024 11:06

llvmbot added vectorization llvm:transforms labels Mar 4, 2024

nikic reviewed Mar 4, 2024

View reviewed changes

fhahn mentioned this pull request Mar 4, 2024

[LV] Fix for miscompile with disjoint or #81922

Closed

fhahn added 3 commits March 18, 2024 14:40

Merge remote-tracking branch 'origin/main' into vplan-replace-disjoin…

df5468e

…t-or-with-add

!fixup fix merge, update test.

f4202fc

!fixup fix formatting

dc92c04

annamthomas approved these changes Mar 18, 2024

View reviewed changes

fhahn merged commit c2c1e6e into llvm:main Mar 19, 2024
4 checks passed

fhahn deleted the vplan-replace-disjoint-or-with-add branch March 19, 2024 19:16

danilaml reviewed Mar 27, 2024

View reviewed changes

ayalz reviewed Mar 27, 2024

View reviewed changes

patrick-rivos mentioned this pull request Apr 2, 2024

[LoopVectorize] Assertion `OpType == OperationType::DisjointOp && "recipe cannot have a disjoing flag"' failed. #87378

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[VPlan] Replace disjoint or with add instead of dropping disjoint. #83821

[VPlan] Replace disjoint or with add instead of dropping disjoint. #83821

fhahn commented Mar 4, 2024 •

edited

llvmbot commented Mar 4, 2024

nikic left a comment

annamthomas commented Mar 4, 2024

annamthomas commented Mar 11, 2024

fhahn commented Mar 18, 2024

github-actions bot commented Mar 18, 2024 •

edited

annamthomas commented Mar 18, 2024

danilaml Mar 27, 2024

annamthomas Mar 27, 2024

annamthomas Mar 27, 2024

fhahn Mar 27, 2024

ayalz left a comment

ayalz Mar 27, 2024

fhahn Mar 27, 2024

[VPlan] Replace disjoint or with add instead of dropping disjoint. #83821

[VPlan] Replace disjoint or with add instead of dropping disjoint. #83821

Conversation

fhahn commented Mar 4, 2024 • edited

llvmbot commented Mar 4, 2024

nikic left a comment

Choose a reason for hiding this comment

annamthomas commented Mar 4, 2024

annamthomas commented Mar 11, 2024

fhahn commented Mar 18, 2024

github-actions bot commented Mar 18, 2024 • edited

annamthomas commented Mar 18, 2024

danilaml Mar 27, 2024

Choose a reason for hiding this comment

annamthomas Mar 27, 2024

Choose a reason for hiding this comment

annamthomas Mar 27, 2024

Choose a reason for hiding this comment

fhahn Mar 27, 2024

Choose a reason for hiding this comment

ayalz left a comment

Choose a reason for hiding this comment

ayalz Mar 27, 2024

Choose a reason for hiding this comment

fhahn Mar 27, 2024

Choose a reason for hiding this comment

fhahn commented Mar 4, 2024 •

edited

github-actions bot commented Mar 18, 2024 •

edited