[OPTIMIZER] Adjusted heuristics (#1001)

triton-lang · Dec 21, 2022 · 19d1831 · 19d1831
1 parent 655b71a
commit 19d1831
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/lib/Dialect/TritonGPU/Transforms/Combine.cpp b/lib/Dialect/TritonGPU/Transforms/Combine.cpp
@@ -461,14 +461,18 @@ class FoldConvertAndReduce : public mlir::RewritePattern {
 public:
   explicit FoldConvertAndReduce(mlir::MLIRContext *context)
       : mlir::RewritePattern(triton::gpu::ConvertLayoutOp::getOperationName(),
-                             2, context) {}
+                             1, context) {}
 
   mlir::LogicalResult
   matchAndRewrite(mlir::Operation *cvtOp,
                   mlir::PatternRewriter &rewriter) const override {
     auto cvt = dyn_cast<triton::gpu::ConvertLayoutOp>(*cvtOp);
     auto srcEncoding =
+        cvt.getOperand().getType().cast<RankedTensorType>().getEncoding();
+    auto dstEncoding =
         cvt.getResult().getType().cast<RankedTensorType>().getEncoding();
+    if (srcEncoding.isa<triton::gpu::SliceEncodingAttr>())
+      return failure();
     SetVector<Operation *> cvtSlices;
     auto filter = [&](Operation *op) {
       return op->getBlock() == cvt->getBlock() &&
@@ -499,7 +503,7 @@ class FoldConvertAndReduce : public mlir::RewritePattern {
         llvm::MapVector<Value, Attribute> toConvert;
         if (argOp && (argOp != cvt) && cvtSlices.count(argOp) == 0 &&
             failed(simulateBackwardRematerialization(argOp, processed, layout,
-                                                     toConvert, srcEncoding))) {
+                                                     toConvert, dstEncoding))) {
           return failure();
         }
       }