add shared loop creation helpers

add shared checked arithmetic helpers
refactor pim passes into Pim/Transforms
more robust memory coalescing pass

src/PIM/Conversion/SpatialToPim/BatchCoreLoweringPatterns.cpp

@@ -10,6 +10,7 @@
 #include "Conversion/SpatialToPim/SpatialToPimPass.hpp"
 #include "src/Accelerators/PIM/Common/IR/BatchCoreUtils.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
+#include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
 #include "src/Accelerators/PIM/Conversion/SpatialToPim/Common.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
 
@@ -25,14 +26,21 @@ static bool isUsedOnlyAsExplicitHostOperand(Value value) {
   });
 }
 
-static SmallVector<int32_t> getPimCoreIdsForBatchOp(spatial::SpatComputeBatch computeBatchOp, size_t& fallbackCoreId) {
+static FailureOr<SmallVector<int32_t>> getPimCoreIdsForBatchOp(spatial::SpatComputeBatch computeBatchOp,
+                                                               size_t& fallbackCoreId) {
   if (auto coreIdsAttr = computeBatchOp->getAttrOfType<DenseI32ArrayAttr>(onnx_mlir::kCoreIdsAttrName))
     return SmallVector<int32_t>(coreIdsAttr.asArrayRef().begin(), coreIdsAttr.asArrayRef().end());
 
   SmallVector<int32_t> coreIds;
   coreIds.reserve(static_cast<size_t>(computeBatchOp.getLaneCount()));
-  for (uint32_t lane = 0; lane < computeBatchOp.getLaneCount(); ++lane)
-    coreIds.push_back(static_cast<int32_t>(fallbackCoreId++));
+  for (uint32_t lane = 0; lane < computeBatchOp.getLaneCount(); ++lane) {
+    auto checkedCoreId =
+      pim::checkedI32(static_cast<uint64_t>(fallbackCoreId), computeBatchOp, "fallback spatial compute_batch core id");
+    if (failed(checkedCoreId))
+      return failure();
+    coreIds.push_back(*checkedCoreId);
+    ++fallbackCoreId;
+  }
   return coreIds;
 }
 
@@ -102,21 +110,24 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeBatchOp(spatial::SpatCompute
       "resultful compute_batch lowering currently requires a spat.in_parallel terminator");
   }
 
-  SmallVector<int32_t> coreIds = getPimCoreIdsForBatchOp(computeBatchOp, coreId);
+  auto coreIds = getPimCoreIdsForBatchOp(computeBatchOp, coreId);
+  if (failed(coreIds))
+    return failure();
   SmallVector<Value> batchWeights(computeBatchOp.getWeights().begin(), computeBatchOp.getWeights().end());
   SmallVector<Value> batchInputs;
   if (!computeBatchOp.getInputs().empty())
     batchInputs.append(computeBatchOp.getInputs().begin(), computeBatchOp.getInputs().end());
 
   rewriter.setInsertionPointAfter(computeBatchOp);
-  auto coreBatchOp = pim::PimCoreBatchOp::create(rewriter,
-                                                 loc,
-                                                 rewriter.getI32IntegerAttr(computeBatchOp.getLaneCount()),
-                                                 ValueRange(batchWeights),
-                                                 ValueRange(batchInputs));
+  auto laneCountAttr = pim::getCheckedI32Attr(
+    rewriter, computeBatchOp, static_cast<uint64_t>(computeBatchOp.getLaneCount()), "pim core_batch lane count");
+  if (failed(laneCountAttr))
+    return failure();
+  auto coreBatchOp =
+    pim::PimCoreBatchOp::create(rewriter, loc, *laneCountAttr, ValueRange(batchWeights), ValueRange(batchInputs));
   coreBatchOp.getProperties().setOperandSegmentSizes(
     {static_cast<int>(batchWeights.size()), static_cast<int>(batchInputs.size())});
-  coreBatchOp->setAttr(onnx_mlir::kCoreIdsAttrName, rewriter.getDenseI32ArrayAttr(coreIds));
+  coreBatchOp->setAttr(onnx_mlir::kCoreIdsAttrName, rewriter.getDenseI32ArrayAttr(*coreIds));
 
   SmallVector<unsigned> returnOperandIndices;
   if (computeBatchOp.getNumResults() != 0) {
@@ -160,14 +171,11 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeBatchOp(spatial::SpatCompute
     auto newArgType = cast<ShapedType>(newArg.getType());
     auto outputBuffer = createEmptyTensorFromShaped(rewriter, loc, newArgType);
     Value zeroOffset = getOrCreateIndexConstant(rewriter, coreBatchOp.getOperation(), 0);
-    auto copied = pim::PimMemCopyHostToDevOp::create(rewriter,
-                                                     loc,
-                                                     outputBuffer.getType(),
-                                                     zeroOffset,
-                                                     zeroOffset,
-                                                     outputBuffer,
-                                                     newArg,
-                                                     getTensorSizeInBytesAttr(rewriter, newArg))
+    auto sizeAttr = getTensorSizeInBytesAttr(rewriter, coreBatchOp.getOperation(), newArg);
+    if (failed(sizeAttr))
+      return failure();
+    auto copied = pim::PimMemCopyHostToDevOp::create(
+                    rewriter, loc, outputBuffer.getType(), zeroOffset, zeroOffset, outputBuffer, newArg, *sizeAttr)
                     .getOutput();
     mapper.map(*oldArg, copied);
   }
@@ -209,6 +217,9 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeBatchOp(spatial::SpatCompute
         auto hostTargetType = cast<ShapedType>(hostTarget.getType());
         Value hostTargetOffset = createHostTargetOffset(rewriter, insertSlice, hostTargetType, mapper);
         Value zeroOffset = getOrCreateIndexConstant(rewriter, coreBatchOp.getOperation(), 0);
+        auto sizeAttr = getTensorSizeInBytesAttr(rewriter, coreBatchOp.getOperation(), mappedSource);
+        if (failed(sizeAttr))
+          return failure();
         pim::PimMemCopyDevToHostOp::create(rewriter,
                                            insertSlice.getLoc(),
                                            hostTarget.getType(),
@@ -216,7 +227,7 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeBatchOp(spatial::SpatCompute
                                            zeroOffset,
                                            hostTarget,
                                            mappedSource,
-                                           getTensorSizeInBytesAttr(rewriter, mappedSource));
+                                           *sizeAttr);
       }
       continue;
     }
@@ -232,15 +243,13 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeBatchOp(spatial::SpatCompute
         auto clonedType = cast<ShapedType>(clonedTensor.getType());
         auto outputBuffer = createEmptyTensorFromShaped(rewriter, loc, clonedType);
         Value zeroOffset = getOrCreateIndexConstant(rewriter, coreBatchOp.getOperation(), 0);
-        auto copied = pim::PimMemCopyHostToDevOp::create(rewriter,
-                                                         loc,
-                                                         outputBuffer.getType(),
-                                                         zeroOffset,
-                                                         zeroOffset,
-                                                         outputBuffer,
-                                                         clonedTensor,
-                                                         getTensorSizeInBytesAttr(rewriter, clonedTensor))
-                        .getOutput();
+        auto sizeAttr = getTensorSizeInBytesAttr(rewriter, coreBatchOp.getOperation(), clonedTensor);
+        if (failed(sizeAttr))
+          return failure();
+        auto copied =
+          pim::PimMemCopyHostToDevOp::create(
+            rewriter, loc, outputBuffer.getType(), zeroOffset, zeroOffset, outputBuffer, clonedTensor, *sizeAttr)
+            .getOutput();
         mapper.map(toTensorOp.getResult(), copied);
         continue;
       }

src/PIM/Conversion/SpatialToPim/Common.cpp

@@ -5,14 +5,18 @@
 #include <cassert>
 
 #include "Common.hpp"
+#include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
 
 using namespace llvm;
 using namespace mlir;
 
 namespace onnx_mlir {
 
-IntegerAttr getTensorSizeInBytesAttr(Builder& builder, mlir::Value value) {
-  return builder.getI32IntegerAttr(static_cast<int32_t>(getShapedTypeSizeInBytes(cast<ShapedType>(value.getType()))));
+FailureOr<IntegerAttr> getTensorSizeInBytesAttr(Builder& builder, Operation* anchor, mlir::Value value) {
+  auto byteSize = pim::getCheckedShapedTypeSizeInBytes(cast<ShapedType>(value.getType()), anchor, "tensor byte size");
+  if (failed(byteSize))
+    return failure();
+  return pim::getCheckedI32Attr(builder, anchor, *byteSize, "tensor byte size");
 }
 
 Operation* getEarliestUserWithinBlock(mlir::Value value) {

src/PIM/Conversion/SpatialToPim/Common.hpp

@@ -1,12 +1,14 @@
 #pragma once
 
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Support/LogicalResult.h"
 
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 
 namespace onnx_mlir {
 
-mlir::IntegerAttr getTensorSizeInBytesAttr(mlir::Builder& builder, mlir::Value value);
+mlir::FailureOr<mlir::IntegerAttr>
+getTensorSizeInBytesAttr(mlir::Builder& builder, mlir::Operation* anchor, mlir::Value value);
 
 template <class T>
 size_t rangeLength(const mlir::iterator_range<T> range) {

src/PIM/Conversion/SpatialToPim/CoreLoweringPatterns.cpp

@@ -9,6 +9,7 @@
 #include "Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "Conversion/SpatialToPim/SpatialToPimPass.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
+#include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
 #include "src/Accelerators/PIM/Conversion/SpatialToPim/Common.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
@@ -54,10 +55,15 @@ cloneMappedHelperOperands(Operation* op, IRMapping& mapping, IRRewriter& rewrite
   }
 }
 
-static int32_t getPimCoreIdForComputeOp(spatial::SpatCompute computeOp, size_t& fallbackCoreId) {
+static FailureOr<int32_t> getPimCoreIdForComputeOp(spatial::SpatCompute computeOp, size_t& fallbackCoreId) {
   if (auto spatialCoreIdAttr = computeOp->getAttrOfType<IntegerAttr>(onnx_mlir::kCoreIdAttrName))
-    return static_cast<int32_t>(spatialCoreIdAttr.getInt());
-  return static_cast<int32_t>(fallbackCoreId++);
+    return pim::checkedI32(spatialCoreIdAttr.getInt(), computeOp, "spatial compute core id");
+  auto checkedCoreId =
+    pim::checkedI32(static_cast<uint64_t>(fallbackCoreId), computeOp, "fallback spatial compute core id");
+  if (failed(checkedCoreId))
+    return failure();
+  ++fallbackCoreId;
+  return *checkedCoreId;
 }
 
 static LogicalResult collectHelperComputeChain(spatial::SpatCompute computeOp,
@@ -163,10 +169,12 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeOp(spatial::SpatCompute comp
       rewriter.setInsertionPoint(getEarliestUserWithinBlock(*blockArg));
       auto outputType = cast<ShapedType>(blockArg->getType());
       auto outputBuffer = createEmptyTensorFromShaped(rewriter, receiveOp.getLoc(), outputType);
-      auto sizeAttr = getTensorSizeInBytesAttr(rewriter, *blockArg);
+      auto sizeAttr = getTensorSizeInBytesAttr(rewriter, computeOp.getOperation(), *blockArg);
+      if (failed(sizeAttr))
+        return failure();
       Value received =
         PimReceiveOp::create(
-          rewriter, receiveOp.getLoc(), outputBuffer.getType(), outputBuffer, sizeAttr, receiveOp.getSourceCoreId())
+          rewriter, receiveOp.getLoc(), outputBuffer.getType(), outputBuffer, *sizeAttr, receiveOp.getSourceCoreId())
           .getOutput();
       blockArg->replaceAllUsesWith(received);
       markOpToRemove(receiveOp);
@@ -206,8 +214,13 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeOp(spatial::SpatCompute comp
   if (!computeOp.getWeights().empty())
     computeWeights.append(computeOp.getWeights().begin(), computeOp.getWeights().end());
   rewriter.setInsertionPointAfter(computeOp);
-  auto coreOp = PimCoreOp::create(
-    rewriter, loc, ValueRange(computeWeights), rewriter.getI32IntegerAttr(getPimCoreIdForComputeOp(computeOp, coreId)));
+  auto checkedCoreId = getPimCoreIdForComputeOp(computeOp, coreId);
+  if (failed(checkedCoreId))
+    return failure();
+  auto coreIdAttr = pim::getCheckedI32Attr(rewriter, computeOp, static_cast<int64_t>(*checkedCoreId), "pim core id");
+  if (failed(coreIdAttr))
+    return failure();
+  auto coreOp = PimCoreOp::create(rewriter, loc, ValueRange(computeWeights), *coreIdAttr);
   rewriter.setInsertionPointToStart(&block);
   auto& coreOpBlocks = coreOp.getBody().getBlocks();
   for (auto [inputIndex, input] : llvm::enumerate(computeOp.getInputs())) {
@@ -226,6 +239,9 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeOp(spatial::SpatCompute comp
     if (!inputType)
       return computeOp.emitOpError("expected shaped compute input during pim.core lowering");
     auto outputBuffer = createEmptyTensorFromShaped(rewriter, loc, inputType);
+    auto sizeAttr = getTensorSizeInBytesAttr(rewriter, computeOp.getOperation(), input);
+    if (failed(sizeAttr))
+      return failure();
     auto copied =
       PimMemCopyHostToDevOp::create(rewriter,
                                     loc,
@@ -234,7 +250,7 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeOp(spatial::SpatCompute comp
                                     getOrCreateIndexConstant(constantFolder, outputBuffer.getOperation(), 0),
                                     outputBuffer,
                                     input,
-                                    getTensorSizeInBytesAttr(rewriter, input))
+                                    *sizeAttr)
         .getOutput();
     blockArg->replaceAllUsesWith(copied);
   }

src/PIM/Conversion/SpatialToPim/Patterns/ChannelLowering.cpp

@@ -14,8 +14,10 @@ struct ChannelSendLowering : OpRewritePattern<spatial::SpatChannelSendOp> {
   using OpRewritePattern::OpRewritePattern;
 
   LogicalResult matchAndRewrite(spatial::SpatChannelSendOp op, PatternRewriter& rewriter) const override {
-    pim::PimSendOp::create(
-      rewriter, op.getLoc(), op.getInput(), getTensorSizeInBytesAttr(rewriter, op.getInput()), op.getTargetCoreId());
+    auto sizeAttr = getTensorSizeInBytesAttr(rewriter, op.getOperation(), op.getInput());
+    if (failed(sizeAttr))
+      return failure();
+    pim::PimSendOp::create(rewriter, op.getLoc(), op.getInput(), *sizeAttr, op.getTargetCoreId());
     rewriter.eraseOp(op);
     return success();
   }
@@ -32,12 +34,11 @@ struct ChannelReceiveLowering : OpRewritePattern<spatial::SpatChannelReceiveOp>
     auto outputType = cast<ShapedType>(op.getResult().getType());
     Value outputBuffer =
       tensor::EmptyOp::create(rewriter, op.getLoc(), outputType.getShape(), outputType.getElementType()).getResult();
-    Value received = pim::PimReceiveOp::create(rewriter,
-                                               op.getLoc(),
-                                               op.getResult().getType(),
-                                               outputBuffer,
-                                               getTensorSizeInBytesAttr(rewriter, op.getResult()),
-                                               op.getSourceCoreId())
+    auto sizeAttr = getTensorSizeInBytesAttr(rewriter, op.getOperation(), op.getResult());
+    if (failed(sizeAttr))
+      return failure();
+    Value received = pim::PimReceiveOp::create(
+                       rewriter, op.getLoc(), op.getResult().getType(), outputBuffer, *sizeAttr, op.getSourceCoreId())
                        .getOutput();
     rewriter.replaceOp(op, received);
     return success();

src/PIM/Conversion/SpatialToPim/ReturnPathNormalization.cpp

@@ -12,6 +12,7 @@
 #include "Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "Conversion/SpatialToPim/SpatialToPimPass.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
+#include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
 #include "src/Accelerators/PIM/Conversion/SpatialToPim/Common.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
@@ -71,6 +72,20 @@ static SmallVector<int64_t> expandFlatElementIndex(int64_t flatIndex, ArrayRef<i
   return indices;
 }
 
+static FailureOr<int64_t>
+getCheckedByteOffset(int64_t elementOffset, size_t elementSize, Operation* anchor, StringRef fieldName) {
+  if (elementOffset < 0) {
+    anchor->emitOpError() << fieldName << " requires a nonnegative element offset";
+    return failure();
+  }
+
+  auto byteOffset =
+    pim::checkedMul(static_cast<uint64_t>(elementOffset), static_cast<uint64_t>(elementSize), anchor, fieldName);
+  if (failed(byteOffset))
+    return failure();
+  return pim::checkedCast<int64_t>(*byteOffset, anchor, fieldName);
+}
+
 static LogicalResult collectHelperComputeChain(spatial::SpatCompute computeOp,
                                                SmallVectorImpl<Operation*>& helperChain) {
   if (computeOp.getInputs().size() != 1 || computeOp.getNumResults() != 1)
@@ -360,18 +375,21 @@ static void cloneHelperChain(Value sourceValue,
   }
 }
 
-static Value emitHostCopy(IRRewriter& rewriter,
-                          Location loc,
-                          Value outputTensor,
-                          Value sourceValue,
-                          int32_t hostTargetOffset,
-                          int32_t deviceSourceOffset,
-                          int32_t sizeInBytes,
-                          OperationFolder& constantFolder) {
+static FailureOr<Value> emitHostCopy(IRRewriter& rewriter,
+                                     Location loc,
+                                     Value outputTensor,
+                                     Value sourceValue,
+                                     int64_t hostTargetOffset,
+                                     int64_t deviceSourceOffset,
+                                     uint64_t sizeInBytes,
+                                     OperationFolder& constantFolder) {
   Operation* anchorOp = sourceValue.getDefiningOp() ? sourceValue.getDefiningOp() : outputTensor.getDefiningOp();
   assert(anchorOp && "expected a concrete op anchor for return-path host copy constants");
   Value hostTargetOffsetValue = getOrCreateIndexConstant(constantFolder, anchorOp, hostTargetOffset);
   Value deviceSourceOffsetValue = getOrCreateIndexConstant(constantFolder, anchorOp, deviceSourceOffset);
+  auto sizeAttr = pim::getCheckedI32Attr(rewriter, anchorOp, sizeInBytes, "return-path host copy byte size");
+  if (failed(sizeAttr))
+    return failure();
   return PimMemCopyDevToHostOp::create(rewriter,
                                        loc,
                                        outputTensor.getType(),
@@ -379,7 +397,7 @@ static Value emitHostCopy(IRRewriter& rewriter,
                                        deviceSourceOffsetValue,
                                        outputTensor,
                                        sourceValue,
-                                       rewriter.getI32IntegerAttr(sizeInBytes))
+                                       *sizeAttr)
     .getOutput();
 }
 
@@ -433,18 +451,15 @@ raptor::SpatialToPimPass::ReturnPathLoweringResult raptor::SpatialToPimPass::low
       markOpToRemove(op);
 
     auto storedType = cast<ShapedType>(currentStoredValue.getType());
-    size_t elementSize = getElementTypeSizeInBytes(storedType.getElementType());
+    auto byteSize = pim::getCheckedShapedTypeSizeInBytes(storedType, producerOp, "return-path host copy byte size");
+    if (failed(byteSize))
+      return ReturnPathLoweringResult::Failure;
     if (auto storedOp = currentStoredValue.getDefiningOp())
       rewriter.setInsertionPointAfter(storedOp);
     Value outputTensor = outputTensors[returnUse->returnIndex](rewriter, loc);
-    emitHostCopy(rewriter,
-                 loc,
-                 outputTensor,
-                 currentStoredValue,
-                 0,
-                 0,
-                 static_cast<int32_t>(storedType.getNumElements() * elementSize),
-                 constantFolder);
+    auto copied = emitHostCopy(rewriter, loc, outputTensor, currentStoredValue, 0, 0, *byteSize, constantFolder);
+    if (failed(copied))
+      return ReturnPathLoweringResult::Failure;
     return ReturnPathLoweringResult::Handled;
   }
 
@@ -455,23 +470,25 @@ raptor::SpatialToPimPass::ReturnPathLoweringResult raptor::SpatialToPimPass::low
 
     if (isa<func::ReturnOp>(resultUser)) {
       size_t resultIndexInReturn = resultUse.getOperandNumber();
-      size_t elementSize = getElementTypeSizeInBytes(storedTensorType.getElementType());
+      auto byteSize =
+        pim::getCheckedShapedTypeSizeInBytes(storedTensorType, producerOp, "return-path host copy byte size");
+      if (failed(byteSize))
+        return ReturnPathLoweringResult::Failure;
       rewriter.setInsertionPointAfterValue(storedValue);
       Value outputTensor = outputTensors[resultIndexInReturn](rewriter, loc);
-      emitHostCopy(rewriter,
-                   loc,
-                   outputTensor,
-                   storedValue,
-                   0,
-                   0,
-                   static_cast<int32_t>(storedTensorType.getNumElements() * elementSize),
-                   constantFolder);
+      auto copied = emitHostCopy(rewriter, loc, outputTensor, storedValue, 0, 0, *byteSize, constantFolder);
+      if (failed(copied))
+        return ReturnPathLoweringResult::Failure;
       return ReturnPathLoweringResult::Handled;
     }
   }
 
   if (auto concatReturnUse = analyzeConcatReturnUse(producedValue)) {
     size_t elementSize = getElementTypeSizeInBytes(storedTensorType.getElementType());
+    auto storedByteSize =
+      pim::getCheckedShapedTypeSizeInBytes(storedTensorType, producerOp, "concat return-path copy byte size");
+    if (failed(storedByteSize))
+      return ReturnPathLoweringResult::Failure;
     for (Operation* concatOp : concatReturnUse->concatChain)
       markOpToRemove(concatOp);
 
@@ -480,14 +497,13 @@ raptor::SpatialToPimPass::ReturnPathLoweringResult raptor::SpatialToPimPass::low
       Value outputTensor = outputTensors[concatReturnUse->returnIndex](rewriter, loc);
       auto outputType = cast<ShapedType>(outputTensor.getType());
       int64_t flatOffset = computeFlatElementIndex(concatReturnUse->sliceOffsets, outputType.getShape());
-      emitHostCopy(rewriter,
-                   loc,
-                   outputTensor,
-                   storedValue,
-                   static_cast<int32_t>(flatOffset * elementSize),
-                   0,
-                   static_cast<int32_t>(storedTensorType.getNumElements() * elementSize),
-                   constantFolder);
+      auto hostOffset = getCheckedByteOffset(flatOffset, elementSize, producerOp, "concat return-path host offset");
+      if (failed(hostOffset))
+        return ReturnPathLoweringResult::Failure;
+      auto copied =
+        emitHostCopy(rewriter, loc, outputTensor, storedValue, *hostOffset, 0, *storedByteSize, constantFolder);
+      if (failed(copied))
+        return ReturnPathLoweringResult::Failure;
       return ReturnPathLoweringResult::Handled;
     }
 
@@ -531,14 +547,18 @@ raptor::SpatialToPimPass::ReturnPathLoweringResult raptor::SpatialToPimPass::low
       rewriter.setInsertionPointAfter(elementSlice);
 
       int64_t destinationFlatOffset = computeFlatElementIndex(destinationIndices, outputType.getShape());
-      outputTensor = emitHostCopy(rewriter,
-                                  loc,
-                                  outputTensor,
-                                  elementSlice.getResult(),
-                                  static_cast<int32_t>(destinationFlatOffset * elementSize),
-                                  0,
-                                  static_cast<int32_t>(elementSize),
-                                  constantFolder);
+      auto hostOffset =
+        getCheckedByteOffset(destinationFlatOffset, elementSize, producerOp, "concat helper return-path host offset");
+      if (failed(hostOffset))
+        return ReturnPathLoweringResult::Failure;
+      auto elementByteSize = pim::checkedCast<uint64_t>(elementSize, producerOp, "return-path scalar copy byte size");
+      if (failed(elementByteSize))
+        return ReturnPathLoweringResult::Failure;
+      auto copied = emitHostCopy(
+        rewriter, loc, outputTensor, elementSlice.getResult(), *hostOffset, 0, *elementByteSize, constantFolder);
+      if (failed(copied))
+        return ReturnPathLoweringResult::Failure;
+      outputTensor = *copied;
     }
     return ReturnPathLoweringResult::Handled;
   }

src/PIM/Conversion/SpatialToPim/SpatialToPimPass.cpp

@@ -25,8 +25,9 @@
 #include <cassert>
 #include <utility>
 
-#include "Common/PimCommon.hpp"
 #include "Common/IR/ConstantUtils.hpp"
+#include "Common/PimCommon.hpp"
+#include "Common/Support/CheckedArithmetic.hpp"
 #include "Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "Conversion/SpatialToPim/Common.hpp"
 #include "Conversion/SpatialToPim/Patterns.hpp"
@@ -75,21 +76,28 @@ static memref::GlobalOp getOrCreateZeroGlobal(IRRewriter& rewriter, Location loc
                                   IntegerAttr {});
 }
 
-static Value createZeroedDeviceHVector(IRRewriter& rewriter,
-                                       Location loc,
-                                       RankedTensorType tensorType,
-                                       OperationFolder& constantFolder) {
+static FailureOr<Value> createZeroedDeviceHVector(IRRewriter& rewriter,
+                                                  Location loc,
+                                                  RankedTensorType tensorType,
+                                                  OperationFolder& constantFolder) {
   auto outputBuffer = createEmptyTensorFromShaped(rewriter, loc, tensorType);
   auto zeroGlobal = getOrCreateZeroGlobal(rewriter, loc, tensorType);
   auto zeroValue = memref::GetGlobalOp::create(rewriter, loc, zeroGlobal.getType(), zeroGlobal.getName());
   auto zeroIndex = getOrCreateIndexConstant(constantFolder, outputBuffer.getOperation(), 0);
-  auto sizeAttr = rewriter.getI32IntegerAttr(static_cast<int32_t>(getShapedTypeSizeInBytes(tensorType)));
+  auto byteSize =
+    pim::getCheckedShapedTypeSizeInBytes(tensorType, outputBuffer.getOperation(), "host-to-device zero copy byte size");
+  if (failed(byteSize))
+    return failure();
+  auto sizeAttr =
+    pim::getCheckedI32Attr(rewriter, outputBuffer.getOperation(), *byteSize, "host-to-device zero copy byte size");
+  if (failed(sizeAttr))
+    return failure();
   return PimMemCopyHostToDevOp::create(
-           rewriter, loc, tensorType, zeroIndex, zeroIndex, outputBuffer, zeroValue, sizeAttr)
+           rewriter, loc, tensorType, zeroIndex, zeroIndex, outputBuffer, zeroValue, *sizeAttr)
     .getOutput();
 }
 
-static Value
+static FailureOr<Value>
 padHVectorInputToCrossbarSize(IRRewriter& rewriter, Location loc, Value vector, OperationFolder& constantFolder) {
   auto vectorType = cast<RankedTensorType>(vector.getType());
   ArrayRef<int64_t> shape = vectorType.getShape();
@@ -101,10 +109,18 @@ padHVectorInputToCrossbarSize(IRRewriter& rewriter, Location loc, Value vector,
 
   auto paddedType = RankedTensorType::get(
     {shape[0], static_cast<int64_t>(crossbarSize)}, vectorType.getElementType(), vectorType.getEncoding());
-  Value zeroed = createZeroedDeviceHVector(rewriter, loc, paddedType, constantFolder);
-  Value zeroIndex = getOrCreateIndexConstant(constantFolder, zeroed.getDefiningOp(), 0);
-  auto sizeAttr = rewriter.getI32IntegerAttr(static_cast<int32_t>(getShapedTypeSizeInBytes(vectorType)));
-  return PimMemCopyOp::create(rewriter, loc, paddedType, zeroIndex, zeroIndex, zeroed, vector, sizeAttr).getOutput();
+  auto zeroed = createZeroedDeviceHVector(rewriter, loc, paddedType, constantFolder);
+  if (failed(zeroed))
+    return failure();
+  Value zeroIndex = getOrCreateIndexConstant(constantFolder, zeroed->getDefiningOp(), 0);
+  auto byteSize =
+    pim::getCheckedShapedTypeSizeInBytes(vectorType, zeroed->getDefiningOp(), "device padding copy byte size");
+  if (failed(byteSize))
+    return failure();
+  auto sizeAttr = pim::getCheckedI32Attr(rewriter, zeroed->getDefiningOp(), *byteSize, "device padding copy byte size");
+  if (failed(sizeAttr))
+    return failure();
+  return PimMemCopyOp::create(rewriter, loc, paddedType, zeroIndex, zeroIndex, *zeroed, vector, *sizeAttr).getOutput();
 }
 
 void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
@@ -234,7 +250,11 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
     }
   }
 
-  enlargeVMMOutTensorsToCrossbarSize(funcOp, rewriter);
+  if (failed(enlargeVMMOutTensorsToCrossbarSize(funcOp, rewriter))) {
+    funcOp.emitOpError("failed to enlarge VMM output tensors to crossbar size");
+    signalPassFailure();
+    return;
+  }
   replaceReturnWithOutputBuffers(returnOp, rewriter);
   eraseOpsToRemove();
 
@@ -271,8 +291,9 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
   dumpModule(moduleOp, "pim0");
 }
 
-void raptor::SpatialToPimPass::enlargeVMMOutTensorsToCrossbarSize(func::FuncOp funcOp, IRRewriter& rewriter) {
+LogicalResult raptor::SpatialToPimPass::enlargeVMMOutTensorsToCrossbarSize(func::FuncOp funcOp, IRRewriter& rewriter) {
   OperationFolder constantFolder(funcOp.getContext());
+  bool hasFailure = false;
   funcOp.walk([&](PimVMMOp vmmOp) {
     auto outputType = cast<RankedTensorType>(vmmOp.getOutput().getType());
     ArrayRef<int64_t> outputShape = outputType.getShape();
@@ -280,19 +301,23 @@ void raptor::SpatialToPimPass::enlargeVMMOutTensorsToCrossbarSize(func::FuncOp f
     assert(outputShape[1] <= static_cast<int64_t>(crossbarSize) && "output width must fit in one crossbar");
 
     rewriter.setInsertionPoint(vmmOp);
-    Value paddedInput = padHVectorInputToCrossbarSize(rewriter, vmmOp.getLoc(), vmmOp.getInput(), constantFolder);
+    auto paddedInput = padHVectorInputToCrossbarSize(rewriter, vmmOp.getLoc(), vmmOp.getInput(), constantFolder);
+    if (failed(paddedInput)) {
+      hasFailure = true;
+      return WalkResult::interrupt();
+    }
     auto paddedOutputType = RankedTensorType::get(
       {outputShape[0], static_cast<int64_t>(crossbarSize)}, outputType.getElementType(), outputType.getEncoding());
     Value paddedOutputBuffer = outputShape[1] == static_cast<int64_t>(crossbarSize)
                                ? vmmOp.getOutputBuffer()
                                : createEmptyTensorFromShaped(rewriter, vmmOp.getLoc(), paddedOutputType).getResult();
-    vmmOp.getInputMutable().assign(paddedInput);
+    vmmOp.getInputMutable().assign(*paddedInput);
     vmmOp.getOutputBufferMutable().assign(paddedOutputBuffer);
 
     vmmOp.getOutput().setType(paddedOutputType);
 
     if (outputShape[1] == static_cast<int64_t>(crossbarSize))
-      return;
+      return WalkResult::advance();
 
     SmallVector<OpFoldResult> offsets = {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
     SmallVector<OpFoldResult> sizes = {rewriter.getIndexAttr(outputShape[0]), rewriter.getIndexAttr(outputShape[1])};
@@ -302,13 +327,16 @@ void raptor::SpatialToPimPass::enlargeVMMOutTensorsToCrossbarSize(func::FuncOp f
       tensor::ExtractSliceOp::create(rewriter, vmmOp.getLoc(), outputType, vmmOp.getOutput(), offsets, sizes, strides);
     SmallPtrSet<Operation*, 2> exceptions = {vmmOp, sliceOp};
     vmmOp.getOutput().replaceAllUsesExcept(sliceOp.getResult(), exceptions);
+    return WalkResult::advance();
   });
+  return success(!hasFailure);
 }
 
 LogicalResult raptor::SpatialToPimPass::allocateAndInitializeCoreLocalVariables(func::FuncOp funcOp,
                                                                                 IRRewriter& rewriter) {
   Location loc = funcOp.getLoc();
   OperationFolder constantFolder(funcOp.getContext());
+  bool hasFailure = false;
 
   auto insertMemCopyHostToDev = [&](Value inputTensor, int64_t elementsOffset) {
     auto tensorType = cast<ShapedType>(inputTensor.getType());
@@ -319,17 +347,28 @@ LogicalResult raptor::SpatialToPimPass::allocateAndInitializeCoreLocalVariables(
     rewriter.setInsertionPointAfter(inputTensor.getDefiningOp());
 
     auto deviceTensor = tensor::EmptyOp::create(rewriter, loc, tensorType.getShape(), elementType);
+    auto offsetBytes = pim::checkedMul(
+      static_cast<size_t>(elementsOffset), elementByteSize, deviceTensor.getOperation(), "host input byte offset");
+    auto byteSize =
+      pim::getCheckedShapedTypeSizeInBytes(tensorType, deviceTensor.getOperation(), "host input copy byte size");
+    auto sizeAttr =
+      succeeded(byteSize)
+        ? pim::getCheckedI32Attr(rewriter, deviceTensor.getOperation(), *byteSize, "host input copy byte size")
+        : FailureOr<IntegerAttr>(failure());
+    if (failed(offsetBytes) || failed(sizeAttr)) {
+      hasFailure = true;
+      return;
+    }
 
     auto memCopyHostToDevOp = PimMemCopyHostToDevOp::create(
       rewriter,
       loc,
       tensorType,
       getOrCreateIndexConstant(constantFolder, deviceTensor.getOperation(), 0),
-      getOrCreateIndexConstant(
-        constantFolder, deviceTensor.getOperation(), static_cast<int64_t>(elementsOffset * elementByteSize)),
+      getOrCreateIndexConstant(constantFolder, deviceTensor.getOperation(), static_cast<int64_t>(*offsetBytes)),
       deviceTensor,
       inputTensor,
-      rewriter.getI32IntegerAttr(static_cast<int32_t>(tensorType.getNumElements() * elementByteSize)));
+      *sizeAttr);
 
     rewriter.replaceAllUsesExcept(inputTensor, memCopyHostToDevOp.getResult(), {memCopyHostToDevOp});
   };
@@ -347,7 +386,7 @@ LogicalResult raptor::SpatialToPimPass::allocateAndInitializeCoreLocalVariables(
       }
     }
 
-  return success();
+  return success(!hasFailure);
 }
 
 void raptor::SpatialToPimPass::markOpToRemove(Operation* op) {

src/PIM/Conversion/SpatialToPim/SpatialToPimPass.hpp

@@ -64,7 +64,7 @@ private:
   void markOpToRemove(mlir::Operation* op);
   void eraseOpsToRemove();
 
-  void enlargeVMMOutTensorsToCrossbarSize(mlir::func::FuncOp funcOp, mlir::IRRewriter& rewriter);
+  mlir::LogicalResult enlargeVMMOutTensorsToCrossbarSize(mlir::func::FuncOp funcOp, mlir::IRRewriter& rewriter);
 };
 
 } // namespace raptor