DeadLock

src/PIM/Conversion/SpatialToPim/SpatialToPimPass.cpp

@@ -25,9 +25,11 @@
 #include <cassert>
 #include <utility>
 
+#include "Common/IR/ShapeUtils.hpp"
 #include "Common/IR/ConstantUtils.hpp"
 #include "Common/PimCommon.hpp"
 #include "Common/Support/CheckedArithmetic.hpp"
+#include "Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp"
 #include "Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "Conversion/SpatialToPim/Common.hpp"
 #include "Conversion/SpatialToPim/Patterns.hpp"
@@ -97,6 +99,64 @@ static FailureOr<Value> createZeroedDeviceHVector(IRRewriter& rewriter,
     .getOutput();
 }
 
+static bool isHostBackedMemRefValue(Value value) {
+  while (Operation* definingOp = value.getDefiningOp()) {
+    if (auto subviewOp = dyn_cast<memref::SubViewOp>(definingOp)) {
+      value = subviewOp.getSource();
+      continue;
+    }
+    if (auto castOp = dyn_cast<memref::CastOp>(definingOp)) {
+      value = castOp.getSource();
+      continue;
+    }
+    if (auto collapseOp = dyn_cast<memref::CollapseShapeOp>(definingOp)) {
+      value = collapseOp.getSrc();
+      continue;
+    }
+    if (auto expandOp = dyn_cast<memref::ExpandShapeOp>(definingOp)) {
+      value = expandOp.getSrc();
+      continue;
+    }
+    return isa<memref::GetGlobalOp>(definingOp);
+  }
+  return false;
+}
+
+static bool isHostBackedTensorValue(Value value) {
+  while (Operation* definingOp = value.getDefiningOp()) {
+    if (auto extractSliceOp = dyn_cast<tensor::ExtractSliceOp>(definingOp)) {
+      auto sourceType = dyn_cast<RankedTensorType>(extractSliceOp.getSource().getType());
+      auto resultType = dyn_cast<RankedTensorType>(extractSliceOp.getResult().getType());
+      if (!sourceType || !resultType || !sourceType.hasStaticShape() || !resultType.hasStaticShape())
+        return false;
+      if (!onnx_mlir::isContiguousSubviewWithDynamicOffsets(sourceType.getShape(),
+                                                            extractSliceOp.getMixedOffsets(),
+                                                            extractSliceOp.getStaticSizes(),
+                                                            extractSliceOp.getStaticStrides())) {
+        return false;
+      }
+      value = extractSliceOp.getSource();
+      continue;
+    }
+    if (auto collapseOp = dyn_cast<tensor::CollapseShapeOp>(definingOp)) {
+      value = collapseOp.getSrc();
+      continue;
+    }
+    if (auto expandOp = dyn_cast<tensor::ExpandShapeOp>(definingOp)) {
+      value = expandOp.getSrc();
+      continue;
+    }
+    if (auto castOp = dyn_cast<tensor::CastOp>(definingOp)) {
+      value = castOp.getSource();
+      continue;
+    }
+    if (auto toTensorOp = dyn_cast<bufferization::ToTensorOp>(definingOp))
+      return isHostBackedMemRefValue(toTensorOp.getBuffer());
+    return false;
+  }
+  return false;
+}
+
 static FailureOr<Value>
 padHVectorInputToCrossbarSize(IRRewriter& rewriter, Location loc, Value vector, OperationFolder& constantFolder) {
   auto vectorType = cast<RankedTensorType>(vector.getType());
@@ -120,6 +180,10 @@ padHVectorInputToCrossbarSize(IRRewriter& rewriter, Location loc, Value vector,
   auto sizeAttr = pim::getCheckedI32Attr(rewriter, zeroed->getDefiningOp(), *byteSize, "device padding copy byte size");
   if (failed(sizeAttr))
     return failure();
+  if (isHostBackedTensorValue(vector)) {
+    return PimMemCopyHostToDevOp::create(rewriter, loc, paddedType, zeroIndex, zeroIndex, *zeroed, vector, *sizeAttr)
+      .getOutput();
+  }
   return PimMemCopyOp::create(rewriter, loc, paddedType, zeroIndex, zeroIndex, *zeroed, vector, *sizeAttr).getOutput();
 }
 
@@ -137,6 +201,12 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
     return;
   }
   func::FuncOp funcOp = *entryFunc;
+  if (failed(verifyScheduledSpatialInvariants(funcOp))) {
+    funcOp.emitOpError(
+      "RAPTOR_PHASE_CHECK scheduled Spatial verification failed at the start of SpatialToPim");
+    signalPassFailure();
+    return;
+  }
 
   IRRewriter rewriter(&getContext());
   OperationFolder constantFolder(&getContext());
@@ -176,19 +246,19 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
     return;
   }
 
-  for (auto computeOp : funcOp.getOps<spatial::SpatCompute>()) {
+  for (auto computeOp : funcOp.getOps<spatial::SpatScheduledCompute>()) {
     markOpToRemove(computeOp);
     if (failed(lowerComputeOp(computeOp, rewriter, constantFolder))) {
-      computeOp.emitOpError("failed to lower spat.compute to pim.core");
+      computeOp.emitOpError("failed to lower spat.scheduled_compute to pim.core");
       signalPassFailure();
       return;
     }
   }
 
-  for (auto computeBatchOp : funcOp.getOps<spatial::SpatComputeBatch>()) {
+  for (auto computeBatchOp : funcOp.getOps<spatial::SpatScheduledComputeBatch>()) {
     markOpToRemove(computeBatchOp);
     if (failed(lowerComputeBatchOp(computeBatchOp, rewriter))) {
-      computeBatchOp.emitOpError("failed to lower spat.compute_batch to pim.core_batch");
+      computeBatchOp.emitOpError("failed to lower spat.scheduled_compute_batch to pim.core_batch");
       signalPassFailure();
       return;
     }
@@ -374,7 +444,7 @@ LogicalResult raptor::SpatialToPimPass::allocateAndInitializeCoreLocalVariables(
   };
 
   for (auto& op : funcOp.getBody().getOps())
-    if (auto computeOp = dyn_cast<spatial::SpatCompute>(op)) {
+    if (auto computeOp = dyn_cast<spatial::SpatScheduledCompute>(op)) {
       if (!computeOp.getInputs().empty() || computeOp.getBody().front().getNumArguments() != 0)
         continue;
       for (auto getGlobal : computeOp.getOps<memref::GetGlobalOp>()) {