Add verification for NVMMASharedEncodingAttr during reshape inference. Implement verifyNVMMASharedEncoding to check compatibility with TMA block shapes and tile dimensions. Integrate this verification into inferMemDescReshapeOpEncoding to ensure that reshapes involving nvmma_shared layouts result in valid encodings. Add a test case to catch issues with invalid reshape outcomes. --- a/lib/Dialect/TritonGPU/IR/Ops.cpp +++ b/lib/Dialect/TritonGPU/IR/Ops.cpp @@ -1,3 +1,5 @@ +#include "llvm/Support/Casting.h" +#include "llvm/Support/LogicalResult.h" #include "mlir/IR/BuiltinTypes.h" #include "mlir/IR/Diagnostics.h" #include "mlir/Support/DebugStringHelper.h" @@ -9,6 +11,7 @@ #include "triton/Dialect/TritonGPU/IR/Types.h" #include "triton/Dialect/TritonGPU/Transforms/Utility.h" #include "triton/Dialect/TritonNvidiaGPU/IR/Dialect.h" +#include "triton/Dialect/TritonNvidiaGPU/Transforms/TMAUtilities.h" #include "triton/Tools/LayoutUtils.h" #include "llvm/Support/Casting.h" #include "llvm/Support/LogicalResult.h" @@ -571,7 +574,45 @@ LogicalResult MemDescReshapeOp::verify() { return OpTrait::impl::verifyEquivalentMemDescType(expectedTy, dstType); } -static LogicalResult inferMemDescReshapeOpEncoding(ArrayRef srcShape, +// Verification copied from nvmmaSharedToLinearLayout(). +LogicalResult verifyNVMMASharedEncoding(std::optional loc, + NVMMASharedEncodingAttr attr, + ArrayRef shape, + int elementBitWidth) { + if (attr.getSwizzlingByteWidth() == 0) return success(); + if (shape.size() < 2) + return emitOptionalError(loc, "nvmma_shared encoding requires rank >= 2"); + + auto shapePerCTA = getShapePerCTA(attr, shape); + auto tmaShape = triton::nvidia_gpu::getTMABlockShape(attr, shapePerCTA, + /*packedSize=*/true, + gpu::TMAMode::Tiled); + std::array collapsedTmaShape{1, tmaShape.back()}; + for (int i = 0; i + 1 < shape.size(); i++) + collapsedTmaShape[0] *= tmaShape[i]; + if (attr.getTransposed()) { + std::swap(collapsedTmaShape[0], collapsedTmaShape[1]); + } + + int tileRows = 8; + int tileCols = 8 * attr.getSwizzlingByteWidth() / elementBitWidth; + if (attr.getFp4Padded()) tileCols /= 2; + + int packingFactor = attr.getFp4Padded() ? 2 : 1; + if (collapsedTmaShape[1] * packingFactor < tileCols || + collapsedTmaShape[0] < tileRows) { + return emitOptionalError( + loc, + "Illegal shared layout; expected collapsed shapePerCTA to " + "be at least [", + tileRows, ", ", (tileCols / packingFactor), "], collapsedTmaShape: [", + collapsedTmaShape[0], ", ", collapsedTmaShape[1], "]"); + } + return success(); +} + +static LogicalResult inferMemDescReshapeOpEncoding(std::optional loc, + ArrayRef srcShape, Attribute srcEnc, ArrayRef dstShape, Attribute &dstEnc) { @@ -592,6 +633,11 @@ static LogicalResult inferMemDescReshapeOpEncoding(ArrayRef srcShape, ctx, mmaEncoding.getSwizzlingByteWidth(), mmaEncoding.getTransposed(), mmaEncoding.getElementBitWidth(), mmaEncoding.getFp4Padded(), CGALayout); + if (failed(verifyNVMMASharedEncoding( + loc, cast(candidateEncoding), dstShape, + mmaEncoding.getElementBitWidth()))) { + return failure(); + } auto srcLL = toLinearLayout(srcShape, srcEnc); auto dstLL = toLinearLayout(dstShape, candidateEncoding); if (reshapeLayout(ctx, srcLL, dstShape) == dstLL) { @@ -631,8 +677,8 @@ LogicalResult MemDescReshapeOp::inferReturnTypes( Attribute dstEncoding; if (Attribute srcEnc = srcTy.getEncoding()) { - if (failed(inferMemDescReshapeOpEncoding(srcTy.getShape(), srcEnc, dstShape, - dstEncoding))) + if (failed(inferMemDescReshapeOpEncoding(loc, srcTy.getShape(), srcEnc, + dstShape, dstEncoding))) return failure(); }