diff --git a/third_party/amd/lib/TritonAMDGPUTransforms/ReorderInstructions.cpp b/third_party/amd/lib/TritonAMDGPUTransforms/ReorderInstructions.cpp index 6981bd31bdb3..f15feb7b252f 100644 --- a/third_party/amd/lib/TritonAMDGPUTransforms/ReorderInstructions.cpp +++ b/third_party/amd/lib/TritonAMDGPUTransforms/ReorderInstructions.cpp @@ -227,22 +227,17 @@ static void scheduleGlobalLoadLocalStore(triton::FuncOp funcOp) { // Gather use-def chain in block. Block *block = op->getBlock(); bool leadsToLoad = false; - bool dontReorder = false; SetVector backwardSet; BackwardSliceOptions options; options.omitBlockArguments = true; options.inclusive = false; + // Slice should inlcude values flowing into op regions + options.omitUsesFromAbove = false; options.filter = [&](Operation *defOp) -> bool { Block *defBlock = defOp->getBlock(); if (!block->findAncestorOpInBlock(*defOp)) return false; - // Don't hoist control flow as we don't track backtraces of ops within - // their regions. - if (isa(defOp)) { - dontReorder = true; - return false; - } // Check for a `load` dependent path. leadsToLoad |= isa(defOp); @@ -252,9 +247,6 @@ static void scheduleGlobalLoadLocalStore(triton::FuncOp funcOp) { mlir::getBackwardSlice(op, &backwardSet, options); backwardSet.insert(op); - // If we found ops in the slice we don't want to hoist. - if (dontReorder) - continue; // Don't move a local_store if its source is a load from // the same iteration. if (isa(op) && leadsToLoad)