From 42fdea56f9cdae92bc764677b4c320b0a8b16d4d Mon Sep 17 00:00:00 2001
From: Nishant Patel <nishant.b.patel@intel.com>
Date: Wed, 6 Nov 2024 15:16:36 -0800
Subject: [PATCH] Use OpTrait to check for elementwise ops (#955)

---
 lib/Dialect/XeTile/Transforms/WgToSg.cpp | 11 +----------
 1 file changed, 1 insertion(+), 10 deletions(-)

diff --git a/lib/Dialect/XeTile/Transforms/WgToSg.cpp b/lib/Dialect/XeTile/Transforms/WgToSg.cpp
index 77ef9f7d2..783cabd52 100644
--- 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) {