chore: import upstream snapshot with attribution

2026-07-13 13:36:25 +08:00
commit 26446540fa
3151 changed files with 974126 additions and 0 deletions
@@ -0,0 +1,110 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=invalid-name, too-many-locals, too-many-arguments
+"""Utility functions for bitserial operators"""
+
+import numpy as np
+
+import tvm
+from tvm import te
+from tvm.topi.transform import concatenate
+
+from ..utils import get_const_int
+
+
+def bitpack(data, bits, pack_axis, bit_axis, pack_type, name="QuantizeInput"):
+    """Packs data into format necessary for bitserial computation
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        The input tvm tensor
+    bits : int
+        Number of bits to use for packing
+    pack_axis : int
+       index of the axis to pack in data
+    bit_axis : int
+       index of axis to place bit axis in resulting packed data
+    pack_type : str
+        Data type for packing, must be one of: ['uint8', 'uint16', 'uint32', 'uint64']
+    name : Optional[str] = "QuantizeInput"
+        Name for the operation
+    """
+    ishape = data.shape
+    n = len(ishape)
+    if pack_type == "uint8":
+        data_width = 8
+    elif pack_type == "uint16":
+        data_width = 16
+    elif pack_type == "uint32":
+        data_width = 32
+    elif pack_type == "uint64":
+        data_width = 64
+
+    # Data must be in multiples of the data_width
+    assert get_const_int(ishape[pack_axis]) % data_width == 0, "Not a multiple of word size"
+
+    shape_vec = list(ishape)
+    shape_vec[pack_axis] = shape_vec[pack_axis] // data_width
+    shape_vec.insert(bit_axis, 1)
+    bitserial_oshape = tuple(shape_vec)
+    masks = np.array([0x1, 0x2, 0x4, 0x8, 0x10, 0x20, 0x40, 0x80])
+
+    # pack axis shifts if bit axis comes before
+    if bit_axis <= pack_axis:
+        pack_axis += 1
+
+    def _bitpack(*indices):
+        packed_data = [tvm.tirx.const(0, pack_type)] * bits
+        for k in range(data_width):
+            # Translate indices for packed data back to original
+            idx = [0] * n
+            j = 0
+            for i in range(n + 1):
+                if i == bit_axis:
+                    continue
+                if i == pack_axis:
+                    idx[j] = indices[i] * data_width + k
+                else:
+                    idx[j] = indices[i]
+                j += 1
+
+            element = data(*idx)
+            for b in range(bits):
+                extracted_bit = ((element & tvm.tirx.const(masks[b], "int32")) >> b).astype(
+                    pack_type
+                )
+                packed_data[b] = packed_data[b] | extracted_bit
+                if k < data_width - 1:
+                    packed_data[b] = packed_data[b] << 1
+
+            if k == data_width - 1:
+                return tuple(packed_data)
+        return tuple(packed_data)
+
+    output_tuple = te.compute(bitserial_oshape, _bitpack, name=name, tag="bitpack")
+
+    if bits > 1:
+        return concatenate(output_tuple, axis=bit_axis)
+    return output_tuple
+
+
+def binary_op_multiplier(pack_dtype):
+    """ "Returns number of bits packed into
+    pack_dtype: string
+        pack type for the operator (must be a uint)"""
+    return int(pack_dtype[4:])