Repository: llvm/llvm-project
File: llvm/lib/Transforms/Vectorize/VPlanConstruction.cpp (C++)

Code before:
    // Find final reduction computation and replace it with an
    // extract.last.active intrinsic.
    auto *RdxResult =
        vputils::findUserOf<VPInstruction::ComputeReductionResult>(
            BackedgeSelect);
    if (!RdxResult)
      return false;

    // Add mask phi.
    VPBuilder Builder = VPBuilder::getToInsertAfter(PhiR);
    auto *MaskPHI = new VPWidenPHIRecipe(nullptr, /*Start=*/Plan.getFalse());
    Builder.insert(MaskPHI);

    // Add select for mask.
    Builder.setInsertPoint(SelectR);

    if (Op1 == PhiR) {
      // Normalize to selecting the data operand when the condition is true by
      // swapping operands and negating the condition.
      std::swap(Op1, Op2);
      Cond = Builder.createNot(Cond);
    }
    assert(Op2 == PhiR && "data value must be selected if Cond is true");

    if (HeaderMask)
      Cond = Builder.createAnd(Cond, HeaderMask);

    VPValue *AnyOf = Builder.createNaryOp(VPInstruction::AnyOf, {Cond});
    VPValue *MaskSelect = Builder.createSelect(AnyOf, Cond, MaskPHI);
    MaskPHI->addOperand(MaskSelect);

    // Replace select for data.
    VPValue *DataSelect =
        Builder.createSelect(AnyOf, Op1, Op2, SelectR->getDebugLoc());
    SelectR->replaceAllUsesWith(DataSelect);
    PhiR->setBackedgeValue(DataSelect);
    SelectR->eraseFromParent();

    Builder.setInsertPoint(RdxResult);
    auto *ExtractLastActive =
        Builder.createNaryOp(VPInstruction::ExtractLastActive,
                             {DataSelect, MaskSelect, PhiR->getStartValue()},
                             RdxResult->getDebugLoc());
    RdxResult->replaceAllUsesWith(ExtractLastActive);
    RdxResult->eraseFromParent();
  }

  return true;
}

/// Given a first argmin/argmax pattern with strict predicate consisting of

Review comment (suggestion):
TODO: I probably need to make this a select. I don't think the "and" handles `poison` correctly (though currently this still passes the test suite). 

Code after:
    // Find final reduction computation and replace it with an
    // extract.last.active intrinsic.
    auto *RdxResult =
        vputils::findUserOf<VPInstruction::ComputeReductionResult>(
            BackedgeSelect);
    assert(RdxResult && "Could not find reduction result");

    // Add mask phi.
    VPBuilder Builder = VPBuilder::getToInsertAfter(PhiR);
    auto *MaskPHI = new VPWidenPHIRecipe(nullptr, /*Start=*/Plan.getFalse());
    Builder.insert(MaskPHI);

    // Add select for mask.
    Builder.setInsertPoint(SelectR);

    if (Op1 == PhiR) {
      // Normalize to selecting the data operand when the condition is true by
      // swapping operands and negating the condition.
      std::swap(Op1, Op2);
      Cond = Builder.createNot(Cond);
    }
    assert(Op2 == PhiR && "data value must be selected if Cond is true");

    if (HeaderMask)
      Cond = Builder.createLogicalAnd(HeaderMask, Cond);

    VPValue *AnyOf = Builder.createNaryOp(VPInstruction::AnyOf, {Cond});
    VPValue *MaskSelect = Builder.createSelect(AnyOf, Cond, MaskPHI);
    MaskPHI->addOperand(MaskSelect);

    // Replace select for data.
    VPValue *DataSelect =
        Builder.createSelect(AnyOf, Op1, Op2, SelectR->getDebugLoc());
    SelectR->replaceAllUsesWith(DataSelect);
    PhiR->setBackedgeValue(DataSelect);
    SelectR->eraseFromParent();

    Builder.setInsertPoint(RdxResult);
    auto *ExtractLastActive =
        Builder.createNaryOp(VPInstruction::ExtractLastActive,
                             {DataSelect, MaskSelect, PhiR->getStartValue()},
                             RdxResult->getDebugLoc());
    RdxResult->replaceAllUsesWith(ExtractLastActive);
    RdxResult->eraseFromParent();
  }

  return true;
}

/// Given a first argmin/argmax pattern with strict predicate consisting of
/// 1) a MinOrMax reduction \p MinOrMaxPhiR producing \p MinOrMaxResult,