tvm: preserve elementwise ranks for fusion

[guillon]

Motivation

TVM elementwise operators were flattened to a single axis by inserting hidden
reshapes. Consumer or producer fusion therefore targeted a reshape rather than
the intended elementwise operation, making fusion ineffective.

Description

Generate TVM ReLU directly with the rank and dimensions propagated from its
input tensor type. This removes the intermediate reshapes, exposes one parallel
scheduling axis per tensor dimension, and allows producer and consumer fusion
to target the ReLU block directly.

Update the matmul and convolution FileCheck tests and add coverage for fusing
ReLU under different matmul loop levels and through the descriptor scheduler.

Additionally, add representative ResNet18 and YOLO9000 multi-node graph
fixtures and allow loop-explore to schedule a named graph node with --node.

Commits

• tests: fix wrongly ordered conv2d axes
• tvm: implement proper fusion, discarding reshapes
• tests: add multi node test graphs from yolo/resnet
• explore: support scheduling a selected graph node

Testing

• Verified that TVM producer fusion targets relu -> ... and consumer fusion
  targets ... -> relu.
• pytest -q tests/pytest/tvm — 14 passed
• Nine affected descriptor and fusion FileCheck tests — 9 passed
• Graph serialization and node-targeted exploration tests — 13 passed

Discussion

Explicit axis coalescing remains possible future scheduling work. This change
intentionally exposes the propagated tensor rank instead of adding implicit
reshapes.

[guillon]

@liamsemeria I basically implemented what I discussed some time ago, i.e. generating elementwise with full rank dimensions instead of reshaping it, giving axis names i to z, i.e. i, j, k, l, ... up to the rank. This avoids reshapes in the middle which were fused instead of the actual relu op.

Is it in line with what you did in MLIR? Should I also look at updating the MLIR side?

[codecov]

Codecov Report

❌ Patch coverage is 88.23529% with 2 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/xtc/search/explore.py	66.66%	1 Missing ⚠️
src/xtc/search/strategies.py	90.90%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!