[Polly][WIP] Restore the initial ordering of dimensions before applying the pattern matching

[gareevroman]

Summary

Dimensions of band nodes can be implicitly permuted by the algorithm applied during the schedule generation.

For example, in case of the following matrix-matrix multiplication,

for (i = 0; i < 1024; i++)

for (k = 0; k < 1024; k++)
  for (j = 0; j < 1024; j++)
    C[i][j] += A[i][k] * B[k][j];

it can produce the following schedule tree

domain: "{ Stmt_for_body6[i0, i1, i2] : 0 <= i0 <= 1023 and 0 <= i1 <= 1023 and 0 <= i2 <= 1023 }"
child:

schedule: "[{ Stmt_for_body6[i0, i1, i2] -> [(i0)] }, { Stmt_for_body6[i0, i1, i2] -> [(i1)] }, { Stmt_for_body6[i0, i1, i2] -> [(i2)] }]"
permutable: 1
coincident: [ 1, 1, 0 ]

The current implementation of the pattern matching optimizations relies on the initial ordering of dimensions. Otherwise, it can produce the miscompilation (e.g., [1]).

This patch helps to restore the initial ordering of dimensions by recreating the band node in case the corresponding conditions are satisfied.

Refs.:

[1] - https://bugs.llvm.org/show_bug.cgi?id=32500

[gareevroman]

lib/Transform/ScheduleOptimizer.cpp

1330 ↗

(On Diff #94320)

The changes related to ScheduleTreeOptimizer::isMatrMultPattern help to ensure that the band node represents all the dimensions of the iteration domain. Otherwise, it's not safe to recreate the band node.

polly/trunk/lib/Transform/ScheduleOptimizer.cpp

static isl_schedule_node *markInterIterationAliasFree(isl_schedule_node *Node,
   if (!BasePtr)
     return Node;
 
   auto *Ctx = isl_schedule_node_get_ctx(Node);
   auto *Id = isl_id_alloc(Ctx, "Inter iteration alias-free", BasePtr);
   return isl_schedule_node_child(isl_schedule_node_insert_mark(Node, Id), 0);
 }
 
+/// Restore the initial ordering of dimensions of the band node
+///
+/// In case the band node represents all the dimensions of the iteration
+/// domain, recreate the band node to restore the initial ordering of the
+/// dimensions.
+///
+/// @param Node The band node to be modified.
+/// @return The modified schedule node.
+namespace {
+isl::schedule_node getBandNodeWithOriginDimOrder(isl::schedule_node Node) {
+  assert(isl_schedule_node_get_type(Node.keep()) == isl_schedule_node_band);
+  if (isl_schedule_node_get_type(Node.child(0).keep()) !=
+      isl_schedule_node_leaf)
+    return Node;
+  auto Domain = isl::manage(isl_schedule_node_get_universe_domain(Node.keep()));
+  assert(isl_union_set_n_set(Domain.keep()) == 1);
+  if (isl_schedule_node_get_schedule_depth(Node.keep()) != 0 ||
+      (isl::set(isl::manage(Domain.copy())).dim(isl::dim::set) !=
+       isl_schedule_node_band_n_member(Node.keep())))
+    return Node;
+  Node = isl::manage(isl_schedule_node_delete(Node.take()));
+  auto PartialSchedulePwAff =
+      isl::manage(isl_union_set_identity_union_pw_multi_aff(Domain.take()));
+  auto PartialScheduleMultiPwAff =
+      isl::multi_union_pw_aff(PartialSchedulePwAff);
+  PartialScheduleMultiPwAff = isl::manage(isl_multi_union_pw_aff_reset_tuple_id(
+      PartialScheduleMultiPwAff.take(), isl_dim_set));
+  return isl::manage(isl_schedule_node_insert_partial_schedule(
+      Node.take(), PartialScheduleMultiPwAff.take()));
+}
+} // namespace
+
 __isl_give isl_schedule_node *ScheduleTreeOptimizer::optimizeMatMulPattern(
     __isl_take isl_schedule_node *Node, const llvm::TargetTransformInfo *TTI,
     MatMulInfoTy &MMI) {
   assert(TTI && "The target transform info should be provided.");
   Node = markInterIterationAliasFree(Node, MMI.WriteToC->getLatestBaseAddr());
   int DimOutNum = isl_schedule_node_band_n_member(Node);
   assert(DimOutNum > 2 && "In case of the matrix multiplication the loop nest "
                           "and, consequently, the corresponding scheduling "
                           "functions have at least three dimensions.");
+  Node = getBandNodeWithOriginDimOrder(isl::manage(Node)).take();
   Node = permuteBandNodeDimensions(Node, MMI.i, DimOutNum - 3);
   int NewJ = MMI.j == DimOutNum - 3 ? MMI.i : MMI.j;
   int NewK = MMI.k == DimOutNum - 3 ? MMI.i : MMI.k;
   Node = permuteBandNodeDimensions(Node, NewJ, DimOutNum - 2);
   NewK = MMI.k == DimOutNum - 2 ? MMI.j : MMI.k;
   Node = permuteBandNodeDimensions(Node, NewK, DimOutNum - 1);
   auto MicroKernelParams = getMicroKernelParams(TTI, MMI);
   auto MacroKernelParams = getMacroKernelParams(MicroKernelParams, MMI);
return optimizeDataLayoutMatrMulPattern(Node, MapOldIndVar, MicroKernelParams,
                                           MacroKernelParams, MMI);
 }
 
 bool ScheduleTreeOptimizer::isMatrMultPattern(
     __isl_keep isl_schedule_node *Node, const Dependences *D,
     MatMulInfoTy &MMI) {
   auto *PartialSchedule =
       isl_schedule_node_band_get_partial_schedule_union_map(Node);
-  if (isl_schedule_node_band_n_member(Node) < 3 ||
+  Node = isl_schedule_node_child(Node, 0);
+  auto LeafType = isl_schedule_node_get_type(Node);
+  Node = isl_schedule_node_parent(Node);
+  if (LeafType != isl_schedule_node_leaf ||
+      isl_schedule_node_band_n_member(Node) < 3 ||
+      isl_schedule_node_get_schedule_depth(Node) != 0 ||
       isl_union_map_n_map(PartialSchedule) != 1) {
     isl_union_map_free(PartialSchedule);
     return false;
   }
   auto *NewPartialSchedule = isl_map_from_union_map(PartialSchedule);
   if (containsMatrMult(NewPartialSchedule, D, MMI)) {
     isl_map_free(NewPartialSchedule);
     return true;

polly/trunk/test/ScheduleOptimizer/pattern-matching-based-opts_4.ll

+; RUN: opt %loadPolly -polly-opt-isl -polly-pattern-matching-based-opts=true \
+; RUN: -debug < %s 2>&1| FileCheck %s
+; RUN: opt %loadPolly -polly-opt-isl -polly-pattern-matching-based-opts=true \
+; RUN: -polly-target-throughput-vector-fma=1 \
+; RUN: -polly-target-latency-vector-fma=8 \
+; RUN: -polly-target-1st-cache-level-size=32768 \
+; RUN: -polly-target-vector-register-bitwidth=256 \
+; RUN: -polly-target-2nd-cache-level-size=262144 -polly-ast \
+; RUN: -analyze < %s | FileCheck %s --check-prefix=PATTERN-MATCHING-OPTS
+; REQUIRES: asserts
+;
+;    C := A * B + C
+;    Check that the pattern matching optimizations can detect different
+;    permutations of GEMM loop and produce the correct ISL AST. In this case,
+;    dimensions of band nodes can be implicitly permuted by the algorithm
+;    applied during the schedule generation. It should be taken into the
+;    account during the pattern matching optimizations.
+;    for (i = 0; i < _PB_NI; i++)
+;      for (k = 0; k < _PB_NK; ++k)
+;        for (j = 0; j < _PB_NJ; j++)
+;    C[i][j] += A[i][k] * B[k][j];
+;
+; CHECK: The matrix multiplication pattern was detected
+;
+; PATTERN-MATCHING-OPTS:    // 1st level tiling - Tiles
+; PATTERN-MATCHING-OPTS-NEXT:    for (int c1 = 0; c1 <= 3; c1 += 1) {
+; PATTERN-MATCHING-OPTS-NEXT:      for (int c3 = 256 * c1; c3 <= 256 * c1 + 255; c3 += 1)
+; PATTERN-MATCHING-OPTS-NEXT:        for (int c4 = 0; c4 <= 1023; c4 += 1)
+; PATTERN-MATCHING-OPTS-NEXT:          CopyStmt_0(0, c3, c4);
+; PATTERN-MATCHING-OPTS-NEXT:      for (int c2 = 0; c2 <= 10; c2 += 1) {
+; PATTERN-MATCHING-OPTS-NEXT:        for (int c3 = 96 * c2; c3 <= min(1023, 96 * c2 + 95); c3 += 1)
+; PATTERN-MATCHING-OPTS-NEXT:          for (int c4 = 256 * c1; c4 <= 256 * c1 + 255; c4 += 1)
+; PATTERN-MATCHING-OPTS-NEXT:            CopyStmt_1(c3, c4, 0);
+; PATTERN-MATCHING-OPTS-NEXT:        // 1st level tiling - Points
+; PATTERN-MATCHING-OPTS-NEXT:        // Register tiling - Tiles
+; PATTERN-MATCHING-OPTS-NEXT:        for (int c3 = 0; c3 <= 127; c3 += 1)
+; PATTERN-MATCHING-OPTS-NEXT:          for (int c4 = 0; c4 <= min(23, -24 * c2 + 255); c4 += 1)
+; PATTERN-MATCHING-OPTS-NEXT:            for (int c5 = 0; c5 <= 255; c5 += 1) {
+; PATTERN-MATCHING-OPTS-NEXT:              // Register tiling - Points
+; PATTERN-MATCHING-OPTS-NEXT:              {
+; PATTERN-MATCHING-OPTS-NEXT:                Stmt_for_body6(96 * c2 + 4 * c4, 256 * c1 + c5, 8 * c3);
+; PATTERN-MATCHING-OPTS-NEXT:                Stmt_for_body6(96 * c2 + 4 * c4, 256 * c1 + c5, 8 * c3 + 1);
+; PATTERN-MATCHING-OPTS-NEXT:                Stmt_for_body6(96 * c2 + 4 * c4, 256 * c1 + c5, 8 * c3 + 2);
+; PATTERN-MATCHING-OPTS-NEXT:                Stmt_for_body6(96 * c2 + 4 * c4, 256 * c1 + c5, 8 * c3 + 3);
+; PATTERN-MATCHING-OPTS-NEXT:                Stmt_for_body6(96 * c2 + 4 * c4, 256 * c1 + c5, 8 * c3 + 4);
+; PATTERN-MATCHING-OPTS-NEXT:                Stmt_for_body6(96 * c2 + 4 * c4, 256 * c1 + c5, 8 * c3 + 5);
+; PATTERN-MATCHING-OPTS-NEXT:                Stmt_for_body6(96 * c2 + 4 * c4, 256 * c1 + c5, 8 * c3 + 6);
+; PATTERN-MATCHING-OPTS-NEXT:                Stmt_for_body6(96 * c2 + 4 * c4, 256 * c1 + c5, 8 * c3 + 7);
+; PATTERN-MATCHING-OPTS-NEXT:                Stmt_for_body6(96 * c2 + 4 * c4 + 1, 256 * c1 + c5, 8 * c3);
+; PATTERN-MATCHING-OPTS-NEXT:                Stmt_for_body6(96 * c2 + 4 * c4 + 1, 256 * c1 + c5, 8 * c3 + 1);
+; PATTERN-MATCHING-OPTS-NEXT:                Stmt_for_body6(96 * c2 + 4 * c4 + 1, 256 * c1 + c5, 8 * c3 + 2);
+; PATTERN-MATCHING-OPTS-NEXT:                Stmt_for_body6(96 * c2 + 4 * c4 + 1, 256 * c1 + c5, 8 * c3 + 3);
+; PATTERN-MATCHING-OPTS-NEXT:                Stmt_for_body6(96 * c2 + 4 * c4 + 1, 256 * c1 + c5, 8 * c3 + 4);
+; PATTERN-MATCHING-OPTS-NEXT:                Stmt_for_body6(96 * c2 + 4 * c4 + 1, 256 * c1 + c5, 8 * c3 + 5);
+; PATTERN-MATCHING-OPTS-NEXT:                Stmt_for_body6(96 * c2 + 4 * c4 + 1, 256 * c1 + c5, 8 * c3 + 6);
+; PATTERN-MATCHING-OPTS-NEXT:                Stmt_for_body6(96 * c2 + 4 * c4 + 1, 256 * c1 + c5, 8 * c3 + 7);
+; PATTERN-MATCHING-OPTS-NEXT:                Stmt_for_body6(96 * c2 + 4 * c4 + 2, 256 * c1 + c5, 8 * c3);
+; PATTERN-MATCHING-OPTS-NEXT:                Stmt_for_body6(96 * c2 + 4 * c4 + 2, 256 * c1 + c5, 8 * c3 + 1);
+; PATTERN-MATCHING-OPTS-NEXT:                Stmt_for_body6(96 * c2 + 4 * c4 + 2, 256 * c1 + c5, 8 * c3 + 2);
+; PATTERN-MATCHING-OPTS-NEXT:                Stmt_for_body6(96 * c2 + 4 * c4 + 2, 256 * c1 + c5, 8 * c3 + 3);
+; PATTERN-MATCHING-OPTS-NEXT:                Stmt_for_body6(96 * c2 + 4 * c4 + 2, 256 * c1 + c5, 8 * c3 + 4);
+; PATTERN-MATCHING-OPTS-NEXT:                Stmt_for_body6(96 * c2 + 4 * c4 + 2, 256 * c1 + c5, 8 * c3 + 5);
+; PATTERN-MATCHING-OPTS-NEXT:                Stmt_for_body6(96 * c2 + 4 * c4 + 2, 256 * c1 + c5, 8 * c3 + 6);
+; PATTERN-MATCHING-OPTS-NEXT:                Stmt_for_body6(96 * c2 + 4 * c4 + 2, 256 * c1 + c5, 8 * c3 + 7);
+; PATTERN-MATCHING-OPTS-NEXT:                Stmt_for_body6(96 * c2 + 4 * c4 + 3, 256 * c1 + c5, 8 * c3);
+; PATTERN-MATCHING-OPTS-NEXT:                Stmt_for_body6(96 * c2 + 4 * c4 + 3, 256 * c1 + c5, 8 * c3 + 1);
+; PATTERN-MATCHING-OPTS-NEXT:                Stmt_for_body6(96 * c2 + 4 * c4 + 3, 256 * c1 + c5, 8 * c3 + 2);
+; PATTERN-MATCHING-OPTS-NEXT:                Stmt_for_body6(96 * c2 + 4 * c4 + 3, 256 * c1 + c5, 8 * c3 + 3);
+; PATTERN-MATCHING-OPTS-NEXT:                Stmt_for_body6(96 * c2 + 4 * c4 + 3, 256 * c1 + c5, 8 * c3 + 4);
+; PATTERN-MATCHING-OPTS-NEXT:                Stmt_for_body6(96 * c2 + 4 * c4 + 3, 256 * c1 + c5, 8 * c3 + 5);
+; PATTERN-MATCHING-OPTS-NEXT:                Stmt_for_body6(96 * c2 + 4 * c4 + 3, 256 * c1 + c5, 8 * c3 + 6);
+; PATTERN-MATCHING-OPTS-NEXT:                Stmt_for_body6(96 * c2 + 4 * c4 + 3, 256 * c1 + c5, 8 * c3 + 7);
+; PATTERN-MATCHING-OPTS-NEXT:              }
+; PATTERN-MATCHING-OPTS-NEXT:            }
+; PATTERN-MATCHING-OPTS-NEXT:      }
+; PATTERN-MATCHING-OPTS-NEXT:    }
+;
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-unknown"
+
+define internal void @kernel_gemm(i32 %ni, i32 %nj, i32 %nk, double %alpha, double %beta, [1024 x double]* %C, [1024 x double]* %A, [1024 x double]* %B) {
+entry:
+  br label %entry.split
+
+entry.split:                                      ; preds = %entry
+  br label %for.cond1.preheader
+
+for.cond1.preheader:                              ; preds = %for.inc20, %entry.split
+  %indvars.iv41 = phi i64 [ 0, %entry.split ], [ %indvars.iv.next42, %for.inc20 ]
+  br label %for.cond4.preheader
+
+for.cond4.preheader:                              ; preds = %for.inc17, %for.cond1.preheader
+  %indvars.iv38 = phi i64 [ 0, %for.cond1.preheader ], [ %indvars.iv.next39, %for.inc17 ]
+  br label %for.body6
+
+for.body6:                                        ; preds = %for.body6, %for.cond4.preheader
+  %indvars.iv = phi i64 [ 0, %for.cond4.preheader ], [ %indvars.iv.next, %for.body6 ]
+  %arrayidx8 = getelementptr inbounds [1024 x double], [1024 x double]* %A, i64 %indvars.iv41, i64 %indvars.iv38
+  %tmp = load double, double* %arrayidx8, align 8
+  %arrayidx12 = getelementptr inbounds [1024 x double], [1024 x double]* %B, i64 %indvars.iv38, i64 %indvars.iv
+  %tmp1 = load double, double* %arrayidx12, align 8
+  %mul = fmul double %tmp, %tmp1
+  %arrayidx16 = getelementptr inbounds [1024 x double], [1024 x double]* %C, i64 %indvars.iv41, i64 %indvars.iv
+  %tmp2 = load double, double* %arrayidx16, align 8
+  %add = fadd double %tmp2, %mul
+  store double %add, double* %arrayidx16, align 8
+  %indvars.iv.next = add nuw nsw i64 %indvars.iv, 1
+  %exitcond = icmp ne i64 %indvars.iv.next, 1024
+  br i1 %exitcond, label %for.body6, label %for.inc17
+
+for.inc17:                                        ; preds = %for.body6
+  %indvars.iv.next39 = add nuw nsw i64 %indvars.iv38, 1
+  %exitcond40 = icmp ne i64 %indvars.iv.next39, 1024
+  br i1 %exitcond40, label %for.cond4.preheader, label %for.inc20
+
+for.inc20:                                        ; preds = %for.inc17
+  %indvars.iv.next42 = add nuw nsw i64 %indvars.iv41, 1
+  %exitcond43 = icmp ne i64 %indvars.iv.next42, 1024
+  br i1 %exitcond43, label %for.cond1.preheader, label %for.end22
+
+for.end22:                                        ; preds = %for.inc20
+  ret void
+}