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

[Nishant Patel]

================
@@ -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;
+    auto newTy = VectorType::get(subShape, vecTy.getElementType());
+    SmallVector<Type> types(count, newTy);
+    (*whileArgTypeMap)[arg] = std::move(types);
+  };
+  topLevelOp->walk([&](scf::WhileOp whileOp) {
----------------
nbpatel wrote:

The pre-compute is purely a workaround for scf::WhileOpConversion since it ends up with a detached "before" block while doing the conversion and ends up with a nullptr...in future we need to fix upstream code to avoid this precomputation..but for now I think its ok since we dont have any use cases with scf.while so we rarely end up here......scf.for / scf.if keep their blocks attached during conversion hence it works fine and doesnt require the pre-compute

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