Also handle non-rank reducing offset operations

Author: mortbopet

lib/polygeist/Passes/LowerPolygeistOps.cpp

@@ -37,19 +37,20 @@ struct SubIndexToReinterpretCast
                   ConversionPatternRewriter &rewriter) const override {
     auto srcMemRefType = op.source().getType().cast<MemRefType>();
     auto resMemRefType = op.result().getType().cast<MemRefType>();
-    auto shape = srcMemRefType.getShape();
+    auto inShape = srcMemRefType.getShape();
+    auto outShape = resMemRefType.getShape();
 
     if (!resMemRefType.hasStaticShape())
       return failure();
 
+    llvm::SmallVector<OpFoldResult> strides, sizes;
     int64_t innerSize = resMemRefType.getNumElements();
     auto offset = rewriter.create<arith::MulIOp>(
         op.getLoc(), op.index(),
         rewriter.create<ConstantIndexOp>(op.getLoc(), innerSize));
 
-    llvm::SmallVector<OpFoldResult> sizes, strides;
     int64_t strideAcc = 1;
-    for (auto dim : llvm::reverse(shape.drop_front())) {
+    for (auto dim : llvm::reverse(outShape)) {
       sizes.insert(sizes.begin(), rewriter.getIndexAttr(dim));
       strides.insert(strides.begin(), rewriter.getIndexAttr(strideAcc));
       strideAcc *= dim;

test/polygeist-opt/lower_polygeist_ops.mlir

@@ -14,7 +14,6 @@ func @main(%arg0 : index) -> memref<30xi32> {
   return %1 : memref<30xi32>
 }
 
-
 // -----
 
 // CHECK-LABEL:   func @main(
@@ -30,3 +29,20 @@ func @main(%arg0 : index) -> memref<42x43x44x45xi32> {
   %1 = "polygeist.subindex"(%0, %arg0) : (memref<41x42x43x44x45xi32>, index) -> memref<42x43x44x45xi32>
   return %1 : memref<42x43x44x45xi32>
 }
+
+// -----
+
+// CHECK-LABEL:   func @main(
+// CHECK-SAME:               %[[VAL_0:.*]]: index) -> memref<29x30xi32> {
+// CHECK:           %[[VAL_1:.*]] = memref.alloca() : memref<30x30xi32>
+// CHECK:           %[[VAL_2:.*]] = arith.constant 870 : index
+// CHECK:           %[[VAL_3:.*]] = arith.muli %[[VAL_0]], %[[VAL_2]] : index
+// CHECK:           %[[VAL_4:.*]] = memref.reinterpret_cast %[[VAL_1]] to offset: {{\[}}%[[VAL_3]]], sizes: [29, 30], strides: [30, 1] : memref<30x30xi32> to memref<29x30xi32>
+// CHECK:           return %[[VAL_4]] : memref<29x30xi32>
+// CHECK:         }
+
+func @main(%arg0 : index) -> memref<29x30xi32> {
+  %0 = memref.alloca() : memref<30x30xi32>
+  %1 = "polygeist.subindex"(%0, %arg0) : (memref<30x30xi32>, index) -> memref<29x30xi32>
+  return %1 : memref<29x30xi32>
+}