[Triton-MLIR][OPTIMIZER] Add ExtElemwiseOp to expensive_to_remat list

2022-11-04 15:23:58 +08:00
parent 4218e68d74
commit a6d672166c
1 changed files with 2 additions and 2 deletions
--- 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>();
-}
+}