Relax unnecessary src/dst storage restriction for scalar bitwise tile ops

[Zhendong404]

Problem

On main branch, pto.tands rejects IR where src and dst refer to the same tile storage:

'pto.tands' op expects src and dst to use different storage

The restriction appears to come from the shared verifier helper used by scalar bitwise tile ops:

static FailureOr<Type> verifyDistinctRowMajorUnaryTileOpCommon(
    Operation *op, Value src, Value dst, StringRef srcName = "src",
    StringRef dstName = "dst") {
  if (src == dst) {
    op->emitOpError("expects src and dst to use different storage");
    return failure();
  }
  ...
}

TAndSOp::verify, TOrSOp::verify, and TXorSOp::verify all call this helper.

Reproducer

module attributes {pto.target_arch = "a5"} {
  func.func @tands_inplace() {
    %tile = pto.alloc_tile
      : !pto.tile_buf<loc=vec, dtype=i32, rows=16, cols=64, v_row=16, v_col=64,
                      blayout=row_major, slayout=none_box, fractal=512, pad=0>
    %scalar = arith.constant 0xFF : i32
    pto.tands ins(%tile, %scalar : !pto.tile_buf<loc=vec, dtype=i32, rows=16, cols=64, v_row=16, v_col=64,
                                      blayout=row_major, slayout=none_box, fractal=512, pad=0>, i32)
              outs(%tile : !pto.tile_buf<loc=vec, dtype=i32, rows=16, cols=64, v_row=16, v_col=64,
                                blayout=row_major, slayout=none_box, fractal=512, pad=0>)
    return
  }
}

Command:

ptoas --pto-arch=a5 --pto-backend=vpto tands_inplace.pto -o -

Actual result:

loc("tands_inplace.pto":6:5): error: 'pto.tands' op expects src and dst to use different storage
Error: Failed to parse MLIR.

Why this restriction seems too strong

The IR-level restriction seems to reflect one possible C++/tile implementation strategy rather than an inherent semantic requirement of tands.

For example, an implementation that expands the scalar into dst first and then runs a binary tile op like:

TEXPANDS_IMPL(dst, scalar);
TAND_IMPL(dst, src, dst);

cannot safely support src == dst, because broadcasting the scalar into dst overwrites the original source before the and operation reads it.

However, tands can be implemented in a src/dst-reuse-safe way by loading source data before storing the result. The VPTO-style lowering naturally has this shape:

vec        = vlds(src)
scalar_vec = vbr(scalar)
result     = vand(vec, scalar_vec, mask)
vsts(result, dst)

With that lowering order, src == dst is well-defined: each vector chunk is read before the corresponding result is written back.

So the current verifier rejects a semantically valid in-place form because one backend implementation strategy is not in-place safe.

Expected behavior

pto.tands should allow src and dst to reuse the same tile storage when the selected lowering/backend can implement it safely.

At minimum, the verifier should not unconditionally reject src == dst for pto.tands at the IR level.

The same question likely applies to pto.tors, and possibly pto.txors depending on the intended role of its tmp operand.

Suggested fix direction

1. Revisit verifyDistinctRowMajorUnaryTileOpCommon usage for TAndSOp, TOrSOp, and TXorSOp.
2. Split the verifier concerns:
   • Keep shape/layout/element-type checks in the common helper.
   • Do not make src != dst a universal IR invariant unless the op semantics truly require it.
3. Update the C++/tile lowering path for tands/tors to use an in-place-safe sequence, or route these ops through the VPTO-style vlds + vbr + vand/vor + vsts lowering where available.
4. Add regression tests for in-place scalar bitwise tile ops, e.g.:
   • pto.tands ins(%tile, %scalar) outs(%tile) should verify and lower on A5 VPTO.
   • pto.tors ins(%tile, %scalar) outs(%tile) should verify and lower if the same reasoning applies.

Impact

The current restriction also interacts badly with memory planning: if PlanMemory legally reuses two short-lived tile buffers by assigning the same local offset, later verification can fail with the same src/dst different storage diagnostic. Relaxing the IR-level restriction, or making lowering explicitly safe for in-place use, would avoid rejecting valid programs and reduce pressure on the memory planner to preserve an implementation-specific non-aliasing constraint.

[Zhendong404]

One extra point from the PlanMemory failure mode: there are two possible repair directions here, and they should probably be discussed separately.

Direction 1: change the PTO ISA / PTOAS legality constraint

If TANDS(src, scalar, dst) is semantically just:

dst = src & scalar

then src == dst should be a legal in-place update. In that model, the current src/dst distinct-storage verifier rule is too strong and should be relaxed once the PTO ISA C++ implementation is made alias-safe.

This is the direction I think is semantically preferable: it avoids exposing a backend implementation detail as an IR-level legality rule, and it lets memory planning reuse tile storage for short-lived values naturally.

The implementation-side fix would be to make the PTO ISA wrapper read src before clobbering dst, or use a temporary/register sequence when needed. The current pattern:

TEXPANDS_IMPL(dst, scalar);
TAND_IMPL(dst, src, dst);

is not safe for src == dst because the broadcast overwrites the original source tile before the AND reads it.

Direction 2: make PlanMemory satisfy the current constraint

If the current distinct-storage rule is intentional or cannot be changed immediately, then PlanMemory must model it as an allocation constraint.

In that model, for ops such as tands/tors/txors, PlanMemory should not assign the input tile and output tile to the same storage when the verifier/backend requires them to be distinct. Otherwise PlanMemory can produce IR that is lifetime-valid but verifier-invalid.

This would be a compatibility fix for the current backend limitation, but it may unnecessarily increase local memory pressure if in-place tile-scalar ops are actually semantically valid.

So the short-term fix could be Direction 2, but the longer-term preferred fix may be Direction 1 if the ISA/spec agrees that in-place TANDS is valid.

[zhangstevenunity]

This feature is not supported by pto isa