CUDA Эффективное сечение по nd-массиву (тензор)

Question

Допустим, у меня есть 3d-массив как сплющенный 1d с размерами [N, M, K]. И я хочу обработать кусок из него как [0:N, 1:M, 0:K]. Я создал вспомогательную функцию, которая обращается к базовому массиву по индексам из нарезанного массива (для простоты я делаю срезы только по второму измерению).

#define N somevalue
#define M somevalue
#define K somevalue
// i is an index in sliced array so we need to translate it into original one
template<class T, int FROM>
 __device__   __forceinline__ T slice(const T * const __restrict__ x, const size_t i) {
    auto batch_size = (M - FROM) * K;
    auto batch_index = i / batch_size;
    auto offset_0 = i % batch_size;
    auto offset_1 = offset_0 / STATES;
    auto offset_2 = offset_0 % STATES;

    return x[batch_index * M * K + (offset_1 + FROM) * K + offset_2];
}

От профилировщика NVidia я вижу, что деление и деление по модулю отнимают много вычислительных ресурсов. Также размеры не являются степенью 2, поэтому я не могу использовать трюк с битами сдвига напрямую.

Что вы можете посоветовать? Как я знаю, нарезка является довольно распространенной операцией в TF, так как они решили ее?

Answer 1

Cuda - это объединенный доступ к памяти и simd.И произвольные срезы являются полной противоположностью.Так что ответ как обычно: это зависит.

Если ваше смещение равно и остается 1, измените расположение вашей памяти в сторону MN K. Если игнорируемые записи действительно очень редки, следуйте традиционному способу и просто переведите нескольконити (да, это больно, но некоторые threadIdx calc без модуля могут быть быстрее).В противном случае вам нужно будет вычислить это биективное отображение потока / идентификатора блока для идентификатора элемента, как вы написали в своем вопросе.

Иногда вам придется проглотить горькую пилюлю.Есть несколько способов представить по модулю некоторые другие операции.Но обычно лучше потратить время на улучшение других частей ядра.