leavelet/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 24 Packages Projects Releases Wiki Activity

v4.5 dev update. (#3153)

Browse Source
This commit is contained in:
Junkai-Wu
2026-04-07 12:16:05 -04:00
committed by GitHub
parent 418d38a5de
commit a221da7ccf
265 changed files with 4913 additions and 1478 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/cute/tensor_zip.hpp
+70 -2
View File
@@ -72,9 +72,12 @@ struct ZipIterator
template <class... Index>
CUTE_HOST_DEVICE constexpr
ZipIterator operator+(cute::tuple<Index...> const& idxs) const {
auto operator+(cute::tuple<Index...> const& idxs) const {
static_assert(sizeof...(Index) == sizeof...(Iters), "Expect same number of offsets as iterators.");
return cute::transform(iters_, idxs, [](auto&& iter, auto&& idx) { return iter + idx; });
return cute::transform_apply(iters_, idxs,
[](auto&& iter, auto&& idx) { return iter + idx; },
[](auto... iter) { return ZipIterator<decltype(iter)...>(iter...); }
);
}
template <class... Index>
@@ -149,6 +152,13 @@ struct ZipLayout
template <class... Layouts>
struct is_layout<ZipLayout<Layouts...>> : true_type {};
template <class Layout>
struct is_zip_layout : false_type {};
template <class... Layouts>
struct is_zip_layout<ZipLayout<Layouts...>> : true_type {};
//
// make_zip_tensor and unzip_tensor
//
@@ -191,6 +201,23 @@ size(ZipLayout<Layouts...> const& layouts)
return size<Is...>(get<0>(layouts.layouts_));
}
template <int... Is, class... Layouts>
CUTE_HOST_DEVICE constexpr
auto
get(ZipLayout<Layouts...> const& layouts)
{
return ZipLayout(cute::transform(layouts.layouts_, [&](auto t){ return get<Is...>(t); }));
}
template <int... Is, class... Layouts>
CUTE_HOST_DEVICE constexpr
auto
layout(ZipLayout<Layouts...> const& layouts)
{
return get<Is...>(layouts);
}
//
// Manipulation
//
@@ -243,4 +270,45 @@ slice_and_offset(Coord const& c, ZipLayout<Layouts...> const& layouts)
return cute::make_tuple(ZipLayout(get<0>(result)), get<1>(result));
}
template <int B, int E, class... Layouts>
CUTE_HOST_DEVICE constexpr
auto
group(ZipLayout<Layouts...> const& layouts)
{
return ZipLayout(cute::transform(layouts.layouts_, [&](auto t){ return group<B,E>(t); }));
}
template <int... Is, class... Layouts>
CUTE_HOST_DEVICE constexpr
auto
shape(ZipLayout<Layouts...> const& layouts) {
return shape<Is...>(get<0>(layouts.layouts_));
}
template <int... Is, class... Layouts>
CUTE_HOST_DEVICE constexpr
auto
coshape(ZipLayout<Layouts...> const& layouts) {
return cute::transform(layouts.layouts_, [&](auto t){ return coshape<Is...>(t); });
}
template <int... Is, class... Layouts>
CUTE_HOST_DEVICE constexpr
auto
cosize(ZipLayout<Layouts...> const& layouts)
{
return size(coshape<Is...>(layouts));
}
template <class... Layouts>
CUTE_HOST_DEVICE constexpr
auto
nullspace(ZipLayout<Layouts...> const& layouts) {
return cute::fold(layouts.layouts_, make_layout(size(layouts)),
[](auto null, auto layout) {
return composition(null, nullspace(composition(layout, null)));
});
}
} // end namespace cute