github/ROCm
1
1
Fork 0
mirror of https://github.com/ROCm/ROCm.git synced 2026-04-05 03:01:17 -04:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #268 from ROCmSoftwarePlatform/improve_reduce_for_fa

Browse Source 
[CHERRY-PICKED FROM UPSTREAM][BACKEND] no longer uses shared mem or barriers for single-warp reductions (openai#1915)
This commit is contained in:
jayfurmanek
2023-08-21 13:29:11 -05:00
committed by GitHub
parent d86b19f7a3 d0b7793935
commit fa429316d4
6 changed files with 83 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
lib/Analysis/Utility.cpp
View File

@@ -60,9 +60,10 @@ SmallVector<SmallVector<unsigned>> ReduceOpHelper::getScratchConfigsFast() {
auto argLayout = getSrcLayout();
auto argLayoutMma = argLayout.dyn_cast<triton::gpu::MmaEncodingAttr>();
// if (argLayoutMma && argLayoutMma.getVersionMajor() == 2 &&
// triton::gpu::getWarpsPerCTA(argLayout)[axis] == 1)
// return {{1, 1}, {1, 1}};
// that case doesn't need inter-warp communication
if (isFastReduction() && triton::gpu::getWarpsPerCTA(argLayout)[axis] == 1)
return {{0, 0}, {0, 0}};
/// shared memory block0
smemShapes[0] = convertType<unsigned>(getSrcShape());