update and clean up

Author: CRobeck

lib/Dialect/TritonGPU/Transforms/Pipeliner/PipelineExpander.cpp

@@ -66,8 +66,8 @@ struct LoopPipelinerInternal {
   bool dynamicLoop;
   bool peelEpilogue;
   bool guardEpilogue;
-  triton::PipeliningOption::PredicateOpFnType predicateFn = nullptr;
   triton::PipeliningOption::AnnotationlFnType annotateFn = nullptr;
+  triton::PipeliningOption::PredicateOpFnType predicateFn = nullptr;
 
   // When peeling the kernel we generate several version of each value for
   // different stage of the prologue. This map tracks the mapping between
@@ -158,7 +158,6 @@ bool LoopPipelinerInternal::initializeLoopInfo(
   peelEpilogue = options.peelEpilogue;
   predicateFn = options.predicateFn;
   guardEpilogue = options.guardEpilogue;
-
   if ((!peelEpilogue || dynamicLoop) && predicateFn == nullptr) {
     LDBG("--no epilogue or predicate set -> BAIL");
     return false;
@@ -776,16 +775,11 @@ LoopPipelinerInternal::emitEpilogue(RewriterBase &rewriter,
           unsigned nextVersion = currentVersion + 1;
           Value pred = predicates[currentVersion];
           Value prevValue = valueMapping[mapVal][currentVersion];
-          Value nextValue = pair.value();
-
-          if (guardEpilogue) {
-            nextValue = rewriter.create<arith::SelectOp>(loc, pred, nextValue,
-                                                         prevValue);
-          }
-
-          returnValues[ri] = nextValue;
+          auto selOp = rewriter.create<arith::SelectOp>(loc, pred, pair.value(),
+                                                        prevValue);
+          returnValues[ri] = selOp;
           if (nextVersion <= maxStage)
-            setValueMapping(mapVal, nextValue, nextVersion);
+            setValueMapping(mapVal, selOp, nextVersion);
         }
       }
     }

third_party/amd/backend/compiler.py

@@ -246,7 +246,7 @@ def make_ttgir(mod, metadata, options):
         global_prefetch = int(os.getenv("TRITON_HIP_GLOBAL_PREFETCH", "0"))
         local_prefetch = int(os.getenv("TRITON_HIP_LOCAL_PREFETCH", "0"))
         use_async_copy = int(os.getenv("TRITON_HIP_USE_ASYNC_COPY", "0")) == 1
-        must_guard_epilogue = 1
+        must_guard_epilogue = int(os.getenv("TRITON_HIP_GUARD_EPILOGUE", "1"))
 
         # The `local-prefetch` scheduling variant requires turning on buffer ops.
         if options.schedule_hint == "local-prefetch":

third_party/amd/lib/TritonAMDGPUTransforms/StreamPipeline.cpp

@@ -127,14 +127,6 @@ class StreamPipeliner {
     stages[SCHED_LOCAL_STORE] = _globalPrefetch;
     stages[SCHED_LOCAL_LOAD] = lastStage - _localPrefetch;
     stages[SCHED_COMPUTE] = lastStage;
-
-    options.supportDynamicLoops = true;
-    options.peelEpilogue = true;
-    options.guardEpilogue = _mustGuardEpilogue;
-    if (_mustGuardEpilogue)
-      options.predicateFn = streamPredication;
-    else
-      options.predicateFn = tt::predicateOp;
   }
 
   LogicalResult pipelineLoop();
@@ -210,9 +202,6 @@ class StreamPipeliner {
 
   // Capture list of new shared memory buffers.
   SmallVector<Value> sharedMemAllocs;
-
-  // Pipelining options for the PipelineExpander
-  tt::PipeliningOption options;
 };
 
 } // namespace
@@ -277,7 +266,8 @@ bool StreamPipeliner::safeDAG(Value v, int index) {
   }
   return false;
 }
-
+// TODO(crobeck): is this valid if we have loop-carried 
+// results needed for the next epilogue stage?
 void StreamPipeliner::checkResultResilience() {
   auto yieldVals = forOp.getYieldedValuesMutable().value();
   for (auto [index, res] : llvm::enumerate(forOp.getResults())) {
@@ -1016,18 +1006,6 @@ LogicalResult StreamPipeliner::preprocessLoopAndBuildSchedule() {
     schedule.dump();
   });
 
-  // Create the final schedule for the kernel loop. This will dictate the
-  // stages and order of operations to the pipeline expander.
-  std::vector<std::pair<Operation *, unsigned>> coarseSchedule =
-      schedule.createFinalSchedule(forOp);
-
-  // Fill out the pipeline options.
-  options.getScheduleFn =
-      [coarseSchedule](scf::ForOp,
-                       std::vector<std::pair<Operation *, unsigned>> &s) {
-        s = std::move(coarseSchedule);
-      };
-
   OpBuilder builder(forOp);
   builder.setInsertionPointAfter(forOp);
   // Explicitly deallocate created allocations.
@@ -1041,6 +1019,27 @@ LogicalResult StreamPipeliner::pipelineLoop() {
   if (failed(preprocessLoopAndBuildSchedule()))
     return failure();
   LDBG("Loop before sending to expander:\n" << *forOp);
+  // Create the final schedule for the kernel loop. This will dictate the
+  // stages and order of operations to the pipeline expander.
+  std::vector<std::pair<Operation *, unsigned>> coarseSchedule =
+      schedule.createFinalSchedule(forOp);
+
+  // Pipelining options for the PipelineExpander
+  tt::PipeliningOption options;
+  options.supportDynamicLoops = true;
+  options.peelEpilogue = true;
+  options.guardEpilogue = mustGuardEpilogue;
+
+  if (mustGuardEpilogue)
+    options.predicateFn = streamPredication;
+  else
+    options.predicateFn = tt::predicateOp;
+
+  options.getScheduleFn =
+      [&coarseSchedule](scf::ForOp,
+                        std::vector<std::pair<Operation *, unsigned>> &s) {
+        s = std::move(coarseSchedule);
+      };
 
   IRRewriter rewriter(forOp->getContext());
   rewriter.setInsertionPoint(forOp);