[Pipeliner] Properly fail instead of assert if cannot predicate op (#4659)

Mirrored from upstream change:
llvm/llvm-project#107442

Author: sjw36

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

@@ -89,7 +89,7 @@ struct LoopPipelinerInternal {
   bool initializeLoopInfo(ForOp op, const triton::PipeliningOption &options);
   /// Emits the prologue, this creates `maxStage - 1` part which will contain
   /// operations from stages [0; i], where i is the part index.
-  void emitPrologue(RewriterBase &rewriter);
+  LogicalResult emitPrologue(RewriterBase &rewriter);
   /// Gather liverange information for Values that are used in a different stage
   /// than its definition.
   llvm::MapVector<Value, LiverangeInfo> analyzeCrossStageValues();
@@ -275,7 +275,7 @@ cloneAndUpdateOperands(RewriterBase &rewriter, Operation *op,
   return clone;
 }
 
-void LoopPipelinerInternal::emitPrologue(RewriterBase &rewriter) {
+LogicalResult LoopPipelinerInternal::emitPrologue(RewriterBase &rewriter) {
   // Initialize the iteration argument to the loop initiale values.
   for (auto [arg, operand] :
        llvm::zip(forOp.getRegionIterArgs(), forOp.getInitsMutable())) {
@@ -323,7 +323,8 @@ void LoopPipelinerInternal::emitPrologue(RewriterBase &rewriter) {
       if (predicates[predicateIdx]) {
         OpBuilder::InsertionGuard insertGuard(rewriter);
         newOp = predicateFn(rewriter, newOp, predicates[predicateIdx]);
-        assert(newOp && "failed to predicate op.");
+        if (newOp == nullptr)
+          return failure();
       }
       if (annotateFn)
         annotateFn(newOp, triton::PipeliningOption::PipelinerPart::Prologue, i);
@@ -351,6 +352,7 @@ void LoopPipelinerInternal::emitPrologue(RewriterBase &rewriter) {
       }
     }
   }
+  return success();
 }
 
 llvm::MapVector<Value, LoopPipelinerInternal::LiverangeInfo>
@@ -787,7 +789,8 @@ mlir::triton::pipelineForLoop(RewriterBase &rewriter, ForOp forOp,
     *modifiedIR = true;
 
   // 1. Emit prologue.
-  pipeliner.emitPrologue(rewriter);
+  if (failed(pipeliner.emitPrologue(rewriter)))
+    return failure();
 
   // 2. Track values used across stages. When a value cross stages it will
   // need to be passed as loop iteration arguments.