[Mlir-commits] [mlir] [MLIR][XeGPU] Use context-aware type converter in WgToSgDistribute and Blocking pass (PR #194685)

[Nishant Patel]

================
@@ -380,51 +383,80 @@ void XeGPUBlockingPass::runOnOperation() {
     return std::make_pair(tileShape, count);
   };
 
-  // Perform type conversion for SCF control folow ops
-  TypeConverter converter;
-  converter.addConversion([](Type type) -> Type { return type; });
-  converter.addConversion(
-      [&](RankedTensorType type,
-          SmallVectorImpl<Type> &result) -> std::optional<LogicalResult> {
-        Type elemTy = type.getElementType();
-        ArrayRef<int64_t> shape = type.getShape();
-
-        auto layout =
-            llvm::dyn_cast_if_present<xegpu::LayoutAttr>(type.getEncoding());
-        if (layout && layout.isForWorkgroup())
-          return failure();
-
-        int count;
-        SmallVector<int64_t> subShape;
-        std::tie(subShape, count) = getTileShapeAndCount(shape, layout);
-        auto newTy = VectorType::get(subShape, elemTy);
-        result.append(count, newTy);
-        return success();
-      });
-  converter.addConversion(
-      [&](xegpu::TensorDescType type,
-          SmallVectorImpl<Type> &result) -> std::optional<LogicalResult> {
-        Type elemTy = type.getElementType();
-        ArrayRef<int64_t> shape = type.getShape();
-
-        xegpu::DistributeLayoutAttr layout = type.getLayoutAttr();
-        if (layout && layout.isForWorkgroup())
-          return failure();
-
-        int count;
-        SmallVector<int64_t> subShape;
-        std::tie(subShape, count) = getTileShapeAndCount(shape, layout);
-
-        if (layout)
-          layout = layout.dropInstData();
-
-        auto newTy = xegpu::TensorDescType::get(
-            type.getContext(), subShape, elemTy, type.getEncoding(), layout);
-        result.append(count, newTy);
-        return success();
-      });
-
-  xegpu::doSCFStructuralTypeConversionWithTensorType(op, converter);
+  // Perform context-aware type conversion for SCF structural ops.
+  // Inspects Values to find inst_data layout information for 1:N conversion.
+  llvm::SmallSetVector<UnrealizedConversionCastOp, 8> existingCasts;
+  op->walk(
+      [&](UnrealizedConversionCastOp castOp) { existingCasts.insert(castOp); });
+
+  {
+    TypeConverter converter;
+    converter.addConversion([](Type type) -> Type { return type; });
+
+    // TensorDescType 1:N converter (type-based, layout is in the type).
+    converter.addConversion(
+        [&](xegpu::TensorDescType type,
+            SmallVectorImpl<Type> &result) -> std::optional<LogicalResult> {
+          Type elemTy = type.getElementType();
+          ArrayRef<int64_t> shape = type.getShape();
+
+          xegpu::DistributeLayoutAttr layout = type.getLayoutAttr();
+          if (layout && layout.isForWorkgroup())
+            return failure();
+
+          int count;
+          SmallVector<int64_t> subShape;
+          std::tie(subShape, count) = getTileShapeAndCount(shape, layout);
+
+          if (layout)
+            layout = layout.dropInstData();
+
+          auto newTy = xegpu::TensorDescType::get(
+              type.getContext(), subShape, elemTy, type.getEncoding(), layout);
+          result.append(count, newTy);
+          return success();
+        });
+
+    // Context-aware 1:N conversion for VectorType based on inst_data.
+    xegpu::addContextAwareVectorTypeConversion(
----------------
nbpatel wrote:

inlined addSCFStructuralMaterializations....renamed addContextAwareVectorTypeConversion to addVectorTypeConversion.....its used in all the three distribution passes hence its abstracted as a utility for reuse

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