Tile.hpp

Go to the documentation of this file.

 
//
// Copyright (c) 2021-2023, University of Colorado Denver. All rights reserved.
//
// This file is part of <T>LAPACK.
// <T>LAPACK is free software: you can redistribute it and/or modify it under
// the terms of the BSD 3-Clause license. See the accompanying LICENSE file.
 
#ifndef TLAPACK_STARPU_TILE_HH
#define TLAPACK_STARPU_TILE_HH
 
#include <starpu.h>
 
#include "tlapack/starpu/filters.hpp"
 
namespace tlapack {
namespace starpu {
 
    struct Tile {
        const starpu_data_handle_t root_handle;  
        const idx_t i, j;  
        const idx_t m, n;  
 
        starpu_data_handle_t handle;  
        bool partition_planned =
            false;  
 
        Tile(starpu_data_handle_t tile_handle,
             idx_t i,
             idx_t j,
             idx_t m,
             idx_t n) noexcept
            : root_handle(tile_handle), i(i), j(j), m(m), n(n)
        {
            // Collect information about the tile
            const idx_t M = starpu_matrix_get_nx(tile_handle);
            const idx_t N = starpu_matrix_get_ny(tile_handle);
 
            // Partition the tile if it is not a full tile
            if (m != M || n != N) {
                idx_t pos[4] = {i, j, m, n};
                struct starpu_data_filter f_tile = {
                    .filter_func = filter_tile,
                    .nchildren = 1,
                    .filter_arg_ptr = (void*)pos};
 
                starpu_data_partition_plan(root_handle, &f_tile, &handle);
                partition_planned = true;
            }
            else {
                handle = root_handle;
            }
 
            assert(starpu_matrix_get_nx(handle) == m &&
                   starpu_matrix_get_ny(handle) == n && "Invalid tile size");
        }
 
        ~Tile() noexcept
        {
            if (partition_planned)
                starpu_data_partition_clean(root_handle, 1, &handle);
        }
 
        static void create_compatible_handles(starpu_data_handle_t handles[2],
                                              const Tile& A,
                                              const Tile& B) noexcept
        {
            if (A.root_handle == B.root_handle) {
                idx_t pos[8] = {A.i, A.j, A.m, A.n, B.i, B.j, B.m, B.n};
 
                struct starpu_data_filter f_ntiles = {
                    .filter_func = filter_ntiles,
                    .nchildren = 2,
                    .filter_arg_ptr = (void*)pos};
 
                starpu_data_partition_plan(A.root_handle, &f_ntiles, handles);
 
                assert(starpu_matrix_get_nx(handles[0]) == A.m &&
                       starpu_matrix_get_ny(handles[0]) == A.n &&
                       "Invalid tile size");
                assert(starpu_matrix_get_nx(handles[1]) == B.m &&
                       starpu_matrix_get_ny(handles[1]) == B.n &&
                       "Invalid tile size");
            }
            else {
                handles[0] = A.handle;
                handles[1] = B.handle;
            }
        }
 
        static void clean_compatible_handles(starpu_data_handle_t handles[2],
                                             const Tile& A,
                                             const Tile& B) noexcept
        {
            if (A.root_handle == B.root_handle)
                starpu_data_partition_clean(A.root_handle, 2, handles);
        }
 
        void create_compatible_inout_handles(starpu_data_handle_t handles[3],
                                             const Tile& A,
                                             const Tile& B) const noexcept
        {
            if (root_handle == A.root_handle) {
                if (root_handle == B.root_handle) {
                    idx_t pos[12] = {i,   j,   m,   n,   A.i, A.j,
                                     A.m, A.n, B.i, B.j, B.m, B.n};
 
                    struct starpu_data_filter f_ntiles = {
                        .filter_func = filter_ntiles,
                        .nchildren = 3,
                        .filter_arg_ptr = (void*)pos};
 
                    starpu_data_partition_plan(root_handle, &f_ntiles, handles);
 
                    assert(starpu_matrix_get_nx(handles[0]) == m &&
                           starpu_matrix_get_ny(handles[0]) == n &&
                           "Invalid tile size");
                    assert(starpu_matrix_get_nx(handles[1]) == A.m &&
                           starpu_matrix_get_ny(handles[1]) == A.n &&
                           "Invalid tile size");
                    assert(starpu_matrix_get_nx(handles[2]) == B.m &&
                           starpu_matrix_get_ny(handles[2]) == B.n &&
                           "Invalid tile size");
                }
                else {
                    create_compatible_handles(handles, *this, A);
                    handles[2] = B.handle;
                }
            }
            else if (root_handle == B.root_handle) {
                create_compatible_handles(handles, *this, B);
                handles[2] = handles[1];
                handles[1] = A.handle;
            }
            else {
                handles[0] = handle;
                handles[1] = A.handle;
                handles[2] = B.handle;
            }
        }
 
        void clean_compatible_inout_handles(starpu_data_handle_t handles[3],
                                            const Tile& A,
                                            const Tile& B) const noexcept
        {
            if (root_handle == A.root_handle) {
                if (root_handle == B.root_handle)
                    starpu_data_partition_clean(root_handle, 3, handles);
                else
                    starpu_data_partition_clean(root_handle, 2, handles);
            }
            else if (root_handle == B.root_handle) {
                starpu_data_handle_t aux = handles[1];
                handles[1] = handles[2];
                starpu_data_partition_clean(root_handle, 2, handles);
                handles[2] = handles[1];
                handles[1] = aux;
            }
        }
    };
 
}  // namespace starpu
}  // namespace tlapack
 
#endif  // TLAPACK_STARPU_TILE_HPP