[BACKEND] Implement generic code to allow for dot_scaled(mmav3) and warp choices (triton-lang#5103)

Even though this refactor allows for sharing the code for choosing warps
and mma version with the regular `DotOp`, we don't activate any of the
two paths yet, given that we still have to fix a couple things for them
to work.

Putting this preliminary PR up to avoid packing too many things in one
PR

Author: lezcano

lib/Conversion/TritonGPUToLLVM/ConvertLayoutOpToLLVM.cpp

@@ -371,10 +371,9 @@ struct ConvertLayoutOpUsingLinearLayoutsConversion
     auto srcTy = op.getSrc().getType();
     auto dstTy = op.getType();
 
-    // TODO (Keren): Currently, we handle general mma/blocked/slice ->
-    // mma/blocked/slice conversions.
-    // The following tasks must be completed before we can remove the layoutIsOK
-    // check:
+    // TODO (Keren): Currently, we handle general mma/blocked/slice/dot(ampere)
+    // -> mma/blocked/slice/dot(ampere) conversions. The following tasks must be
+    // completed before we can remove the layoutIsOK check:
     // 1. Support for AMD's MFMA and WMMA
     std::function<bool(Attribute)> layoutIsOK = [&](Attribute layout) {
       if (auto nvidiaMma = dyn_cast<NvidiaMmaEncodingAttr>(layout)) {

lib/Conversion/TritonGPUToLLVM/MemoryOpToLLVM.cpp

@@ -140,8 +140,11 @@ struct LocalLoadOpConversion : public ConvertOpToLLVMPattern<LocalLoadOp> {
   // Do for all DotOperandEncodingAttr once we have LLs for all of them
   static bool isSupportedDotOpLayout(Attribute layout) {
     if (auto dot = dyn_cast<DotOperandEncodingAttr>(layout)) {
+      // Use when the SharedToDotOperandMMAv2OrV3 is known to be buggy:
+      // - kWidth == 8
       if (auto mma = dyn_cast<NvidiaMmaEncodingAttr>(dot.getParent())) {
-        return mma.isAmpere() && dot.getKWidth() == 8;
+        bool legacyLoweringIsBuggy = dot.getKWidth() >= 8;
+        return legacyLoweringIsBuggy && mma.isAmpere();
       }
       if (isa<AMDMfmaEncodingAttr>(dot.getParent()))
         return true;

lib/Dialect/TritonGPU/IR/Ops.cpp

@@ -52,14 +52,23 @@ LogicalResult UpcastMXFPOp::verify() {
         "all dimensions except the last must match between operands");
   }
 
-  auto dotEncoding =
-      dyn_cast_or_null<DotOperandEncodingAttr>(xTy.getEncoding());
+  auto layoutX = xTy.getEncoding();
+  auto layoutScale = scaleTy.getEncoding();
+  if (bool(layoutX) != bool(layoutScale)) {
+    return emitOpError(
+        "Expected either both or neither operands to have an encoding");
+  }
+  // Nothing to check if no encoding. This is used to infer the return type in
+  // AccelerateMatmul.cpp
+  if (!layoutX) {
+    return success();
+  }
+
+  auto dotEncoding = dyn_cast<DotOperandEncodingAttr>(layoutX);
   if (!dotEncoding) {
     return emitOpError("Expected a DotOperandEncodingAttr for values");
   }
-
-  auto blockedScale =
-      dyn_cast_or_null<BlockedEncodingAttr>(scaleTy.getEncoding());
+  auto blockedScale = dyn_cast<BlockedEncodingAttr>(layoutScale);
   if (!blockedScale) {
     return emitOpError("Expected a BlockOperandEncoding for scales");
   }
@@ -86,22 +95,23 @@ LogicalResult UpcastMXFPOp::inferReturnTypes(
   auto xShape = xTy.getShape();
 
   auto encoding = xTy.getEncoding();
-  if (!encoding) {
-    return emitOptionalError(loc, "expected an encoding");
-  }
-  if (!mlir::isa<DotOperandEncodingAttr>(encoding)) {
-    return emitOptionalError(loc, "expected a dotOperand encoding");
-  }
 
   if (typeEncoded == ScaleDotElemType::E2M1) {
-    auto oldEncoding = cast<DotOperandEncodingAttr>(encoding);
-    auto newVEncoding = DotOperandEncodingAttr::get(
-        ctx, oldEncoding.getOpIdx(), oldEncoding.getParent(),
-        oldEncoding.getKWidth() * 2);
+    RankedTensorType retTy;
+
     auto newShape = SmallVector<int64_t>(xShape);
     newShape.back() *= 2;
-    inferredReturnTypes.push_back(
-        RankedTensorType::get(newShape, FloatType::getBF16(ctx), newVEncoding));
+    if (!encoding) {
+      retTy = RankedTensorType::get(xShape, FloatType::getBF16(ctx));
+    } else {
+      auto oldEncoding = cast<DotOperandEncodingAttr>(encoding);
+      auto newVEncoding = DotOperandEncodingAttr::get(
+          ctx, oldEncoding.getOpIdx(), oldEncoding.getParent(),
+          oldEncoding.getKWidth() * 2);
+      retTy = RankedTensorType::get(newShape, FloatType::getBF16(ctx),
+                                    newVEncoding);
+    }
+    inferredReturnTypes.push_back(retTy);
   } else {
     inferredReturnTypes.push_back(xTy);
   }