Token arrays needed to supplement channel arrays

mikex-oss commented 2 months ago

What's hard to do? (limit 100 words)

Trying to take advantage of channel arrays now that #704 is fixed.

However, I think I'm blocked by the fact that I can't keep track of tokens across the channel array separately. The test example as part of delivering the feature only shows sequential tokens:

https://github.com/google/xls/blob/b117a11ef67534bba97f8562b024278fb9ba7ee5/xls/dslx/ir_convert/proc_config_ir_converter_test.cc#L305-L311

However, if I'd like each channel in the array to operate in parallel, I want to save all those tokens so that I can sequence downstream sends/receives based on the specific token element.

I tried out token arrays and it first fails on:

F0911 15:12:26.916626    2400 statusor.cc:99] Attempting to fetch value instead of handling error INVALID_ARGUMENT: Type "token" has no defined signedness.
*** Check failure stack trace: ***
    @     0x564c153d0304  absl::log_internal::LogMessage::SendToLog()
    @     0x564c153cfe76  absl::log_internal::LogMessage::Flush()
    @     0x564c155a1b4e  absl::log_internal::LogMessage::~LogMessage()
    @     0x564c1539fa9f  Google3AbseilInternalLog::Hook()
    @     0x564c153d4618  absl::internal_statusor::ThrowBadStatusOrAccess()
    @     0x564c1352e896  xls::dslx::BuiltinTypeAnnotation::GetSignedness()
    @     0x564c1221ba73  xls::dslx::(anonymous namespace)::DeduceVisitor::HandleArrayTypeAnnotation()
    @     0x564c122165a1  xls::dslx::Deduce()
    @     0x564c121eb0d4  std::__u::__function::__policy_invoker<>::__call_impl<>()
    @     0x564c122b4c7b  xls::dslx::DeduceCtx::Deduce()
    @     0x564c1221d9e5  xls::dslx::(anonymous namespace)::DeduceVisitor::HandleTupleTypeAnnotation()
    @     0x564c122165a1  xls::dslx::Deduce()
    @     0x564c121eb0d4  std::__u::__function::__policy_invoker<>::__call_impl<>()
    @     0x564c122b4c7b  xls::dslx::DeduceCtx::Deduce()
    @     0x564c12216982  xls::dslx::DeduceAndResolve()
    @     0x564c12226e31  xls::dslx::(anonymous namespace)::DeduceVisitor::HandleUnrollFor()

I tried hacking BuiltinTypeAnnotation to define a signedness for tokens, but then it just trips over thinking that it is a sN or uN type:

TypeInferenceError: uN[8] Want argument 0 to 'update' to be an array; got uN[8]

Current best alternative workaround (limit 100 words)

I'm not aware of any workaround aside from keeping the sends/receives manually unrolled.

Your view of the "best case XLS enhancement" (limit 100 words)

Would like to be able to do something like:

        let send_toks = unroll_for!(i, send_toks): (u32, token[N]) in u32:0..N {
            update(send_toks, i, send(tok, chan[i], x))
        }(zero!<token[N]>());

Could something like how we handle channel arrays be applied to tokens as well?

mikex-oss commented 4 weeks ago

The notable additional feature mentioned there is the ability to join on such a token array.

richmckeever commented 3 weeks ago

We've talked about this a bit, and even played around with a possible implementation, but I think the conclusion is we are not likely to support a syntax like update(send_toks, i, send(tok, chan[i], x)), because tokens can't be aliased. With today's token model, I think there would need to be a map construct like:

let tokens = const map channel: chan<u32> in enumerate(channels) {
  send(join(), channel, data)
};

which would have an expansion like:

let token__0 = send(channels__0, data);
let token__1 = send(channels__1, data);
etc.

and no underlying arrays at all.

However, I think fundamentally the desire to put tokens in an array kind of pushes against their nature as more code labels (or IR node labels) than real variables. Either we should turn them into variables that can be aliased, or we should turn them into something like parameterizable code labels, as in:

unroll_for!(i, _): (u32, ()) in u32:0..N {
    send_of_x(i):: send(chan[i], x);
}()

unroll_for!(i, _): (u32, ()) in u32:0..N {
    after send_of_x(i):: let (data, _) = recv(chan[i]);
}()

Either way, because there is not a strong immediate need (and the above example just begs the question of why not use a unified loop), I think this will be deferred until we look at how implicit tokens and the like might be done.

ericastor commented 3 weeks ago

I think an alternative approach would just be to make tokens valid subjects for array operations. As currently implemented, if XLS were to allow this, the result will be semantically equivalent to aliasing tokens in circumstances where XLS's optimizations can see through the operations at compile time, with a fallback to effectively joining the tokens that could have been aliased when it's not resolvable at compile time.

Pros:

Easy to implement
Equivalent to full aliasing in many circumstances
Semantically consistent; tokens express "non-data" dependencies (impacting execution sequencing), and XLS tries to reduce the impact of unnecessary dependencies where it can.

Cons:

Potentially misleading (users might expect full aliasing at execution time)

richmckeever commented 3 weeks ago

We have played around with allowing lowering of update(send_toks, i, send(tok, chan[i], x)) and send_toks[i]. Without the notion that send_toks[i] evaluates to a join of everything it might be, appropriate front-end guard rails are hard to install. Even with that idea, it feels to me like a broken metaphor that we have time to think about how to fix, i.e. this is far from being the biggest problem with tokens.

google / xls