Multiple Layers of For-Loops

[MichaelElexon]

Hey all.

I'm trying to do some iterative optimization and I'm struggling to figure out how I can nest for-loops.
My cost function has a for-loop in it (it calculate the cost by doing an operation over a vector field).
Each iteration needs to run the cost function with 27 different parameter combinations (thus I'm using a for-loop to iteratively feed these into the cost functions)
I want to perform 80000 iterations (thus I'm using a for-loop here too).

EG, the python code looks like this:

for i in iterations: # 80000
   for j in parameters: # 27
      objective_function_taichi_kernel(parameter[j]) # size based on data set

I get the following error when trying to implement these loops directly in taichi:
"struct_for cannot be nested inside a kernel"

As a result, I've put the outer loops in python context, but but I did some profiling and 85% of my computation time is coming from doing this.

I don't need the outer iterations to be parallelized (i'm choosing a new set of parameters based upon results, so I have to do iterations serially).

Is there a common solution for nesting these loops? Or another taichi-esque way of solving this problem?

[hughperkins]

Could you provide a minimum reproducible example of your code please? (try to reduce it down to ~10-20 lines idealy)

[MichaelElexon]

Hey,

Here's a minimum reproducible example:

import taichi as ti

ti.init(arch=ti.cpu) # Use CPU is fine for showing a compile-time error

# Arbitrary data
data = ti.field(dtype=ti.i32, shape=4)
data.fill(1)

result = ti.field(dtype=ti.i32, shape=4)
result.fill(0)

# This is a function which does arbitrary operations using a struct-for loop (ti.ndrange)
@ti.func
def a_struct_for_evaluation_function():
    for i, j in ti.ndrange(4, 4):
        result[i] += data[j]

# This is a TI kernel which calls the function in a simple for loop for a given number of iterations.
@ti.kernel
def perform_optimization(iterations: ti.i32):

    for i in range(iterations):
        a_struct_for_evaluation_function()

# Error occurs when calling this function:
#   Calling the function compiles the taichi kernel, but it won't compile.
#   The error states that `struct_for cannot be nested inside a kernel`.
perform_optimization(1000)

The problem is that there's a struct-for inside another for loop.

AFAIK the way the compiler works is that parallelizes only the top-most for loop of any kernel.
Moreover, it is illegal for a struct-for loop (I.E. a ti.ndrange loop) to be nested inside a loop.
This is because the ti.ndrange loop is specifically for parallelizing every execution, but nested loops aren't parallelized.

The question is, therefore:
What is the idiomatic Taichi way of handling doing iterative calls to a ti.ndrange loop?

The simplest solution is to use python to supervise the outer most loop, but this is computationally heavy:

for i in range(1000)
    a_struct_for_evaluation_function() # Note, this would need to be converted to a ti.kernel

So what is the idiomatic taichi way of making this kind of iteration efficient?

Thanks for the help