feat: bitwise-ops-for-tensors

[quinton11]

Pull Request Template

Checklist

• Confirmed that run-checks all script has been executed.
• Made sure the book is up to date with changes in this PR.

Related Issues/PRs

#2234
Blocked by CubeCL

Changes

Bitwise Operations for Tensors

• Bitwise_And
• Bitwise_Or
• Bitwise_Xor
• Bitwise_And_Scalar
• Bitwise_Or_Scalar
• Bitwise_Xor_Scalar
• Bitwise_Not
• Bitwise_left_shift
• Bitwise_right_shift

Testing

The corresponding tests for the ops were included under the burn_tensor/tensor/tests directory.
Candle seems to not have bitwise operations so as it stands the implementation for the candle backend
is replaced with the todo macro.

[laggui]

Great start! Some minor comments but here is my preliminary review:

I think the bitwise ops should only be added as an int tensor operation, it doesn't make sense for floats.

We could eventually extend the ops to boolean tensors with their logical counterpart (would be applied on a single bit represented by the bool), but this can be left for another PR.

We can leave the candle ops as unimplemented, but for the JIT backends we should wait to merge once it's implemented with cubecl.

[quinton11]

The idea behind the whole float implementation was to convert to int before computing and back to float after, you're right, doesn't really make sense, will make the changes.

Raised an issue in the cubecl repo for adding the bitwise op support, will look at that first before implementing here

[github-actions]

This PR has been marked as stale because it has not been updated for over a month

[0x7CFE]

Amazing! This is the PR I was looking for!

Is there a way to support bit counting on integer tensors as well?

P.S.: Also, what is the status of the PR? Does somebody work on it? Maybe I should pick it up?

[quinton11]

Amazing! This is the PR I was looking for!

Is there a way to support bit counting on integer tensors as well?

P.S.: Also, what is the status of the PR? Does somebody work on it? Maybe I should pick it up?

@0x7CFE Yes, this implementation will allow for bit counting. Unless you mean having direct tensor op methods where you can just call and have the bits counted, ex: tensor::count_leading_ones or tensor::count_zeros. With that, existing tensor libraries like torch or numpy don't have such direct methods. But this implementation allows for using rust's bitwise operators with the tensors to acheive whatever bit counting you want to do. @laggui what do you think, should there be direct methods?

Also this PR is pending a release in cubecl with this implementation Cubecl

[0x7CFE]

With that, existing tensor libraries like torch or numpy don't have such direct methods.

Just in case, numpy has bitwise_count, wgsl has countOneBits, CUDA has __popc. Apparently, pytorch does not have a ready solution, but there are some tricks to do the counting efficiently given that boolean ops are available.

So I believe it would be possible to efficiently implement this at least for some of the backends.

P.S.: Since the blocker mentioned above was already resolved, what are the next steps for this PR?

[quinton11]

So I believe it would be possible to efficiently implement this at least for some of the backends.

Possible yes, what bit counting ops exactly are you looking to include? As for if its something the burn team would like to move forward with @laggui would have to comment on that

P.S.: Since the blocker mentioned above was already resolved, what are the next steps for this PR?

Waiting on an official release for CubeCl containing that change so we can implement it in burn

[0x7CFE]

what bit counting ops exactly are you looking to include?

I described my use case in #2641, so essentially I need uXX::count_ones().

[wingertge]

I described my use case in #2641, so essentially I need uXX::count_ones().

I opened a PR to add the required backend ops to cubecl (tracel-ai/cubecl#391), so this should be easy enough to add at least for CPU and JIT backends.

[laggui]

Sorry all I didn't access my computer during the holidays 😅 github app is not as great for reviews and long discussions

But this implementation allows for using rust's bitwise operators with the tensors to acheive whatever bit counting you want to do. @laggui what do you think, should there be direct methods?

I think something like bitwise_count

Also this PR is pending a release in cubecl with this implementation Cubecl

The main branch of burn is actually kept up to date as much as possible with the latest cubecl revision, so your changes are already available 🙂

[laggui]

P.S.: Since the blocker mentioned above was already resolved, what are the next steps for this PR?

I just need to review the requested changes for the current PR, but this one should land pretty soon.

Regarding your request for bit counting, I think once the linked cubecl PR lands (thanks Genna 🙏) we could add a bitwise_count op. Feel free to open a draft PR if you want to get the ball rolling 🙂

[laggui]

Thanks for addressing the changes 🙂

I have a couple of comments, most of them just pointing out some commented code you forgot to remove.

Also, bitwise tests are failing but that's probably just because you're not using the correct operations (as pointed out in the review below).

[codecov]

Codecov Report

Attention: Patch coverage is 53.96114% with 616 lines in your changes missing coverage. Please review.

Project coverage is 83.56%. Comparing base (e40c69b) to head (d0a4658).
Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
crates/burn-router/src/ops/op_int.rs	0.00%	186 Missing ⚠️
crates/burn-tch/src/ops/base.rs	0.00%	103 Missing ⚠️
crates/burn-jit/src/kernel/binary_int.rs	47.12%	101 Missing ⚠️
crates/burn-ndarray/src/ops/int_tensor.rs	0.00%	56 Missing ⚠️
crates/burn-tch/src/ops/int_tensor.rs	0.00%	48 Missing ⚠️
crates/burn-jit/src/kernel/unary_int.rs	54.08%	45 Missing ⚠️
crates/burn-autodiff/src/ops/int_tensor.rs	0.00%	33 Missing ⚠️
crates/burn-candle/src/ops/int_tensor.rs	0.00%	22 Missing ⚠️
crates/burn-router/src/runner.rs	0.00%	22 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2498      +/-   ##
==========================================
- Coverage   83.91%   83.56%   -0.36%     
==========================================
  Files         821      824       +3     
  Lines      106890   108225    +1335     
==========================================
+ Hits        89702    90440     +738     
- Misses      17188    17785     +597     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[laggui]

Thanks for all the additions 🙏

Should be good with one last round of minor changes. See my comments below.

[laggui]

Fixed the remaining dead code and added the ops to the book while I was at it.

Good to go now! 🎉

[quinton11]

Add bitwise int ops to book + remove dead code

Haha I was about pushing for the book till I saw you merged it

[laggui]

Haha I was about pushing for the book till I saw you merged it

Figured for the last minor changes missed I'd do it myself instead of going through the request + change phase again 😄

Thanks again!

[quinton11]

Figured for the last minor changes missed I'd do it myself instead of going through the request + change phase again 😄

Fair Enough 😆

Thanks again!

My pleasure