[LAYOUTS] Generic stmatrix lowering

[lezcano]

We use divideLeft to lower a generic local_store using stmatrix
whenever possible.

We implement ColumnAction as a helper that allows us to permute
the bases and values and remove broadcasting in a generic way.

The current codegen for stmatrix has some limitations, so we just
return early in those cases. We'll fix those in a future PR.

[peterbell10]

What's going on here, seems like some abuse is going on?

[lezcano]

abuse? This is just implementing the address map that ldmatrix asks for:

In particular, the lower three bits of the lane should map to the columns, while the top two bits should map to the 4 different matrices (given by the first 2 basis of the reps of quot). I'll write a comment.

[peterbell10]

But you're setting the register index based on the lane id so clearly you're abusing the labels kReg and kLane to mean something else. That's pretty confusing.

[peterbell10]

I'm a bit concerned that this new feature that is supposed to simplify all our lowerings is resulting in (for me) unreadable code.

[peterbell10]

Okay, I understand now. When using .x4 you pass the pointer corresponding to the row for register T0:r[i] in T(0 + 8*i):p so the lane really is giving the register index. Makes sense. I think I would restate your comment though as:

  // Here we implement the stmatrix.x4 addressing. In particular, the row pointers
  // for each submatrix r in thread t are communicated by the stmatrix call in 
  // laneId = (t // 8) + (8 * r), so r = laneId // 8 and t = laneId % 8.

[lezcano]

I rewrote it in a third way, tell me WDYT.

[lezcano]

fwiw, before merging, I'll implement something so that we just lower via this path if the lowering would have no bank conflicts. Just to make sure we don't over-use this function. I'm also missing to run benchmarks.

[Jokeren]

Looks good to me. The only comment I have is that permuteInDimToFront can be very useful in other places where we need register permutation so probably better to put in LayoutUtils.h

[Jokeren]

Also I think stmatrix use in convertlayout op lowering hasn't been replaced yet. Could be done in another PR though

[peterbell10]

IIUC this is more like findRegPermutationThatDividesLeft?

[lezcano]

sure, I was just looking for a shorter name as this function is going to be used quite a bit. I was also thinking of packing this + the division and action on the regs into a helper function. Might do that tomorrow.

[lezcano]

chose a shorter but representative name

[peterbell10]

I don't find this very helpful because it relates thread id to columns of the matrix, but not to register and lane id which is what you're giving to the linear layout. So this comment is correct, but doesn't explain the code at all.

[lezcano]

The mapping to the 2 bits into the registers is explained in the parens in this line you comment.
The only thing that's missing is noting that each column i starts with thread t[4*i], and since the quotient has already removed the 4 threads, then on this layout hte map is i -> t[i].
The second part is a bit redundant once you have intuition for what divideLeft does tho.

Would you want me to add that?

[peterbell10]

Ah so kLane refers to the lane id with two bases removed such that it's not really a lane id any more. This is why you refer to them by different names now in the comment, but not in the code.

Perhaps It would be less confusing to always act on the "reps" layout? idk

[lezcano]

In this case that'd be equivalent to the current state, generating the same number of ops.

Now, the better codegen would be to simply create a new layout composed of:

• 3 zero bases and he first vec non-zero bases of quot[kReg] (2 int his case)
• The first 3 non-zero basas of `quot[kLane]
• All the warp bases
  We can call these dimensions kVec, kCol, kWarp respectively

This would make sure that you can just pass

applyLinearLayout(loc, rewriter, quot,
                                   {{kVec, laneId},
                                    {kCol, laneId},
                                    {kWarp, warpId}})[0]

This linear layout has just only 5 ones in the columns of kReg and kLane so it would generate pretty low opcount (note that now there are more ones in the kReg dim as we didn't trim them. LLVM may be able to optimise these, but who knows)

[lezcano]

I'm leaving things as-is for now tho. The main thing to note is that, after dividing by a layout, the labels on the resulting layout do not refer to the hardware coords (neither the offsets on the output refer to the actual offset) but they refer to equivalence classes.

In CS terms, this would be like when you have a float ptr and doing increments of 1 with operator[] moves you 4 bytes. Then a division would be like casting it to float4 (grouping a tile of 4 together). Then, the result of the division tells you how to move around in the float4 world, where you have collapsed 4 floats into one element. In particular, in this world, moving 1 moves you 16 bytes.

This is what we are doing but with a more complex tile, where moving 1 over a thread moves us 8 offsets (i.e. one full tile, as per the line regBase = b.shl(regBase, b.i32_val(tile.getTotalOutDimSizeLog2()));)

[lezcano]

nvm, I implemented it. Now the codegen should be better (on our end at least) and hopefully the explanation for why we create the layout should be clearer. WDYT

[lezcano]

cc @peterbell10
There was a test in test_tensor_descriptor.py that was not passing with 2CTAs. Now with this line passes. This line says that we just bail out if there's any interCTA business going on. It seems to work just fine as long as you are addressing things within your own CTA, even without any map or anything.

[lezcano]

This is ready for review:

• It's missing benchmarking, will do tomorrow
• The convertlayout port I'll do it in a different PR, it should be easy.
• The generalisation of the stmatrix to non bf16/f16, other vectorizations, transpose, I'll do in a different PR

[peterbell10]

My remaining comments are non-blocking, we can chat offline.