Extend polygeist subindex lowering to multidim memrefs

llvm · Jan 29, 2022 · afc0b95 · afc0b95
1 parent f7afb69
commit afc0b95
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/lib/polygeist/Passes/LowerPolygeistOps.cpp b/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>(