Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[LinearLayouts] Faster pext algorithm #5621

Merged
merged 4 commits into from
Jan 15, 2025
Merged

[LinearLayouts] Faster pext algorithm #5621

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
d3d1d32
[LinearLayouts] Faster pext algorithm
lezcano Jan 15, 2025
df66fe4
Skip HIP for Reduction(LinearLayouts)
lezcano Jan 15, 2025
ee37a63
Change the only non-AMD use of getWarpSize and getNumWarpsPerCTA
lezcano Jan 15, 2025
3596620
not this function the other one
lezcano Jan 15, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
33 changes: 22 additions & 11 deletions lib/Conversion/TritonGPUToLLVM/Utility.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,11 +12,18 @@
// from https://gist.github.com/pps83/3210a2f980fd02bb2ba2e5a1fc4a2ef0
#include <intrin.h>

static int __builtin_clz(unsigned x) {
unsigned long r;
_BitScanReverse(&r, x);
return static_cast<int>(r);
}

static int __builtin_ctz(unsigned x) {
unsigned long r;
_BitScanForward(&r, x);
return static_cast<int>(r);
}

#endif

namespace mlir {
Expand Down Expand Up @@ -601,18 +608,22 @@ Value pext_i32(RewriterBase &rewriter, Location loc, Value a, uint32_t mask) {
if (mask == 0xFFFFFFFF)
return a;

// We implement a blocked algorithm to avoid generating too many instructions
// Implements the blocked algorithm from
// https://forums.developer.nvidia.com/t/pdep-and-pext-functionality-for-cuda/270973
uint32_t mskConst = mask;
uint32_t extcnt = 0;
Value result = i32_val(0);
int resultPos = 0;
while (mask) {
int start = __builtin_ctz(mask);
int width = __builtin_ctz(~(mask >> start));
Value shifted = lshr(a, i32_val(start));
Value widthMask = i32_val(((1u << width) - 1));
Value blockVal = and_(shifted, widthMask);
result = or_(result, shl(blockVal, i32_val(resultPos)));
resultPos += width;
mask &= ~(((1u << width) - 1) << start);
while (mskConst) {
uint32_t oldmsk = mskConst;
uint32_t bitgrplsb = mskConst & (-mskConst);
mskConst &= bitgrplsb + mskConst;
uint32_t bitgrp = mskConst ^ oldmsk;
uint32_t lsbpos = 31 - __builtin_clz(bitgrplsb);
// like popcount for a number 0..01..1..0 but portable
uint32_t grplen = __builtin_ctz(~(bitgrp >> lsbpos));
uint32_t shift = lsbpos - extcnt;
extcnt += grplen;
result = or_(result, lshr(and_(i32_val(bitgrp), a), i32_val(shift)));
}
return result;
}
Expand Down
3 changes: 2 additions & 1 deletion lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -447,7 +447,8 @@ class DecomposeScaledBlocked
mmaEnc.getInstrShape()[versionMajor == 3
? 0
: mmaEnc.getInstrShape().size() - 2];
auto warpSize = getWarpSize(newAEncoding);
auto mod = scaledDotOp->getParentOfType<ModuleOp>();
int warpSize = triton::gpu::TritonGPUDialect::getThreadsPerWarp(mod);
assert(instrShapeM <= warpSize);
// Necessary choice to leave all the scales of the tile in that given warp
auto threadsPerWarp =
Expand Down
2 changes: 2 additions & 0 deletions python/test/unit/language/test_core.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2793,6 +2793,8 @@ def test_reduce_layouts(M, N, src_layout, axis, epilogue_kind, dtype_str, add_ov
pytest.skip("Skipping test because it runs out of shared memory")
if reduce_op == "sum" and dtype_str == "float16" and M * N > 1024:
pytest.skip("Skipping sum reduction on float16 due to accuracy issues")
if is_hip() and isinstance(src_layout, LinearLayout):
pytest.skip("FIXME: LinearLayout not supported on HIP")
Comment on lines +2796 to +2797
Copy link
Contributor Author

@lezcano lezcano Jan 15, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

cc @antiagainst regarding the HIP skip. See also the warpSize / numWarp comment in the OP.


if isinstance(src_layout, MmaLayout) and src_layout.version == 3:
src_layout[2] = 16 if dtype_str == "float16" else 8
Expand Down
Loading