[0035] Alignment guarantee for Matrix Base, Offset, Stride and Size

[anupamachandra]

Which document does this relate to?
https://github.com/microsoft/hlsl-specs/blob/main/proposals/0035-linalg-matrix.md

Describe the issue you see with the spec
Alignment guarantees for Matrix and Vector Accumulate buffers missing in the linAlg spec.

Additional context
In the earlier cooperative-vector spec buffers had the following alignment guarantees:

Matrix Buffer(for Mul/MulAdd/OuterProduct):
The matrix base address and offset guaranteed to be aligned to at least 128 bytes, and the matrix stride is guaranteed to be aligned to at least 16 bytes and the matrix size was guaranteed to be a multiple of 16 bytes ensuring that the *16 bytes access of the last row/column of the matrix is valid memory.

OutputArray (Vector Accumulate):
The base address and output array offset must be 64 byte aligned. Also note that the size of the underlying allocation is guaranteed to be a multiple of 16 bytes, ensuring that there is valid memory between the end of the array and the 16 byte boundary. Implementations may write to the contents of the padding between the end of the matrix and the 16 byte boundary, so developers should not use this padding space for anything else.

This allows for efficient martrix loads/stores and cases where implementations might require padding the data due to h/w constraints. This guanratee is missing in the current linAlg spec. What we have in its place are align parameters for the relevant HLSL/DXIL ops, while this can be used by users to inform the driver/compiler of the alignment in conjuction with proper allocation of these resources.. just having the restrictions from the older spec without the flexibility seems adequate for most uses, provides portable performance and is sufficiently futureproof.

[tex3d]

Clarifications under "Matrix Buffer":

• Shouldn't these apply to the [RW]ByteAddressBuffer operations: linAlgMatrixLoadFromDescriptor, linAlgMatrixStoreToDescriptor, and linAlgMatrixAccumulateToDescriptor operations? If not, can you specify the exact DXIL operations?
• Perhaps instead of "matrix base address and offset" you mean to say "absolute GPUVA of the matrix Offset"?
• "ensuring that the 16 bytes access of the last row/column of the matrix is valid memory." - Perhaps "ensuring that the entire 16-byte aligned size of every row/column of the matrix is in-bounds."? (and resident?) Or would it require the entire stride of the last row/column to be in-bounds (which may be larger than the aligned row/column size)?
• For linAlgMatrixStoreToDescriptor should it note that the 16-byte aligned size of rows/columns may be overwritten, even if the row/column size is smaller and there's padding (like for the groupshared case)?

Clarifications under "OutputArray":

• By "OutputArray" do you mean linAlgMatrixAccumulateToMemory for groupshared? If so, then why did you write Vector Accumulate?
• The driver can ensure that the base address of a particular groupshared global is 64-byte aligned in the groupshared allocation for any array accessed by the operation, so it would seem that the requirement is that the Offset * component size is aligned by 64-bytes.
• There's a wrinkle to requiring base array alignment of 64-bytes, if this creates scenarios where alignment would require gaps between allocated groupshared globals. We could require that the array byte size is also aligned by 64-bytes, then as long as the driver can move/sort groupshared globals, gaps can be avoided. Otherwise, if there are gaps, we won't account for these gaps in our groupshared memory usage computation for runtime validation.

[anupamachandra]

Clarifications under "Matrix Buffer":

Shouldn't these apply to the [RW]ByteAddressBuffer operations: linAlgMatrixLoadFromDescriptor, linAlgMatrixStoreToDescriptor, and linAlgMatrixAccumulateToDescriptor operations? If not, can you specify the exact DXIL operations?
Yes, you are right. I was thinking of the CoopVec operations with loads included as a part of the operation. The operations you have above would cover all the cases for Thread/Wave/TG scope. Just one clarification, the linAlgMatrixAccumulateToDescriptor op description doesn't call out component-wise atomicity. That's just an oversight? The interlockedAccumulate does call out atomicity, but misses component-wise which might be good to add back to avoid confusion.

Perhaps instead of "matrix base address and offset" you mean to say "absolute GPUVA of the matrix Offset"?
Since we are going to add this restriction to the runtime spec, I thought it made sense to call out the BaseAdress alignment and Offset alignment separately. Programmers should ensure that the resource allocated is 128B aligned and the offsets in HLSL are a multiple of 128B. What say?

"ensuring that the 16 bytes access of the last row/column of the matrix is valid memory." - Perhaps "ensuring that the entire 16-byte aligned size of every row/column of the matrix is in-bounds."? (and resident?) Or would it require the entire stride of the last row/column to be in-bounds (which may be larger than the aligned row/column size)?
How is this: "Each row or column occupies at least its logical byte size rounded up to 16 bytes, and that whole rounded-up region must be valid and resident; for stores/accumulates, implementations may overwrite bytes in that rounded-up padding region." What we want to guarantee are efficient loads/stores which we get with the 16B guarantees and in cases like fp16 matrices with odd dimensions and the h/w can only perform fp32 atomics, we can overwrite the padding region.

For linAlgMatrixStoreToDescriptor should it note that the 16-byte aligned size of rows/columns may be overwritten, even if the row/column size is smaller and there's padding (like for the groupshared case)?
Yeah, for the reason above we should have that restriction.

Clarifications under "OutputArray":

By "OutputArray" do you mean linAlgMatrixAccumulateToMemory for groupshared? If so, then why did you write Vector Accumulate?
The cooperative vector had a VectorAccumulate function to atomically accumulate an array into memory, to support bias vector accumulation. This was missing from the linAlg spec and I had filed an issue for it and it was verbally approved but looks like the spec hasn't been updated. It's an important piece we'd like added for the preview.

The driver can ensure that the base address of a particular groupshared global is 64-byte aligned in the groupshared allocation for any array accessed by the operation, so it would seem that the requirement is that the Offset x component size is aligned by 64-bytes.
There's a wrinkle to requiring base array alignment of 64-bytes, if this creates scenarios where alignment would require gaps between allocated groupshared globals. We could require that the array byte size is also aligned by 64-bytes, then as long as the driver can move/sort groupshared globals, gaps can be avoided. Otherwise, if there are gaps, we won't account for these gaps in our groupshared memory usage computation for runtime validation.
Yep, The driver compiler can control the base array alignment in Group shared arrays (for most implementations I think), so we don't the restriction there. A 16B offset and stride requirement for optimizing loads will be good. Another point here, relevant tangentially, is that Wave and ThreadGroup Scope matrix processing will likely require some internal shared memory usage by the driver compier for base perf, so developers will have to account for that when using shared memory in the shader with linAlg ops. This in conjunction with the variable shared memory spec would mean that the maximum published shared memory might not be available to developers. This would be a good device property check in the runtime query, so an implementation can publish how much GS memory they need for internal linAlg implementation and developers can adjust their shared memory usage based on that to avoid perf degradation. @jenatali for viz

[anupamachandra]

Also tagging @pow2clk for viz.

[tex3d]

the linAlgMatrixAccumulateToDescriptor op description doesn't call out component-wise atomicity.

Sounds like a spec oversight. Can you file an issue?

Programmers should ensure that the resource allocated is 128B aligned and the offsets in HLSL are a multiple of 128B. What say?

That sounds fine.

How is this: "Each row or column occupies at least its logical byte size rounded up to 16 bytes, and ...

This sounds good, though we probably need to note what happens when any part is out of bounds for a bounds-checked resource, right? Which is likely: for reads: reads all zeros, for writes/atomic: writes nothing is written (entire operation is predicated by whole bounds-check). For non-bounds checked resources, behavior is undefined if not entirely in-bounds and resident as specified.

an implementation can publish how much GS memory they need for internal linAlg implementation

Won't this depend on the linalg operations performed (and perhaps dimensions/component types)? How would this be exposed through/to the runtime? I guess this is additional detail to work out. We should have an issue to track it.

[tex3d]

Another point discussed, but not mentioned here: Should we remove alignment fields from the API and DXIL ops if the required alignment is sufficient for any implementation (no implementation would benefit from higher alignment values)?

Initially, I thought there wouldn't be any purpose to keeping the alignment parameter in the DXIL op, but if we were to relax the strict alignment requirements for any reason in the future, we would want to have the parameter.

So, I'm leaning in the direction of keeping this immediate parameter in the DXIL op, removing it from the HLSL API which would automatically supply the required minimum alignment to the builtin, and enforcing the minimum alignment for the DXIL op in DxilValidation.