Home | History | Annotate | Line # | Download | only in Utils
      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