[FRONTEND] Relax calling function inside a loop (#7002)

We should allow calling functions inside a loop even if they return a
value as this doesn't cause early exit.
If the function has unstructured control flow it may not be inlined, we
should have a separate check about that.

Author: ThomasRaoux

python/test/unit/language/test_frontend.py

@@ -109,3 +109,20 @@ def test_list_of_functions():
     # CHECK-NEXT: call @anchor
     # CHECK-NEXT: call @forward
     list_of_functions_constexpr(tl.arange(0, 4), [anchor, forward])
+
+
+@triton.jit
+def accumulate(a, b):
+    return a + b
+
+
+# Check that we can call a function returning a value from a loop.
+@filecheck_test
+@triton.jit
+def test_call_in_loop():
+    # CHECK-LABEL: test_call_in_loop
+    acc = 0
+    # CHECK: scf.for
+    # CHECK:   call @accumulate
+    for i in range(10):
+        acc = accumulate(acc, i)

python/triton/compiler/code_generator.py

@@ -135,10 +135,9 @@ def _visit_stmts(self, body) -> bool:
         return any(self.visit(s) for s in body)
 
     def _visit_function(self, fn) -> bool:
-        # Currently we only support JITFunctions defined in the global scope
-        if isinstance(fn, JITFunction) and not fn.noinline:
-            fn_node = fn.parse()
-            return ContainsReturnChecker(self.gscope).visit(fn_node)
+        # 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.
         return False
 
     def generic_visit(self, node) -> bool: