cache.hpp

Go to the documentation of this file.

//=======================================================================
// Copyright (c) 2014-2023 Baptiste Wicht
// Distributed under the terms of the MIT License.
// (See accompanying file LICENSE or copy at
//  http://opensource.org/licenses/MIT)
//=======================================================================

#pragma once

#ifdef ETL_CUDNN_MODE

#include "etl/impl/cublas/cuda.hpp"
#include "etl/impl/cudnn/cudnn.hpp"

#endif

namespace etl::impl::cudnn {

#ifdef ETL_CUDNN_MODE

template <etl_1d M>
bool fast_compare(M& lhs, M& rhs) {
    return etl::dim<0>(lhs) == etl::dim<0>(rhs);
}

template <etl_2d M>
bool fast_compare(M& lhs, M& rhs) {
    return etl::dim<0>(lhs) == etl::dim<0>(rhs) && etl::dim<1>(lhs) == etl::dim<1>(rhs);
}

template <etl_3d M>
bool fast_compare(M& lhs, M& rhs) {
    return etl::dim<0>(lhs) == etl::dim<0>(rhs) && etl::dim<1>(lhs) == etl::dim<1>(rhs) && etl::dim<2>(lhs) == etl::dim<2>(rhs);
}

template <etl_4d M>
bool fast_compare(M& lhs, M& rhs) {
    return etl::dim<0>(lhs) == etl::dim<0>(rhs) && etl::dim<1>(lhs) == etl::dim<1>(rhs) && etl::dim<2>(lhs) == etl::dim<2>(rhs)
           && etl::dim<3>(lhs) == etl::dim<3>(rhs);
}

template <typename M, bool F, size_t D>
struct mat_cache_key_impl;

template <typename M>
struct mat_cache_key_impl<M, false, 1> {
    size_t a; 

    mat_cache_key_impl() {
        // Nothing else to init
    }

    explicit mat_cache_key_impl(M& mat) : a(etl::dim<0>(mat)) {
        // Nothing else to init
    }

    bool operator==(M& rhs) {
        return a == etl::dim<0>(rhs);
    }
};

template <typename M>
struct mat_cache_key_impl<M, false, 2> {
    size_t a; 
    size_t b; 

    mat_cache_key_impl() {
        // Nothing else to init
    }

    explicit mat_cache_key_impl(M& mat) : a(etl::dim<0>(mat)), b(etl::dim<1>(mat)) {
        // Nothing else to init
    }

    bool operator==(M& rhs) {
        return a == etl::dim<0>(a) && b == etl::dim<1>(rhs);
    }
};

template <typename M>
struct mat_cache_key_impl<M, false, 3> {
    size_t a; 
    size_t b; 
    size_t c; 

    mat_cache_key_impl() {
        // Nothing else to init
    }

    explicit mat_cache_key_impl(M& mat) : a(etl::dim<0>(mat)), b(etl::dim<1>(mat)), c(etl::dim<2>(mat)) {
        // Nothing else to init
    }

    bool operator==(M& rhs) {
        return a == etl::dim<0>(rhs) && b == etl::dim<1>(rhs) && c == etl::dim<2>(rhs);
    }
};

template <typename M>
struct mat_cache_key_impl<M, false, 4> {
    size_t a; 
    size_t b; 
    size_t c; 
    size_t d; 

    mat_cache_key_impl() {
        // Nothing else to init
    }

    explicit mat_cache_key_impl(M& mat) : a(etl::dim<0>(mat)), b(etl::dim<1>(mat)), c(etl::dim<2>(mat)), d(etl::dim<3>(mat)) {
        // Nothing else to init
    }

    bool operator==(M& rhs) {
        return a == etl::dim<0>(rhs) && b == etl::dim<1>(rhs) && c == etl::dim<2>(rhs) && d == etl::dim<3>(rhs);
    }
};

template <typename M>
using mat_cache_key = mat_cache_key_impl<M, is_fast<M>, decay_traits<M>::dimensions()>;

template <typename A, typename B, typename C>
struct ternary_cache_key {
    mat_cache_key<A> key_a; 
    mat_cache_key<B> key_b; 
    mat_cache_key<C> key_c; 

    ternary_cache_key() {
        // Nothing else to init
    }

    ternary_cache_key(A& a, B& b, C& c) : key_a(a), key_b(b), key_c(c) {
        // Nothing else to init
    }

    bool equals(A& a, B& b, C& c) {
        return key_a == a && key_b == b && key_c == c;
    }
};

template <typename K, typename V, size_t L = 16>
struct ternary_static_cache {
    std::array<K, L> keys;   
    std::array<V, L> values; 

    size_t size = 0; 

    static constexpr size_t last = L; 

    template <typename A, typename B, typename C>
    size_t find(A& a, B& b, C& c) {
        for (size_t i = 0; i < size; ++i) {
            if (keys[i].equals(a, b, c)) {
                return i;
            }
        }

        return last;
    }

    template <typename A, typename B, typename C>
    size_t insert(A& a, B& b, C& c) {
        if (size == last - 1) {
            return last;
        }

        ++size;

        new (&keys[size - 1]) K(a, b, c);

        return size - 1;
    }

    auto& operator[](size_t i) {
        return values[i];
    }
};

struct conv4_descriptor {
    cudnnTensorDescriptor_t input_tensor;     
    cudnnTensorDescriptor_t output_tensor;    
    cudnnFilterDescriptor_t filter;           
    cudnnConvolutionDescriptor_t convolution; 
    cudnnConvolutionFwdAlgo_t conv_algo;      

    size_t workspace_size = 0; 
};

#endif

} //end of namespace etl::impl::cudnn