xls icon indicating copy to clipboard operation
xls copied to clipboard

[enhancement] Treat for-loop iterator value as constant when iteration over constant values

Open rw1nkler opened this issue 1 year ago • 2 comments

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

As for today, the value of DSLX for-loop iterator is not treated as constant when iterating over constant values. Having the ability to use the for-loop iterator in this way can improve the expressiveness of the language. In all these cases, when iterating over constant values it should be possible to unroll the for-loop.

These small examples show the current limitations of the language:


let data = u64:0x00_11_22_33_44_55_66_77;
for (i, ()) in s32:0..s32:8 {
    // works if rewritten with width slices
    let slice = data[(i * s32:8) : (i + s32:1) * s32:8];
    trace_fmt!("slice is {}", slice);
}(());

// for (i, ()) in s32:0..s32:8 {
//     let slice = data[(i * s32:8) : (i + s32:1) * s32:8];
// ~~~~~~~~~~~~~~~~~~~~~~~~^ TypeInferenceError: Unable to resolve slice start to a compile-time constant.

// the unrolled version works
let slice = data[(s32:0 * s32:8) : (s32:0 + s32:1) * s32:8];
let slice = data[(s32:1 * s32:8) : (s32:1 + s32:1) * s32:8];
let slice = data[(s32:2 * s32:8) : (s32:2 + s32:1) * s32:8];
//...
let slice = data[(s32:7 * s32:8) : (s32:7 + s32:1) * s32:8];
fn add_const<Y: s32>(x: s32) -> s32 { x + Y }
for (i, acc) in s32:0..s32:8 {
    add_const<i>(acc);
}((u32:3));

// for (i, acc) in s32:0..s32:8 {
//     add_const<i>(acc);
// ~~~~~~~~~~~~~~^ TypeInferenceError: sN[32] Parametric expression `i` was not constexpr -- parametric values must be compile-time constants
  • An example of data concatenation. We tried to use similar logic as a part of a buffering proc. The proc accumulated incoming data in the state and sent back the requested number of bytes of data.
//This data may come from channels
let s = s32:4;
let data1 = u64:0x00_11_22_33_44_55_66_77;
let data2 = u64:0x88_99_AA_BB_CC_DD_EE_FF;

let slice = for (i, slice) in s32:0..s32:8 {
   if (s == i) {
       data1[i * s32:8 : s32:64] ++ data2[0 : i * s32:8]
   } else {
       slice
   }
}(u64:0);
trace_fmt!("slice is {}", slice);

let slice = if  s == s32:0 {data1[s32:0 * s32:8 : s32:64] ++ data2[0 : s32:0 * s32:8]} else {u64:0};
let slice = if  s == s32:1 {data1[s32:1 * s32:8 : s32:64] ++ data2[0 : s32:1 * s32:8]} else {slice};
let slice = if  s == s32:2 {data1[s32:2 * s32:8 : s32:64] ++ data2[0 : s32:2 * s32:8]} else {slice};
let slice = if  s == s32:3 {data1[s32:3 * s32:8 : s32:64] ++ data2[0 : s32:3 * s32:8]} else {slice};
let slice = if  s == s32:4 {data1[s32:4 * s32:8 : s32:64] ++ data2[0 : s32:4 * s32:8]} else {slice};
let slice = if  s == s32:5 {data1[s32:5 * s32:8 : s32:64] ++ data2[0 : s32:5 * s32:8]} else {slice};
let slice = if  s == s32:6 {data1[s32:6 * s32:8 : s32:64] ++ data2[0 : s32:6 * s32:8]} else {slice};
let slice = if  s == s32:7 {data1[s32:7 * s32:8 : s32:64] ++ data2[0 : s32:7 * s32:8]} else {slice};
trace_fmt!("slice: {:#x}", slice);

Current best alternative workaround (limit 100 words)

Unroll the for-loop as in the examples above

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

In cases when the for-loop iterates over constant values, the iterator can be treated as a constant and XLS can treat the loop as if it were unrolled

rw1nkler avatar Apr 19 '24 14:04 rw1nkler

Note that this is what unroll_for! is for in the DSL -- in a normal type system that doesn't macro expand it doesn't make sense to have one binding (i.e. name in the program) with different types in different loop iterations, which is why having i be constexpr doesn't make as much sense.

I'm not sure how fully supported it is at the moment: https://github.com/google/xls/blob/164fe8a8e320ce6c32a4a9548ae0c783f70503f6/xls/dslx/frontend/parser_test.cc#L1706

It's something @RobSpringer was working on that we didn't get to pick up and run with in the time since.

cdleary avatar Apr 19 '24 18:04 cdleary

@richmckeever @meheff is this something we still want to leave open given unroll_for! from https://github.com/google/xls/issues/1387#issuecomment-2067090587 is resolved?

mikex-oss avatar Oct 17 '24 16:10 mikex-oss

I found one scenario where unroll_for! is insufficient because the DSLX to IR conversion does not guarantee any ordering in the unrolled for.

Take this contrived example:

import std;

fn foo() -> u32 {
    unroll_for! (j, result): (u32, u32) in range(u32:0, u32:10) {
        for (k, _): (u32, u32) in range(u32:0, std::upow(u32:2, j)) {
            k
        }(result)
    }(u32:0)
}

#[test]
fn foo_test() { assert_eq(foo(), u32:511); }

This passes in DSLX but generates a JIT mismatch.

mikex-oss avatar Oct 31 '24 17:10 mikex-oss