[Polly][MatMul] Make MatMul detection independent of internal isl representations.

[annanay25]

Summary

The pattern recognition for MatMul is restrictive.

The number of "disjuncts" in the isl_map containing constraint information was previously required to be 1 (as per isl_coalesce - which should ideally produce a domain map with a single disjunct, but does not under some circumstances).

This was changed and made more flexible.

[grosser]

Nice. Could you possibly add a test case?

lib/Transform/ScheduleOptimizer.cpp

650

Can you pull the MemAccess->getStatement()->getDomain() out?

[Meinersbur]

lib/Transform/ScheduleOptimizer.cpp

570–576

[Suggestion] I think we don't need that check that there are exactly 2 constraints per basic_map anymore.

584–607

[Suggestion] I think it is possible to make a loop for this:

int FirstDims[] = {0,0,1,1,2,2};
int SecondDims[] = {1,2,2,0,0,1};
for (int i = 0; i <= 6; i+=1)
    PossibleMatMul[i] = Universe.equate(isl::dim::in, FirstDims[i], isl::dim::out, 0).equate(isl::dim::in, SecondDims[i], isl::dim::out, 1);

609

[Style] remove commented code.

[gareevroman]

The pattern recognition for MatMul is restrictive.

The number of "disjuncts" in the isl_map containing constraint information was previously required to be 1 (as per isl_coalesce - which should ideally produce a domain map with a single disjunct, but does not under some circumstances).

Could you please show an example of such an isl_map?

lib/Transform/ScheduleOptimizer.cpp

604

Could we just drop the constraints not involving the output dimensions:

auto *NewAccMap = isl_map_drop_constraints_not_involving_dims(

isl_map_copy(AccMap), isl_dim_out, 0, 2);

and investigate the new map, using the old code?

if (NewAccMap.foreach_basic_map(Lambda) != isl::stat::ok || DimInPos[0] < 0 ||

[Meinersbur]

lib/Transform/ScheduleOptimizer.cpp

604

This is not about constraints of uninvolved dims.

isl does not have unique representation for representing the same sets. An example are the sets

{ Stmt[0] -> MemRef[0]; Stmt[1] -> MemRef[1]; Stmt[2] -> MemRef[2] }
{ Stmt[i0] -> MemRef[i0] : 0 <= i0 < 3 }

which both represent the same logical set. There is isl_map_coalesce which unfortunately does not guaranttee to get to any cannonical form.

It is more robust to use operations that operate on logical sets instead to investigate the internal structure.

[annanay25]

In D36460#835110, @grosser wrote:

Nice. Could you possibly add a test case?

Yes, will do. Thanks!

lib/Transform/ScheduleOptimizer.cpp

584–607

Aha! Thanks! Thought I could individually comment the values of FirstPos and SecondPos for each domain.

604

Thanks for this discussion. I would also point to my post on the isl-development group for more information (and an example) on this.

Link: https://groups.google.com/forum/#!topic/isl-development/iURXrsSVC28

609

Oops.

650

Will do.

[annanay25]

lib/Transform/ScheduleOptimizer.cpp

584–607

@Meinersbur , should I change it to the loop according to your suggestion? Or should I leave the comments behind for readability?

[gareevroman]

lib/Transform/ScheduleOptimizer.cpp

604

This is not about constraints of uninvolved dims.

isl does not have unique representation for representing the same sets. An example are the sets

{ Stmt[0] -> MemRef[0]; Stmt[1] -> MemRef[1]; Stmt[2] -> MemRef[2] }
{ Stmt[i0] -> MemRef[i0] : 0 <= i0 < 3 }

which both represent the same logical set. There is isl_map_coalesce which unfortunately does not guaranttee to get to any cannonical form.

Right. Dropping the constraints would work only if Stmt[0] -> MemRef[0]; is represented as Stmt[i0] -> MemRef[o0] : o0 = i0 and i0 = 1

P.S.: Btw, I'm not sure that it's safe to use optimizeDataLayoutMatrMulPattern in case A and B have constraints, since we build the corresponding AccessRelation from scratch and don't take old constraints into the account. Furthermore, probably, it'd require to modify the isolation algorithm.

[Meinersbur]

lib/Transform/ScheduleOptimizer.cpp

584–607

This is up to you, I am ok with both versions. The more important is to have a test case such that we can commit this.

For the long term, I think we should avoid code duplication and make single loop for this and the following code:

// Permutations of the three matmul 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);
    if (computeMatchedDimensions(AccMap, PossibleMatMul, FirstDims[i], SecondDims[i], FirstPos, SecondPos)
      return true;
}
return false;

Some minor speed-up might be possible:

int FirstDims[] = {0,0,1,1,2,2};
int SecondDims[] = {1,2,2,0,0,1};
for (int i = 0; i <= 6; i+=1) {
    if (FirstPos >= 0 && FirstPos != FirstDims[i])
      continue;
    if (SecondPos >= 0 && SecondPos != SecondDims[i])
      continue;

    auto PossibleMatMul = Universe.equate(isl::dim::in, FirstDims[i], isl::dim::out, 0).equate(isl::dim::in, SecondDims[i], isl::dim::out, 1);
    if (computeMatchedDimensions(AccMap, PossibleMatMul, FirstDims[i], SecondDims[i], FirstPos, SecondPos)
      return true;
}
return false;

[annanay25]

lib/Transform/ScheduleOptimizer.cpp

570–576

@Meinersbur ;
IMO it is necessary to ensure that there are exactly 2 out-dims. In the sense that, the access is of the form -

Stmt[i0, i1, i2] -> MemRef[o0, o1]

and should have two constraints, one each for o0 and o1. Ex:
o0 = i0 and o1 = i1

604

I am not familiar with the usage of optimizeDataLayoutMatrMulPattern function. If you think it is a cleaner way of doing this I will look into it, thanks!

[Meinersbur]

lib/Transform/ScheduleOptimizer.cpp

570–576

isl_map_is_equal will test whether "o0 = i0" and "o1 = i1" are true.

Having exactly two constraints is neither a sufficient nor necessary condition. E.g.
{ Stmt[i0, i1, i2] -> MemRef[o0, o1] : o0 = i0 and i1 = 5 and o1= 5 }
fulfills the condition, and
[p0] -> { Stmt[i0, i1, i2] -> MemRef[o0, o1] : o0 = i0 and i1 = p0 }
does not.

[annanay25]

lib/Transform/ScheduleOptimizer.cpp

570–576

Yes, okay I agree.

When I remove this check and run check-polly , I get a failing test case (polly/test/ScheduleOptimizer/2012-04-16-Trivially-vectorizable-loops.ll) with AccMap as follows -

{ Stmt_for_body8[i0, i1, i2] -> MemRef_C[o0] : o0 = 1536i0 + i1 }

This is because of the Universe.equate(isl::dim::in, FirstDims[i], isl::dim::out, 0).equate(isl::dim::in, SecondDims[i], isl::dim::out, 1).

[annanay25]

@Meinersbur ;

So the reason I haven't been able to come up with a test case is

(lldb) po isl_map_dump(AccMap.ptr)                                                                                                                                                
{ Stmt_loop_body_reduction[i0, i1, i2] -> MemRef_retval[o0, o1] : (i2 = 783 and o0 = i0 and o1 = i1 and i1 <= 1 - i0 and i0 >= 0 and i0 <= 1 and i1 >= 0 and i1 <= 1) or (i0 = 1 and i1 = 1 and i2 = 783 and o0 = 1 and o1 = 1) }
(lldb) po isl_map_dump(PossibleMatMul.ptr)                                                                                                                                           
{ Stmt_loop_body_reduction[i0, i1, i2] -> MemRef_retval[o0, o1] : o0 = i0 and o1 = i1 and i0 >= 0 and i0 <= 1 and i1 >= 0 and i1 <= 1 and i2 >= 0 and i2 <= 783 }

These two maps are not detected as equal - AccMap.is_equal(PossibleMatMul) is false.

Tried removing extra dims / retaining constraints having only two input dims / etc.

Reason: o0 = 1 and o1 = 1 and i0 = 1 and i1 = 1 != o0 = i0 and o1 = i1.

[annanay25]

[todo] Test case.

[Meinersbur]

Can you add a test case using -polly-import=jscop?

lib/Transform/ScheduleOptimizer.cpp

582

nice!

601

[Compiler error] RedAccMap is not defined anywhere. Did you mean AccMap?

[annanay25]

Added test case.

[annanay25]

Addressed @Meinersbur's comments.

Removed test case dependency on DeLICM.

[gareevroman]

lib/Transform/ScheduleOptimizer.cpp

604

Sorry, for some reason, I haven't received this comment. However, optimizeDataLayoutMatrMulPattern is used by Polly. I'll try to check and, probably, update optimizeDataLayoutMatrMulPattern, when I have time.

P.S.: Btw, I'm not sure that it's safe to use optimizeDataLayoutMatrMulPattern in case A and B have constraints, since we build the corresponding AccessRelation from scratch and don't take old constraints into the account. Furthermore, probably, it'd require to modify the isolation algorithm.

@Meinersbur are you sure that it's always safe to use this detection?

[Meinersbur]

lib/Transform/ScheduleOptimizer.cpp

604

Can you elaborate on your concerns? The check should do the same as before, except not requiring inspection of the isl_map's internal representation, thus also accepting internal representations that are equivalent.

If you mean constraints on parameters or domains, then we'd have the same problems as before, because such constraints are gist'ed away for the access relation.

lib/Transform/ScheduleOptimizer.cpp

for (int i = Dims - 1; i >= 0; i--)
     if (Node.band_member_get_coincident(i)) {
       Node = prevectSchedBand(Node, i, PrevectorWidth);
       break;
     }
 
   return Node;
 }
 
-/// Get the position of a dimension with a non-zero coefficient.
-///
-/// Check that isl constraint @p Constraint has only one non-zero
-/// coefficient for dimensions that have type @p DimType. If this is true,
-/// return the position of the dimension corresponding to the non-zero
-/// coefficient and negative value, otherwise.
-///
-/// @param Constraint The isl constraint to be checked.
-/// @param DimType    The type of the dimensions.
-/// @return           The position of the dimension in case the isl
-///                   constraint satisfies the requirements, a negative
-///                   value, otherwise.
-static int getMatMulConstraintDim(isl::constraint Constraint,
-                                  isl::dim DimType) {
-  int DimPos = -1;
-  auto LocalSpace = Constraint.get_local_space();
-  int LocalSpaceDimNum = LocalSpace.dim(DimType);
-  for (int i = 0; i < LocalSpaceDimNum; i++) {
-    auto Val = Constraint.get_coefficient_val(DimType, i);
-    if (Val.is_zero())
-      continue;
-    if (DimPos >= 0 || (DimType == isl::dim::out && !Val.is_one()) ||
-        (DimType == isl::dim::in && !Val.is_negone()))
-      return -1;
-    DimPos = i;
-  }
-  return DimPos;
-}
-
-/// Check the form of the isl constraint.
-///
-/// Check that the @p DimInPos input dimension of the isl constraint
-/// @p Constraint has a coefficient that is equal to negative one, the @p
-/// DimOutPos has a coefficient that is equal to one and others
-/// have coefficients equal to zero.
-///
-/// @param Constraint The isl constraint to be checked.
-/// @param DimInPos   The input dimension of the isl constraint.
-/// @param DimOutPos  The output dimension of the isl constraint.
-/// @return           isl_stat_ok in case the isl constraint satisfies
-///                   the requirements, isl_stat_error otherwise.
-static isl_stat isMatMulOperandConstraint(isl::constraint Constraint,
-                                          int &DimInPos, int &DimOutPos) {
-  auto Val = Constraint.get_constant_val();
-  if (!isl_constraint_is_equality(Constraint.get()) || !Val.is_zero())
-    return isl_stat_error;
-  DimInPos = getMatMulConstraintDim(Constraint, isl::dim::in);
-  if (DimInPos < 0)
-    return isl_stat_error;
-  DimOutPos = getMatMulConstraintDim(Constraint, isl::dim::out);
-  if (DimOutPos < 0)
-    return isl_stat_error;
-  return isl_stat_ok;
-}
-
 /// Permute the two dimensions of the isl map.
 ///
 /// Permute @p DstPos and @p SrcPos dimensions of the isl map @p Map that
 /// have type @p DimType.
 ///
 /// @param Map     The isl map to be modified.
 /// @param DimType The type of the dimensions.
 /// @param DstPos  The first dimension.
isl::map permuteDimensions(isl::map Map, isl::dim DimType, unsigned DstPos,
   Map = Map.move_dims(DimType, MaxDim, FreeDim, 0, 1);
   if (DimId)
     Map = Map.set_tuple_id(DimType, DimId);
   if (FreeDimId)
     Map = Map.set_tuple_id(FreeDim, FreeDimId);
   return Map;
 }
 
+/// Compute values of FirstPos and SecondPos.
+///
+/// @param AccMap            The access relation to be checked.
+/// @param PossibleMatMul    One of the many(6) possible access maps
+///                          corresponding to MatMul pattern.
+/// @param ExpectedFirstPos  The expected value of FirstPos (if already
+///                          assigned).
+/// @param ExpectedSecondPos The expected value of SecondPos (if already
+///                          assigned).
+/// @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 FirstPos, @p SecondPos have
+///                          the expected values, and false otherwise.
+static bool computeMatchedDimensions(isl::map AccMap, isl::map PossibleMatMul,
+                                     int ExpectedFirstPos,
+                                     int ExpectedSecondPos, int &FirstPos,
+                                     int &SecondPos) {
+  if (AccMap.is_equal(PossibleMatMul)) {
+    if (FirstPos != -1 && FirstPos != ExpectedFirstPos)
+      return false;
+    FirstPos = ExpectedFirstPos;
+    if (SecondPos != -1 && SecondPos != ExpectedSecondPos)
+      return false;
+    SecondPos = ExpectedSecondPos;
+    return true;
+  }
+
+  // If AccMap is not the same as PossibleMatMul here (the two maps have
+  // been gisted at this point), it means that their domains are not equal.
+  // This means that the writes are not complete, or in other words,
+  // it is a Partial write and Partial writes must be rejected.
+  return false;
+}
+
 /// Check the form of the access relation.
 ///
 /// Check that the access relation @p AccMap has the form M[i][j], where i
 /// is a @p FirstPos and j is a @p SecondPos.
 ///
 /// @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::map AccMap, int &FirstPos, int &SecondPos) {
-  int DimInPos[] = {FirstPos, SecondPos};
-  auto Lambda = [=, &DimInPos](isl::basic_map BasicMap) -> isl::stat {
+static bool isMatMulOperandAcc(isl::set Domain, isl::map AccMap, int &FirstPos,
+                               int &SecondPos) {
+
+  auto Lambda = [](isl::basic_map BasicMap) -> isl::stat {
     auto Constraints = BasicMap.get_constraint_list();
     if (isl_constraint_list_n_constraint(Constraints.get()) != 2)
       return isl::stat::error;
-    for (int i = 0; i < 2; i++) {
-      auto Constraint =
-          isl::manage(isl_constraint_list_get_constraint(Constraints.get(), i));
-      int InPos, OutPos;
-      if (isMatMulOperandConstraint(Constraint, InPos, OutPos) ==
-              isl_stat_error ||
-          OutPos > 1 || (DimInPos[OutPos] >= 0 && DimInPos[OutPos] != InPos))
-        return isl::stat::error;
-      DimInPos[OutPos] = InPos;
-    }
     return isl::stat::ok;
   };
-  if (AccMap.foreach_basic_map(Lambda) != isl::stat::ok || DimInPos[0] < 0 ||
-      DimInPos[1] < 0)
+  if (AccMap.foreach_basic_map(Lambda) != isl::stat::ok)
     return false;
-  FirstPos = DimInPos[0];
-  SecondPos = DimInPos[1];
-  return true;
+
+  isl::map PossibleMatMul[6];
+
+  isl::space Space = AccMap.get_space();
+  isl::map Universe = isl::map::universe(Space);
+
+  // { Stmt[i0, i1, i2] -> MemRef_A[i0, i1] }
+  // FirstPos = 0, SecondPos = 1
+  PossibleMatMul[0] = Universe.equate(isl::dim::in, 0, isl::dim::out, 0)
+                          .equate(isl::dim::in, 1, isl::dim::out, 1);
+  // { Stmt[i0, i1, i2] -> MemRef_A[i0, i2] }
+  // FirstPos = 0, SecondPos = 2
+  PossibleMatMul[1] = Universe.equate(isl::dim::in, 0, isl::dim::out, 0)
+                          .equate(isl::dim::in, 2, isl::dim::out, 1);
+  // { Stmt[i0, i1, i2] -> MemRef_A[i1, i2] }
+  // FirstPos = 1, SecondPos = 2
+  PossibleMatMul[2] = Universe.equate(isl::dim::in, 1, isl::dim::out, 0)
+                          .equate(isl::dim::in, 2, isl::dim::out, 1);
+  // { Stmt[i0, i1, i2] -> MemRef_A[i1, i0] }
+  // FirstPos = 1, SecondPos = 0
+  PossibleMatMul[3] = Universe.equate(isl::dim::in, 1, isl::dim::out, 0)
+                          .equate(isl::dim::in, 0, isl::dim::out, 1);
+  // { Stmt[i0, i1, i2] -> MemRef_A[i2, i0] }
+  // FirstPos = 2, SecondPos = 0
+  PossibleMatMul[4] = Universe.equate(isl::dim::in, 2, isl::dim::out, 0)
+                          .equate(isl::dim::in, 0, isl::dim::out, 1);
+  // { Stmt[i0, i1, i2] -> MemRef_A[i2, i1] }
+  // FirstPos = 2, SecondPos = 1
+  PossibleMatMul[5] = Universe.equate(isl::dim::in, 2, isl::dim::out, 0)
+                          .equate(isl::dim::in, 1, isl::dim::out, 1);
+
+  // AccMap = isl_map_intersect_domain(AccMap, isl_set_copy(Domain));
+  AccMap = AccMap.intersect_domain(Domain);
+  PossibleMatMul[0] = PossibleMatMul[0].intersect_domain(Domain);
+  PossibleMatMul[1] = PossibleMatMul[1].intersect_domain(Domain);
+  PossibleMatMul[2] = PossibleMatMul[2].intersect_domain(Domain);
+  PossibleMatMul[3] = PossibleMatMul[3].intersect_domain(Domain);
+  PossibleMatMul[4] = PossibleMatMul[4].intersect_domain(Domain);
+  PossibleMatMul[5] = PossibleMatMul[5].intersect_domain(Domain);
+
+  // Relying on short-circuiting of expression.
+  return (computeMatchedDimensions(AccMap, PossibleMatMul[0], 0, 1, FirstPos,
+                                   SecondPos) ||
+          computeMatchedDimensions(AccMap, PossibleMatMul[1], 0, 2, FirstPos,
+                                   SecondPos) ||
+          computeMatchedDimensions(AccMap, PossibleMatMul[2], 1, 2, FirstPos,
+                                   SecondPos) ||
+          computeMatchedDimensions(AccMap, PossibleMatMul[3], 1, 0, FirstPos,
+                                   SecondPos) ||
+          computeMatchedDimensions(AccMap, PossibleMatMul[4], 2, 0, FirstPos,
+                                   SecondPos) ||
+          computeMatchedDimensions(AccMap, PossibleMatMul[5], 2, 1, FirstPos,
+                                   SecondPos));
 }
 
 /// Does the memory access represent a non-scalar operand of the matrix
 /// multiplication.
 ///
 /// Check that the memory access @p MemAccess is the read access to a non-scalar
 /// operand of the matrix multiplication or its result.
 ///
 /// @param MemAccess The memory access to be checked.
 /// @param MMI       Parameters of the matrix multiplication operands.
 /// @return          True in case the memory access represents the read access
 ///                  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();
-  if (isMatMulOperandAcc(AccMap, MMI.i, MMI.j) && !MMI.ReadFromC &&
-      isl_map_n_basic_map(AccMap.get()) == 1) {
+  if (isMatMulOperandAcc(MemAccess->getStatement()->getDomain(), AccMap, MMI.i,
+                         MMI.j) &&
+      !MMI.ReadFromC) {
     MMI.ReadFromC = MemAccess;
     return true;
   }
-  if (isMatMulOperandAcc(AccMap, MMI.i, MMI.k) && !MMI.A &&
-      isl_map_n_basic_map(AccMap.get()) == 1) {
+  if (isMatMulOperandAcc(MemAccess->getStatement()->getDomain(), AccMap, MMI.i,
+                         MMI.k) &&
+      !MMI.A) {
     MMI.A = MemAccess;
     return true;
   }
-  if (isMatMulOperandAcc(AccMap, MMI.k, MMI.j) && !MMI.B &&
-      isl_map_n_basic_map(AccMap.get()) == 1) {
+  if (isMatMulOperandAcc(MemAccess->getStatement()->getDomain(), AccMap, MMI.k,
+                         MMI.j) &&
+      !MMI.B) {
     MMI.B = MemAccess;
     return true;
   }
   return false;
 }
 
 /// Check accesses to operands of the matrix multiplication.
 ///
if (Stmt->size() <= 1)
     return false;
   for (auto *MemA = Stmt->end() - 1; MemA != Stmt->begin(); MemA--) {
     auto *MemAccessPtr = *MemA;
     if (!MemAccessPtr->isLatestArrayKind())
       continue;
     if (!MemAccessPtr->isWrite())
       return false;
     auto AccMap = MemAccessPtr->getLatestAccessRelation();
-    if (isl_map_n_basic_map(AccMap.get()) != 1 ||
-        !isMatMulOperandAcc(AccMap, MMI.i, MMI.j))
+    if (!isMatMulOperandAcc(Stmt->getDomain(), AccMap, MMI.i, MMI.j))
       return false;
     MMI.WriteToC = MemAccessPtr;
     break;
   }
 
   if (!containsOnlyMatMulDep(PartialSchedule, D, MMI.k))
     return false;