[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:53 PDT 2025
================
@@ -328,6 +331,63 @@ struct WgToSgPrefetchNdOp : public OpConversionPattern<xegpu::PrefetchNdOp> {
}
};
+// This pattern transforms elementwise ops in math/arith dialect
+struct WgToSgElementwiseOp : public ConversionPattern {
+ WgToSgElementwiseOp(MLIRContext *ctx)
+ : ConversionPattern(MatchAnyOpTypeTag(), /*benefit=*/1, ctx) {}
+
+ LogicalResult
+ matchAndRewrite(Operation *op, ArrayRef<ValueRange> operands,
+ ConversionPatternRewriter &rewriter) const override {
+ // Only match ops with elementwise trait and single result.
+ if (!OpTrait::hasElementwiseMappableTraits(op) || op->getNumResults() != 1)
+ return failure();
+
+ auto resultType = dyn_cast<VectorType>(op->getResult(0).getType());
----------------
chencha3 wrote:
maybe there is a need to check whether resultType is null. I remember some ops work on scalars. Could you help to double check?
https://github.com/llvm/llvm-project/pull/142797
More information about the Mlir-commits
mailing list