[intel] define barrier in AssertOp' last block

Signed-off-by: Anatoly Myachev <[email protected]>

Author: anmyachev

lib/Conversion/TritonGPUToLLVM/AssertOpToLLVM.cpp

@@ -34,22 +34,23 @@ struct AssertOpConversion : public ConvertOpToLLVMPattern<triton::AssertOp> {
         return failure();
       }
     }
-    llAssert(op, condition, adaptor.getMessage(), rewriter);
+    Block *thenBlock = llAssert(op, condition, adaptor.getMessage(), rewriter);
     if (isa<RankedTensorType>(op.getCondition().getType())) {
       // Add a barrier to avoid a race condition in case an assert is followed
       // by an op that may trap if the assert condition is true. Since the
       // tensor in those two operations may have different layout we need to
       // make sure all the threads are done executing the assert before going to
       // the next op.
+      rewriter.setInsertionPointToStart(thenBlock);
       barrier();
     }
     rewriter.eraseOp(op);
     return success();
   }
   // op: the op at which the assert is inserted. Unlike printf, we need to
   // know about the op to split the block.
-  void llAssert(Operation *op, Value condition, StringRef message,
-                ConversionPatternRewriter &rewriter) const {
+  Block *llAssert(Operation *op, Value condition, StringRef message,
+                  ConversionPatternRewriter &rewriter) const {
     ConversionPatternRewriter::InsertionGuard guard(rewriter);
 
     auto ctx = rewriter.getContext();
@@ -87,6 +88,7 @@ struct AssertOpConversion : public ConvertOpToLLVMPattern<triton::AssertOp> {
     rewriter.create<cf::BranchOp>(loc, thenBlock);
     rewriter.setInsertionPointToEnd(prevBlock);
     rewriter.create<cf::CondBranchOp>(loc, condition, ifBlock, thenBlock);
+    return thenBlock;
   }
 
 protected: