a new new beginning phase 1

2026-07-08 22:53:53 +02:00
parent 1f4f58de1c
commit 871fcfa832
22 changed files with 2917 additions and 1012 deletions
@@ -60,6 +60,56 @@ struct SpatComputeBatchBodyArgs {
  mlir::ValueRange outputs;
 };

+inline mlir::SmallVector<mlir::Type> getGraphComputeBlockArgTypes(mlir::ValueRange weights, mlir::ValueRange inputs) {
+  mlir::SmallVector<mlir::Type> blockArgTypes;
+  blockArgTypes.reserve(weights.size() + inputs.size());
+  for (mlir::Value weight : weights)
+    blockArgTypes.push_back(weight.getType());
+  for (mlir::Value input : inputs)
+    blockArgTypes.push_back(input.getType());
+  return blockArgTypes;
+}
+
+inline mlir::SmallVector<mlir::Location> getGraphComputeBlockArgLocs(mlir::Location defaultLoc,
+                                                                     mlir::ValueRange weights,
+                                                                     mlir::ValueRange inputs) {
+  mlir::SmallVector<mlir::Location> blockArgLocs;
+  blockArgLocs.reserve(weights.size() + inputs.size());
+  for (mlir::Value weight : weights)
+    blockArgLocs.push_back(weight.getLoc());
+  for (mlir::Value input : inputs)
+    blockArgLocs.push_back(input.getLoc());
+  return blockArgLocs;
+}
+
+inline mlir::SmallVector<mlir::Type> getGraphComputeBatchBlockArgTypes(mlir::OpBuilder& builder,
+                                                                       mlir::TypeRange resultTypes,
+                                                                       mlir::ValueRange weights,
+                                                                       mlir::ValueRange inputs) {
+  mlir::SmallVector<mlir::Type> blockArgTypes {builder.getIndexType()};
+  blockArgTypes.reserve(1 + weights.size() + inputs.size() + resultTypes.size());
+  for (mlir::Value weight : weights)
+    blockArgTypes.push_back(weight.getType());
+  for (mlir::Value input : inputs)
+    blockArgTypes.push_back(input.getType());
+  llvm::append_range(blockArgTypes, resultTypes);
+  return blockArgTypes;
+}
+
+inline mlir::SmallVector<mlir::Location> getGraphComputeBatchBlockArgLocs(mlir::Location defaultLoc,
+                                                                          mlir::TypeRange resultTypes,
+                                                                          mlir::ValueRange weights,
+                                                                          mlir::ValueRange inputs) {
+  mlir::SmallVector<mlir::Location> blockArgLocs {defaultLoc};
+  blockArgLocs.reserve(1 + weights.size() + inputs.size() + resultTypes.size());
+  for (mlir::Value weight : weights)
+    blockArgLocs.push_back(weight.getLoc());
+  for (mlir::Value input : inputs)
+    blockArgLocs.push_back(input.getLoc());
+  blockArgLocs.append(resultTypes.size(), defaultLoc);
+  return blockArgLocs;
+}
+
 } // namespace detail

 template <typename RewriterT>
@@ -87,6 +137,31 @@ inline mlir::Value createSpatConcat(RewriterT& rewriter, mlir::Location loc, int
  return spatial::SpatConcatOp::create(rewriter, loc, outputType, rewriter.getI64IntegerAttr(axis), inputs).getOutput();
 }

+template <typename RewriterT>
+spatial::SpatGraphCompute createEmptySpatGraphCompute(RewriterT& rewriter,
+                                                      mlir::Location loc,
+                                                      mlir::TypeRange resultTypes,
+                                                      mlir::ValueRange weights,
+                                                      mlir::ValueRange inputs,
+                                                      mlir::TypeRange blockArgTypes,
+                                                      llvm::ArrayRef<mlir::Location> blockArgLocs) {
+  auto computeOp = spatial::SpatGraphCompute::create(rewriter, loc, resultTypes, weights, inputs);
+  rewriter.createBlock(&computeOp.getBody(), computeOp.getBody().end(), blockArgTypes, blockArgLocs);
+  rewriter.setInsertionPointToStart(&computeOp.getBody().front());
+  return computeOp;
+}
+
+template <typename RewriterT>
+spatial::SpatGraphCompute createEmptySpatGraphCompute(RewriterT& rewriter,
+                                                      mlir::Location loc,
+                                                      mlir::TypeRange resultTypes,
+                                                      mlir::ValueRange weights,
+                                                      mlir::ValueRange inputs) {
+  auto blockArgTypes = detail::getGraphComputeBlockArgTypes(weights, inputs);
+  auto blockArgLocs = detail::getGraphComputeBlockArgLocs(loc, weights, inputs);
+  return createEmptySpatGraphCompute(rewriter, loc, resultTypes, weights, inputs, blockArgTypes, blockArgLocs);
+}
+
 /// Builds a `spat.graph_compute` with a fixed number of SSA inputs and erases it if
 /// the body callback reports failure.
 template <size_t NumInputs, typename RewriterT, typename BodyFn>
@@ -97,16 +172,8 @@ auto createSpatGraphCompute(RewriterT& rewriter,
                            mlir::ValueRange inputs,
                            BodyFn&& body) {
  assert(inputs.size() == NumInputs && "NumInputs must match the number of input values");
-  auto computeOp = spatial::SpatGraphCompute::create(rewriter, loc, resultTypes, weights, inputs);
-
-  auto* block = new mlir::Block();
-  for (mlir::Value weight : weights)
-    block->addArgument(weight.getType(), loc);
-  for (mlir::Value input : inputs)
-    block->addArgument(input.getType(), loc);
-
-  computeOp.getBody().push_back(block);
-  rewriter.setInsertionPointToStart(block);
+  auto computeOp = createEmptySpatGraphCompute(rewriter, loc, resultTypes, weights, inputs);
+  auto* block = &computeOp.getBody().front();

  using BodyResult = detail::InvokeWithBlockArgsResultT<std::decay_t<BodyFn>, std::make_index_sequence<NumInputs>>;
  if constexpr (std::is_same_v<BodyResult, void>) {
@@ -140,16 +207,8 @@ auto createSpatGraphCompute(RewriterT& rewriter,
                            mlir::ValueRange weights,
                            mlir::ValueRange inputs,
                            BodyFn&& body) {
-  auto computeOp = spatial::SpatGraphCompute::create(rewriter, loc, resultTypes, weights, inputs);
-
-  auto* block = new mlir::Block();
-  for (mlir::Value weight : weights)
-    block->addArgument(weight.getType(), loc);
-  for (mlir::Value input : inputs)
-    block->addArgument(input.getType(), loc);
-
-  computeOp.getBody().push_back(block);
-  rewriter.setInsertionPointToStart(block);
+  auto computeOp = createEmptySpatGraphCompute(rewriter, loc, resultTypes, weights, inputs);
+  auto* block = &computeOp.getBody().front();

  using BodyResult = detail::InvokeWithValueRangeResultT<std::decay_t<BodyFn>>;
  if constexpr (std::is_same_v<BodyResult, void>) {
@@ -170,14 +229,15 @@ auto createSpatGraphCompute(RewriterT& rewriter,
  }
 }

-template <typename RewriterT, typename BodyFn>
-auto createSpatGraphComputeBatch(RewriterT& rewriter,
-                                 mlir::Location loc,
-                                 mlir::TypeRange resultTypes,
-                                 int64_t laneCount,
-                                 mlir::ValueRange weights,
-                                 mlir::ValueRange inputs,
-                                 BodyFn&& body) {
+template <typename RewriterT>
+auto createEmptySpatGraphComputeBatch(RewriterT& rewriter,
+                                      mlir::Location loc,
+                                      mlir::TypeRange resultTypes,
+                                      int64_t laneCount,
+                                      mlir::ValueRange weights,
+                                      mlir::ValueRange inputs,
+                                      mlir::TypeRange blockArgTypes,
+                                      llvm::ArrayRef<mlir::Location> blockArgLocs) {
  if (laneCount <= 0 || laneCount > std::numeric_limits<int32_t>::max())
    return mlir::FailureOr<spatial::SpatGraphComputeBatch>(mlir::failure());

@@ -186,27 +246,36 @@ auto createSpatGraphComputeBatch(RewriterT& rewriter,
    return mlir::FailureOr<spatial::SpatGraphComputeBatch>(mlir::failure());

  auto batchOp = spatial::SpatGraphComputeBatch::create(rewriter, loc, resultTypes, *laneCountAttr, weights, inputs);
+  rewriter.createBlock(&batchOp.getBody(), batchOp.getBody().end(), blockArgTypes, blockArgLocs);
+  rewriter.setInsertionPointToStart(&batchOp.getBody().front());
+  return mlir::FailureOr<spatial::SpatGraphComputeBatch>(batchOp);
+}

-  mlir::SmallVector<mlir::Type> blockArgTypes {rewriter.getIndexType()};
-  mlir::SmallVector<mlir::Location> blockArgLocs {loc};
-  blockArgTypes.reserve(1 + weights.size() + inputs.size() + resultTypes.size());
-  blockArgLocs.reserve(1 + weights.size() + inputs.size() + resultTypes.size());
-  for (mlir::Value weight : weights) {
-    blockArgTypes.push_back(weight.getType());
-    blockArgLocs.push_back(weight.getLoc());
-  }
-  for (mlir::Value input : inputs) {
-    blockArgTypes.push_back(input.getType());
-    blockArgLocs.push_back(input.getLoc());
-  }
-  for (mlir::Type resultType : resultTypes) {
-    blockArgTypes.push_back(resultType);
-    blockArgLocs.push_back(loc);
-  }
+template <typename RewriterT>
+auto createEmptySpatGraphComputeBatch(RewriterT& rewriter,
+                                      mlir::Location loc,
+                                      mlir::TypeRange resultTypes,
+                                      int64_t laneCount,
+                                      mlir::ValueRange weights,
+                                      mlir::ValueRange inputs) {
+  auto blockArgTypes = detail::getGraphComputeBatchBlockArgTypes(rewriter, resultTypes, weights, inputs);
+  auto blockArgLocs = detail::getGraphComputeBatchBlockArgLocs(loc, resultTypes, weights, inputs);
+  return createEmptySpatGraphComputeBatch(
+    rewriter, loc, resultTypes, laneCount, weights, inputs, blockArgTypes, blockArgLocs);
+}

-  auto* block =
-    rewriter.createBlock(&batchOp.getBody(), batchOp.getBody().end(), mlir::TypeRange(blockArgTypes), blockArgLocs);
-  rewriter.setInsertionPointToStart(block);
+template <typename RewriterT, typename BodyFn>
+auto createSpatGraphComputeBatch(RewriterT& rewriter,
+                                 mlir::Location loc,
+                                 mlir::TypeRange resultTypes,
+                                 int64_t laneCount,
+                                 mlir::ValueRange weights,
+                                 mlir::ValueRange inputs,
+                                 BodyFn&& body) {
+  auto batchOp = createEmptySpatGraphComputeBatch(rewriter, loc, resultTypes, laneCount, weights, inputs);
+  if (failed(batchOp))
+    return mlir::FailureOr<spatial::SpatGraphComputeBatch>(mlir::failure());
+  auto* block = &(*batchOp).getBody().front();

  detail::SpatComputeBatchBodyArgs args {
    block->getArgument(0),
@@ -217,18 +286,18 @@ auto createSpatGraphComputeBatch(RewriterT& rewriter,
  using BodyResult = std::invoke_result_t<BodyFn, detail::SpatComputeBatchBodyArgs>;
  if constexpr (std::is_same_v<BodyResult, void>) {
    std::forward<BodyFn>(body)(args);
-    rewriter.setInsertionPointAfter(batchOp);
-    return mlir::FailureOr<spatial::SpatGraphComputeBatch>(batchOp);
+    rewriter.setInsertionPointAfter(*batchOp);
+    return batchOp;
  }
  else {
    auto bodyResult = std::forward<BodyFn>(body)(args);
    if (mlir::failed(bodyResult)) {
-      rewriter.setInsertionPointAfter(batchOp);
-      rewriter.eraseOp(batchOp);
+      rewriter.setInsertionPointAfter(*batchOp);
+      rewriter.eraseOp(*batchOp);
      return mlir::FailureOr<spatial::SpatGraphComputeBatch>(mlir::failure());
    }
-    rewriter.setInsertionPointAfter(batchOp);
-    return mlir::FailureOr<spatial::SpatGraphComputeBatch>(batchOp);
+    rewriter.setInsertionPointAfter(*batchOp);
+    return batchOp;
  }
 }