Wrap Innermost Loop as neura.kernel

[ShangkunLi]

Hi~ @HobbitQia, I enhance the convert-taskflow-to-neura pass with two modes:

• 'hyperblock' mode: this mode runs the flow with counters
• 'innermost' mode: this mode wraps each innermost loop as the neura.kernel; this could be a starting point for your fusion

[tancheng]

Can you provide IR example for convert-taskflow-to-neura?

I thought hyperblock would always be there when we generate task, and will be lowered to neura.kernel after the convert-taskflow-to-neura. Then, the neura.kernel would be the entry point for mapper. So why do we need two modes? The neura.kernel with counters can also be a target of our mapper, right?

[ShangkunLi]

Can you provide IR example for convert-taskflow-to-neura?

I thought hyperblock would always be there when we generate task, and will be lowered to neura.kernel after the convert-taskflow-to-neura. Then, the neura.kernel would be the entry point for mapper. So why do we need two modes? The neura.kernel with counters can also be a target of our mapper, right?

The innermost mode is to provide a starting point for @HobbitQia, because there are no counter ops in his architecture.

Here is a transformed ir using --convert-taskflow-to-neura="mode=innermost":

%write_outputs = taskflow.task @Task_0 read_memrefs(%arg0 : memref<?x8x6xi32>) write_memrefs(%arg5 : memref<?xi32>) [original_read_memrefs(%arg0 : memref<?x8x6xi32>), original_write_memrefs(%arg5 : memref<?xi32>)] : (memref<?x8x6xi32>, memref<?xi32>) -> (memref<?xi32>) {
    ^bb0(%arg10: memref<?x8x6xi32>, %arg11: memref<?xi32>):
      affine.for %arg12 = 0 to 4 {
        affine.for %arg13 = 0 to 8 {
          neura.kernel inputs(%arg10, %arg12, %arg13, %arg11 : memref<?x8x6xi32>, index, index, memref<?xi32>) attributes {kernel_name = "kernel_0"} {
          ^bb0(%arg14: memref<?x8x6xi32>, %arg15: index, %arg16: index, %arg17: memref<?xi32>):
            %c0 = arith.constant 0 : index
            %c6 = arith.constant 6 : index
            %c1 = arith.constant 1 : index
            scf.for %arg18 = %c0 to %c6 step %c1 {
              %1 = memref.load %arg14[%arg15, %arg16, %arg18] : memref<?x8x6xi32>
              memref.store %1, %arg17[%arg18] : memref<?xi32>
            }
            neura.yield
          }
        }
      }
      taskflow.yield writes(%arg11 : memref<?xi32>)
    }

If we use the architecture that has counter, the transformation process should use the original one -- --convert-taskflow-to-neura="mode=hyperblock". The transfromed ir:

    %write_outputs = taskflow.task @Task_0 read_memrefs(%arg0 : memref<?x8x6xi32>) write_memrefs(%arg5 : memref<?xi32>) [original_read_memrefs(%arg0 : memref<?x8x6xi32>), original_write_memrefs(%arg5 : memref<?xi32>)] : (memref<?x8x6xi32>, memref<?xi32>) -> (memref<?xi32>) {
    ^bb0(%arg10: memref<?x8x6xi32>, %arg11: memref<?xi32>):
      %1 = taskflow.counter attributes {counter_id = 0 : i32, counter_type = "root", lower_bound = 0 : index, step = 1 : index, upper_bound = 4 : index} : index
      %2 = taskflow.counter parent(%1 : index) attributes {counter_id = 1 : i32, counter_type = "relay", lower_bound = 0 : index, step = 1 : index, upper_bound = 8 : index} : index
      %3 = taskflow.counter parent(%2 : index) attributes {counter_id = 2 : i32, counter_type = "leaf", lower_bound = 0 : index, step = 1 : index, upper_bound = 6 : index} : index
      neura.kernel inputs(%arg10, %arg11 : memref<?x8x6xi32>, memref<?xi32>) {
      ^bb0(%arg12: memref<?x8x6xi32>, %arg13: memref<?xi32>):
        %4 = neura.counter {counter_id = 0 : i32, counter_type = "root", lower_bound = 0 : index, step = 1 : index, upper_bound = 4 : index} : index
        %5 = neura.counter {counter_id = 1 : i32, counter_type = "relay", lower_bound = 0 : index, step = 1 : index, upper_bound = 8 : index} : index
        %6 = neura.counter {counter_id = 2 : i32, counter_type = "leaf", lower_bound = 0 : index, step = 1 : index, upper_bound = 6 : index} : index
        %7 = memref.load %arg12[%4, %5, %6] : memref<?x8x6xi32>
        memref.store %7, %arg13[%6] : memref<?xi32>
        neura.yield
      }
      taskflow.yield writes(%arg11 : memref<?xi32>)
    }

We can also support architectures with counter FUs by utilizing this ir.

[tancheng]

I still think mode is unnecessary. We should have support_counter in arch_spec.yaml to determine which mode we are targeting. Instead of using different flag. How does this sound?

[ShangkunLi]

I still think mode is unnecessary. We should have support_counter in arch_spec.yaml to determine which mode we are targeting. Instead of using different flag. How does this sound?

Hmmm, but this should diverges from higher level. Because for counter-based neura.kernel, it should be transformed through:

--convert-affine-to-taskflow
--contruct-hyperblock-from-task
--convert-taskflow-to-neura="mode=hyperblock"

For the innermost-loop-based neura.kernel, it should be transformed through:

--convert-affine-to-taskflow
--convert-taskflow-to-neura="mode=innermost"

They actually correspond to two different pass pipelines, which should be handled by the neura-compiler instead of the mlir-neura-opt?

[tancheng]

So you mean mode=innermost will have no hyperblock exist?

The principle should be:

• neura-compiler and mlir-neura-opt should have the same functionality. opt is just pick some passes for verification. we should have single compiler no matter what use case. All passes should be applied on different inputs, just some passes might not make impact on the IRs.

So both opt and compiler should have all 3 passes enabled:

--convert-affine-to-taskflow
--contruct-hyperblock-from-task
--convert-taskflow-to-neura

And hyperblock should exist in both cases. Is this possible?

[ShangkunLi]

So you mean mode=innermost will have no hyperblock exist?

The principle should be:

• neura-compiler and mlir-neura-opt should have the same functionality. opt is just pick some passes for verification. we should have single compiler no matter what use case. All passes should be applied on different inputs, just some passes might not make impact on the IRs.

So both opt and compiler should have all 3 passes enabled:

--convert-affine-to-taskflow
--contruct-hyperblock-from-task
--convert-taskflow-to-neura

And hyperblock should exist in both cases. Is this possible?

Good idea~!

I have already modified the conversion pipeline to this pipeline.