Skip to content

Carry layout/placement decisions to StableHLO as untyped metadata #1147

Description

@michalharakal

Follow-up of #1134. The #1134 resolution is decide-and-carry: core resolvers (WeightFormResolver, AllocationResolver) decide; the recorded path's job is to carry those decisions to the MLIR so the downstream compiler can honor them. This issue is the carry slice — deliberately deferred until an in-repo IREE consumer exists (today it would be write-only metadata).

Scope (when unblocked)

Follow the TensorSpecEncoding precedent — untyped TensorSpec.metadata + typed accessor extension; sk.ainet.lang.memory types are never imported into skainet-compile (rule recorded in the SKEEP resolution doc):

  • Tape→graph: extend TensorRef (id/encoding, optionally layout) in TraceSession.refOf; stop dropping metadata in TraceToGraphBuilder.buildInputSpecs/buildOutputSpecs.
  • Graph: TensorSpecLayout.kt mirroring TensorSpecEncoding.kt; a decision pass in skainet-compile-opt (after TransposeEliminationPass, before LLMFusionPass), registered via TargetOptimizers; populate the pre-built dead seams ResolvedComputeGraph.resolvedLayout(edgeId) / backendAssignment(nodeId).
  • Emission: branch at the designated point in dag2hlo ; module attribute skainet.tensor_layouts serialized structurally (strides/blockOrder/blockSize — not name-strings, which lose blockSize today); later per-type encoding attrs via TypeMapper (requires auditing the ~102 raw "tensor<…>" literals in converters).
  • Boundary: ExternalParameterRef gains layout/blockOrder so kernel-feed-order weights can be blitted verbatim into .irpa with their order declared in the MLIR reference.
  • Thread target for real (ConversionContext.target is currently never set and never read).

Gate

Blocked on an in-repo IREE consumer (see the .vmfb parity issue). Do not start before that exists.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    codingImplementation task (DARC: C)compute-backendCompute backends and graphs

    Type

    No type

    Projects

    No projects

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions