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

[Jianhui Li]

================
@@ -986,3 +842,169 @@ bool xegpu::matchSplitDimExpansion(
   }
   return srcIdx == src.size();
 }
+
+//===----------------------------------------------------------------------===//
+// Context-aware type conversion utilities
+//===----------------------------------------------------------------------===//
+
+void xegpu::addSCFStructuralMaterializations(TypeConverter &converter) {
+  auto materializeCast = [](OpBuilder &builder, Type type, ValueRange inputs,
+                            Location loc) -> Value {
+    return UnrealizedConversionCastOp::create(builder, loc, type, inputs)
+        .getResult(0);
+  };
+  // Source materialization: N:1 (N converted values -> 1 original value).
+  converter.addSourceMaterialization(materializeCast);
+  // Target materialization: 1:1 (single value type conversion).
+  converter.addTargetMaterialization(materializeCast);
+}
+
+void xegpu::addContextAwareVectorTypeConversion(
+    TypeConverter &converter, Operation *topLevelOp,
+    SubShapeAndCountFn getSubShapeAndCount) {
+  // Pre-compute 1:N type mappings for scf.while block arguments only.
+  // During scf.while structural conversion, blocks are detached from their
+  // parent region before convertBlockSignature is called. Block::getParent()
+  // crashes on detached blocks (LLVM ilist assertion), so we cannot look up
+  // layout attributes at that point. Other SCF ops (scf.for, scf.if) keep
+  // blocks attached during conversion.
+  auto whileArgTypeMap = std::make_shared<DenseMap<Value, SmallVector<Type>>>();
+  auto recordBlockArgTypes = [&](Value init, BlockArgument arg) {
+    auto vecTy = dyn_cast<VectorType>(init.getType());
+    if (!vecTy)
+      return;
+    auto layout = xegpu::getDistributeLayoutAttr(init);
+    if (!layout)
+      return;
+    auto [subShape, count] = getSubShapeAndCount(vecTy, layout);
+    if (count <= 0)
+      return;
----------------
Jianhui-Li wrote:

It is possible count is 0? Maybe should assert here or inside the getSubShapeAndCount. 

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