include_starpu/skepu/matrix.h

#ifndef MATRIX_H
#define MATRIX_H

#include <iostream>
#include <fstream>
#include <sstream>
#include <cstdlib>

#include <vector>

#include <map>



#include <starpu.h>

#include "skepu/src/environment.h"

#include "skepu/src/malloc_allocator.h"





namespace skepu
{
  
 enum AccessType
 {
  ROW_WISE, //C style iterating from rows
  COL_WISE // fortran style iterating from columns 
 };

 
 
template<typename T>
class Matrix
{

private:

    void unpartitionMatrix() const
    {
        DEBUG_TEXT_LEVEL1("*****  MATRIX UNPARTIOTIINING xparts: "<< xparts <<",  yparts: "<< yparts <<" *****\n")
        if(xparts>1 && yparts >1)
          starpu_data_unpartition(matrix_handle, 0); // need to check as it might only unpartition only one filter.
        else if(xparts>1 || yparts >1)
          starpu_data_unpartition(matrix_handle, 0); // no matter x or y filter, unpartitioning logic is same.
        xparts=1;
        yparts=1;     
    }
    
    void release_acquire()
    {
        if(isAcquire)
        {
            if(xparts>1 || yparts>1) // partitioning
            {
                DEBUG_TEXT_LEVEL1("*****  MATRIX RELEASE **** xparts: "<< xparts << ",  yparts: "<< yparts <<"\n")
                
                if(xparts > 1 && yparts > 1)  // 2d partitioning
                {
                    for (int y = 0; y < yparts; y++) 
                        for (int x = 0; x < xparts; x++) 
                        {
                            starpu_data_release(starpu_data_get_sub_data(matrix_handle, 2, x, y));
                        }
                }
                else // 1D partitioning
                {
                    int parts = (xparts>1) ? xparts : yparts;
                    // parts should be equal to what parts variable points to
                    assert (parts == starpu_data_get_nb_children(matrix_handle));
                    for (int x = 0; x < parts; x++) 
                    {
                        starpu_data_release(starpu_data_get_sub_data(matrix_handle, 1, x));
                    }
                }
            }
            else // no partitioning
            {
                DEBUG_TEXT_LEVEL1("*****  MATRIX RELEASE ****\n")
                starpu_data_release(matrix_handle);
            }
            isAcquire = false;
        }
    }

  
  public:
  

    starpu_data_handle_t matrix_handle;
    mutable bool isOnStarPU;
    mutable bool isReadBack;
    mutable bool isWriteBack;
    mutable int xparts; // horizontal partitioning
    mutable int yparts; // vertical partitioning
    
    mutable bool isAcquire;
    mutable enum starpu_access_mode mode;
    
    struct starpu_data_filter matrix_filter_x;
    struct starpu_data_filter matrix_filter_y;
    //enum enum starpu_access_mode accessMode;
    
    starpu_data_handle_t registerMatrix()
    {
        release_acquire();
        
        if(!isOnStarPU)
        {
            starpu_matrix_data_register(&matrix_handle, 0, (uintptr_t)&m_data[0], m_cols, m_cols, m_rows, sizeof(m_data[0]));
            isOnStarPU=true;
            DEBUG_TEXT_LEVEL1("*****  MATRIX REGISTERING **** rows: "<< m_rows <<" cols: "<< m_cols <<" \n")
        }
        else if(xparts>1 || yparts >1)
        {
            unpartitionMatrix();
        }
        return matrix_handle;
    }
    
    void unregisterMatrix(bool update=true)
    {
        if(isOnStarPU)
        {
            release_acquire(); // Should we release before unregister? Yes!
                
            if(xparts>1 || yparts >1) // do we need to unpartition before unregistering? Yes!
            {
                unpartitionMatrix();
            }
            
            if(update)  
                starpu_data_unregister(matrix_handle);
                
            else // dont need to get updated data back
                starpu_data_unregister_no_coherency(matrix_handle);
                
            isOnStarPU = false;
            isAcquire = false;      
        }
    }
    
    starpu_data_handle_t registerPartitions(int _xparts=1, int _yparts=1)
    {
        release_acquire();
        
        if(!isOnStarPU) // if not registered, register it first
        {
            starpu_matrix_data_register(&matrix_handle, 0, (uintptr_t)&m_data[0], m_cols, m_cols, m_rows, sizeof(m_data[0]));
            isOnStarPU=true;
            DEBUG_TEXT_LEVEL1("*****   MATRIX REGISTERING **** rows: "<< m_rows <<" cols: "<< m_cols <<" \n")
        }
        
        if(_xparts == xparts && _yparts==yparts)
            return matrix_handle;   
        if(xparts>1 || yparts>1)
        {
            unpartitionMatrix();
        }
        if(_xparts>1 && _yparts>1)
        {
            matrix_filter_x.nchildren = _xparts;
            matrix_filter_y.nchildren = _yparts;    
            starpu_data_map_filters(matrix_handle, 2, &matrix_filter_x, &matrix_filter_y);
            DEBUG_TEXT_LEVEL1("*****  MATRIX PARTIOTIINING xparts: "<< _xparts <<",  yparts: "<< _yparts <<" *****\n")
            xparts = _xparts;
            yparts = _yparts;
        }
        else if(_xparts>1)
        {
            matrix_filter_x.nchildren = _xparts;
            starpu_data_partition(matrix_handle, &matrix_filter_x);
            DEBUG_TEXT_LEVEL1("*****  MATRIX PARTIOTIINING xparts: "<< _xparts <<" *****\n")
            xparts = _xparts;
        }
        else if(_yparts>1)
        {
            matrix_filter_y.nchildren = _yparts;
            starpu_data_partition(matrix_handle, &matrix_filter_y);
            DEBUG_TEXT_LEVEL1("*****  MATRIX PARTIOTIINING yparts: "<< _yparts <<" *****\n")
            yparts = _yparts;
        }
        
        return matrix_handle;
    }
    
    
   
   // typedefs
#ifdef USE_PINNED_MEMORY
    typedef std::vector<T, malloc_allocator<T> > container_type;
    typedef typename std::vector<T, malloc_allocator<T> >::size_type size_type;
    typedef typename std::vector<T, malloc_allocator<T> >::value_type value_type;
    typedef typename std::vector<T, malloc_allocator<T> >::difference_type difference_type;
    typedef typename std::vector<T, malloc_allocator<T> >::pointer pointer;
    typedef typename std::vector<T, malloc_allocator<T> >::reference reference;
    typedef typename std::vector<T, malloc_allocator<T> >::const_reference const_reference;
    typedef typename std::vector<T, malloc_allocator<T> >::const_iterator const_iterator;
    typedef typename std::vector<T, malloc_allocator<T> >::const_reverse_iterator const_reverse_iterator;
#else
    typedef std::vector<T> container_type;
    typedef typename std::vector<T>::size_type size_type;
    typedef typename std::vector<T>::value_type value_type;
    typedef typename std::vector<T>::difference_type difference_type;
    typedef typename std::vector<T>::pointer pointer;
    typedef typename std::vector<T>::reference reference;
    typedef typename std::vector<T>::const_reference const_reference;
    typedef typename std::vector<T>::const_iterator const_iterator;
    typedef typename std::vector<T>::const_reverse_iterator const_reverse_iterator;
#endif
  
  public: //-- For Testing --//
    
     T* GetArrayRep()
     {
       return &m_data[0];
     }
   
    friend std::ostream& operator<<(std::ostream &os, Matrix<T>& matrix)
    {
        matrix.acquireRead();
        
        os << "Matrix: ("<< matrix.total_rows() <<" X "<<matrix.total_cols()<<")\n";
        for(int i=0;i<matrix.size();i++)
        {
          os<<(matrix(i))<<" ";
          if((i+1)%(matrix.total_cols())==0)
             os << "\n";
        }
        os<<"\n";
        return os;
    }   

    
    void randomize(int min = 0, int max = RAND_MAX)
    {
        unregisterMatrix(false);

        for(unsigned int i = 0; i < size(); i++)
        {
            m_data.at(i) = (T)( rand() % max + min);
        }
    }

    void save(const std::string& filename)
    {
        std::ofstream file(filename.c_str());

        if (file.is_open())
        {
            for(unsigned int i = 0; i < m_data.size(); ++i)
            {
                file<<m_data.at(i) <<" ";
            }
            file.close();
        }
        else
        {
            std::cout<<"Unable to open file\n";
        }
    }

    void load(const std::string& filename, int rowWidth, int numRows = 0)
    {
        std::ifstream file(filename.c_str());

        if (file.is_open())
        {
            std::string line;
            getline (file,line);
            std::istringstream ss(line);
            T num;
            clear();

            //Load all elements
            if(numRows == 0)
            {
                while(ss >> num)
                {
                    push_back(num);
                }
            }
            // Load only numElements elements
            else
            {
                for(int i = 0; i < (numRows*rowWidth); ++i)
                {
                    ss >> num;
                    push_back(num);
                }
            }
        
        m_cols = rowWidth;
        m_rows = (size()/rowWidth);

            file.close();
        }
        else
        {
            std::cout<<"Unable to open file\n";
        }
    }
    
    

  
// Constructors, destructors
public:
      
    ~Matrix()
    {
        DEBUG_TEXT_LEVEL1("*****  MATRIX DESTRUCTOR **** rows: "<< total_rows()<<",  cols: "<< total_cols() <<"\n")
        if(isOnStarPU)
        {
            release_acquire();
            
            if(xparts>1 || yparts>1)
                unpartitionMatrix();
            
//          starpu_data_unregister(matrix_handle);
        }  
     }


 Matrix(size_type _rows, size_type _cols): m_rows(_rows), m_cols(_cols), m_data(m_rows * m_cols), transpose_matrix(0)
    {
        isOnStarPU = false;
        isAcquire = false;
        xparts=1;
        yparts=1;
        memset(&matrix_filter_x, 0, sizeof(matrix_filter_x));
        matrix_filter_x.filter_func=starpu_vertical_block_filter_func;
        
        memset(&matrix_filter_y, 0, sizeof(matrix_filter_y));
        matrix_filter_y.filter_func=starpu_block_filter_func;
    }
    
 Matrix(size_type _rows, size_type _cols, const T& val): m_rows(_rows), m_cols(_cols),m_data(m_rows * m_cols, val), transpose_matrix(0)
    {
        isOnStarPU = false;
        isAcquire = false;
        xparts=1;
        yparts=1;
        memset(&matrix_filter_x, 0, sizeof(matrix_filter_x));
        matrix_filter_x.filter_func=starpu_vertical_block_filter_func;
        
        memset(&matrix_filter_y, 0, sizeof(matrix_filter_y));
        matrix_filter_y.filter_func=starpu_block_filter_func;
    }
    
    
    Matrix(const Matrix<T>& copy)
    {
       copy.acquireRead();
      
       this->m_rows = copy.total_rows();
       this->m_cols = copy.total_cols();
       this->m_data= copy.m_data;
       this->transpose_matrix = copy.transpose_matrix;
       
       isOnStarPU = copy.isOnStarPU;
       isAcquire = copy.isAcquire;
       
       matrix_filter_x = copy.matrix_filter_x;
       matrix_filter_y = copy.matrix_filter_y;
       
       xparts=copy.xparts;
       yparts=copy.yparts;
    }

 public:

    size_type size() const { return m_data.size();}
    
    size_type total_rows() const {return m_rows; }
    
    size_type total_cols() const {return m_cols; }

    void change_layout()
    {  
      int tmp = m_rows;
      m_rows=m_cols;
      m_cols = tmp;
    }
    


 private:
    size_type m_rows, m_cols;
    
// A custom memory allocator is written to support pinned memory allocation for CUDA. Enabled by defining USE_PINNED_MEMORY macro
#ifdef USE_PINNED_MEMORY
    mutable std::vector<T, malloc_allocator<T> > m_data;
#else
    mutable std::vector<T> m_data;
#endif

    
    // for col_iterator, not tested 
    mutable Matrix<T> *transpose_matrix;

    template<typename Type>
    void item_swap(Type &t1, Type &t2);
    
    
    // External classes
public:
        //template <AccessType accessType=ROW_WISE>
    class iterator;
    
    class col_iterator;
    
    class proxy_elem;
    
    
    

public: //-- Operators --//

    void resize(size_type _rows, size_type _cols, T val = T());

    Matrix<T>& operator=(Matrix<T>& other);
    Matrix<T>& operator=(const T& elem);

    bool operator==(const Matrix<T>& c1);
    bool operator!=(const Matrix<T>& c1);
    bool operator<(const Matrix<T>& c1);
    bool operator>(const Matrix<T>& c1);
    bool operator<=(const Matrix<T>& c1);
    bool operator>=(const Matrix<T>& c1);

    Matrix<T>& subsection(size_type row, size_type col, size_type rowWidth, size_type colWidth);
    

public: //-- STL vector regular interface --//
    //Iterators
    iterator begin();
    const_iterator begin() const;
    iterator begin(unsigned row);
    const_iterator begin(unsigned row) const;

    iterator end();
    const_iterator end() const;
    iterator end(unsigned row);
    const_iterator end(unsigned row) const;
    
    //Column Iterators, taking transpose and using its iterator, work only for const_iterators
    // iterator col_begin();
    std::pair<const_iterator, const_iterator> col_iterator_range();
    col_iterator col_begin();

   // iterator col_end();
    col_iterator col_end();


    //Capacity
    size_type capacity() const;

    

    void flush();
    
    
    
    bool empty() const;

    //Element access
    proxy_elem at(size_type row, size_type col);
    const T& at(size_type row, size_type col) const;

    size_type row_back(size_type row);
    const T& row_back(size_type row) const;

    size_type row_front(size_type row);
    const T& row_front(size_type row) const;

    proxy_elem col_back(size_type col);
    const T& col_back(size_type col) const;

    proxy_elem col_front(size_type col);
    const T& col_front(size_type col) const;

    void clear();

    iterator erase( iterator loc );
    iterator erase( iterator start, iterator end );

    void swap(Matrix<T>& from);

public: //-- Additions to interface --//   

    // Does care about device data
//    proxy_elem operator()(const size_type row, const size_type col) const;
    
    // Does care about device data
    const T& operator()(const size_type row, const size_type col) const;
    
    T& operator()(const size_type row, const size_type col);
    
    // Does not care about device data, use with care
    T& operator()(const size_type index);
    
    // Does not care about device data, use with care
    T& operator[](const size_type index);
    
    // unary transpose operator
  inline Matrix<T> operator~()
  {
     Matrix<T> temp(total_cols(),total_rows());
     for (size_t i=0; i < total_rows(); i++)
      for (size_t j=0; j < total_cols(); j++)
      {
         T x = (*this)(i,j);
     temp(j,i) = x;
      }
     return temp;
  }
  
  Matrix<T>& get_transpose_matrix(bool updated = false)
  {
    if(transpose_matrix==0 || updated)
    {
      transpose_matrix = new Matrix<T>(total_cols(),total_rows());
      for (size_t i=0; i < total_rows(); i++)
      for (size_t j=0; j < total_cols(); j++)
      {
      T x = (*this)(i,j);
      (*transpose_matrix)(j,i) = x;
      }
    }
    return (*transpose_matrix);
  }

    

    // To be able to explicitly force updates without flushing entire matrix.
    // Could be used with operator () above to avoid unneccesary function calls
    // due to implicit synch.
    
    void acquireRead() const;
    void acquireReadWrite();
    
//    void updateHost() const;
//    void invalidateDevice();
//    void updateHostAndInvalidateDevice();
    
    
    const Matrix<T>& operator+=(const Matrix<T>& rhs);
    const Matrix<T>& operator+=(const T& rhs);
    
    const Matrix<T>& operator-=(const Matrix<T>& rhs);
    const Matrix<T>& operator-=(const T& rhs);
    
    const Matrix<T>& operator*=(const Matrix<T>& rhs);
    const Matrix<T>& operator*=(const T& rhs);
    
    const Matrix<T>& operator/=(const Matrix<T>& rhs);
    const Matrix<T>& operator/=(const T& rhs);
    
    const Matrix<T>& operator%=(const Matrix<T>& rhs);
    const Matrix<T>& operator%=(const T& rhs);
};

}

#include "src/matrix_iterator.inl"

#include "src/matrix_col_iterator.inl"

#include "src/matrix_proxy.inl"
#include "src/matrix.inl"

#endif