[Mlir-commits] [mlir] [MLIR][XeGPU] Add unrolling/blocking support for 3D+ batched operations (PR #201725)

[Charitha Saumya]

================
@@ -44,36 +44,57 @@ resolveUnrealizedConversionCastOp(UnrealizedConversionCastOp castOp) {
   ValueRange inputs = castOp.getInputs();
   ValueRange outputs = castOp.getOutputs();
 
-  auto hasIdenticalVectorTypes = [](ValueRange values) {
+  auto hasIdenticalVectorOrTdescTypes = [](ValueRange values) {
     auto types = values.getTypes();
     return llvm::all_of(types, [&](Type type) {
-      return isa<VectorType>(type) && type == types.front();
+      return (isa<VectorType>(type) || isa<xegpu::TensorDescType>(type)) &&
+             type == types.front();
     });
   };
 
   // We only interest in the case where all inputs and outputs have the
   // identical VectorTypes
-  if (!hasIdenticalVectorTypes(inputs) || !hasIdenticalVectorTypes(outputs)) {
+  if (!hasIdenticalVectorOrTdescTypes(inputs) ||
+      !hasIdenticalVectorOrTdescTypes(outputs)) {
     LDBG() << "skip unrealized conversion cast op not emulating pack/unpack.";
     return;
   }
 
   VectorType outputTy = dyn_cast<VectorType>(outputs[0].getType());
-  OpBuilder builder(castOp);
-  if (inputs.size() > 1 && outputs.size() == 1) {
-    // the castOp is emulating an unpack op
-    ArrayRef<int64_t> shape = outputTy.getShape();
-    Value result = xegpu::createVectorWithShapeFromValues(
-        builder, castOp.getLoc(), inputs, shape);
-    castOp->replaceAllUsesWith(ValueRange(result));
-    castOp->erase();
-  } else if (castOp.getNumResults() > 1 && castOp.getNumOperands() == 1) {
-    // the castOp is emulating a pack op
-    ArrayRef<int64_t> tileShape = outputTy.getShape();
-    SmallVector<Value> results = xegpu::extractVectorsWithShapeFromValue(
-        builder, castOp.getLoc(), inputs[0], tileShape);
-    castOp->replaceAllUsesWith(results);
-    castOp->erase();
+  if (outputTy) {
+    OpBuilder builder(castOp);
+    if (inputs.size() > 1 && outputs.size() == 1) {
+      // the castOp is emulating an unpack op
+      ArrayRef<int64_t> shape = outputTy.getShape();
+      Value result = xegpu::createVectorWithShapeFromValues(
+          builder, castOp.getLoc(), inputs, shape);
+      castOp->replaceAllUsesWith(ValueRange(result));
+      castOp->erase();
+    } else if (castOp.getNumResults() > 1 && castOp.getNumOperands() == 1) {
+      // the castOp is emulating a pack op
+      ArrayRef<int64_t> tileShape = outputTy.getShape();
+      SmallVector<Value> results = xegpu::extractVectorsWithShapeFromValue(
+          builder, castOp.getLoc(), inputs[0], tileShape);
+      castOp->replaceAllUsesWith(results);
+      castOp->erase();
+    }
+  } else {
----------------
charithaintc wrote:

in what case there is no vector output? is this the tensor desc case? worth a comment on what is handled inside else branch

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