Reland the block store lowering changes. #4646
Conversation
chengjunlu
force-pushed
the
chengjun/re_land_block_store
branch
from
261e487 to
ed0a771
Compare
third_party/intel/lib/TritonIntelGPUToLLVM/LoadStoreOpToLLVM.cpp
Outdated
Show resolved
Hide resolved
third_party/intel/lib/TritonIntelGPUToLLVM/LoadStoreOpToLLVM.cpp
Outdated
Show resolved
Hide resolved
chengjunlu
force-pushed
the
chengjun/re_land_block_store
branch
3 times, most recently
from
fac026c to
26ab3b4
Compare
Signed-off-by: Lu,Chengjun <[email protected]>
|
This PR looks too large. We will work on slitting it. |
There was a problem hiding this comment.
Pull Request Overview
This PR relands the 2D block store lowering changes, fixes a block store issue, and enhances StoreOp conversion with LinearLayout-based tile size computation.
- Introduces
getBlockIOTileSizeand extendsisMemoryRowMajorto handleStoreOp - Unifies and refactors StoreOp lowering into a single
matchAndRewrite, supporting both block pointers and regular pointers via a newMatrix2DBlockStoreOp - Updates tests (
blockptr_store.mlir) for expected 2D block store patterns and adds new kernel cases; adds a new env var inGetEnv.hpp
Reviewed Changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 4 comments.
| File | Description |
|---|---|
LoadStoreOpToLLVM.cpp |
Added includes, getBlockIOTileSize, extended StoreOp lowering |
test/TritonIntelGPU/blockptr_store.mlir |
Updated expected checks, added new block store test scenarios |
include/triton/Tools/Sys/GetEnv.hpp |
Added TRITON_INTEL_ENABLE_BLOCK_IO_STORE_ON_REGULAR_PTR to set |
| llvm_unreachable(("Could not find the input dim:" + inDim + | ||
| ", on the ll:" + ll.toString()) | ||
| .c_str()); |
There was a problem hiding this comment.
Passing c_str() of a temporary std::string into llvm_unreachable risks a dangling pointer. Construct the message in a local std::string variable or use llvm::formatv to ensure the string data remains valid.
| llvm_unreachable(("Could not find the input dim:" + inDim + | |
| ", on the ll:" + ll.toString()) | |
| .c_str()); | |
| std::string errorMessage = "Could not find the input dim:" + inDim + | |
| ", on the ll:" + ll.toString(); | |
| llvm_unreachable(errorMessage.c_str()); |
| size_t rank = tensorShape.size(); | ||
| // 2D block store supports 64 bytes per row at most. | ||
| unsigned totalBytesPerRowPerMatrix = tileWidth * elemSizeInBits / 8; | ||
| if (totalBytesPerRowPerMatrix > 64) |
There was a problem hiding this comment.
[nitpick] The literal 64 appears as a magic number here. Consider defining a named constant (e.g. kMaxBytesPerRow) to clarify its meaning and ease future updates.
| if (totalBytesPerRowPerMatrix > 64) | |
| if (totalBytesPerRowPerMatrix > kMaxBytesPerRow) |
|
|
||
| // Only lower StoreOp with dpas layout encoding. | ||
| if (!hasDpasEncoding(tensorType)) | ||
| matchAndRewrite(triton::StoreOp op, OpAdaptor adaptor, |
There was a problem hiding this comment.
[nitpick] This matchAndRewrite function has grown quite large. Consider refactoring block-pointer and regular-pointer handling into separate helper methods to improve readability and testability.
| assert(!llMask && "The masks is expected to be used with regular tensor " | ||
| "pointer type, but got a block pointer type."); |
There was a problem hiding this comment.
Using assert for input validation can crash in release builds. Prefer returning failure() with an informative diagnostic rather than aborting, so consumers get a graceful rewrite failure.
| assert(!llMask && "The masks is expected to be used with regular tensor " | |
| "pointer type, but got a block pointer type."); | |
| if (llMask) { | |
| rewriter.emitError(loc, "The masks are expected to be used with regular tensor " | |
| "pointer type, but got a block pointer type."); | |
| return failure(); | |
| } |
|
I created a new PR #4666 to enable the block store for regular pointer first. |
|
Created #4667 for some easier changes. |
Reland the block store lowering changes with fix of an issue in block store.