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

Add conversion from ElementBias to ElementCompute (#961)

Browse Source
This commit is contained in:
Jack Kosaian
2023-05-26 23:08:36 -04:00
committed by GitHub
parent 6f47420213
commit 7dbf423763
4 changed files with 203 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
include/cutlass/epilogue/collective/sm90_epilogue_tma_warpspecialized_bias_elementwise.hpp
View File

@@ -475,6 +475,12 @@ public:
Tensor tRS_rT_frg = recast<typename ThreadEpilogueOp::FragmentT>(tRS_rT);
Tensor tRS_rBias_frg = recast<typename ThreadEpilogueOp::FragmentBias>(tRS_rBias);
// thread::LinearCombinationBiasElementwise expects that the bias passed in is of
// type ElementCompute. Therefore, conversion from type ElementBias to ElementCompute
// is needed before calling the thread-level epilogue.
cutlass::NumericArrayConverter<ElementCompute, ElementBias,
ThreadEpilogueOp::FragmentBias::kElements> bias_converter;
// Partition for smem to register copy (tSR_)
TiledCopy tiled_s2r = make_tiled_copy_S(Copy_Atom<CopyOpS2R,InternalElementC>{}, tiled_r2s);
ThrCopy thread_s2r = tiled_s2r.get_slice(thread_idx);
@@ -538,13 +544,15 @@ public:
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(tRS_rD_frg); ++i) {
epilogue_op(tRS_rD_frg(i), tRS_rT_frg(i), tRS_rAcc_frg_mn(r2s_v + i), tRS_rC_frg(i), tRS_rBias_frg(i));
typename ThreadEpilogueOp::FragmentCompute converted_bias = bias_converter(tRS_rBias_frg(i));
epilogue_op(tRS_rD_frg(i), tRS_rT_frg(i), tRS_rAcc_frg_mn(r2s_v + i), tRS_rC_frg(i), converted_bias);
}
}
else {
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(tRS_rD_frg); ++i) {
epilogue_op(tRS_rD_frg(i), tRS_rT_frg(i), tRS_rAcc_frg_mn(r2s_v + i), tRS_rBias_frg(i));
typename ThreadEpilogueOp::FragmentCompute converted_bias = bias_converter(tRS_rBias_frg(i));
epilogue_op(tRS_rD_frg(i), tRS_rT_frg(i), tRS_rAcc_frg_mn(r2s_v + i), converted_bias);
}
}

2
include/cutlass/epilogue/thread/linear_combination_bias_elementwise.h
View File

@@ -95,7 +95,7 @@ public:
using FragmentSource = FragmentC;
using FragmentOutput = FragmentZ;
using ElementBias = ElementVector;
using FragmentBias = FragmentCompute;
using FragmentBias = Array<ElementBias, kElementsPerAccess>;
using ActivationFunctor = ElementwiseOp;
static const ScaleType::Kind kScale = ScaleType::Default;