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

[Jianhui Li]

================
@@ -991,3 +847,190 @@ bool xegpu::matchSplitDimExpansion(
   }
   return srcIdx == src.size();
 }
+
+//===----------------------------------------------------------------------===//
+// Context-aware type conversion utilities
+//===----------------------------------------------------------------------===//
+
+// Pre-computes block argument type mappings for SCF loops (scf.while,
+// scf.for).
+//
+// Block-arg layouts ARE available in the IR (layout recovery propagates
+// them onto the loop op as `layout_operand_N`). The reason we cannot rely
+// on the regular `getDistributeLayoutAttr(v)` lookup during structural
+// conversion is structural, not informational:
+//   - For `scf.while`, `scf::WhileOpConversion` detaches the before/after
+//     blocks from their parent region before invoking
+//     `convertSignatureBlock`. Looking up a detached BlockArgument's layout
+//     walks `v.getParentBlock()->getParent()` and trips an LLVM ilist
+//     assertion.
+//   - For `scf.for`, `scf::ForOpConverter` builds a new `scf.for` and moves
+//     the body block into it. The new op does NOT inherit the temporary
+//     `layout_operand_N` attributes that layout recovery set on the old
+//     op, so any post-move query of a body block argument's layout (e.g.
+//     when a pattern that consumes the iter_arg via a non-anchor op like
+//     `vector.insert_strided_slice` runs after the move) returns null.
+// Caching the distributed types by `Value` identity sidesteps both failure
+// modes. `scf.if` has no block arguments and is therefore not covered here.
+DenseMap<Value, SmallVector<Type>>
+xegpu::precomputeLoopBlockArgTypes(Operation *topLevelOp,
+                                   SubShapeAndCountFn getSubShapeAndCount) {
+  DenseMap<Value, SmallVector<Type>> loopArgTypes;
+  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);
+    loopArgTypes[arg] = std::move(types);
+  };
+  topLevelOp->walk([&](Operation *op) {
+    if (auto whileOp = dyn_cast<scf::WhileOp>(op)) {
+      // "before" region block arguments correspond to the `inits` operands.
+      for (auto [init, arg] :
+           llvm::zip(whileOp.getInits(), whileOp.getBeforeArguments()))
+        recordBlockArgTypes(init, arg);
+      // "after" region block arguments correspond to the operands of the
+      // embedded `scf.condition` op (not the `inits`). In general the two
+      // type lists may differ.
+      scf::ConditionOp condOp = whileOp.getConditionOp();
+      for (auto [condArg, arg] :
+           llvm::zip(condOp.getArgs(), whileOp.getAfterArguments()))
+        recordBlockArgTypes(condArg, arg);
+      return;
+    }
+    if (auto forOp = dyn_cast<scf::ForOp>(op)) {
+      // Body block args (excluding the induction variable) correspond to
+      // the `initArgs` operands.
+      for (auto [init, arg] :
+           llvm::zip(forOp.getInitArgs(), forOp.getRegionIterArgs()))
+        recordBlockArgTypes(init, arg);
+      return;
+    }
+  });
+  return loopArgTypes;
+}
+
+void xegpu::addVectorTypeConversion(
+    TypeConverter &converter, SubShapeAndCountFn getSubShapeAndCount,
+    DenseMap<Value, SmallVector<Type>> loopArgTypes) {
+  // Context-aware VectorType conversion (1:1 shape-changing or 1:N). For
+  // SCF loop block arguments (scf.while, scf.for), uses the pre-computed
+  // map. For all other Values, retrieves the layout directly via
+  // getDistributeLayoutAttr.
+  auto loopArgTypeMap = std::make_shared<DenseMap<Value, SmallVector<Type>>>(
+      std::move(loopArgTypes));
+  converter.addConversion(
+      [loopArgTypeMap, getSubShapeAndCount](
+          Value v,
+          SmallVectorImpl<Type> &result) -> std::optional<LogicalResult> {
+        if (!isa<VectorType>(v.getType()))
+          return std::nullopt;
+
+        // Check pre-computed map first (for SCF loop block args).
+        if (isa<BlockArgument>(v)) {
+          auto it = loopArgTypeMap->find(v);
+          if (it != loopArgTypeMap->end()) {
+            result.append(it->second.begin(), it->second.end());
+            return success();
+          }
+        }
+
+        // For OpResults and other block arguments (e.g. region args of
----------------
Jianhui-Li wrote:

can these layout be precomputed and so these code be removed, if you have to precompute and cached anyway? 

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