[LoopInterchange] Prevent Loop Interchange for non-affine value store to affine access

[geng]

Summary

In the following case
  #include <stdio.h>
  int a;
  int c[4][4];

  void test_deps() {
      for (int i = 0; i <= 3; i++) {
          for (int j = 0; j <= 3; j++) {
              a ^= 0x1000;
              c[j][i] = a;
          }
      }
  }

  main() {
      test_deps();
      for (int k = 0; k < 4; k++)
      printf("%d\n", c[0][k]);
  }

If we do the loop interchange, there will be correctness issues. However, the loop interchange can not be prevent by dependence matrix. The reason is that we have only one multi-dimensional array access c[j][i] within the loop, other accesses are scalar access. This patch is supposed to prevent loop interchange from the above case. We think the root cause of correctness issue is that the non-affine value (scalar) is stored to affine address (multi-dimensional array access).

This patch closes https://bugs.llvm.org/show_bug.cgi?id=47915

[fhahn]

Thanks for the patch. Some initial comments inline.

llvm/lib/Transforms/Scalar/LoopInterchange.cpp

951

nit: re-flow comments

960

Nit: start comment with a capital letter, end with period (.).

961

Usually early returns are used to limit the indentation level. Could this just be something like

if (!isa<Instruction>(*Op) || OutLoop->isLoopInvariant(Op))
  return false;

and drop the else { block?

983

just return the expression? It would also be good to add a brief comment explaining the condition.

991

could this just be something like return any_of(I->operands(), [](Use &U) {});?

llvm/test/Transforms/LoopInterchange/non-affine-store-to-affine.ll

2

you need to add ; REQUIRES: asserts when using -debug, so the test is not run if assertions. are disabled.

42

A few small suggestions for the test case.

it would be helpful to update the names of the BBs to include more information, e.g. outer.header, inner.body or something like that. And re-order them in the file so they are in a more natural order like preheader, outer header, inner, outer.latch, exit.

Might also be good to have an exit block for the loop nest that returns and maybe clean up the names of the values (e.g. use %j for %indvars.iv.i %j.next for the step).

[geng]

Thanks for the comments! Changed the code according to the initial comments

[fhahn]

Thanks for the update!

llvm/lib/Transforms/Scalar/LoopInterchange.cpp

969

I'm not sure I entirely follow the logic here. Can we just check if the PHI is affine by getting the SCEV and checking if it is an affine addrec? also, we might check sibling loops that are not related by using subloops? (granted, sibling sub-loops are not supported at the moment.) If we would start at InnerLoop and traverse the getParent chain until we reach OuterLoop, we should only visit the relevant loops?

Would it make sense to flip the logic to instead of looking backwards to just iterate over a worklist starting at affine PHIs in inner and outer and then propagate the affine property to all uses and add newly discovered affine nodes to the worklist? Might need a bit of extra logic to deal with nodes where the operands are still pending.

Then we just need to check the set of visited node to see if we are dealing with affine operands for the stores?

[jdoerfert]

I don't understand the reasoning here, even if I assume for a moment the underlying idea is sound, the patch doesn't seem to be.
For example, it will only check stores of a certain shape, why can we store anything to non-affine accesses? What about non-store accesses?
This looks too much like ad-hoc pattern matching.

The problem, as I see it, is not about "storing" but the fact that the loop carried dependence through a (caused by a ^= 0x1000;) is not taken into account properly. While by itself it looks like a reduction dependency (see for example [0]), it cannot be ignored as intermediate values escape. I haven't read all of loop interchange to determine where/why the dependency is broken, but that is the reason this breaks.

A sketch for an example without a "store of a non-affine value" would look something like this:

void test_deps() {
    for (int i = 0; i <= 3; i++) {
        for (int j = 0; j <= 3; j++) {
            a ^= 0x1000;
            if (a == ...)
              c[j][i] = 42;
        }
    }
}

[0] https://arxiv.org/abs/1505.07716

llvm/lib/Transforms/Scalar/LoopInterchange.cpp

1020

This breaks out of this loop with both true and false return value. That cannot be correct.

llvm/lib/Transforms/Scalar/LoopInterchange.cpp

for (PHINode &PHI : LoopNestExit->phis()) {
      //        well.
      if (OuterLoop->getLoopLatch()->getUniquePredecessor() == nullptr)
        return false;
     }
   }
   return true;
 }
 
+static bool containsInductionOfLoop(Loop *L, const SCEV *Expr) {
+  const SCEVAddRecExpr *AddRec = dyn_cast<SCEVAddRecExpr>(Expr);
+  if (AddRec == nullptr || !AddRec->isAffine())
+    return false;
+  if (AddRec->getLoop() == L)
+    return true;
+  else
+    return containsInductionOfLoop(L, AddRec->getStart());
+}
+
+// If a value/variable is loaded from non-affine access and did finite number
+// of calculation, we call it non-affine value/variable.  If a value/variable
+// is loaded from affine access and did finite number of calculation, we call
+// it affine value/variable.
+static bool checkNonAffineValue(Value *Op, ScalarEvolution *SE, Loop *OuterLoop,
+                                Loop *InnerLoop) {
+  if (!isa<Instruction>(*Op) || OuterLoop->isLoopInvariant(Op))
+    return false;
+  // Check if Op is the induction variable of any loop contained in OuterLoop.
+  // If the Op is the induction variable of the OuterLoop or InnerLoop. It is
+  // safe to store to the affine value. Otherwise, the assignment of the
+  // induction variable to the affine variable is similar to assign a
+  // non-affine value to the affine variable.
+  if (PHINode *PHI = dyn_cast<PHINode>(Op)) {
+    return !(containsInductionOfLoop(InnerLoop, SE->getSCEV(PHI)) ||
+             containsInductionOfLoop(OuterLoop, SE->getSCEV(PHI)));
+  }
+
+  if (LoadInst *Ld = dyn_cast<LoadInst>(Op)) {
+    Value *LdAddr = Ld->getPointerOperand();
+    // If the value is loaded from a non-affine access, we need to prevent
+    // interchanging when it is assigned to an affine variable.
+    return !(containsInductionOfLoop(InnerLoop, SE->getSCEV(LdAddr)) ||
+             containsInductionOfLoop(OuterLoop, SE->getSCEV(LdAddr)));
+  } else {
+    return any_of(cast<Instruction>(Op)->operands(), [&](const Use &U) {
+      return checkNonAffineValue(U.get(), SE, OuterLoop, InnerLoop);
+    });
+  }
+}
+
+static bool nonAffineValueAssignedToAffine(Loop *InnerLoop, Loop *OuterLoop,
+                                           ScalarEvolution *SE) {
+
+  using ValueVector = SmallVector<Value *, 16>;
+  ValueVector AffineStInstr;
+
+  for (BasicBlock *BB : OuterLoop->blocks()) {
+    for (Instruction &I : *BB) {
+      if (auto *St = dyn_cast<StoreInst>(&I)) {
+        Value *StAddr = St->getPointerOperand();
+        const SCEVAddRecExpr *StAddrAddRec =
+            dyn_cast<SCEVAddRecExpr>(SE->getSCEV(StAddr));
+        if (StAddrAddRec != nullptr && StAddrAddRec->isAffine())
+          AffineStInstr.push_back(&I);
+      }
+    }
+  }
+
+  // For each affine store instruction, we need to check if the store value is
+  // non-affine.
+  for (ValueVector::iterator I = AffineStInstr.begin(),
+                             IE = AffineStInstr.end();
+       I != IE; ++I) {
+    Value *Op = cast<StoreInst>(*I)->getOperand(0);
+    if (checkNonAffineValue(Op, SE, OuterLoop, InnerLoop))
+      return true;
+  }
+  return false;
+}
+
 bool LoopInterchangeLegality::canInterchangeLoops(unsigned InnerLoopId,
                                                   unsigned OuterLoopId,
                                                   CharMatrix &DepMatrix) {
   if (!isLegalToInterChangeLoops(DepMatrix, InnerLoopId, OuterLoopId)) {
     LLVM_DEBUG(dbgs() << "Failed interchange InnerLoopId = " << InnerLoopId
                       << " and OuterLoopId = " << OuterLoopId
                       << " due to dependence\n");
     ORE->emit([&]() {
ORE->emit([&]() {
       return OptimizationRemarkMissed(DEBUG_TYPE, "UnsupportedExitPHI",
                                       OuterLoop->getStartLoc(),
                                       OuterLoop->getHeader())
              << "Found unsupported PHI node in loop exit.";
     });
     return false;
   }
 
+  if (nonAffineValueAssignedToAffine(InnerLoop, OuterLoop, SE)) {
+    LLVM_DEBUG(dbgs() << "Found non-affine value assigned by affine access.\n");
+    ORE->emit([&]() {
+      return OptimizationRemarkMissed(
+                 DEBUG_TYPE, "NonAffineValueAssignedToAffine",
+                 OuterLoop->getStartLoc(), OuterLoop->getHeader())
+             << "Found non-affine value assigned by affine access.";
+    });
+    return false;
+  }
+
   return true;
 }
 
 int LoopInterchangeProfitability::getInstrOrderCost() {
   unsigned GoodOrder, BadOrder;
   BadOrder = GoodOrder = 0;
   for (BasicBlock *BB : InnerLoop->blocks()) {
     for (Instruction &Ins : *BB) {

llvm/test/Transforms/LoopInterchange/non-affine-store-to-affine.ll

+; REQUIRES: asserts
+; RUN: opt < %s -basicaa -loop-interchange -S -debug 2>&1 | FileCheck %s
+
+target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
+target triple = "aarch64-unknown-linux-gnu"
+
+@f = dso_local local_unnamed_addr global i32* undef, align 8
+@c = common dso_local local_unnamed_addr global [4 x [4 x i32]] zeroinitializer, align 4
+
+;; The following case should not be interchanged.
+;; But interchange can not be prevent by dependence matrix.
+;;
+;; int* f;
+;; int c[4][4];
+;;
+;; void foo() {
+;;   for (int i = 0; i <= 3; i++) {
+;;     for (int j = 0; j <= 3; j++) {
+;;       *f ^= 0x1000;
+;;       c[j][i] = *f;
+;;     }
+;;   }
+;; }
+
+; CHECK: Found non-affine value assigned by affine access.
+; CHECK: Not interchanging loops. Cannot prove legality.
+
+define void @foo() {
+entry:
+  %0 = load i32*, i32** @f, align 8
+  br label %outer.header
+
+outer.header:                            ; preds = %outer.latch, %entry
+  %indvars.iv.i = phi i64 [ 0, %entry ], [ %indvars.iv.next.i, %outer.latch ]
+  br label %inner.header
+
+inner.header:                                      ; preds = %for.body4.i, %outer.header
+  %indvars.iv.j = phi i64 [ 0, %outer.header ], [ %indvars.iv.next.j, %inner.header ]
+  %1 = load i32, i32* %0, align 4
+  %xor.i = xor i32 %1, 4096
+  store i32 %xor.i, i32* %0, align 4
+  %arrayidx6.i = getelementptr inbounds [4 x [4 x i32]], [4 x [4 x i32]]* @c, i64 0, i64 %indvars.iv.j, i64 %indvars.iv.i
+  store i32 %xor.i, i32* %arrayidx6.i, align 4
+  %indvars.iv.next.j = add nuw nsw i64 %indvars.iv.j, 1
+  %exitcond.i = icmp eq i64 %indvars.iv.next.j, 4
+  br i1 %exitcond.i, label %outer.latch, label %inner.header
+
+outer.latch:                              ; preds = %inner.header
+  %indvars.iv.next.i = add nuw nsw i64 %indvars.iv.i, 1
+  %exitcond20.i = icmp eq i64 %indvars.iv.next.i, 4
+  br i1 %exitcond20.i, label %outer.exit, label %outer.header
+
+outer.exit:                                         ; preds = %outer.latch
+  br label %exit
+
+exit:                                        ; preds = %outer.exit
+  ret void
+}
+