[NFC] Separate back memdesc reshape encoding inference (#7544)

This fixes extra error message emitted for non error kind of failure.
This avoids having op errors being emitted when the memdesc pattern
fails.

Author: ThomasRaoux

lib/Dialect/TritonGPU/IR/Dialect.cpp

@@ -2469,39 +2469,6 @@ struct TritonGPUInferLayoutInterface
                                              Attribute srcEnc,
                                              ArrayRef<int64_t> dstShape,
                                              Attribute &dstEnc) const {
-    if (auto mmaEncoding = dyn_cast<NVMMASharedEncodingAttr>(srcEnc)) {
-      // TODO: supporting reshape of CTA layouts is non-trivial.
-      if (getNumCTAs(mmaEncoding) > 1)
-        return failure();
-      int innerDimDst =
-          mmaEncoding.getTransposed() ? dstShape.front() : dstShape.back();
-      int innerDimSrc =
-          mmaEncoding.getTransposed() ? srcShape.front() : srcShape.back();
-      // For now disallow reshape of the inner dimension.
-      if (innerDimDst != innerDimSrc)
-        return failure();
-      auto *ctx = srcEnc.getContext();
-
-      // CTALayout can be all 1's because we bailed on multi-CTA layouts above.
-      auto CTALayout = CTALayoutAttr::get(
-          ctx,
-          /*CTAsPerCGA=*/SmallVector<unsigned>(dstShape.size(), 1),
-          /*CTASplitNum=*/SmallVector<unsigned>(dstShape.size(), 1),
-          /*CTAOrder=*/llvm::to_vector(llvm::seq<unsigned>(dstShape.size())));
-      dstEnc = NVMMASharedEncodingAttr::get(
-          ctx, mmaEncoding.getSwizzlingByteWidth(), mmaEncoding.getTransposed(),
-          mmaEncoding.getElementBitWidth(), mmaEncoding.getFp4Padded(),
-          CTALayout);
-      // Big guns, check linear layouts are equivalent
-      // We disallow reshaping memdesc_subviews in the verifier
-      // We disallow reshaping memdesc_subviews in the verifier
-      auto srcLL = toLinearLayout(srcShape, srcEnc, srcShape);
-      auto dstLL = toLinearLayout(dstShape, dstEnc, dstShape);
-      if (reshapeLayout(ctx, srcLL, dstShape) != dstLL) {
-        return failure();
-      }
-      return success();
-    }
     auto src = mlir::dyn_cast<BlockedEncodingAttr>(srcEnc);
     if (!src) {
       return failure();

lib/Dialect/TritonGPU/IR/Ops.cpp

@@ -494,6 +494,46 @@ LogicalResult MemDescReshapeOp::verify() {
   return success();
 }
 
+static LogicalResult inferMemDescReshapeOpEncoding(ArrayRef<int64_t> srcShape,
+                                                   Attribute srcEnc,
+                                                   ArrayRef<int64_t> dstShape,
+                                                   Attribute &dstEnc) {
+  if (auto mmaEncoding = dyn_cast<NVMMASharedEncodingAttr>(srcEnc)) {
+    // TODO: supporting reshape of CTA layouts is non-trivial.
+    if (getNumCTAs(mmaEncoding) > 1)
+      return failure();
+    int innerDimDst =
+        mmaEncoding.getTransposed() ? dstShape.front() : dstShape.back();
+    int innerDimSrc =
+        mmaEncoding.getTransposed() ? srcShape.front() : srcShape.back();
+    // For now disallow reshape of the inner dimension.
+    if (innerDimDst != innerDimSrc)
+      return failure();
+    auto *ctx = srcEnc.getContext();
+
+    // CTALayout can be all 1's because we bailed on multi-CTA layouts above.
+    auto CTALayout = CTALayoutAttr::get(
+        ctx,
+        /*CTAsPerCGA=*/SmallVector<unsigned>(dstShape.size(), 1),
+        /*CTASplitNum=*/SmallVector<unsigned>(dstShape.size(), 1),
+        /*CTAOrder=*/llvm::to_vector(llvm::seq<unsigned>(dstShape.size())));
+    dstEnc = NVMMASharedEncodingAttr::get(
+        ctx, mmaEncoding.getSwizzlingByteWidth(), mmaEncoding.getTransposed(),
+        mmaEncoding.getElementBitWidth(), mmaEncoding.getFp4Padded(),
+        CTALayout);
+    // Big guns, check linear layouts are equivalent
+    // We disallow reshaping memdesc_subviews in the verifier
+    // We disallow reshaping memdesc_subviews in the verifier
+    auto srcLL = toLinearLayout(srcShape, srcEnc, srcShape);
+    auto dstLL = toLinearLayout(dstShape, dstEnc, dstShape);
+    if (reshapeLayout(ctx, srcLL, dstShape) != dstLL) {
+      return failure();
+    }
+    return success();
+  }
+  return failure();
+}
+
 LogicalResult MemDescReshapeOp::inferReturnTypes(
     MLIRContext *context, std::optional<Location> loc, MemDescType srcTy,
     ArrayRef<int64_t> dstShape, MemDescType &inferredReturnType) {
@@ -503,9 +543,8 @@ LogicalResult MemDescReshapeOp::inferReturnTypes(
 
   Attribute dstEncoding;
   if (Attribute srcEnc = srcTy.getEncoding()) {
-    auto *inferLayout = cast<DialectInferLayoutInterface>(&srcEnc.getDialect());
-    if (failed(inferLayout->inferReshapeOpEncoding(srcTy.getShape(), srcEnc,
-                                                   dstShape, dstEncoding, loc)))
+    if (failed(inferMemDescReshapeOpEncoding(srcTy.getShape(), srcEnc, dstShape,
+                                             dstEncoding)))
       return failure();
   }