Spaces:
Sleeping
Sleeping
File size: 4,604 Bytes
e45d058 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 |
#################################################################################################
#
# Copyright (c) 2023 - 2024 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.
#
#################################################################################################
"""
Methods for layout swizzling
"""
from .layout import *
def shiftr(a, s):
return a >> s if s > 0 else shiftl(a, -s)
def shiftl(a, s):
return a << s if s > 0 else shiftr(a, -s)
## A generic Swizzle functor
# 0bxxxxxxxxxxxxxxxYYYxxxxxxxZZZxxxx
# ^--^ Base is the number of least-sig bits to keep constant
# ^-^ ^-^ Bits is the number of bits in the mask
# ^---------^ Shift is the distance to shift the YYY mask
# (pos shifts YYY to the right, neg shifts YYY to the left)
#
# e.g. Given
# 0bxxxxxxxxxxxxxxxxYYxxxxxxxxxZZxxx
# the result is
# 0bxxxxxxxxxxxxxxxxYYxxxxxxxxxAAxxx where AA = ZZ xor YY
#
class Swizzle:
def __init__(self, bits, base, shift):
assert bits >= 0
assert base >= 0
assert abs(shift) >= bits
self.bits = bits
self.base = base
self.shift = shift
bit_msk = (1 << bits) - 1
self.yyy_msk = bit_msk << (base + max(0,shift))
self.zzz_msk = bit_msk << (base - min(0,shift))
# operator () (transform integer)
def __call__(self, offset):
return offset ^ shiftr(offset & self.yyy_msk, self.shift)
# Size of the domain
def size(self):
return 1 << (bits + base + abs(shift))
# Size of the codomain
def cosize(self):
return self.size()
# print and str
def __str__(self):
return f"SW_{self.bits}_{self.base}_{self.shift}"
# error msgs and representation
def __repr__(self):
return f"Swizzle({self.bits},{self.base},{self.shift})"
class ComposedLayout(LayoutBase):
def __init__(self, layoutB, offset, layoutA):
self.layoutB = layoutB
self.offset = offset
self.layoutA = layoutA
# operator ==
def __eq__(self, other):
return self.layoutB == other.layoutB and self.offset == other.offset and self.layoutA == other.layoutA
# operator len(L) (len [rank] like tuples)
def __len__(self):
return len(self.layoutA)
# operator () (map coord to idx)
def __call__(self, *args):
return self.layoutB(self.offset + self.layoutA(*args))
# operator [] (get-i like tuples)
def __getitem__(self, i):
return ComposedLayout(self.layoutB, self.offset, self.layoutA[i])
# size(layout) Size of the domain
def size(self):
return size(self.layoutA)
# cosize(layout) Size of the codomain
def cosize(self):
return cosize(self.layoutB)
# print and str
def __str__(self):
return f"{self.layoutB} o {self.offset} o {self.layoutA}"
# error msgs and representation
def __repr__(self):
return f"ComposedLayout({repr(self.layoutB)},{repr(self.offset)},{repr(self.layoutA)})"
|