Affine.dma_start and affine_maps

MLIR
1

Hi folks, i’m looking for some help with using affine.dma_start to represent a DMA operation that does a layout transformation as part of the memory transfer.

The DMA operation i’m attempting to model takes a tensor (let’s say AxB) and does a transpose as part of the copy (to BxA). My approach is to set the src_map attribute to affine map that includes this transpose (e.g. (d0, d1) → (d1, d0)). this allows me to later read the map information when lowering the affine.dma_start to do the right thing.

I believe i may be misusing this property though, or at least missing something, as the src_map (and dst_map) don’t print in the output of the instruction (which means i cannot read in MLIR that sets this attribute).

The code and documentation both conflate the indexing and the affine_map, such that they are displayed and parsed as one value. If my indexing begins at a constant offset (e.g. (0, 0)), applying the map is a no-op.

Is there something i’m missing here? or am i misusing this operation?

thanks in advance,

ian Bearman
Principal Software Engineering Manager
Microsoft Visual C++ Team: ML Optimization & Code Generation
#BlackLivesMatter
/* Making your code faster, smaller, smarter! */

2

@manbearian The src_map and dest_map are (like in the case of affine.load/store ops) applied on the indices and would appear in the subscripts when you print the operation. As an example from test/Dialect/affine/dma.mlir:

 affine.dma_start %0[%i0 + symbol(%arg0), %i1],
                       %1[%i0, %i1 + symbol(%arg1) + 7],
                       %2[%i0 + %i1 + 11], %c64
        : memref<100x100xf32>, memref<100x100xf32, 2>, memref<1xi32>

It’s not the map you’d use to say that you want to be copying the specified source or a part of it transposed. In order to do the latter, you should be encoding the permutation map in the layout map of the source memref.

3

Thanks, that makes a lot of sense.