google / google/xls

Improve optimization of masking and updating of slices

Open
#273 0 comments 0 reactions 1 assignee Claimed by @meheff View on GitHub
optimizer
Dominant language
C++
Stars
1.9k
Forks
283
Avg merge
2d 10h
Merged PRs (30d)
135

Description

Sample input DSLX:

```
pub fn make_mask
B else B>
(num_ones : uN[B])
-> uN[N] {
let num_bits_clamped =
(N as uN[MAX_N_B])
if (num_ones as uN[MAX_N_B]) > (N as uN[MAX_N_B])
else num_ones as uN[MAX_N_B];
let wider = (uN[N_PLUS_1]:1 << (num_bits_clamped as uN[N_PLUS_1]))
- uN[N_PLUS_1]:1;
wider as uN[N]
}

pub fn update_slice(
original : uN[N],
offset : u32,
length : u32,
update : uN[NUpdate])
-> uN[N] {
let ms_shift = offset + length;
let ms_mask = make_mask(N - ms_shift) << (ms_shift as uN[N]);
let msbits = original & ms_mask;

let lsbits = original & make_mask(offset);

let update_clamped = (update as uN[N]) & make_mask(length);
let update_shifted = update_clamped << (offset as uN[N]);

msbits | update_shifted | lsbits
}
```

We generally do a poor job of optimizing this with unnecessary shifts and adds (from the make_mask) remaining in the optimized IR. Of particular interest in update_slice is a form which an input is "written" to one of several mutually exclusive partitions of the output. Imagine, a 64-bit value being written into one of 4 locations (dynamically selected) in a 256bit wide vector. This can be represented simply as a 2-way mux at each bit, but we generate some mostrosity of shifts and adds and terribleness.

To fix this issue we should have some DSLX-level tests added to xls/tests/dslx_optimization_test which have these DSLX constructs.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.