Skip to content

Add support for instanced subgraphs in ktdf_arch. - #37

Open
KFAFSP wants to merge 15 commits into
torch-spyre:mainfrom
KFAFSP:kfaf/ktdf-arch-neighborhood
Open

KFAFSP wants to merge 15 commits into
torch-spyre:mainfrom
KFAFSP:kfaf/ktdf-arch-neighborhood

Conversation

@KFAFSP

@KFAFSP KFAFSP commented Sep 1, 2026

Copy link
Copy Markdown
Contributor

We're starting to run into the issue that the device declaration files are becoming quite large to fit all the copies of groups of a kind. Since these copies must all be identical in the IR (required by having the same kind), these ops carry little information. Still, they must all exist as unique objects in memory to allow resource allocation.

This PR makes a few backwards compatible changes to ktdf_arch, including turning its containers into graph regions. It then adds the two new ktdf_arch.neighbor(hood) operations, which can be used to define regular recurrences of subgraphs.

This allows the user to write an architecture specification that lists every unique topology only once, and then parameterize over its repetitions. There is an example that features a ring topology. The resulting graph must be instantiated before it is consumed by the scheduler. This is done via the ktdfarch-instantiate-neighborhoods pass.

EDIT: As a side note, neighborhoods are regular subgraphs. That means existing analyses such as ResourceKinds will observe the same structure regardless of whether the neighborhoods were instantiated or not. They just won't be traversed by endpoint lookups, making each value that would cross their boundary opaque.

Comment thread include/dataflow-scheduler/Dialect/KTDFArch/KTDFArch.td Outdated
Comment thread test/Dialect/KTDFArch/instantiate-neighborhoods.mlir
Comment thread lib/Dialect/KTDFArch/Transforms/InstantiateNeighborhoods.cpp
Signed-off-by: Karl F. A. Friebel <karl.friebel@ibm.com>
This relaxes the operand types of `ktdf_arch.yield` to accept any kind
of endpoint type, but preserves the format where all types are omitted,
in which case `!ktdf_arch.exec_unit` is assumed.

Signed-off-by: Karl F. A. Friebel <karl.friebel@ibm.com>
Signed-off-by: Karl F. A. Friebel <karl.friebel@ibm.com>
Signed-off-by: Karl F. A. Friebel <karl.friebel@ibm.com>
This pass instantiates `ktdf_arch.neighborhood` tiles and resolves the
references to them via `ktdf_arch.neighbor` to concrete values, where
possible. Note that the pass may fail for certain recursive constructs.

Signed-off-by: Karl F. A. Friebel <karl.friebel@ibm.com>
The documentation said that the first result dimension is the one over
the operands, but the code was using the last dimension instead. This
commit changes the code to match the documented behavior.

Signed-off-by: Karl F. A. Friebel <karl.friebel@ibm.com>
Signed-off-by: Karl F. A. Friebel <karl.friebel@ibm.com>
Instantiations of neighborhoods are independent of references to them.

Signed-off-by: Karl F. A. Friebel <karl.friebel@ibm.com>
Review determined that the tuple indexing form of `ktdf_arch.neighbor`
is confusing because the number of result dimensions of the map does not
match the declared type.

This commit changes the printer / parser for this op to introduce a
separate form with a local neighborhood that clarifies this and prints
the appropriate type. This also improves the documentation.

Unfortunately, there still is a degenerate form where only one argument
is provided but has to be treated as a tuple. This appears when the user
has specified an invalid index into a neighborhood of dimensions [1].
While the new printer handles this case without surprises, the verifier
still has to accept this special degenerate form.

Signed-off-by: Karl F. A. Friebel <karl.friebel@ibm.com>
@KFAFSP
KFAFSP force-pushed the kfaf/ktdf-arch-neighborhood branch from d456c7d to 33bd5f7 Compare September 11, 2026 09:00
As discussed in the review, this stand-alone tool allows us to anchor
more thorough architecture verification (still TODO) and baking into
a self-contained executable with a stable CLI, suitable for CI.

This also introduces the bytecode support, which is already provided
by the `DeviceManager` import implementation. Using `ktdf-archgen` with
the `--emit-bytecode` option, we can optimize baked files.

Signed-off-by: Karl F. A. Friebel <karl.friebel@ibm.com>
Signed-off-by: Karl F. A. Friebel <karl.friebel@ibm.com>
Signed-off-by: Karl F. A. Friebel <karl.friebel@ibm.com>
@KFAFSP
KFAFSP force-pushed the kfaf/ktdf-arch-neighborhood branch from 4637c16 to 7520772 Compare September 11, 2026 12:15
Signed-off-by: Karl F. A. Friebel <karl.friebel@ibm.com>
@KFAFSP

KFAFSP commented Sep 11, 2026

Copy link
Copy Markdown
Contributor Author

@bmahjour Ive added ktdf-archgen and tested bytcode, and it actually already works out of the box. Import declarations can already import bytecode, and the tool can already write and recover it. Parsing it is also much faster - I did not add skipping the verifier to the importer though, that could be for later.

Comment thread lib/Dialect/KTDFArch/KTDFArchOps.cpp
Comment thread include/dataflow-scheduler/Dialect/KTDFArch/KTDFArch.td Outdated

// CHECK-NEXT: datapath #skip_one %[[EXEC0]] to %[[EXEC2]]
datapath {kind = "skip_one"} %exec_0 to %exec_2 : exec_unit, exec_unit
}

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please add a more complicated test involving nesting. I get a crash when I try the following:

ktdf_arch.device @nested {
  %outer = neighborhood %self_o : (exec_unit)[2] {
    %inner = neighborhood %self_i : (exec_unit)[3] {
      %e = exec_unit @exec
      yield %e
    }
    %e0 = neighbor affine_map<(d0) -> (d0+1)> in %inner : (exec_unit)[3]
    yield %e0
  }
  %a = neighbor affine_map<() -> (0)> in %outer : (exec_unit)[2]
  %b = neighbor affine_map<() -> (1)> in %outer : (exec_unit)[2]

  datapath %a to %b : exec_unit, exec_unit
}

It has to do with the (d0) -> (d0+1) map to select from the %inner. It goes through if I map (d0) -> (0).

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, I will. I must have messed up the out of bounds check there, thanks for catching it!

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I was unable to reproduce a crash and instead observed only the expected behavior. I have added a test regardless. Maybe you were running an old commit?

Signed-off-by: Karl F. A. Friebel <karl.friebel@ibm.com>
Signed-off-by: Karl F. A. Friebel <karl.friebel@ibm.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants