[RFC] Compile-time memref.alloc Scheduling/Merging optimization

MLIR
21

I have seen memory plan pass implemented in many mlir projects, it will be better if milr support it natively.

22

The term “allocation scope” is introduced to make it easier to handle the parallel loops which is a special case for the allocation.

Note that in a simple serial for-loop, you can alloca a single buffer for a buffer allocated within the loop:

for(...) {
   float a[100];
   float b[100];
   run(a,b);
}

can be optimized to

float merged[100+100];
for(...) {
   run(merged, merged+100);
}

However the optimization does not apply for the parallel-for.

paralle_for(...) {
   float a[100];
   float b[100];
   run(a,b);
}

// the below is incorrect:
float merged[100+100];
parallel_for(...) {
   run(merged, merged+100);
}

It is because temp buffer float a[100] (and b also) is defined as “thread-local”. You cannot lift it out of the loop into a simple merged buffer like float merged[100+100]. Otherwise, data races may occur - multiple threads may operate on the same range of memory, which is not expected in the original code.

What we propose is that, we define a “scope” for the parallel-for, and buffers inside the parallel-for are merged to that scope, instead of the one outside of the parallel-for.

parallel_for(...) {
   float merged[100+100];
   run(merged, merged+100);
}

You can further optimize this code, if the number of threads is known at compile-time, say num_threads=8. You can lift the merged buffer out of the parallel for loop, and even merge buffers with other buffers outside of the parallel-for.

float c[100]
last_use_of_c(c);
paralle_for(...) {
   float a[100];
   float b[100];
   run(a,b);
}

// the below is correct for num_thread=8:
float merged[(100+100)*8];
float *c = &merged[0];
last_use_of_c(c);
parallel_for(...) {
   float* a = merged[tid*200]
   float* b = merged[tid*200+100]
   run(a, b);
}

Note that the feature “lifting marged buffer outside of the paralle-for” is not included in the PR.

23

thanks for your reply.

paralle_for(…) {
float a[100];
float b[100];
run(a,b);
}
// the below is incorrect:
float merged[100+100];
parallel_for(…) {
run(merged, merged+100);
}

i’m sorry, but it still confuses me about data race.
it is because each thread will access same memory area?
if so, the reason should be wrong IR is generated ?
the correct IR should be:
float merged[(100+100)NumOfThread];
parallel_for(…) {
thread_offset = IdOfThread(100 + 100)
run(merged+thread_offset, merged+thread_offset+100)
}
is my understanding correct?

24

Yes. Your understanding is correct.

25

many thanks.