Example 13: Non-uniform Block Sizes

This example demonstrates creating a matrix with non-uniform tile sizes.

Key Concepts

  1. Lambda Constructors: Using lambda functions to define tile properties (tileNb, tileRank, tileDevice) instead of fixed values.

  2. Custom Block Sizes: Defining a function that returns the block size for a given block index \(j\).

  3. Process/Device Mapping: Using helper functions like slate::func::process_2d_grid and slate::func::device_1d_grid to define distribution.

C++ Example

Custom Block Size Function (Lines 30-34)

std::function< int64_t (int64_t j) >
tileNb = [n, nb_](int64_t j)
{
    return (j % 2 != 0 ? nb_/2 : nb_);
};

Instead of passing a constant nb, we define a lambda function tileNb. - Input: Block index j (0, 1, 2, …). - Output: The number of columns in block column j. - This example alternates block sizes between nb and nb/2.

Distribution Functions (Lines 36-38)

auto tileRank = slate::func::process_2d_grid( slate::GridOrder::Col, p_, q_ );
auto tileDevice = slate::func::device_1d_grid( slate::GridOrder::Col, p_, num_devices_ );

We can also customize how tiles map to MPI ranks and devices. Here we use standard helper functions from slate::func, but you could write your own lambdas (e.g., to force a specific block cyclic pattern or a custom distribution).

Matrix Construction (Line 40)

slate::Matrix<scalar_type> A( n, n, tileNb, tileNb, tileRank, tileDevice, MPI_COMM_WORLD );

The constructor takes these functions as arguments. - tileNb is passed twice: once for row heights tileMb and once for column widths tileNb. This creates a square-tiled matrix (though the tiles themselves can vary in size). - tileRank determines the MPI rank for tile (i, j). - tileDevice determines the GPU device ID for tile (i, j).

Verification (Lines 53-59)

The code iterates through the matrix to verify that the block sizes match the logic defined in the lambda. Note that A.tileNb(j) handles the boundary condition at the end of the matrix automatically (clamping to n).

  1// ex13_non_uniform_block_size.cc
  2// create 1000 x 1000 matrix on 2 x 2 MPI process grid, with non-uniform tile size
  3
  4/// !!!   Lines between `//---------- begin label`          !!!
  5/// !!!             and `//---------- end label`            !!!
  6/// !!!   are included in the SLATE Users' Guide.           !!!
  7
  8#include <slate/slate.hh>
  9
 10#include "util.hh"
 11
 12int mpi_size = 0;
 13int mpi_rank = 0;
 14int grid_p = 0;
 15int grid_q = 0;
 16
 17//------------------------------------------------------------------------------
 18template <typename scalar_type>
 19void test_matrix_lambda()
 20{
 21    print_func( mpi_rank );
 22
 23    int64_t n=1000, nb=256;
 24
 25    int nb_ = nb;
 26    int p_ = grid_p;
 27    int q_ = grid_q;
 28    int num_devices_ = 0;
 29
 30    std::function< int64_t (int64_t j) >
 31    tileNb = [n, nb_](int64_t j)
 32    {
 33        return (j % 2 != 0 ? nb_/2 : nb_);
 34    };
 35
 36    auto tileRank = slate::func::process_2d_grid( slate::GridOrder::Col, p_, q_ );
 37    auto tileDevice = slate::func::device_1d_grid( slate::GridOrder::Col,
 38                                                   p_, num_devices_ );
 39
 40    slate::Matrix<scalar_type> A( n, n, tileNb, tileNb, tileRank, tileDevice, MPI_COMM_WORLD );
 41    A.insertLocalTiles();
 42
 43    for (int64_t j = 0; j < A.nt(); ++j) {
 44        for (int64_t i = 0; i < A.mt(); ++i) {
 45            if (A.tileIsLocal( i, j )) {
 46                slate::Tile<scalar_type> T = A( i, j );
 47                random_matrix( T.mb(), T.nb(), T.data(), T.stride() );
 48            }
 49        }
 50    }
 51
 52    // verify nt, tileNb(i), and sum tileNb(i) == n
 53    int nt = A.nt();
 54    int jj = 0;
 55    for (int j = 0; j < nt; ++j) {
 56        assert( A.tileNb(j) == blas::min( tileNb(j), n - jj ) );
 57        jj += A.tileNb( j );
 58    }
 59    assert( jj == n );
 60}
 61
 62//------------------------------------------------------------------------------
 63int main( int argc, char** argv )
 64{
 65    try {
 66        // Parse command line to set types for s, d, c, z precisions.
 67        bool types[ 4 ];
 68        parse_args( argc, argv, types );
 69
 70        int provided = 0;
 71        slate_mpi_call(
 72            MPI_Init_thread( &argc, &argv, MPI_THREAD_MULTIPLE, &provided ) );
 73        assert( provided == MPI_THREAD_MULTIPLE );
 74
 75        slate_mpi_call(
 76            MPI_Comm_size( MPI_COMM_WORLD, &mpi_size ) );
 77
 78        slate_mpi_call(
 79            MPI_Comm_rank( MPI_COMM_WORLD, &mpi_rank ) );
 80
 81        // Determine p-by-q grid for this MPI size.
 82        grid_size( mpi_size, &grid_p, &grid_q );
 83        if (mpi_rank == 0) {
 84            printf( "mpi_size %d, grid_p %d, grid_q %d\n",
 85                    mpi_size, grid_p, grid_q );
 86        }
 87
 88        // so random_matrix is different on different ranks.
 89        srand( 100 * mpi_rank );
 90
 91        if (types[ 0 ]) {
 92            test_matrix_lambda< float >();
 93        }
 94
 95        if (types[ 1 ]) {
 96            test_matrix_lambda< double >();
 97        }
 98
 99        if (types[ 2 ]) {
100            test_matrix_lambda< std::complex<float> >();
101        }
102
103        if (types[ 3 ]) {
104            test_matrix_lambda< std::complex<double> >();
105        }
106
107        slate_mpi_call(
108            MPI_Finalize() );
109    }
110    catch (std::exception const& ex) {
111        fprintf( stderr, "%s", ex.what() );
112        return 1;
113    }
114    return 0;
115}