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

Скажем, у меня есть трехмерный массив в виде плоского 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 я вижу, что деление и деление по модулю требуют больших вычислительных мощностей. Также размеры не являются степенью двойки, поэтому я не могу напрямую использовать трюк с битами сдвига.

Что посоветуете? Насколько я знаю, нарезка - довольно распространенная операция в TF, так как же они ее решили?