Merge with fast resnet

2026-07-20 18:01:58 +02:00
parent 5c94e00f43 ab54243fda
commit 961fd613bd
85 changed files with 4475 additions and 4455 deletions
@@ -79,6 +79,57 @@ materializeRowStripToDense(const RowStripPhysicalValue& rowStripValue, Location
  return createRowStripAssemblyBlueprint(rowStripValue.storage, rowStripValue.logicalType, rewriter, loc);
 }

+static FailureOr<Value> lowerDenseBatchBiasAdd(Value input, Value bias, RankedTensorType resultType,
+                                               PatternRewriter& rewriter, Location loc) {
+  auto producer = input.getDefiningOp<spatial::SpatGraphComputeBatch>();
+  auto inputType = dyn_cast<RankedTensorType>(input.getType());
+  auto biasType = dyn_cast<RankedTensorType>(bias.getType());
+  if (!producer || !inputType || !biasType || !inputType.hasStaticShape() || !biasType.hasStaticShape()
+      || !resultType.hasStaticShape() || inputType.getDimSize(0) != producer.getLaneCount()
+      || biasType.getDimSize(0) != producer.getLaneCount() || resultType.getDimSize(0) != producer.getLaneCount())
+    return failure();
+  auto inputFragmentType = spatial::getGraphBatchFragmentType(inputType, producer.getLaneCount());
+  auto outputFragmentType = spatial::getGraphBatchFragmentType(resultType, producer.getLaneCount());
+  if (failed(inputFragmentType) || failed(outputFragmentType) || inputFragmentType->getRank() != biasType.getRank()
+      || inputFragmentType->getDimSize(0) != 1 || inputFragmentType->getShape().drop_front() != biasType.getShape().drop_front()
+      || inputFragmentType->getRank() != outputFragmentType->getRank() + 1)
+    return failure();
+  for (auto [inputDim, outputDim] : llvm::zip(inputFragmentType->getShape().drop_front(), outputFragmentType->getShape()))
+    if (outputDim > inputDim)
+      return failure();
+
+  auto batch = createSpatComputeBatch(rewriter, loc, TypeRange {resultType}, producer.getLaneCount(), {}, ValueRange {input, bias},
+    [&](detail::SpatComputeBatchBodyArgs args) -> LogicalResult {
+      FailureOr<Value> fragment = extractGraphBatchPhysicalFragment(rewriter, loc, args.inputs[0], args.lane, *inputFragmentType);
+      if (failed(fragment))
+        return failure();
+      MixedSliceGeometry biasSlice;
+      for (int64_t dim : inputFragmentType->getShape()) {
+        biasSlice.offsets.push_back(biasSlice.offsets.empty() ? OpFoldResult(args.lane) : rewriter.getIndexAttr(0));
+        biasSlice.sizes.push_back(rewriter.getIndexAttr(dim));
+        biasSlice.strides.push_back(rewriter.getIndexAttr(1));
+      }
+      Value biasFragment = extractMixedSliceOrIdentity(rewriter, loc, args.inputs[1], *inputFragmentType, biasSlice);
+      if (!biasFragment)
+        return failure();
+      Value added = spatial::SpatVAddOp::create(rewriter, loc, *inputFragmentType, *fragment, biasFragment);
+      MixedSliceGeometry outputSlice;
+      outputSlice.offsets.assign(inputFragmentType->getRank(), rewriter.getIndexAttr(0));
+      outputSlice.sizes.push_back(rewriter.getIndexAttr(1));
+      outputSlice.strides.assign(inputFragmentType->getRank(), rewriter.getIndexAttr(1));
+      for (int64_t dim : outputFragmentType->getShape())
+        outputSlice.sizes.push_back(rewriter.getIndexAttr(dim));
+      Value output = extractMixedSliceOrIdentity(rewriter, loc, added, *outputFragmentType, outputSlice);
+      if (!output)
+        return failure();
+      publishGraphBatchPhysicalFragment(rewriter, loc, output, args.outputs.front(), args.lane);
+      return success();
+    });
+  if (failed(batch))
+    return failure();
+  return batch->getResult(0);
+}
+
 struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, OperationPass<ModuleOp>> {
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LowerSpatialPlansPass)

@@ -267,6 +318,13 @@ struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, Operatio
          signalPassFailure();
          return;
        }
+        if (planOp.getInput().getDefiningOp<spatial::SpatGraphComputeBatch>()) {
+          FailureOr<Value> lowered = lowerDenseBatchBiasAdd(planOp.getInput(), *denseBias, resultType, rewriter, planOp.getLoc());
+          if (succeeded(lowered)) {
+            rewriter.replaceOp(planOp, *lowered);
+            continue;
+          }
+        }
        auto computeOp = createSpatCompute<2>(rewriter,
                                              planOp.getLoc(),
                                              planOp.getOutput().getType(),