Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets
This commit is contained in:
Kolp
@@ -0,0 +1,74 @@
|
||||
#################################################################################################
|
||||
#
|
||||
# Copyright (c) 2023 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
#
|
||||
# Redistribution and use in source and binary forms, with or without
|
||||
# modification, are permitted provided that the following conditions are met:
|
||||
#
|
||||
# 1. Redistributions of source code must retain the above copyright notice, this
|
||||
# list of conditions and the following disclaimer.
|
||||
#
|
||||
# 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
# this list of conditions and the following disclaimer in the documentation
|
||||
# and/or other materials provided with the distribution.
|
||||
#
|
||||
# 3. Neither the name of the copyright holder nor the names of its
|
||||
# contributors may be used to endorse or promote products derived from
|
||||
# this software without specific prior written permission.
|
||||
#
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
#
|
||||
#################################################################################################
|
||||
|
||||
from .int_tuple import (
|
||||
as_tuple,
|
||||
crd2crd,
|
||||
crd2idx,
|
||||
elem_scale,
|
||||
flatten,
|
||||
has_none,
|
||||
idx2crd,
|
||||
inner_product,
|
||||
IntTuple,
|
||||
is_int,
|
||||
is_tuple,
|
||||
match_structure,
|
||||
product,
|
||||
shape_div,
|
||||
signum,
|
||||
slice_,
|
||||
suffix_product,
|
||||
tuple_max,
|
||||
)
|
||||
from .layout import (
|
||||
coalesce,
|
||||
complement,
|
||||
composition,
|
||||
cosize,
|
||||
filter,
|
||||
is_layout,
|
||||
Layout,
|
||||
LayoutBase,
|
||||
left_inverse,
|
||||
logical_divide,
|
||||
logical_product,
|
||||
make_layout,
|
||||
right_inverse,
|
||||
size,
|
||||
slice_and_offset,
|
||||
tiled_divide,
|
||||
tiled_product,
|
||||
zipped_divide,
|
||||
zipped_product,
|
||||
)
|
||||
from .typing import Integer
|
||||
@@ -0,0 +1,285 @@
|
||||
#################################################################################################
|
||||
#
|
||||
# Copyright (c) 2023 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
#
|
||||
# Redistribution and use in source and binary forms, with or without
|
||||
# modification, are permitted provided that the following conditions are met:
|
||||
#
|
||||
# 1. Redistributions of source code must retain the above copyright notice, this
|
||||
# list of conditions and the following disclaimer.
|
||||
#
|
||||
# 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
# this list of conditions and the following disclaimer in the documentation
|
||||
# and/or other materials provided with the distribution.
|
||||
#
|
||||
# 3. Neither the name of the copyright holder nor the names of its
|
||||
# contributors may be used to endorse or promote products derived from
|
||||
# this software without specific prior written permission.
|
||||
#
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
#
|
||||
#################################################################################################
|
||||
|
||||
"""
|
||||
Functions for manipulating IntTuples
|
||||
"""
|
||||
|
||||
from functools import reduce
|
||||
from itertools import chain
|
||||
from typing import TypeAlias
|
||||
from typing_extensions import TypeIs
|
||||
|
||||
from .typing import Integer
|
||||
|
||||
|
||||
# Type aliases for better readability
|
||||
IntTuple: TypeAlias = int | tuple["IntTuple", ...]
|
||||
|
||||
|
||||
def is_int(x: object) -> TypeIs[int]:
|
||||
return isinstance(x, Integer)
|
||||
|
||||
|
||||
def is_tuple(x: object) -> TypeIs[tuple]:
|
||||
return isinstance(x, tuple)
|
||||
|
||||
|
||||
def as_tuple(x: IntTuple) -> tuple[IntTuple, ...]:
|
||||
if is_int(x):
|
||||
return (x,)
|
||||
return x
|
||||
|
||||
|
||||
def match_structure(a: IntTuple, b: IntTuple) -> bool:
|
||||
if is_int(a) and is_int(b):
|
||||
return True
|
||||
if is_tuple(a) and is_tuple(b):
|
||||
return len(a) == len(b) and all(match_structure(x, y) for x, y in zip(a, b))
|
||||
return False
|
||||
|
||||
|
||||
def flatten(t: IntTuple) -> tuple[int, ...]:
|
||||
if is_tuple(t):
|
||||
if len(t) == 0:
|
||||
return ()
|
||||
else:
|
||||
return tuple(i for a in t for i in flatten(a))
|
||||
else:
|
||||
return (t,)
|
||||
|
||||
|
||||
def signum(a: int) -> int:
|
||||
return bool(a > 0) - bool(a < 0)
|
||||
|
||||
|
||||
def product(a: IntTuple) -> int:
|
||||
if is_tuple(a):
|
||||
return reduce(lambda val, elem: val * product(elem), a, 1)
|
||||
else:
|
||||
return a
|
||||
|
||||
|
||||
def inner_product(a: IntTuple, b: IntTuple) -> int:
|
||||
if is_tuple(a) and is_tuple(b): # tuple tuple
|
||||
if len(a) != len(b):
|
||||
raise AssertionError
|
||||
return sum(inner_product(x, y) for x, y in zip(a, b))
|
||||
else: # "int" "int"
|
||||
if is_tuple(a) or is_tuple(b):
|
||||
raise AssertionError
|
||||
return a * b
|
||||
|
||||
|
||||
def tuple_max(a: IntTuple) -> int:
|
||||
if is_tuple(a):
|
||||
return max(tuple_max(x) for x in a)
|
||||
else:
|
||||
return a
|
||||
|
||||
|
||||
def elem_scale(a: IntTuple, b: IntTuple) -> IntTuple:
|
||||
if is_tuple(a):
|
||||
if is_tuple(b): # tuple tuple
|
||||
if len(a) != len(b):
|
||||
raise AssertionError
|
||||
return tuple(elem_scale(x, y) for x, y in zip(a, b))
|
||||
else: # tuple "int"
|
||||
raise AssertionError("Invalid combination: tuple with int")
|
||||
else:
|
||||
if is_tuple(b): # "int" tuple
|
||||
return elem_scale(a, product(b))
|
||||
else: # "int" "int"
|
||||
return a * b
|
||||
|
||||
|
||||
# Inclusive prefix ceil div with output congruent to input a
|
||||
def shape_div(a: IntTuple, b: IntTuple) -> IntTuple:
|
||||
if is_tuple(a):
|
||||
if is_tuple(b): # tuple tuple
|
||||
if len(a) != len(b):
|
||||
raise AssertionError
|
||||
return tuple(shape_div(x, y) for x, y in zip(a, b))
|
||||
else: # tuple "int"
|
||||
# r = [shape_div(a[0],b)] + [shape_div(a[i],b := shape_div(b, product(a[i-1]))) for i in range(1,len(a))]
|
||||
r = []
|
||||
for v in a:
|
||||
r.append(shape_div(v, b))
|
||||
b = shape_div(b, product(v))
|
||||
return tuple(r)
|
||||
else:
|
||||
if is_tuple(b): # "int" tuple
|
||||
return shape_div(a, product(b))
|
||||
else: # "int" "int"
|
||||
if not (a % b == 0 or b % a == 0):
|
||||
raise AssertionError
|
||||
return (a + b - 1) // b
|
||||
|
||||
|
||||
# Exclusive suffix product with output congruent to input a (lexicographic)
|
||||
def suffix_product(a: IntTuple, init: IntTuple = 1) -> IntTuple:
|
||||
# TODO: With all these length asserts, may want to create a zip_strict wrapper.
|
||||
if is_tuple(a):
|
||||
if is_tuple(init): # tuple tuple
|
||||
if len(a) != len(init):
|
||||
raise AssertionError
|
||||
return tuple(suffix_product(x, i) for x, i in zip(a, init))
|
||||
else: # tuple "int"
|
||||
# Process from right to left for lexicographic ordering
|
||||
# r = [prefix_product(a[len(a)-1],init)] +
|
||||
# [prefix_product(a[i],init := init * product(a[i+1])) for i in range(len(a)-1,0)].reverse()
|
||||
r = []
|
||||
|
||||
# Calculate products from right to left, appending to list
|
||||
for i in range(len(a) - 1, -1, -1):
|
||||
r.append(suffix_product(a[i], init))
|
||||
init = init * product(a[i])
|
||||
|
||||
# Reverse to get correct lexicographic order
|
||||
r.reverse()
|
||||
return tuple(r)
|
||||
else:
|
||||
if is_tuple(init): # "int" tuple
|
||||
raise AssertionError("Invalid combination: int with tuple init")
|
||||
else: # "int" "int"
|
||||
return init
|
||||
|
||||
|
||||
def idx2crd(idx: IntTuple, shape: IntTuple, stride: IntTuple | None = None) -> IntTuple:
|
||||
if stride is None:
|
||||
stride = suffix_product(shape)
|
||||
|
||||
if is_tuple(idx):
|
||||
if is_tuple(shape) and is_tuple(stride): # tuple tuple tuple
|
||||
if not (len(idx) == len(shape) and len(stride) == len(shape)):
|
||||
raise AssertionError
|
||||
return tuple(idx2crd(i, s, d) for i, s, d in zip(idx, shape, stride))
|
||||
else: # tuple "int" "int"
|
||||
raise AssertionError("Invalid combination: tuple with int stride")
|
||||
else:
|
||||
if is_tuple(shape) and is_tuple(stride): # "int" tuple tuple
|
||||
if len(shape) != len(stride):
|
||||
raise AssertionError
|
||||
return tuple(idx2crd(idx, s, d) for s, d in zip(shape, stride))
|
||||
else: # "int" "int" "int"
|
||||
if is_tuple(shape) or is_tuple(stride):
|
||||
raise AssertionError
|
||||
return (idx // stride) % shape # all are ints after type checks
|
||||
|
||||
|
||||
def crd2idx(
|
||||
crd: IntTuple | None, shape: IntTuple, stride: IntTuple | None = None
|
||||
) -> int:
|
||||
if stride is None:
|
||||
stride = suffix_product(shape)
|
||||
|
||||
if is_tuple(crd):
|
||||
if is_tuple(shape) and is_tuple(stride): # tuple tuple tuple
|
||||
if not (len(crd) == len(shape) and len(stride) == len(shape)):
|
||||
raise AssertionError
|
||||
return sum(crd2idx(c, s, d) for c, s, d in zip(crd, shape, stride))
|
||||
else: # tuple "int" "int"
|
||||
raise AssertionError(f"Invalid combination: crd={crd}, shape={shape}")
|
||||
else:
|
||||
if crd is None:
|
||||
crd = 0
|
||||
|
||||
if is_tuple(shape) and is_tuple(stride): # "int" tuple tuple
|
||||
if len(shape) != len(stride):
|
||||
raise AssertionError
|
||||
result = 0
|
||||
# Process from right to left for lexicographic ordering
|
||||
for i in range(len(shape) - 1, 0, -1):
|
||||
result += crd2idx(crd % product(shape[i]), shape[i], stride[i])
|
||||
crd = crd // product(shape[i])
|
||||
if len(shape) > 0:
|
||||
result += crd2idx(crd, shape[0], stride[0])
|
||||
return result
|
||||
else: # "int" "int" "int"
|
||||
if is_tuple(shape) or is_tuple(stride):
|
||||
raise AssertionError
|
||||
return crd * stride # all are ints after type checks
|
||||
|
||||
|
||||
# Transform crd into the dst_shape's iteration space
|
||||
def crd2crd(
|
||||
crd: IntTuple, dst_shape: IntTuple, src_shape: IntTuple | None = None
|
||||
) -> IntTuple:
|
||||
if is_tuple(crd):
|
||||
if is_tuple(dst_shape): # tuple tuple
|
||||
if len(crd) != len(dst_shape):
|
||||
raise AssertionError
|
||||
return tuple(crd2crd(x, y) for x, y in zip(crd, dst_shape))
|
||||
else: # tuple "int"
|
||||
# Ambiguous unless we have src_shape
|
||||
if src_shape is None:
|
||||
raise AssertionError
|
||||
return crd2idx(crd, src_shape)
|
||||
else:
|
||||
if is_tuple(dst_shape): # "int" tuple
|
||||
return idx2crd(crd, dst_shape)
|
||||
else: # "int" "int"
|
||||
if crd >= dst_shape:
|
||||
raise AssertionError
|
||||
return crd
|
||||
|
||||
|
||||
# Filter trg according to crd: keep only elements of trg that are paired with None
|
||||
def slice_(crd: tuple | int | None, trg: tuple | int) -> tuple | int:
|
||||
if is_tuple(crd):
|
||||
if is_tuple(trg): # tuple tuple
|
||||
if len(crd) != len(trg):
|
||||
raise AssertionError
|
||||
# match C++ behavior of `filter_tuple` using `tuple_cat(...)`
|
||||
return tuple(
|
||||
chain(
|
||||
*filter( # type: ignore[arg-type] # filter returns Iterator which is compatible
|
||||
lambda x: x != (),
|
||||
[slice_(c, s) for c, s in zip(crd, trg)],
|
||||
)
|
||||
)
|
||||
)
|
||||
else:
|
||||
raise AssertionError("Invalid combination: tuple crd with int trg")
|
||||
elif crd is None:
|
||||
# match C++ behavior `return cute::tuple<B>{b};`
|
||||
return (trg,)
|
||||
else:
|
||||
return ()
|
||||
|
||||
|
||||
# Determine if None appears at any of an int_tuples' terminals
|
||||
def has_none(a: tuple | int | None) -> bool:
|
||||
if is_tuple(a):
|
||||
return any(has_none(v) for v in a)
|
||||
else:
|
||||
return a is None
|
||||
@@ -0,0 +1,484 @@
|
||||
#################################################################################################
|
||||
#
|
||||
# Copyright (c) 2023 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
#
|
||||
# Redistribution and use in source and binary forms, with or without
|
||||
# modification, are permitted provided that the following conditions are met:
|
||||
#
|
||||
# 1. Redistributions of source code must retain the above copyright notice, this
|
||||
# list of conditions and the following disclaimer.
|
||||
#
|
||||
# 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
# this list of conditions and the following disclaimer in the documentation
|
||||
# and/or other materials provided with the distribution.
|
||||
#
|
||||
# 3. Neither the name of the copyright holder nor the names of its
|
||||
# contributors may be used to endorse or promote products derived from
|
||||
# this software without specific prior written permission.
|
||||
#
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
#
|
||||
#################################################################################################
|
||||
|
||||
"""
|
||||
Definition of CuTe Layouts and functions to manipulate them which works with the order
|
||||
of lexicographic instead of co-lexicographic as implemented in the original layout.py
|
||||
"""
|
||||
|
||||
from itertools import chain
|
||||
from typing import TypeAlias
|
||||
from typing_extensions import Self, TypeIs
|
||||
|
||||
from .int_tuple import (
|
||||
crd2idx,
|
||||
flatten,
|
||||
has_none,
|
||||
IntTuple,
|
||||
is_int,
|
||||
is_tuple,
|
||||
product,
|
||||
slice_,
|
||||
suffix_product,
|
||||
)
|
||||
|
||||
|
||||
# Type aliases
|
||||
CoordinateType: TypeAlias = (
|
||||
int | IntTuple | tuple[object, ...] | None
|
||||
) # Input for slice_ and crd2idx functions
|
||||
|
||||
|
||||
class LayoutBase:
|
||||
pass
|
||||
|
||||
|
||||
def is_layout(x: object) -> TypeIs["Layout"]:
|
||||
return isinstance(x, LayoutBase)
|
||||
|
||||
|
||||
class Layout(LayoutBase):
|
||||
def __init__(self, _shape: IntTuple, _stride: IntTuple | None = None) -> None:
|
||||
self.shape = _shape
|
||||
if _stride is None:
|
||||
self.stride = suffix_product(self.shape)
|
||||
else:
|
||||
self.stride = _stride
|
||||
|
||||
# operator ==
|
||||
def __eq__(self, other: object) -> bool:
|
||||
if not isinstance(other, Layout):
|
||||
return False
|
||||
return self.shape == other.shape and self.stride == other.stride
|
||||
|
||||
# operator len(L) (len [rank] like tuples)
|
||||
def __len__(self) -> int:
|
||||
if is_tuple(self.shape):
|
||||
return len(self.shape)
|
||||
else:
|
||||
return 1
|
||||
|
||||
# operator () (map coord to idx)
|
||||
def __call__(self, *args: CoordinateType) -> Self | int:
|
||||
"""
|
||||
Map a logical coordinate to a linear index (Coord has no Underscore slice operators)
|
||||
OR
|
||||
Slice the layout and return the sublayout (Coord has an Underscore slice op)
|
||||
|
||||
Follow the same behavior of `Layout::operator(Coord const&)` in cute C++
|
||||
"""
|
||||
if has_none(args):
|
||||
if len(args) == 1:
|
||||
return Layout(slice_(args[0], self.shape), slice_(args[0], self.stride))
|
||||
else:
|
||||
return Layout(slice_(args, self.shape), slice_(args, self.stride))
|
||||
else:
|
||||
if len(args) == 1:
|
||||
return crd2idx(args[0], self.shape, self.stride) # type: ignore[arg-type]
|
||||
else:
|
||||
return crd2idx(args, self.shape, self.stride) # type: ignore[arg-type]
|
||||
|
||||
# operator [] (get-i like tuples)
|
||||
def __getitem__(self, i: int) -> Self:
|
||||
if is_tuple(self.shape):
|
||||
return Layout(self.shape[i], self.stride[i]) # type: ignore[index]
|
||||
else:
|
||||
if i != 0:
|
||||
raise AssertionError
|
||||
return Layout(self.shape, self.stride)
|
||||
|
||||
# size(layout) Size of the domain
|
||||
def size(self) -> int:
|
||||
return product(self.shape)
|
||||
|
||||
# cosize(layout) Size of the codomain
|
||||
def cosize(self) -> int:
|
||||
return self(self.size() - 1) + 1 # type: ignore[operator]
|
||||
|
||||
# print and str
|
||||
def __str__(self) -> str:
|
||||
return f"{self.shape}:{self.stride}"
|
||||
|
||||
# error msgs and representation
|
||||
def __repr__(self) -> str:
|
||||
return f"Layout({self.shape},{self.stride})"
|
||||
|
||||
|
||||
# Type aliases
|
||||
LayoutOrIntTuple: TypeAlias = Layout | IntTuple
|
||||
LayoutProfile: TypeAlias = tuple[object, ...] | Layout | None
|
||||
LayoutInput: TypeAlias = Layout | IntTuple | tuple[object, ...] | None
|
||||
|
||||
|
||||
# Make Layout from a list of layouts (each layout it's own mode in the result)
|
||||
def make_layout(*layouts: Layout | tuple[Layout, ...]) -> Layout:
|
||||
if len(layouts) == 1 and not is_layout(layouts[0]):
|
||||
layouts = layouts[0]
|
||||
|
||||
shape, stride = zip(*((a.shape, a.stride) for a in layouts)) # type: ignore[union-attr]
|
||||
return Layout(shape, stride)
|
||||
|
||||
|
||||
# Size of the domain
|
||||
def size(layout: LayoutOrIntTuple) -> int:
|
||||
if is_layout(layout):
|
||||
return layout.size()
|
||||
return product(layout)
|
||||
|
||||
|
||||
# Size of the codomain
|
||||
def cosize(layout: Layout) -> int:
|
||||
return layout.cosize()
|
||||
|
||||
|
||||
# Layout coalesce -- flatten and combine as many modes as possible while preserving the int-to-int function
|
||||
def coalesce(layout: Layout, profile: LayoutProfile = None) -> Layout:
|
||||
if is_tuple(profile):
|
||||
if len(layout) < len(profile):
|
||||
raise AssertionError
|
||||
return make_layout(
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
chain(
|
||||
(coalesce(layout[i], profile[i]) for i in range(len(profile))), # type: ignore[arg-type]
|
||||
(layout[i] for i in range(len(profile), len(layout))),
|
||||
)
|
||||
)
|
||||
|
||||
result_shape = [1]
|
||||
result_stride = [0]
|
||||
# Since we now follow lexicographic order, we need to process from right to left.
|
||||
# And to make implementation more efficient, we append to the end of list and reverse it in the end.
|
||||
for shape, stride in zip(
|
||||
reversed(flatten(layout.shape)), reversed(flatten(layout.stride))
|
||||
):
|
||||
# skip their shape-1s
|
||||
if shape == 1:
|
||||
continue
|
||||
# replace our shape-1 with anything
|
||||
elif result_shape[-1] == 1:
|
||||
result_shape[-1] = shape
|
||||
result_stride[-1] = stride
|
||||
# merge modes if the shape*stride match
|
||||
elif result_shape[-1] * result_stride[-1] == stride:
|
||||
result_shape[-1] = result_shape[-1] * shape
|
||||
# append a new mode
|
||||
else:
|
||||
result_shape.append(shape)
|
||||
result_stride.append(stride)
|
||||
|
||||
if len(result_shape) == 1:
|
||||
return Layout(result_shape[0], result_stride[0])
|
||||
else:
|
||||
result_shape.reverse()
|
||||
result_stride.reverse()
|
||||
return Layout(tuple(result_shape), tuple(result_stride))
|
||||
|
||||
|
||||
# Layout filter -- replace all stride-0 modes with size-1 and then coalesce to remove them
|
||||
def filter(layout: Layout, profile: LayoutProfile = None) -> Layout:
|
||||
if is_tuple(profile):
|
||||
if len(layout) < len(profile):
|
||||
raise AssertionError
|
||||
return make_layout(
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
chain(
|
||||
(filter(layout[i], profile[i]) for i in range(len(profile))), # type: ignore[arg-type]
|
||||
(layout[i] for i in range(len(profile), len(layout))),
|
||||
)
|
||||
)
|
||||
|
||||
result_shape = []
|
||||
result_stride = []
|
||||
for shape, stride in zip(flatten(layout.shape), flatten(layout.stride)):
|
||||
# skip their shape-1s and stride-0s
|
||||
if not (shape == 1 or stride == 0):
|
||||
result_shape.append(shape)
|
||||
result_stride.append(stride)
|
||||
|
||||
if len(result_shape) == 0:
|
||||
return Layout(1, 0)
|
||||
else:
|
||||
return coalesce(Layout(tuple(result_shape), tuple(result_stride)))
|
||||
|
||||
|
||||
# Layout composition
|
||||
# Use tuples-of-layouts to perform this operation by-mode and None as no-op
|
||||
def composition(layoutA: Layout, layoutB: LayoutInput) -> Layout:
|
||||
if layoutB is None:
|
||||
return layoutA
|
||||
elif is_int(layoutB):
|
||||
return composition(layoutA, Layout(layoutB))
|
||||
elif is_tuple(layoutB):
|
||||
if len(layoutA) < len(layoutB):
|
||||
raise AssertionError
|
||||
return make_layout(
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
chain(
|
||||
(composition(layoutA[i], layoutB[i]) for i in range(len(layoutB))), # type: ignore[arg-type]
|
||||
(layoutA[i] for i in range(len(layoutB), len(layoutA))),
|
||||
)
|
||||
)
|
||||
elif is_tuple(layoutB.shape):
|
||||
return make_layout(composition(layoutA, layoutB_i) for layoutB_i in layoutB) # type: ignore[arg-type, attr-defined]
|
||||
|
||||
if layoutB.stride == 0:
|
||||
return Layout(layoutB.shape, 0)
|
||||
else:
|
||||
result_shape = []
|
||||
result_stride = []
|
||||
rest_shape = layoutB.shape
|
||||
rest_stride = layoutB.stride
|
||||
flat_A = coalesce(layoutA)
|
||||
# when left layout is multi-dimensional sublayout, aka, self = (a,b,...,c):(x,y,...,z), layout = s:d,
|
||||
# for integral s and d means that we want:
|
||||
# (1) “remove” the first d elements from left, starting from rightmost. (This will increase the stride.)
|
||||
# (2) “keep” the first s of those strided elements. (This does not affect the stride.)
|
||||
# For example, if self = (6,2):(2,1), layout = (3:2)
|
||||
# Step 1: remove the first 2 elements from self with stride increase, i.e., (6,2):(2,1) -> (6,1):(2,2)
|
||||
# Step 2: keep the first 3 of those strided elements, i.e., (6,1):(2,2) -> (3,1):(2,2)
|
||||
# Because we are going lexicographically, we go through left layout from right to left.
|
||||
for curr_shape, curr_stride in zip(
|
||||
reversed(flatten(flat_A.shape)[1:]), reversed(flatten(flat_A.stride)[1:])
|
||||
):
|
||||
if not (curr_shape % rest_stride == 0 or rest_stride % curr_shape == 0): # type: ignore[operator]
|
||||
raise AssertionError
|
||||
new_shape = min(max(1, curr_shape // rest_stride), rest_shape) # type: ignore[operator]
|
||||
|
||||
if new_shape != 1:
|
||||
result_shape.append(new_shape) # Append to end, will reverse later
|
||||
result_stride.append(rest_stride * curr_stride)
|
||||
|
||||
rest_shape = rest_shape // new_shape # type: ignore[operator]
|
||||
rest_stride = -(
|
||||
-rest_stride // curr_shape # type: ignore[operator]
|
||||
) # Python exclusive impl: "//" is always floor div so == ceil_div(abs(rest_stride), curr_shape) * signum(rest_stride)
|
||||
|
||||
# When left has single-size sublayout or reach the last sublayout, aka, left = a:b, layout = s:d,
|
||||
# the result is rather trivial: left o layout = a:b o s:d = s:(b*d).
|
||||
# For example, if self = (6:2), layout = (3:2), the result is (3:(2*2)) = (3:4).
|
||||
if rest_shape != 1 or len(result_shape) == 0:
|
||||
result_shape.append(rest_shape) # Append to end, will reverse later
|
||||
result_stride.append(rest_stride * flatten(flat_A.stride)[0])
|
||||
|
||||
# Reverse the lists because we build lists in reverse order (append to end), this way it is more efficient.
|
||||
result_shape.reverse()
|
||||
result_stride.reverse()
|
||||
|
||||
if len(result_shape) == 1:
|
||||
return Layout(result_shape[0], result_stride[0]) # type: ignore[arg-type]
|
||||
else:
|
||||
return Layout(tuple(result_shape), tuple(result_stride)) # type: ignore[arg-type]
|
||||
|
||||
|
||||
# Layout complement
|
||||
def complement(layout: LayoutOrIntTuple, max_idx: int = 1) -> Layout:
|
||||
if is_int(layout):
|
||||
return complement(Layout(layout))
|
||||
|
||||
result_shape = []
|
||||
result_stride = []
|
||||
current_idx = 1
|
||||
|
||||
sorted_DS = sorted(zip(flatten(layout.stride), flatten(layout.shape))) # type: ignore[union-attr]
|
||||
for stride, shape in sorted_DS:
|
||||
if stride == 0 or shape == 1:
|
||||
continue
|
||||
|
||||
in_bound = current_idx <= shape * stride
|
||||
# To support symbolic value which can't be evaluated now
|
||||
if (type(in_bound) is bool) and not in_bound:
|
||||
raise AssertionError
|
||||
|
||||
result_shape.append(stride // current_idx)
|
||||
result_stride.append(current_idx)
|
||||
current_idx = shape * stride
|
||||
|
||||
result_shape.append((max_idx + current_idx - 1) // current_idx) # ceil_div
|
||||
result_stride.append(current_idx)
|
||||
# This is different from original pycute implementation, because we want to follow the lexicographic order here
|
||||
# where the right-most dimension is the innermost dimension (smallest stride).
|
||||
result_shape.reverse()
|
||||
result_stride.reverse()
|
||||
|
||||
return coalesce(Layout(tuple(result_shape), tuple(result_stride)))
|
||||
|
||||
|
||||
# Layout right inverse
|
||||
def right_inverse(layout: LayoutOrIntTuple | None) -> Layout | None:
|
||||
if layout is None:
|
||||
return None
|
||||
elif is_int(layout):
|
||||
return Layout(layout)
|
||||
|
||||
result_shape = []
|
||||
result_stride = []
|
||||
current_idx = 1
|
||||
|
||||
flat_shape = flatten(layout.shape) # type: ignore[union-attr]
|
||||
flat_stride = flatten(layout.stride) # type: ignore[union-attr]
|
||||
sorted_DSA = sorted(zip(flat_stride, flat_shape, suffix_product(flat_shape))) # type: ignore[arg-type]
|
||||
for stride, shape, rstride in sorted_DSA:
|
||||
if shape == 1:
|
||||
continue
|
||||
if current_idx != stride:
|
||||
break
|
||||
|
||||
result_shape.append(shape)
|
||||
result_stride.append(rstride)
|
||||
current_idx = shape * stride
|
||||
|
||||
result_shape.reverse()
|
||||
result_stride.reverse()
|
||||
return coalesce(Layout(tuple(result_shape), tuple(result_stride)))
|
||||
|
||||
|
||||
# Layout left inverse
|
||||
def left_inverse(layout: LayoutOrIntTuple | None) -> Layout | None:
|
||||
if layout is None:
|
||||
return None
|
||||
elif is_int(layout):
|
||||
return Layout(layout)
|
||||
return right_inverse(make_layout(complement(layout), layout)) # type: ignore[arg-type]
|
||||
|
||||
|
||||
# Split a layout by the composition of B and the "rest"
|
||||
# Use tuples-of-layouts to perform this operation by-mode and None as no-op
|
||||
def logical_divide(layoutA: Layout, layoutB: LayoutInput) -> Layout:
|
||||
if layoutB is None:
|
||||
return layoutA
|
||||
elif is_int(layoutB):
|
||||
return logical_divide(layoutA, Layout(layoutB))
|
||||
elif is_tuple(layoutB):
|
||||
if len(layoutA) < len(layoutB):
|
||||
raise AssertionError
|
||||
return make_layout(
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
chain(
|
||||
(
|
||||
logical_divide(layoutA[i], layoutB[i]) # type: ignore[arg-type]
|
||||
for i in range(len(layoutB))
|
||||
),
|
||||
(layoutA[i] for i in range(len(layoutB), len(layoutA))),
|
||||
)
|
||||
)
|
||||
|
||||
return composition(
|
||||
layoutA,
|
||||
make_layout(layoutB, complement(layoutB, size(layoutA))),
|
||||
)
|
||||
|
||||
|
||||
# Reproduce a layoutA over a layoutB
|
||||
# Use tuples-of-layouts to perform this operation by-mode and None as no-op
|
||||
def logical_product(layoutA: Layout, layoutB: LayoutInput) -> Layout:
|
||||
if layoutB is None:
|
||||
return layoutA
|
||||
elif is_int(layoutB):
|
||||
return logical_divide(layoutA, Layout(layoutB))
|
||||
elif is_tuple(layoutB):
|
||||
if len(layoutA) < len(layoutB):
|
||||
raise AssertionError
|
||||
return make_layout(
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
chain(
|
||||
(
|
||||
logical_product(layoutA[i], layoutB[i]) # type: ignore[arg-type]
|
||||
for i in range(len(layoutB))
|
||||
),
|
||||
(layoutA[i] for i in range(len(layoutB), len(layoutA))),
|
||||
)
|
||||
)
|
||||
|
||||
return make_layout(
|
||||
layoutA,
|
||||
composition(complement(layoutA, size(layoutA) * cosize(layoutB)), layoutB),
|
||||
)
|
||||
|
||||
|
||||
# Gather the modes from a hierarchical logical_divide or logical_product
|
||||
def hier_unzip(
|
||||
splitter: object,
|
||||
layoutA: Layout,
|
||||
layoutB: LayoutInput,
|
||||
) -> Layout:
|
||||
if layoutB is None:
|
||||
return make_layout(Layout(1, 0), layoutA)
|
||||
elif is_tuple(layoutB):
|
||||
if len(layoutA) < len(layoutB):
|
||||
raise AssertionError
|
||||
# A layout with shape ((A,a),(B,b),(C,c))
|
||||
split = make_layout(
|
||||
hier_unzip(splitter, layoutA[i], layoutB[i]) # type: ignore[arg-type]
|
||||
for i in range(len(layoutB))
|
||||
)
|
||||
# Gather to shape ((A,B,C,...),(a,b,c,...,y,z))
|
||||
return make_layout(
|
||||
make_layout(split[i][0] for i in range(len(layoutB))), # type: ignore[arg-type]
|
||||
make_layout(
|
||||
chain( # type: ignore[arg-type]
|
||||
(split[i][1] for i in range(len(layoutB))),
|
||||
(layoutA[i] for i in range(len(layoutB), len(layoutA))),
|
||||
)
|
||||
),
|
||||
)
|
||||
|
||||
# splitter must return a rank-2 layout
|
||||
return splitter(layoutA, layoutB) # type: ignore[operator]
|
||||
|
||||
|
||||
# Apply logical divide hierarchically and gather the split modes into two modes
|
||||
def zipped_divide(layoutA: Layout, layoutB: LayoutInput) -> Layout:
|
||||
return hier_unzip(logical_divide, layoutA, layoutB)
|
||||
|
||||
|
||||
# Perform logical divide hierarchically and gather tiles (B-layouts) into a new mode
|
||||
def tiled_divide(layoutA: Layout, layoutB: LayoutInput) -> Layout:
|
||||
result = zipped_divide(layoutA, layoutB)
|
||||
return make_layout([result[0]] + [result[1][i] for i in range(len(result[1]))]) # type: ignore[arg-type]
|
||||
|
||||
|
||||
# Apply logical product hierarchically and gather the split modes into two modes
|
||||
def zipped_product(layoutA: Layout, layoutB: LayoutInput) -> Layout:
|
||||
return hier_unzip(logical_product, layoutA, layoutB)
|
||||
|
||||
|
||||
# Perform logical product hierarchically and gather tiles (B-layouts) into a new mode
|
||||
def tiled_product(layoutA: Layout, layoutB: LayoutInput) -> Layout:
|
||||
result = zipped_product(layoutA, layoutB)
|
||||
return make_layout([result[0]] + [result[1][i] for i in range(len(result[1]))]) # type: ignore[arg-type]
|
||||
|
||||
|
||||
def slice_and_offset(crd: tuple[object, ...], layout: Layout) -> tuple[Layout, int]:
|
||||
return (
|
||||
Layout(slice_(crd, layout.shape), slice_(crd, layout.stride)),
|
||||
crd2idx(crd, layout.shape, layout.stride), # type: ignore[arg-type]
|
||||
)
|
||||
@@ -0,0 +1,42 @@
|
||||
#################################################################################################
|
||||
#
|
||||
# Copyright (c) 2023 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
#
|
||||
# Redistribution and use in source and binary forms, with or without
|
||||
# modification, are permitted provided that the following conditions are met:
|
||||
#
|
||||
# 1. Redistributions of source code must retain the above copyright notice, this
|
||||
# list of conditions and the following disclaimer.
|
||||
#
|
||||
# 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
# this list of conditions and the following disclaimer in the documentation
|
||||
# and/or other materials provided with the distribution.
|
||||
#
|
||||
# 3. Neither the name of the copyright holder nor the names of its
|
||||
# contributors may be used to endorse or promote products derived from
|
||||
# this software without specific prior written permission.
|
||||
#
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
#
|
||||
#################################################################################################
|
||||
|
||||
from abc import ABC
|
||||
|
||||
|
||||
class Integer(ABC): # noqa: B024 # Uses __subclasshook__ instead of abstract methods
|
||||
@classmethod
|
||||
def __subclasshook__(cls, c: type) -> bool:
|
||||
if c in [bool, float]:
|
||||
return False
|
||||
|
||||
return issubclass(c, int)
|
||||
Reference in New Issue
Block a user