Allow same memory space for SRC and DST of dma_start operations

[manbearian]

Summary

This change allows the SRC and DST of dma_start operations to be located in the
same memory space. This applies to both the Affine dialect and Memref dialect
versions of these Ops. The documention has been updated to reflect this by
explicitly stating overlapping memory locations are not supported (undefined
behavior).

[manbearian]

See this discussion for more context: https://llvm.discourse.group/t/allowing-dma-start-between-same-memory-space/3444

[manbearian]

updated "invalid op" test to model new behavior

[manbearian]

@bondhugula , thanks for the sign-off. What are the next steps to get this code merged, this is my first contribution to LLVM or MLIR?

thanks,
ian

mlir/include/mlir/Dialect/Affine/IR/AffineOps.h

 /// destination memref, and the tag memref have the same restrictions as any
 /// affine.load/store. In particular, index for each memref dimension must be an
 /// affine expression of loop induction variables and symbols.
 /// The optional stride arguments should be of 'index' type, and specify a
 /// stride for the slower memory space (memory space with a lower memory space
 /// id), transferring chunks of number_of_elements_per_stride every stride until
 /// %num_elements are transferred. Either both or no stride arguments should be
 /// specified. The value of 'num_elements' must be a multiple of
-/// 'number_of_elements_per_stride'.
+/// 'number_of_elements_per_stride'. If the source and destination locations
+/// overlap the behavior of this operation is not defined.
 //
 // For example, an AffineDmaStartOp operation that transfers 256 elements of a
 // memref '%src' in memory space 0 at indices [%i + 3, %j] to memref '%dst' in
 // memory space 1 at indices [%k + 7, %l], would be specified as follows:
 //
 //   %num_elements = constant 256
 //   %idx = constant 0 : index
 //   %tag = alloc() : memref<1xi32, 4>

mlir/include/mlir/Dialect/MemRef/IR/MemRef.h

 // elemental type of the memref), a tag memref with its indices, and optionally
 // at the end, a stride and a number_of_elements_per_stride arguments. The tag
 // location is used by a DmaWaitOp to check for completion. The indices of the
 // source memref, destination memref, and the tag memref have the same
 // restrictions as any load/store. The optional stride arguments should be of
 // 'index' type, and specify a stride for the slower memory space (memory space
 // with a lower memory space id), transferring chunks of
 // number_of_elements_per_stride every stride until %num_elements are
-// transferred. Either both or no stride arguments should be specified.
+// transferred. Either both or no stride arguments should be specified. If the
+// source and destination locations overlap the behavior of this operation is
+// not defined.
 //
 // For example, a DmaStartOp operation that transfers 256 elements of a memref
 // '%src' in memory space 0 at indices [%i, %j] to memref '%dst' in memory space
 // 1 at indices [%k, %l], would be specified as follows:
 //
 //   %num_elements = constant 256
 //   %idx = constant 0 : index
 //   %tag = alloc() : memref<1 x i32, (d0) -> (d0), 4>

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

 LogicalResult AffineDmaStartOp::verify() {
   if (!getOperand(getSrcMemRefOperandIndex()).getType().isa<MemRefType>())
     return emitOpError("expected DMA source to be of memref type");
   if (!getOperand(getDstMemRefOperandIndex()).getType().isa<MemRefType>())
     return emitOpError("expected DMA destination to be of memref type");
   if (!getOperand(getTagMemRefOperandIndex()).getType().isa<MemRefType>())
     return emitOpError("expected DMA tag to be of memref type");
 
-  // DMAs from different memory spaces supported.
-  if (getSrcMemorySpace() == getDstMemorySpace()) {
-    return emitOpError("DMA should be between different memory spaces");
-  }
   unsigned numInputsAllMaps = getSrcMap().getNumInputs() +
                               getDstMap().getNumInputs() +
                               getTagMap().getNumInputs();
   if (getNumOperands() != numInputsAllMaps + 3 + 1 &&
       getNumOperands() != numInputsAllMaps + 3 + 1 + 2) {
     return emitOpError("incorrect number of operands");
   }
 

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

LogicalResult DmaStartOp::verify() {
   if (numOperands < numExpectedOperands)
     return emitOpError() << "expected at least " << numExpectedOperands
                          << " operands";
   if (!getTagIndices().empty() &&
       !llvm::all_of(getTagIndices().getTypes(),
                     [](Type t) { return t.isIndex(); }))
     return emitOpError("expected tag indices to be of index type");
 
-  // DMAs from different memory spaces supported.
-  if (getSrcMemorySpace() == getDstMemorySpace())
-    return emitOpError("DMA should be between different memory spaces");
-
   // Optional stride-related operands must be either both present or both
   // absent.
   if (numOperands != numExpectedOperands &&
       numOperands != numExpectedOperands + 2)
     return emitOpError("incorrect number of operands");
 
   // 5. Strides.
   if (isStrided()) {

mlir/test/IR/invalid-ops.mlir

func @dma_start_dst_index_wrong_type(
     %tag: memref<2xi32,2>, %flt: f32) {
   // expected-error@+1 {{expected tag indices to be of index type}}
   "memref.dma_start"(%src, %idx, %idx, %dst, %idx, %idx, %tag, %flt)
       : (memref<2x2xf32>, index, index, memref<2xf32,1>, index, index, memref<2xi32,2>, f32) -> ()
 }
 
 // -----
 
-func @dma_start_same_space(
-    %src: memref<2x2xf32>, %idx: index, %dst: memref<2xf32>,
-    %tag: memref<i32,2>) {
-  // expected-error@+1 {{DMA should be between different memory spaces}}
-  memref.dma_start %src[%idx, %idx], %dst[%idx], %idx, %tag[] : memref<2x2xf32>, memref<2xf32>, memref<i32,2>
-}
-
-// -----
-
 func @dma_start_too_many_operands(
     %src: memref<2x2xf32>, %idx: index, %dst: memref<2xf32,1>,
     %tag: memref<i32,2>) {
   // expected-error@+1 {{incorrect number of operands}}
   "memref.dma_start"(%src, %idx, %idx, %dst, %idx, %idx, %tag, %idx, %idx, %idx)
       : (memref<2x2xf32>, index, index, memref<2xf32,1>, index, index, memref<i32,2>, index, index, index) -> ()
 }