[CuTe] [Xe] Allow general layouts in make_subgroup_tensor

petercad · petercad · commit eb7ea6b7344b · 2025-10-22T20:32:53.000-07:00
diff --git a/include/cute/tensor_sg.hpp b/include/cute/tensor_sg.hpp
@@ -115,7 +115,11 @@ CUTE_HOST_DEVICE
 constexpr auto
 make_subgroup_tensor(Layout<Shape,Stride> const& sg_layout)
 {
-  return make_subgroup_tensor(make_fragment_like<T>(sg_layout(0,_)), sg_layout);
+  using _SG = intel::_SGSize;
+  auto ilayout = make_layout(make_shape(_SG{}, size(sg_layout) / _SG{}),
+                             make_stride(_1{}, _16{}));
+  auto sv_layout = sg_layout.compose(ilayout);
+  return make_subgroup_tensor(make_fragment_like<T>(sv_layout(0,_)), sv_layout);
 }
 
 template <typename T, class... Args>

Original file line number	Diff line number	Diff line change
`@@ -115,7 +115,11 @@ CUTE_HOST_DEVICE`
`115`	`115`	`constexpr auto`
`116`	`116`	`make_subgroup_tensor(Layout<Shape,Stride> const& sg_layout)`
`117`	`117`	`{`
`118`		`- return make_subgroup_tensor(make_fragment_like<T>(sg_layout(0,_)), sg_layout);`
	`118`	`+ using _SG = intel::_SGSize;`
	`119`	`+ auto ilayout = make_layout(make_shape(_SG{}, size(sg_layout) / _SG{}),`
	`120`	`+ make_stride(_1{}, _16{}));`
	`121`	`+ auto sv_layout = sg_layout.compose(ilayout);`
	`122`	`+ return make_subgroup_tensor(make_fragment_like<T>(sv_layout(0,_)), sv_layout);`
`119`	`123`	`}`
`120`	`124`
`121`	`125`	`template <typename T, class... Args>`