[mlir][vector] Handle memory space conflicts in VectorTransferSplit patterns

[qedawkins]

Summary

Currently the transfer splitting patterns will generate an invalid cast
when the source memref for a transfer op has a non-default memory space.
This is handled by first introducing a memref.memory_space_cast in
such cases.

mlir/lib/Dialect/Vector/Transforms/VectorTransferSplitRewritePatterns.cpp

resStrides[idx] =
         (aStrides[idx] == bStrides[idx]) ? aStrides[idx] : ShapedType::kDynamic;
   }
   resOffset = (aOffset == bOffset) ? aOffset : ShapedType::kDynamic;
   return MemRefType::get(
       resShape, aT.getElementType(),
       StridedLayoutAttr::get(aT.getContext(), resOffset, resStrides));
 }
 
+/// Casts the given memref to a compatible memref type. If the source memref has
+/// a different address space than the target type, a `memref.memory_space_cast`
+/// is first inserted, followed by a `memref.cast`.
+static Value castToCompatibleMemRefType(OpBuilder &b, Value memref,
+                                        MemRefType compatibleMemRefType) {
+  MemRefType sourceType = memref.getType().cast<MemRefType>();
+  Value res = memref;
+  if (sourceType.getMemorySpace() != compatibleMemRefType.getMemorySpace()) {
+    sourceType = MemRefType::get(
+        sourceType.getShape(), sourceType.getElementType(),
+        sourceType.getLayout(), compatibleMemRefType.getMemorySpace());
+    res = b.create<memref::MemorySpaceCastOp>(memref.getLoc(), sourceType, res);
+  }
+  if (sourceType == compatibleMemRefType)
+    return res;
+  return b.create<memref::CastOp>(memref.getLoc(), compatibleMemRefType, res);
+}
+
 /// Operates under a scoped context to build the intersection between the
 /// view `xferOp.source()` @ `xferOp.indices()` and the view `alloc`.
 // TODO: view intersection/union/differences should be a proper std op.
 static std::pair<Value, Value>
 createSubViewIntersection(RewriterBase &b, VectorTransferOpInterface xferOp,
                           Value alloc) {
   Location loc = xferOp.getLoc();
   int64_t memrefRank = xferOp.getShapedType().getRank();

 }
 
 /// Given an `xferOp` for which:
 ///   1. `inBoundsCond` and a `compatibleMemRefType` have been computed.
 ///   2. a memref of single vector `alloc` has been allocated.
 /// Produce IR resembling:
 /// ```
 ///    %1:3 = scf.if (%inBounds) {
+///      (memref.memory_space_cast %A: memref<A..., addr_space> to memref<A...>)
 ///      %view = memref.cast %A: memref<A...> to compatibleMemRefType
 ///      scf.yield %view, ... : compatibleMemRefType, index, index
 ///    } else {
 ///      %2 = linalg.fill(%pad, %alloc)
 ///      %3 = subview %view [...][...][...]
 ///      %4 = subview %alloc [0, 0] [...] [...]
 ///      linalg.copy(%3, %4)
 ///      %5 = memref.cast %alloc: memref<B...> to compatibleMemRefType
 ///      scf.yield %5, ... : compatibleMemRefType, index, index
 ///   }
 /// ```
 /// Return the produced scf::IfOp.
 static scf::IfOp
 createFullPartialLinalgCopy(RewriterBase &b, vector::TransferReadOp xferOp,
                             TypeRange returnTypes, Value inBoundsCond,
                             MemRefType compatibleMemRefType, Value alloc) {
   Location loc = xferOp.getLoc();
   Value zero = b.create<arith::ConstantIndexOp>(loc, 0);
   Value memref = xferOp.getSource();
   return b.create<scf::IfOp>(
       loc, inBoundsCond,
       [&](OpBuilder &b, Location loc) {
-        Value res = memref;
-        if (compatibleMemRefType != xferOp.getShapedType())
-          res = b.create<memref::CastOp>(loc, compatibleMemRefType, memref);
+        Value res = castToCompatibleMemRefType(b, memref, compatibleMemRefType);
         scf::ValueVector viewAndIndices{res};
         viewAndIndices.insert(viewAndIndices.end(), xferOp.getIndices().begin(),
                               xferOp.getIndices().end());
         b.create<scf::YieldOp>(loc, viewAndIndices);
       },
       [&](OpBuilder &b, Location loc) {
         b.create<linalg::FillOp>(loc, ValueRange{xferOp.getPadding()},
                                  ValueRange{alloc});
         // Take partial subview of memref which guarantees no dimension
         // overflows.
         IRRewriter rewriter(b);
         std::pair<Value, Value> copyArgs = createSubViewIntersection(
             rewriter, cast<VectorTransferOpInterface>(xferOp.getOperation()),
             alloc);
         b.create<memref::CopyOp>(loc, copyArgs.first, copyArgs.second);
         Value casted =
-            b.create<memref::CastOp>(loc, compatibleMemRefType, alloc);
+            castToCompatibleMemRefType(b, alloc, compatibleMemRefType);
         scf::ValueVector viewAndIndices{casted};
         viewAndIndices.insert(viewAndIndices.end(), xferOp.getTransferRank(),
                               zero);
         b.create<scf::YieldOp>(loc, viewAndIndices);
       });
 }
 
 /// Given an `xferOp` for which:
 ///   1. `inBoundsCond` and a `compatibleMemRefType` have been computed.
 ///   2. a memref of single vector `alloc` has been allocated.
 /// Produce IR resembling:
 /// ```
 ///    %1:3 = scf.if (%inBounds) {
+///      (memref.memory_space_cast %A: memref<A..., addr_space> to memref<A...>)
 ///      memref.cast %A: memref<A...> to compatibleMemRefType
 ///      scf.yield %view, ... : compatibleMemRefType, index, index
 ///    } else {
 ///      %2 = vector.transfer_read %view[...], %pad : memref<A...>, vector<...>
 ///      %3 = vector.type_cast %extra_alloc :
 ///        memref<...> to memref<vector<...>>
 ///      store %2, %3[] : memref<vector<...>>
 ///      %4 = memref.cast %alloc: memref<B...> to compatibleMemRefType
 ///      scf.yield %4, ... : compatibleMemRefType, index, index
 ///   }
 /// ```
 /// Return the produced scf::IfOp.
 static scf::IfOp createFullPartialVectorTransferRead(
     RewriterBase &b, vector::TransferReadOp xferOp, TypeRange returnTypes,
     Value inBoundsCond, MemRefType compatibleMemRefType, Value alloc) {
   Location loc = xferOp.getLoc();
   scf::IfOp fullPartialIfOp;
   Value zero = b.create<arith::ConstantIndexOp>(loc, 0);
   Value memref = xferOp.getSource();
   return b.create<scf::IfOp>(
       loc, inBoundsCond,
       [&](OpBuilder &b, Location loc) {
-        Value res = memref;
-        if (compatibleMemRefType != xferOp.getShapedType())
-          res = b.create<memref::CastOp>(loc, compatibleMemRefType, memref);
+        Value res = castToCompatibleMemRefType(b, memref, compatibleMemRefType);
         scf::ValueVector viewAndIndices{res};
         viewAndIndices.insert(viewAndIndices.end(), xferOp.getIndices().begin(),
                               xferOp.getIndices().end());
         b.create<scf::YieldOp>(loc, viewAndIndices);
       },
       [&](OpBuilder &b, Location loc) {
         Operation *newXfer = b.clone(*xferOp.getOperation());
         Value vector = cast<VectorTransferOpInterface>(newXfer).vector();
         b.create<memref::StoreOp>(
             loc, vector,
             b.create<vector::TypeCastOp>(
                 loc, MemRefType::get({}, vector.getType()), alloc));
 
         Value casted =
-            b.create<memref::CastOp>(loc, compatibleMemRefType, alloc);
+            castToCompatibleMemRefType(b, alloc, compatibleMemRefType);
         scf::ValueVector viewAndIndices{casted};
         viewAndIndices.insert(viewAndIndices.end(), xferOp.getTransferRank(),
                               zero);
         b.create<scf::YieldOp>(loc, viewAndIndices);
       });
 }
 
 /// Given an `xferOp` for which:
getLocationToWriteFullVec(RewriterBase &b, vector::TransferWriteOp xferOp,
                           MemRefType compatibleMemRefType, Value alloc) {
   Location loc = xferOp.getLoc();
   Value zero = b.create<arith::ConstantIndexOp>(loc, 0);
   Value memref = xferOp.getSource();
   return b
       .create<scf::IfOp>(
           loc, inBoundsCond,
           [&](OpBuilder &b, Location loc) {
-            Value res = memref;
-            if (compatibleMemRefType != xferOp.getShapedType())
-              res = b.create<memref::CastOp>(loc, compatibleMemRefType, memref);
+            Value res =
+                castToCompatibleMemRefType(b, memref, compatibleMemRefType);
             scf::ValueVector viewAndIndices{res};
             viewAndIndices.insert(viewAndIndices.end(),
                                   xferOp.getIndices().begin(),
                                   xferOp.getIndices().end());
             b.create<scf::YieldOp>(loc, viewAndIndices);
           },
           [&](OpBuilder &b, Location loc) {
             Value casted =
-                b.create<memref::CastOp>(loc, compatibleMemRefType, alloc);
+                castToCompatibleMemRefType(b, alloc, compatibleMemRefType);
             scf::ValueVector viewAndIndices{casted};
             viewAndIndices.insert(viewAndIndices.end(),
                                   xferOp.getTransferRank(), zero);
             b.create<scf::YieldOp>(loc, viewAndIndices);
           })
       ->getResults();
 }
 

mlir/test/Dialect/Vector/vector-transfer-full-partial-split.mlir

func.func @split_vector_transfer_read_strided_2d(
   // CHECK-SAME:   memref<?x8xf32, strided<[?, 1], offset: ?>>, vector<4x8xf32>
   %1 = vector.transfer_read %A[%i, %j], %f0 :
     memref<7x8xf32, strided<[?, 1], offset: ?>>, vector<4x8xf32>
 
   // CHECK: return %[[res]] : vector<4x8xf32>
   return %1 : vector<4x8xf32>
 }
 
+func.func @split_vector_transfer_read_mem_space(%A: memref<?x8xf32, 3>, %i: index, %j: index) -> vector<4x8xf32> {
+  %c0 = arith.constant 0 : index
+  %f0 = arith.constant 0.0 : f32
+
+  //      CHECK: scf.if {{.*}} -> (memref<?x8xf32, strided<[8, 1]>>, index, index) {
+  //               inBounds with a different memory space
+  //      CHECK:   %[[space_cast:.*]] = memref.memory_space_cast %{{.*}} :
+  // CHECK-SAME:     memref<?x8xf32, 3> to memref<?x8xf32>
+  //      CHECK:   %[[cast:.*]] = memref.cast %[[space_cast]] :
+  // CHECK-SAME:     memref<?x8xf32> to memref<?x8xf32, strided<[8, 1]>>
+  //      CHECK:   scf.yield %[[cast]], {{.*}} : memref<?x8xf32, strided<[8, 1]>>, index, index
+  //      CHECK: } else {
+  //               slow path, fill tmp alloc and yield a memref_casted version of it
+  //      CHECK:   %[[slow:.*]] = vector.transfer_read %[[A]][%[[i]], %[[j]]], %cst :
+  // CHECK-SAME:     memref<?x8xf32, 3>, vector<4x8xf32>
+  //      CHECK:   %[[cast_alloc:.*]] = vector.type_cast %[[alloc]] :
+  // CHECK-SAME:     memref<4x8xf32> to memref<vector<4x8xf32>>
+  //      CHECK:   store %[[slow]], %[[cast_alloc]][] : memref<vector<4x8xf32>>
+  //      CHECK:   %[[yielded:.*]] = memref.cast %[[alloc]] :
+  // CHECK-SAME:     memref<4x8xf32> to memref<?x8xf32, strided<[8, 1]>>
+  //      CHECK:   scf.yield %[[yielded]], %[[c0]], %[[c0]] :
+  // CHECK-SAME:     memref<?x8xf32, strided<[8, 1]>>, index, index
+  //      CHECK: }
+  //      CHECK: %[[res:.*]] = vector.transfer_read %[[ifres]]#0[%[[ifres]]#1, %[[ifres]]#2], %cst
+  // CHECK-SAME:   {in_bounds = [true, true]} : memref<?x8xf32, strided<[8, 1]>>, vector<4x8xf32>
+
+  %1 = vector.transfer_read %A[%i, %j], %f0 : memref<?x8xf32, 3>, vector<4x8xf32>
+
+  return %1: vector<4x8xf32>
+}
+
 transform.sequence failures(propagate) {
 ^bb1(%func_op: !transform.op<"func.func">):
   transform.apply_patterns to %func_op {
     transform.apply_patterns.vector.split_transfer_full_partial split_transfer_strategy = "vector-transfer"
   } : !transform.op<"func.func">
 }
 
 // -----

 ^bb1(%func_op: !transform.op<"func.func">):
   transform.apply_patterns to %func_op {
     transform.apply_patterns.vector.split_transfer_full_partial split_transfer_strategy = "vector-transfer"
   } : !transform.op<"func.func">
 }
 
 // -----
 
+func.func @split_vector_transfer_write_mem_space(%V: vector<4x8xf32>, %A: memref<?x8xf32, 3>, %i: index, %j: index) {
+  vector.transfer_write %V, %A[%i, %j] :
+    vector<4x8xf32>, memref<?x8xf32, 3>
+  return
+}
+
+// CHECK:     func @split_vector_transfer_write_mem_space(
+// CHECK:           scf.if {{.*}} -> (memref<?x8xf32, strided<[8, 1]>>, index, index) {
+// CHECK:             %[[space_cast:.*]] = memref.memory_space_cast %{{.*}} :
+// CHECK-SAME:          memref<?x8xf32, 3> to memref<?x8xf32>
+// CHECK:             %[[cast:.*]] = memref.cast %[[space_cast]] :
+// CHECK-SAME:          memref<?x8xf32> to memref<?x8xf32, strided<[8, 1]>>
+// CHECK:             scf.yield %[[cast]], {{.*}} : memref<?x8xf32, strided<[8, 1]>>, index, index
+// CHECK:           } else {
+// CHECK:             %[[VAL_15:.*]] = memref.cast %[[TEMP]]
+// CHECK-SAME:            : memref<4x8xf32> to memref<?x8xf32, strided<[8, 1]>>
+// CHECK:             scf.yield %[[VAL_15]], %[[C0]], %[[C0]]
+// CHECK-SAME:            : memref<?x8xf32, strided<[8, 1]>>, index, index
+// CHECK:           }
+// CHECK:           vector.transfer_write %[[VEC]],
+// CHECK-SAME:           %[[IN_BOUND_DEST:.*]]#0[%[[IN_BOUND_DEST]]#1, %[[IN_BOUND_DEST]]#2]
+// CHECK-SAME:           {in_bounds = [true, true]} : vector<4x8xf32>, memref<?x8xf32, strided<[8, 1]>>
+
+
+transform.sequence failures(propagate) {
+^bb1(%func_op: !transform.op<"func.func">):
+  transform.apply_patterns to %func_op {
+    transform.apply_patterns.vector.split_transfer_full_partial split_transfer_strategy = "vector-transfer"
+  } : !transform.op<"func.func">
+}
+
+
+// -----
+
 func.func private @fake_side_effecting_fun(%0: vector<2x2xf32>) -> ()
 
 // CHECK-LABEL: transfer_read_within_async_execute
 func.func @transfer_read_within_async_execute(%A : memref<?x?xf32>) -> !async.token {
   %c0 = arith.constant 0 : index
   %f0 = arith.constant 0.0 : f32
   // CHECK-NOT: alloca
   //     CHECK: async.execute