Plumb pooling ops through Linalg on tensors path.

[hanhanW]

1. Set pooling ops as root op.
2. Convert linalg.init_tensor to a AllocOp.
3. Set pooling ops launch config for Linalg on tensors path.

The (2) is only created for shaped operand. The alloc op will be deleted
once lower to loops. This is the same behavior as Linalg on buffers
path.

Part of #5043

[hanhanW]

This will need to replace AllocOp with memref::AllocOp after the google -> main happens.

[hanhanW]

To be honest, I don't really understand how to set this and what the constraints are...

Should the size of ts should be as same as the number of parallel loops in Linalg on tensors path?
What is the relationship between this and workloadPerWorkgroup? I tried 0, tileSizeY, tileSizeX, 0, but it would hit an assertion:

Assertion `workloadPerWorkgroup.size() <= 3 && "workloadPerWorkgroup size greater than 3 not handled"`

Now I'm just luckily having a config to make it work. I'm happy to learn more about this.

[MaheshRavishankar]

Yes, Hanhan, you are right about the fact that this is fairly haphazard. Its on my list of things to clean up.

The overall goal here is that you only want to partition upto 3 loops at the Flow level. The HAL layer can only support 3 levels of processor x,y,z. We can have Flow in theory handle any number of parallelism mode, but it gets tricky then to lower from Flow -> HAL. So for now the assumption is to try to keep 3 levels of parallelism at flow level.

W.R.T to how to decide which loops to parallelize and which not to parallelize, and what tile size to use. We should uniformly try to say tile size of 0 implies it is not parallelized, both in the Flow level when the distrubution happens, and here when setting the configuration.

Let me look through your change, in a bit. Its the thing I want to get to next, so I will have to see whats done currently and then give you feedback here.

[antiagainst]

It's complicated because of three main reasons: 1) existing buffer vs. new tensor path, 2) convolution/pooling ops has more than 3 parallel dimensions but right now we can only support <= 3 Flow tiling, 3) the tile sizes here matches linalg op implicit loop order, which is typically the reverse of distribution order.

In the existing buffer path by default we always tile along the first three parallel dimensions, namely N, OH, OW. So that's why you see we only have three numbers here for the existing buffer path: they are for the first 3 implicit loop of conv/pooling ops, namely N, OH, OW. That's not very good because N is typically 1 in inference so we are effectively dropping one dimension of parallelism. That's why we want to tile along the second three parallel dimensions, namely OH, OW, OC. The new tensor based approach adopts this scheme for conv, by having the special check in DispatchLinalgOnTensorsPass to force tiling along the second three parallel dimensions, namely OH, OW, OC. Now remember that these tiling configurations need to match with implicit loop order in linalg ops. So we need a special leading 0 to mean disabling tiling on the first parallel dimension, namely N. Given you made pooling ops be treated similarly like convolution ops, you need the same thing here. Hope this clarifies. :)

It will improve a bit once we can remove the existing buffer path. But the special checks in DispatchLinalgOnTensorsPass will remain. I would actually think making flow level workgroups really n-D (n being arbitrary number) good to address the special checks. I was doing that in the initial prototype. Will send out a pull request later to show how it looks and whether we should proceed with it.

[hanhanW]

Now I understand more about these, thanks for the detailed explanantion!

[MaheshRavishankar]

Yes, Hanhan, you are right about the fact that this is fairly haphazard. Its on my list of things to clean up.

The overall goal here is that you only want to partition upto 3 loops at the Flow level. The HAL layer can only support 3 levels of processor x,y,z. We can have Flow in theory handle any number of parallelism mode, but it gets tricky then to lower from Flow -> HAL. So for now the assumption is to try to keep 3 levels of parallelism at flow level.

W.R.T to how to decide which loops to parallelize and which not to parallelize, and what tile size to use. We should uniformly try to say tile size of 0 implies it is not parallelized, both in the Flow level when the distrubution happens, and here when setting the configuration.

Let me look through your change, in a bit. Its the thing I want to get to next, so I will have to see whats done currently and then give you feedback here.

[MaheshRavishankar]

You probably need to check if the result is already mapped to a buffer here. There is a pre-processing step for Linalg operations that tries to map the result buffer (if it is already mapped) to the outs operand. So here this will be undone.

[hanhanW]

I see, thanks!

[MaheshRavishankar]

Hey Hanhan, this is probably needed before this can land.

[hanhanW]

Does if (bvm.contains(initTensorOp.getResult())) return success(); address the comment? I thought bvm contains the information of mapping between outs and buffers.

[MaheshRavishankar]

Right. That works. Missed that.

[antiagainst]

It's complicated because of three main reasons: 1) existing buffer vs. new tensor path, 2) convolution/pooling ops has more than 3 parallel dimensions but right now we can only support <= 3 Flow tiling, 3) the tile sizes here matches linalg op implicit loop order, which is typically the reverse of distribution order.

In the existing buffer path by default we always tile along the first three parallel dimensions, namely N, OH, OW. So that's why you see we only have three numbers here for the existing buffer path: they are for the first 3 implicit loop of conv/pooling ops, namely N, OH, OW. That's not very good because N is typically 1 in inference so we are effectively dropping one dimension of parallelism. That's why we want to tile along the second three parallel dimensions, namely OH, OW, OC. The new tensor based approach adopts this scheme for conv, by having the special check in DispatchLinalgOnTensorsPass to force tiling along the second three parallel dimensions, namely OH, OW, OC. Now remember that these tiling configurations need to match with implicit loop order in linalg ops. So we need a special leading 0 to mean disabling tiling on the first parallel dimension, namely N. Given you made pooling ops be treated similarly like convolution ops, you need the same thing here. Hope this clarifies. :)

It will improve a bit once we can remove the existing buffer path. But the special checks in DispatchLinalgOnTensorsPass will remain. I would actually think making flow level workgroups really n-D (n being arbitrary number) good to address the special checks. I was doing that in the initial prototype. Will send out a pull request later to show how it looks and whether we should proceed with it.

[MaheshRavishankar]

It will improve a bit once we can remove the existing buffer path. But the special checks in DispatchLinalgOnTensorsPass will remain. I would actually think making flow level workgroups really n-D (n being arbitrary number) good to address the special checks. I was doing that in the initial prototype. Will send out a pull request later to show how it looks and whether we should proceed with it.

@antiagainst I would push back against going down this path. There is really no reason we should have some non-trivial logic to convert nD distribution to 3D distribution in the Flow to HAL conversion. It is less surprising to just expose only three levels of parallelism at the Flow level itself. As we have discussed before, eventually everything needs to map to a 3D grid cause that is what hardware supports. We know what op we are targeting and would just be cleaner if we did not partition the operation to more than 3D to begin with by using the tile size of 0. It makes things less surprising. The only use case I can think of for considering more than 3D partitioning is if a HAL backends needs more than 3D partitioning.

[MaheshRavishankar]

Hey Hanhan, this is probably needed before this can land.

[antiagainst]

It will improve a bit once we can remove the existing buffer path. But the special checks in DispatchLinalgOnTensorsPass will remain. I would actually think making flow level workgroups really n-D (n being arbitrary number) good to address the special checks. I was doing that in the initial prototype. Will send out a pull request later to show how it looks and whether we should proceed with it.

@antiagainst I would push back against going down this path. There is really no reason we should have some non-trivial logic to convert nD distribution to 3D distribution in the Flow to HAL conversion. It is less surprising to just expose only three levels of parallelism at the Flow level itself. As we have discussed before, eventually everything needs to map to a 3D grid cause that is what hardware supports. We know what op we are targeting and would just be cleaner if we did not partition the operation to more than 3D to begin with by using the tile size of 0. It makes things less surprising. The only use case I can think of for considering more than 3D partitioning is if a HAL backends needs more than 3D partitioning.

Understood your points. Not saying I'm super attached to that and must go with it. (I didn't implement that way when plumbing through conv support. ;-P) Note that my points were just to send a pull request to show how it looks like, given I initially have it already and I think it's not that complicated.

It is less surprising to just expose only three levels of parallelism at the Flow level itself.

+1. We claim flow workgroups can be arbitrary-D but in reality it can only use <= 3; so that's just a confusing point. If we aren't gonna use it, it's better to drop that. :)

[antiagainst]

Clear from me. But please address Mahesh's comments.

[hanhanW]

Rebase and delete special logic that is handled in #5172

[hanhanW]

@MaheshRavishankar please take another look. This rebased on a commit after 9f968d0. This needs your approval because you requested changes. Thanks!

[MaheshRavishankar]

Hey Hanhan, AFAICS, there is still one comment w.r.t to InitTensor initialization that still needs to be addressed?

[hanhanW]

I probably misunderstood something. Does if (bvm.contains(initTensorOp.getResult())) return success(); address the comment? I thought bvm contains the information of mapping between outs and buffers.

[MaheshRavishankar]

Right. That works. Missed that.