[MLIR] Add static offset support for ViewOp's source memref.

[timshen]

Summary

Discussion is here: https://llvm.discourse.group/t/viewop-isnt-expressive-enough/991/9.

The fix is not the same as discussed in the thread. It instead partially support non-identity affine maps for the source memref. Currently, only static offset support is added.

[nicolasvasilache]

I started reading the revision but I believe that this is only solving a part of the problem, which happens to be your particular blocker.
I believe the required changes are much deeper and after experimenting a bit I find that a good chunk of the ViewOp semantics is unnecessary.
See: https://reviews.llvm.org/D79541 and in particular the update to Ops.td.

Also @andydavis1 for visibility since this also spawns discussions we had a while back.

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

2625

if(!ShapedType::isDynamicOffsetOrStride(baseOffset))

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

2712

<< baseType or you'll get wrong error messages.

2715

<< baseType or you'll get wrong error messages.

mlir/test/IR/invalid-ops.mlir

960

wrong error message

[nicolasvasilache]

I think this is now obsolete but please add me back if it isn't

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

struct ViewOpLowering : public ConvertOpToLLVMPattern<ViewOp> {
 
   LogicalResult
   matchAndRewrite(Operation *op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     auto loc = op->getLoc();
     auto viewOp = cast<ViewOp>(op);
     ViewOpOperandAdaptor adaptor(operands);
 
+    auto sourceMemRefType = viewOp.source().getType().cast<MemRefType>();
+    auto sourceElementTy =
+        typeConverter.convertType(sourceMemRefType.getElementType())
+            .dyn_cast<LLVM::LLVMType>();
     auto viewMemRefType = viewOp.getType();
     auto targetElementTy =
         typeConverter.convertType(viewMemRefType.getElementType())
             .dyn_cast<LLVM::LLVMType>();
     auto targetDescTy =
         typeConverter.convertType(viewMemRefType).dyn_cast<LLVM::LLVMType>();
     if (!targetDescTy)
       return op->emitWarning("Target descriptor type not converted to LLVM"),
              failure();
 
+    int64_t baseOffset;
+    {
+      SmallVector<int64_t, 4> strides;
+      if (failed(getStridesAndOffset(sourceMemRefType, strides, baseOffset)))
+        return op->emitWarning("cannot cast base type to non-strided shape"),
+               failure();
+    }
+
     int64_t offset;
     SmallVector<int64_t, 4> strides;
     auto successStrides = getStridesAndOffset(viewMemRefType, strides, offset);
     if (failed(successStrides))
       return op->emitWarning("cannot cast to non-strided shape"), failure();
 
     // Create the descriptor.
     MemRefDescriptor sourceMemRef(adaptor.source());
     auto targetMemRef = MemRefDescriptor::undef(rewriter, loc, targetDescTy);
 
     // Field 1: Copy the allocated pointer, used for malloc/free.
     Value extracted = sourceMemRef.allocatedPtr(rewriter, loc);
     Value bitcastPtr = rewriter.create<LLVM::BitcastOp>(
         loc, targetElementTy.getPointerTo(), extracted);
     targetMemRef.setAllocatedPtr(rewriter, loc, bitcastPtr);
 
     // Field 2: Copy the actual aligned pointer to payload.
-    extracted = sourceMemRef.alignedPtr(rewriter, loc);
+    Value alignedPtr = sourceMemRef.alignedPtr(rewriter, loc);
+
+    if (baseOffset > 0)
+      alignedPtr = rewriter.create<LLVM::GEPOp>(
+          loc, sourceElementTy.getPointerTo(), alignedPtr,
+          ValueRange{createIndexConstant(rewriter, loc, baseOffset)});
+
     bitcastPtr = rewriter.create<LLVM::BitcastOp>(
-        loc, targetElementTy.getPointerTo(), extracted);
+        loc, targetElementTy.getPointerTo(), alignedPtr);
+
     targetMemRef.setAlignedPtr(rewriter, loc, bitcastPtr);
 
     // Field 3: Copy the offset in aligned pointer.
     unsigned numDynamicSizes = llvm::size(viewOp.getDynamicSizes());
     (void)numDynamicSizes;
     bool hasDynamicOffset = offset == MemRefType::getDynamicStrideOrOffset();
     auto sizeAndOffsetOperands = adaptor.operands();
     assert(llvm::size(sizeAndOffsetOperands) ==
            numDynamicSizes + (hasDynamicOffset ? 1 : 0));
-    Value baseOffset = !hasDynamicOffset
+    Value targetOffset = !hasDynamicOffset
                              ? createIndexConstant(rewriter, loc, offset)
                              // TODO(ntv): better adaptor.
                              : sizeAndOffsetOperands.front();
-    targetMemRef.setOffset(rewriter, loc, baseOffset);
+    targetMemRef.setOffset(rewriter, loc, targetOffset);
 
     // Early exit for 0-D corner case.
     if (viewMemRefType.getRank() == 0)
       return rewriter.replaceOp(op, {targetMemRef}), success();
 
     // Fields 4 and 5: Update sizes and strides.
     if (strides.back() != 1)
       return op->emitWarning("cannot cast to non-contiguous shape"), failure();

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

static LogicalResult verifyDynamicStrides(MemRefType memrefType,
   }
   return success();
 }
 
 static LogicalResult verify(ViewOp op) {
   auto baseType = op.getOperand(0).getType().cast<MemRefType>();
   auto viewType = op.getResult().getType().cast<MemRefType>();
 
-  // The base memref should have identity layout map (or none).
-  if (baseType.getAffineMaps().size() > 1 ||
-      (baseType.getAffineMaps().size() == 1 &&
-       !baseType.getAffineMaps()[0].isIdentity()))
+  // The base memref should have at most one affine map with stride 1.
+  if (baseType.getAffineMaps().size() > 1)
     return op.emitError("unsupported map for base memref type ") << baseType;
 
+  if (baseType.getAffineMaps().size() == 1) {
+    int64_t offset;
+    SmallVector<int64_t, 4> strides;
+    if (failed(getStridesAndOffset(baseType, strides, offset)))
+      return op.emitError("base type ") << viewType << " is not strided";
+
+    if (offset == MemRefType::getDynamicStrideOrOffset())
+      return op.emitError("base type ") << viewType << " has dynamic offset";
+
+    if (strides.size() != 1 || strides.back() != 1)
+      return op.emitError("base type ")
+             << viewType << " has non-trivial strides";
+  }
+
   // The base memref and the view memref should be in the same memory space.
   if (baseType.getMemorySpace() != viewType.getMemorySpace())
     return op.emitError("different memory spaces specified for base memref "
                         "type ")
            << baseType << " and view memref type " << viewType;
 
   // Verify that the result memref type has a strided layout map.
   int64_t offset;

mlir/test/Conversion/StandardToLLVM/convert-to-llvmir.mlir

 
 // CHECK-LABEL: func @mlir_cast_from_llvm
 // CHECK-SAME: %[[ARG:.*]]:
 func @mlir_cast_from_llvm(%0 : !llvm.half) -> f16 {
   %1 = llvm.mlir.cast %0 : !llvm.half to f16
   // CHECK-NEXT: llvm.return %[[ARG]]
   return %1 : f16
 }
+
+// -----
+
+// CHECK-LABEL: func @view_base_offset(
+func @view_base_offset(%0: memref<2045xi8, offset: 3, strides: [1]>) {
+  // Test static base offset
+
+  // CHECK: %[[DESC:[0-9]+]] = llvm.insertvalue %{{.*}}, %{{.*}}[0] : !llvm<"{ float*, float*, i64, [2 x i64], [2 x i64] }">
+  // CHECK: %[[BASE:[0-9]+]] = llvm.extractvalue %{{.*}}[1] : !llvm<"{ i8*, i8*, i64, [1 x i64], [1 x i64] }">
+  // CHECK: %[[OFFSET:[0-9]+]] = llvm.mlir.constant(3 : index) : !llvm.i64
+  // CHECK: %[[PTR:[0-9]+]] = llvm.getelementptr %[[BASE]][%[[OFFSET]]] : (!llvm<"i8*">, !llvm.i64) -> !llvm<"i8*">
+  // CHECK: %[[CASTED:[0-9]+]] = llvm.bitcast %[[PTR]] : !llvm<"i8*"> to !llvm<"float*">
+  // CHECK: llvm.insertvalue %[[CASTED]], %{{.*}}[1] : !llvm<"{ float*, float*, i64, [2 x i64], [2 x i64] }">
+  %8 = view %0[][] : memref<2045xi8, offset: 3, strides: [1]> to memref<2x3xf32>
+
+  return
+}

mlir/test/IR/core-ops.mlir

 // RUN: mlir-opt -allow-unregistered-dialect %s | FileCheck %s
 // Verify the printed output can be parsed.
 // RUN: mlir-opt -allow-unregistered-dialect %s | mlir-opt -allow-unregistered-dialect | FileCheck %s
 // Verify the generic form can be parsed.
 // RUN: mlir-opt -allow-unregistered-dialect -mlir-print-op-generic %s | mlir-opt -allow-unregistered-dialect | FileCheck %s
 
 // CHECK: #map0 = affine_map<(d0) -> (d0 + 1)>
 
 // CHECK: #map1 = affine_map<()[s0] -> (s0 + 1)>
 
 // CHECK-DAG: #[[VIEW_MAP1:map[0-9]+]] = affine_map<(d0, d1) -> (d0 * 4 + d1)>
 // CHECK-DAG: #[[VIEW_MAP2:map[0-9]+]] = affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + d1 + s0)>
 // CHECK-DAG: #[[VIEW_MAP3:map[0-9]+]] = affine_map<(d0, d1)[s0] -> (d0 * s0 + d1)>
+// CHECK-DAG: #[[VIEW_MAP4:map[0-9]+]] = affine_map<(d0) -> (d0 + 3)>
 
 // CHECK-DAG: #[[BASE_MAP0:map[0-9]+]] = affine_map<(d0, d1, d2) -> (d0 * 64 + d1 * 4 + d2)>
 // CHECK-DAG: #[[BASE_MAP3:map[0-9]+]] = affine_map<(d0, d1, d2)[s0, s1, s2, s3] -> (d0 * s1 + s0 + d1 * s2 + d2 * s3)>
 // CHECK-DAG: #[[SUBVIEW_MAP0:map[0-9]+]] = affine_map<(d0, d1, d2)[s0, s1, s2, s3] -> (d0 * s1 + d1 * s2 + d2 * s3 + s0)>
 
 // CHECK-DAG: #[[BASE_MAP1:map[0-9]+]] = affine_map<(d0)[s0] -> (d0 + s0)>
 // CHECK-DAG: #[[SUBVIEW_MAP1:map[0-9]+]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
 
func @memref_view(%arg0 : index, %arg1 : index, %arg2 : index) {
   // CHECK: %{{.*}} = std.view %0[][%arg0] : memref<2048xi8> to memref<?x4xf32, #[[VIEW_MAP1]]>
   %4 = view %0[][%arg0]
     : memref<2048xi8> to memref<?x4xf32, affine_map<(d0, d1) -> (d0 * 4 + d1)>>
 
   // Test static sizes and static offset.
   // CHECK: %{{.*}} = std.view %0[][] : memref<2048xi8> to memref<64x4xf32, #[[VIEW_MAP1]]>
   %5 = view %0[][]
     : memref<2048xi8> to memref<64x4xf32, affine_map<(d0, d1) -> (d0 * 4 + d1)>>
+
+  // Test static sizes and static offset with base offset.
+  %6 = subview %0[][][] : memref<2048xi8> to memref<2045xi8, offset: 3, strides: [1]>
+  // CHECK: %{{.*}} = std.view %6[][] : memref<2045xi8, #[[VIEW_MAP4]]> to memref<64x4xf32, #[[VIEW_MAP1]]>
+  %7 = view %6[][]
+    : memref<2045xi8, offset: 3, strides: [1]> to memref<64x4xf32, affine_map<(d0, d1) -> (d0 * 4 + d1)>>
   return
 }
 
 // CHECK-LABEL: func @memref_subview(%arg0
 func @memref_subview(%arg0 : index, %arg1 : index, %arg2 : index) {
   %c0 = constant 0 : index
   %c1 = constant 1 : index
 

mlir/test/IR/invalid-ops.mlir

%1 = view %0[][%arg0, %arg1]
     : memref<2048xf32> to memref<?x?xf32, affine_map<(d0, d1)[s0] -> (d0 * 4 + d1 + s0)>>
   return
 }
 
 // -----
 
 func @invalid_view(%arg0 : index, %arg1 : index, %arg2 : index) {
   %0 = alloc() : memref<2048xi8, affine_map<(d0) -> (d0 floordiv 8, d0 mod 8)>>
-  // expected-error@+1 {{unsupported map for base memref}}
+  // expected-error@+1 {{base type 'memref<?x?xf32, affine_map<(d0, d1)[s0] -> (d0 * 4 + d1 + s0)>>' is not strided}}
   %1 = view %0[][%arg0, %arg1]
     : memref<2048xi8, affine_map<(d0) -> (d0 floordiv 8, d0 mod 8)>> to
       memref<?x?xf32, affine_map<(d0, d1)[s0] -> (d0 * 4 + d1 + s0)>>
   return
 }
 
 // -----