Loop Transform Using ISL

[Man Xavier]

I'm new to ISL.

I am trying to transform code
for i in range(10):

    for j in range (10):

        S0: A1[0,j] = max(A1[0,j], A0[i,j])

for i in range(10):

    for j in range (10):

        S1 : A2[i,j] = A0[i,j] - A1[0,j]

TO code

for j in range(10):

    for i in range (10):

        S0: A1[0,j] = max(A1[0,j], A0[i,j])

    for i in range (10):

        S1: A2[i,j] = A0[i,j] - A1[0,j]

Here are my codes:

1. import isl
2. 
   
3. context = isl.set("{ : }")
4. 
   
5. schedule_space = isl.set("{[t0,t1,t2]:}").get_space()
6. precedes = isl.map.lex_lt(schedule_space)
7. 
   
8. # iteration domain
9. domain = isl.union_set(" {S0[i,j] : 0<=i<10 and 0<=j<10; S1[i,j] : 0<=i<10 and 0<=j<10; }")
10. 
    
11. schedule = isl.union_map("{S0[i,j] -> [t0,t1,t2] : t0 = 1 and t1 = i and t2 = j; S1[i,j] -> [t0,t1,t2] : t0 = 2 and  t1 = i and t2 = j}")
12. 
    
13. build = isl.ast_build.from_context(context)
14. ast = build.node_from_schedule_map(schedule.intersect_domain(domain))
15. print(ast.to_C_str())
16. 
    
17. reads = isl.union_map("{S0[i,j] -> A0[a=i,b=j]; S0[i,j] -> A1[a=0,b=j];S1[i,j] -> A1[a=0,b=j];}")
18. reads= reads.intersect_domain(domain)
19. 
    
20. writes = isl.union_map("{S0[i,j] -> A1[a=0,b=j]; S1[i,j] -> A2[a=i,b=j]; S[j] -> A1[0,b=j]}")
21. writes= writes.intersect_domain(domain)
22. 
    
23. raw = writes.apply_range(reads.reverse())
24. raw = raw.apply_domain(schedule).apply_range(schedule)
25. raw = raw.intersect(precedes)
26. raw = raw.apply_domain(schedule.reverse()).apply_range(schedule.reverse())
27. 
    
28. 
    
29. war = reads.apply_range(writes.reverse())
30. war = war.apply_domain(schedule).apply_range(schedule)
31. war = war.intersect(precedes)
32. war = war.apply_domain(schedule.reverse()).apply_range(schedule.reverse())
33. 
    
34. 
    
35. waw = writes.apply_range(writes.reverse())
36. waw = waw.apply_domain(schedule).apply_range(schedule)
37. waw = waw.intersect(precedes)
38. waw = waw.apply_domain(schedule.reverse()).apply_range(schedule.reverse())
39. 
    
40. 
    
41. sc = isl.schedule_constraints.on_domain(domain)
42. dep = war.union(raw).union(waw)
43. 
    
44. sc = sc.set_validity(dep)
45. sc = sc.set_proximity(dep)
46. sched = sc.compute_schedule()
47. build = isl.ast_build.from_context(context)
48. ast = build.node_from(sched)
49. print(ast.to_C_str())

The experimental result is that : I can get the expected AST WITHOUT line 18 and line 21.

However,  WITH line 18 and line 21, I just get the AST like this: 
 for (int c0 = 0; c0 <= 9; c0 += 1)
    for (int c1 = 0; c1 <= 9; c1 += 1)
      S0(c1, c0);
  for (int c0 = 0; c0 <= 9; c0 += 1)
    for (int c1 = 0; c1 <= 9; c1 += 1)
      S1(c0, c1);
Why does this happen? What is the correct way to get the expected AST?

[Sven Verdoolaege]

On Wed, Nov 12, 2025 at 06:03:23AM -0800, Man Xavier wrote:
> *I'm new to ISL.*
> *I am trying to transform code*

> for i in range(10):
> for j in range (10):
> S0: A1[0,j] = max(A1[0,j], A0[i,j])
> for i in range(10):
> for j in range (10):
> S1 : A2[i,j] = A0[i,j] - A1[0,j]

> *TO code*

> for j in range(10):
> for i in range (10):
> S0: A1[0,j] = max(A1[0,j], A0[i,j])
> for i in range (10):
> S1: A2[i,j] = A0[i,j] - A1[0,j]
>

[..]
>
> The experimental result is that : I can get the expected AST *WITHOUT *line
> 18 and line 21.
> However, *WITH* line 18 and line 21, I just get the AST like this:

> for (int c0 = 0; c0 <= 9; c0 += 1)
> for (int c1 = 0; c1 <= 9; c1 += 1)
> S0(c1, c0);
> for (int c0 = 0; c0 <= 9; c0 += 1)
> for (int c1 = 0; c1 <= 9; c1 += 1)
> S1(c0, c1);
> Why does this happen? What is the correct way to get the expected AST?

You'll have to explain in a bit more detail what you are trying to do and
why you expect this to result in the desired schedule.

All I can say from a quick glance is that your proximity constraints
look very strange. You are trying to put too many pairs of statement
instances close to each other, so the scheduler is likely just giving
up on them.

skimo

[Man Xavier]

Thanks for your instant reply!

The code I'm trying to transform is part of safe softmax. 

(Actually, I'm doing optimizations targeting NPU, but the it can be simplified to this problem.)
I'm trying to optimize the time locality of access A1. In other words, A1[0,j] should be used as soon as A1[0,j] is computed completely.

However, it is complex to do do this optimization:

We can inspire that A1[0,j] is computed completely only after traversing iterator i.

Therefore, we should interchange loop i and loop j to traverse iterator i first.

After loop interchange, the outermost loops can be fused, ensuring the data dependencies.

[Sven Verdoolaege]

On Wed, Nov 12, 2025 at 05:31:46PM -0800, Man Xavier wrote:
> Thanks for your instant reply!
> The code I'm trying to transform is part of safe softmax.
> (Actually, I'm doing optimizations targeting NPU, but the it can
> be simplified to this problem.)
> I'm trying to optimize the time locality of access A1. In other words,
> A1[0,j] should be used as soon as A1[0,j] is computed completely.

If you want to use the scheduler to compute a schedule for you
based on this locality, then you should set the proximity constraints
accordingly.
Your proximity constaints are

{ S0[i, j] -> S1[i', j' = j] : 0 <= i <= 9 and 0 <= j <= 9 and 0 <= i' <= 9; S0[i, j] -> S0[i', j' = j] : i >= 0 and 0 <= j <= 9 and i < i' <= 9 }

but, AFAICT, based on your description you should use something more like

{ S0[9, j] -> S1[i, j] : 0 <= j <= 9 and 0 <= i <= 9; S0[i, j] -> S0[i', j] : i >= 0 and 0 <= j <= 9 and i < i' <= 9 }

Like so:

$ ./isl_schedule | ./isl_codegen
{
domain: "{ S0[i, j] : 0 <= i <= 9 and 0 <= j <= 9; S1[i, j] : 0 <= i <= 9 and 0 <= j <= 9 }",
validity: "{ S0[i, j] -> S1[i', j' = j] : 0 <= i <= 9 and 0 <= j <= 9 and 0 <= i' <= 9; S0[i, j] -> S0[i', j' = j] : i >= 0 and 0 <= j <= 9 and i < i' <= 9 }",
proximity: "{ S0[9, j] -> S1[i, j] : 0 <= i <= 9 and 0 <= j <= 9; S0[i, j] -> S0[i', j] : i >= 0 and 0 <= j <= 9 and i < i' <= 9 }"

}
for (int c0 = 0; c0 <= 9; c0 += 1) {
for (int c1 = 0; c1 <= 9; c1 += 1)
S0(c1, c0);

for (int c1 = 9; c1 <= 18; c1 += 1)
S1(c1 - 9, c0);
}

You can compute these flow dependencies using the functions described
in section "Dependence Analysis" (use reads as sinks and (must-)writes
as must-sources).

If, on the other hand, you want to construct a schedule manually,
you can do that using the functions in "Schedule Trees".

skimo