Add support for instanced subgraphs in ktdf_arch.

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

@@ -24,8 +24,10 @@
 #define DATAFLOW_SCHEDULER_DIALECT_KTDFARCH_KTDFARCH_H_
 
 #include <mlir/IR/OpDefinition.h>
+#include <mlir/IR/RegionKindInterface.h>
 #include <mlir/IR/SymbolTable.h>
 #include <mlir/Interfaces/ControlFlowInterfaces.h>
+#include <mlir/Interfaces/InferTypeOpInterface.h>
 #include <mlir/Interfaces/SideEffectInterfaces.h>
 
 #include "dataflow-scheduler/Dialect/KTDFArch/KTDFArchAttributes.h"  // IWYU pragma: export

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

@@ -26,8 +26,10 @@
 
 include "mlir/IR/OpAsmInterface.td"
 include "mlir/IR/OpBase.td"
+include "mlir/IR/RegionKindInterface.td"
 include "mlir/IR/SymbolInterfaces.td"
 include "mlir/Interfaces/ControlFlowInterfaces.td"
+include "mlir/Interfaces/InferTypeOpInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 
 include "dataflow-scheduler/Dialect/KTDFArch/KTDFArchTypes.td"
@@ -50,7 +52,11 @@ def KTDFArch_IsSubgraph : NativeOpTrait<"IsSubgraph">, StructuralOpTrait {
   }];
 }
 
-def KTDFArch_Subgraph : TraitList<[OpAsmOpInterface, KTDFArch_IsSubgraph]>;
+def KTDFArch_Subgraph : TraitList<[
+  OpAsmOpInterface, 
+  KTDFArch_IsSubgraph,
+  RegionKindInterface,
+  HasOnlyGraphRegion]>;
 
 def KTDFArch_IsMeta : NativeOpTrait<"IsMeta"> {
   let cppNamespace = "::mlir::ktdf_arch";
@@ -321,9 +327,9 @@ def KTDFArch_YieldOp : KTDFArch_Op<"yield", [
   Pure, 
   ReturnLike, 
   Terminator, 
-  HasParent<"GroupOp">
+  ParentOneOf<["GroupOp", "NeighborhoodOp"]>
 ]> {
-  let summary = "Yields execution unit resources from a group.";
+  let summary = "Yields resources from a subgraph.";
   let description = [{
     Example:
     ```
@@ -341,7 +347,7 @@ def KTDFArch_YieldOp : KTDFArch_Op<"yield", [
     ```
   }];
 
-  let arguments = (ins Variadic<KTDFArch_ExecutionUnitType>:$operands);
+  let arguments = (ins Variadic<KTDFArch_EndpointType>:$operands);
 
   let skipDefaultBuilders = 1;
   let builders = [
@@ -350,7 +356,7 @@ def KTDFArch_YieldOp : KTDFArch_Op<"yield", [
     }]>
   ];
 
-  let assemblyFormat = "$operands attr-dict";
+  let hasCustomAssemblyFormat = 1;
 }
 
 //===----------------------------------------------------------------------===//
@@ -667,4 +673,164 @@ def KTDFArch_PatternsOp : KTDFArch_Op<"patterns", [
   }];
 }
 
-#endif // DATAFLOW_SCHEDULER_DIALECT_KTDFARCH_KTDFARCH_TD
+def KTDFArch_NeighborhoodOp : KTDFArch_Op<"neighborhood", [
+  SingleBlock,
+  KTDFArch_IsMeta,
+  KTDFArch_Subgraph,
+  SingleBlockImplicitTerminator<"YieldOp">,
+  RecursiveMemoryEffects,
+  AlwaysSpeculatable
+]> {
+  let summary = "Declares a tiled repetition.";
+  let description = [{
+    The `ktdf_arch.neighborhood` operation defines a subgraph that is repeated 
+    for a hyperrectangular index domain.
+
+    Each tile in the domain is an instance of the subgraph defined within the
+    region, which is not isolated from above, and which will be unrolled at
+    the same nesting level of IR. Each tile may yield a number of endpoints,
+    which can be queried using the `ktdf_arch.neighbor` op. Regardless of its
+    users, every `ktdf_arch.neighborhood` is unrolled for each point in its
+    domain in lexicographic order.
+
+    Example:
+    ```
+    ktdf_arch.device @device {
+      %ring = neighborhood %self : (port, exec_unit)[3] {
+        %exec = group share() {
+          %local = memory {kind = "local"}
+          %exec = exec_unit
+          yield %exec
+        } -> exec_unit
+
+        %sw:3 = switch[3]
+        datapath %sw#2 to %exec : port, exec_unit
+        datapath %exec to %sw#2 : exec_unit, port
+
+        %prev_out, %prev_exec = neighbor affine_map<(d0) -> (d0 - 1 mod 3)> in %self : (port, exec_unit)[3]
+        datapath %prev_out to %sw#0 : port, port
+
+        yield %sw#1, %exec : port, exec_unit
+      }
+
+      %port_0, %exec_0 = neighbor affine_map<() -> (0)> in %ring : (port, exec_unit)[3]
+      %port_1, %exec_1 = neighbor affine_map<() -> (1)> in %ring : (port, exec_unit)[3]
+    }
+    ```
+  }];
+
+  let results = (outs KTDFArch_NeighborhoodType:$result);
+  let regions = (region SizedRegion<1>:$region);
+
+  let hasVerifier = 1;
+  let hasRegionVerifier = 1;
+  let hasCanonicalizer = 1;
+
+  let skipDefaultBuilders = 1;
+  let builders = [
+    OpBuilder<(ins 
+        CArg<"NeighborhoodType">:$neighborhood_type,
+        CArg<"function_ref<void(OpBuilder&, Location, Value)>", "nullptr">:$body_builder
+    )>
+  ];
+
+  let hasCustomAssemblyFormat = 1;
+
+  let extraClassDeclaration = [{
+    /// Gets the block argument that carries recursive self references.
+    [[nodiscard]] auto getSelfArgument() -> BlockArgument {
+      return getBody()->getArgument(0);
+    }
+    /// Gets the type of the defined neighborhood.
+    [[nodiscard]] auto getNeighborhoodType() -> NeighborhoodType {
+      return getResult().getType();
+    }
+
+    /// Determines whether there are any recursive self references.
+    [[nodiscard]] auto isRecursive() -> bool {
+      return !getSelfArgument().use_empty();
+    }
+    /// Determines whether this is a singleton domain.
+    [[nodiscard]] auto isSingleton() -> bool {
+      return getNeighborhoodType().getDimensions().empty();
+    }
+    /// Determines whether there are no nested neighborhoods below this one.
+    [[nodiscard]] auto isLeaf() -> bool;
+
+    /// Appends the effective domain to @p domain .
+    ///
+    /// The effective domain is the concatenation of all surrounding
+    /// NeighborhoodOp dimensions, including this operation, from outermost
+    /// to innermost.
+    void getDomain(SmallVectorImpl<int64_t>& domain);
+    /// Gets the effective domain.
+    ///
+    /// See getDomain(SmallVectorImpl<int64_t>&).
+    [[nodiscard]] auto getDomain() -> SmallVector<int64_t> {
+      SmallVector<int64_t> result;
+      getDomain(result);
+      return result;
+    }
+  }];
+}
+
+def KTDFArch_NeighborOp : KTDFArch_Op<"neighbor", [
+  KTDFArch_IsMeta,
+  DeclareOpInterfaceMethods<InferTypeOpInterface>,
+  SameTypeOperands,
+  Pure
+]> {
+  let summary = "Indexes into a tiled repetition.";
+  let description = [{
+    The `ktdf_arch.neighbor` operation allows querying the yielded endpoints
+    of a particular tile instance from a (tuple of) `ktdf_arch.neighborhood`s.
+
+    The input dimensions of the indexing map correspond to the concatenation of
+    the domains of the enclosing `ktdf_arch.neighborhood` operations, while the
+    output dimensions correspond to the domain of the targeted neighborhood. The 
+    results always correspond to the yielded operands of the targeted 
+    neighborhood's tile instance.
+
+    When a neighborhood is unrolled into its instances, it turns from one SSA
+    value into many. For this particular use case, this operation supports an
+    additional syntax, where an additional leading result dimension can select
+    from a local grouping of neighborhoods.
+
+    Example:
+    ```
+    ktdf_arch.device @device {
+      %outer = neighborhood %arg0 : ()[1, 2] {
+        %inner = neighborhood %arg1 : ()[3] {
+          neighbor affine_map<(d0, d1, d2) -> (d2)> in %arg1 : ()[3]
+        }
+        neighbor affine_map<(d0, d1) -> (d0 + d1)> in %inner : ()[3]
+      }
+      neighbor affine_map<() -> (0, 0)> in %outer : ()[1, 2]
+      neighbor affine_map<() -> (1, 0, 0)> in neighborhood(%outer, %outer) : ()[2, 1, 2]
+    }
+    ```
+  }];
+
+  let arguments = (ins 
+    Variadic<KTDFArch_NeighborhoodType>:$neighborhood,
+    ConfinedAttr<AffineMapAttr, [HasNoSymbols]>:$map);
+  let results = (outs Variadic<KTDFArch_EndpointType>:$results);
+
+  let hasVerifier = 1;
+  let hasCanonicalizer = 1;
+
+  let assemblyFormat = [{
+    $map attr-dict 
+    `in` custom<LocalNeighborhood>($neighborhood, type($neighborhood))  
+  }];
+
+  let extraClassDeclaration = [{
+    /// Determines whether this neighbor indexes into a 1-tuple.
+    [[nodiscard]] auto isDegenerate() -> bool;
+
+    /// Gets the type of the queried neighborhood.
+    [[nodiscard]] auto getNeighborhoodType() -> NeighborhoodType;
+  }];
+}
+
+#endif // DATAFLOW_SCHEDULER_DIALECT_KTDFARCH_KTDFARCH_TD

include/dataflow-scheduler/Dialect/KTDFArch/KTDFArchAttributes.td

@@ -56,6 +56,11 @@ def KTDFArch_AdjacencyMatrixAttr : Attr<
   let constBuilderCall = "?";
 }
 
+def HasNoSymbols : AttrConstraint<
+  CPred<"::llvm::cast<::mlir::AffineMapAttr>($_self).getValue().getNumSymbols() == 0">, 
+  "without symbols"
+> {}
+
 // Base class for KTDF dialect attributes
 class KTDFArch_Attr<string name, string attrMnemonic, list<Trait> traits = []>
     : AttrDef<KTDFArch_Dialect, name, traits> {

include/dataflow-scheduler/Dialect/KTDFArch/KTDFArchTypes.h

@@ -25,6 +25,12 @@
 
 #include "dataflow-scheduler/Dialect/KTDFArch/KTDFArchAttributes.h"  // IWYU pragma: keep
 
+namespace mlir::ktdf_arch {
+
+struct EndpointType;
+
+}  // namespace mlir::ktdf_arch
+
 /// Auto-generated includes.
 #define GET_TYPEDEF_CLASSES
 #include "dataflow-scheduler/Dialect/KTDFArch/KTDFArchTypes.h.inc"  // IWYU pragma: export
@@ -34,13 +40,42 @@ namespace mlir::ktdf_arch {
 /// Parses a `ktdf_arch` type mnemonic, or falls back to a qualified type.
 ///
 /// Note that builtin type mnemonics can not be parsed with this function.
-auto parseShortType(OpAsmParser& parser, Type& type) -> ParseResult;
+auto parseShortType(AsmParser& parser, Type& type) -> ParseResult;
 
 /// Prints a `ktdf_arch` type mnemonic, or falls back to a qualified type.
 ///
 /// @pre @p type is not a builtin type.
+void printShortType(AsmPrinter& printer, Type type);
+/// @copydoc printShortType(AsmPrinter&, Type)
 void printShortType(OpAsmPrinter& printer, Operation* op, Type type);
 
+//===----------------------------------------------------------------------===//
+// EndpointType
+//===----------------------------------------------------------------------===//
+
+/// Named constraint for the EndpointType in TableGen.
+struct EndpointType : Type {
+  [[nodiscard]] static auto classof(Type type) -> bool {
+    return isa<MemoryType, ExecutionUnitType, PortType>(type);
+  }
+  [[nodiscard]] static auto classof(MemoryType /*type*/) -> bool {
+    return true;
+  }
+  [[nodiscard]] static auto classof(ExecutionUnitType /*type*/) -> bool {
+    return true;
+  }
+  [[nodiscard]] static auto classof(PortType /*type*/) -> bool { return true; }
+
+  using Type::Type;
+
+  /*implicit*/ EndpointType(MemoryType type)
+      : Type(static_cast<ImplType*>(type.getImpl())) {}
+  /*implicit*/ EndpointType(ExecutionUnitType type)
+      : Type(static_cast<ImplType*>(type.getImpl())) {}
+  /*implicit*/ EndpointType(PortType type)
+      : Type(static_cast<ImplType*>(type.getImpl())) {}
+};
+
 }  // namespace mlir::ktdf_arch
 
 #endif  // DATAFLOW_SCHEDULER_DIALECT_KTDFARCH_KTDFARCHTYPES_H_

include/dataflow-scheduler/Dialect/KTDFArch/KTDFArchTypes.td

@@ -44,6 +44,32 @@ def KTDFArch_PortType : KTDFArch_Type<"Port", "port"> {
   let summary = "Type of a switch port resource.";
 }
 
-def KTDFArch_EndpointType : AnyTypeOf<[KTDFArch_MemoryType, KTDFArch_ExecutionUnitType, KTDFArch_PortType]>;
+def KTDFArch_EndpointType : Type<
+    CPred<"::llvm::isa<::mlir::ktdf_arch::EndpointType>($_self)">, 
+    "resource endpoint type",
+    "::mlir::ktdf_arch::EndpointType"> {}
+
+def KTDFArch_NeighborhoodType : KTDFArch_Type<"Neighborhood", "neighborhood"> {
+  let summary = "";
+  let description = [{
+  }];
+
+  let parameters = (ins 
+    ArrayRefParameter<"Type">:$results,
+    ArrayRefParameter<"int64_t">:$dimensions
+  );
+
+  let genVerifyDecl  = 1;
+
+  let assemblyFormat = [{
+    `<` custom<ShortTypeTuple>($results) `[` $dimensions `]` `>`
+  }];
+
+  let extraClassDeclaration = [{
+    [[nodiscard]] auto cloneWith(
+      std::optional<ArrayRef<Type>> results,
+      std::optional<ArrayRef<int64_t>> dimensions) const -> NeighborhoodType;
+  }];
+}
 
 #endif // DATAFLOW_SCHEDULER_DIALECT_KTDFARCH_KTDFARCHTYPES_TD

include/dataflow-scheduler/Dialect/KTDFArch/Transforms/Passes.td

@@ -46,4 +46,19 @@ def ApplyPatternsPass : Pass<"ktdfarch-apply-patterns", "func::FuncOp"> {
   ];
 }
 
+def InstantiateNeighborhoodsPass : Pass<
+  "ktdfarch-instantiate-neighborhoods", 
+  "ktdf_arch::DeviceOp"
+> {
+  let summary = "Instantiates `ktdf_arch.neighborhood` tiles.";
+  let description = [{
+    Given a `ktdf_arch.device`, instantiates all `ktdf_arch.neighborhood` tiles
+    inside out.
+  }];
+
+  let dependentDialects = [
+    "mlir::ktdf_arch::KTDFArchDialect"
+  ];
+}
+
 #endif // DATAFLOW_SCHEDULER_DIALECT_KTDFARCH_TRANSFORMS_PASSES_TD

lib/Dialect/KTDFArch/Analysis/NodeLinks.cpp

@@ -45,7 +45,7 @@ auto visitLinksImpl(Value value,
     // need to ascend from the GroupOp to find more links.
     if (auto yield = dyn_cast<YieldOp>(use.getOwner()); yield) {
       if (!visitLinksImpl(
-              yield.getParentOp()->getResult(use.getOperandNumber()),
+              yield->getParentOp()->getResult(use.getOperandNumber()),
               callback)) {
         return false;
       }

lib/Dialect/KTDFArch/CMakeLists.txt

@@ -23,4 +23,5 @@ add_dataflow_scheduler_dialects_library(DataflowSchedulerDialectsKTDFArchDialect
     MLIRIR
     MLIRControlFlowInterfaces
     MLIRPDLDialect
+    MLIRInferTypeOpInterface
 )

lib/Dialect/KTDFArch/KTDFArchOpInterfaces.cpp

@@ -34,7 +34,7 @@ using namespace mlir::ktdf_arch;
 auto mlir::ktdf_arch::verifyResource(Operation* op) -> LogicalResult {
   // All resources must be nested under other resources (or devices).
   auto* const parent = op->getParentOp();
-  if (!isa<DeviceOp, Resource>(parent)) {
+  if (!isa<DeviceOp, Resource>(parent) && !parent->hasTrait<IsMeta>()) {
     return op->emitOpError("expects parent op to be '")
            << DeviceOp::getOperationName() << "' or another resource";
   }
@@ -45,7 +45,7 @@ auto mlir::ktdf_arch::verifyResource(Operation* op) -> LogicalResult {
 auto mlir::ktdf_arch::verifyLink(Operation* op) -> LogicalResult {
   // All links must be nested under a resource (or device).
   auto* const parent = op->getParentOp();
-  if (!isa<DeviceOp, Resource>(parent)) {
+  if (!isa<DeviceOp, Resource>(parent) && !parent->hasTrait<IsMeta>()) {
     return op->emitOpError("expects parent op to be '")
            << DeviceOp::getOperationName() << "' or a resource";
   }