Add ktdf.via and the mlir::ktdf::PipelineBuilder.

include/dataflow-scheduler/Dialect/KTDF/Analysis/StageDependency.h

@@ -0,0 +1,68 @@
+//===-- StageDependency.h ---------------------------------------*- c++ -*-===//
+//
+// Part of the Dataflow Scheduler MLIR Dialects project.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef DATAFLOW_SCHEDULER_DIALECT_KTDF_ANALYSIS_STAGEDEPENDENCY_H_
+#define DATAFLOW_SCHEDULER_DIALECT_KTDF_ANALYSIS_STAGEDEPENDENCY_H_
+
+#include <llvm/ADT/DenseMap.h>
+#include <llvm/ADT/SmallPtrSet.h>
+
+#include "dataflow-scheduler/Dialect/KTDF/KTDF.h"
+
+namespace mlir::ktdf {
+
+/// Cache that stores dependency edges between pipeline stages.
+class StageDependency {
+ public:
+  /// Initializes a new StageDependency cache.
+  /*implicit*/ StageDependency() = default;
+
+  /// Gets the producers that @p consumer depends on.
+  static void getProducers(StageOp consumer,
+                           SmallPtrSetImpl<StageOp>& producers,
+                           bool transitive = true);
+  /// Gets the consumers that depend on @p producer .
+  static void getConsumers(StageOp producer,
+                           SmallPtrSetImpl<StageOp>& consumers,
+                           bool transitive = true);
+
+  /// Determines whether @p consumer depends on @p producer .
+  auto contains(StageOp consumer, StageOp producer, bool transitive = false)
+      -> bool;
+
+  /// Updates the cache that @p consumer depends on @p producer .
+  ///
+  /// @warning  This does not modify the IR.
+  auto insert(StageOp consumer, StageOp producer) -> bool;
+
+  /// Erases @p consumer from the cache.
+  ///
+  /// @return   Whether the dependencies changed as a result.
+  auto erase(StageOp consumer) -> bool;
+  /// Erases the dependency of @p consumer on @p producer from the cache.
+  ///
+  /// @return   Whether the dependencies changed as a result.
+  auto erase(StageOp consumer, StageOp producer) -> bool;
+
+ private:
+  DenseMap<StageOp, SmallPtrSet<StageOp, 4>> cache_;
+};
+
+}  // namespace mlir::ktdf
+
+#endif  // DATAFLOW_SCHEDULER_DIALECT_KTDF_ANALYSIS_STAGEDEPENDENCY_H_

include/dataflow-scheduler/Dialect/KTDF/KTDF.h

@@ -23,11 +23,13 @@
 #ifndef DATAFLOW_SCHEDULER_DIALECT_KTDF_KTDF_H_
 #define DATAFLOW_SCHEDULER_DIALECT_KTDF_KTDF_H_
 
+#include <llvm/Support/PointerLikeTypeTraits.h>
 #include <mlir/Dialect/Affine/IR/AffineMemoryOpInterfaces.h>
 #include <mlir/Dialect/Utils/StaticValueUtils.h>
 #include <mlir/IR/OpDefinition.h>
 #include <mlir/Interfaces/ControlFlowInterfaces.h>
 #include <mlir/Interfaces/DestinationStyleOpInterface.h>
+#include <mlir/Interfaces/InferTypeOpInterface.h>
 #include <mlir/Interfaces/LoopLikeInterface.h>
 #include <mlir/Interfaces/SideEffectInterfaces.h>
 
@@ -46,32 +48,55 @@ struct FifoResource : public mlir::SideEffects::Resource::Base<FifoResource> {
 #define GET_OP_CLASSES
 #include "dataflow-scheduler/Dialect/KTDF/KTDF.h.inc"
 
+template <>
+struct llvm::PointerLikeTypeTraits<mlir::ktdf::StageOp>
+    : PointerLikeTypeTraits<mlir::Operation*> {
+  [[nodiscard]] static auto getFromVoidPointer(void* ptr)
+      -> mlir::ktdf::StageOp {
+    return mlir::ktdf::StageOp::getFromOpaquePointer(ptr);
+  }
+};
+
 namespace mlir::ktdf {
 
 /// RAII helper that allows moving code to a PrivateOp in a PipelineOp.
 ///
 /// Users may call `makePrivate` on operations to attempt making them private
 /// to the pipeline. The move is deferred until the helper is destroyed. The
 /// helper allows for both hoisting and sinking of ops into the PrivateOp.
-class PipelinePrivatizer {
+class PipelinePrivatizer : protected RewriterBase {
  public:
+  using Token = TypedValue<TokenType>;
+  using FifoSlot = TypedValue<FifoSlotType>;
+
   /// Canonicalizes the PrivateOp of @p op .
   ///
   /// - Results without users are dropped.
   /// - Values yielded multiple times are coalesced into one result.
   /// - External yielded values replace their results.
   /// - If the resulting PrivateOp is empty, it is erased.
-  static void canonicalize(RewriterBase& rewriter, PipelineOp op) {
+  static void canonicalize(PipelineOp op,
+                           OpBuilder::Listener* listener = nullptr) {
+    PipelinePrivatizer(op, true, listener);
+  }
+  /// @copydoc canonicalize(PipelineOp, OpBuilder::Listener*)
+  static void canonicalize(const RewriterBase& rewriter, PipelineOp op) {
     PipelinePrivatizer(rewriter, op, true);
   }
 
   /// Creates a PipelinePrivatizer for @p pipeline .
   ///
   /// If @p force_recreate is `true`, any existing PrivateOp will be re-created
   /// in its canonical form, even if no modifications are made.
-  explicit PipelinePrivatizer(RewriterBase& rewriter, PipelineOp pipeline,
-                              bool force_recreate = false);
-  ~PipelinePrivatizer();
+  explicit PipelinePrivatizer(PipelineOp pipeline, bool force_recreate = false,
+                              OpBuilder::Listener* listener = nullptr);
+  /// @copydoc PipelinePrivatizer(PipelineOp, bool, OpBuilder::Listener*)
+  explicit PipelinePrivatizer(const OpBuilder& builder, PipelineOp pipeline,
+                              bool force_recreate = false)
+      : PipelinePrivatizer(pipeline, force_recreate, builder.getListener()) {}
+
+  /// Finalizes the outstanding modifications to the pipeline.
+  ~PipelinePrivatizer() override { finalize(); }
 
   PipelinePrivatizer(PipelinePrivatizer&&) = delete;
   PipelinePrivatizer(const PipelinePrivatizer&) = delete;
@@ -80,12 +105,19 @@ class PipelinePrivatizer {
   auto operator=(const PipelinePrivatizer&) = delete;
 
   /// Determines whether @p block will be within the PrivateOp.
-  [[nodiscard]] auto isPrivate(Block* block) -> bool;
+  [[nodiscard]] auto isPrivate(Block* block) const -> bool;
   /// Determines whether @p op will be within the PrivateOp.
-  [[nodiscard]] auto isPrivate(Operation* op) -> bool {
+  [[nodiscard]] auto isPrivate(Operation* op) const -> bool {
     return isPrivate(op->getBlock());
   }
 
+  /// Gets the MLIRContext.
+  [[nodiscard]] auto getContext() const -> MLIRContext* {
+    return OpBuilder::getContext();
+  }
+  /// Gets the underlying pipeline.
+  [[nodiscard]] auto getPipeline() const -> PipelineOp { return pipeline_; }
+
   /// Attempts to make @p op a private result.
   ///
   /// Privating fails if the SSA property would be broken by moving @p op :
@@ -97,12 +129,27 @@ class PipelinePrivatizer {
   /// @return Whether @p op was privated.
   auto makePrivate(Operation* op) -> LogicalResult;
 
+  /// Creates a new token inside the private region.
+  [[nodiscard]] auto createToken(std::optional<Location> loc = std::nullopt)
+      -> Token;
+  /// Creates a new FIFO inside the private region.
+  [[nodiscard]] auto createFifo(ArrayRef<FifoSlotType> slots,
+                                ValueRange dynamic_sizes = {},
+                                std::optional<Location> loc = std::nullopt)
+      -> ValueRange;
+
+  /// Finalizes the outstanding modifications to the pipeline.
+  ///
+  /// If there are no modifications to perform, does nothing. After finalizing,
+  /// the PipelinePrivatizer will be ready again to queue more modifications to
+  /// the same pipeline.
+  virtual auto finalize() -> PipelineOp;
+
  private:
-  RewriterBase& rewriter_;
   PipelineOp pipeline_;
   PrivateOp existing_;
 
-  /// Unlinked accumulator for operations to be privated on destruction.
+  /// Unlinked accumulator for operations to be privated on finalize().
   Block private_;
 };
 

include/dataflow-scheduler/Dialect/KTDF/KTDF.td

@@ -26,6 +26,7 @@
 include "mlir/Dialect/Affine/IR/AffineMemoryOpInterfaces.td"
 include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/Interfaces/DestinationStyleOpInterface.td"
+include "mlir/Interfaces/InferTypeOpInterface.td"
 include "mlir/Interfaces/LoopLikeInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 include "mlir/IR/OpAsmInterface.td"
@@ -561,6 +562,24 @@ def KTDF_StageOp : KTDF_Op<"stage", [
         "ValueRange":$depends_out, 
         CArg<"function_ref<void(OpBuilder&, Location)>", "{}">:$body_builder)>
   ];
+
+  let extraClassDeclaration = [{
+    /// Determines whether this stage has an in dependency on @p operand .
+    [[nodiscard]] auto isInDependency(OpOperand& operand) -> bool;
+
+    /// Modifies the op in-place to add an input dependency token.
+    ///
+    /// @return Whether a new operand was added.
+    auto addInDependency(TypedValue<TokenType> token) -> bool;
+
+    /// Determines whether this stage has an out dependency on @p operand .
+    [[nodiscard]] auto isOutDependency(OpOperand& operand) -> bool;
+
+    /// Modifies the op in-place to add an output dependency token.
+    ///
+    /// @return Whether a new operand was added.
+    auto addOutDependency(TypedValue<TokenType> token) -> bool;
+  }];
 }
 
 //===----------------------------------------------------------------------===//
@@ -664,7 +683,12 @@ def KTDF_DataTransferOp : KTDF_Op<"data_transfer", [
       "ArrayRef<int64_t>":$source_sizes,
       "Value":$destination, "AffineMap":$dest_map,
       "ValueRange":$dest_indices,
-      "ArrayRef<int64_t>":$dest_sizes)>
+      "ArrayRef<int64_t>":$dest_sizes)>,
+    OpBuilder<(ins "Value":$source, "Value":$destination), [{
+      build($_builder, $_state, 
+        source, AffineMap{}, ValueRange{}, ArrayRef<int64_t>{},
+        destination, AffineMap{}, ValueRange{}, ArrayRef<int64_t>{});
+    }]>
   ];
 
   let extraClassDeclaration = [{
@@ -779,6 +803,46 @@ def KTDF_DataTransferOp : KTDF_Op<"data_transfer", [
   }];
 }
 
+def KTDF_ViaOp : KTDF_Op<"via", [Pure, SameOperandsAndResultType]> {
+  let summary = "Places a routing constraint on a value";
+  let description = [{
+    The `ktdf.via` operation represents an implicit data movement constraint in
+    the dataflow pipeline. It indicates that the value has to pass through the
+    specified unit or memory.
+
+    Adjacent hops via the same unit are canonicalized away, but there is no
+    architecture aware canonicalization. These hops have to name legal endpoints
+    in the corresponding architecture, between which a legal path must exist, or
+    the program is ill-formed.
+
+    Example:
+    ```mlir
+    %via = ktdf.via["L1"] %tile : tensor<4x64xf16>
+    ```
+  }];
+
+  let arguments = (ins 
+    AnyTypeOf<[AnyRankedTensor, AnyMemRef], "", "::mlir::ShapedType">:$operand, 
+    ArrayAttr:$hops);
+  let results = (outs AnyShaped:$result);
+
+  let assemblyFormat = "``$hops $operand attr-dict `:` type($operand)";
+
+  let builders = [
+    OpBuilder<(ins "Value":$operand, "Attribute":$hop), [{
+      build($_builder, $_state, operand, $_builder.getArrayAttr({hop}));
+    }]>
+  ];
+
+  let hasFolder = 1;
+
+  let extraClassDeclaration = [{
+    /// Follows this chain of `via` operations until a source is found.
+    [[nodiscard]] auto collectHops(SmallVectorImpl<Attribute>& hops) 
+      -> TypedValue<ShapedType>;
+  }];
+}
+
 //===----------------------------------------------------------------------===//
 // Indirect Data Transfer Operation
 //===----------------------------------------------------------------------===//