[mlir][memref] Introduce memref.offset and memref.stride ops

[Hardcode84]

Summary

These ops allow to extract offset/strides from memref as ssa values.

This patch only introduces ops definiton with basic printing/parsing tests.
Additional folding and llvm lowering will be introduced in separate patches.

[Hardcode84]

Discussion https://discourse.llvm.org/t/extracting-dynamic-offsets-strides-from-memref/64170

[Hardcode84]

fix linux build

[mravishankar]

I think this has been attempted before. I dont know what the context of that discussion is. So I will leave that for more expreienced folks to weigh in on. Have some comments on the patch though.

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

759

THe naming is now confusing cause the same function is used to verify dim ops, stride ops and offset ops. I'd rather not templatize this and just copy the same logic into the individual operation verify methods.

1000

I think the stride op is valid if getStridesAndOffset method here returns success. Otherwise it might not be possible to actually compute this operation. Same is true for the offset op.

[Hardcode84]

getStridesAndOffset

[Hardcode84]

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

759

This code is too big for copypase IMO. Tried to get better name

[Hardcode84]

fix test func name

[chelini]

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

546

I would drop this line I don't think add much. Also here we have a single operand.

[mravishankar]

Dropping my self from reviewers list. This needs some one from core to approve.

[Hardcode84]

Abandoning in favor of https://reviews.llvm.org/D132243

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

def MemRef_DimOp : MemRef_Op<"dim", [NoSideEffect, MemRefsNormalizable]> {
   }];
 
   let hasCanonicalizer = 1;
   let hasFolder = 1;
   let hasVerifier = 1;
 }
 
 //===----------------------------------------------------------------------===//
+// StrideOp
+//===----------------------------------------------------------------------===//
+
+def MemRef_StrideOp : MemRef_Op<"stride", [NoSideEffect, MemRefsNormalizable]> {
+  let summary = "dimension stride operation";
+  let description = [{
+    The `stride` operation takes a memref and a dimension operand of type `index`.
+    It returns the stride of the requested dimension of the given memref.
+    If the dimension index is out of bounds the behavior is undefined.
+
+    The specified memref type is that of the first operand.
+
+    Example:
+
+    ```mlir
+    %c0 = arith.constant 0 : index
+    %x = memref.stride %A, %c0 : memref<4 x ? x f32>
+    ```
+  }];
+
+  let arguments = (ins AnyRankedOrUnrankedMemRef:$source,
+                       Index:$index);
+  let results = (outs Index:$result);
+
+  let assemblyFormat = [{
+    attr-dict $source `,` $index `:` type($source)
+  }];
+
+  let builders = [
+    OpBuilder<(ins "Value":$source, "int64_t":$index)>,
+    OpBuilder<(ins "Value":$source, "Value":$index)>
+  ];
+
+  let extraClassDeclaration = [{
+    /// Helper function to get the index as a simple integer if it is constant.
+    Optional<int64_t> getConstantIndex();
+  }];
+
+  let hasVerifier = 1;
+}
+
+//===----------------------------------------------------------------------===//
+// OffsetOp
+//===----------------------------------------------------------------------===//
+
+def MemRef_OffsetOp : MemRef_Op<"offset", [NoSideEffect, MemRefsNormalizable]> {
+  let summary = "offset operation";
+  let description = [{
+    The `offset` operation takes a memref. It returns the offset of the given
+    memref.
+
+    The specified memref type is that of the first operand.
+
+    Example:
+
+    ```mlir
+    %x = memref.offset %A : memref<4 x ? x f32>
+    ```
+  }];
+
+  let arguments = (ins AnyRankedOrUnrankedMemRef:$source);
+  let results = (outs Index:$result);
+
+  let assemblyFormat = [{
+    attr-dict $source `:` type($source)
+  }];
+
+  let builders = [
+    OpBuilder<(ins "Value":$source)>,
+  ];
+
+  let hasVerifier = 1;
+}
+
+//===----------------------------------------------------------------------===//
 // DmaStartOp
 //===----------------------------------------------------------------------===//
 
 def MemRef_DmaStartOp : MemRef_Op<"dma_start"> {
   let summary = "non-blocking DMA operation that starts a transfer";
   let description = [{
     DmaStartOp starts a non-blocking DMA operation that transfers data from a
     source memref to a destination memref. The source and destination memref

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

 }
 
 Optional<int64_t> DimOp::getConstantIndex() {
   if (auto constantOp = getIndex().getDefiningOp<arith::ConstantOp>())
     return constantOp.getValue().cast<IntegerAttr>().getInt();
   return {};
 }
 
-LogicalResult DimOp::verify() {
+template <typename Op>
+static LogicalResult verifyOpDimIndex(Op op) {
   // Assume unknown index to be in range.
-  Optional<int64_t> index = getConstantIndex();
+  Optional<int64_t> index = op.getConstantIndex();
   if (!index)
     return success();
 
   // Check that constant index is not knowingly out of range.
-  auto type = getSource().getType();
-  if (auto memrefType = type.dyn_cast<MemRefType>()) {
+  auto type = op.getSource().getType();
+  if (auto memrefType = type.template dyn_cast<MemRefType>()) {
     if (*index >= memrefType.getRank())
-      return emitOpError("index is out of range");
-  } else if (type.isa<UnrankedMemRefType>()) {
+      return op.emitOpError("index is out of range");
+  } else if (type.template isa<UnrankedMemRefType>()) {
     // Assume index to be in range.
   } else {
     llvm_unreachable("expected operand with memref type");
   }
   return success();
 }
 
+LogicalResult DimOp::verify() { return verifyOpDimIndex(*this); }
+
 /// Return a map with key being elements in `vals` and data being number of
 /// occurences of it. Use std::map, since the `vals` here are strides and the
 /// dynamic stride value is the same as the tombstone value for
 /// `DenseMap<int64_t>`.
 static std::map<int64_t, unsigned> getNumOccurences(ArrayRef<int64_t> vals) {
   std::map<int64_t, unsigned> numOccurences;
   for (auto val : vals)
     numOccurences[val]++;

 
 } // namespace
 
 void DimOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                         MLIRContext *context) {
   results.add<DimOfMemRefReshape>(context);
 }
 
+//===----------------------------------------------------------------------===//
+// StrideOp
+//===----------------------------------------------------------------------===//
+
+void StrideOp::build(OpBuilder &builder, OperationState &result, Value source,
+                     int64_t index) {
+  auto loc = result.location;
+  Value indexValue = builder.create<arith::ConstantIndexOp>(loc, index);
+  build(builder, result, source, indexValue);
+}
+
+void StrideOp::build(OpBuilder &builder, OperationState &result, Value source,
+                     Value index) {
+  auto indexTy = builder.getIndexType();
+  build(builder, result, indexTy, source, index);
+}
+
+Optional<int64_t> StrideOp::getConstantIndex() {
+  if (auto constantOp = getIndex().getDefiningOp<arith::ConstantOp>())
+    return constantOp.getValue().cast<IntegerAttr>().getInt();
+  return {};
+}
+
+LogicalResult StrideOp::verify() {
+  int64_t offset;
+  SmallVector<int64_t, 4> strides;
+  auto memrefType = getSource().getType().cast<MemRefType>();
+  if (failed(getStridesAndOffset(memrefType, strides, offset))) {
+    emitOpError("invalid memref layout");
+    return failure();
+  }
+
+  return verifyOpDimIndex(*this);
+}
+
+//===----------------------------------------------------------------------===//
+// OffsetOp
+//===----------------------------------------------------------------------===//
+
+void OffsetOp::build(OpBuilder &builder, OperationState &result, Value source) {
+  auto indexTy = builder.getIndexType();
+  build(builder, result, indexTy, source);
+}
+
+LogicalResult OffsetOp::verify() {
+  int64_t offset;
+  SmallVector<int64_t, 4> strides;
+  auto memrefType = getSource().getType().cast<MemRefType>();
+  if (failed(getStridesAndOffset(memrefType, strides, offset))) {
+    emitOpError("invalid memref layout");
+    return failure();
+  }
+
+  return success();
+}
+
 // ---------------------------------------------------------------------------
 // DmaStartOp
 // ---------------------------------------------------------------------------
 
 void DmaStartOp::build(OpBuilder &builder, OperationState &result,
                        Value srcMemRef, ValueRange srcIndices, Value destMemRef,
                        ValueRange destIndices, Value numElements,
                        Value tagMemRef, ValueRange tagIndices, Value stride,

mlir/test/Dialect/MemRef/invalid.mlir

func.func @atomic_yield_type_mismatch(%I: memref<10xf32>, %i : index) {
   // expected-error@+4 {{op types mismatch between yield op: 'i32' and its parent: 'f32'}}
   %x = memref.generic_atomic_rmw %I[%i] : memref<10xf32> {
     ^bb0(%old_value : f32):
       %c1 = arith.constant 1 : i32
       memref.atomic_yield %c1 : i32
   }
   return
 }
+
+// -----
+
+func.func @stride_invalid_index(%t : memref<4x4x?xf32>) {
+  %c5 = arith.constant 5 : index
+  // expected-error@+1 {{index is out of range}}
+  %1 = memref.stride %t, %c5 : memref<4x4x?xf32>
+}
+
+// -----
+
+func.func @stride_invalid_layout(%t : memref<?xf32, affine_map<(d0)[s0] -> (d0 mod 5)>>) {
+  %c0 = arith.constant 0 : index
+  // expected-error@+1 {{invalid memref layout}}
+  %1 = memref.stride %t, %c0 : memref<?xf32, affine_map<(d0)[s0] -> (d0 mod 5)>>
+}
+
+// -----
+
+func.func @offset_invalid_layout(%t : memref<?xf32, affine_map<(d0)[s0] -> (d0 mod 5)>>) {
+  // expected-error@+1 {{invalid memref layout}}
+  %1 = memref.offset %t : memref<?xf32, affine_map<(d0)[s0] -> (d0 mod 5)>>
+}

mlir/test/Dialect/MemRef/ops.mlir

// CHECK-NEXT: memref.generic_atomic_rmw [[BUF]]{{\[}}[[I]], [[J]]] : memref
     ^bb0(%old_value : f32):
       %c1 = arith.constant 1.0 : f32
       %out = arith.addf %c1, %old_value : f32
       memref.atomic_yield %out : f32
   // CHECK: index_attr = 8 : index
   } { index_attr = 8 : index }
   return
 }
+
+// -----
+
+// CHECK-LABEL: func @stride
+// CHECK-SAME: (%[[MEMREF:.*]]: memref<4x4x?xf32>, %[[D:.*]]: index)
+func.func @stride(%t : memref<4x4x?xf32>, %d : index) {
+  // CHECK: %{{.*}} = memref.stride %[[MEMREF]], %[[D]] : memref<4x4x?xf32>
+  %0 = "memref.stride"(%t, %d) : (memref<4x4x?xf32>, index) -> index
+
+  // CHECK: %{{.*}} = memref.stride %[[MEMREF]], %[[D]] : memref<4x4x?xf32>
+  %1 = memref.stride %t, %d  : memref<4x4x?xf32>
+  return
+}
+
+// -----
+
+// CHECK-LABEL: func @offset
+// CHECK-SAME: (%[[MEMREF:.*]]: memref<4x4x?xf32>)
+func.func @offset(%t : memref<4x4x?xf32>) {
+  // CHECK: %{{.*}} = memref.offset %[[MEMREF]] : memref<4x4x?xf32>
+  %0 = "memref.offset"(%t) : (memref<4x4x?xf32>) -> index
+
+  // CHECK: %{{.*}} = memref.offset %[[MEMREF]] : memref<4x4x?xf32>
+  %1 = memref.offset %t : memref<4x4x?xf32>
+  return
+}