[Polly][GEMM] Include parameter constraints in isMatMulOperandAcc

[ttheodor]

Summary

Polly can fail to detect a GEMM pattern due to constraints imposed by
valid integer values (e.g. p0 >=0 and p0 < 2^64-1). These constraints
are not taken into account in isMatMulOperandAcc when "possible" access
relations are constructed.

Fix by intersecting the candidate access relations with the Scop's
context which includes these constraints.

[Meinersbur]

Thanks for the patch. Could run polly-update-format and add a test case?

lib/Transform/ScheduleOptimizer.cpp

1159

unrelated change

[ttheodor]

Add test and fix formatting.

[Meinersbur]

LGTM, but I cannot apply the patch. I assume you uploaded a diff relative to your own version. E.g. there is no newline in at line 1159 in trunk.

It would be nice to have a commit message describing why this is changes, what the intended effect is, and how the change facilitates the effect.

[Meinersbur]

The additional test pattern-matching-based-opts_15.ll still succeeds when I remove intersect_params(Context)?

[ttheodor]

In D43932#1028149, @Meinersbur wrote:

The additional test pattern-matching-based-opts_15.ll still succeeds when I remove intersect_params(Context)?

I generated pattern-matching-based-opts_15.ll with create_ll.sh. If I run clang++ gemm.cpp -mllvm -polly -O3 -c -mllvm -polly-pattern-matching-based-opts -mllvm -stats 2>&1 | grep matrix (with gemm.cpp containing the C++ code embedded in the test) then it only prints something if the intersect_params(Context) is present.

[Meinersbur]

What is your gemm.cpp? The source in pattern-matching-based-opts_15.ll does not correspond to the source in your email: "GEMM pattern matching fails with parametric array sizes". The latter uses unit64_t while the former is from Polybench.

[ttheodor]

gemm.cpp is:

/* C := alpha*A*B + beta*C */
void kernel_gemm(int _PB_NI, int _PB_NJ, int _PB_NK, double alpha, double beta,
                 double *A, double *B, double *C) {
    for (int i = 0; i < _PB_NI; i++)
        for (int j = 0; j < _PB_NJ; j++) {
            C[i * _PB_NJ + j] *= beta;
            for (int k = 0; k < _PB_NK; ++k)
                C[i * _PB_NJ + j] +=
                    alpha * A[i * _PB_NK + k] * B[k * _PB_NJ + j];
        }
}

The one I posted in the e-mail is actually wrong (I didn't linearize properly the accesses) but incidentally exhibits the same behavior when it comes to pattern matching.

[Meinersbur]

Could you try to make pattern-matching-based-opts_15.ll fail without this intersect_params? Instead of create_ll.sh, try using -mllvm -polly-dump-before-file=file.ll to get a more accurate .ll file as Polly sees it when invoked withing clang.

lib/Transform/ScheduleOptimizer.cpp

 ///
 /// @param AccMap    The access relation to be checked.
 /// @param FirstPos  The index of the input dimension that is mapped to
 ///                  the first output dimension.
 /// @param SecondPos The index of the input dimension that is mapped to the
 ///                  second output dimension.
 /// @return          True in case @p AccMap has the expected form and false,
 ///                  otherwise.
-static bool isMatMulOperandAcc(isl::set Domain, isl::map AccMap, int &FirstPos,
-                               int &SecondPos) {
+static bool isMatMulOperandAcc(isl::set Domain, isl::set Context,
+                               isl::map AccMap, int &FirstPos, int &SecondPos) {
   isl::space Space = AccMap.get_space();
   isl::map Universe = isl::map::universe(Space);
 
   if (Space.dim(isl::dim::out) != 2)
     return false;
 
   // MatMul has the form:
   // for (i = 0; i < N; i++)
   //   for (j = 0; j < M; j++)
   //     for (k = 0; k < P; k++)
   //       C[i, j] += A[i, k] * B[k, j]
   //
   // Permutation of three outer loops: 3! = 6 possibilities.
   int FirstDims[] = {0, 0, 1, 1, 2, 2};
   int SecondDims[] = {1, 2, 2, 0, 0, 1};
   for (int i = 0; i < 6; i += 1) {
     auto PossibleMatMul =
         Universe.equate(isl::dim::in, FirstDims[i], isl::dim::out, 0)
             .equate(isl::dim::in, SecondDims[i], isl::dim::out, 1);
 
-    AccMap = AccMap.intersect_domain(Domain);
-    PossibleMatMul = PossibleMatMul.intersect_domain(Domain);
+    AccMap = AccMap.intersect_domain(Domain).intersect_params(Context);
+    PossibleMatMul = PossibleMatMul.intersect_domain(Domain).intersect_params(Context);
+
 
     // If AccMap spans entire domain (Non-partial write),
     // compute FirstPos and SecondPos.
     // If AccMap != PossibleMatMul here (the two maps have been gisted at
     // this point), it means that the writes are not complete, or in other
     // words, it is a Partial write and Partial writes must be rejected.
     if (AccMap.is_equal(PossibleMatMul)) {
       if (FirstPos != -1 && FirstPos != FirstDims[i])

 ///                  to a non-scalar operand of the matrix multiplication and
 ///                  false, otherwise.
 static bool isMatMulNonScalarReadAccess(MemoryAccess *MemAccess,
                                         MatMulInfoTy &MMI) {
   if (!MemAccess->isLatestArrayKind() || !MemAccess->isRead())
     return false;
   auto AccMap = MemAccess->getLatestAccessRelation();
   isl::set StmtDomain = MemAccess->getStatement()->getDomain();
-  if (isMatMulOperandAcc(StmtDomain, AccMap, MMI.i, MMI.j) && !MMI.ReadFromC) {
+  isl::set Context = MemAccess->getStatement()->getParent()->getContext();
+  if (isMatMulOperandAcc(StmtDomain, Context, AccMap, MMI.i, MMI.j) &&
+      !MMI.ReadFromC) {
     MMI.ReadFromC = MemAccess;
     return true;
   }
-  if (isMatMulOperandAcc(StmtDomain, AccMap, MMI.i, MMI.k) && !MMI.A) {
+  if (isMatMulOperandAcc(StmtDomain, Context, AccMap, MMI.i, MMI.k) && !MMI.A) {
     MMI.A = MemAccess;
     return true;
   }
-  if (isMatMulOperandAcc(StmtDomain, AccMap, MMI.k, MMI.j) && !MMI.B) {
+  if (isMatMulOperandAcc(StmtDomain, Context, AccMap, MMI.k, MMI.j) && !MMI.B) {
     MMI.B = MemAccess;
     return true;
   }
   return false;
 }
 
 /// Check accesses to operands of the matrix multiplication.
 ///
static bool containsMatrMult(isl::map PartialSchedule, const Dependences *D,
   auto Accesses = getAccessesInOrder(*Stmt);
   for (auto *MemA = Accesses.end() - 1; MemA != Accesses.begin(); MemA--) {
     auto *MemAccessPtr = *MemA;
     if (!MemAccessPtr->isLatestArrayKind())
       continue;
     if (!MemAccessPtr->isWrite())
       return false;
     auto AccMap = MemAccessPtr->getLatestAccessRelation();
-    if (!isMatMulOperandAcc(Stmt->getDomain(), AccMap, MMI.i, MMI.j))
+    if (!isMatMulOperandAcc(Stmt->getDomain(), Stmt->getParent()->getContext(),
+                            AccMap, MMI.i, MMI.j))
       return false;
     MMI.WriteToC = MemAccessPtr;
     break;
   }
 
   if (!containsOnlyMatMulDep(PartialSchedule, D, MMI.k))
     return false;
 
NewStmt = Stmt->getParent()->addScopStmt(
       OldAcc, MMI.A->getLatestAccessRelation(), Domain);
 
   // Restrict the domains of the copy statements to only execute when also its
   // originating statement is executed.
   ExtMap = ExtMap.set_tuple_id(isl::dim::out, DomainId);
   ExtMap = ExtMap.intersect_range(Domain);
   ExtMap = ExtMap.set_tuple_id(isl::dim::out, NewStmt->getDomainId());
   Node = createExtensionNode(Node, ExtMap);
+
   return Node.child(0).child(0).child(0).child(0).child(0);
 }
 
 /// Get a relation mapping induction variables produced by schedule
 /// transformations to the original ones.
 ///
 /// @param Node The schedule node produced as the result of creation
 ///        of the BLIS kernels.