Transforms/Utils/MatrixUtils.h

0001 //===- MatrixUtils.h - Utilities to lower matrix intrinsics -----*- C++ -*-===//
0002 //
0003 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
0004 // See https://llvm.org/LICENSE.txt for license information.
0005 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
0006 //
0007 //===----------------------------------------------------------------------===//
0008 //
0009 // Utilities for generating tiled loops for matrix operations.
0010 //
0011 //===----------------------------------------------------------------------===//
0012
0013 #ifndef LLVM_TRANSFORMS_UTILS_MATRIXUTILS_H
0014 #define LLVM_TRANSFORMS_UTILS_MATRIXUTILS_H
0015
0016 #include "llvm/ADT/StringRef.h"
0017
0018 namespace llvm {
0019 class DomTreeUpdater;
0020 class BasicBlock;
0021 class Value;
0022 class Loop;
0023 class LoopInfo;
0024 class IRBuilderBase;
0025
0026 /// A helper struct to create IR loop nests for tiling in IR of the following
0027 /// form:
0028 ///   for ColumnLoop.Index = 0..NumColumns
0029 ///     for RowLoop.Index = 0..NumRows
0030 ///       for KLoop.Index = 0..NumInner
0031 struct TileInfo {
0032   /// Number of rows of the matrix.
0033   unsigned NumRows;
0034
0035   /// Number of columns of the matrix.
0036   unsigned NumColumns;
0037
0038   /// Number of columns of the first matrix of a multiply /
0039   /// number of rows of the second matrix of a multiply.
0040   unsigned NumInner;
0041
0042   /// Number of rows/columns in a tile.
0043   unsigned TileSize = -1;
0044
0045   /// Properties of a single loop used when generating the tiled loop nest.
0046   struct MatrixLoop {
0047     /// The index updated on every iteration.
0048     Value *Index = nullptr;
0049     /// The header and latch of the loop.
0050     BasicBlock *Header = nullptr;
0051     BasicBlock *Latch = nullptr;
0052   };
0053
0054   /// The loop iterating on the rows.
0055   MatrixLoop RowLoop;
0056   /// The loop iterating on the columns.
0057   MatrixLoop ColumnLoop;
0058   /// The loop iterating on k (inner dimension).
0059   MatrixLoop KLoop;
0060
0061   TileInfo(unsigned NumRows, unsigned NumColumns, unsigned NumInner,
0062            unsigned TileSize)
0063       : NumRows(NumRows), NumColumns(NumColumns), NumInner(NumInner),
0064         TileSize(TileSize) {}
0065
0066   /// Creates an IR loop nests for tiling of the form below. Returns the block
0067   /// for the inner loop body and sets {Column,Row,Inner}LoopHeader/Latch
0068   /// fields.
0069   ///
0070   /// for ColumnLoop.Index = 0..NumColumns
0071   ///   for RowLoop.Index = 0..NumRows
0072   ///     for InnerLoop.Index = 0..NumInner
0073   BasicBlock *CreateTiledLoops(BasicBlock *Start, BasicBlock *End,
0074                                IRBuilderBase &B, DomTreeUpdater &DTU,
0075                                LoopInfo &LI);
0076
0077 private:
0078   /// Creates a new loop with header, body and latch blocks that iterates from
0079   /// [0, Bound). Updates \p Preheader to branch to the new header and uses \p
0080   /// Exit as exit block.  Adds the new loop blocks to \L and applies dominator
0081   /// tree updates to \p DTU.
0082   static BasicBlock *CreateLoop(BasicBlock *Preheader, BasicBlock *Exit,
0083                                 Value *Bound, Value *Step, StringRef Name,
0084                                 IRBuilderBase &B, DomTreeUpdater &DTU, Loop *L,
0085                                 LoopInfo &LI);
0086 };
0087 } // namespace llvm
0088
0089 #endif