Skip to content

Closes #5150: Try to fix bigint performance again #5151

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
1RyanK wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Closes #5150: Try to fix bigint performance again #5151

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
52 changes: 7 additions & 45 deletions arkouda/numpy/pdarraycreation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,6 @@
from arkouda.numpy._typing._typing import _NumericLikeDType, _StringDType
from arkouda.numpy.dtypes import (
NUMBER_FORMAT_STRINGS,
DTypes,
NumericDTypes,
SeriesDTypes,
bigint,
Expand Down Expand Up @@ -371,43 +370,9 @@ def array(
raise TypeError("a must be a pdarray, np.ndarray, or convertible to a numpy array")
if dtype is not None and dtype not in [bigint, "bigint"]:
a = a.astype(dtype)
if a.dtype == object and all(isinstance(x, (int, np.integer)) for x in a) and dtype is None:
if a.dtype == object and dtype is None and all(isinstance(x, (int, np.integer)) for x in a.flat):
dtype = bigint

# Special case, to get around error when putting negative numbers in a bigint

if (
not isinstance(a, Strings)
and not np.issubdtype(a.dtype, np.str_)
and dtype == bigint
and (a < 0).any()
):
neg_mask = array((a < 0) * -2 + 1)
abs_a = array(abs(a), dtype, max_bits=max_bits)
assert isinstance(abs_a, pdarray) # This is for mypy reasons
return neg_mask * abs_a

if a.dtype == bigint or a.dtype.name not in DTypes or dtype == bigint:
# We need this array whether the number of dimensions is 1 or greater.
uint_arrays: List[Union[pdarray, Strings]] = []

if (a.dtype == bigint or dtype == bigint) and a.ndim in get_array_ranks() and a.ndim > 1:
sh = a.shape
try:
# attempt to break bigint into multiple uint64 arrays
# early out if we would have more uint arrays than can fit in max_bits
early_out = (max_bits // 64) + (max_bits % 64 != 0) if max_bits != -1 else float("inf")
while (a != 0).any() and len(uint_arrays) < early_out:
low, a = a % 2**64, a // 2**64
uint_arrays.append(array(np.array(low, dtype=np.uint), dtype=akuint64))
# If uint_arrays is empty, this will create an empty ak array and reshape it.
if not uint_arrays:
return zeros(size=sh, dtype=bigint, max_bits=max_bits)
else:
return bigint_from_uint_arrays(uint_arrays[::-1], max_bits=max_bits)
except TypeError:
raise RuntimeError(f"Unhandled dtype {a.dtype}")

if a.ndim == 0 and a.dtype.name not in NumericDTypes:
raise TypeError("Must be an iterable or have a numeric DType")

Expand All @@ -416,9 +381,6 @@ def array(

if a.ndim not in get_array_ranks():
raise ValueError(f"array rank {a.ndim} not in compiled ranks {get_array_ranks()}")

if a.dtype == object and all(isinstance(x, (int, np.integer)) for x in a.ravel()) and dtype is None:
dtype = bigint
if a.dtype.kind in ("i", "u", "f", "O") and (dtype == bigint or dtype == "bigint"):
return _bigint_from_numpy(a, max_bits)
elif not (
Expand Down Expand Up @@ -501,7 +463,7 @@ def _bigint_from_numpy(np_a: np.ndarray, max_bits: int) -> pdarray:

a = np.asarray(np_a)
out_shape = a.shape
flat = a.ravel()
flat = a.ravel() # ndarray (view when possible)

# Empty → empty bigint (shape-preserving)
if a.size == 0:
Expand All @@ -512,10 +474,10 @@ def _bigint_from_numpy(np_a: np.ndarray, max_bits: int) -> pdarray:
raise TypeError(f"bigint requires numeric input, got dtype={a.dtype}")

if a.dtype.kind == "O":
# Only allow Python ints
if not all(isinstance(x, (int, float)) for x in flat):
# Only allow Python ints and floats
if not all(isinstance(x, (int, float)) for x in a.flat):
raise TypeError("bigint requires numeric input, got non-numeric object")
if any(isinstance(x, float) for x in flat):
if any(isinstance(x, float) for x in a.flat):
a = a.astype(np.float64, copy=False)

# Fast all-zero path or single limb path
Expand All @@ -542,7 +504,7 @@ def _bigint_from_numpy(np_a: np.ndarray, max_bits: int) -> pdarray:
)
)
else:
if all(int(x) == 0 for x in flat):
if all(int(x) == 0 for x in a.flat):
return zeros(size=a.shape, dtype=bigint, max_bits=max_bits)
any_neg = np.any(flat < 0)
req_bits: int
Expand All @@ -554,7 +516,7 @@ def _bigint_from_numpy(np_a: np.ndarray, max_bits: int) -> pdarray:
mask = (1 << 64) - 1
uint_arrays: List[Union[pdarray, Strings]] = []
# attempt to break bigint into multiple uint64 arrays
while (flat != 0).any() and len(uint_arrays) < req_limbs:
for _ in range(req_limbs):
low = flat & mask
flat = flat >> 64 # type: ignore
# low, flat = flat & mask, flat >> 64
Expand Down
2 changes: 1 addition & 1 deletion benchmarks/array_transfer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@ def time_ak_array_transfer(N, trials, dtype, seed, max_bits=-1):
end = time.time()
to_ndarray_times.append(end - start)
start = time.time()
aka = ak.array(npa, max_bits=max_bits)
aka = ak.array(npa, max_bits=max_bits, dtype=dtype)
end = time.time()
to_pdarray_times.append(end - start)
gc.collect()
Expand Down
59 changes: 33 additions & 26 deletions src/BigIntMsg.chpl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -169,7 +169,7 @@ module BigIntMsg {
st: borrowed SymTab,
type array_dtype,
param array_nd: int): MsgTuple throws
where ((array_dtype == int(64) || array_dtype == uint(64)) && array_nd == 1)
where (array_dtype == uint(64) && array_nd == 1)
{
var repMsg: string;

Expand All @@ -195,23 +195,19 @@ module BigIntMsg {
mask -= 1;
}

const dom = bigIntArray.domain;

if signed {
var i = num_arrays - 1;
const name = arrayNames[i];
const entry = st[name]: SymEntry(array_dtype, array_nd);
ref limbA = entry.a;

// Per-element fold; keep everything local
if doMask {
forall (u, b, s) in zip(limbA, bigIntArray, signs) {
s = (u >> 63): bool;
b += (u & ((1 << 63) - 1));
}
} else {
forall (u, b, s) in zip(limbA, bigIntArray, signs) {
s = (u >> 63): bool;
b += (u & ((1 << 63) - 1));
}
forall idx in dom {
const u = limbA[idx];
signs[idx] = (u >> 63): bool;
bigIntArray[idx] += (u & ((1 << 63) - 1));
}
}

Expand All @@ -225,17 +221,23 @@ module BigIntMsg {

// Per-element fold; keep everything local
if doMask {
forall (u, b) in zip(limbA, bigIntArray) with (var tmp: bigint, const localMask = mask) {
forall idx in dom with (var tmp: bigint, const ref localMask = mask) {
const u = limbA[idx]; // magnitude limb
ref b = bigIntArray[idx];

b <<= 64;
tmp = (u: uint(64)): bigint; // treat limb as magnitude
b += tmp;
b &= localMask;
tmp = (u:uint(64)): bigint;
b += tmp;
b &= localMask;
}
} else {
forall (u, b) in zip(limbA, bigIntArray) with (var tmp: bigint) {
forall idx in dom with (var tmp: bigint) {
const u = limbA[idx];
ref b = bigIntArray[idx];

b <<= 64;
tmp = (u: uint(64)): bigint;
b += tmp;
tmp = (u:uint(64)): bigint;
b += tmp;
}
}
}
Expand All @@ -246,16 +248,21 @@ module BigIntMsg {
// --------------------------------------------------------------------------

if signed {

const negAmount = (1: bigint) << (64 * num_arrays - 1);

forall (b, s) in zip(bigIntArray, signs) {
if s {
if doMask {
b = b - negAmount;
b &= mask;
} else {
b = b - negAmount;
if doMask {
forall idx in dom with (const ref localMask = mask) {
if signs[idx] {
ref b = bigIntArray[idx];
b -= negAmount;
b &= localMask;
}
}
} else {
forall idx in dom {
if signs[idx] {
ref b = bigIntArray[idx];
b -= negAmount;
}
}
}
Expand Down