[WiP] SRHTs #122
[WiP] SRHTs #122
Conversation
|
|
There was a problem hiding this comment.
It's exciting to see progress here!
I left several comments. Please resolve them, but then take a step back. Part of what you're doing is adding FHT support. This already takes effort without thinking about randomization.
Building up your suite of FHT kernels
Make a folder RandBLAS/trig/ and a file RandBLAS/trig/hadamard.hh. Have your various FHT implementations go here. You already have one for applying to column-major data from the left and storing the result in a column-major sketch. You could easily write implementation that takes in row-major data and writes to a row-major sketch (although the only way you can parallelize this is with BLAS1, I think).
It would be great to have implementations for both the Hadamard transform itself and for the transpose of the Hadamard transform (equivalently, the inverse of the Hadamard transform). If you do this then w.l.o.g. you can always assume you're applying the transformation from the left.
Once you have those functions working, add in the ability to implicitly scale the rows of the input matrix by a vector of coefficients. In sketching we end up setting this vector to a Rademacher random vector, but from an implementation standpoint these functions don't need to care where the vector comes from.
Once you have those cases sorted out, you can allow conflicting layouts for the input matrix and the output matrix (i.e., input is row-major and output is column-major). This is the trick I use for resolving transposition in the sparse-times-dense matrix kernels.
Note: if you feel like you need to allocate a temporary matrix for workspace in order to do anything useful, you can definitely try that.
Writing tests
Make a folder test/test_matmul_cores/test_trig/ and a file test/test_matmul_cores/test_trig/test_hadamard.cc. This will handle tests only for your FHT.
You can take some inspiration from
RandBLAS/test/test_matmul_wrappers/test_sketch_sparse.cc
Lines 53 to 190 in 36da117
|
@aryamanjeendgar, for our reference, here's the code you mentioned from the test suite of FFHT: void fht(double *buf, int log_n) {
int n = 1 << log_n;
for (int i = 0; i < log_n; ++i) {
int s1 = 1 << i;
int s2 = s1 << 1;
for (int j = 0; j < n; j += s2) {
for (int k = 0; k < s1; ++k) {
double u = buf[j + k];
double v = buf[j + k + s1];
buf[j + k] = u + v;
buf[j + k + s1] = u - v;
}
}
}
}Step 1 to figuring out how we'll change it: replace the bit manipulations with standard integer arithmetic (or just give explanatory comments). You can also create a GitHub issue to facilitate the discussion, or make an Overleaf project. |
There was a problem hiding this comment.
Please do a detailed read of my notes in GitHub Issue #99.
I didn't review the whole PR since my plane is landing and I need to put away my laptop.
Ping @aryamanjeendgar
RandBLAS/trig_skops.hh
Outdated
| template<typename T, SignedInteger sint_t> | ||
| void applyDiagonalRademacher( | ||
| bool left, | ||
| blas::Layout layout, | ||
| int64_t rows, | ||
| int64_t cols, | ||
| T* A, | ||
| sint_t* diag | ||
| ) | ||
| { | ||
| if(left && layout == blas::Layout::ColMajor) { | ||
| for(int col=0; col < cols; col++) { | ||
| if(diag[col] > 0) | ||
| continue; | ||
| blas::scal(rows, diag[col], &A[col * rows], 1); | ||
| } | ||
| } | ||
| else if(left && layout == blas::Layout::RowMajor) { | ||
| for(int col=0; col < cols; col++) { | ||
| if(diag[col] > 0) | ||
| continue; | ||
| blas::scal(rows, diag[col], &A[col], cols); | ||
| } | ||
| } | ||
| else if(!left && layout == blas::Layout::ColMajor) { | ||
| for(int row = 0; row < rows; row++) { | ||
| if(diag[row] > 0) | ||
| continue; | ||
| blas::scal(cols, diag[row], &A[row], rows); | ||
| } | ||
| } | ||
| else { | ||
| for(int row = 0; row < rows; row++) { | ||
| if(diag[row] > 0) | ||
| continue; | ||
| blas::scal(cols, diag[row], &A[row * cols], 1); | ||
| } | ||
| } | ||
| } |
There was a problem hiding this comment.
This is fine for now, but it can and should be much more efficient. It's also probably better suited to RandBLAS/util.hh.
|
A description of the tests: All of the tests right now (unfortunately) use Eigen --- I use Eigen extensively to be able to consistently produce I also ended up using Much of this infrastructure (with norm computation, scaling, computing products etc.) can be ported (with a lot more LoC XD) to BLAS with the exception of the permutation tests which tests against Eigen's The correctness tests are fairly straightforward to understand:
The tests are simple, but they cover all of the potential code paths in my code (since the code paths are chosen around two inputs: sketching direction + layout of the input matrix) Let me know how we should proceed next, @rileyjmurray! UPDATE: The latest commit also adds in |
There was a problem hiding this comment.
- You don't need the
RandBLAS::namespace qualifications within RandBLAS itself. - Don't explicitly reference r123::Philox4x32. Reference RandBLAS::DefaultRNG (But, again, without the namespace qualification).
- I don't like people having to explicitly pass in
log_nto the various FHT functions. Do you have a good reason for that? - Please don't use a single monolithic "correctness" test.
- Tests should template numerical precision, not hard-code double. The tolerance for the test would ideally be based on some known error bounds for the operations you're trying to perform. The error bounds you're working with should be very simple since every part of an RHT is an orthogonal matrix, and there are special bounds for those.
- There should be functions in RandBLAS proper (not just tests) for inverting an RHT.
This is just a partial review. I have to jump to a meeting.
|
@rileyjmurray, the recent commit should address all of your comments in the review! |
| float rand_value = r123::u01fixedpt<float>(r.v[0]); | ||
|
|
||
| buff[i] = rand_value < 0.5 ? -1 : 1; |
There was a problem hiding this comment.
Conversion to float isn't necessary. Just use sint_t rand_value = (r.v[0] % 2 == 0) -1 : 1;.
| auto [ctr, key] = seed_state; | ||
|
|
||
| for (int64_t i = 0; i < n; ++i) { | ||
| typename DefaultRNG::ctr_type r = rng(ctr, key); |
There was a problem hiding this comment.
This might raise a compiler error if state_t isn't RNGState<DefaultRNG>.
| template<typename T, SignedInteger sint_t = int64_t> | ||
| void permute_rows_to_top( | ||
| blas::Layout layout, | ||
| int64_t rows, | ||
| int64_t cols, | ||
| sint_t* selected_rows, | ||
| int64_t d, // size of `selected_rows` | ||
| T* A | ||
| ) { | ||
| int64_t top = 0; // Keeps track of the topmost unselected row | ||
|
|
||
| //TODO: discuss precise semantics of `selected_rows` in this function | ||
| if(layout == blas::Layout::ColMajor) { | ||
| for (int64_t i=0; i < d; i++) { | ||
| if (selected_rows[i] != top) { | ||
| // Use BLAS swap to swap the entire rows | ||
| // Swapping row 'selected' with row 'top' | ||
| blas::swap(cols, &A[top], rows, &A[selected_rows[i]], rows); | ||
| } // else, continue; | ||
| } | ||
| } | ||
| else { | ||
| // For `RowMajor` ordering | ||
| for (int64_t i=0; i < d; i++) { | ||
| if (selected_rows[i] != top) { | ||
| blas::swap(cols, &A[cols * selected_rows[i]], 1, &A[cols * top], 1); | ||
| } // else, continue; | ||
| } | ||
| } | ||
| } |
There was a problem hiding this comment.
This is missing the leading dimension parameter for A.
| template<typename T, SignedInteger sint_t = int64_t> | ||
| void permute_cols_to_left( | ||
| blas::Layout layout, | ||
| int64_t rows, | ||
| int64_t cols, | ||
| sint_t* selected_cols, | ||
| int64_t d, // size of `selectedRows` | ||
| T* A | ||
| ) { | ||
| int64_t left = 0; // Keeps track of the topmost unselected column | ||
|
|
||
| if(layout == blas::Layout::ColMajor) { | ||
| for (int64_t i=0; i < d; i++) { | ||
| if (selected_cols[i] != left) { | ||
| // Use BLAS::swap to swap entire columns at once | ||
| // Swapping col 'selected' with col 'top' | ||
| blas::swap(rows, &A[rows * selected_cols[i]], 1, &A[rows * left], 1); | ||
| } | ||
| } | ||
| } | ||
| else { | ||
| // For `RowMajor` ordering | ||
| for (int64_t i=0; i < d; i++) { | ||
| if (selected_cols[i] != left) { | ||
| blas::swap(rows, &A[selected_cols[i]], cols, &A[left], cols); | ||
| } | ||
| } | ||
| } | ||
| } |
There was a problem hiding this comment.
As above, this is missing lda.
| int64_t left = 0; // Keeps track of the topmost unselected column | ||
|
|
||
| if(layout == blas::Layout::ColMajor) { | ||
| for (int64_t i=0; i < d; i++) { | ||
| if (selected_cols[i] != left) { | ||
| // Use BLAS::swap to swap entire columns at once | ||
| // Swapping col 'selected' with col 'top' | ||
| blas::swap(rows, &A[rows * selected_cols[i]], 1, &A[rows * left], 1); | ||
| } | ||
| } | ||
| } | ||
| else { | ||
| // For `RowMajor` ordering | ||
| for (int64_t i=0; i < d; i++) { | ||
| if (selected_cols[i] != left) { | ||
| blas::swap(rows, &A[selected_cols[i]], cols, &A[left], cols); | ||
| } | ||
| } | ||
| } |
| template <typename T> | ||
| void fht_dispatch( | ||
| bool left, // Pre-multiplying? | ||
| blas::Layout layout, | ||
| int64_t num_rows, | ||
| int64_t num_cols, | ||
| int64_t log_n, | ||
| T* A, | ||
| int64_t workspace_ld=0, // leading dimension of the workspace buffer | ||
| T* workspace_buf=nullptr | ||
| ) { | ||
| if(left && layout == blas::Layout::ColMajor) | ||
| fht_left_col_major(A, workspace_buf, workspace_ld, log_n, num_rows, num_cols); | ||
| else if(left && layout == blas::Layout::RowMajor) | ||
| fht_left_row_major(A, workspace_buf, log_n, num_rows, num_cols); | ||
| else if(!left && layout == blas::Layout::ColMajor) | ||
| fht_right_col_major(A, workspace_buf, log_n, num_rows, num_cols); | ||
| else | ||
| fht_right_row_major(A, workspace_buf, workspace_ld, log_n, num_rows, num_cols); | ||
| } | ||
| } |
There was a problem hiding this comment.
This kernel (and all kernels it calls) need to accept an lda parameter.
| int64_t workspace_ld=0, // leading dimension of the workspace buffer | ||
| T* workspace_buf=nullptr |
There was a problem hiding this comment.
If workspace_buf is nullptr then this can segfault. Either make it a required argument or have special logic to handle when it's nullptr.
| //NOTE: I am also going to assume that the user is kind enough to | ||
| // populate this fully with zeros |
There was a problem hiding this comment.
I'd prefer that we not make this assumption. We should zero things out as the first step of applying the RHT.
| template <typename T> | ||
| struct PaddedMatrix { | ||
| int64_t n_rows; | ||
| int64_t n_cols; | ||
| blas::Layout layout; | ||
| T* data; // Data-matrix | ||
| int64_t lda; | ||
| T* work; // Extra space of sz: (2 ** n_closest - lda) * cols/rows | ||
| // Memory **has** to be managed by the user | ||
| //NOTE: I am also going to assume that the user is kind enough to | ||
| // populate this fully with zeros | ||
| //NOTE: We could (should?) explore the possibility of managing the buffer by ourselves | ||
| int64_t work_ld; | ||
| bool left; | ||
|
|
||
| // Trivial constructor | ||
| PaddedMatrix( | ||
| bool left, | ||
| blas::Layout layout, | ||
| int64_t rows, | ||
| int64_t cols, | ||
| T* data, | ||
| int64_t lda, | ||
| T* work | ||
| ) : left(left), layout(layout), n_rows(rows), n_cols(cols), data(data), lda(lda), work(work) {}; | ||
| }; |
There was a problem hiding this comment.
This needs more comprehensive documentation. Compare to SparseSkOp. Add web docs, per DevNotes.md in the screenshot below.
| /* | ||
| * A `pure` description of the \PiRHT as a «data-oblvious» transform | ||
| */ | ||
| template <typename T, SignedInteger sint_t = int64_t, typename state_t = RNGState<DefaultRNG>> | ||
| struct HadamardMixingOp { |
There was a problem hiding this comment.
Comprehensive documentation needed.
| // Populating `perm_idx` | ||
| next_state = repeated_fisher_yates<sint_t>( | ||
| hmo.d, | ||
| hmo.dim, | ||
| 1, | ||
| hmo.perm_idx, | ||
| next_state | ||
| ); |
| if(!hmo.filled()) { | ||
| hmo.diag_scale = new sint_t[hmo.dim]; | ||
|
|
||
| hmo.perm_idx = new sint_t[hmo.d]; | ||
|
|
||
| auto [ctr, key] = hmo.seed_state; | ||
|
|
||
| // Populating `diag` | ||
| auto next_state = RandBLAS::trig::kernels::generate_rademacher_vector_r123(hmo.diag_scale, hmo.dim, hmo.seed_state); | ||
|
|
||
| // Populating `perm_idx` | ||
| next_state = repeated_fisher_yates<sint_t>( |
There was a problem hiding this comment.
Remove excessive blank lines. We use blank lines to separate blocks of code that are both nontrivial and conceptually distinct from one another.
| * a very similar signature to the driver itself | ||
| */ | ||
| template <typename T, typename state_t = RNGState<DefaultRNG>, SignedInteger sint_t = int64_t> | ||
| inline int64_t miget_workspace_sz( |
There was a problem hiding this comment.
Naming convention: we don't use "sz" as an abbreviation for "size."
| // populate this fully with zeros | ||
| //NOTE: We could (should?) explore the possibility of managing the buffer by ourselves | ||
| int64_t work_ld; | ||
| bool left; |
There was a problem hiding this comment.
I don't like the name left in this context. I suggest you define an enum (perhaps a static member of this type) that you can use to indicate if you're padding with extra rows or extra columns. As a matter of practice, you'll pad with rows when applying an RHT from the left and you'll pad with columns when applying an RHT from the right.
| template <typename T, SignedInteger sint_t = int64_t, typename state_t = RNGState<DefaultRNG>> | ||
| state_t compute_next_state(state_t state, int64_t dim, int64_t d) { | ||
| state.counter.incr(d * 1 + dim); | ||
|
|
| bool filled() { | ||
| if(this -> diag_scale == nullptr \ | ||
| && this -> perm_idx == nullptr) | ||
| return false; | ||
| else | ||
| return true; | ||
| } |
There was a problem hiding this comment.
This function can be removed and in-lined with equivalent code
bool filled = (hmo.diag_scale != nullptr) && (hmo.perm_idx != nullptr);
| template <typename T, typename state_t = RNGState<DefaultRNG>, SignedInteger sint_t = int64_t> | ||
| void fill_hadamard( | ||
| HadamardMixingOp<T, sint_t> &hmo | ||
| ) { |
There was a problem hiding this comment.
Indent the contents of this function.
|
|
||
| //Step 2: Apply the Hadamard transform (invH = H.T = H) | ||
| // This has to be a `double` because `blas::scal` really doesn't like being passed integer types | ||
| double scal_factor = std::pow(2, int(std::log2(hmo.dim)) + 1); |
There was a problem hiding this comment.
Use T scal_factor, not double scal_factor.
|
@aryamanjeendgar I've reviewed and left more comments. |
This is a draft PR for integrating SRHTs into
RandBLAS