[mlir][vector] let transfer_read and transfer_write take non-zero addrspace.

[whchung]

Summary

Enhance lowering logic and tests so vector.transfer_read and
vector.transfer_write take memrefs on non-zero addrspaces.

[whchung]

Rebase to tip.

[aartbik]

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

837–840

This step 1. can use a bit more documentation now that it has become more eleborate

[whchung]

Revise comment.

[whchung]

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

837–840

@aartbik just added more comments.

[ftynse]

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

847

Please drop trivial braces

[whchung]

Address review comment.

[awpr]

Is it important for the memory accesses here to be done via addrspace 0, or is that just an artifact of making the pointer type match up? It seems that for many architectures with multiple memory spaces, it won't be possible to address all memory via addrspace 0, and memory accesses should be done in the addrspace their memref lives in (as in https://reviews.llvm.org/D96380); but I'm not sure whether making that change would affect the use case this was meant for.

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

bool isMinorIdentity(AffineMap map, unsigned 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.
+/// 1. Bitcast or addrspacecast 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)
if (!isMinorIdentity(xferOp.permutation_map(),
       return failure();
 
     auto toLLVMTy = [&](Type t) { return typeConverter.convertType(t); };
 
     Location loc = op->getLoc();
     Type i64Type = rewriter.getIntegerType(64);
     MemRefType memRefType = xferOp.getMemRefType();
 
     // 1. Get the source/dst address as an LLVM vector pointer.
+    //    The vector pointer would always be on address space 0, therefore
+    //    addrspacecast shall be used when source/dst memrefs are not on
+    //    address space 0.
     // TODO: support alignment when possible.
     Value dataPtr = getDataPtr(loc, memRefType, adaptor.memref(),
                                adaptor.indices(), rewriter, getModule());
     auto vecTy =
         toLLVMTy(xferOp.getVectorType()).template cast<LLVM::LLVMType>();
-    auto vectorDataPtr =
+    Value vectorDataPtr;
+    if (memRefType.getMemorySpace() == 0) {
+      vectorDataPtr =
           rewriter.create<LLVM::BitcastOp>(loc, vecTy.getPointerTo(), dataPtr);
+    } else {
+      vectorDataPtr = rewriter.create<LLVM::AddrSpaceCastOp>(
+          loc, vecTy.getPointerTo(), dataPtr);
+    }
 
     // 2. Create a vector with linear indices [ 0 .. vector_length - 1 ].
     unsigned vecWidth = vecTy.getVectorNumElements();
     VectorType vectorCmpType = VectorType::get(vecWidth, i64Type);
     SmallVector<int64_t, 8> indices;
     indices.reserve(vecWidth);
     for (unsigned i = 0; i < vecWidth; ++i)
       indices.push_back(i);

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

 //       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>">
+
+func @transfer_read_1d_non_zero_addrspace(%A : memref<?xf32, 3>, %base: index) -> vector<17xf32> {
+  %f7 = constant 7.0: f32
+  %f = vector.transfer_read %A[%base], %f7
+      {permutation_map = affine_map<(d0) -> (d0)>} :
+    memref<?xf32, 3>, vector<17xf32>
+  vector.transfer_write %f, %A[%base]
+      {permutation_map = affine_map<(d0) -> (d0)>} :
+    vector<17xf32>, memref<?xf32, 3>
+  return %f: vector<17xf32>
+}
+// CHECK-LABEL: func @transfer_read_1d_non_zero_addrspace
+//  CHECK-SAME: %[[BASE:[a-zA-Z0-9]*]]: !llvm.i64) -> !llvm<"<17 x float>">
+//
+// 1. Check address space for GEP is correct.
+//       CHECK: %[[gep:.*]] = llvm.getelementptr {{.*}} :
+//  CHECK-SAME: (!llvm<"float addrspace(3)*">, !llvm.i64) -> !llvm<"float addrspace(3)*">
+//       CHECK: %[[vecPtr:.*]] = llvm.addrspacecast %[[gep]] :
+//  CHECK-SAME: !llvm<"float addrspace(3)*"> to !llvm<"<17 x float>*">
+//
+// 2. Check address space of the memref is correct.
+//       CHECK: %[[DIM:.*]] = llvm.extractvalue %{{.*}}[3, 0] :
+//  CHECK-SAME: !llvm<"{ float addrspace(3)*, float addrspace(3)*, i64, [1 x i64], [1 x i64] }">
+//
+// 3. Check address apce for GEP is correct.
+//       CHECK: %[[gep_b:.*]] = llvm.getelementptr {{.*}} :
+//  CHECK-SAME: (!llvm<"float addrspace(3)*">, !llvm.i64) -> !llvm<"float addrspace(3)*">
+//       CHECK: %[[vecPtr_b:.*]] = llvm.addrspacecast %[[gep_b]] :
+//  CHECK-SAME: !llvm<"float addrspace(3)*"> to !llvm<"<17 x float>*">