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<int64_t> srcShape,
+// Verification copied from nvmmaSharedToLinearLayout().
+LogicalResult verifyNVMMASharedEncoding(std::optional<Location> loc,
+                                        NVMMASharedEncodingAttr attr,
+                                        ArrayRef<int64_t> 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<int64_t, 2> 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<Location> loc,
+                                                   ArrayRef<int64_t> srcShape,
                                                    Attribute srcEnc,
                                                    ArrayRef<int64_t> dstShape,
                                                    Attribute &dstEnc) {
@@ -592,6 +633,11 @@ static LogicalResult inferMemDescReshapeOpEncoding(ArrayRef<int64_t> srcShape,
             ctx, mmaEncoding.getSwizzlingByteWidth(),
             mmaEncoding.getTransposed(), mmaEncoding.getElementBitWidth(),
             mmaEncoding.getFp4Padded(), CGALayout);
+        if (failed(verifyNVMMASharedEncoding(
+            loc, cast<NVMMASharedEncodingAttr>(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();
   }
 

