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

[Jianhui Li]

================
@@ -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(
----------------
Jianhui-Li wrote:

I am not sure that these two abstraction are needed: addContextAwareVectorTypeConversion() and addSCFStructuralMaterializations().

Other than the naming issue, it breaks the readability flow of the overall code structure here in the use points. Here everything is a sequence of converter.addSomthing() before you add these two functions, which is easier to follow. But these two functions breaks the protocol, and so I have to look back and forth and figure out what the code is doing. Can we keep the high level convert.addSomthing() structure while adding utility if there is reuse?   

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