[Mlir-commits] [mlir] [MLIR][XeGPU] Add support for elementwise ops in Wg to Sg distribute pass [1/N] (PR #142797)
Chao Chen
llvmlistbot at llvm.org
Tue Jun 17 07:49:54 PDT 2025
================
@@ -518,6 +579,29 @@ void XeGPUWgToSgDistributePass::runOnOperation() {
return isLegal(layout);
});
+ target.addDynamicallyLegalDialect<math::MathDialect, arith::ArithDialect>(
+ [=](Operation *op) -> std::optional<bool> {
+ // Only handle elementwise mappable ops
+ if (!OpTrait::hasElementwiseMappableTraits(op))
+ return true;
+
+ VectorType resultType =
+ dyn_cast<VectorType>(op->getResult(0).getType());
+ if (!resultType)
+ return true;
+
+ // Check if all operands are vectors of the same shape
+ for (Value operand : op->getOperands()) {
----------------
chencha3 wrote:
this loop is equivalent to
```
if (llvm::any_of(op->getOperandTypes(), [&](Type type) { return type != resultType; }))
return true;
```
https://github.com/llvm/llvm-project/pull/142797
More information about the Mlir-commits
mailing list