DistributedPartition.cc

// -*- tab-width: 2; indent-tabs-mode: nil; coding: utf-8-with-signature -*-
//-----------------------------------------------------------------------------
// Copyright 2000-2026 CEA (www.cea.fr) IFPEN (www.ifpenergiesnouvelles.com)
// See the top-level COPYRIGHT file for details.
// SPDX-License-Identifier: Apache-2.0
//-----------------------------------------------------------------------------
/*---------------------------------------------------------------------------*/
/*---------------------------------------------------------------------------*/
/*
 * This file is based on the work on AMGCL library (version march 2026)
 * which can be found at https://github.com/ddemidov/amgcl.
 *
 * Copyright (c) 2012-2022 Denis Demidov <dennis.demidov@gmail.com>
 * SPDX-License-Identifier: MIT
 */
/*---------------------------------------------------------------------------*/
/*---------------------------------------------------------------------------*/
 
#include <iostream>
#include <vector>
#include <string>
#include <algorithm>
#include <numeric>
#include <cassert>
 
#include "arccore/alina/AlinaUtils.h"
#include "arccore/alina/IO.h"
#include "arccore/common/internal/ProgramOptions.h"
 
extern "C" {
#include <metis.h>
}
 
using namespace Arcane;
using Alina::precondition;
 
//---------------------------------------------------------------------------
void pointwise_graph(int n, int block_size,
                     const std::vector<int>& ptr,
                     const std::vector<int>& col,
                     std::vector<idx_t>& pptr,
                     std::vector<idx_t>& pcol)
{
  int np = n / block_size;
 
  assert(np * block_size == n);
 
  // Create pointwise matrix
  std::vector<int> ptr1(np + 1, 0);
  std::vector<int> marker(np, -1);
  for (int ip = 0, i = 0; ip < np; ++ip) {
    for (int k = 0; k < block_size; ++k, ++i) {
      for (int j = ptr[i]; j < ptr[i + 1]; ++j) {
        int cp = col[j] / block_size;
        if (marker[cp] != ip) {
          marker[cp] = ip;
          ++ptr1[ip + 1];
        }
      }
    }
  }
 
  std::partial_sum(ptr1.begin(), ptr1.end(), ptr1.begin());
  std::fill(marker.begin(), marker.end(), -1);
 
  std::vector<int> col1(ptr1.back());
 
  for (int ip = 0, i = 0; ip < np; ++ip) {
    int row_beg = ptr1[ip];
    int row_end = row_beg;
 
    for (int k = 0; k < block_size; ++k, ++i) {
      for (int j = ptr[i]; j < ptr[i + 1]; ++j) {
        int cp = col[j] / block_size;
 
        if (marker[cp] < row_beg) {
          marker[cp] = row_end;
          col1[row_end++] = cp;
        }
      }
    }
  }
 
  // Transpose pointwise matrix
  int nnz = ptr1.back();
 
  std::vector<int> ptr2(np + 1, 0);
  std::vector<int> col2(nnz);
 
  for (int i = 0; i < nnz; ++i)
    ++(ptr2[col1[i] + 1]);
 
  std::partial_sum(ptr2.begin(), ptr2.end(), ptr2.begin());
 
  for (int i = 0; i < np; ++i)
    for (int j = ptr1[i]; j < ptr1[i + 1]; ++j)
      col2[ptr2[col1[j]]++] = i;
 
  std::rotate(ptr2.begin(), ptr2.end() - 1, ptr2.end());
  ptr2.front() = 0;
 
  // Merge both matrices.
  std::fill(marker.begin(), marker.end(), -1);
  pptr.resize(np + 1, 0);
 
  for (int i = 0; i < np; ++i) {
    for (int j = ptr1[i]; j < ptr1[i + 1]; ++j) {
      int c = col1[j];
      if (marker[c] != i) {
        marker[c] = i;
        ++pptr[i + 1];
      }
    }
 
    for (int j = ptr2[i]; j < ptr2[i + 1]; ++j) {
      int c = col2[j];
      if (marker[c] != i) {
        marker[c] = i;
        ++pptr[i + 1];
      }
    }
  }
 
  std::partial_sum(pptr.begin(), pptr.end(), pptr.begin());
  std::fill(marker.begin(), marker.end(), -1);
 
  pcol.resize(pptr.back());
 
  for (int i = 0; i < np; ++i) {
    int row_beg = pptr[i];
    int row_end = row_beg;
 
    for (int j = ptr1[i]; j < ptr1[i + 1]; ++j) {
      int c = col1[j];
 
      if (marker[c] < row_beg) {
        marker[c] = row_end;
        pcol[row_end++] = c;
      }
    }
 
    for (int j = ptr2[i]; j < ptr2[i + 1]; ++j) {
      int c = col2[j];
 
      if (marker[c] < row_beg) {
        marker[c] = row_end;
        pcol[row_end++] = c;
      }
    }
  }
}
 
//---------------------------------------------------------------------------
std::vector<idx_t>
pointwise_partition(idx_t npart,
                    const std::vector<idx_t>& ptr,
                    const std::vector<idx_t>& col)
{
  idx_t nrows = ptr.size() - 1;
 
  std::vector<idx_t> part(nrows);
 
  if (npart == 1) {
    std::fill(part.begin(), part.end(), 0);
  }
  else {
    idx_t edgecut;
 
#if defined(METIS_VER_MAJOR) && (METIS_VER_MAJOR >= 5)
    idx_t nconstraints = 1;
    METIS_PartGraphKway(
    &nrows, //nvtxs
    &nconstraints, //ncon -- new
    const_cast<idx_t*>(ptr.data()), //xadj
    const_cast<idx_t*>(col.data()), //adjncy
    NULL, //vwgt
    NULL, //vsize -- new
    NULL, //adjwgt
    &npart,
    NULL, //real t *tpwgts,
    NULL, // real t ubvec
    NULL,
    &edgecut,
    part.data());
#else
    int wgtflag = 0;
    int numflag = 0;
    int options = 0;
 
    METIS_PartGraphKway(
    &nrows,
    const_cast<int*>(ptr.data()),
    const_cast<int*>(col.data()),
    NULL,
    NULL,
    &wgtflag,
    &numflag,
    &npart,
    &options,
    &edgecut,
    part.data());
#endif
  }
 
  return part;
}
 
//---------------------------------------------------------------------------
std::vector<int>
partition(int n, int nparts, int block_size,
          const std::vector<int>& ptr, const std::vector<int>& col)
{
  // Pointwise graph
  std::vector<idx_t> pptr;
  std::vector<idx_t> pcol;
  pointwise_graph(n, block_size, ptr, col, pptr, pcol);
 
  // Pointwise partition
  std::vector<idx_t> ppart = pointwise_partition(nparts, pptr, pcol);
 
  std::vector<int> part(n);
  for (int i = 0; i < n; ++i)
    part[i] = ppart[i / block_size];
 
  return part;
}
 
//---------------------------------------------------------------------------
int main(int argc, char* argv[])
{
  namespace po = Arcane::ProgramOptions;
 
  try {
    std::string ifile;
    std::string ofile = "partition.mtx";
 
    int nparts, block_size;
 
    po::options_description desc("Options");
 
    desc.add_options()("help,h", "show help");
    desc.add_options()("input,i", po::value<std::string>(&ifile)->required(), "Input matrix");
    desc.add_options()("output,o", po::value<std::string>(&ofile)->default_value(ofile), "Output file");
    desc.add_options()("binary,B",
                       po::bool_switch()->default_value(false),
                       "When specified, treat input files as binary instead of as MatrixMarket. ");
    desc.add_options()("nparts,n", po::value<int>(&nparts)->required(), "Number of parts");
    desc.add_options()("block_size,b", po::value<int>(&block_size)->default_value(1), "Block size");
 
    po::positional_options_description pd;
    pd.add("input", 1);
 
    po::variables_map vm;
    po::store(po::command_line_parser(argc, argv).options(desc).positional(pd).run(), vm);
 
    if (vm.count("help")) {
      std::cout << desc << std::endl;
      return 0;
    }
 
    po::notify(vm);
 
    size_t rows;
    std::vector<int> ptr, col;
 
    bool binary = vm["binary"].as<bool>();
 
    if (binary) {
      std::ifstream f(ifile, std::ios::binary);
      precondition(f.read((char*)&rows, sizeof(rows)), "Wrong file format?");
      ptr.resize(rows + 1);
      for (size_t i = 0; i <= rows; ++i) {
        ptrdiff_t p;
        precondition(f.read((char*)&p, sizeof(p)), "Wrong file format?");
        ptr[i] = p;
      }
      col.resize(ptr.back());
      for (ptrdiff_t i = 0; i < ptr.back(); ++i) {
        ptrdiff_t p;
        precondition(f.read((char*)&p, sizeof(p)), "Wrong file format?");
        col[i] = p;
      }
    }
    else {
      std::vector<double> val;
      size_t cols;
      std::tie(rows, cols) = Alina::IO::mm_reader(ifile)(ptr, col, val);
      precondition(rows == cols, "Non-square system matrix");
    }
 
    std::vector<int> part = partition(rows, nparts, block_size, ptr, col);
 
    if (binary) {
      std::ofstream p(ofile.c_str(), std::ios::binary);
 
      Alina::IO::write(p, rows);
      Alina::IO::write(p, size_t(1));
      Alina::IO::write(p, part);
    }
    else {
      Alina::IO::mm_write(ofile, &part[0], part.size());
    }
  }
  catch (const std::exception& e) {
    std::cerr << "Error: " << e.what() << std::endl;
    return 1;
  }
}