Hugging Face's logo

Kernels:
flashrt
/
fp4-gemm

Kernel card Files
xet
 Community
fp4-gemm / csrc /cutlass /util /packed_stride.hpp
liangsu9988's picture
liangsu9988
Publish GROOT N1.7 Thor kernel sources
88b8ef2 verified
Raw
History Blame
20.9 kB
/***************************************************************************************************
 * Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * SPDX-License-Identifier: BSD-3-Clause
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 * 1. Redistributions of source code must retain the above copyright notice, this
 * list of conditions and the following disclaimer.
 *
 * 2. Redistributions in binary form must reproduce the above copyright notice,
 * this list of conditions and the following disclaimer in the documentation
 * and/or other materials provided with the distribution.
 *
 * 3. Neither the name of the copyright holder nor the names of its
 * contributors may be used to endorse or promote products derived from
 * this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
 * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
 * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
 * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *
 **************************************************************************************************/
/*! \file
 \brief Utilities for packing constructing canonical CuTe stride types for 3.x mainloop params.
*/
#pragma once
#include "cute/layout.hpp"
#include "cute/container/array.hpp" // cute::array
#include "cutlass/conv/convolution.h" // cutlass::conv::Operator
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass {
/////////////////////////////////////////////////////////////////////////////////////////////////
// Strides without batch mode
template <class IntT>
CUTLASS_HOST_DEVICE
cute::Stride<IntT, cute::Int<1>>
make_cute_packed_stride(cute::Stride<IntT, cute::Int<1>> s, cute::Shape<int,int,int> shape_MKL) {
 static_assert(std::is_integral_v<IntT>,
 "Stride must have an integral type so it can be set dynamically. Static strides not supported.");
 auto s_copy = s;
 cute::get<0>(s_copy) = static_cast<IntT>(cute::get<1>(shape_MKL));
 return s_copy;
}
template <class IntT>
CUTLASS_HOST_DEVICE
cute::Stride<cute::Int<1>, IntT>
make_cute_packed_stride(cute::Stride<cute::Int<1>, IntT> s, cute::Shape<int,int,int> shape_MKL) {
 static_assert(std::is_integral_v<IntT>,
 "Stride must have an integral type so it can be set dynamically. Static strides not supported.");
 auto s_copy = s;
 cute::get<1>(s_copy) = static_cast<IntT>(cute::get<0>(shape_MKL));
 return s_copy;
}
/////////////////////////////////////////////////////////////////////////////////////////////////
// Strides with batch mode
template <class IntT>
CUTLASS_HOST_DEVICE
cute::Stride<IntT, cute::Int<1>, int64_t>
make_cute_packed_stride(cute::Stride<IntT, cute::Int<1>, int64_t> s, cute::Shape<int,int,int> shape_MKL) {
 static_assert(std::is_integral_v<IntT>,
 "Stride must have an integral type so it can be set dynamically. Static strides not supported.");
 auto s_copy = s;
 cute::get<0>(s_copy) = static_cast<IntT>(cute::get<1>(shape_MKL));
 int batch_count = cute::get<2>(shape_MKL);
 if (batch_count > 1) {
 cute::get<2>(s_copy) = static_cast<IntT>(cute::get<0>(shape_MKL) * cute::get<1>(shape_MKL));
 }
 else {
 cute::get<2>(s_copy) = static_cast<IntT>(0);
 }
 return s_copy;
}
template <class IntT>
CUTLASS_HOST_DEVICE
cute::Stride<cute::Int<1>, IntT, int64_t>
make_cute_packed_stride(cute::Stride<cute::Int<1>, IntT, int64_t> s, cute::Shape<int,int,int> shape_MKL) {
 static_assert(std::is_integral_v<IntT>,
 "Stride must have an integral type so it can be set dynamically. Static strides not supported.");
 auto s_copy = s;
 cute::get<1>(s_copy) = static_cast<IntT>(cute::get<0>(shape_MKL));
 int batch_count = cute::get<2>(shape_MKL);
 if (batch_count > 1) {
 cute::get<2>(s_copy) = static_cast<IntT>(cute::get<0>(shape_MKL) * cute::get<1>(shape_MKL));
 }
 else {
 cute::get<2>(s_copy) = static_cast<IntT>(0);
 }
 return s_copy;
}
/////////////////////////////////////////////////////////////////////////////////////////////////
// Strides with group mode
template <class StrideIntT>
CUTLASS_HOST_DEVICE
cute::Stride<StrideIntT, cute::Int<1>, cute::Int<0>>
make_cute_packed_stride(cute::Stride<StrideIntT, cute::Int<1>, cute::Int<0>> s, cute::Shape<int,int,int> shape_MKL) {
 static_assert(std::is_integral_v<StrideIntT>,
 "Stride must have an integral type so it can be set dynamically. Static strides not supported.");
 auto s_copy = s;
 cute::get<0>(s_copy) = static_cast<StrideIntT>(cute::get<1>(shape_MKL));
 return s_copy;
}
template <class StrideIntT>
CUTLASS_HOST_DEVICE
cute::Stride<cute::Int<1>, StrideIntT, cute::Int<0>>
make_cute_packed_stride(cute::Stride<cute::Int<1>, StrideIntT, cute::Int<0>> s, cute::Shape<int,int,int> shape_MKL) {
 static_assert(std::is_integral_v<StrideIntT>,
 "Stride must have an integral type so it can be set dynamically. Static strides not supported.");
 auto s_copy = s;
 cute::get<1>(s_copy) = static_cast<StrideIntT>(cute::get<0>(shape_MKL));
 return s_copy;
}
/////////////////////////////////////////////////////////////////////////////////////////////////
// Strides for convolutions
// Output cutlass::layout::TensorNDHWC -> rank-3 stride (InT,_1,_0)
// Note: For fprop/dgrad kernel, strides are assumed to be layout right in NZPQK/NDHWC order
// and therefore can be coalesced to just q/w. For wgrad kernel, strides are assumed to be layout
// right in KTRSC order and can be coalesced to just k.
// We enforce this condition here with asserts.
template <class IntT, size_t RankT_>
CUTLASS_HOST_DEVICE
cute::Stride<IntT, cute::Int<1>, cute::Int<0>>
make_cute_packed_stride(
 cute::Stride<IntT, cute::Int<1>, cute::Int<0>> s,
 cute::array<int32_t, RankT_> shape_output,
 cute::array<IntT, RankT_> stride_output,
 cutlass::conv::Operator conv_op) {
 static_assert(std::is_integral_v<IntT>,
 "Stride must have an integral type so it can be set dynamically. Static strides not supported.");
 static_assert(RankT_ >= 3u);
 constexpr static int RankT = static_cast<int>(RankT_);
assert(stride_output[RankT-1] == 1);
 cute::for_each(cute::make_seq<RankT-2>{}, [&](auto i) {
 assert(stride_output[i] == shape_output[i+1] * stride_output[i+1]);
 });
auto s_copy = s;
 cute::get<0>(s_copy) = (conv_op == cutlass::conv::Operator::kWgrad) ?
 stride_output[0] :
 stride_output[RankT-2];
 return s_copy;
}
//
// Activation tensor ((w, h, d, n), _1) for fprop kernel
//
// Activation cutlass::layout::TensorNWC -> rank-2 stride ((W,N),_1)
template <class IntT>
CUTLASS_HOST_DEVICE
cute::Stride<cute::Stride<IntT, IntT>, cute::Int<1>>
make_cute_packed_stride(
 cute::Stride<cute::Stride<IntT, IntT>, cute::Int<1>> s,
 cute::array<IntT, 3> stride_nwc,
 conv::Operator ConvOp) {
 static_assert(std::is_integral_v<IntT>,
 "Stride must have an integral type so it can be set dynamically. Static strides not supported.");
 assert(stride_nwc[2] == 1);
 auto s_copy = s;
 cute::get<0,0>(s_copy) = stride_nwc[1];
 cute::get<0,1>(s_copy) = stride_nwc[0];
 return s_copy;
}
// Activation cutlass::layout::TensorNHWC -> rank-2 stride ((W,H,N),_1)
template <class IntT>
CUTLASS_HOST_DEVICE
cute::Stride<cute::Stride<IntT, IntT, IntT>, cute::Int<1>>
make_cute_packed_stride(
 cute::Stride<cute::Stride<IntT, IntT, IntT>, cute::Int<1>> s,
 cute::array<IntT, 4> stride_nhwc,
 conv::Operator ConvOp) {
 static_assert(std::is_integral_v<IntT>,
 "Stride must have an integral type so it can be set dynamically. Static strides not supported.");
 assert(stride_nhwc[3] == 1);
 auto s_copy = s;
 cute::for_each(cute::make_seq<3>{}, [&](auto i) {
 cute::get<0,i>(s_copy) = stride_nhwc[2-i];
 });
 return s_copy;
}
// Activation cutlass::layout::TensorNDHWC -> rank-2 stride ((W,H,D,N),_1)
template <class IntT>
CUTLASS_HOST_DEVICE
cute::Stride<cute::Stride<IntT, IntT, IntT, IntT>, cute::Int<1>>
make_cute_packed_stride(
 cute::Stride<cute::Stride<IntT, IntT, IntT, IntT>, cute::Int<1>> s,
 cute::array<IntT, 5> stride_ndhwc,
 conv::Operator ConvOp) {
 static_assert(std::is_integral_v<IntT>,
 "Stride must have an integral type so it can be set dynamically. Static strides not supported.");
assert(stride_ndhwc[4] == 1);
 auto s_copy = s;
 cute::for_each(cute::make_seq<4>{}, [&](auto i) {
 cute::get<0,i>(s_copy) = stride_ndhwc[3-i];
 });
 return s_copy;
}
//
// Filter tensor (k, (_1, s, r, t)) for fprop kernel
//
// Filter cutlass::layout::TensorNWC -> rank-2 stride (k, (_1, s))
template <class IntT>
CUTLASS_HOST_DEVICE
cute::Stride<IntT, cute::Stride<cute::Int<1>, IntT>>
make_cute_packed_stride(
 cute::Stride<IntT, cute::Stride<cute::Int<1>, IntT>> s,
 cute::array<IntT, 3> stride_ksc,
 conv::Operator ConvOp) {
 static_assert(std::is_integral_v<IntT>,
 "Stride must have an integral type so it can be set dynamically. Static strides not supported.");
assert(stride_ksc[2] == 1);
 auto s_copy = s;
 cute::get<0,0>(s_copy) = stride_ksc[0];
 cute::get<1,1>(s_copy) = stride_ksc[1];
 return s_copy;
}
// Filter cutlass::layout::TensorNHWC -> rank-2 stride (k, (_1, s, r))
template <class IntT>
CUTLASS_HOST_DEVICE
cute::Stride<IntT, cute::Stride<cute::Int<1>, IntT, IntT>>
make_cute_packed_stride(
 cute::Stride<IntT, cute::Stride<cute::Int<1>, IntT, IntT>> s,
 cute::array<IntT, 4> stride_krsc,
 conv::Operator ConvOp) {
 static_assert(std::is_integral_v<IntT>,
 "Stride must have an integral type so it can be set dynamically. Static strides not supported.");
assert(stride_krsc[3] == 1);
 auto s_copy = s;
 cute::get<0,0>(s_copy) = stride_krsc[0];
 cute::for_each(cute::make_seq<2>{}, [&](auto i) {
 cute::get<1,2-i>(s_copy) = stride_krsc[i+1];
 });
 return s_copy;
}
// Filter cutlass::layout::TensorNDHWC -> rank-2 stride (k, (_1, s, r, t))
template <class IntT>
CUTLASS_HOST_DEVICE
cute::Stride<IntT, cute::Stride<cute::Int<1>, IntT, IntT, IntT>>
make_cute_packed_stride(
 cute::Stride<IntT, cute::Stride<cute::Int<1>, IntT, IntT, IntT>> s,
 cute::array<IntT, 5> stride_ktrsc,
 conv::Operator ConvOp) {
 static_assert(std::is_integral_v<IntT>,
 "Stride must have an integral type so it can be set dynamically. Static strides not supported.");
assert(stride_ktrsc[4] == 1);
 auto s_copy = s;
 cute::get<0,0>(s_copy) = stride_ktrsc[0];
 cute::for_each(cute::make_seq<3>{}, [&](auto i) {
 cute::get<1,3-i>(s_copy) = stride_ktrsc[i+1];
 });
 return s_copy;
}
//
// Activation tensor (_1, (w, h, d, n)) for wgrad kernel
//
// It is also Filter tensor ((_1), (k, s, r, t)) for dgrad kernel
//
// Activation cutlass::layout::TensorNWC -> rank-2 stride (_1, (W,N)) in wgrad
// Filter cutlass::layout::TensorNWC -> rank-2 stride ((_1), (k, s)) in dgrad
template <class IntT>
CUTLASS_HOST_DEVICE
cute::Stride<cute::Int<1>, cute::Stride<IntT, IntT>>
make_cute_packed_stride(
 cute::Stride<cute::Int<1>, cute::Stride<IntT, IntT>> s,
 cute::array<IntT, 3> stride_nwc,
 conv::Operator ConvOp) {
 static_assert(std::is_integral_v<IntT>,
 "Stride must have an integral type so it can be set dynamically. Static strides not supported.");
assert(stride_nwc[2] == 1);
 auto s_copy = s;
 if (ConvOp == cutlass::conv::Operator::kWgrad) {
 cute::get<1,0>(s_copy) = stride_nwc[1];
 cute::get<1,1>(s_copy) = stride_nwc[0];
 }
 else if (ConvOp == cutlass::conv::Operator::kDgrad) {
 // stride_nwc in dgrad is ksc.
 cute::get<1,0>(s_copy) = stride_nwc[0];
 cute::get<1,1>(s_copy) = stride_nwc[1];
 }
 return s_copy;
}
// Activation cutlass::layout::TensorNHWC -> rank-2 stride (_1, (W,H,N)) in wgrad
// Filter cutlass::layout::TensorNHWC -> rank-2 stride ((_1), (k, s, r)) in dgrad
template <class IntT>
CUTLASS_HOST_DEVICE
cute::Stride<cute::Int<1>, cute::Stride<IntT, IntT, IntT>>
make_cute_packed_stride(
 cute::Stride<cute::Int<1>, cute::Stride<IntT, IntT, IntT>> s,
 cute::array<IntT, 4> stride_nhwc,
 conv::Operator ConvOp) {
 static_assert(std::is_integral_v<IntT>,
 "Stride must have an integral type so it can be set dynamically. Static strides not supported.");
assert(stride_nhwc[3] == 1);
 auto s_copy = s;
 if (ConvOp == cutlass::conv::Operator::kWgrad) {
 cute::for_each(cute::make_seq<3>{}, [&](auto i) {
 cute::get<1,i>(s_copy) = stride_nhwc[2-i];
 });
 }
 else if (ConvOp == cutlass::conv::Operator::kDgrad) {
 // stride_nhwc in dgrad is krsc.
 cute::get<1,0>(s_copy) = stride_nhwc[0];
 cute::for_each(cute::make_seq<2>{}, [&](auto i) {
 cute::get<1,2-i>(s_copy) = stride_nhwc[i+1];
 });
 }
 return s_copy;
}
// Activation cutlass::layout::TensorNDHWC -> rank-2 stride (_1, (W,H,D,N)) in wgrad
// Filter cutlass::layout::TensorNDHWC -> rank-2 stride ((_1), (k, s, r, t)) in dgrad
template <class IntT>
CUTLASS_HOST_DEVICE
cute::Stride<cute::Int<1>, cute::Stride<IntT, IntT, IntT, IntT>>
make_cute_packed_stride(
 cute::Stride<cute::Int<1>, cute::Stride<IntT, IntT, IntT, IntT>> s,
 cute::array<IntT, 5> stride_ndhwc,
 conv::Operator ConvOp) {
 static_assert(std::is_integral_v<IntT>,
 "Stride must have an integral type so it can be set dynamically. Static strides not supported.");
assert(stride_ndhwc[4] == 1);
 auto s_copy = s;
 if (ConvOp == cutlass::conv::Operator::kWgrad) {
 cute::for_each(cute::make_seq<4>{}, [&](auto i) {
 cute::get<1,i>(s_copy) = stride_ndhwc[3-i];
 });
 }
 else if (ConvOp == cutlass::conv::Operator::kDgrad) {
 // stride_ndhwc in dgrad is ktrsc.
 cute::get<1,0>(s_copy) = stride_ndhwc[0];
 cute::for_each(cute::make_seq<3>{}, [&](auto i) {
 cute::get<1,3-i>(s_copy) = stride_ndhwc[i+1];
 });
 }
 return s_copy;
}
//
// NZPQ tensor (_1, nzpq) for wgrad kernel
//
// cutlass::layout::TensorNWC -> rank-2 stride (_1, nzpq)
template <class IntT>
CUTLASS_HOST_DEVICE
cute::Stride<cute::Int<1>, IntT>
make_cute_packed_stride(
 cute::Stride<cute::Int<1>, IntT> s,
 cute::array<IntT, 3> stride_nqk,
 conv::Operator ConvOp) {
 static_assert(std::is_integral_v<IntT>,
 "Stride must have an integral type so it can be set dynamically. Static strides not supported.");
assert(stride_nqk[2] == 1);
 auto s_copy = s;
 cute::get<1>(s_copy) = stride_nqk[1];
 return s_copy;
}
// cutlass::layout::TensorNHWC -> rank-2 stride (_1, nzpq)
template <class IntT>
CUTLASS_HOST_DEVICE
cute::Stride<cute::Int<1>, IntT>
make_cute_packed_stride(
 cute::Stride<cute::Int<1>, IntT> s,
 cute::array<IntT, 4> stride_npqk,
 conv::Operator ConvOp) {
 static_assert(std::is_integral_v<IntT>,
 "Stride must have an integral type so it can be set dynamically. Static strides not supported.");
assert(stride_npqk[3] == 1);
 auto s_copy = s;
 cute::get<1>(s_copy) = stride_npqk[2];
 return s_copy;
}
// cutlass::layout::TensorNDHWC -> rank-2 stride (_1, nzpq)
template <class IntT>
CUTLASS_HOST_DEVICE
cute::Stride<cute::Int<1>, IntT>
make_cute_packed_stride(
 cute::Stride<cute::Int<1>, IntT> s,
 cute::array<IntT, 5> stride_nzpqk,
 conv::Operator ConvOp) {
 static_assert(std::is_integral_v<IntT>,
 "Stride must have an integral type so it can be set dynamically. Static strides not supported.");
assert(stride_nzpqk[4] == 1);
 auto s_copy = s;
 cute::get<1>(s_copy) = stride_nzpqk[3];
 return s_copy;
}
//
// Wgrad output tensor (k, (_1, s, r, t), _0)
//
// Filter cutlass::layout::TensorKCS -> rank-3 stride (k, (_1, s), _0)
template <class IntT>
CUTLASS_HOST_DEVICE
cute::Stride<IntT, cute::Stride<cute::Int<1>, IntT>, cute::Int<0>>
make_cute_packed_stride(
 cute::Stride<IntT, cute::Stride<cute::Int<1>, IntT>, cute::Int<0>> s,
 [[maybe_unused]] cute::array<int32_t, 3> shape_output,
 cute::array<IntT, 3> stride_ksc,
 conv::Operator ConvOp) {
 static_assert(std::is_integral_v<IntT>,
 "Stride must have an integral type so it can be set dynamically. Static strides not supported.");
assert(stride_ksc[2] == 1);
 auto s_copy = s;
 cute::get<0,0>(s_copy) = stride_ksc[0];
 cute::get<1,1>(s_copy) = stride_ksc[1];
 return s_copy;
}
// Filter cutlass::layout::TensorKCSR -> rank-3 stride (k, (_1, s, r), _0)
template <class IntT>
CUTLASS_HOST_DEVICE
cute::Stride<IntT, cute::Stride<cute::Int<1>, IntT, IntT>, cute::Int<0>>
make_cute_packed_stride(
 cute::Stride<IntT, cute::Stride<cute::Int<1>, IntT, IntT>, cute::Int<0>> s,
 [[maybe_unused]] cute::array<int32_t, 4> shape_output,
 cute::array<IntT, 4> stride_krsc,
 conv::Operator ConvOp) {
 static_assert(std::is_integral_v<IntT>,
 "Stride must have an integral type so it can be set dynamically. Static strides not supported.");
assert(stride_krsc[3] == 1);
 auto s_copy = s;
 cute::get<0,0>(s_copy) = stride_krsc[0];
 cute::for_each(cute::make_seq<2>{}, [&](auto i) {
 cute::get<1,2-i>(s_copy) = stride_krsc[i+1];
 });
 return s_copy;
}
// Filter cutlass::layout::TensorKCSRT -> rank-3 stride (k, (_1, s, r, t), _0)
template <class IntT>
CUTLASS_HOST_DEVICE
cute::Stride<IntT, cute::Stride<cute::Int<1>, IntT, IntT, IntT>, cute::Int<0>>
make_cute_packed_stride(
 cute::Stride<IntT, cute::Stride<cute::Int<1>, IntT, IntT, IntT>, cute::Int<0>> s,
 [[maybe_unused]] cute::array<int32_t, 5> shape_output,
 cute::array<IntT, 5> stride_ktrsc,
 conv::Operator ConvOp) {
 static_assert(std::is_integral_v<IntT>,
 "Stride must have an integral type so it can be set dynamically. Static strides not supported.");
assert(stride_ktrsc[4] == 1);
 auto s_copy = s;
 cute::get<0,0>(s_copy) = stride_ktrsc[0];
 cute::for_each(cute::make_seq<3>{}, [&](auto i) {
 cute::get<1,3-i>(s_copy) = stride_ktrsc[i+1];
 });
 return s_copy;
}
//
// Wgrad output tensor ((_1, s, r, t), k, _0)
//
// Filter cutlass::layout::TensorCSK -> rank-3 stride ((_1, s), k, _0)
template <class IntT>
CUTLASS_HOST_DEVICE
cute::Stride<cute::Stride<cute::Int<1>, IntT>, IntT, cute::Int<0>>
make_cute_packed_stride(
 cute::Stride<cute::Stride<cute::Int<1>, IntT>, IntT, cute::Int<0>> s,
 [[maybe_unused]] cute::array<int32_t, 3> shape_output,
 cute::array<IntT, 3> stride_ksc,
 conv::Operator ConvOp) {
 static_assert(std::is_integral_v<IntT>,
 "Stride must have an integral type so it can be set dynamically. Static strides not supported.");
assert(stride_ksc[2] == 1);
 auto s_copy = s;
 cute::get<1,0>(s_copy) = stride_ksc[0];
 cute::get<0,1>(s_copy) = stride_ksc[1];
 return s_copy;
}
// Filter cutlass::layout::TensorCSRK -> rank-3 stride ((_1, s, r), k, _0)
template <class IntT>
CUTLASS_HOST_DEVICE
cute::Stride<cute::Stride<cute::Int<1>, IntT, IntT>, IntT, cute::Int<0>>
make_cute_packed_stride(
 cute::Stride<cute::Stride<cute::Int<1>, IntT, IntT>, IntT, cute::Int<0>> s,
 [[maybe_unused]] cute::array<int32_t, 4> shape_output,
 cute::array<IntT, 4> stride_krsc,
 conv::Operator ConvOp) {
 static_assert(std::is_integral_v<IntT>,
 "Stride must have an integral type so it can be set dynamically. Static strides not supported.");
assert(stride_krsc[3] == 1);
 auto s_copy = s;
 cute::get<1,0>(s_copy) = stride_krsc[0];
 cute::for_each(cute::make_seq<2>{}, [&](auto i) {
 cute::get<0,2-i>(s_copy) = stride_krsc[i+1];
 });
 return s_copy;
}
// Filter cutlass::layout::TensorCSRTK -> rank-3 stride ((_1, s, r, t), k, _0)
template <class IntT>
CUTLASS_HOST_DEVICE
cute::Stride<cute::Stride<cute::Int<1>, IntT, IntT, IntT>, IntT, cute::Int<0>>
make_cute_packed_stride(
 cute::Stride<cute::Stride<cute::Int<1>, IntT, IntT, IntT>, IntT, cute::Int<0>> s,
 [[maybe_unused]] cute::array<int32_t, 5> shape_output,
 cute::array<IntT, 5> stride_ktrsc,
 conv::Operator ConvOp) {
 static_assert(std::is_integral_v<IntT>,
 "Stride must have an integral type so it can be set dynamically. Static strides not supported.");
assert(stride_ktrsc[4] == 1);
 auto s_copy = s;
 cute::get<1,0>(s_copy) = stride_ktrsc[0];
 cute::for_each(cute::make_seq<3>{}, [&](auto i) {
 cute::get<0,3-i>(s_copy) = stride_ktrsc[i+1];
 });
 return s_copy;
}
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace cutlass