Address comments, add tests for additional dependences.

huntergr-arm · huntergr-arm · commit 6d7bb66f56e2 · 2024-08-29T15:42:49.000Z
diff --git a/llvm/include/llvm/Transforms/Vectorize/LoopVectorizationLegality.h b/llvm/include/llvm/Transforms/Vectorize/LoopVectorizationLegality.h
@@ -414,9 +414,7 @@ class LoopVectorizationLegality {
   }
 
   /// Returns a list of all known histogram operations in the loop.
-  const SmallVectorImpl<HistogramInfo> &getHistograms() const {
-    return Histograms;
-  }
+  bool hasHistograms() const { return !Histograms.empty(); }
 
   PredicatedScalarEvolution *getPredicatedScalarEvolution() const {
     return &PSE;
@@ -467,9 +465,9 @@ class LoopVectorizationLegality {
   bool canVectorizeMemory();
 
   /// If LAA cannot determine whether all dependences are safe, we may be able
-  /// to further analyse some unknown dependences and if they match a certain
-  /// pattern (like a histogram) then we may still be able to vectorize.
-  bool canVectorizeUncheckedDependences();
+  /// to further analyse some IndirectUnsafe dependences and if they match a
+  /// certain pattern (like a histogram) then we may still be able to vectorize.
+  bool canVectorizeIndirectUnsafeDependences();
 
   /// Return true if we can vectorize this loop using the IF-conversion
   /// transformation.
diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp
@@ -1129,29 +1129,38 @@ static bool findHistogram(LoadInst *LI, StoreInst *HSt, Loop *TheLoop,
   return true;
 }
 
-bool LoopVectorizationLegality::canVectorizeUncheckedDependences() {
-  // For now, we only support an unknown dependency that calculates a histogram
+bool LoopVectorizationLegality::canVectorizeIndirectUnsafeDependences() {
+  // For now, we only support an IndirectUnsafe dependency that calculates
+  // a histogram
   if (!EnableHistogramVectorization)
     return false;
 
-  // FIXME: Support more than one unchecked dependence, and check to see if some
-  //        are handled by runtime checks before looking for histograms.
+  // Find a single IndirectUnsafe dependency.
   LAI = &LAIs.getInfo(*TheLoop);
+  const MemoryDepChecker::Dependence *IUDep = nullptr;
   const MemoryDepChecker &DepChecker = LAI->getDepChecker();
   const auto *Deps = DepChecker.getDependences();
-  if (!Deps || Deps->size() != 1)
-    return false;
+  for (const MemoryDepChecker::Dependence &Dep : *Deps) {
+    // Ignore dependencies that are either known to be safe or can be
+    // checked at runtime.
+    if (MemoryDepChecker::Dependence::isSafeForVectorization(Dep.Type) !=
+        MemoryDepChecker::VectorizationSafetyStatus::Unsafe)
+      continue;
 
-  const MemoryDepChecker::Dependence &Dep = (*Deps).front();
+    // We're only interested in IndirectUnsafe dependencies here, where the
+    // address might come from a load from memory. We also only want to handle
+    // one such dependency, at least for now.
+    if (Dep.Type != MemoryDepChecker::Dependence::IndirectUnsafe || IUDep)
+      return false;
 
-  // We're only interested in Unknown or IndirectUnsafe dependences.
-  if (Dep.Type != MemoryDepChecker::Dependence::Unknown &&
-      Dep.Type != MemoryDepChecker::Dependence::IndirectUnsafe)
+    IUDep = &Dep;
+  }
+  if (!IUDep)
     return false;
 
   // For now only normal loads and stores are supported.
-  LoadInst *LI = dyn_cast<LoadInst>(Dep.getSource(DepChecker));
-  StoreInst *SI = dyn_cast<StoreInst>(Dep.getDestination(DepChecker));
+  LoadInst *LI = dyn_cast<LoadInst>(IUDep->getSource(DepChecker));
+  StoreInst *SI = dyn_cast<StoreInst>(IUDep->getDestination(DepChecker));
 
   if (!LI || !SI)
     return false;
@@ -1171,7 +1180,7 @@ bool LoopVectorizationLegality::canVectorizeMemory() {
   }
 
   if (!LAI->canVectorizeMemory())
-    return canVectorizeUncheckedDependences();
+    return canVectorizeIndirectUnsafeDependences();
 
   if (LAI->hasLoadStoreDependenceInvolvingLoopInvariantAddress()) {
     reportVectorizationFailure("We don't allow storing to uniform addresses",
diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -4653,7 +4653,7 @@ bool LoopVectorizationPlanner::isCandidateForEpilogueVectorization(
     return false;
 
   // Loops containing histograms are not currently supported.
-  return Legal->getHistograms().empty();
+  return !Legal->hasHistograms();
 }
 
 bool LoopVectorizationCostModel::isEpilogueVectorizationProfitable(
@@ -10056,7 +10056,7 @@ bool LoopVectorizePass::processLoop(Loop *L) {
   // If there is a histogram in the loop, do not just interleave without
   // vectorizing. The order of operations will be incorrect without the
   // histogram intrinsics, which are only used for recipes with VF > 1.
-  if (!VectorizeLoop && InterleaveLoop && !LVL.getHistograms().empty()) {
+  if (!VectorizeLoop && InterleaveLoop && LVL.hasHistograms()) {
     LLVM_DEBUG(dbgs() << "LV: Not interleaving without vectorization due "
                       << "to histogram operations.\n");
     IntDiagMsg = std::make_pair(
diff --git a/llvm/lib/Transforms/Vectorize/VPlan.h b/llvm/lib/Transforms/Vectorize/VPlan.h
@@ -1619,9 +1619,7 @@ class VPHistogramRecipe : public VPRecipeBase {
   /// Return the mask operand if one was provided, or a null pointer if all
   /// lanes should be executed unconditionally.
   VPValue *getMask() const {
-    if (getNumOperands() == 3)
-      return getOperand(2);
-    return nullptr;
+    return getNumOperands() == 3 ? getOperand(2) : nullptr;
   }
 
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
diff --git a/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp b/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
@@ -967,7 +967,7 @@ void VPHistogramRecipe::execute(VPTransformState &State) {
 
   for (unsigned Part = 0; Part < State.UF; ++Part) {
     Value *Address = State.get(getOperand(0), Part);
-    Value *IncVec = State.get(getOperand(1), Part);
+    Value *IncAmt = State.get(getOperand(1), Part, /*IsScalar=*/true);
     VectorType *VTy = cast<VectorType>(Address->getType());
 
     // The histogram intrinsic requires a mask even if the recipe doesn't;
@@ -980,18 +980,12 @@ void VPHistogramRecipe::execute(VPTransformState &State) {
       Mask = Builder.CreateVectorSplat(
           VTy->getElementCount(), ConstantInt::getTrue(Builder.getInt1Ty()));
 
-    // Not sure how to make IncAmt stay scalar yet. For now just extract the
-    // first element and tidy up later.
-    // FIXME: Do we actually want this to be scalar? We just splat it in the
-    //        backend anyway...
-    Value *IncAmt = Builder.CreateExtractElement(IncVec, Builder.getInt64(0));
-
     // If this is a subtract, we want to invert the increment amount. We may
     // add a separate intrinsic in future, but for now we'll try this.
     if (Opcode == Instruction::Sub)
       IncAmt = Builder.CreateNeg(IncAmt);
     else
-      assert(Opcode == Instruction::Add);
+      assert(Opcode == Instruction::Add && "only add or sub supported for now");
 
     State.Builder.CreateIntrinsic(Intrinsic::experimental_vector_histogram_add,
                                   {VTy, IncAmt->getType()},
diff --git a/llvm/test/Transforms/LoopVectorize/AArch64/sve2-histcnt.ll b/llvm/test/Transforms/LoopVectorize/AArch64/sve2-histcnt.ll
@@ -208,6 +208,17 @@ target triple = "aarch64-unknown-linux-gnu"
 ; CHECK-NEXT: No successors
 ; CHECK-NEXT: }
 
+;; Confirm cost calculation for runtime checks
+; CHECK-LABEL: LV: Checking a loop in 'simple_histogram_rtdepcheck'
+; CHECK: Calculating cost of runtime checks:
+; CHECK: Total cost of runtime checks:
+; CHECK: LV: Minimum required TC for runtime checks to be profitable:
+
+;; Confirm inability to vectorize with potential alias to buckets
+; CHECK-LABEL: LV: Checking a loop in 'simple_histogram_unsafe_alias'
+; CHECK: LV: Can't vectorize due to memory conflicts
+; CHECK-NEXT: LV: Not vectorizing: Cannot prove legality.
+
 define void @simple_histogram(ptr noalias %buckets, ptr readonly %indices, i64 %N) #0 {
 ; CHECK-LABEL: define void @simple_histogram(
 ; CHECK-SAME: ptr noalias [[BUCKETS:%.*]], ptr readonly [[INDICES:%.*]], i64 [[N:%.*]]) #[[ATTR0:[0-9]+]] {
@@ -817,6 +828,140 @@ for.exit:
   ret void
 }
 
+;; Check that we can still vectorize a histogram when LAA found another dependency
+;; that doesn't conflict with the buckets.
+define void @simple_histogram_rtdepcheck(ptr noalias %buckets, ptr %array, ptr %indices, i64 %N) #0 {
+; CHECK-LABEL: define void @simple_histogram_rtdepcheck(
+; CHECK-SAME: ptr noalias [[BUCKETS:%.*]], ptr [[ARRAY:%.*]], ptr [[INDICES:%.*]], i64 [[N:%.*]]) #[[ATTR0]] {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = call i64 @llvm.vscale.i64()
+; CHECK-NEXT:    [[TMP1:%.*]] = shl nuw nsw i64 [[TMP0]], 2
+; CHECK-NEXT:    [[TMP2:%.*]] = call i64 @llvm.umax.i64(i64 [[TMP1]], i64 8)
+; CHECK-NEXT:    [[MIN_ITERS_CHECK:%.*]] = icmp ugt i64 [[TMP2]], [[N]]
+; CHECK-NEXT:    br i1 [[MIN_ITERS_CHECK]], label [[SCALAR_PH:%.*]], label [[VECTOR_MEMCHECK:%.*]]
+; CHECK:       vector.memcheck:
+; CHECK-NEXT:    [[ARRAY1:%.*]] = ptrtoint ptr [[ARRAY]] to i64
+; CHECK-NEXT:    [[INDICES2:%.*]] = ptrtoint ptr [[INDICES]] to i64
+; CHECK-NEXT:    [[TMP3:%.*]] = call i64 @llvm.vscale.i64()
+; CHECK-NEXT:    [[TMP4:%.*]] = shl nuw nsw i64 [[TMP3]], 4
+; CHECK-NEXT:    [[TMP5:%.*]] = sub i64 [[ARRAY1]], [[INDICES2]]
+; CHECK-NEXT:    [[DIFF_CHECK:%.*]] = icmp ult i64 [[TMP5]], [[TMP4]]
+; CHECK-NEXT:    br i1 [[DIFF_CHECK]], label [[SCALAR_PH]], label [[VECTOR_PH:%.*]]
+; CHECK:       vector.ph:
+; CHECK-NEXT:    [[TMP6:%.*]] = call i64 @llvm.vscale.i64()
+; CHECK-NEXT:    [[DOTNEG:%.*]] = mul nsw i64 [[TMP6]], -4
+; CHECK-NEXT:    [[N_VEC:%.*]] = and i64 [[DOTNEG]], [[N]]
+; CHECK-NEXT:    [[TMP7:%.*]] = call i64 @llvm.vscale.i64()
+; CHECK-NEXT:    [[TMP8:%.*]] = shl nuw nsw i64 [[TMP7]], 2
+; CHECK-NEXT:    [[TMP9:%.*]] = call <vscale x 4 x i32> @llvm.experimental.stepvector.nxv4i32()
+; CHECK-NEXT:    [[TMP10:%.*]] = call i32 @llvm.vscale.i32()
+; CHECK-NEXT:    [[TMP11:%.*]] = shl nuw nsw i32 [[TMP10]], 2
+; CHECK-NEXT:    [[DOTSPLATINSERT:%.*]] = insertelement <vscale x 4 x i32> poison, i32 [[TMP11]], i64 0
+; CHECK-NEXT:    [[DOTSPLAT:%.*]] = shufflevector <vscale x 4 x i32> [[DOTSPLATINSERT]], <vscale x 4 x i32> poison, <vscale x 4 x i32> zeroinitializer
+; CHECK-NEXT:    br label [[VECTOR_BODY:%.*]]
+; CHECK:       vector.body:
+; CHECK-NEXT:    [[INDEX:%.*]] = phi i64 [ 0, [[VECTOR_PH]] ], [ [[INDEX_NEXT:%.*]], [[VECTOR_BODY]] ]
+; CHECK-NEXT:    [[VEC_IND:%.*]] = phi <vscale x 4 x i32> [ [[TMP9]], [[VECTOR_PH]] ], [ [[VEC_IND_NEXT:%.*]], [[VECTOR_BODY]] ]
+; CHECK-NEXT:    [[TMP12:%.*]] = getelementptr inbounds i32, ptr [[INDICES]], i64 [[INDEX]]
+; CHECK-NEXT:    [[WIDE_LOAD:%.*]] = load <vscale x 4 x i32>, ptr [[TMP12]], align 4
+; CHECK-NEXT:    [[TMP13:%.*]] = zext <vscale x 4 x i32> [[WIDE_LOAD]] to <vscale x 4 x i64>
+; CHECK-NEXT:    [[TMP14:%.*]] = getelementptr inbounds i32, ptr [[BUCKETS]], <vscale x 4 x i64> [[TMP13]]
+; CHECK-NEXT:    call void @llvm.experimental.vector.histogram.add.nxv4p0.i32(<vscale x 4 x ptr> [[TMP14]], i32 1, <vscale x 4 x i1> shufflevector (<vscale x 4 x i1> insertelement (<vscale x 4 x i1> poison, i1 true, i64 0), <vscale x 4 x i1> poison, <vscale x 4 x i32> zeroinitializer))
+; CHECK-NEXT:    [[TMP15:%.*]] = getelementptr inbounds i32, ptr [[ARRAY]], i64 [[INDEX]]
+; CHECK-NEXT:    store <vscale x 4 x i32> [[VEC_IND]], ptr [[TMP15]], align 4
+; CHECK-NEXT:    [[INDEX_NEXT]] = add nuw i64 [[INDEX]], [[TMP8]]
+; CHECK-NEXT:    [[VEC_IND_NEXT]] = add <vscale x 4 x i32> [[VEC_IND]], [[DOTSPLAT]]
+; CHECK-NEXT:    [[TMP16:%.*]] = icmp eq i64 [[INDEX_NEXT]], [[N_VEC]]
+; CHECK-NEXT:    br i1 [[TMP16]], label [[MIDDLE_BLOCK:%.*]], label [[VECTOR_BODY]], !llvm.loop [[LOOP16:![0-9]+]]
+; CHECK:       middle.block:
+; CHECK-NEXT:    [[CMP_N:%.*]] = icmp eq i64 [[N_VEC]], [[N]]
+; CHECK-NEXT:    br i1 [[CMP_N]], label [[FOR_EXIT:%.*]], label [[SCALAR_PH]]
+; CHECK:       scalar.ph:
+; CHECK-NEXT:    [[BC_RESUME_VAL:%.*]] = phi i64 [ [[N_VEC]], [[MIDDLE_BLOCK]] ], [ 0, [[ENTRY:%.*]] ], [ 0, [[VECTOR_MEMCHECK]] ]
+; CHECK-NEXT:    br label [[FOR_BODY:%.*]]
+; CHECK:       for.body:
+; CHECK-NEXT:    [[IV:%.*]] = phi i64 [ [[BC_RESUME_VAL]], [[SCALAR_PH]] ], [ [[IV_NEXT:%.*]], [[FOR_BODY]] ]
+; CHECK-NEXT:    [[ARRAYIDX:%.*]] = getelementptr inbounds i32, ptr [[INDICES]], i64 [[IV]]
+; CHECK-NEXT:    [[TMP17:%.*]] = load i32, ptr [[ARRAYIDX]], align 4
+; CHECK-NEXT:    [[IDXPROM1:%.*]] = zext i32 [[TMP17]] to i64
+; CHECK-NEXT:    [[ARRAYIDX2:%.*]] = getelementptr inbounds i32, ptr [[BUCKETS]], i64 [[IDXPROM1]]
+; CHECK-NEXT:    [[TMP18:%.*]] = load i32, ptr [[ARRAYIDX2]], align 4
+; CHECK-NEXT:    [[INC:%.*]] = add nsw i32 [[TMP18]], 1
+; CHECK-NEXT:    store i32 [[INC]], ptr [[ARRAYIDX2]], align 4
+; CHECK-NEXT:    [[IDX_ADDR:%.*]] = getelementptr inbounds i32, ptr [[ARRAY]], i64 [[IV]]
+; CHECK-NEXT:    [[IV_TRUNC:%.*]] = trunc i64 [[IV]] to i32
+; CHECK-NEXT:    store i32 [[IV_TRUNC]], ptr [[IDX_ADDR]], align 4
+; CHECK-NEXT:    [[IV_NEXT]] = add nuw nsw i64 [[IV]], 1
+; CHECK-NEXT:    [[EXITCOND:%.*]] = icmp eq i64 [[IV_NEXT]], [[N]]
+; CHECK-NEXT:    br i1 [[EXITCOND]], label [[FOR_EXIT]], label [[FOR_BODY]], !llvm.loop [[LOOP17:![0-9]+]]
+; CHECK:       for.exit:
+; CHECK-NEXT:    ret void
+;
+entry:
+  br label %for.body
+
+for.body:
+  %iv = phi i64 [ 0, %entry ], [ %iv.next, %for.body ]
+  %arrayidx = getelementptr inbounds i32, ptr %indices, i64 %iv
+  %0 = load i32, ptr %arrayidx, align 4
+  %idxprom1 = zext i32 %0 to i64
+  %arrayidx2 = getelementptr inbounds i32, ptr %buckets, i64 %idxprom1
+  %1 = load i32, ptr %arrayidx2, align 4
+  %inc = add nsw i32 %1, 1
+  store i32 %inc, ptr %arrayidx2, align 4
+  %idx.addr = getelementptr inbounds i32, ptr %array, i64 %iv
+  %iv.trunc = trunc i64 %iv to i32
+  store i32 %iv.trunc, ptr %idx.addr, align 4
+  %iv.next = add nuw nsw i64 %iv, 1
+  %exitcond = icmp eq i64 %iv.next, %N
+  br i1 %exitcond, label %for.exit, label %for.body
+
+for.exit:
+  ret void
+}
+
+;; Make sure we don't vectorize if there's a potential alias between buckets
+;; and indices.
+define void @simple_histogram_unsafe_alias(ptr %buckets, ptr %indices, i64 %N) #0 {
+; CHECK-LABEL: define void @simple_histogram_unsafe_alias(
+; CHECK-SAME: ptr [[BUCKETS:%.*]], ptr [[INDICES:%.*]], i64 [[N:%.*]]) #[[ATTR0]] {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label [[FOR_BODY:%.*]]
+; CHECK:       for.body:
+; CHECK-NEXT:    [[IV:%.*]] = phi i64 [ 0, [[ENTRY:%.*]] ], [ [[IV_NEXT:%.*]], [[FOR_BODY]] ]
+; CHECK-NEXT:    [[ARRAYIDX:%.*]] = getelementptr inbounds i32, ptr [[INDICES]], i64 [[IV]]
+; CHECK-NEXT:    [[TMP12:%.*]] = load i32, ptr [[ARRAYIDX]], align 4
+; CHECK-NEXT:    [[IDXPROM1:%.*]] = zext i32 [[TMP12]] to i64
+; CHECK-NEXT:    [[ARRAYIDX2:%.*]] = getelementptr inbounds i32, ptr [[BUCKETS]], i64 [[IDXPROM1]]
+; CHECK-NEXT:    [[TMP13:%.*]] = load i32, ptr [[ARRAYIDX2]], align 4
+; CHECK-NEXT:    [[INC:%.*]] = add nsw i32 [[TMP13]], 1
+; CHECK-NEXT:    store i32 [[INC]], ptr [[ARRAYIDX2]], align 4
+; CHECK-NEXT:    [[IV_NEXT]] = add nuw nsw i64 [[IV]], 1
+; CHECK-NEXT:    [[EXITCOND:%.*]] = icmp eq i64 [[IV_NEXT]], [[N]]
+; CHECK-NEXT:    br i1 [[EXITCOND]], label [[FOR_EXIT:%.*]], label [[FOR_BODY]]
+; CHECK:       for.exit:
+; CHECK-NEXT:    ret void
+;
+entry:
+  br label %for.body
+
+for.body:
+  %iv = phi i64 [ 0, %entry ], [ %iv.next, %for.body ]
+  %arrayidx = getelementptr inbounds i32, ptr %indices, i64 %iv
+  %0 = load i32, ptr %arrayidx, align 4
+  %idxprom1 = zext i32 %0 to i64
+  %arrayidx2 = getelementptr inbounds i32, ptr %buckets, i64 %idxprom1
+  %1 = load i32, ptr %arrayidx2, align 4
+  %inc = add nsw i32 %1, 1
+  store i32 %inc, ptr %arrayidx2, align 4
+  %iv.next = add nuw nsw i64 %iv, 1
+  %exitcond = icmp eq i64 %iv.next, %N
+  br i1 %exitcond, label %for.exit, label %for.body
+
+for.exit:
+  ret void
+}
+
 attributes #0 = { "target-features"="+sve2" vscale_range(1,16) }
 
 !0 = distinct !{!0, !1}
