heat.pow() speed-up when exponent is int

[ClaudiaComito]

Description

This PR exploits torch.pow(a,b)'s performance when b is an integer. For a torch/heat performance comparison (on 1 process) for a few binary operations, see table below.

The problem was identified by @coquelin77 in #789 and #793. This PR replaces #793 which had gotten quite stale and difficult to update. Thanks again @coquelin77 for all the work.

Issue/s resolved: #789

Changes proposed:

• In _operations.__binary_op In arithmetics.pow: introduce an early check for integer scalar terms, perform the torch operation immediately without further checks and return results.

Type of change

• Bug fix (non-breaking change which fixes an issue)

Memory requirements

NA

Performance

a = randn(10000, 10000) (float64 torch tensor or float64 DNDarray respectively)

operation	torch 1 proc	main 1 proc	789-pow 1 proc
a * a	0.14	0.14	0.14
a + a	0.14	0.14	0.14
a + 2	0.14	0.14	0.14
a / a	0.14	0.14	0.14
a ** a	1.16	1.16	1.16
a ** 2	0.14	1.3	0.14	<---
2 ** a	0.56	0.56	0.56
a ** 2.5	1.03	1.03	1.03

Due Diligence

• All split configurations tested
• Multiple dtypes tested in relevant functions
• Documentation updated (if needed)
• Title of PR is suitable for corresponding CHANGELOG entry

Does this change modify the behaviour of other functions? If so, which?

no

[ghost]

👇 Click on the image for a new way to code review

Legend

[github-actions]

Thank you for the PR!

[codecov]

Codecov Report

Merging #1141 (ed242bc) into main (eaa90d6) will increase coverage by 0.01%.
The diff coverage is 100.00%.

@@            Coverage Diff             @@
##             main    #1141      +/-   ##
==========================================
+ Coverage   91.77%   91.79%   +0.01%     
==========================================
  Files          72       72              
  Lines       10485    10497      +12     
==========================================
+ Hits         9623     9636      +13     
+ Misses        862      861       -1     

Flag	Coverage Δ
unit	91.79% <100.00%> (+0.01%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

Impacted Files	Coverage Δ
heat/core/arithmetics.py	99.11% <100.00%> (+0.51%)	⬆️

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more

[github-actions]

Thank you for the PR!

[mtar]

• Why are floats excluded? torch.pow supports them.
• Such an early exit would be a good fit for binary_op itself. All calling functions would benefit from the speedup.

[ClaudiaComito]

Thanks @mtar.

• Why are floats excluded? torch.pow supports them.

floats exponents will be dealt with in the general __binary_op(), torch.pow() optimization for integer exponents makes for a significant speedup, see table above.

* Such an early exit would be a good fit for `binary_op` itself. All calling functions would benefit from the speedup.

Turns out many binary operations in torch require both inputs to be tensors, so we need to wrap the ints into DNDarrays in those cases.

[github-actions]

Thank you for the PR!

[github-actions]

Thank you for the PR!

[github-actions]

Thank you for the PR!

[github-actions]

Thank you for the PR!

[mtar]

Fix broken link