[GSoC 2016] [Polly] [WIP] Apply all necessary tilings and interchangings to get a macro-kernel

[gareevroman]

Summary

This is the second patch to apply the BLIS matmul optimization pattern on matmul kernels (http://www.cs.utexas.edu/users/flame/pubs/TOMS-BLIS-Analytical.pdf). BLIS implements gemm as three nested loops around a macro-kernel, plus two packing routines. The macro-kernel is implemented in terms of two additional loops around a micro-kernel. The micro-kernel is a loop around a rank-1 (i.e., outer product) update. In this change we create the BLIS macro-kernel by applying a combination of tiling and interchanging. In subsequent changes we will implement the packing transformation.

[gareevroman]

This is a draft of the second patch to apply the BLIS matmul optimization pattern on matmul kernels. I'm not sure how it is better to implement interchanging of loops using schedule trees. That's why I would be very grateful for your comments, feedback and ideas.

I think that there are two ways to do it. The first way is to split a band node two times and perform sink to get the desired order of dimensions. In this case, the subtree of the band node would probably cause issues, if, for example, it contains mark nodes. As far as I know, they can't be sinked and should be manually reinserted.

The second way is to manually interchange dimensions of a partial schedule, delete the band node and insert a band node, which contains a new partial schedule. I've implemented it in the patch. However, there is a possible issue, because isl_schedule_node_delete can't be applied to a band node with an anchored subtree and isl_schedule_node_insert_partial_schedule can't be be applied to the schedule tree, which has anchored nodes.

[Meinersbur]

Thanks for the patch and meaningful comments.

include/polly/ScheduleOptimizer.h

134

vice versa

261–270

Please don't copy-paste the same text. There is a @see tag to refer to other parts.

lib/Transform/ScheduleOptimizer.cpp

539

also check for CacheLevelAssociativityDegrees.size() >=2 to avoid out-of-range errors?

553

You can use

int MacroKernelParams[] = {Mc, Nc, Kc};

here. It is automatically converted to ArrayRef<int>.

560–562

Style: No single return in an else branch. The following style seems to be established.

  if (c) {
    ...
    return X;
  }
  return Y;
}

579

Add an empty line befire to make clear that what is coming is another section?

[gareevroman]

In D21491#462093, @Meinersbur wrote:

Thanks for the patch and meaningful comments.

Thank you for the comments! I'll try to address them soon.

[gareevroman]

Changed according to the comments.

[gareevroman]

lib/Transform/ScheduleOptimizer.cpp

539

As far as I understand, it requires information about the first two levels of cache to determine all the parameters (Kc, Mc, Nc) of a macro-kernel. I've added a corresponding comment in the new version of the patch.

560–562

Thanks. In the new version I've tried to also take into account "Use Early Exits and continue to Simplify Code"

578

Yes, it's the reason. Probably, it's not good. I wasn't sure wether we should create a function to create a micro-kernel that contains only "std::vector<int> MicroKernelParams{Mr, Nr};" and "Node = applyRegisterTiling(Node, MicroKernelParams, 1);"

[grosser]

Hi Roman,

thanks for this patch and Michael thank you for your review! I added a couple of comments.

Best,
Tobias

include/polly/ScheduleOptimizer.h

128–129

Not part of this commit, but please add in a separate commit (no review needed) full title/author information to this paper, in case this link becomes dead at some point.

129

"interchanging The"

This seems incomplete.

131

This paper is already referenced earlier. No need to reference it again.

I personally would suggest to use @see to reference your implementation of the BLIS macro kernel, which than contains the macro-kernel specific information -- in sufficient detail to be self-contained without the BLIS paper.

132

Drop "it should be noticed". It does not add any additional information.

133

Drop "should". It does not add any specific information.

139

Make this a "TODO:"

139

Make this a TODO:

265

Can you document what "Mr" describes?

Also, the documentation of this function requires to have read the BLIS paper to understand what this parameter means. That seems not ideal. I suggest to expand this documentation to be (mostly) self contained.

(It seems you may have had more documentation here before. I suggest to move the documentation you have from optimizeMatMulPattern and add the @see in the other function).

lib/Transform/ScheduleOptimizer.cpp

138

What about" -polly-target-cache-level-associativity - The associativity of each cache level"

and CacheLevelAssociativity

144

What about "-polly-target-cache-level-sizes - The size of each cache level"?

and CacheLevelSizes

531

Nice!

540

than two.

554

In the subsequent patch you pull out the parameter calculation in a helper function. I suggest to do this already in this patch as otherwise code is just unnecessarily moved around.

559

Instead of just storing these parameters in a list of integers, I suggest to introduce a struct that allows us to encapsulate these parameters and access them by name.

struct MacroKernelParamsTy {                                                     
  int Mc;                                                                        
  int Nc;                                                                        
  int Kc;                                                                        
};                                                                               
                                                                                 
void foo() {                                                                     
  int Mc, Nc, Kc;                                                                
  struct MacroKernelParamsTy MacroKernelParams = {Mc, Nc, Kc};                   
}

560–562

Nice!

578

I am in favor of a new function. It might be mostly empty, but it helps to the code.

[gareevroman]

Hi Tobias,

Thank you for the comments! I've tried to address them. Should we pull out creation of the BLIS micro-kernel and its parameters in helper functions in a separate patch?

[gareevroman]

include/polly/ScheduleOptimizer.h

128–129

I'll try to do it soon.

[gareevroman]

The test case was missed.

[grosser]

Hi Roman,

this patch looks very good to me now. Please commit changes that just improve the code structure or documentation of existing code separately.

We can probably still improve the documentation of the cost functions further, but for now this looks sufficient. After all the infrastructure is in place, I expect that we perform further experiments and performance tuning. The results of these efforts will -- at the end of your summer of code -- certainly help us to extend the documentation.

include/polly/ScheduleOptimizer.h

48

Nice!

128–157

Very nice documentation!

144

"account in creation" -> "account during the creation of the BLIS of the BLIS kernels and the computation of their parameters."

155

The node is required to successfully pass ScheduleTreeOptimizer::isMatrMultPattern

263

The values of ...

277

The values passed in MicroKernelParam are used ...

lib/Transform/ScheduleOptimizer.cpp

538

Nice. This change and some of the changes to improve documentation can be committed separately.

548

Its 2nd for "second" or 1st for "first, "1nd" does not make sense.

557

We choose the Mr and Nr parameters of the micro kernel to be large enough such that no stalls

559

the updates _of_ the resulting

561

be as small as possible

583

During the computation of the matrix multiplication,

586

"It" sounds not right. What is it referring to?

[Meinersbur]

lib/Transform/ScheduleOptimizer.cpp

535

MicroKernelParamsTy,MacroKernelParamsTy are small enough to be passed by value.

[gareevroman]

lib/Transform/ScheduleOptimizer.cpp

535

MicroKernelParamsTy,MacroKernelParamsTy are small enough to be passed by value.

Hi Michael,

shouldn't we follow advices from "Effective C++", according to LLVM Coding Standards?

As far as I understand, in item 20 it is stated that the only types for which we can reasonably assume that pass-by-value is inexpensive are build-in types and STL iterator and function object types.

[Meinersbur]

lib/Transform/ScheduleOptimizer.cpp

535

LLVM has its own types that are neither STL or built-ins that are meant to be passes by value: ArrayRef, StringRef, PointerIntPair, CallSite, etc.

"Effective C++" in the LLVM Coding Standards is referenced as the source of some recommendations, not as implicit import of rules.

We can reasonable assume that passing these types by value is inexpensive as we declared them ourselves as POD with few members in its header.

Since C++11 it can be better to pass even heavy objects by value if the function is going to consume/take ownership using std::move().

However, please feel free to not change it if you prefer this style.

[mreisinger]

test/ScheduleOptimizer/pattern-matching-based-opts_3.ll

66 ↗

(On Diff #63846)

"1nd" - see Tobias' note from above

70 ↗

(On Diff #63846)

"1nd" - see Tobias' note from above

[gareevroman]

lib/Transform/ScheduleOptimizer.cpp

535

Thanks for the explanation.

LLVM has its own types that are neither STL or built-ins that are meant to be passes by value: ArrayRef, StringRef, PointerIntPair, CallSite, etc.
"Effective C++" in the LLVM Coding Standards is referenced as the source of some recommendations, not as implicit import of rules.

As far as I understand, according to the description, LLVM Coding Standards should not be regarded as an absolute requirement to be followed in all instances and "Effective C++" is one of the two books, which are proposed as particularly important books for work. I just tried to understand how much they are important for us.

We can reasonable assume that passing these types by value is inexpensive as we declared them ourselves as POD with few members in its header.
Since C++11 it can be better to pass even heavy objects by value if the function is going to consume/take ownership using std::move().

Yes, I agree with it. My understanding is that in this case passing by reference can be a precaution, because implementation of a user-defined type can be changed in future.

However, please feel free to not change it if you prefer this style.

I prefer to follow LLVM Coding Standards. OK, I'll fix it so they are passed by value.

[gareevroman]

@grosser, @Meinersbur, @mreisinger, thank you for the comments! I've tried to address them in r276616 and r276627.

include/polly/ScheduleOptimizer.h

 struct isl_schedule_node;
 struct isl_union_map;
 
 namespace polly {
 extern bool DisablePollyTiling;
 class Scop;
 } // namespace polly
 
+/// @brief Parameters of the micro kernel.
+///
+/// Parameters, which determine sizes of rank-1 (i.e., outer product) update
+/// used in the optimized matrix multiplication.
+///
+struct MicroKernelParamsTy {
+  int Mr;
+  int Nr;
+};
+
+/// @brief Parameters of the macro kernel.
+///
+/// Parameters, which determine sizes of blocks of partitioned matrices
+/// used in the optimized matrix multiplication.
+///
+struct MacroKernelParamsTy {
+  int Mc;
+  int Nc;
+  int Kc;
+};
+
 class ScheduleTreeOptimizer {
 public:
   /// @brief Apply schedule tree transformations.
   ///
   /// This function takes an (possibly already optimized) schedule tree and
   /// applies a set of additional optimizations on the schedule tree. The
   /// transformations applied include:
   ///
private:
   ///
   /// @param Node            The node to register tile.
   /// @param TileSizes       A vector of tile sizes that should be used for
   ///                        tiling.
   /// @param DefaultTileSize A default tile size that is used for dimensions
   static __isl_give isl_schedule_node *
   applyRegisterTiling(__isl_take isl_schedule_node *Node,
                       llvm::ArrayRef<int> TileSizes, int DefaultTileSize);
 
-  /// @brief Apply the BLIS matmul optimization pattern
+  /// @brief Apply the BLIS matmul optimization pattern.
   ///
-  /// Apply the BLIS matmul optimization pattern
-  /// (http://www.cs.utexas.edu/users/flame/pubs/TOMS-BLIS-Analytical.pdf).
-  /// BLIS implements gemm as three nested loops around a macro-kernel,
-  /// plus two packing routines. The macro-kernel is implemented in terms
-  /// of two additional loops around a micro-kernel. The micro-kernel
-  /// is a loop around a rank-1 (i.e., outer product) update.
-  ///
-  /// We create the BLIS micro-kernel by applying a combination of tiling
-  /// and unrolling. In subsequent changes we will add the extraction
-  /// of the BLIS macro-kernel and implement the packing transformation.
-  ///
-  /// It is assumed that the Node is successfully checked
-  /// by ScheduleTreeOptimizer::isMatrMultPattern. Consequently
-  /// in case of matmul kernels the application of optimizeMatMulPattern
-  /// can lead to close-to-peak performance. Maybe it can be generalized
-  /// to effectively optimize the whole class of successfully checked
-  /// statements.
-  ///
-  /// @param Node the node that contains a band to be optimized.
+  /// Apply the BLIS matmul optimization pattern. BLIS implements gemm as three
+  /// nested loops around a macro-kernel, plus two packing routines.
+  /// The macro-kernel is implemented in terms of two additional loops around
+  /// a micro-kernel. The micro-kernel is a loop around a rank-1
+  /// (i.e., outer product) update.
+  ///
+  /// For a detailed description please see:
+  /// Analytical Modeling is Enough for High Performance BLIS
+  /// Tze Meng Low, Francisco D Igual, Tyler M Smith, Enrique S Quintana-Orti
+  /// Technical Report, 2014
+  /// http://www.cs.utexas.edu/users/flame/pubs/TOMS-BLIS-Analytical.pdf
+  ///
+  /// In our case matrices are stored in row-major order, which is taken into
+  /// account in creation of BLIS kernels and determination of their
+  /// parameters.
+  ///
+  /// @see ScheduleTreeOptimizer::createMicroKernel
+  /// @see ScheduleTreeOptimizer::createMacroKernel
+  /// @see getMicroKernelParams
+  /// @see getMacroKernelParams
+  ///
+  /// TODO: Implement the packing transformation.
+  ///
+  /// @param Node The node that contains a band to be optimized. It is assumed
+  ///             that it is successfully checked.
+  ///             by ScheduleTreeOptimizer::isMatrMultPattern.
   /// @return Modified isl_schedule_node.
   static __isl_give isl_schedule_node *
   optimizeMatMulPattern(__isl_take isl_schedule_node *Node,
                         const llvm::TargetTransformInfo *TTI);
 
   /// @brief Check if this node is a band node we want to tile.
   ///
   /// We look for innermost band nodes where individual dimensions are marked as
   /// permutable.
private:
   ///    read memory access and the last one is write memory access.
   /// 5. all subscripts of the last memory access of the statement don't
   ///    contain the variable used in the inner loop.
   /// If this is the case, we could try to use an approach that is similar to
   /// the one used to get close-to-peak performance of matrix multiplications.
   ///
   /// @param Node The node to check.
   static bool isMatrMultPattern(__isl_keep isl_schedule_node *Node);
+
+  /// @brief Create the BLIS macro-kernel.
+  ///
+  /// We create the BLIS macro-kernel by applying a combination of tiling
+  /// of dimensions of the band node and interchanging of two innermost
+  /// modified dimensions. Values of MacroKernelParams's fields are used
+  /// as tile sizes.
+  ///
+  /// @param Node The schedule node to be modified.
+  /// @param MacroKernelParams Parameters of the macro kernel
+  ///                          to be used as tile sizes.
+  static __isl_give isl_schedule_node *
+  createMacroKernel(__isl_take isl_schedule_node *Node,
+                    const MacroKernelParamsTy &MacroKernelParams);
+
+  /// @brief Create the BLIS micro-kernel
+  ///
+  /// We create the BLIS micro-kernel by applying a combination of tiling
+  /// of two outer dimensions of the band node and unrolling of the modified
+  /// schedule dimensions. Values of MicroKernelParams's fields are used
+  /// as tile sizes.
+  ///
+  /// @param Node The schedule node to be modified.
+  /// @param MicroKernelParams Parameters of the micro kernel
+  ///                          to be used as tile sizes.
+  /// @see MicroKernelParamsTy
+  static __isl_give isl_schedule_node *
+  createMicroKernel(__isl_take isl_schedule_node *Node,
+                    const MicroKernelParamsTy &MicroKernelParams);
 };
 
 #endif

lib/Transform/ScheduleOptimizer.cpp

 
 static cl::opt<int> ThrougputVectorFma(
     "polly-target-througput-vector-fma",
     cl::desc("A throughput of the processor floating-point arithmetic units "
              "expressed in the number of vector fused multiply-add "
              "instructions per clock cycle."),
     cl::Hidden, cl::init(1), cl::ZeroOrMore, cl::cat(PollyCategory));
 
+static cl::list<int>
+    CacheLevelAssociativity("polly-target-cache-level-associativity",
+                            cl::desc("The associativity of each cache level."),
+                            cl::Hidden, cl::ZeroOrMore, cl::CommaSeparated,
+                            cl::cat(PollyCategory));
+
+static cl::list<int> CacheLevelSizes(
+    "polly-target-cache-level-sizes",
+    cl::desc("The size of each cache level specified in bytes."), cl::Hidden,
+    cl::ZeroOrMore, cl::CommaSeparated, cl::cat(PollyCategory));
+
 static cl::opt<int> FirstLevelDefaultTileSize(
     "polly-default-tile-size",
     cl::desc("The default tile size (if not enough were provided by"
              " --polly-tile-sizes)"),
     cl::Hidden, cl::init(32), cl::ZeroOrMore, cl::cat(PollyCategory));
 
 static cl::list<int> FirstLevelTileSizes(
     "polly-tile-sizes", cl::desc("A tile size for each loop dimension, filled "
static __isl_give isl_map *circularShiftOutputDims(__isl_take isl_map *IslMap) {
   if (DimNum == 0)
     return IslMap;
   auto InputDimsId = isl_map_get_tuple_id(IslMap, isl_dim_in);
   IslMap = isl_map_move_dims(IslMap, isl_dim_in, 0, isl_dim_out, DimNum - 1, 1);
   IslMap = isl_map_move_dims(IslMap, isl_dim_out, 0, isl_dim_in, 0, 1);
   return isl_map_set_tuple_id(IslMap, isl_dim_in, InputDimsId);
 }
 
-__isl_give isl_schedule_node *ScheduleTreeOptimizer::optimizeMatMulPattern(
-    __isl_take isl_schedule_node *Node, const llvm::TargetTransformInfo *TTI) {
+/// @brief Permute two dimensions of the band node.
+///
+/// Permute FirstDim and SecondDim dimensions of the Node.
+///
+/// @param Node The band node to be modified.
+/// @param FirstDim The first dimension to be permuted.
+/// @param SecondDim The second dimension to be permuted.
+static __isl_give isl_schedule_node *
+permuteBandNodeDimensions(__isl_take isl_schedule_node *Node, unsigned FirstDim,
+                          unsigned SecondDim) {
+  assert(isl_schedule_node_get_type(Node) == isl_schedule_node_band &&
+         isl_schedule_node_band_n_member(Node) > std::max(FirstDim, SecondDim));
+  auto PartialSchedule = isl_schedule_node_band_get_partial_schedule(Node);
+  auto PartialScheduleFirstDim =
+      isl_multi_union_pw_aff_get_union_pw_aff(PartialSchedule, FirstDim);
+  auto PartialScheduleSecondDim =
+      isl_multi_union_pw_aff_get_union_pw_aff(PartialSchedule, SecondDim);
+  PartialSchedule = isl_multi_union_pw_aff_set_union_pw_aff(
+      PartialSchedule, SecondDim, PartialScheduleFirstDim);
+  PartialSchedule = isl_multi_union_pw_aff_set_union_pw_aff(
+      PartialSchedule, FirstDim, PartialScheduleSecondDim);
+  Node = isl_schedule_node_delete(Node);
+  Node = isl_schedule_node_insert_partial_schedule(Node, PartialSchedule);
+  return Node;
+}
+
+__isl_give isl_schedule_node *ScheduleTreeOptimizer::createMicroKernel(
+    __isl_take isl_schedule_node *Node,
+    const MicroKernelParamsTy &MicroKernelParams) {
+  return applyRegisterTiling(Node, {MicroKernelParams.Mr, MicroKernelParams.Nr},
+                             1);
+}
+
+__isl_give isl_schedule_node *ScheduleTreeOptimizer::createMacroKernel(
+    __isl_take isl_schedule_node *Node,
+    const MacroKernelParamsTy &MacroKernelParams) {
+  assert(isl_schedule_node_get_type(Node) == isl_schedule_node_band);
+  if (MacroKernelParams.Mc == 1 && MacroKernelParams.Nc == 1 &&
+      MacroKernelParams.Kc == 1)
+    return Node;
+  Node = tileNode(
+      Node, "1nd level tiling",
+      {MacroKernelParams.Mc, MacroKernelParams.Nc, MacroKernelParams.Kc}, 1);
+  Node = isl_schedule_node_parent(isl_schedule_node_parent(Node));
+  Node = permuteBandNodeDimensions(Node, 1, 2);
+  return isl_schedule_node_child(isl_schedule_node_child(Node, 0), 0);
+}
+
+/// Get parameters of the BLIS micro kernel.
+///
+/// We try to determine Mr and Nr parameters of the micro kernel that
+/// are large enough so that no stalls caused by the combination
+/// of latencies and dependencies are introduced during the updates
+/// the resulting matrix of the matrix multiplication. However,
+/// they should also be smallest as possible to release more registers
+/// for entries of multiplied matrices.
+///
+/// @param TTI Target Transform Info.
+/// @return The structure of type MicroKernelParamsTy.
+/// @see MicroKernelParamsTy
+static struct MicroKernelParamsTy
+getMicroKernelParams(const llvm::TargetTransformInfo *TTI) {
   assert(TTI && "The target transform info should be provided.");
-  // Get a micro-kernel.
   // Nvec - Number of double-precision floating-point numbers that can be hold
   // by a vector register. Use 2 by default.
   auto Nvec = TTI->getRegisterBitWidth(true) / 64;
   if (Nvec == 0)
     Nvec = 2;
   int Nr =
       ceil(sqrt(Nvec * LatencyVectorFma * ThrougputVectorFma) / Nvec) * Nvec;
   int Mr = ceil(Nvec * LatencyVectorFma * ThrougputVectorFma / Nr);
-  std::vector<int> MicroKernelParams{Mr, Nr};
-  Node = applyRegisterTiling(Node, MicroKernelParams, 1);
+  return {Mr, Nr};
+}
+
+/// Get parameters of the BLIS macro kernel.
+///
+/// During computation of matrix multiplication, blocks of partitioned
+/// matrices are mapped to different layers of the memory hierarchy.
+/// To optimize data reuse, blocks should be ideally kept in cache between
+/// iterations. It, along with parameters of the micro-kernel, imposes bounds
+/// on parameters of the macro kernel.
+///
+/// @param MicroKernelParams Parameters of the micro-kernel
+///                          to be taken into account
+/// @return The structure of type MacroKernelParamsTy
+/// @see MacroKernelParamsTy
+/// @see MicroKernelParamsTy
+static struct MacroKernelParamsTy
+getMacroKernelParams(const MicroKernelParamsTy &MicroKernelParams) {
+  // According to www.cs.utexas.edu/users/flame/pubs/TOMS-BLIS-Analytical.pdf,
+  // it requires information about the first two levels of a cache to determine
+  // all the parameters of a macro-kernel. It also checks that an associativity
+  // degree of a cache level is greater than two. Otherwise, another algorithm
+  // for determination of the parameters should be used.
+  if (!(MicroKernelParams.Mr > 0 && MicroKernelParams.Nr > 0 &&
+        CacheLevelSizes.size() >= 2 && CacheLevelAssociativity.size() >= 2 &&
+        CacheLevelSizes[0] > 0 && CacheLevelSizes[1] > 0 &&
+        CacheLevelAssociativity[0] > 2 && CacheLevelAssociativity[1] > 2))
+    return {1, 1, 1};
+  int Cbr = floor(
+      (CacheLevelAssociativity[0] - 1) /
+      (1 + static_cast<double>(MicroKernelParams.Mr) / MicroKernelParams.Nr));
+  int Kc = (Cbr * CacheLevelSizes[0]) /
+           (MicroKernelParams.Nr * CacheLevelAssociativity[0] * 8);
+  double Cac = static_cast<double>(MicroKernelParams.Mr * Kc * 8 *
+                                   CacheLevelAssociativity[1]) /
+               CacheLevelSizes[1];
+  double Cbc = static_cast<double>(MicroKernelParams.Nr * Kc * 8 *
+                                   CacheLevelAssociativity[1]) /
+               CacheLevelSizes[1];
+  int Mc = floor(MicroKernelParams.Mr / Cac);
+  int Nc =
+      floor((MicroKernelParams.Nr * (CacheLevelAssociativity[1] - 2)) / Cbc);
+  return {Mc, Nc, Kc};
+}
+
+__isl_give isl_schedule_node *ScheduleTreeOptimizer::optimizeMatMulPattern(
+    __isl_take isl_schedule_node *Node, const llvm::TargetTransformInfo *TTI) {
+  assert(TTI && "The target transform info should be provided.");
+  auto MicroKernelParams = getMicroKernelParams(TTI);
+  auto MacroKernelParams = getMacroKernelParams(MicroKernelParams);
+  Node = createMacroKernel(Node, MacroKernelParams);
+  Node = createMicroKernel(Node, MicroKernelParams);
   return Node;
 }
 
 bool ScheduleTreeOptimizer::isMatrMultPattern(
     __isl_keep isl_schedule_node *Node) {
   auto *PartialSchedule =
       isl_schedule_node_band_get_partial_schedule_union_map(Node);
   if (isl_schedule_node_band_n_member(Node) != 3 ||