[Mlir-commits] [mlir] [MLIR][NVVM] Update TMA tensor prefetch Op (PR #153464)

[Durgadoss R]

================
@@ -99,10 +99,41 @@ LogicalResult CpAsyncOp::verify() {
 }
 
 LogicalResult CpAsyncBulkTensorPrefetchOp::verify() {
-  size_t numIm2ColOffsets = getIm2colOffsets().size();
-  bool isIm2Col = numIm2ColOffsets > 0;
-  return cpAsyncBulkTensorCommonVerifier(getCoordinates().size(), isIm2Col,
-                                         numIm2ColOffsets, getLoc());
+  size_t tensorDims = getCoordinates().size();
+  if (tensorDims < 1 || tensorDims > 5)
+    return emitError("expects coordinates between 1 to 5 dimension");
+
+  auto checkTMALoadParams = [&](TMALoadMode mode, bool isIm2col,
+                                size_t expectedNumOffsets) -> LogicalResult {
+    if (isIm2col && (tensorDims < 3))
+      return emitError()
+             << "to use " << stringifyEnum(mode)
+             << " mode, the tensor has to be at least 3-dimensional";
+
+    if (getIm2colOffsets().size() != expectedNumOffsets)
+      return emitError() << " im2col offsets expected " << expectedNumOffsets
+                         << " (provided " << getIm2colOffsets().size() << ")";
+
+    return success();
+  };
+
+  auto mode = getMode();
+  switch (mode) {
+  case TMALoadMode::TILE:
+    return checkTMALoadParams(mode, false, 0);
+  case TMALoadMode::IM2COL:
+    return checkTMALoadParams(mode, true, tensorDims - 2);
+  case TMALoadMode::IM2COL_W:
+  case TMALoadMode::IM2COL_W_128:
+    return checkTMALoadParams(mode, true, 2);
+  case TMALoadMode::TILE_GATHER4:
+    return (tensorDims == 5) ? checkTMALoadParams(mode, false, 0)
+                             : emitError("Gather4 mode expects 5 coordinates");
+  default:
+    return emitError("Invalid LoadMode in CpAsyncBulkTensorPrefetchOp.");
+  }
----------------
durga4github wrote:

The intention is to make it common for both the TMA Load and TMA Prefetch Ops.
For now, I have kept it here. I plan to make it common with my TMA Load Op changes. (since there will be more than one usages of this function then).

https://github.com/llvm/llvm-project/pull/153464