Ukernel lowering for data-tiled multi_mma with mfma_i32_16x16x32_i8

[bjacob]

This finishes implementing an initial ukernel for multi_mma for DataTiledMMAAttr with kind = mfma_i32_16x16x32_i8.

The ukernel takes unroll and subgroup parameters as function parameters. The idea is that once inlining works as intended, these function parameters will be constants and the optimized code will be the same as if we had hardcoded specific values. This inlining isn't happening at the moment, but that is a bug that we should fix first. It is happening in LLVMCPU, so that's probably something missing in LLVMGPU.

The ukernel file has a comment with a few TODOs to get from this initial naive ukernel to something faster. The first step is to fix the above-mentioned inlining problem, then get shared memory, then get better instruction scheduling.

[bjacob]

Just thought of a problem: while it takes a int unroll_k parameter, the only value that works is 2, because it uses a fixed vector type with that size. I'll fix that tomorrow.

[MaheshRavishankar]

Just thought of a problem: while it takes a int unroll_k parameter, the only value that works is 2, because it uses a fixed vector type with that size. I'll fix that tomorrow.

Drive by comment, aren't the unroll_k etc. implementation details of the compiler? Do they need to cross the ukernel api boundary? I would expect the ukernel to only worry about the problem size/architecture and not these details.

[bjacob]

These unroll_{m,n,k} parameters control the tile shape and layout, of which the corresponding unrolling effect on the kernel code is just an implementation detail. Setting unroll_k = 2 causes the tile to have a 2x larger K-dimension size. See the effect here in the TileSwizzle calculation (which is where both tile shape and layout are decided):

iree/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/GPUTileSwizzleUtils.cpp

Lines 150 to 155 in e553425

	if (mma.getUnrollK() > 1) {
	expand(swizzle, 1, {Kind::CrossIntrinsic, mma.getUnrollK()});
	int interleavingIdx =
	getInnermostNonInternalDimIdx(swizzle.expandShape[1]);
	interleave(swizzle, 1, interleavingIdx);
	}

[hanhanW]

nit: use kConfigAttrName.

iree/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h

Line 47 in e553425

constexpr StringLiteral kConfigAttrName = "lowering_config";

Question: do we preserve all the discardable attributes (i.e., additional attributes that are not defined by the op itself)? If so, you can do something like

newMmaOp->setDiscardableAttrs(mmaOp->getDiscardableAttrDictionary());

[bjacob]

Thanks, that worked! Earlier I had tried something with setAttrs / getAttrs and that caused other tests to fail. I didn't know about setDiscardableAttrs.

[hanhanW]

I did not pay much attention on the changes of DataTiledMMAAttr. Why do we drop the newKind here? Does it impact codegen path? Or it is handled by attribute interface implementation? Thanks for your explanation in advance, and perhaps we can put the information to the PR description.

[bjacob]

We are now transporting the old kind, unchanged. The code being deleted here was creating a new kind that only preserved the unroll_* parameters, but had the subgroups_* parameters set to 1. I thought that those parameters were inherently not needed after thread-distribution. That changed with ukernels. While it remains true that the subgroups_* parameters should not be needed after thread-distribution, in order to avoid needing them, codegen makes use of all the stride information for all the expanded dimensions. We could pass all these strides to the ukernels, but that woud be cumbersome, particularly as the number of dimension varies as unit dimensions are omitted. So in this case, passing the original DataTiledMMAAttr parameters and letting the ukernel infer the strides, results in much simpler code. The drawback is an interaction-at-a-distance in the layouts implied by these parameters, but I think that's OK.

[bjacob]

Just thought of a problem: while it takes a int unroll_k parameter, the only value that works is 2, because it uses a fixed vector type with that size. I'll fix that tomorrow.

Resolved.