compact memory contiguity with for loops

src/PIM/Common/IR/AddressAnalysis.cpp

@@ -69,6 +69,16 @@ mlir::Value resolveLoopCarriedAliasImpl(mlir::Value value, const StaticValueKnow
 llvm::FailureOr<int64_t> resolveOpFoldResult(mlir::OpFoldResult ofr, const StaticValueKnowledge* knowledge);
 llvm::FailureOr<int64_t> resolveIndexValueImpl(mlir::Value value, const StaticValueKnowledge* knowledge);
 
+static llvm::FailureOr<llvm::SmallVector<int64_t>> getStaticMemRefStrides(mlir::MemRefType type) {
+  llvm::SmallVector<int64_t> strides;
+  int64_t offset = 0;
+  if (failed(type.getStridesAndOffset(strides, offset)))
+    return mlir::failure();
+  if (llvm::any_of(strides, mlir::ShapedType::isDynamic))
+    return mlir::failure();
+  return strides;
+}
+
 static llvm::FailureOr<int64_t> resolveConstantGlobalLoad(mlir::memref::LoadOp loadOp,
                                                           const StaticValueKnowledge* knowledge) {
   auto getGlobalOp = loadOp.getMemRef().getDefiningOp<mlir::memref::GetGlobalOp>();
@@ -539,8 +549,10 @@ llvm::FailureOr<ResolvedContiguousAddress> resolveContiguousAddressImpl(mlir::Va
       if (!isMemoryContiguous(sourceType.getShape(), offsets, sizes, strides))
         return mlir::failure();
 
-      auto sourceStrides = computeRowMajorStrides(sourceType.getShape());
-      byteOffset += linearizeIndex(offsets, sourceStrides) * getElementTypeSizeInBytes(subviewType.getElementType());
+      auto sourceStrides = getStaticMemRefStrides(sourceType);
+      if (failed(sourceStrides))
+        return mlir::failure();
+      byteOffset += linearizeIndex(offsets, *sourceStrides) * getElementTypeSizeInBytes(subviewType.getElementType());
       value = resolveAlias(subviewOp.getSource(), knowledge);
       continue;
     }
@@ -651,12 +663,16 @@ llvm::FailureOr<CompiledAddressExpr> compileContiguousAddressExprImpl(mlir::Valu
         if (!isMemoryContiguous(sourceType.getShape(), staticOffsets, staticSizes, staticStrides))
           return mlir::failure();
 
-        auto sourceStrides = computeRowMajorStrides(sourceType.getShape());
+        auto sourceStrides = getStaticMemRefStrides(sourceType);
+        if (failed(sourceStrides))
+          return mlir::failure();
         constantByteOffset +=
-          linearizeIndex(staticOffsets, sourceStrides) * getElementTypeSizeInBytes(subviewType.getElementType());
+          linearizeIndex(staticOffsets, *sourceStrides) * getElementTypeSizeInBytes(subviewType.getElementType());
       }
       else {
-        llvm::SmallVector<int64_t> sourceStrides = computeRowMajorStrides(sourceType.getShape());
+        auto sourceStrides = getStaticMemRefStrides(sourceType);
+        if (failed(sourceStrides))
+          return mlir::failure();
         CompiledIndexExpr offsetExpr;
         {
           CompiledIndexExprNode expr;
@@ -665,7 +681,7 @@ llvm::FailureOr<CompiledAddressExpr> compileContiguousAddressExprImpl(mlir::Valu
           offsetExpr = makeCompiledIndexExpr(std::move(expr));
         }
 
-        for (auto [mixedOffset, sourceStride] : llvm::zip_equal(subviewOp.getMixedOffsets(), sourceStrides)) {
+        for (auto [mixedOffset, sourceStride] : llvm::zip_equal(subviewOp.getMixedOffsets(), *sourceStrides)) {
           CompiledIndexExpr operandExpr;
           if (auto attr = mlir::dyn_cast<mlir::Attribute>(mixedOffset)) {
             CompiledIndexExprNode expr;

src/PIM/Common/IR/ConstantUtils.cpp

@@ -5,8 +5,6 @@
 #include "mlir/IR/Dialect.h"
 
 #include "ConstantUtils.hpp"
-#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
-#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 
 using namespace mlir;
 
@@ -15,16 +13,12 @@ namespace onnx_mlir {
 Block* getConstantInsertionBlock(Operation* anchorOp) {
   assert(anchorOp && "expected a valid anchor operation");
 
-  for (Operation* current = anchorOp; current; current = current->getParentOp())
-    if (isa<spatial::SpatCompute, spatial::SpatComputeBatch, pim::PimCoreOp, pim::PimCoreBatchOp>(current))
-      return current->getBlock();
-
   if (auto funcOp = dyn_cast<func::FuncOp>(anchorOp))
     return &funcOp.getBody().front();
+  if (auto funcOp = anchorOp->getParentOfType<func::FuncOp>())
+    return &funcOp.getBody().front();
   if (auto moduleOp = dyn_cast<ModuleOp>(anchorOp))
     return moduleOp.getBody();
-  if (auto funcOp = anchorOp->getParentOfType<func::FuncOp>())
-    return &funcOp.getBody().front();
   if (auto moduleOp = anchorOp->getParentOfType<ModuleOp>())
     return moduleOp.getBody();
   return anchorOp->getBlock();