[BACKEND] Move lowering of CF as the last step of conversion to LLVM (#7213)

This prevents problem with analysis picking up argument from blocks that
have been removed.

Author: ThomasRaoux

lib/Conversion/TritonGPUToLLVM/AssertOpToLLVM.cpp

@@ -84,9 +84,9 @@ struct AssertOpConversion : public ConvertOpToLLVMPattern<triton::AssertOp> {
     // Split a block after the call.
     Block *thenBlock = rewriter.splitBlock(ifBlock, op->getIterator());
     rewriter.setInsertionPointToEnd(ifBlock);
-    rewriter.create<cf::BranchOp>(loc, thenBlock);
+    rewriter.create<LLVM::BrOp>(loc, thenBlock);
     rewriter.setInsertionPointToEnd(prevBlock);
-    rewriter.create<cf::CondBranchOp>(loc, condition, ifBlock, thenBlock);
+    rewriter.create<LLVM::CondBrOp>(loc, condition, ifBlock, thenBlock);
     rewriter.setInsertionPointToStart(thenBlock);
   }
 

lib/Conversion/TritonGPUToLLVM/ReduceScanCommon.h

@@ -97,12 +97,12 @@ inline SmallVector<Value> applyCombineOp(Location loc,
     thenBlockArgs.push_back(undef);
     thenBlock->addArgument(ty, loc);
   }
-  rewriter.create<cf::CondBranchOp>(loc, pred, &newCombine, combineArgs,
-                                    thenBlock, thenBlockArgs);
+  rewriter.create<LLVM::CondBrOp>(loc, pred, &newCombine, combineArgs,
+                                  thenBlock, thenBlockArgs);
 
   // Split a block after the call.
   rewriter.setInsertionPointToEnd(&newCombine);
-  rewriter.replaceOpWithNewOp<cf::BranchOp>(returnOp, thenBlock, results);
+  rewriter.replaceOpWithNewOp<LLVM::BrOp>(returnOp, results, thenBlock);
   rewriter.setInsertionPointToStart(thenBlock);
   return SmallVector<Value>(thenBlock->getArguments());
 }

test/Conversion/tritongpu_to_llvm.mlir

@@ -2397,3 +2397,18 @@ tt.func private @memdesc_reinterpret(%arg0: !ttg.memdesc<4x1024xi8, #shared0, #t
 }
 
 }
+
+// -----
+
+#blocked = #ttg.blocked<{sizePerThread = [1, 1], threadsPerWarp = [1, 32], warpsPerCTA = [1, 4], order = [1, 0]}>
+module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
+  // CHECK-LABEL: load_br
+  tt.func @load_br(%arg0: tensor<16x4x!tt.ptr<i8>, #blocked>) {
+    // CHECK: llvm.br
+    cf.br ^bb1(%arg0 : tensor<16x4x!tt.ptr<i8>, #blocked>)
+    ^bb1(%arg1: tensor<16x4x!tt.ptr<i8>, #blocked>):
+    // CHECK: ld.global.b8
+      %0 = tt.load %arg1 : tensor<16x4x!tt.ptr<i8>, #blocked>
+      tt.return
+  }
+}

third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp

@@ -95,8 +95,6 @@ struct ConvertTritonGPUToLLVM
     RewritePatternSet funcPatterns(context);
     mlir::triton::populateFuncOpConversionPattern(
         typeConverter, funcPatterns, targetInfo, patternBenefitDefault);
-    mlir::cf::populateControlFlowToLLVMConversionPatterns(typeConverter,
-                                                          funcPatterns);
     if (failed(
             applyPartialConversion(mod, funcTarget, std::move(funcPatterns))))
       return signalPassFailure();
@@ -152,8 +150,6 @@ struct ConvertTritonGPUToLLVM
     mlir::arith::populateArithToLLVMConversionPatterns(typeConverter, patterns);
     mlir::populateMathToLLVMConversionPatterns(typeConverter, patterns);
     mlir::populateGpuToNVVMConversionPatterns(typeConverter, patterns);
-    mlir::cf::populateControlFlowToLLVMConversionPatterns(typeConverter,
-                                                          patterns);
     mlir::ub::populateUBToLLVMConversionPatterns(typeConverter, patterns);
     mlir::triton::populateViewOpToLLVMPatterns(typeConverter, patterns,
                                                benefit);
@@ -173,6 +169,18 @@ struct ConvertTritonGPUToLLVM
     if (failed(applyPartialConversion(mod, convTarget, std::move(patterns))))
       return signalPassFailure();
 
+    // Lower CF ops separately to avoid breaking analysis.
+    TritonLLVMFunctionConversionTarget cfTarget(*context);
+    cfTarget.markUnknownOpDynamicallyLegal([&](Operation *op) {
+      return op->getDialect() !=
+             context->getLoadedDialect<cf::ControlFlowDialect>();
+    });
+    RewritePatternSet cfPatterns(context);
+    mlir::cf::populateControlFlowToLLVMConversionPatterns(typeConverter,
+                                                          cfPatterns);
+    if (failed(applyPartialConversion(mod, cfTarget, std::move(cfPatterns))))
+      return signalPassFailure();
+
     // Fold CTAId when there is only 1 CTA.
     int numCTAs = triton::gpu::TritonGPUDialect::getNumCTAs(mod);
     if (numCTAs == 1) {