// Copyright (c) 2020 Graphcore Ltd. All rights reserved. #ifndef popsparse_HyperGraphBlock_hpp #define popsparse_HyperGraphBlock_hpp #include "HyperGraph.hpp" #include "ZoltanPartitioner.hpp" #include #include // empirical constant const int TARGET_V_NODES_PER_TILE = 2; namespace popsparse { namespace experimental { /* This class is a base class for block-based matrix multiplication partitioning algoriothms. Derived classes uses different partitioning algorithms. */ class HyperGraphBlock : public HyperGraph { public: HyperGraphBlock(BlockMatrix &A, BlockMatrix &B, poplar::Type inDataTypeIn, poplar::Type outDataTypeIn, poplar::Type partialDataTypeIn, int nTileIn, int nTargetNodesVPerTileIn = TARGET_V_NODES_PER_TILE); virtual ~HyperGraphBlock() = default; protected: /* DataNode class describes an node in a Hypergraph that references input or output values. In our case the reference is blockId field. One DataNode corresponds to one input/output block tensor. */ class DataNode { public: unsigned int id; unsigned int blockId; unsigned int blockRow; unsigned int blockCol; float w; DataNode(unsigned int idIn, unsigned int blockIdIn, unsigned int blockRowIn, unsigned int blockColIn) : id(idIn), blockId(blockIdIn), blockRow(blockRowIn), blockCol(blockColIn), w(1.0f) {} }; /* ComputeNode class describes an node in a Hypergraph where ops happen. In our case it is muliply and add operation that computes partial result of matrix A block-row by matrix B block-column multiplicatoin. Fields idxA, inxB refer to block Ids of matrices A and B accordingly. Blocks of idxA are pairwise multiplied with blocks of idxB. Multiplication results are summed as partials to be further reduced. One ComputeNode corresponds to one 1x1 convolution vertex. */ class ComputeNode { public: unsigned int id; std::vector idxA; // block Ids in matrix A std::vector idxB; // block Ids in matrix B unsigned int blockRow; unsigned int blockCol; unsigned int partition; float w; ComputeNode(unsigned int idIn, std::vector &idxAIn, std::vector &idxBIn, unsigned int blockRowIn, unsigned int blockColIn, unsigned int partitionIn) : id(idIn), idxA(std::move(idxAIn)), idxB(std::move(idxBIn)), blockRow(blockRowIn), blockCol(blockColIn), partition(partitionIn), w(1.0f) {} }; /* HyperEdge class describes directional connections to or from a node. Each node can have multiple input and output connections. In our design we take several connections flowing to/from a compute node from/to data nodes of the same type (A, B or C) and put them into a single Edge class. */ class HyperEdge { public: unsigned int id; float w; std::vector in; // Input (source) node Id std::vector out; // Output (target) node Id HyperEdge(unsigned int idIn) : id(idIn), w(1.0f) {} }; // These 4 nodes describe a partial expression: // A * B = C // Partial of matrix A // Corresponds to a block tensor std::vector nodeA; // Partial of matrix B // Corresponds to a block tensor std::vector nodeB; // Partial of matrix C // Corresponds to a block tensor std::vector nodeC; // Partial of A * B // Corresponds to a vertex std::vector nodeV; // Node V inverted index std::unordered_map nodeVIdxById; // Node V layout std::vector>> nodeVLayout; // HyperEdges connected to nodes A std::vector edgeA; // HyperEdges connected to nodes B std::vector edgeB; // HyperEdges connected to nodes C std::vector edgeC; // Edgs A id layout by block row, block column std::vector> hyperEdgeIdA; // Edgs B id layout by block row, block column std::vector> hyperEdgeIdB; // Edgs C id layout by block row, block column std::vector> hyperEdgeIdC; // block id look up matrix for A std::vector> blockIdMatrixA; // block id look up matrix for B std::vector> blockIdMatrixB; // Global node Id counter unsigned int gNodeId; // Global edge Id counter unsigned int gEdgeId; // Number of block matrix multiplications std::size_t numMuls; // Optimization parameter // This is the desired number of V nodes per tile int nTargetNodesVPerTile; // Tile assignment map for all nodes sorted by their id std::vector tileAssignment; // Tile assignment map for nodes C sorted by node index std::vector nodeCTileId; public: const std::vector &getNodeA() const { return nodeA; } const std::vector &getNodeB() const { return nodeB; } const std::vector &getNodeC() const { return nodeC; } const std::vector &getNodeV() const { return nodeV; } const std::vector &getEdgeA() const { return edgeA; } const std::vector &getEdgeB() const { return edgeB; } const std::vector &getEdgeC() const { return edgeC; } // Creates a graph for (sparse) matrix multiplication void createGraphMatMul(poplar::Graph &graph, const poplar::DebugNameAndId &dnai) override; // Creates a graph for matrix multiplication and sparsifies the result void createGraphMatMulSparsifyResult(poplar::Graph &graph, const unsigned char *sparsity, const poplar::DebugNameAndId &dnai) override; // Creates a program to perform matmul void createProgramMatMul(poplar::Graph &graph, SubBlockMask subBlockMask, poplar::program::Sequence &prog, const poplar::DebugNameAndId &dnai) override; // Adds a logic to perform matmul to existing compusets void createProgramMatMul(poplar::Graph &graph, poplar::ComputeSet *transposeCS, poplar::ComputeSet &mulCS, poplar::ComputeSet &reduceCS, poplar::program::Sequence &prog, const poplar::DebugNameAndId &dnai) override; // Creates a compute set to perform a partial matmul void createComputeSetMatMul( poplar::Graph &graph, std::map &partialData, poplar::program::Sequence &prog, const poplar::DebugNameAndId &dnai); // Creates a compute set to perform a partial matmul void createComputeSetMatMul( poplar::Graph &graph, std::map &partialData, poplar::ComputeSet &mulCS, poplar::ComputeSet *transposeCS, poplar::program::Sequence &prog, const poplar::DebugNameAndId &dnai); // Creates a compute set to perform reduce void createComputeSetReduce( poplar::Graph &graph, const std::map &partialDataIn, poplar::program::Sequence &prog, const poplar::DebugNameAndId &dnai); // Creates a compute set to perform reduce void createComputeSetReduce( poplar::Graph &graph, const std::map &partialDataIn, poplar::ComputeSet &reduceCS, const poplar::DebugNameAndId &dnai); // Set the tile mapping for left hand matrix void setTileMappingLHS(poplar::Graph &graph, poplar::Tensor &lhsTensor) override; // Set the tile mapping for right hand matrix void setTileMappingRHS(poplar::Graph &graph, poplar::Tensor &rhsTensor) override; unsigned int getTotalNodes() const { return gNodeId; } void setTileAssignment(std::vector &tileAssignmentIn) { tileAssignment = tileAssignmentIn; } // Adds necessary codelets // This method must be called once per graph static void addCodelets(poplar::Graph &graph); protected: // Performs graph patitioning virtual void partitionGraph() = 0; // Creates a graph for dense * sparse = dense matrix multiplication void createGraphMatMulDSD(poplar::Graph &graph, const poplar::DebugNameAndId &dnai); // Creates a graph for dense * dense -> sparse matrix multiplication void createGraphMatMulDDSSparsiryResult(poplar::Graph &graph, const unsigned char *sparsity, const poplar::DebugNameAndId &dnai); // Helper methods // Populates nodes and edges for a node group A, B or C std::vector> populateDataNodes( int nRow, int nCol, const std::vector> &blockIdMatrix, std::vector &nodeGroup, std::vector &edgeGroup); // Populates data nodes and edges for non zero blocks of matrix A void populateNodesA(); // Populates data nodes and edges for non zero blocks of matrix B void populateNodesB(); // Populates data nodes and edges for non zero blocks of matrix C void populateNodesC(int nRowC, int nColC, const std::vector> &blockIdMatrixC); // Populates compute nodes V virtual void populateNodesV(int nRowC, int nColC, const std::vector> &blockIdMatrixC); // Adds node V and all corresponding connections void populateNodeV( int row, int col, unsigned int p, const std::vector> &aList, const std::vector> &bList); // Helper fuinction. Gets number of muls per V node // Used in some partitioning algorithms unsigned int getMulsPerVNode() const; virtual void setupWeights(const poplar::Graph &graph) {} // Map C nodes to tiles after partitioning virtual void mapCNodes(poplar::Graph &graph); // Logs tile assignment virtual void logTileAssignment(const std::vector &tileAssignment); struct MemoryStatistics { std::size_t maxBytesOnTile = 0; std::size_t minBytesOnTile = 0; std::size_t idxMinOccupiedTile = 0; std::size_t maxBytesOnTileMatmul = 0; std::size_t idxMaxOccupiedTileMatmul = 0; std::size_t maxBytesOnTileReduce = 0; std::size_t idxMaxOccupiedTileReduce = 0; std::vector bytesAByTiles; std::vector bytesBByTiles; std::vector bytesVaByTiles; std::vector bytesVbByTiles; std::vector bytesVpByTiles; std::vector bytesCByTiles; std::vector bytesCpByTiles; // For matmul we count A,B,C,P,a,b std::vector bytesByTilesMatmul; // For reduce we count A,B,C,P,p,B std::vector bytesByTilesReduce; }; // Computes estimated data breakdowh per each tile void computeBytesByTile(MemoryStatistics &stats); }; static const float KBf = 1024.0f; } // namespace experimental } // namespace popsparse #endif