Use OpTrait to check for elementwise ops (#955)

intel · Nov 6, 2024 · 42fdea5 · 42fdea5
1 parent b2742c3
commit 42fdea5
Showing 1 changed file with 1 addition and 10 deletions.
diff --git a/lib/Dialect/XeTile/Transforms/WgToSg.cpp b/lib/Dialect/XeTile/Transforms/WgToSg.cpp
@@ -750,14 +750,6 @@ class WGToSGVectorBroadcast
   }
 };
 
-
-// TODO: Add more pre-ops
-bool isElementWiseOp(mlir::Operation *op) {
-  return llvm::isa<mlir::arith::AddFOp>(op) ||
-         llvm::isa<mlir::arith::MulFOp>(op) ||
-         llvm::isa<mlir::math::ExpOp>(op);
-}
-
 // Helper function to analyze the def-use chain of initTileOps. Currently we
 // pattern match the following def-use chain as a candidate for
 // load + tranpose optimization.
@@ -813,7 +805,7 @@ void analyzeInitTileOps(mlir::Operation *op) {
     // Check if vector.transpose is consumed by TileMMA directly or
     // is consumed by some pre-op and then TileMMA.
     if(!llvm::isa<imex::xetile::TileMMAOp>(consumerOp)){
-      if(!isElementWiseOp(consumerOp) &&
+      if(!OpTrait::hasElementwiseMappableTraits(consumerOp) &&
          !(llvm::isa<mlir::vector::BroadcastOp>(consumerOp))) {
         return mlir::WalkResult::skip();
       }
@@ -834,7 +826,6 @@ void analyzeInitTileOps(mlir::Operation *op) {
   });
 }
 
-
 void populateXeTileWgToSgPatterns(imex::XeOneToNTypeConverter &converter,
                                   mlir::RewritePatternSet &patterns,
                                   TileUsageAnalysis &analysis) {