[FRONTEND] Remove incorrect and misleading error message from the return-in-while check (#7551)

We actually don't check `return`s in functions being called within
`while` or `for`.

```
    def _visit_function(self, fn) -> bool:
        # No need to check within the function as it won't cause an early return.
        # If the function itself has unstructured control flow we may not be able to inline it causing poor performance,
        # we should check for this and emit a warning.
        return False
```

Author: Jokeren

python/test/unit/language/test_frontend.py

@@ -1,6 +1,8 @@
 import triton
 import triton.language as tl
 from triton._filecheck import filecheck_test, run_filecheck_test
+from triton.compiler.errors import CompilationError
+import pytest
 
 # ===-----------------------------------------------------------------------===#
 # Unit Tests
@@ -528,3 +530,36 @@ def test_specialized_recursion():
 
     # CHECK: func {{.*}}recursive_reduce__i32S4S
     # CHECK-COUNT-2: call {{.*}}recursive_reduce__i32S2S
+
+
+@triton.jit
+def trivial_return():
+    return
+
+
+@filecheck_test
+@triton.jit
+def test_call_in_while():
+    # CHECK-LABEL: test_call_in_while
+    i = 0
+    while i < 10:
+        if i == 5:
+            trivial_return()
+        else:
+            trivial_return()
+
+
+def test_return_in_while():
+
+    @triton.jit
+    def kernel():
+        i = 0
+        while i < 10:
+            if i == 5:
+                return
+            i += 1
+
+    with pytest.raises(CompilationError) as e:
+        kernel.warmup(grid=(1, ))
+
+    assert "Cannot have `return` statements inside `while` or `for` statements in triton" in str(e.value)

python/triton/compiler/code_generator.py

@@ -149,9 +149,9 @@ def _visit_stmts(self, body) -> bool:
         return any(self.visit(s) for s in body)
 
     def _visit_function(self, fn) -> bool:
-        # no need to check within the function as it won't cause an early return.
-        # If the function itself has unstructured control flow we may not be able to inline it causing poor performance.
-        # We should check for this and fail or emit a warning.
+        # No need to check within the function as it won't cause an early return.
+        # If the function itself has unstructured control flow we may not be able to inline it causing poor performance,
+        # we should check for this and emit a warning.
         return False
 
     def generic_visit(self, node) -> bool:
@@ -857,13 +857,10 @@ def visit_If(self, node):
                     % ast.unparse(node.test))
                 cond = language.core._unsplat(cond, _semantic=self.semantic, _generator=self)
             cond = cond.to(language.int1, _semantic=self.semantic)
-            contains_return = ContainsReturnChecker(self.gscope).visit(node)
-            if contains_return:
+            if ContainsReturnChecker(self.gscope).visit(node):
                 if self.scf_stack:
                     raise self._unsupported(
-                        node, "Cannot have `return` statements inside `while` or `for` statements in triton "
-                        "(note that this also applies to `return` statements that are inside functions "
-                        "transitively called from within `while`/`for` statements)")
+                        node, "Cannot have `return` statements inside `while` or `for` statements in triton.")
                 self.visit_if_top_level(cond, node)
             else:
                 self.visit_if_scf(cond, node)