part_descriptor.hpp

Go to the documentation of this file.

/*
 * SPDX-FileCopyrightText: Copyright (c) 2019-2025, NVIDIA CORPORATION.
 * SPDX-License-Identifier: Apache-2.0
 */
 
#pragma once
 
#include "data.hpp"
 
#include <stdint.h>
 
#include <ostream>
#include <set>
#include <vector>
 
namespace MLCommon {
namespace Matrix {
 
enum Layout {
  LayoutRowMajor = 0,
  LayoutColMajor
};
 
struct RankSizePair {
  RankSizePair() : rank(-1), size(0) {}
 
  RankSizePair(int _rank, size_t _size) : rank(_rank), size(_size) {}
 
  int rank;
 
  size_t size;
};
 
struct PartDescriptor {
  size_t M;
  size_t N;
 
  int rank;
 
  Layout layout;
  std::vector<RankSizePair*> partsToRanks;
 
  PartDescriptor(size_t _M,
                 size_t _N,
                 const std::vector<RankSizePair*>& _partsToRanks,
                 int rank,
                 Layout _layout = LayoutColMajor);
 
  int totalBlocks() const { return partsToRanks.size(); }
 
  int totalBlocksOwnedBy(int rank) const;
 
  std::set<int> uniqueRanks();
 
  std::vector<size_t> startIndices() const;
 
  std::vector<size_t> startIndices(int rank) const;
 
  std::vector<RankSizePair*> blocksOwnedBy(int rank) const;
 
  size_t totalElementsOwnedBy(int rank) const;
 
  friend std::ostream& operator<<(std::ostream& os, const PartDescriptor& desc);
  friend bool operator==(const PartDescriptor& a, const PartDescriptor& b);
};
 
std::ostream& operator<<(std::ostream& os, const PartDescriptor& desc);
 
bool operator==(const PartDescriptor& a, const PartDescriptor& b);
 
};  // end namespace Matrix
};  // end namespace MLCommon