1 //===- MatrixUtils.h - Utilities to lower matrix intrinsics -----*- C++ -*-===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 //
9 // Utilities for generating tiled loops for matrix operations.
10 //
11 //===----------------------------------------------------------------------===//
12 
13 #ifndef LLVM_TRANSFORMS_UTILS_MATRIXUTILS_H
14 #define LLVM_TRANSFORMS_UTILS_MATRIXUTILS_H
15 
16 #include "llvm/ADT/StringRef.h"
17 
18 namespace llvm {
19 class DomTreeUpdater;
20 class BasicBlock;
21 class Value;
22 class Loop;
23 class LoopInfo;
24 class IRBuilderBase;
25 
26 /// A helper struct to create IR loop nests for tiling in IR of the following
27 /// form:
28 ///   for CurrentColumn = 0..NumColumns
29 ///     for CurrentRow = 0..NumRows
30 ///       for CurrentInner = 0..NumInner
31 struct TileInfo {
32   /// Number of rows of the matrix.
33   unsigned NumRows;
34 
35   /// Number of columns of the matrix.
36   unsigned NumColumns;
37 
38   /// Number of columns of the first matrix of a multiply /
39   /// number of rows of the second matrix of a multiply.
40   unsigned NumInner;
41 
42   /// Number of rows/columns in a tile.
43   unsigned TileSize = -1;
44 
45   /// Start row of the current tile to compute.
46   Value *CurrentRow;
47 
48   /// Start column of the current tile to compute.
49   Value *CurrentCol;
50 
51   /// Current tile offset during the tile computation.
52   Value *CurrentK;
53 
54   /// Header of the outermost loop iterating from 0..NumColumns.
55   BasicBlock *ColumnLoopHeader = nullptr;
56 
57   /// Header of the second loop iterating from 0..NumRows.
58   BasicBlock *RowLoopHeader = nullptr;
59   /// Latch of the second loop iterating from 0..NumRows.
60   BasicBlock *RowLoopLatch = nullptr;
61   /// Header of the innermost loop iterating from 0..NumInner.
62   BasicBlock *InnerLoopHeader = nullptr;
63   /// Latch of the innermost loop iterating from 0..NumInner.
64   BasicBlock *InnerLoopLatch = nullptr;
65 
66   TileInfo(unsigned NumRows, unsigned NumColumns, unsigned NumInner,
67            unsigned TileSize)
68       : NumRows(NumRows), NumColumns(NumColumns), NumInner(NumInner),
69         TileSize(TileSize) {}
70 
71   /// Creates an IR loop nests for tiling of the form below. Returns the block
72   /// for the inner loop body and sets {Column,Row,Inner}LoopHeader/Latch
73   /// fields.
74   ///
75   /// for CurrentColumn = 0..NumColumns
76   ///   for CurrentRow = 0..NumRows
77   ///     for CurrentInner = 0..NumInner
78   BasicBlock *CreateTiledLoops(BasicBlock *Start, BasicBlock *End,
79                                IRBuilderBase &B, DomTreeUpdater &DTU,
80                                LoopInfo &LI);
81 
82 private:
83   /// Creates a new loop with header, body and latch blocks that iterates from
84   /// [0, Bound). Updates \p Preheader to branch to the new header and uses \p
85   /// Exit as exit block.  Adds the new loop blocks to \L and applies dominator
86   /// tree updates to \p DTU.
87   static BasicBlock *CreateLoop(BasicBlock *Preheader, BasicBlock *Exit,
88                                 Value *Bound, Value *Step, StringRef Name,
89                                 IRBuilderBase &B, DomTreeUpdater &DTU, Loop *L,
90                                 LoopInfo &LI);
91 };
92 } // namespace llvm
93 
94 #endif
95