[BACKEND] Fix convert to triton gpu for CF ops (#6909)

Functions with unstructured control flow passing tensors were not
working as the function conversion pattern was converting all the blocks
and creating illegal IR.

Author: ThomasRaoux

lib/Conversion/TritonToTritonGPU/TritonGPUConversion.cpp

@@ -101,6 +101,15 @@ TritonGPUConversionTarget::TritonGPUConversionTarget(
       return true;
     return false;
   });
+  addDynamicallyLegalOp<triton::FuncOp>([](triton::FuncOp funcOp) -> bool {
+    for (auto arg : funcOp.getArguments()) {
+      if (auto tensor = dyn_cast<RankedTensorType>(arg.getType())) {
+        if (!tensor.getEncoding())
+          return false;
+      }
+    }
+    return true;
+  });
 }
 
 bool TritonGPUConversionTarget::isDynamicallyLegal(

lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp

@@ -481,14 +481,17 @@ class TritonFuncOpPattern : public OpConversionPattern<triton::FuncOp> {
   matchAndRewrite(triton::FuncOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     auto converter = getTypeConverter();
+    TypeConverter::SignatureConversion result(op.getNumArguments());
     auto newOp = rewriter.replaceOpWithNewOp<triton::FuncOp>(
         op, op.getName(), op.getFunctionType());
     addNamedAttrs(newOp, adaptor.getAttributes());
     rewriter.inlineRegionBefore(op.getBody(), newOp.getBody(),
                                 newOp.getBody().end());
-    if (failed(rewriter.convertRegionTypes(&newOp.getBody(), *converter)))
-      return failure();
-
+    // Convert just the entry block. The remaining unstructured control flow is
+    // converted by br patterns.
+    if (!newOp.getBody().empty())
+      rewriter.applySignatureConversion(&newOp.getBody().front(), result,
+                                        converter);
     return success();
   }
 };

test/Conversion/triton_to_tritongpu.mlir

@@ -155,3 +155,16 @@ tt.func @ub_poison() {
   %0 = ub.poison : tensor<128x64xf16>
   tt.return
 }
+
+// -----
+
+// CHECK-LABEL: @cf_br
+tt.func @cf_br(%ptr: !tt.ptr<i32>) {
+  %cst = arith.constant dense<1> : tensor<128xi32>
+  // cf.br ^bb1(%{{.+}} : tensor<128xi32, #{{.+}}>)
+  cf.br ^bb1(%cst : tensor<128xi32>)
+^bb1(%arg0: tensor<128xi32>):
+  %ptrs = tt.splat %ptr : !tt.ptr<i32> -> tensor<128x!tt.ptr<i32>>
+  tt.store %ptrs, %arg0 : tensor<128x!tt.ptr<i32>>
+  tt.return
+}