google / xls

XLS: Accelerated HW Synthesis
http://google.github.io/xls/
Apache License 2.0
1.21k stars 179 forks source link

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

Open rw1nkler opened 7 months ago

rw1nkler commented 7 months ago

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
//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

cdleary commented 7 months ago

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.

mikex-oss commented 1 month ago

@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 commented 2 weeks ago

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.