Lower polygeist.subindex through memref.reinterpret_cast

This should be a (hopefully) foolproof method of performing indexing into a memref. A reintrepret_cast is inserted with a dynamic index calculated from the subindex index operand + the producto of the sizes of the target type.

Also re-enables the existing canonicalization tests.

Author: mortbopet

lib/polygeist/Ops.cpp

@@ -176,43 +176,33 @@ class SubToCast final : public OpRewritePattern<SubIndexOp> {
   }
 };
 
-// Simplify polygeist.subindex to memref.subview.
-class SubToSubView final : public OpRewritePattern<SubIndexOp> {
+// Simplify polygeist.subindex to a memref.reinterpret_cast.
+class SubToReinterpretCast final : public OpRewritePattern<SubIndexOp> {
 public:
   using OpRewritePattern<SubIndexOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(SubIndexOp op,
                                 PatternRewriter &rewriter) const override {
     auto srcMemRefType = op.source().getType().cast<MemRefType>();
     auto resMemRefType = op.result().getType().cast<MemRefType>();
-    auto dims = srcMemRefType.getShape().size();
+    auto shape = srcMemRefType.getShape();
 
-    // For now, restrict subview lowering to statically defined memref's
-    if (!srcMemRefType.hasStaticShape() | !resMemRefType.hasStaticShape())
+    if (!resMemRefType.hasStaticShape())
       return failure();
 
-    // For now, restrict to simple rank-reducing indexing
-    if (srcMemRefType.getShape().size() <= resMemRefType.getShape().size())
-      return failure();
+    int64_t innerSize = resMemRefType.getNumElements();
+    auto offset = rewriter.create<arith::MulIOp>(
+        op.getLoc(), op.index(),
+        rewriter.create<ConstantIndexOp>(op.getLoc(), innerSize));
 
-    // Build offset, sizes and strides
-    SmallVector<OpFoldResult> sizes(dims, rewriter.getIndexAttr(0));
-    sizes[0] = op.index();
-    SmallVector<OpFoldResult> offsets(dims);
-    for (auto dim : llvm::enumerate(srcMemRefType.getShape())) {
-      if (dim.index() == 0)
-        offsets[0] = rewriter.getIndexAttr(1);
-      else
-        offsets[dim.index()] = rewriter.getIndexAttr(dim.value());
+    llvm::SmallVector<OpFoldResult> sizes, strides;
+    for (auto dim : shape.drop_front()) {
+      sizes.push_back(rewriter.getIndexAttr(dim));
+      strides.push_back(rewriter.getIndexAttr(1));
     }
-    SmallVector<OpFoldResult> strides(dims, rewriter.getIndexAttr(1));
-
-    // Generate the appropriate return type:
-    auto subMemRefType = MemRefType::get(srcMemRefType.getShape().drop_front(),
-                                         srcMemRefType.getElementType());
 
-    rewriter.replaceOpWithNewOp<memref::SubViewOp>(
-        op, subMemRefType, op.source(), sizes, offsets, strides);
+    rewriter.replaceOpWithNewOp<memref::ReinterpretCastOp>(
+        op, resMemRefType, op.source(), offset.getResult(), sizes, strides);
 
     return success();
   }
@@ -677,8 +667,8 @@ void SubIndexOp::getCanonicalizationPatterns(OwningRewritePatternList &results,
                                              MLIRContext *context) {
   results.insert<CastOfSubIndex, SubIndex2, SubToCast, SimplifySubViewUsers,
                  SimplifySubIndexUsers, SelectOfCast, SelectOfSubIndex,
-                 RedundantDynSubIndex>(context);
-  // Disabled: SubToSubView
+                 RedundantDynSubIndex, SubToReinterpretCast>(context);
+  // Disabled:
 }
 
 /// Simplify pointer2memref(memref2pointer(x)) to cast(x)

test/polygeist-opt/canonicalization.mlir

@@ -1,10 +1,13 @@
 // RUN: polygeist-opt --canonicalize --split-input-file %s | FileCheck %s
-// XFAIL: *
-// CHECK: func @main(%arg0: index) -> memref<30xi32> {
-// CHECK:   %0 = memref.alloca() : memref<30x30xi32>
-// CHECK:   %1 = memref.subview %0[%arg0, 0] [1, 30] [1, 1] : memref<30x30xi32> to memref<30xi32>
-// CHECK:   return %1 : memref<30xi32>
-// CHECK: }
+
+// CHECK-LABEL:   func @main(
+// CHECK-SAME:               %[[VAL_0:.*]]: index) -> memref<30xi32> {
+// CHECK:           %[[VAL_1:.*]] = arith.constant 30 : index
+// CHECK:           %[[VAL_2:.*]] = memref.alloca() : memref<30x30xi32>
+// CHECK:           %[[VAL_3:.*]] = arith.muli %[[VAL_0]], %[[VAL_1]] : index
+// CHECK:           %[[VAL_4:.*]] = memref.reinterpret_cast %[[VAL_2]] to offset: {{\[}}%[[VAL_3]]], sizes: [30], strides: [1] : memref<30x30xi32> to memref<30xi32>
+// CHECK:           return %[[VAL_4]] : memref<30xi32>
+// CHECK:         }
 module {
   func @main(%arg0 : index) -> memref<30xi32> {
     %0 = memref.alloca() : memref<30x30xi32>
@@ -15,11 +18,14 @@ module {
 
 // -----
 
-// CHECK:  func @main(%arg0: index) -> memref<1000xi32> {
-// CHECK:    %0 = memref.alloca() : memref<2x1000xi32>
-// CHECK:    %1 = memref.subview %0[%arg0, 0] [1, 1000] [1, 1] : memref<2x1000xi32> to memref<1000xi32>
-// CHECK:    return %1 : memref<1000xi32>
-// CHECK:  }
+// CHECK-LABEL:   func @main(
+// CHECK-SAME:               %[[VAL_0:.*]]: index) -> memref<1000xi32> {
+// CHECK:           %[[VAL_1:.*]] = arith.constant 1000 : index
+// CHECK:           %[[VAL_2:.*]] = memref.alloca() : memref<2x1000xi32>
+// CHECK:           %[[VAL_3:.*]] = arith.muli %[[VAL_0]], %[[VAL_1]] : index
+// CHECK:           %[[VAL_4:.*]] = memref.reinterpret_cast %[[VAL_2]] to offset: {{\[}}%[[VAL_3]]], sizes: [1000], strides: [1] : memref<2x1000xi32> to memref<1000xi32>
+// CHECK:           return %[[VAL_4]] : memref<1000xi32>
+// CHECK:         }
 func @main(%arg0 : index) -> memref<1000xi32> {
   %c0 = arith.constant 0 : index
   %1 = memref.alloca() : memref<2x1000xi32>