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

[Nishant Patel]

================
@@ -986,3 +842,155 @@ 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()
----------------
nbpatel wrote:

> is this called inside the scf structural conversions? worth adding a fix in upstream if this is the case rather than fixing it on our side (for now its ok). but we should fix the issue in future and removed this pre-computing.

Yes, the crash originates in scf::WhileOpConversion which detaches the before/after blocks before invoking convertSignatureBlock..so its an upstream issue specific to scf.while

> this still bit confusing. in the crashing scenario who is looking up the layouts? is it a block arg? it can not be a an op inside the blocks if the blocks are detached. if it is an arg shouldn't that layout still available through the layout prop analysis?

The lookup is on the scf.while block argument itself; getDistributeLayoutAttr(v) walks up via v.getParentBlock()->getParent() to find the layout, and that traversal hits an LLVM assertion on detached blocks 

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