MCPcopy Create free account
hub / github.com/NVIDIA/cutlass / batch_reduce

Function batch_reduce

include/cute/algorithm/tensor_reduce.hpp:62–75  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

60 class DstEngine, class DstLayout,
61 class BinaryOp = cute::plus>
62CUTE_HOST_DEVICE constexpr
63void
64batch_reduce(Tensor<SrcEngine, SrcLayout> const& src, // (RedMode, BatchMode)
65 Tensor<DstEngine, DstLayout> & dst, // (BatchMode)
66 BinaryOp op = {})
67{
68 // Precondition
69 CUTE_STATIC_ASSERT_V(rank(src) == Int<2>{});
70 assert(size<1>(src) == size(dst));
71
72 for (int i = 0; i < size(dst); ++i) {
73 dst(i) = reduce(src(_,i), dst(i), op);
74 }
75}
76
77
78// Reduce @src tensor along selected modes specified in @target_profile using binary reduction operator @op

Callers 2

TESTFunction · 0.85
logical_reduceFunction · 0.85

Calls 3

reduceFunction · 0.85
rankFunction · 0.50
sizeFunction · 0.50

Tested by

no test coverage detected