From a6d672166ce1a2d3f7986568afbf511c557cefb5 Mon Sep 17 00:00:00 2001
From: Qingyi Liu <qingyil@nvidia.com>
Date: Fri, 4 Nov 2022 15:23:58 +0800
Subject: [PATCH] [Triton-MLIR][OPTIMIZER] Add ExtElemwiseOp to
 expensive_to_remat list

---
 lib/Dialect/TritonGPU/Transforms/Combine.cpp | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/lib/Dialect/TritonGPU/Transforms/Combine.cpp b/lib/Dialect/TritonGPU/Transforms/Combine.cpp
index 456ce1200..0135a1dca 100644
--- a/lib/Dialect/TritonGPU/Transforms/Combine.cpp
+++ b/lib/Dialect/TritonGPU/Transforms/Combine.cpp
@@ -200,7 +200,7 @@ inline bool expensive_to_remat(Operation *op) {
     return true;
   if (isa<triton::gpu::ExtractSliceOp, triton::gpu::AllocTensorOp,
           triton::gpu::InsertSliceAsyncOp, triton::LoadOp, triton::StoreOp,
-          triton::DotOp>(op))
+          triton::DotOp, triton::ExtElemwiseOp>(op))
     return true;
   if (isa<scf::YieldOp, scf::ForOp>(op))
     return true;
@@ -623,4 +623,4 @@ public:
 
 std::unique_ptr<Pass> mlir::createTritonGPUCombineOpsPass() {
   return std::make_unique<TritonGPUCombineOpsPass>();
-}
\ No newline at end of file
+}