Extend polygeist subindex lowering to multidim memrefs

Author: mortbopet

lib/polygeist/Passes/LowerPolygeistOps.cpp

@@ -46,9 +46,11 @@ struct SubIndexToReinterpretCast
         rewriter.create<ConstantIndexOp>(op.getLoc(), innerSize));
 
     llvm::SmallVector<OpFoldResult> sizes, strides;
-    for (auto dim : shape.drop_front()) {
-      sizes.push_back(rewriter.getIndexAttr(dim));
-      strides.push_back(rewriter.getIndexAttr(1));
+    int64_t strideAcc = 1;
+    for (auto dim : llvm::reverse(shape.drop_front())) {
+      sizes.insert(sizes.begin(), rewriter.getIndexAttr(dim));
+      strides.insert(strides.begin(), rewriter.getIndexAttr(strideAcc));
+      strideAcc *= dim;
     }
 
     rewriter.replaceOpWithNewOp<memref::ReinterpretCastOp>(

test/polygeist-opt/lower_polygeist_ops.mlir

@@ -8,10 +8,25 @@
 // CHECK:           %[[VAL_4:.*]] = memref.reinterpret_cast %[[VAL_1]] 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>
-    %1 = "polygeist.subindex"(%0, %arg0) : (memref<30x30xi32>, index) -> memref<30xi32>
-    return %1 : memref<30xi32>
-  }
+func @main(%arg0 : index) -> memref<30xi32> {
+  %0 = memref.alloca() : memref<30x30xi32>
+  %1 = "polygeist.subindex"(%0, %arg0) : (memref<30x30xi32>, index) -> memref<30xi32>
+  return %1 : memref<30xi32>
+}
+
+
+// -----
+
+// CHECK-LABEL:   func @main(
+// CHECK-SAME:               %[[VAL_0:.*]]: index) -> memref<42x43x44x45xi32> {
+// CHECK:           %[[VAL_1:.*]] = memref.alloca() : memref<41x42x43x44x45xi32>
+// CHECK:           %[[VAL_2:.*]] = arith.constant 3575880 : index
+// CHECK:           %[[VAL_3:.*]] = arith.muli %[[VAL_0]], %[[VAL_2]] : index
+// CHECK:           %[[VAL_4:.*]] = memref.reinterpret_cast %[[VAL_1]] to offset: {{\[}}%[[VAL_3]]], sizes: [42, 43, 44, 45], strides: [85140, 1980, 45, 1] : memref<41x42x43x44x45xi32> to memref<42x43x44x45xi32>
+// CHECK:           return %[[VAL_4]] : memref<42x43x44x45xi32>
+// CHECK:         }
+func @main(%arg0 : index) -> memref<42x43x44x45xi32> {
+  %0 = memref.alloca() : memref<41x42x43x44x45xi32>
+  %1 = "polygeist.subindex"(%0, %arg0) : (memref<41x42x43x44x45xi32>, index) -> memref<42x43x44x45xi32>
+  return %1 : memref<42x43x44x45xi32>
 }