distributed_matrix.h

Go to the documentation of this file.

#ifndef MADNESS_DISTRIBUTED_MATRIX_H
#define MADNESS_DISTRIBUTED_MATRIX_H

/*
  This file is part of MADNESS.

  Copyright (C) 2007,2010 Oak Ridge National Laboratory

  This program is free software; you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation; either version 2 of the License, or
  (at your option) any later version.

  This program is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
  GNU General Public License for more details.

  You should have received a copy of the GNU General Public License
  along with this program; if not, write to the Free Software
  Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA

  For more information please contact:

  Robert J. Harrison
  Oak Ridge National Laboratory
  One Bethel Valley Road
  P.O. Box 2008, MS-6367

  email: harrisonrj@ornl.gov
  tel:   865-241-3937
  fax:   865-572-0680

  $Id$
*/

// THE STUFF IN THIS FILE IS IN TRANSITION!  THE API AND
// IMPLEMENTATION WILL BOTH SHIFT RAPIDLY AS WE TRANSITION FROM
// REPLICATED TO DISTRIBUTED MATRIX ALGORITHMS, AND SUBSEQUENTLY
// REFINE THE DESIGN AND INTERFACE TO 3RD PARTY PACKAGES.

#include <madness/world/world.h>
#include <utility>
#include <madness/tensor/tensor.h>

namespace madness {

    // If in a fit of misplaced enthusiasm you desire to change
    // int64_t to either long or std::size_t you should be aware that
    // some uses below may contain quantities greater than may be
    // represented in a 32-bit integer and may also be negative.
    // I.e., a simple global replace will fail, though the existing
    // test suite may not detect that.  Also, large skinny matrices
    // could easily need more than 32 bit integers to address.


    // Forward declarations for friends
    class DistributedMatrixDistribution;
    template <typename T> class DistributedMatrix;
    
    static inline DistributedMatrixDistribution column_distributed_matrix_distribution(World& world, int64_t n, int64_t m, int64_t coltile=0);
    static inline DistributedMatrixDistribution row_distributed_matrix_distribution(World& world, int64_t n, int64_t m, int64_t rowtile=0);

    template <typename T>
    DistributedMatrix<T> concatenate_rows(const DistributedMatrix<T>& a, const DistributedMatrix<T>& b);

    template <typename T>
    DistributedMatrix<T> interleave_rows(const DistributedMatrix<T>& a, const DistributedMatrix<T>& b);

    class DistributedMatrixDistribution {
        friend DistributedMatrixDistribution column_distributed_matrix_distribution(World& world, int64_t n, int64_t m, int64_t coltile);
        friend DistributedMatrixDistribution row_distributed_matrix_distribution(World& world, int64_t n, int64_t m, int64_t rowtile);
        template <typename T> friend class DistributedMatrix;

    protected:
        World* pworld;
        int64_t P;                //< No. of processors
        ProcessID rank;           //< My processor rank
        int64_t n;                //< Column dimension of A(n,m)
        int64_t m;                //< Row dimension of A(n,m)
        int64_t tilen;            //< Tile size for column
        int64_t tilem;            //< Tile size for row
        int64_t Pcoldim;          //< Column dimension of processor grid
        int64_t Prowdim;          //< Row dimension of processor grid
        int64_t Pcol;             //< Column of processor grid for this processor
        int64_t Prow;             //< Row of processor grid for this processor
        int64_t ilo,ihi;          //< Range of column indices on this processor
        int64_t jlo,jhi;          //< Range of row indices on this processor
        int64_t idim,jdim;        //< Dimension of data on this processor



        DistributedMatrixDistribution(World& world, int64_t n, int64_t m, int64_t coltile, int64_t rowtile)
            : pworld(&world)
            , P(world.size())
            , rank(world.rank())
            , n(n)
            , m(m)
            , tilen(coltile)
            , tilem(rowtile)
            , Pcoldim((n-1)/tilen+1)
            , Prowdim((m-1)/tilem+1)
            , Pcol(rank/Prowdim)
            , Prow(rank - Pcol*Prowdim)
            , ilo(Pcol*tilen)
            , ihi(std::min(ilo+tilen-1,n-1))
            , jlo(Prow*tilem)
            , jhi(std::min(jlo+tilem-1,m-1))
            , idim(std::max(ihi-ilo+1,int64_t(0)))
            , jdim(std::max(jhi-jlo+1,int64_t(0)))
        {
            if (ilo > ihi || jlo > jhi) {
                ilo = jlo = 0;
                ihi = jhi = -1;
            }
        }


    public:

        DistributedMatrixDistribution() 
            : pworld(0)
            , P(0)
            , rank(0)
            , n(0)
            , m(0)
            , tilen(0)
            , tilem(0)
            , Pcoldim(0)
            , Prowdim(0)
            , Pcol(0)
            , Prow(0)
            , ilo(0)
            , ihi(-1)
            , jlo(0)
            , jhi(-1)
            , idim(0)
            , jdim(0)
        {}


        void clear() {
            pworld = (World*)(0);
            P = rank = n = m = tilen = tilem = Pcoldim = Prowdim = Pcol = Prow = ilo = ihi = jlo = jhi = idim = jdim = 0;
        }

        bool operator==(const DistributedMatrixDistribution& d) const {
            return 
                pworld  == d.pworld &&
                P       == d.P      &&
                rank    == d.rank   &&
                n       == d.n      &&
                m       == d.m      &&
                tilen   == d.tilen  &&
                tilem   == d.tilem  &&
                Pcoldim == d.Pcoldim &&
                Prowdim == d.Prowdim &&
                Pcol    == d.Pcol   &&
                Prow    == d.Prow   &&
                ilo     == d.ilo    &&
                ihi     == d.ihi    &&
                jlo     == d.jlo    &&
                jhi     == d.jhi    &&
                idim    == d.idim   &&
                jdim    == d.jdim;
        }



        int64_t coldim() const {
            return n;
        }



        int64_t rowdim() const {
            return m;
        }



        int64_t coltile() const {
            return tilen;
        }



        int64_t rowtile() const {
            return tilem;
        }



        int64_t process_coldim() const {return Pcoldim;}



        int64_t process_rowdim() const {return Prowdim;}



        int64_t local_size() const {return idim*jdim;}



        int64_t local_coldim() const {return idim;}



        int64_t local_rowdim() const {return jdim;}


        void local_colrange(int64_t& ilow, int64_t& ihigh) const {
            ilow = ilo;
            ihigh = ihi;
        }



        void local_rowrange(int64_t& jlow, int64_t& jhigh) const {
            jlow = jlo;
            jhigh = jhi;
        }


        int64_t local_ilow() const {
            return ilo;
        }

        int64_t local_ihigh() const {
            return ihi;
        }


        int64_t local_jlow() const  {
            return jlo;
        }


        int64_t local_jhigh() const  {
            return jhi;
        }



        void get_range(int p, int64_t& ilow, int64_t& ihigh, int64_t& jlow, int64_t& jhigh) const {
            int pi = p/Prowdim;
            int pj = p - pi*Prowdim;
            if (pi >= process_coldim() || pj >= process_rowdim()) {
                ilow = jlow = 0;
                ihigh = jhigh = -1;
            }
            else {
                ilow = pi*tilen;
                jlow = pj*tilem;
                ihigh= std::min(ilow+tilen-1,n-1);
                jhigh= std::min(jlow+tilem-1,m-1);
            }

            return;
        }



        void get_colrange(int p, int64_t& ilow, int64_t& ihigh) const {
            int64_t jlow, jhigh;
            get_range(p, ilow, ihigh, jlow, jhigh);

            return;
        }



        void get_rowrange(int p, int64_t& jlow, int64_t& jhigh) const {
            int64_t ilow, ihigh;
            get_range(p, ilow, ihigh, jlow, jhigh);

            return;
        }



        World& get_world() const {return *pworld;}



        bool is_column_distributed() const {return process_rowdim()==1;}



        bool is_row_distributed() const {return process_coldim()==1;}


        const DistributedMatrixDistribution& distribution() const {return *this;}


        ProcessID owner(int64_t i, int64_t j) const {
            int pcol = i/coltile();
            int prow = j/rowtile();
            
            return pcol*process_rowdim() + prow;
        }

        virtual ~DistributedMatrixDistribution() {}
    };



    template <typename T>
    class DistributedMatrix : public DistributedMatrixDistribution {
        friend DistributedMatrix<T> interleave_rows<T>(const DistributedMatrix<T>& a, const DistributedMatrix<T>& b);
        friend DistributedMatrix<T> concatenate_rows<T>(const DistributedMatrix<T>& a, const DistributedMatrix<T>& b);

        Tensor<T> t;            //< The data

        static T idij(const int64_t i, const int64_t j) {return (i==j) ?  T(1) : T(0);}

    protected:


        DistributedMatrix(World& world, int64_t n, int64_t m, int64_t coltile, int64_t rowtile)
            : DistributedMatrixDistribution(world, n, m, coltile, rowtile)
        {
            if (idim>0 && jdim>0) t = Tensor<T>(idim,jdim);
        }

    public:

        DistributedMatrix()
            : DistributedMatrixDistribution()
            , t()
        {}


        DistributedMatrix(const DistributedMatrixDistribution& d)
            : DistributedMatrixDistribution(d)
        {
            if (idim>0 && jdim>0) t = Tensor<T>(idim,jdim);
        }


        DistributedMatrix(const DistributedMatrix<T>& A, bool deepcopy=false)
            : DistributedMatrixDistribution(A)
            , t(deepcopy ? copy(A.t) : A.t) 
        {}


        DistributedMatrix<T>& operator=(const DistributedMatrix<T>& A) {
            if (this != &A) {
                DistributedMatrixDistribution::operator=(A);
                t = A.t;
            }
            return *this;
        }

        virtual ~DistributedMatrix() {}


        void clear() {
            DistributedMatrixDistribution::clear();
            t.clear();
        }



        template <typename funcT>
        void fill(const funcT& f) {
            for (int64_t i=ilo; i<=ihi; i++) {
                for (int64_t j=jlo; j<=jhi; j++) {
                    t(i-ilo,j-jlo) = f(i,j);
                }
            }
        }



        void fill(T value) {
            t.fill(value);
        }


        void fill_identity() {
            fill(DistributedMatrix<T>::idij);
        }



        Tensor<T>& data() {return t;}


        const Tensor<T>& data() const {return t;}



        void copy_from_replicated(const Tensor<T>& s) {
            if (local_size() > 0) t(___) = s(Slice(ilo,ihi),Slice(jlo,jhi));
        }


        void copy_to_replicated(Tensor<T>& s) const {
            MADNESS_ASSERT(s.iscontiguous());
            s = 0.0;
            if (local_size() > 0) s(Slice(ilo,ihi),Slice(jlo,jhi)) = t(___);
            get_world().gop.sum(s.ptr(), s.size());
        }


        void copy_from_replicated_patch(int64_t ilow, int64_t ihigh, int64_t jlow, int64_t jhigh, const Tensor<T>& s) {
            int64_t i0 = std::max(ilo,ilow);
            int64_t j0 = std::max(jlo,jlow);
            int64_t i1 = std::min(ihi,ihigh);
            int64_t j1 = std::min(jhi,jhigh);
            if (i0<=i1 && j0<=j1) {
                t(Slice(i0-ilo,i1-ilo),Slice(j0-jlo,j1-jlo)) = s(Slice(i0-ilow,i1-ilow),Slice(j0-jlow,j1-jlow));
            }
        }


        void copy_to_replicated_patch(int64_t ilow, int64_t ihigh, int64_t jlow, int64_t jhigh, Tensor<T>& s) const {
            MADNESS_ASSERT(s.iscontiguous());
            s = 0;
            int64_t i0 = std::max(ilo,ilow);
            int64_t j0 = std::max(jlo,jlow);
            int64_t i1 = std::min(ihi,ihigh);
            int64_t j1 = std::min(jhi,jhigh);
            if (i0<=i1 && j0<=j1) {
                t(Slice(i0-ilo,i1-ilo),Slice(j0-jlo,j1-jlo)) = s(Slice(i0-ilow,i1-ilow),Slice(j0-jlow,j1-jlow));
            }
            get_world().gop.sum(s.ptr(), s.size());
        }

        void extract_columns(int64_t jlow, int64_t jhigh, DistributedMatrix<T>& U) const {
            int newrowdim = jhigh - jlow + 1;
            MADNESS_ASSERT(jlow >= 0);
            MADNESS_ASSERT(jhigh < rowdim());
            MADNESS_ASSERT(newrowdim == U.rowdim());
            MADNESS_ASSERT(coldim() == U.coldim());
            MADNESS_ASSERT(is_column_distributed());
            MADNESS_ASSERT(U.is_column_distributed());
            MADNESS_ASSERT(coltile() == U.coltile());

            int64_t i0 = ilo;
            int64_t j0 = std::max(jlo,jlow);
            int64_t i1 = ihi;
            int64_t j1 = std::min(jhi,jhigh);
            if (i0<=i1 && j0<=j1) {
                U.data()(___) = t(Slice(i0-ilo,i1-ilo),Slice(j0-jlo,j1-jlo));
            }
        }

        template <typename R> 
        bool has_same_dimension_and_distribution(const DistributedMatrix<R>& A) {
            return DistributedMatrixDistribution::operator==(A);
        }


        DistributedMatrix<T>& operator+=(const DistributedMatrix<T>& A) {
            MADNESS_ASSERT(has_same_dimension_and_distribution(A));
            t += A.t;
            return *this;
        }



        DistributedMatrix<T> operator+(const DistributedMatrix<T>& A) const {
            MADNESS_ASSERT(has_same_dimension_and_distribution(A));
            return copy(*this)+=A;
        }



        DistributedMatrix<T>& operator*=(const T s) {
            t.scale(s);
            return *this;
        }

        void set(int64_t i, int64_t j, const T x) {
            MADNESS_ASSERT(i>=ilo && i<=ihi && j>=jlo && j<=jhi);
            t(i-ilo,j-jlo) = x;
        }

        T get(int64_t i, int64_t j) const {
            MADNESS_ASSERT(i>=ilo && i<=ihi && j>=jlo && j<=jhi);
            return t(i-ilo,j-jlo);
        }
    };



    template <typename T>
    DistributedMatrix<T> copy(const DistributedMatrix<T>& A) {
        return DistributedMatrix<T>(A,true);
    }
    


    static inline DistributedMatrixDistribution
    column_distributed_matrix_distribution(World& world, int64_t n, int64_t m, int64_t coltile) { // default coltile=0 above
        if (world.size()*coltile < n) coltile = (n-1)/world.size() + 1;
        coltile = std::min(coltile,n);
        if ((coltile&0x1)) ++coltile; // ??? Was before the previous statement

        return DistributedMatrixDistribution(world, n, m, coltile, m);
    }



    template <typename T>
    DistributedMatrix<T> column_distributed_matrix(World& world, int64_t n, int64_t m, int64_t coltile=0) {
        return DistributedMatrix<T>(column_distributed_matrix_distribution(world, n, m, coltile));
    }



    static inline DistributedMatrixDistribution
    row_distributed_matrix_distribution(World& world, int64_t n, int64_t m, int64_t rowtile) { // default rowtile=0 above
        if (world.size()*rowtile < m) rowtile = (m-1)/world.size() + 1;
        rowtile = std::min(rowtile,m);

        return DistributedMatrixDistribution(world, n, m, n, rowtile);
    }


    template <typename T>
    DistributedMatrix<T> row_distributed_matrix(World& world, int64_t n, int64_t m, int64_t rowtile=0) {
        return DistributedMatrix<T>(row_distributed_matrix_distribution(world, n, m, rowtile));
    }



    template <typename T>
    DistributedMatrix<T> interleave_rows(const DistributedMatrix<T>& a, const DistributedMatrix<T>& b) {
        MADNESS_ASSERT(a.rowdim()==b.rowdim() && a.coldim()==b.coldim() && a.coltile()==b.coltile() && a.rowtile()==b.rowtile());

        DistributedMatrix<T> c(a.get_world(), a.coldim()*2, a.rowdim(), a.coltile()*2, a.rowtile());
        c.data()(Slice(0,-1,2),_) = a.data()(___);
        c.data()(Slice(1,-1,2),_) = b.data()(___);
    }



    template <typename T>
    DistributedMatrix<T> concatenate_rows(const DistributedMatrix<T>& a, const DistributedMatrix<T>& b) {
        MADNESS_ASSERT(a.coldim()==b.coldim() && a.coltile()==b.coltile() && a.is_column_distributed() && b.is_column_distributed());

        int64_t ma = a.rowdim();
        int64_t mb = b.rowdim();

        DistributedMatrix<T> c(a.get_world(), a.coldim(), ma+mb, a.coltile(), ma+mb);
        
        int64_t ilow, ihigh;
        a.local_colrange(ilow, ihigh);
        if (ilow <= ihigh) {
            c.data()(_,Slice(0,ma-1)) = a.data()(___);
            c.data()(_,Slice(ma,-1))  = b.data()(___);
        }

        return c;
    }


    template <typename T>
    DistributedMatrix<T> concatenate_rows( const DistributedMatrix<T>& a, const DistributedMatrix<T>& b, const DistributedMatrix<T>& c, const DistributedMatrix<T>& d) {
        MADNESS_ASSERT(a.coldim()==b.coldim() && b.coldim()==c.coldim() && c.coldim()==d.coldim());
        MADNESS_ASSERT(a.coltile()==b.coltile() && b.coltile()==c.coltile() && c.coltile()==d.coltile());
        MADNESS_ASSERT(a.is_column_distributed() && b.is_column_distributed() && c.is_column_distributed() && d.is_column_distributed());
        
        int64_t ma = a.rowdim();
        int64_t mb = b.rowdim();
        int64_t mc = c.rowdim();
        int64_t md = d.rowdim();
        
        DistributedMatrix<T> result(a.get_world(), a.coldim(), ma+mb+mc+md, a.coltile(), ma+mb+mc+md);

        if(a.local_size() > 0) result.data()( _ , Slice(0,ma-1) ) = a.data()(___);
        if(b.local_size() > 0) result.data()( _ , Slice(ma, ma+mb-1) ) = b.data()(___);
        if(c.local_size() > 0) result.data()( _ , Slice(ma+mb, ma+mb+mc-1) ) = c.data()(___);
        if(d.local_size() > 0) result.data()( _ , Slice(ma+mb+mc, -1) ) = d.data()(___);

        return result;
    }



    template <typename T>
    DistributedMatrix<T> concatenate_columns(const DistributedMatrix<T>& a, const DistributedMatrix<T>& b) {
        MADNESS_ASSERT(a.rowdim()==b.rowdim() && a.rowtile()==b.rowtile() && a.is_row_distributed() && b.is_row_distributed());

        int64_t ma = a.coldim();
        int64_t mt = ma + b.coldim();

        DistributedMatrix<T> c(a.get_world(), mt, a.rowdim(), b.rowtile(), mt);

        if(a.local_size() > 0) c.data()( Slice(0,ma-1), _ ) = a.data()(___);
        if(a.local_size() > 0) c.data()( Slice(ma,-1), _ ) = b.data()(___);

        return c;
    }
}

#endif