[mlir][VectorOps] Extend VectorTransfer lowering to n-D memref with minor identity map

[nicolasvasilache]

Summary

This revision extends the lowering of vector transfers to work with n-D memref and 1-D vector where the permutation map is an identity on the most minor dimensions (1 for now).

[ftynse]

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

858

Nit: I don't remember if the transfer_op verifier checks for the list of indices being non-empty

[nicolasvasilache]

Address review comment.

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

getTransferOpAdapter(TransferReadOp xferOp, ArrayRef<Value> operands) {
   return TransferReadOpOperandAdaptor(operands);
 }
 
 static TransferWriteOpOperandAdaptor
 getTransferOpAdapter(TransferWriteOp xferOp, ArrayRef<Value> operands) {
   return TransferWriteOpOperandAdaptor(operands);
 }
 
+bool isMinorIdentity(AffineMap map, unsigned rank) {
+  if (map.getNumResults() < rank)
+    return false;
+  unsigned startDim = map.getNumDims() - rank;
+  for (unsigned i = 0; i < rank; ++i)
+    if (map.getResult(i) != getAffineDimExpr(startDim + i, map.getContext()))
+      return false;
+  return true;
+}
+
 /// Conversion pattern that converts a 1-D vector transfer read/write op in a
 /// sequence of:
 /// 1. Bitcast to vector form.
 /// 2. Create an offsetVector = [ offset + 0 .. offset + vector_length - 1 ].
 /// 3. Create a mask where offsetVector is compared against memref upper bound.
 /// 4. Rewrite op as a masked read or write.
 template <typename ConcreteOp>
 class VectorTransferConversion : public ConvertToLLVMPattern {
 public:
   explicit VectorTransferConversion(MLIRContext *context,
                                     LLVMTypeConverter &typeConv)
       : ConvertToLLVMPattern(ConcreteOp::getOperationName(), context,
                              typeConv) {}
 
   LogicalResult
   matchAndRewrite(Operation *op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     auto xferOp = cast<ConcreteOp>(op);
     auto adaptor = getTransferOpAdapter(xferOp, operands);
-    if (xferOp.getMemRefType().getRank() != 1)
-      return failure();
-    if (!xferOp.permutation_map().isIdentity())
+    if (!isMinorIdentity(xferOp.permutation_map(),
+                         xferOp.getVectorType().getRank()))
       return failure();
 
     auto toLLVMTy = [&](Type t) { return typeConverter.convertType(t); };
 
     Location loc = op->getLoc();
     Type i64Type = rewriter.getIntegerType(64);
     MemRefType memRefType = xferOp.getMemRefType();
 
for (unsigned i = 0; i < vecWidth; ++i)
       indices.push_back(i);
     Value linearIndices = rewriter.create<ConstantOp>(
         loc, vectorCmpType,
         DenseElementsAttr::get(vectorCmpType, ArrayRef<int64_t>(indices)));
     linearIndices = rewriter.create<LLVM::DialectCastOp>(
         loc, toLLVMTy(vectorCmpType), linearIndices);
 
     // 3. Create offsetVector = [ offset + 0 .. offset + vector_length - 1 ].
-    Value offsetIndex = *(xferOp.indices().begin());
-    offsetIndex = rewriter.create<IndexCastOp>(
-        loc, vectorCmpType.getElementType(), offsetIndex);
+    // TODO(ntv, ajcbik): when the leaf transfer rank is k > 1 we need the last
+    // `k` dimensions here.
+    unsigned lastIndex = llvm::size(xferOp.indices()) - 1;
+    Value offsetIndex = *(xferOp.indices().begin() + lastIndex);
+    offsetIndex = rewriter.create<IndexCastOp>(loc, i64Type, offsetIndex);
     Value base = rewriter.create<SplatOp>(loc, vectorCmpType, offsetIndex);
     Value offsetVector = rewriter.create<AddIOp>(loc, base, linearIndices);
 
     // 4. Let dim the memref dimension, compute the vector comparison mask:
     //   [ offset + 0 .. offset + vector_length - 1 ] < [ dim .. dim ]
-    Value dim = rewriter.create<DimOp>(loc, xferOp.memref(), 0);
-    dim =
-        rewriter.create<IndexCastOp>(loc, vectorCmpType.getElementType(), dim);
+    Value dim = rewriter.create<DimOp>(loc, xferOp.memref(), lastIndex);
+    dim = rewriter.create<IndexCastOp>(loc, i64Type, dim);
     dim = rewriter.create<SplatOp>(loc, vectorCmpType, dim);
     Value mask =
         rewriter.create<CmpIOp>(loc, CmpIPredicate::slt, offsetVector, dim);
     mask = rewriter.create<LLVM::DialectCastOp>(loc, toLLVMTy(mask.getType()),
                                                 mask);
 
     // 5. Rewrite as a masked read / write.
     replaceTransferOp<ConcreteOp>(rewriter, typeConverter, loc, op, operands,

mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir

 //  CHECK-SAME:  0 : i32, 0 : i32, 0 : i32, 0 : i32, 0 : i32, 0 : i32] :
 //  CHECK-SAME: !llvm<"<17 x i64>">, !llvm<"<17 x i64>">
 //       CHECK: %[[mask_b:.*]] = llvm.icmp "slt" {{.*}} : !llvm<"<17 x i64>">
 //
 // 5. Rewrite as a masked write.
 //       CHECK: llvm.intr.masked.store %[[loaded]], %[[vecPtr_b]], %[[mask_b]]
 //  CHECK-SAME: {alignment = 1 : i32} :
 //  CHECK-SAME: !llvm<"<17 x float>">, !llvm<"<17 x i1>"> into !llvm<"<17 x float>*">
+
+func @transfer_read_2d_to_1d(%A : memref<?x?xf32>, %base0: index, %base1: index) -> vector<17xf32> {
+  %f7 = constant 7.0: f32
+  %f = vector.transfer_read %A[%base0, %base1], %f7
+      {permutation_map = affine_map<(d0, d1) -> (d1)>} :
+    memref<?x?xf32>, vector<17xf32>
+  return %f: vector<17xf32>
+}
+// CHECK-LABEL: func @transfer_read_2d_to_1d
+//  CHECK-SAME: %[[BASE_0:[a-zA-Z0-9]*]]: !llvm.i64, %[[BASE_1:[a-zA-Z0-9]*]]: !llvm.i64) -> !llvm<"<17 x float>">
+//
+// Create offsetVector = [ offset + 0 .. offset + vector_length - 1 ].
+//       CHECK: %[[offsetVec:.*]] = llvm.mlir.undef : !llvm<"<17 x i64>">
+//       CHECK: %[[c0:.*]] = llvm.mlir.constant(0 : i32) : !llvm.i32
+// Here we check we properly use %BASE_1
+//       CHECK: %[[offsetVec2:.*]] = llvm.insertelement %[[BASE_1]], %[[offsetVec]][%[[c0]] :
+//  CHECK-SAME: !llvm.i32] : !llvm<"<17 x i64>">
+//       CHECK: %[[offsetVec3:.*]] = llvm.shufflevector %[[offsetVec2]], %{{.*}} [
+//  CHECK-SAME:  0 : i32, 0 : i32, 0 : i32, 0 : i32, 0 : i32, 0 : i32, 0 : i32,
+//  CHECK-SAME:  0 : i32, 0 : i32, 0 : i32, 0 : i32, 0 : i32, 0 : i32, 0 : i32,
+//  CHECK-SAME:  0 : i32, 0 : i32, 0 : i32] :
+//
+// Let dim the memref dimension, compute the vector comparison mask:
+//    [ offset + 0 .. offset + vector_length - 1 ] < [ dim .. dim ]
+// Here we check we properly use %DIM[1]
+//       CHECK: %[[DIM:.*]] = llvm.extractvalue %{{.*}}[3, 1] :
+//  CHECK-SAME: !llvm<"{ float*, float*, i64, [2 x i64], [2 x i64] }">
+//       CHECK: %[[dimVec:.*]] = llvm.mlir.undef : !llvm<"<17 x i64>">
+//       CHECK: %[[c01:.*]] = llvm.mlir.constant(0 : i32) : !llvm.i32
+//       CHECK: %[[dimVec2:.*]] = llvm.insertelement %[[DIM]], %[[dimVec]][%[[c01]] :
+//  CHECK-SAME:  !llvm.i32] : !llvm<"<17 x i64>">
+//       CHECK: %[[dimVec3:.*]] = llvm.shufflevector %[[dimVec2]], %{{.*}} [
+//  CHECK-SAME:  0 : i32, 0 : i32, 0 : i32, 0 : i32, 0 : i32, 0 : i32, 0 : i32,
+//  CHECK-SAME:  0 : i32, 0 : i32, 0 : i32, 0 : i32, 0 : i32, 0 : i32, 0 : i32,
+//  CHECK-SAME:  0 : i32, 0 : i32, 0 : i32] :
+//  CHECK-SAME: !llvm<"<17 x i64>">, !llvm<"<17 x i64>">