mirror of
https://github.com/microsoft/BitNet.git
synced 2026-09-01 13:55:40 +00:00
initial commit
This commit is contained in:
potassiummmm
+46
@@ -0,0 +1,46 @@
|
||||
# Extensions
|
||||
|
||||
*.a
|
||||
*.bat
|
||||
*.bin
|
||||
*.dll
|
||||
*.dot
|
||||
*.etag
|
||||
*.exe
|
||||
*.gcda
|
||||
*.gcno
|
||||
*.gcov
|
||||
*.gguf
|
||||
*.gguf.json
|
||||
*.lastModified
|
||||
*.log
|
||||
*.metallib
|
||||
*.o
|
||||
*.so
|
||||
*.tmp
|
||||
|
||||
# IDE / OS
|
||||
|
||||
.cache/
|
||||
.ccls-cache/
|
||||
.direnv/
|
||||
.DS_Store
|
||||
.envrc
|
||||
.idea/
|
||||
.swiftpm
|
||||
.vs/
|
||||
.vscode/
|
||||
nppBackup
|
||||
|
||||
# Models
|
||||
models/*
|
||||
|
||||
# Python
|
||||
|
||||
/.venv
|
||||
__pycache__/
|
||||
*/poetry.lock
|
||||
poetry.toml
|
||||
|
||||
build/
|
||||
logs/
|
||||
@@ -0,0 +1,4 @@
|
||||
[submodule "3rdparty/llama.cpp"]
|
||||
path = 3rdparty/llama.cpp
|
||||
url = https://github.com/Eddie-Wang1120/llama.cpp.git
|
||||
branch = merge-dev
|
||||
+1
Submodule 3rdparty/llama.cpp added at 5371710215
@@ -0,0 +1,73 @@
|
||||
cmake_minimum_required(VERSION 3.14) # for add_link_options and implicit target directories.
|
||||
project("bitnet.cpp" C CXX)
|
||||
include(CheckIncludeFileCXX)
|
||||
|
||||
set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
|
||||
|
||||
if (NOT XCODE AND NOT MSVC AND NOT CMAKE_BUILD_TYPE)
|
||||
set(CMAKE_BUILD_TYPE Release CACHE STRING "Build type" FORCE)
|
||||
set_property(CACHE CMAKE_BUILD_TYPE PROPERTY STRINGS "Debug" "Release" "MinSizeRel" "RelWithDebInfo")
|
||||
endif()
|
||||
|
||||
set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/bin)
|
||||
|
||||
# option list
|
||||
option(BITNET_ARM_TL1 "bitnet.cpp: use tl1 on arm platform" OFF)
|
||||
option(BITNET_X86_TL2 "bitnet.cpp: use tl2 on x86 platform" OFF)
|
||||
|
||||
|
||||
set(CMAKE_CXX_STANDARD_REQUIRED true)
|
||||
set(CMAKE_C_STANDARD 11)
|
||||
set(CMAKE_C_STANDARD_REQUIRED true)
|
||||
set(THREADS_PREFER_PTHREAD_FLAG ON)
|
||||
|
||||
# override ggml options
|
||||
set(GGML_BITNET_ARM_TL1 ${BITNET_ARM_TL1})
|
||||
set(GGML_BITNET_X86_TL2 ${BITNET_X86_TL2})
|
||||
|
||||
if (GGML_BITNET_ARM_TL1)
|
||||
add_compile_definitions(GGML_BITNET_ARM_TL1)
|
||||
endif()
|
||||
if (GGML_BITNET_X86_TL2)
|
||||
add_compile_definitions(GGML_BITNET_X86_TL2)
|
||||
endif()
|
||||
|
||||
find_package(Threads REQUIRED)
|
||||
|
||||
add_subdirectory(src)
|
||||
add_subdirectory(3rdparty/llama.cpp)
|
||||
|
||||
# install
|
||||
|
||||
include(GNUInstallDirs)
|
||||
include(CMakePackageConfigHelpers)
|
||||
|
||||
set(LLAMA_INCLUDE_INSTALL_DIR ${CMAKE_INSTALL_INCLUDEDIR}
|
||||
CACHE PATH "Location of header files")
|
||||
set(LLAMA_LIB_INSTALL_DIR ${CMAKE_INSTALL_LIBDIR}
|
||||
CACHE PATH "Location of library files")
|
||||
set(LLAMA_BIN_INSTALL_DIR ${CMAKE_INSTALL_BINDIR}
|
||||
CACHE PATH "Location of binary files")
|
||||
set(LLAMA_BUILD_NUMBER ${BUILD_NUMBER})
|
||||
set(LLAMA_BUILD_COMMIT ${BUILD_COMMIT})
|
||||
set(LLAMA_INSTALL_VERSION 0.0.${BUILD_NUMBER})
|
||||
|
||||
get_target_property(GGML_DIRECTORY ggml SOURCE_DIR)
|
||||
get_directory_property(GGML_DIR_DEFINES DIRECTORY ${GGML_DIRECTORY} COMPILE_DEFINITIONS)
|
||||
get_target_property(GGML_TARGET_DEFINES ggml COMPILE_DEFINITIONS)
|
||||
set(GGML_TRANSIENT_DEFINES ${GGML_TARGET_DEFINES} ${GGML_DIR_DEFINES})
|
||||
get_target_property(GGML_LINK_LIBRARIES ggml LINK_LIBRARIES)
|
||||
|
||||
get_directory_property(LLAMA_TRANSIENT_DEFINES COMPILE_DEFINITIONS)
|
||||
|
||||
write_basic_package_version_file(
|
||||
${CMAKE_CURRENT_BINARY_DIR}/LlamaConfigVersion.cmake
|
||||
VERSION ${LLAMA_INSTALL_VERSION}
|
||||
COMPATIBILITY SameMajorVersion)
|
||||
|
||||
install(FILES ${CMAKE_CURRENT_BINARY_DIR}/LlamaConfig.cmake
|
||||
${CMAKE_CURRENT_BINARY_DIR}/LlamaConfigVersion.cmake
|
||||
DESTINATION ${CMAKE_INSTALL_LIBDIR}/cmake/Llama)
|
||||
|
||||
set_target_properties(llama PROPERTIES PUBLIC_HEADER ${CMAKE_CURRENT_SOURCE_DIR}/llama.h)
|
||||
install(TARGETS llama LIBRARY PUBLIC_HEADER)
|
||||
@@ -0,0 +1,9 @@
|
||||
# Microsoft Open Source Code of Conduct
|
||||
|
||||
This project has adopted the [Microsoft Open Source Code of Conduct](https://opensource.microsoft.com/codeofconduct/).
|
||||
|
||||
Resources:
|
||||
|
||||
- [Microsoft Open Source Code of Conduct](https://opensource.microsoft.com/codeofconduct/)
|
||||
- [Microsoft Code of Conduct FAQ](https://opensource.microsoft.com/codeofconduct/faq/)
|
||||
- Contact [opencode@microsoft.com](mailto:opencode@microsoft.com) with questions or concerns
|
||||
@@ -0,0 +1,21 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) Microsoft Corporation.
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE
|
||||
@@ -0,0 +1,228 @@
|
||||
# bitnet.cpp
|
||||
[](https://opensource.org/licenses/MIT)
|
||||
|
||||
|
||||
bitnet.cpp is the official inference framework for BitNet models (e.g., BitNet b1.58), optimized for CPU devices. It offers a suite of optimized kernels, that support lossless inference of 1.58-bit models on both x86 and ARM architectures.
|
||||
|
||||
## Demo
|
||||
|
||||
A demo of bitnet.cpp runing a BitNet b1.58 3B model on Apple M2:
|
||||
|
||||
https://github.com/user-attachments/assets/7f46b736-edec-4828-b809-4be780a3e5b1
|
||||
|
||||
## Timeline
|
||||
|
||||
- 10/17/2024 bitnet.cpp 1.0 released.
|
||||
- 02/27/2024 [The Era of 1-bit LLMs: All Large Language Models are in 1.58 Bits](https://arxiv.org/abs/2402.17764)
|
||||
- 10/17/2023 [BitNet: Scaling 1-bit Transformers for Large Language Models](https://arxiv.org/abs/2310.11453)
|
||||
|
||||
## Supported Models
|
||||
|
||||
bitnet.cpp supports a list of 1-bit models available on [Hugging Face](https://huggingface.co/)
|
||||
|
||||
|
||||
<table>
|
||||
</tr>
|
||||
<tr>
|
||||
<th rowspan="2">Model</th>
|
||||
<th rowspan="2">Parameters</th>
|
||||
<th rowspan="2">CPU</th>
|
||||
<th colspan="3">Kernel</th>
|
||||
</tr>
|
||||
<tr>
|
||||
<th>I2_S</th>
|
||||
<th>TL1</th>
|
||||
<th>TL2</th>
|
||||
</tr>
|
||||
<tr>
|
||||
<td rowspan="2"><a href="https://huggingface.co/1bitLLM/bitnet_b1_58-large">bitnet_b1_58-large</a></td>
|
||||
<td rowspan="2">0.7B</td>
|
||||
<td>x86</td>
|
||||
<td>✔</td>
|
||||
<td>✘</td>
|
||||
<td>✔</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>ARM</td>
|
||||
<td>✔</td>
|
||||
<td>✔</td>
|
||||
<td>✘</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td rowspan="2"><a href="https://huggingface.co/1bitLLM/bitnet_b1_58-3B">bitnet_b1_58-3B</a></td>
|
||||
<td rowspan="2">3.3B</td>
|
||||
<td>x86</td>
|
||||
<td>✘</td>
|
||||
<td>✘</td>
|
||||
<td>✔</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>ARM</td>
|
||||
<td>✘</td>
|
||||
<td>✔</td>
|
||||
<td>✘</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td rowspan="2"><a href="https://huggingface.co/HF1BitLLM/Llama3-8B-1.58-100B-tokens">Llama3-8B-1.58-100B-tokens</a></td>
|
||||
<td rowspan="2">8.0B</td>
|
||||
<td>x86</td>
|
||||
<td>✔</td>
|
||||
<td>✘</td>
|
||||
<td>✔</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>ARM</td>
|
||||
<td>✔</td>
|
||||
<td>✔</td>
|
||||
<td>✘</td>
|
||||
</tr>
|
||||
</table>
|
||||
|
||||
|
||||
|
||||
## Installation
|
||||
|
||||
### Requirements
|
||||
- python>=3.9
|
||||
- cmake>=3.22
|
||||
- clang>=18
|
||||
- For Windows users, install [Visual Studio 2022](https://visualstudio.microsoft.com/downloads/). In the installer, toggle on at least the following options(this also automatically installs the required additional tools like CMake):
|
||||
- Desktop-development with C++
|
||||
- C++-CMake Tools for Windows
|
||||
- Git for Windows
|
||||
- C++-Clang Compiler for Windows
|
||||
- MS-Build Support for LLVM-Toolset (clang)
|
||||
- For Debian/Ubuntu users, you can download with [Automatic installation script](https://apt.llvm.org/)
|
||||
|
||||
` bash -c "$(wget -O - https://apt.llvm.org/llvm.sh)"`
|
||||
- conda (highly recommend)
|
||||
|
||||
### Build from source
|
||||
|
||||
> [!IMPORTANT]
|
||||
> If you are using Windows, please remember to always use a Developer Command Prompt / PowerShell for VS2022 for the following commands
|
||||
|
||||
1. Clone the repo
|
||||
```bash
|
||||
git clone --recursive https://github.com/microsoft/BitNet.git
|
||||
cd BitNet
|
||||
```
|
||||
2. Install the dependencies
|
||||
```bash
|
||||
# (Recommended) Create a new conda environment
|
||||
conda create -n bitnet-cpp python=3.9
|
||||
conda activate bitnet-cpp
|
||||
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
3. Build the project
|
||||
```bash
|
||||
# Download the model from Hugging Face, convert it to quantized gguf format, and build the project
|
||||
python setup_env.py --hf-repo HF1BitLLM/Llama3-8B-1.58-100B-tokens -q i2_s
|
||||
|
||||
# Or you can manually download the model and run with local path
|
||||
huggingface-cli download HF1BitLLM/Llama3-8B-1.58-100B-tokens --local-dir models/Llama3-8B-1.58-100B-tokens
|
||||
python setup_env.py -md models/Llama3-8B-1.58-100B-tokens -q i2_s
|
||||
```
|
||||
<pre>
|
||||
usage: setup_env.py [-h] [--hf-repo {1bitLLM/bitnet_b1_58-large,1bitLLM/bitnet_b1_58-3B,HF1BitLLM/Llama3-8B-1.58-100B-tokens}] [--model-dir MODEL_DIR] [--log-dir LOG_DIR] [--quant-type {i2_s,tl1}] [--quant-embd]
|
||||
[--use-pretuned]
|
||||
|
||||
Setup the environment for running inference
|
||||
|
||||
optional arguments:
|
||||
-h, --help show this help message and exit
|
||||
--hf-repo {1bitLLM/bitnet_b1_58-large,1bitLLM/bitnet_b1_58-3B,HF1BitLLM/Llama3-8B-1.58-100B-tokens}, -hr {1bitLLM/bitnet_b1_58-large,1bitLLM/bitnet_b1_58-3B,HF1BitLLM/Llama3-8B-1.58-100B-tokens}
|
||||
Model used for inference
|
||||
--model-dir MODEL_DIR, -md MODEL_DIR
|
||||
Directory to save/load the model
|
||||
--log-dir LOG_DIR, -ld LOG_DIR
|
||||
Directory to save the logging info
|
||||
--quant-type {i2_s,tl1}, -q {i2_s,tl1}
|
||||
Quantization type
|
||||
--quant-embd Quantize the embeddings to f16
|
||||
--use-pretuned, -p Use the pretuned kernel parameters
|
||||
</pre>
|
||||
## Usage
|
||||
### Basic usage
|
||||
```bash
|
||||
# Run inference with the quantized model
|
||||
python run_inference.py -m models/Llama3-8B-1.58-100B-tokens/ggml-model-i2_s.gguf -p "Daniel went back to the the the garden. Mary travelled to the kitchen. Sandra journeyed to the kitchen. Sandra went to the hallway. John went to the bedroom. Mary went back to the garden. Where is Mary?\nAnswer:" -n 6 -temp 0
|
||||
|
||||
# Output:
|
||||
# Daniel went back to the the the garden. Mary travelled to the kitchen. Sandra journeyed to the kitchen. Sandra went to the hallway. John went to the bedroom. Mary went back to the garden. Where is Mary?
|
||||
# Answer: Mary is in the garden.
|
||||
|
||||
```
|
||||
<pre>
|
||||
usage: run_inference.py [-h] [-m MODEL] [-n N_PREDICT] -p PROMPT [-t THREADS] [-c CTX_SIZE] [-temp TEMPERATURE]
|
||||
|
||||
Run inference
|
||||
|
||||
optional arguments:
|
||||
-h, --help show this help message and exit
|
||||
-m MODEL, --model MODEL
|
||||
Path to model file
|
||||
-n N_PREDICT, --n-predict N_PREDICT
|
||||
Number of tokens to predict when generating text
|
||||
-p PROMPT, --prompt PROMPT
|
||||
Prompt to generate text from
|
||||
-t THREADS, --threads THREADS
|
||||
Number of threads to use
|
||||
-c CTX_SIZE, --ctx-size CTX_SIZE
|
||||
Size of the prompt context
|
||||
-temp TEMPERATURE, --temperature TEMPERATURE
|
||||
Temperature, a hyperparameter that controls the randomness of the generated text
|
||||
</pre>
|
||||
|
||||
### Benchmark
|
||||
We provide scripts to run the inference benchmark providing a model.
|
||||
|
||||
```
|
||||
usage: e2e_benchmark.py -m MODEL [-n N_TOKEN] [-p N_PROMPT] [-t THREADS]
|
||||
|
||||
Setup the environment for running the inference
|
||||
|
||||
required arguments:
|
||||
-m MODEL, --model MODEL
|
||||
Path to the model file.
|
||||
|
||||
optional arguments:
|
||||
-h, --help
|
||||
Show this help message and exit.
|
||||
-n N_TOKEN, --n-token N_TOKEN
|
||||
Number of generated tokens.
|
||||
-p N_PROMPT, --n-prompt N_PROMPT
|
||||
Prompt to generate text from.
|
||||
-t THREADS, --threads THREADS
|
||||
Number of threads to use.
|
||||
```
|
||||
|
||||
Here's a brief explanation of each argument:
|
||||
|
||||
- `-m`, `--model`: The path to the model file. This is a required argument that must be provided when running the script.
|
||||
- `-n`, `--n-token`: The number of tokens to generate during the inference. It is an optional argument with a default value of 128.
|
||||
- `-p`, `--n-prompt`: The number of prompt tokens to use for generating text. This is an optional argument with a default value of 512.
|
||||
- `-t`, `--threads`: The number of threads to use for running the inference. It is an optional argument with a default value of 2.
|
||||
- `-h`, `--help`: Show the help message and exit. Use this argument to display usage information.
|
||||
|
||||
For example:
|
||||
|
||||
```sh
|
||||
python utils/e2e_benchmark.py -m /path/to/model -n 200 -p 256 -t 4
|
||||
```
|
||||
|
||||
This command would run the inference benchmark using the model located at `/path/to/model`, generating 200 tokens from a 256 token prompt, utilizing 4 threads.
|
||||
|
||||
For the model layout that do not supported by any public model, we provide scripts to generate a dummy model with the given model layout, and run the benchmark on your machine:
|
||||
|
||||
```bash
|
||||
python utils/generate-dummy-bitnet-model.py models/bitnet_b1_58-large --outfile models/dummy-bitnet-125m.tl1.gguf --outtype tl1 --model-size 125M
|
||||
|
||||
# Run benchmark with the generated model, use -m to specify the model path, -p to specify the prompt processed, -n to specify the number of token to generate
|
||||
python utils/e2e_benchmark.py -m models/dummy-bitnet-125m.tl1.gguf -p 512 -n 128
|
||||
```
|
||||
|
||||
## Acknowledgements
|
||||
|
||||
This project is based on the [llama.cpp](https://github.com/ggerganov/llama.cpp) framework. We would like to thank all the authors for their contributions to the open-source community. We also thank [T-MAC](https://github.com/microsoft/T-MAC/) team for the helpful discussion on the LUT method for low-bit LLM inference.
|
||||
+41
@@ -0,0 +1,41 @@
|
||||
<!-- BEGIN MICROSOFT SECURITY.MD V0.0.9 BLOCK -->
|
||||
|
||||
## Security
|
||||
|
||||
Microsoft takes the security of our software products and services seriously, which includes all source code repositories managed through our GitHub organizations, which include [Microsoft](https://github.com/Microsoft), [Azure](https://github.com/Azure), [DotNet](https://github.com/dotnet), [AspNet](https://github.com/aspnet) and [Xamarin](https://github.com/xamarin).
|
||||
|
||||
If you believe you have found a security vulnerability in any Microsoft-owned repository that meets [Microsoft's definition of a security vulnerability](https://aka.ms/security.md/definition), please report it to us as described below.
|
||||
|
||||
## Reporting Security Issues
|
||||
|
||||
**Please do not report security vulnerabilities through public GitHub issues.**
|
||||
|
||||
Instead, please report them to the Microsoft Security Response Center (MSRC) at [https://msrc.microsoft.com/create-report](https://aka.ms/security.md/msrc/create-report).
|
||||
|
||||
If you prefer to submit without logging in, send email to [secure@microsoft.com](mailto:secure@microsoft.com). If possible, encrypt your message with our PGP key; please download it from the [Microsoft Security Response Center PGP Key page](https://aka.ms/security.md/msrc/pgp).
|
||||
|
||||
You should receive a response within 24 hours. If for some reason you do not, please follow up via email to ensure we received your original message. Additional information can be found at [microsoft.com/msrc](https://www.microsoft.com/msrc).
|
||||
|
||||
Please include the requested information listed below (as much as you can provide) to help us better understand the nature and scope of the possible issue:
|
||||
|
||||
* Type of issue (e.g. buffer overflow, SQL injection, cross-site scripting, etc.)
|
||||
* Full paths of source file(s) related to the manifestation of the issue
|
||||
* The location of the affected source code (tag/branch/commit or direct URL)
|
||||
* Any special configuration required to reproduce the issue
|
||||
* Step-by-step instructions to reproduce the issue
|
||||
* Proof-of-concept or exploit code (if possible)
|
||||
* Impact of the issue, including how an attacker might exploit the issue
|
||||
|
||||
This information will help us triage your report more quickly.
|
||||
|
||||
If you are reporting for a bug bounty, more complete reports can contribute to a higher bounty award. Please visit our [Microsoft Bug Bounty Program](https://aka.ms/security.md/msrc/bounty) page for more details about our active programs.
|
||||
|
||||
## Preferred Languages
|
||||
|
||||
We prefer all communications to be in English.
|
||||
|
||||
## Policy
|
||||
|
||||
Microsoft follows the principle of [Coordinated Vulnerability Disclosure](https://aka.ms/security.md/cvd).
|
||||
|
||||
<!-- END MICROSOFT SECURITY.MD BLOCK -->
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 13 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 16 KiB |
@@ -0,0 +1,49 @@
|
||||
Codegen for TL1 and TL2
|
||||
------------------------
|
||||
|
||||
codegen_tl1.py and codegen_tl2.py are using params to generate kernel codes in different devices to achieve fastest performance for TL1 and TL2.
|
||||
|
||||
We cutting weight into multiple compute blocks to best utilize hardware capabilities.
|
||||
|
||||
### Example
|
||||
bitnet_b1_58-large:
|
||||
|
||||
- Make sure Mamtul kernels shapes \
|
||||
For example, bitnet_b1_58-large Matmul kernel shapes are:\
|
||||
[1536, 4096]\
|
||||
[1536, 1536]\
|
||||
[4096, 1536]
|
||||
|
||||
- Make sure each BM, BK, bm for each kernel to meet the requirements below
|
||||
- Generate codes\
|
||||
For example, for bitnet_b1_58-large, we can gencode like:
|
||||
|
||||
```bash
|
||||
# For TL1
|
||||
python utils/codegen_tl1.py --model bitnet_b1_58-large --BM 256,128,256 --BK 128,64,128 --bm 32,64,32
|
||||
|
||||
# For TL2
|
||||
python utils/codegen_tl2.py --model bitnet_b1_58-large --BM 256,128,256 --BK 96,192,96 --bm 32,32,32
|
||||
```
|
||||
|
||||
### TL1:
|
||||
|
||||
|
||||
For TL1, we cut weight into M / BM weights, each weight shape is (BM, K). Then we cut weight into K / BK weights, each weight shape is (BM, BK). As for (BM, BK) weight, we cut it the same way into (bm, compute_num / bm) compute blocks, and finish computing in it.
|
||||
|
||||
Thus, we need to make sure
|
||||
- M % BM == 0
|
||||
- K % BK == 0
|
||||
- BM % bm == 0
|
||||
- bm choose in [32, 64]
|
||||
|
||||
### TL2:
|
||||
|
||||
|
||||
For TL2, things got a little more complicated. Due to TL2 needs BK % 6 == 0, we need to split K into threeK and twoK, in which compute in TL2 for (M, threeK), compute in TL1 for (M, two_K).
|
||||
|
||||
Thus, we needs to make sure
|
||||
- M % BM == 0
|
||||
- K % BK % 32 == 0
|
||||
- BM % bm == 0
|
||||
- bm choose in \[32\]
|
||||
@@ -0,0 +1,49 @@
|
||||
#pragma once
|
||||
|
||||
#include "ggml.h"
|
||||
#include "ggml-backend.h"
|
||||
|
||||
#ifdef __ARM_NEON
|
||||
#include <arm_neon.h>
|
||||
typedef float32_t bitnet_float_type;
|
||||
#else
|
||||
typedef float bitnet_float_type;
|
||||
#endif
|
||||
|
||||
#ifdef __cplusplus
|
||||
extern "C" {
|
||||
#endif
|
||||
|
||||
struct bitnet_tensor_extra {
|
||||
int lut_scales_size;
|
||||
int BK;
|
||||
int n_tile_num;
|
||||
uint8_t * qweights;
|
||||
bitnet_float_type * scales;
|
||||
};
|
||||
|
||||
GGML_API void ggml_bitnet_init(void);
|
||||
GGML_API void ggml_bitnet_free(void);
|
||||
// src0->type == Q4_0/IQ2_XXS/IQ3_XXS
|
||||
// bitnet.cpp currently only supports BitNet quantization or GPTQ-like quantization (only scales, without zeros)
|
||||
// If use i-quantization gguf models, the results will be wrong
|
||||
// TODO: add customized block types Q2_0/Q3_0
|
||||
GGML_API bool ggml_bitnet_can_mul_mat(const struct ggml_tensor * src0, const struct ggml_tensor * src1, const struct ggml_tensor * dst);
|
||||
GGML_API size_t ggml_bitnet_mul_mat_get_wsize(const struct ggml_tensor * src0, const struct ggml_tensor * src1, const struct ggml_tensor * dst);
|
||||
GGML_API void ggml_bitnet_mul_mat_task_init(void * src1, void * qlut, void * lut_scales, void * lut_biases, int n, int k, int m, int bits);
|
||||
GGML_API void ggml_bitnet_mul_mat_task_compute(void * src0, void * scales, void * qlut, void * lut_scales, void * lut_biases, void * dst, int n, int k, int m, int bits);
|
||||
GGML_API void ggml_bitnet_transform_tensor(struct ggml_tensor * tensor);
|
||||
GGML_API int ggml_bitnet_get_type_bits(enum ggml_type type);
|
||||
GGML_API void ggml_bitnet_set_n_threads(int n_threads);
|
||||
#if defined(GGML_BITNET_ARM_TL1)
|
||||
GGML_API void ggml_qgemm_lut(int m, int k, void* A, void* LUT, void* Scales, void* LUT_Scales, void* C);
|
||||
GGML_API void ggml_preprocessor(int m, int k, void* B, void* LUT_Scales, void* QLUT);
|
||||
#endif
|
||||
#if defined(GGML_BITNET_X86_TL2)
|
||||
GGML_API void ggml_qgemm_lut(int bs, int m, int k, int BK, void* A, void* sign, void* LUT, void* Scales, void* LUT_Scales, void* C);
|
||||
GGML_API void ggml_preprocessor(int bs, int m, int three_k, int two_k, void* B, void* LUT_Scales, void* Three_QLUT, void* Two_QLUT);
|
||||
#endif
|
||||
|
||||
#ifdef __cplusplus
|
||||
}
|
||||
#endif
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 36 KiB |
Binary file not shown.
@@ -0,0 +1,771 @@
|
||||
#if defined(GGML_BITNET_ARM_TL1)
|
||||
#include "ggml-bitnet.h"
|
||||
#define GGML_BITNET_MAX_NODES 8192
|
||||
static bool initialized = false;
|
||||
static bitnet_tensor_extra * bitnet_tensor_extras = nullptr;
|
||||
static size_t bitnet_tensor_extras_index = 0;
|
||||
static void * aligned_malloc(size_t size) {{
|
||||
#if defined(_WIN32)
|
||||
return _aligned_malloc(size, 64);
|
||||
#else
|
||||
void * ptr = nullptr;
|
||||
posix_memalign(&ptr, 64, size);
|
||||
return ptr;
|
||||
#endif
|
||||
}}
|
||||
static void aligned_free(void * ptr) {{
|
||||
#if defined(_WIN32)
|
||||
_aligned_free(ptr);
|
||||
#else
|
||||
free(ptr);
|
||||
#endif
|
||||
}}
|
||||
|
||||
void per_tensor_quant(int k, void* lut_scales_, void* b_) {{
|
||||
bitnet_float_type* lut_scales = (bitnet_float_type*)lut_scales_;
|
||||
bitnet_float_type* b = (bitnet_float_type*)b_;
|
||||
#ifdef __ARM_NEON
|
||||
float32x4_t temp_max = vdupq_n_f32(0);
|
||||
for (int i=0; i < k / 4; i++) {{
|
||||
float32x4_t vec_bs = vld1q_f32(b + 4 * i);
|
||||
float32x4_t abssum = vabsq_f32(vec_bs);
|
||||
temp_max = vmaxq_f32(abssum, temp_max);
|
||||
}}
|
||||
float32_t scales = 127 / vmaxvq_f32(temp_max);
|
||||
*lut_scales = scales;
|
||||
#elif defined __AVX2__
|
||||
__m256 max_vec = _mm256_set1_ps(0.f);
|
||||
const __m256 vec_sign = _mm256_set1_ps(-0.0f);
|
||||
// #pragma unroll
|
||||
for (int i = 0; i < k / 8; i++) {{
|
||||
__m256 vec_b = _mm256_loadu_ps(b + i * 8);
|
||||
__m256 vec_babs = _mm256_andnot_ps(vec_sign, vec_b);
|
||||
max_vec = _mm256_max_ps(vec_babs, max_vec);
|
||||
}}
|
||||
__m128 max1 = _mm_max_ps(_mm256_extractf128_ps(max_vec, 1), _mm256_castps256_ps128(max_vec));
|
||||
max1 = _mm_max_ps(max1, _mm_movehl_ps(max1, max1));
|
||||
max1 = _mm_max_ss(max1, _mm_movehdup_ps(max1));
|
||||
float scales = 127 / _mm_cvtss_f32(max1);
|
||||
*lut_scales = scales;
|
||||
#endif
|
||||
}}
|
||||
|
||||
void partial_max_reset(void* lut_scales_) {{
|
||||
bitnet_float_type* lut_scales = (bitnet_float_type*)lut_scales_;
|
||||
*lut_scales = 0.0;
|
||||
}}
|
||||
|
||||
#ifdef __ARM_NEON
|
||||
inline void Transpose_8_8(
|
||||
int16x8_t *v0,
|
||||
int16x8_t *v1,
|
||||
int16x8_t *v2,
|
||||
int16x8_t *v3,
|
||||
int16x8_t *v4,
|
||||
int16x8_t *v5,
|
||||
int16x8_t *v6,
|
||||
int16x8_t *v7)
|
||||
{{
|
||||
int16x8x2_t q04 = vzipq_s16(*v0, *v4);
|
||||
int16x8x2_t q15 = vzipq_s16(*v1, *v5);
|
||||
int16x8x2_t q26 = vzipq_s16(*v2, *v6);
|
||||
int16x8x2_t q37 = vzipq_s16(*v3, *v7);
|
||||
|
||||
int16x8x2_t q0246_0 = vzipq_s16(q04.val[0], q26.val[0]);
|
||||
int16x8x2_t q0246_1 = vzipq_s16(q04.val[1], q26.val[1]);
|
||||
int16x8x2_t q1357_0 = vzipq_s16(q15.val[0], q37.val[0]);
|
||||
int16x8x2_t q1357_1 = vzipq_s16(q15.val[1], q37.val[1]);
|
||||
|
||||
int16x8x2_t q_fin_0 = vzipq_s16(q0246_0.val[0], q1357_0.val[0]);
|
||||
int16x8x2_t q_fin_1 = vzipq_s16(q0246_0.val[1], q1357_0.val[1]);
|
||||
int16x8x2_t q_fin_2 = vzipq_s16(q0246_1.val[0], q1357_1.val[0]);
|
||||
int16x8x2_t q_fin_3 = vzipq_s16(q0246_1.val[1], q1357_1.val[1]);
|
||||
|
||||
*v0 = q_fin_0.val[0];
|
||||
*v1 = q_fin_0.val[1];
|
||||
*v2 = q_fin_1.val[0];
|
||||
*v3 = q_fin_1.val[1];
|
||||
*v4 = q_fin_2.val[0];
|
||||
*v5 = q_fin_2.val[1];
|
||||
*v6 = q_fin_3.val[0];
|
||||
*v7 = q_fin_3.val[1];
|
||||
}}
|
||||
#endif
|
||||
|
||||
template<int act_k>
|
||||
inline void lut_ctor(int8_t* qlut, bitnet_float_type* b, bitnet_float_type* lut_scales) {{
|
||||
#ifdef __ARM_NEON
|
||||
int16x8_t vec_lut[16];
|
||||
float32_t scales = *lut_scales;
|
||||
uint8_t tbl_mask[16];
|
||||
tbl_mask[0] = 0;
|
||||
tbl_mask[1] = 2;
|
||||
tbl_mask[2] = 4;
|
||||
tbl_mask[3] = 6;
|
||||
tbl_mask[4] = 8;
|
||||
tbl_mask[5] = 10;
|
||||
tbl_mask[6] = 12;
|
||||
tbl_mask[7] = 14;
|
||||
tbl_mask[8] = 1;
|
||||
tbl_mask[9] = 3;
|
||||
tbl_mask[10] = 5;
|
||||
tbl_mask[11] = 7;
|
||||
tbl_mask[12] = 9;
|
||||
tbl_mask[13] = 11;
|
||||
tbl_mask[14] = 13;
|
||||
tbl_mask[15] = 15;
|
||||
uint8x16_t tbl_mask_q = vld1q_u8(tbl_mask);
|
||||
#pragma unroll
|
||||
for (int k = 0; k < act_k / 16; ++k) {{
|
||||
float32x4x2_t vec_bs_x0 = vld2q_f32(b + k * 16);
|
||||
float32x4x2_t vec_bs_x1 = vld2q_f32(b + k * 16 + 8);
|
||||
float32x4_t vec_f_0 = vmulq_n_f32(vec_bs_x0.val[0], scales);
|
||||
float32x4_t vec_f_1 = vmulq_n_f32(vec_bs_x0.val[1], scales);
|
||||
float32x4_t vec_f_2 = vmulq_n_f32(vec_bs_x1.val[0], scales);
|
||||
float32x4_t vec_f_3 = vmulq_n_f32(vec_bs_x1.val[1], scales);
|
||||
int32x4_t vec_b_0 = vcvtnq_s32_f32(vec_f_0);
|
||||
int32x4_t vec_b_1 = vcvtnq_s32_f32(vec_f_1);
|
||||
int32x4_t vec_b_2 = vcvtnq_s32_f32(vec_f_2);
|
||||
int32x4_t vec_b_3 = vcvtnq_s32_f32(vec_f_3);
|
||||
int16x4_t vec_b16_0 = vmovn_s32(vec_b_0);
|
||||
int16x4_t vec_b16_1 = vmovn_s32(vec_b_1);
|
||||
int16x4_t vec_b16_2 = vmovn_s32(vec_b_2);
|
||||
int16x4_t vec_b16_3 = vmovn_s32(vec_b_3);
|
||||
int16x8_t vec_bs_0 = vcombine_s16(vec_b16_0, vec_b16_2);
|
||||
int16x8_t vec_bs_1 = vcombine_s16(vec_b16_1, vec_b16_3);
|
||||
vec_lut[0] = vdupq_n_s16(0);
|
||||
vec_lut[0] = vec_lut[0] - vec_bs_0;
|
||||
vec_lut[0] = vec_lut[0] - vec_bs_1;
|
||||
vec_lut[1] = vdupq_n_s16(0);
|
||||
vec_lut[1] = vec_lut[1] - vec_bs_0;
|
||||
vec_lut[2] = vdupq_n_s16(0);
|
||||
vec_lut[2] = vec_lut[2] - vec_bs_0;
|
||||
vec_lut[2] = vec_lut[2] + vec_bs_1;
|
||||
vec_lut[3] = vdupq_n_s16(0);
|
||||
vec_lut[3] = vec_lut[3] - vec_bs_1;
|
||||
vec_lut[4] = vdupq_n_s16(0);
|
||||
vec_lut[5] = vec_bs_1;
|
||||
vec_lut[6] = vec_bs_0;
|
||||
vec_lut[6] = vec_lut[6] - vec_bs_1;
|
||||
vec_lut[7] = vec_bs_0;
|
||||
vec_lut[8] = vec_bs_0;
|
||||
vec_lut[8] = vec_lut[8] + vec_bs_1;
|
||||
Transpose_8_8(&(vec_lut[0]), &(vec_lut[1]), &(vec_lut[2]), &(vec_lut[3]),
|
||||
&(vec_lut[4]), &(vec_lut[5]), &(vec_lut[6]), &(vec_lut[7]));
|
||||
Transpose_8_8(&(vec_lut[8]), &(vec_lut[9]), &(vec_lut[10]), &(vec_lut[11]),
|
||||
&(vec_lut[12]), &(vec_lut[13]), &(vec_lut[14]), &(vec_lut[15]));
|
||||
#pragma unroll
|
||||
for (int idx = 0; idx < 8; idx++) {{
|
||||
int8x16_t q0_s = vqtbl1q_s8(vreinterpretq_s8_s16(vec_lut[idx]), tbl_mask_q);
|
||||
int8x8_t q0_low = vget_low_s8(q0_s);
|
||||
int8x8_t q0_high = vget_high_s8(q0_s);
|
||||
int8x16_t q1_s = vqtbl1q_s8(vreinterpretq_s8_s16(vec_lut[idx + 8]), tbl_mask_q);
|
||||
int8x8_t q1_low = vget_low_s8(q1_s);
|
||||
int8x8_t q1_high = vget_high_s8(q1_s);
|
||||
vst1_s8(qlut + k * 16 * 8 * 2 + idx * 16 * 2, q0_high);
|
||||
vst1_s8(qlut + k * 16 * 8 * 2 + idx * 16 * 2 + 8, q1_high);
|
||||
vst1_s8(qlut + k * 16 * 8 * 2 + idx * 16 * 2 + 16, q0_low);
|
||||
vst1_s8(qlut + k * 16 * 8 * 2 + idx * 16 * 2 + 24, q1_low);
|
||||
}}
|
||||
}}
|
||||
#endif
|
||||
}}
|
||||
|
||||
static bool is_type_supported(enum ggml_type type) {{
|
||||
if (type == GGML_TYPE_Q4_0 ||
|
||||
type == GGML_TYPE_TL1) {{
|
||||
return true;
|
||||
}} else {{
|
||||
return false;
|
||||
}}
|
||||
}}
|
||||
#include <arm_neon.h>
|
||||
|
||||
#define BM14336_4096 256
|
||||
#define BBK14336_4096 128
|
||||
inline void tbl_impl_14336_4096(int32_t* c, int8_t* lut, uint8_t* a) {
|
||||
#ifdef __ARM_NEON
|
||||
const int KK = BBK14336_4096 / 2;
|
||||
const uint8x16_t vec_mask = vdupq_n_u8(0x0f);
|
||||
const int8x16_t vec_zero = vdupq_n_s16(0x0000);
|
||||
int8x16_t vec_lut[2 * KK];
|
||||
int16x8_t vec_c[8];
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 2 * KK; k++) {
|
||||
vec_lut[k] = vld1q_s8(lut + k * 16);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < BM14336_4096; i += 64) {
|
||||
#pragma unroll
|
||||
for (int i=0; i<8; i++) {
|
||||
vec_c[i] = vandq_s16(vec_c[i], vec_zero);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int k = 0; k < KK / 2; k++) {
|
||||
|
||||
uint8x16_t vec_a_0 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 0 * 16);
|
||||
uint8x16_t vec_a0_top = vshrq_n_u8(vec_a_0, 4);
|
||||
uint8x16_t vec_a0_bot = vandq_u8(vec_a_0, vec_mask);
|
||||
int8x16_t vec_v_0_left_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 0], vec_a0_top);
|
||||
int8x16_t vec_v_0_left_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 1], vec_a0_top);
|
||||
int8x16_t vec_v_0_right_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 2], vec_a0_bot);
|
||||
int8x16_t vec_v_0_right_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 3], vec_a0_bot);
|
||||
int8x16x2_t vec_v_left_0 = vzipq_s8(vec_v_0_left_tmp1, vec_v_0_left_tmp0);
|
||||
int8x16x2_t vec_v_right_0 = vzipq_s8(vec_v_0_right_tmp1, vec_v_0_right_tmp0);
|
||||
vec_c[0] += vec_v_left_0.val[0];
|
||||
vec_c[0] += vec_v_right_0.val[0];
|
||||
vec_c[1] += vec_v_left_0.val[1];
|
||||
vec_c[1] += vec_v_right_0.val[1];
|
||||
|
||||
uint8x16_t vec_a_1 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 1 * 16);
|
||||
uint8x16_t vec_a1_top = vshrq_n_u8(vec_a_1, 4);
|
||||
uint8x16_t vec_a1_bot = vandq_u8(vec_a_1, vec_mask);
|
||||
int8x16_t vec_v_1_left_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 0], vec_a1_top);
|
||||
int8x16_t vec_v_1_left_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 1], vec_a1_top);
|
||||
int8x16_t vec_v_1_right_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 2], vec_a1_bot);
|
||||
int8x16_t vec_v_1_right_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 3], vec_a1_bot);
|
||||
int8x16x2_t vec_v_left_1 = vzipq_s8(vec_v_1_left_tmp1, vec_v_1_left_tmp0);
|
||||
int8x16x2_t vec_v_right_1 = vzipq_s8(vec_v_1_right_tmp1, vec_v_1_right_tmp0);
|
||||
vec_c[2] += vec_v_left_1.val[0];
|
||||
vec_c[2] += vec_v_right_1.val[0];
|
||||
vec_c[3] += vec_v_left_1.val[1];
|
||||
vec_c[3] += vec_v_right_1.val[1];
|
||||
|
||||
uint8x16_t vec_a_2 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 2 * 16);
|
||||
uint8x16_t vec_a2_top = vshrq_n_u8(vec_a_2, 4);
|
||||
uint8x16_t vec_a2_bot = vandq_u8(vec_a_2, vec_mask);
|
||||
int8x16_t vec_v_2_left_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 0], vec_a2_top);
|
||||
int8x16_t vec_v_2_left_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 1], vec_a2_top);
|
||||
int8x16_t vec_v_2_right_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 2], vec_a2_bot);
|
||||
int8x16_t vec_v_2_right_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 3], vec_a2_bot);
|
||||
int8x16x2_t vec_v_left_2 = vzipq_s8(vec_v_2_left_tmp1, vec_v_2_left_tmp0);
|
||||
int8x16x2_t vec_v_right_2 = vzipq_s8(vec_v_2_right_tmp1, vec_v_2_right_tmp0);
|
||||
vec_c[4] += vec_v_left_2.val[0];
|
||||
vec_c[4] += vec_v_right_2.val[0];
|
||||
vec_c[5] += vec_v_left_2.val[1];
|
||||
vec_c[5] += vec_v_right_2.val[1];
|
||||
|
||||
uint8x16_t vec_a_3 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 3 * 16);
|
||||
uint8x16_t vec_a3_top = vshrq_n_u8(vec_a_3, 4);
|
||||
uint8x16_t vec_a3_bot = vandq_u8(vec_a_3, vec_mask);
|
||||
int8x16_t vec_v_3_left_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 0], vec_a3_top);
|
||||
int8x16_t vec_v_3_left_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 1], vec_a3_top);
|
||||
int8x16_t vec_v_3_right_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 2], vec_a3_bot);
|
||||
int8x16_t vec_v_3_right_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 3], vec_a3_bot);
|
||||
int8x16x2_t vec_v_left_3 = vzipq_s8(vec_v_3_left_tmp1, vec_v_3_left_tmp0);
|
||||
int8x16x2_t vec_v_right_3 = vzipq_s8(vec_v_3_right_tmp1, vec_v_3_right_tmp0);
|
||||
vec_c[6] += vec_v_left_3.val[0];
|
||||
vec_c[6] += vec_v_right_3.val[0];
|
||||
vec_c[7] += vec_v_left_3.val[1];
|
||||
vec_c[7] += vec_v_right_3.val[1];
|
||||
|
||||
}
|
||||
|
||||
int32x4_t vec_v_bot_low_low_0 = vmovl_s16(vget_low_s16(vec_c[0]));
|
||||
int32x4_t vec_v_bot_low_high_0 = vmovl_high_s16(vec_c[0]);
|
||||
vst1q_s32(c + i + 0, vld1q_s32(c + i + 0) + vec_v_bot_low_low_0);
|
||||
vst1q_s32(c + i + 4, vld1q_s32(c + i + 4) + vec_v_bot_low_high_0);
|
||||
int32x4_t vec_v_bot_low_low_1 = vmovl_s16(vget_low_s16(vec_c[1]));
|
||||
int32x4_t vec_v_bot_low_high_1 = vmovl_high_s16(vec_c[1]);
|
||||
vst1q_s32(c + i + 8, vld1q_s32(c + i + 8) + vec_v_bot_low_low_1);
|
||||
vst1q_s32(c + i + 12, vld1q_s32(c + i + 12) + vec_v_bot_low_high_1);
|
||||
int32x4_t vec_v_bot_low_low_2 = vmovl_s16(vget_low_s16(vec_c[2]));
|
||||
int32x4_t vec_v_bot_low_high_2 = vmovl_high_s16(vec_c[2]);
|
||||
vst1q_s32(c + i + 16, vld1q_s32(c + i + 16) + vec_v_bot_low_low_2);
|
||||
vst1q_s32(c + i + 20, vld1q_s32(c + i + 20) + vec_v_bot_low_high_2);
|
||||
int32x4_t vec_v_bot_low_low_3 = vmovl_s16(vget_low_s16(vec_c[3]));
|
||||
int32x4_t vec_v_bot_low_high_3 = vmovl_high_s16(vec_c[3]);
|
||||
vst1q_s32(c + i + 24, vld1q_s32(c + i + 24) + vec_v_bot_low_low_3);
|
||||
vst1q_s32(c + i + 28, vld1q_s32(c + i + 28) + vec_v_bot_low_high_3);
|
||||
int32x4_t vec_v_bot_low_low_4 = vmovl_s16(vget_low_s16(vec_c[4]));
|
||||
int32x4_t vec_v_bot_low_high_4 = vmovl_high_s16(vec_c[4]);
|
||||
vst1q_s32(c + i + 32, vld1q_s32(c + i + 32) + vec_v_bot_low_low_4);
|
||||
vst1q_s32(c + i + 36, vld1q_s32(c + i + 36) + vec_v_bot_low_high_4);
|
||||
int32x4_t vec_v_bot_low_low_5 = vmovl_s16(vget_low_s16(vec_c[5]));
|
||||
int32x4_t vec_v_bot_low_high_5 = vmovl_high_s16(vec_c[5]);
|
||||
vst1q_s32(c + i + 40, vld1q_s32(c + i + 40) + vec_v_bot_low_low_5);
|
||||
vst1q_s32(c + i + 44, vld1q_s32(c + i + 44) + vec_v_bot_low_high_5);
|
||||
int32x4_t vec_v_bot_low_low_6 = vmovl_s16(vget_low_s16(vec_c[6]));
|
||||
int32x4_t vec_v_bot_low_high_6 = vmovl_high_s16(vec_c[6]);
|
||||
vst1q_s32(c + i + 48, vld1q_s32(c + i + 48) + vec_v_bot_low_low_6);
|
||||
vst1q_s32(c + i + 52, vld1q_s32(c + i + 52) + vec_v_bot_low_high_6);
|
||||
int32x4_t vec_v_bot_low_low_7 = vmovl_s16(vget_low_s16(vec_c[7]));
|
||||
int32x4_t vec_v_bot_low_high_7 = vmovl_high_s16(vec_c[7]);
|
||||
vst1q_s32(c + i + 56, vld1q_s32(c + i + 56) + vec_v_bot_low_low_7);
|
||||
vst1q_s32(c + i + 60, vld1q_s32(c + i + 60) + vec_v_bot_low_high_7);
|
||||
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
int32_t qgemm_lut_14336_4096(void* A, void* LUT, void* Scales, void* LUT_Scales, void* C) {
|
||||
alignas(32) uint32_t CBits[BM14336_4096];
|
||||
memset(&(CBits[0]), 0, BM14336_4096 * sizeof(int32_t));
|
||||
#pragma unroll
|
||||
for (int32_t k_outer = 0; k_outer < 4096 / BBK14336_4096; ++k_outer) {
|
||||
tbl_impl_14336_4096((&(((int32_t*)CBits)[0])), (&(((int8_t*)LUT)[(k_outer * BBK14336_4096 / 2 * 32)])), (&(((uint8_t*)A)[(k_outer * BBK14336_4096 / 2 / 2 * BM14336_4096)])));
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i = 0; i < BM14336_4096; i++) {
|
||||
((bitnet_float_type*)C)[i] = (((int32_t*)CBits)[i]) / ((bitnet_float_type*)LUT_Scales)[0] * ((bitnet_float_type*)Scales)[0];
|
||||
}
|
||||
return 0;
|
||||
};
|
||||
#include <arm_neon.h>
|
||||
|
||||
#define BM4096_14336 256
|
||||
#define BBK4096_14336 128
|
||||
inline void tbl_impl_4096_14336(int32_t* c, int8_t* lut, uint8_t* a) {
|
||||
#ifdef __ARM_NEON
|
||||
const int KK = BBK4096_14336 / 2;
|
||||
const uint8x16_t vec_mask = vdupq_n_u8(0x0f);
|
||||
const int8x16_t vec_zero = vdupq_n_s16(0x0000);
|
||||
int8x16_t vec_lut[2 * KK];
|
||||
int16x8_t vec_c[4];
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 2 * KK; k++) {
|
||||
vec_lut[k] = vld1q_s8(lut + k * 16);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < BM4096_14336; i += 32) {
|
||||
#pragma unroll
|
||||
for (int i=0; i<4; i++) {
|
||||
vec_c[i] = vandq_s16(vec_c[i], vec_zero);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int k = 0; k < KK / 4; k++) {
|
||||
|
||||
uint8x16_t vec_a_0 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 0 * 16);
|
||||
uint8x16_t vec_a0_top = vshrq_n_u8(vec_a_0, 4);
|
||||
uint8x16_t vec_a0_bot = vandq_u8(vec_a_0, vec_mask);
|
||||
int8x16_t vec_v_0_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 0], vec_a0_top);
|
||||
int8x16_t vec_v_0_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 1], vec_a0_top);
|
||||
int8x16_t vec_v_0_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 2], vec_a0_bot);
|
||||
int8x16_t vec_v_0_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 3], vec_a0_bot);
|
||||
int8x16x2_t vec_v_left_0 = vzipq_s8(vec_v_0_left_tmp1, vec_v_0_left_tmp0);
|
||||
int8x16x2_t vec_v_right_0 = vzipq_s8(vec_v_0_right_tmp1, vec_v_0_right_tmp0);
|
||||
vec_c[0] += vec_v_left_0.val[0];
|
||||
vec_c[0] += vec_v_right_0.val[0];
|
||||
vec_c[1] += vec_v_left_0.val[1];
|
||||
vec_c[1] += vec_v_right_0.val[1];
|
||||
|
||||
uint8x16_t vec_a_1 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 1 * 16);
|
||||
uint8x16_t vec_a1_top = vshrq_n_u8(vec_a_1, 4);
|
||||
uint8x16_t vec_a1_bot = vandq_u8(vec_a_1, vec_mask);
|
||||
int8x16_t vec_v_1_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 4], vec_a1_top);
|
||||
int8x16_t vec_v_1_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 5], vec_a1_top);
|
||||
int8x16_t vec_v_1_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 6], vec_a1_bot);
|
||||
int8x16_t vec_v_1_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 7], vec_a1_bot);
|
||||
int8x16x2_t vec_v_left_1 = vzipq_s8(vec_v_1_left_tmp1, vec_v_1_left_tmp0);
|
||||
int8x16x2_t vec_v_right_1 = vzipq_s8(vec_v_1_right_tmp1, vec_v_1_right_tmp0);
|
||||
vec_c[0] += vec_v_left_1.val[0];
|
||||
vec_c[0] += vec_v_right_1.val[0];
|
||||
vec_c[1] += vec_v_left_1.val[1];
|
||||
vec_c[1] += vec_v_right_1.val[1];
|
||||
|
||||
uint8x16_t vec_a_2 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 2 * 16);
|
||||
uint8x16_t vec_a2_top = vshrq_n_u8(vec_a_2, 4);
|
||||
uint8x16_t vec_a2_bot = vandq_u8(vec_a_2, vec_mask);
|
||||
int8x16_t vec_v_2_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 0], vec_a2_top);
|
||||
int8x16_t vec_v_2_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 1], vec_a2_top);
|
||||
int8x16_t vec_v_2_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 2], vec_a2_bot);
|
||||
int8x16_t vec_v_2_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 3], vec_a2_bot);
|
||||
int8x16x2_t vec_v_left_2 = vzipq_s8(vec_v_2_left_tmp1, vec_v_2_left_tmp0);
|
||||
int8x16x2_t vec_v_right_2 = vzipq_s8(vec_v_2_right_tmp1, vec_v_2_right_tmp0);
|
||||
vec_c[2] += vec_v_left_2.val[0];
|
||||
vec_c[2] += vec_v_right_2.val[0];
|
||||
vec_c[3] += vec_v_left_2.val[1];
|
||||
vec_c[3] += vec_v_right_2.val[1];
|
||||
|
||||
uint8x16_t vec_a_3 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 3 * 16);
|
||||
uint8x16_t vec_a3_top = vshrq_n_u8(vec_a_3, 4);
|
||||
uint8x16_t vec_a3_bot = vandq_u8(vec_a_3, vec_mask);
|
||||
int8x16_t vec_v_3_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 4], vec_a3_top);
|
||||
int8x16_t vec_v_3_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 5], vec_a3_top);
|
||||
int8x16_t vec_v_3_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 6], vec_a3_bot);
|
||||
int8x16_t vec_v_3_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 7], vec_a3_bot);
|
||||
int8x16x2_t vec_v_left_3 = vzipq_s8(vec_v_3_left_tmp1, vec_v_3_left_tmp0);
|
||||
int8x16x2_t vec_v_right_3 = vzipq_s8(vec_v_3_right_tmp1, vec_v_3_right_tmp0);
|
||||
vec_c[2] += vec_v_left_3.val[0];
|
||||
vec_c[2] += vec_v_right_3.val[0];
|
||||
vec_c[3] += vec_v_left_3.val[1];
|
||||
vec_c[3] += vec_v_right_3.val[1];
|
||||
|
||||
}
|
||||
|
||||
int32x4_t vec_v_bot_low_low_0 = vmovl_s16(vget_low_s16(vec_c[0]));
|
||||
int32x4_t vec_v_bot_low_high_0 = vmovl_high_s16(vec_c[0]);
|
||||
vst1q_s32(c + i + 0, vld1q_s32(c + i + 0) + vec_v_bot_low_low_0);
|
||||
vst1q_s32(c + i + 4, vld1q_s32(c + i + 4) + vec_v_bot_low_high_0);
|
||||
int32x4_t vec_v_bot_low_low_1 = vmovl_s16(vget_low_s16(vec_c[1]));
|
||||
int32x4_t vec_v_bot_low_high_1 = vmovl_high_s16(vec_c[1]);
|
||||
vst1q_s32(c + i + 8, vld1q_s32(c + i + 8) + vec_v_bot_low_low_1);
|
||||
vst1q_s32(c + i + 12, vld1q_s32(c + i + 12) + vec_v_bot_low_high_1);
|
||||
int32x4_t vec_v_bot_low_low_2 = vmovl_s16(vget_low_s16(vec_c[2]));
|
||||
int32x4_t vec_v_bot_low_high_2 = vmovl_high_s16(vec_c[2]);
|
||||
vst1q_s32(c + i + 16, vld1q_s32(c + i + 16) + vec_v_bot_low_low_2);
|
||||
vst1q_s32(c + i + 20, vld1q_s32(c + i + 20) + vec_v_bot_low_high_2);
|
||||
int32x4_t vec_v_bot_low_low_3 = vmovl_s16(vget_low_s16(vec_c[3]));
|
||||
int32x4_t vec_v_bot_low_high_3 = vmovl_high_s16(vec_c[3]);
|
||||
vst1q_s32(c + i + 24, vld1q_s32(c + i + 24) + vec_v_bot_low_low_3);
|
||||
vst1q_s32(c + i + 28, vld1q_s32(c + i + 28) + vec_v_bot_low_high_3);
|
||||
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
int32_t qgemm_lut_4096_14336(void* A, void* LUT, void* Scales, void* LUT_Scales, void* C) {
|
||||
alignas(32) uint32_t CBits[BM4096_14336];
|
||||
memset(&(CBits[0]), 0, BM4096_14336 * sizeof(int32_t));
|
||||
#pragma unroll
|
||||
for (int32_t k_outer = 0; k_outer < 14336 / BBK4096_14336; ++k_outer) {
|
||||
tbl_impl_4096_14336((&(((int32_t*)CBits)[0])), (&(((int8_t*)LUT)[(k_outer * BBK4096_14336 / 2 * 32)])), (&(((uint8_t*)A)[(k_outer * BBK4096_14336 / 2 / 2 * BM4096_14336)])));
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i = 0; i < BM4096_14336; i++) {
|
||||
((bitnet_float_type*)C)[i] = (((int32_t*)CBits)[i]) / ((bitnet_float_type*)LUT_Scales)[0] * ((bitnet_float_type*)Scales)[0];
|
||||
}
|
||||
return 0;
|
||||
};
|
||||
#include <arm_neon.h>
|
||||
|
||||
#define BM1024_4096 128
|
||||
#define BBK1024_4096 64
|
||||
inline void tbl_impl_1024_4096(int32_t* c, int8_t* lut, uint8_t* a) {
|
||||
#ifdef __ARM_NEON
|
||||
const int KK = BBK1024_4096 / 2;
|
||||
const uint8x16_t vec_mask = vdupq_n_u8(0x0f);
|
||||
const int8x16_t vec_zero = vdupq_n_s16(0x0000);
|
||||
int8x16_t vec_lut[2 * KK];
|
||||
int16x8_t vec_c[8];
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 2 * KK; k++) {
|
||||
vec_lut[k] = vld1q_s8(lut + k * 16);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < BM1024_4096; i += 64) {
|
||||
#pragma unroll
|
||||
for (int i=0; i<8; i++) {
|
||||
vec_c[i] = vandq_s16(vec_c[i], vec_zero);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int k = 0; k < KK / 2; k++) {
|
||||
|
||||
uint8x16_t vec_a_0 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 0 * 16);
|
||||
uint8x16_t vec_a0_top = vshrq_n_u8(vec_a_0, 4);
|
||||
uint8x16_t vec_a0_bot = vandq_u8(vec_a_0, vec_mask);
|
||||
int8x16_t vec_v_0_left_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 0], vec_a0_top);
|
||||
int8x16_t vec_v_0_left_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 1], vec_a0_top);
|
||||
int8x16_t vec_v_0_right_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 2], vec_a0_bot);
|
||||
int8x16_t vec_v_0_right_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 3], vec_a0_bot);
|
||||
int8x16x2_t vec_v_left_0 = vzipq_s8(vec_v_0_left_tmp1, vec_v_0_left_tmp0);
|
||||
int8x16x2_t vec_v_right_0 = vzipq_s8(vec_v_0_right_tmp1, vec_v_0_right_tmp0);
|
||||
vec_c[0] += vec_v_left_0.val[0];
|
||||
vec_c[0] += vec_v_right_0.val[0];
|
||||
vec_c[1] += vec_v_left_0.val[1];
|
||||
vec_c[1] += vec_v_right_0.val[1];
|
||||
|
||||
uint8x16_t vec_a_1 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 1 * 16);
|
||||
uint8x16_t vec_a1_top = vshrq_n_u8(vec_a_1, 4);
|
||||
uint8x16_t vec_a1_bot = vandq_u8(vec_a_1, vec_mask);
|
||||
int8x16_t vec_v_1_left_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 0], vec_a1_top);
|
||||
int8x16_t vec_v_1_left_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 1], vec_a1_top);
|
||||
int8x16_t vec_v_1_right_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 2], vec_a1_bot);
|
||||
int8x16_t vec_v_1_right_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 3], vec_a1_bot);
|
||||
int8x16x2_t vec_v_left_1 = vzipq_s8(vec_v_1_left_tmp1, vec_v_1_left_tmp0);
|
||||
int8x16x2_t vec_v_right_1 = vzipq_s8(vec_v_1_right_tmp1, vec_v_1_right_tmp0);
|
||||
vec_c[2] += vec_v_left_1.val[0];
|
||||
vec_c[2] += vec_v_right_1.val[0];
|
||||
vec_c[3] += vec_v_left_1.val[1];
|
||||
vec_c[3] += vec_v_right_1.val[1];
|
||||
|
||||
uint8x16_t vec_a_2 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 2 * 16);
|
||||
uint8x16_t vec_a2_top = vshrq_n_u8(vec_a_2, 4);
|
||||
uint8x16_t vec_a2_bot = vandq_u8(vec_a_2, vec_mask);
|
||||
int8x16_t vec_v_2_left_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 0], vec_a2_top);
|
||||
int8x16_t vec_v_2_left_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 1], vec_a2_top);
|
||||
int8x16_t vec_v_2_right_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 2], vec_a2_bot);
|
||||
int8x16_t vec_v_2_right_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 3], vec_a2_bot);
|
||||
int8x16x2_t vec_v_left_2 = vzipq_s8(vec_v_2_left_tmp1, vec_v_2_left_tmp0);
|
||||
int8x16x2_t vec_v_right_2 = vzipq_s8(vec_v_2_right_tmp1, vec_v_2_right_tmp0);
|
||||
vec_c[4] += vec_v_left_2.val[0];
|
||||
vec_c[4] += vec_v_right_2.val[0];
|
||||
vec_c[5] += vec_v_left_2.val[1];
|
||||
vec_c[5] += vec_v_right_2.val[1];
|
||||
|
||||
uint8x16_t vec_a_3 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 3 * 16);
|
||||
uint8x16_t vec_a3_top = vshrq_n_u8(vec_a_3, 4);
|
||||
uint8x16_t vec_a3_bot = vandq_u8(vec_a_3, vec_mask);
|
||||
int8x16_t vec_v_3_left_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 0], vec_a3_top);
|
||||
int8x16_t vec_v_3_left_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 1], vec_a3_top);
|
||||
int8x16_t vec_v_3_right_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 2], vec_a3_bot);
|
||||
int8x16_t vec_v_3_right_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 3], vec_a3_bot);
|
||||
int8x16x2_t vec_v_left_3 = vzipq_s8(vec_v_3_left_tmp1, vec_v_3_left_tmp0);
|
||||
int8x16x2_t vec_v_right_3 = vzipq_s8(vec_v_3_right_tmp1, vec_v_3_right_tmp0);
|
||||
vec_c[6] += vec_v_left_3.val[0];
|
||||
vec_c[6] += vec_v_right_3.val[0];
|
||||
vec_c[7] += vec_v_left_3.val[1];
|
||||
vec_c[7] += vec_v_right_3.val[1];
|
||||
|
||||
}
|
||||
|
||||
int32x4_t vec_v_bot_low_low_0 = vmovl_s16(vget_low_s16(vec_c[0]));
|
||||
int32x4_t vec_v_bot_low_high_0 = vmovl_high_s16(vec_c[0]);
|
||||
vst1q_s32(c + i + 0, vld1q_s32(c + i + 0) + vec_v_bot_low_low_0);
|
||||
vst1q_s32(c + i + 4, vld1q_s32(c + i + 4) + vec_v_bot_low_high_0);
|
||||
int32x4_t vec_v_bot_low_low_1 = vmovl_s16(vget_low_s16(vec_c[1]));
|
||||
int32x4_t vec_v_bot_low_high_1 = vmovl_high_s16(vec_c[1]);
|
||||
vst1q_s32(c + i + 8, vld1q_s32(c + i + 8) + vec_v_bot_low_low_1);
|
||||
vst1q_s32(c + i + 12, vld1q_s32(c + i + 12) + vec_v_bot_low_high_1);
|
||||
int32x4_t vec_v_bot_low_low_2 = vmovl_s16(vget_low_s16(vec_c[2]));
|
||||
int32x4_t vec_v_bot_low_high_2 = vmovl_high_s16(vec_c[2]);
|
||||
vst1q_s32(c + i + 16, vld1q_s32(c + i + 16) + vec_v_bot_low_low_2);
|
||||
vst1q_s32(c + i + 20, vld1q_s32(c + i + 20) + vec_v_bot_low_high_2);
|
||||
int32x4_t vec_v_bot_low_low_3 = vmovl_s16(vget_low_s16(vec_c[3]));
|
||||
int32x4_t vec_v_bot_low_high_3 = vmovl_high_s16(vec_c[3]);
|
||||
vst1q_s32(c + i + 24, vld1q_s32(c + i + 24) + vec_v_bot_low_low_3);
|
||||
vst1q_s32(c + i + 28, vld1q_s32(c + i + 28) + vec_v_bot_low_high_3);
|
||||
int32x4_t vec_v_bot_low_low_4 = vmovl_s16(vget_low_s16(vec_c[4]));
|
||||
int32x4_t vec_v_bot_low_high_4 = vmovl_high_s16(vec_c[4]);
|
||||
vst1q_s32(c + i + 32, vld1q_s32(c + i + 32) + vec_v_bot_low_low_4);
|
||||
vst1q_s32(c + i + 36, vld1q_s32(c + i + 36) + vec_v_bot_low_high_4);
|
||||
int32x4_t vec_v_bot_low_low_5 = vmovl_s16(vget_low_s16(vec_c[5]));
|
||||
int32x4_t vec_v_bot_low_high_5 = vmovl_high_s16(vec_c[5]);
|
||||
vst1q_s32(c + i + 40, vld1q_s32(c + i + 40) + vec_v_bot_low_low_5);
|
||||
vst1q_s32(c + i + 44, vld1q_s32(c + i + 44) + vec_v_bot_low_high_5);
|
||||
int32x4_t vec_v_bot_low_low_6 = vmovl_s16(vget_low_s16(vec_c[6]));
|
||||
int32x4_t vec_v_bot_low_high_6 = vmovl_high_s16(vec_c[6]);
|
||||
vst1q_s32(c + i + 48, vld1q_s32(c + i + 48) + vec_v_bot_low_low_6);
|
||||
vst1q_s32(c + i + 52, vld1q_s32(c + i + 52) + vec_v_bot_low_high_6);
|
||||
int32x4_t vec_v_bot_low_low_7 = vmovl_s16(vget_low_s16(vec_c[7]));
|
||||
int32x4_t vec_v_bot_low_high_7 = vmovl_high_s16(vec_c[7]);
|
||||
vst1q_s32(c + i + 56, vld1q_s32(c + i + 56) + vec_v_bot_low_low_7);
|
||||
vst1q_s32(c + i + 60, vld1q_s32(c + i + 60) + vec_v_bot_low_high_7);
|
||||
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
int32_t qgemm_lut_1024_4096(void* A, void* LUT, void* Scales, void* LUT_Scales, void* C) {
|
||||
alignas(32) uint32_t CBits[BM1024_4096];
|
||||
memset(&(CBits[0]), 0, BM1024_4096 * sizeof(int32_t));
|
||||
#pragma unroll
|
||||
for (int32_t k_outer = 0; k_outer < 4096 / BBK1024_4096; ++k_outer) {
|
||||
tbl_impl_1024_4096((&(((int32_t*)CBits)[0])), (&(((int8_t*)LUT)[(k_outer * BBK1024_4096 / 2 * 32)])), (&(((uint8_t*)A)[(k_outer * BBK1024_4096 / 2 / 2 * BM1024_4096)])));
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i = 0; i < BM1024_4096; i++) {
|
||||
((bitnet_float_type*)C)[i] = (((int32_t*)CBits)[i]) / ((bitnet_float_type*)LUT_Scales)[0] * ((bitnet_float_type*)Scales)[0];
|
||||
}
|
||||
return 0;
|
||||
};
|
||||
#include <arm_neon.h>
|
||||
|
||||
#define BM4096_4096 128
|
||||
#define BBK4096_4096 64
|
||||
inline void tbl_impl_4096_4096(int32_t* c, int8_t* lut, uint8_t* a) {
|
||||
#ifdef __ARM_NEON
|
||||
const int KK = BBK4096_4096 / 2;
|
||||
const uint8x16_t vec_mask = vdupq_n_u8(0x0f);
|
||||
const int8x16_t vec_zero = vdupq_n_s16(0x0000);
|
||||
int8x16_t vec_lut[2 * KK];
|
||||
int16x8_t vec_c[4];
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 2 * KK; k++) {
|
||||
vec_lut[k] = vld1q_s8(lut + k * 16);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < BM4096_4096; i += 32) {
|
||||
#pragma unroll
|
||||
for (int i=0; i<4; i++) {
|
||||
vec_c[i] = vandq_s16(vec_c[i], vec_zero);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int k = 0; k < KK / 4; k++) {
|
||||
|
||||
uint8x16_t vec_a_0 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 0 * 16);
|
||||
uint8x16_t vec_a0_top = vshrq_n_u8(vec_a_0, 4);
|
||||
uint8x16_t vec_a0_bot = vandq_u8(vec_a_0, vec_mask);
|
||||
int8x16_t vec_v_0_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 0], vec_a0_top);
|
||||
int8x16_t vec_v_0_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 1], vec_a0_top);
|
||||
int8x16_t vec_v_0_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 2], vec_a0_bot);
|
||||
int8x16_t vec_v_0_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 3], vec_a0_bot);
|
||||
int8x16x2_t vec_v_left_0 = vzipq_s8(vec_v_0_left_tmp1, vec_v_0_left_tmp0);
|
||||
int8x16x2_t vec_v_right_0 = vzipq_s8(vec_v_0_right_tmp1, vec_v_0_right_tmp0);
|
||||
vec_c[0] += vec_v_left_0.val[0];
|
||||
vec_c[0] += vec_v_right_0.val[0];
|
||||
vec_c[1] += vec_v_left_0.val[1];
|
||||
vec_c[1] += vec_v_right_0.val[1];
|
||||
|
||||
uint8x16_t vec_a_1 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 1 * 16);
|
||||
uint8x16_t vec_a1_top = vshrq_n_u8(vec_a_1, 4);
|
||||
uint8x16_t vec_a1_bot = vandq_u8(vec_a_1, vec_mask);
|
||||
int8x16_t vec_v_1_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 4], vec_a1_top);
|
||||
int8x16_t vec_v_1_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 5], vec_a1_top);
|
||||
int8x16_t vec_v_1_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 6], vec_a1_bot);
|
||||
int8x16_t vec_v_1_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 7], vec_a1_bot);
|
||||
int8x16x2_t vec_v_left_1 = vzipq_s8(vec_v_1_left_tmp1, vec_v_1_left_tmp0);
|
||||
int8x16x2_t vec_v_right_1 = vzipq_s8(vec_v_1_right_tmp1, vec_v_1_right_tmp0);
|
||||
vec_c[0] += vec_v_left_1.val[0];
|
||||
vec_c[0] += vec_v_right_1.val[0];
|
||||
vec_c[1] += vec_v_left_1.val[1];
|
||||
vec_c[1] += vec_v_right_1.val[1];
|
||||
|
||||
uint8x16_t vec_a_2 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 2 * 16);
|
||||
uint8x16_t vec_a2_top = vshrq_n_u8(vec_a_2, 4);
|
||||
uint8x16_t vec_a2_bot = vandq_u8(vec_a_2, vec_mask);
|
||||
int8x16_t vec_v_2_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 0], vec_a2_top);
|
||||
int8x16_t vec_v_2_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 1], vec_a2_top);
|
||||
int8x16_t vec_v_2_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 2], vec_a2_bot);
|
||||
int8x16_t vec_v_2_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 3], vec_a2_bot);
|
||||
int8x16x2_t vec_v_left_2 = vzipq_s8(vec_v_2_left_tmp1, vec_v_2_left_tmp0);
|
||||
int8x16x2_t vec_v_right_2 = vzipq_s8(vec_v_2_right_tmp1, vec_v_2_right_tmp0);
|
||||
vec_c[2] += vec_v_left_2.val[0];
|
||||
vec_c[2] += vec_v_right_2.val[0];
|
||||
vec_c[3] += vec_v_left_2.val[1];
|
||||
vec_c[3] += vec_v_right_2.val[1];
|
||||
|
||||
uint8x16_t vec_a_3 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 3 * 16);
|
||||
uint8x16_t vec_a3_top = vshrq_n_u8(vec_a_3, 4);
|
||||
uint8x16_t vec_a3_bot = vandq_u8(vec_a_3, vec_mask);
|
||||
int8x16_t vec_v_3_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 4], vec_a3_top);
|
||||
int8x16_t vec_v_3_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 5], vec_a3_top);
|
||||
int8x16_t vec_v_3_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 6], vec_a3_bot);
|
||||
int8x16_t vec_v_3_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 7], vec_a3_bot);
|
||||
int8x16x2_t vec_v_left_3 = vzipq_s8(vec_v_3_left_tmp1, vec_v_3_left_tmp0);
|
||||
int8x16x2_t vec_v_right_3 = vzipq_s8(vec_v_3_right_tmp1, vec_v_3_right_tmp0);
|
||||
vec_c[2] += vec_v_left_3.val[0];
|
||||
vec_c[2] += vec_v_right_3.val[0];
|
||||
vec_c[3] += vec_v_left_3.val[1];
|
||||
vec_c[3] += vec_v_right_3.val[1];
|
||||
|
||||
}
|
||||
|
||||
int32x4_t vec_v_bot_low_low_0 = vmovl_s16(vget_low_s16(vec_c[0]));
|
||||
int32x4_t vec_v_bot_low_high_0 = vmovl_high_s16(vec_c[0]);
|
||||
vst1q_s32(c + i + 0, vld1q_s32(c + i + 0) + vec_v_bot_low_low_0);
|
||||
vst1q_s32(c + i + 4, vld1q_s32(c + i + 4) + vec_v_bot_low_high_0);
|
||||
int32x4_t vec_v_bot_low_low_1 = vmovl_s16(vget_low_s16(vec_c[1]));
|
||||
int32x4_t vec_v_bot_low_high_1 = vmovl_high_s16(vec_c[1]);
|
||||
vst1q_s32(c + i + 8, vld1q_s32(c + i + 8) + vec_v_bot_low_low_1);
|
||||
vst1q_s32(c + i + 12, vld1q_s32(c + i + 12) + vec_v_bot_low_high_1);
|
||||
int32x4_t vec_v_bot_low_low_2 = vmovl_s16(vget_low_s16(vec_c[2]));
|
||||
int32x4_t vec_v_bot_low_high_2 = vmovl_high_s16(vec_c[2]);
|
||||
vst1q_s32(c + i + 16, vld1q_s32(c + i + 16) + vec_v_bot_low_low_2);
|
||||
vst1q_s32(c + i + 20, vld1q_s32(c + i + 20) + vec_v_bot_low_high_2);
|
||||
int32x4_t vec_v_bot_low_low_3 = vmovl_s16(vget_low_s16(vec_c[3]));
|
||||
int32x4_t vec_v_bot_low_high_3 = vmovl_high_s16(vec_c[3]);
|
||||
vst1q_s32(c + i + 24, vld1q_s32(c + i + 24) + vec_v_bot_low_low_3);
|
||||
vst1q_s32(c + i + 28, vld1q_s32(c + i + 28) + vec_v_bot_low_high_3);
|
||||
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
int32_t qgemm_lut_4096_4096(void* A, void* LUT, void* Scales, void* LUT_Scales, void* C) {
|
||||
alignas(32) uint32_t CBits[BM4096_4096];
|
||||
memset(&(CBits[0]), 0, BM4096_4096 * sizeof(int32_t));
|
||||
#pragma unroll
|
||||
for (int32_t k_outer = 0; k_outer < 4096 / BBK4096_4096; ++k_outer) {
|
||||
tbl_impl_4096_4096((&(((int32_t*)CBits)[0])), (&(((int8_t*)LUT)[(k_outer * BBK4096_4096 / 2 * 32)])), (&(((uint8_t*)A)[(k_outer * BBK4096_4096 / 2 / 2 * BM4096_4096)])));
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i = 0; i < BM4096_4096; i++) {
|
||||
((bitnet_float_type*)C)[i] = (((int32_t*)CBits)[i]) / ((bitnet_float_type*)LUT_Scales)[0] * ((bitnet_float_type*)Scales)[0];
|
||||
}
|
||||
return 0;
|
||||
};
|
||||
|
||||
template<int K>
|
||||
void preprocessor_k(void* B, void* LUT_Scales, void* QLUT) {{
|
||||
partial_max_reset((&(((bitnet_float_type*)LUT_Scales)[0])));
|
||||
per_tensor_quant(K, (&(((bitnet_float_type*)LUT_Scales)[0])), (&(((bitnet_float_type*)B)[0])));
|
||||
|
||||
lut_ctor<K>((&(((int8_t*)QLUT)[0])), (&(((bitnet_float_type*)B)[0])), (&(((bitnet_float_type*)LUT_Scales)[0])));
|
||||
}}
|
||||
void ggml_preprocessor(int m, int k, void* B, void* LUT_Scales, void* QLUT) {
|
||||
if (m == 14336 && k == 4096) {
|
||||
preprocessor_k<4096>(B, LUT_Scales, QLUT);
|
||||
}
|
||||
else if (m == 4096 && k == 14336) {
|
||||
preprocessor_k<14336>(B, LUT_Scales, QLUT);
|
||||
}
|
||||
else if (m == 1024 && k == 4096) {
|
||||
preprocessor_k<4096>(B, LUT_Scales, QLUT);
|
||||
}
|
||||
else if (m == 4096 && k == 4096) {
|
||||
preprocessor_k<4096>(B, LUT_Scales, QLUT);
|
||||
}
|
||||
}
|
||||
void ggml_qgemm_lut(int m, int k, void* A, void* LUT, void* Scales, void* LUT_Scales, void* C) {
|
||||
if (m == 14336 && k == 4096) {
|
||||
qgemm_lut_14336_4096(A, LUT, Scales, LUT_Scales, C);
|
||||
}
|
||||
else if (m == 4096 && k == 14336) {
|
||||
qgemm_lut_4096_14336(A, LUT, Scales, LUT_Scales, C);
|
||||
}
|
||||
else if (m == 1024 && k == 4096) {
|
||||
qgemm_lut_1024_4096(A, LUT, Scales, LUT_Scales, C);
|
||||
}
|
||||
else if (m == 4096 && k == 4096) {
|
||||
qgemm_lut_4096_4096(A, LUT, Scales, LUT_Scales, C);
|
||||
}
|
||||
}
|
||||
|
||||
void ggml_bitnet_transform_tensor(struct ggml_tensor * tensor) {
|
||||
if (!(is_type_supported(tensor->type) && tensor->backend == GGML_BACKEND_TYPE_CPU && tensor->extra == nullptr)) {
|
||||
return;
|
||||
}
|
||||
|
||||
int k = tensor->ne[0];
|
||||
int m = tensor->ne[1];
|
||||
const int lut_scales_size = 1;
|
||||
const int scales_size = 1;
|
||||
int bk = 0;
|
||||
int bm = 0;
|
||||
|
||||
if (m == 14336 && k == 4096) {
|
||||
bm = BM14336_4096;
|
||||
bk = BBK14336_4096;
|
||||
}
|
||||
else if (m == 4096 && k == 14336) {
|
||||
bm = BM4096_14336;
|
||||
bk = BBK4096_14336;
|
||||
}
|
||||
else if (m == 1024 && k == 4096) {
|
||||
bm = BM1024_4096;
|
||||
bk = BBK1024_4096;
|
||||
}
|
||||
else if (m == 4096 && k == 4096) {
|
||||
bm = BM4096_4096;
|
||||
bk = BBK4096_4096;
|
||||
}
|
||||
|
||||
const int n_tile_num = m / bm;
|
||||
const int BK = bk;
|
||||
uint8_t * qweights;
|
||||
bitnet_float_type * scales;
|
||||
|
||||
scales = (bitnet_float_type *) aligned_malloc(sizeof(bitnet_float_type));
|
||||
qweights = (uint8_t *) tensor->data;
|
||||
float * i2_scales = (float * )(qweights + k * m / 4);
|
||||
scales[0] = (bitnet_float_type) i2_scales[0];
|
||||
|
||||
tensor->extra = bitnet_tensor_extras + bitnet_tensor_extras_index;
|
||||
bitnet_tensor_extras[bitnet_tensor_extras_index++] = {
|
||||
/* .lut_scales_size = */ lut_scales_size,
|
||||
/* .scales_size = */ scales_size,
|
||||
/* .n_tile_num = */ n_tile_num,
|
||||
/* .qweights = */ qweights,
|
||||
/* .scales = */ scales
|
||||
};
|
||||
}
|
||||
#endif
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,28 @@
|
||||
[Kernels_0]
|
||||
m = 14336
|
||||
k = 4096
|
||||
bm = 256
|
||||
bk = 128
|
||||
bmm = 64
|
||||
|
||||
[Kernels_1]
|
||||
m = 4096
|
||||
k = 14336
|
||||
bm = 256
|
||||
bk = 128
|
||||
bmm = 32
|
||||
|
||||
[Kernels_2]
|
||||
m = 1024
|
||||
k = 4096
|
||||
bm = 128
|
||||
bk = 64
|
||||
bmm = 64
|
||||
|
||||
[Kernels_3]
|
||||
m = 4096
|
||||
k = 4096
|
||||
bm = 128
|
||||
bk = 64
|
||||
bmm = 32
|
||||
|
||||
@@ -0,0 +1,28 @@
|
||||
[Kernels_0]
|
||||
m = 14336
|
||||
k = 4096
|
||||
bm = 256
|
||||
bk = 96
|
||||
bmm = 32
|
||||
|
||||
[Kernels_1]
|
||||
m = 4096
|
||||
k = 14336
|
||||
bm = 128
|
||||
bk = 96
|
||||
bmm = 32
|
||||
|
||||
[Kernels_2]
|
||||
m = 1024
|
||||
k = 4096
|
||||
bm = 256
|
||||
bk = 96
|
||||
bmm = 32
|
||||
|
||||
[Kernels_3]
|
||||
m = 4096
|
||||
k = 4096
|
||||
bm = 128
|
||||
bk = 96
|
||||
bmm = 32
|
||||
|
||||
@@ -0,0 +1,627 @@
|
||||
#if defined(GGML_BITNET_ARM_TL1)
|
||||
#include "ggml-bitnet.h"
|
||||
#define GGML_BITNET_MAX_NODES 8192
|
||||
static bool initialized = false;
|
||||
static bitnet_tensor_extra * bitnet_tensor_extras = nullptr;
|
||||
static size_t bitnet_tensor_extras_index = 0;
|
||||
static void * aligned_malloc(size_t size) {{
|
||||
#if defined(_WIN32)
|
||||
return _aligned_malloc(size, 64);
|
||||
#else
|
||||
void * ptr = nullptr;
|
||||
posix_memalign(&ptr, 64, size);
|
||||
return ptr;
|
||||
#endif
|
||||
}}
|
||||
static void aligned_free(void * ptr) {{
|
||||
#if defined(_WIN32)
|
||||
_aligned_free(ptr);
|
||||
#else
|
||||
free(ptr);
|
||||
#endif
|
||||
}}
|
||||
|
||||
void per_tensor_quant(int k, void* lut_scales_, void* b_) {{
|
||||
bitnet_float_type* lut_scales = (bitnet_float_type*)lut_scales_;
|
||||
bitnet_float_type* b = (bitnet_float_type*)b_;
|
||||
#ifdef __ARM_NEON
|
||||
float32x4_t temp_max = vdupq_n_f32(0);
|
||||
for (int i=0; i < k / 4; i++) {{
|
||||
float32x4_t vec_bs = vld1q_f32(b + 4 * i);
|
||||
float32x4_t abssum = vabsq_f32(vec_bs);
|
||||
temp_max = vmaxq_f32(abssum, temp_max);
|
||||
}}
|
||||
float32_t scales = 127 / vmaxvq_f32(temp_max);
|
||||
*lut_scales = scales;
|
||||
#elif defined __AVX2__
|
||||
__m256 max_vec = _mm256_set1_ps(0.f);
|
||||
const __m256 vec_sign = _mm256_set1_ps(-0.0f);
|
||||
// #pragma unroll
|
||||
for (int i = 0; i < k / 8; i++) {{
|
||||
__m256 vec_b = _mm256_loadu_ps(b + i * 8);
|
||||
__m256 vec_babs = _mm256_andnot_ps(vec_sign, vec_b);
|
||||
max_vec = _mm256_max_ps(vec_babs, max_vec);
|
||||
}}
|
||||
__m128 max1 = _mm_max_ps(_mm256_extractf128_ps(max_vec, 1), _mm256_castps256_ps128(max_vec));
|
||||
max1 = _mm_max_ps(max1, _mm_movehl_ps(max1, max1));
|
||||
max1 = _mm_max_ss(max1, _mm_movehdup_ps(max1));
|
||||
float scales = 127 / _mm_cvtss_f32(max1);
|
||||
*lut_scales = scales;
|
||||
#endif
|
||||
}}
|
||||
|
||||
void partial_max_reset(void* lut_scales_) {{
|
||||
bitnet_float_type* lut_scales = (bitnet_float_type*)lut_scales_;
|
||||
*lut_scales = 0.0;
|
||||
}}
|
||||
|
||||
#ifdef __ARM_NEON
|
||||
inline void Transpose_8_8(
|
||||
int16x8_t *v0,
|
||||
int16x8_t *v1,
|
||||
int16x8_t *v2,
|
||||
int16x8_t *v3,
|
||||
int16x8_t *v4,
|
||||
int16x8_t *v5,
|
||||
int16x8_t *v6,
|
||||
int16x8_t *v7)
|
||||
{{
|
||||
int16x8x2_t q04 = vzipq_s16(*v0, *v4);
|
||||
int16x8x2_t q15 = vzipq_s16(*v1, *v5);
|
||||
int16x8x2_t q26 = vzipq_s16(*v2, *v6);
|
||||
int16x8x2_t q37 = vzipq_s16(*v3, *v7);
|
||||
|
||||
int16x8x2_t q0246_0 = vzipq_s16(q04.val[0], q26.val[0]);
|
||||
int16x8x2_t q0246_1 = vzipq_s16(q04.val[1], q26.val[1]);
|
||||
int16x8x2_t q1357_0 = vzipq_s16(q15.val[0], q37.val[0]);
|
||||
int16x8x2_t q1357_1 = vzipq_s16(q15.val[1], q37.val[1]);
|
||||
|
||||
int16x8x2_t q_fin_0 = vzipq_s16(q0246_0.val[0], q1357_0.val[0]);
|
||||
int16x8x2_t q_fin_1 = vzipq_s16(q0246_0.val[1], q1357_0.val[1]);
|
||||
int16x8x2_t q_fin_2 = vzipq_s16(q0246_1.val[0], q1357_1.val[0]);
|
||||
int16x8x2_t q_fin_3 = vzipq_s16(q0246_1.val[1], q1357_1.val[1]);
|
||||
|
||||
*v0 = q_fin_0.val[0];
|
||||
*v1 = q_fin_0.val[1];
|
||||
*v2 = q_fin_1.val[0];
|
||||
*v3 = q_fin_1.val[1];
|
||||
*v4 = q_fin_2.val[0];
|
||||
*v5 = q_fin_2.val[1];
|
||||
*v6 = q_fin_3.val[0];
|
||||
*v7 = q_fin_3.val[1];
|
||||
}}
|
||||
#endif
|
||||
|
||||
template<int act_k>
|
||||
inline void lut_ctor(int8_t* qlut, bitnet_float_type* b, bitnet_float_type* lut_scales) {{
|
||||
#ifdef __ARM_NEON
|
||||
int16x8_t vec_lut[16];
|
||||
float32_t scales = *lut_scales;
|
||||
uint8_t tbl_mask[16];
|
||||
tbl_mask[0] = 0;
|
||||
tbl_mask[1] = 2;
|
||||
tbl_mask[2] = 4;
|
||||
tbl_mask[3] = 6;
|
||||
tbl_mask[4] = 8;
|
||||
tbl_mask[5] = 10;
|
||||
tbl_mask[6] = 12;
|
||||
tbl_mask[7] = 14;
|
||||
tbl_mask[8] = 1;
|
||||
tbl_mask[9] = 3;
|
||||
tbl_mask[10] = 5;
|
||||
tbl_mask[11] = 7;
|
||||
tbl_mask[12] = 9;
|
||||
tbl_mask[13] = 11;
|
||||
tbl_mask[14] = 13;
|
||||
tbl_mask[15] = 15;
|
||||
uint8x16_t tbl_mask_q = vld1q_u8(tbl_mask);
|
||||
#pragma unroll
|
||||
for (int k = 0; k < act_k / 16; ++k) {{
|
||||
float32x4x2_t vec_bs_x0 = vld2q_f32(b + k * 16);
|
||||
float32x4x2_t vec_bs_x1 = vld2q_f32(b + k * 16 + 8);
|
||||
float32x4_t vec_f_0 = vmulq_n_f32(vec_bs_x0.val[0], scales);
|
||||
float32x4_t vec_f_1 = vmulq_n_f32(vec_bs_x0.val[1], scales);
|
||||
float32x4_t vec_f_2 = vmulq_n_f32(vec_bs_x1.val[0], scales);
|
||||
float32x4_t vec_f_3 = vmulq_n_f32(vec_bs_x1.val[1], scales);
|
||||
int32x4_t vec_b_0 = vcvtnq_s32_f32(vec_f_0);
|
||||
int32x4_t vec_b_1 = vcvtnq_s32_f32(vec_f_1);
|
||||
int32x4_t vec_b_2 = vcvtnq_s32_f32(vec_f_2);
|
||||
int32x4_t vec_b_3 = vcvtnq_s32_f32(vec_f_3);
|
||||
int16x4_t vec_b16_0 = vmovn_s32(vec_b_0);
|
||||
int16x4_t vec_b16_1 = vmovn_s32(vec_b_1);
|
||||
int16x4_t vec_b16_2 = vmovn_s32(vec_b_2);
|
||||
int16x4_t vec_b16_3 = vmovn_s32(vec_b_3);
|
||||
int16x8_t vec_bs_0 = vcombine_s16(vec_b16_0, vec_b16_2);
|
||||
int16x8_t vec_bs_1 = vcombine_s16(vec_b16_1, vec_b16_3);
|
||||
vec_lut[0] = vdupq_n_s16(0);
|
||||
vec_lut[0] = vec_lut[0] - vec_bs_0;
|
||||
vec_lut[0] = vec_lut[0] - vec_bs_1;
|
||||
vec_lut[1] = vdupq_n_s16(0);
|
||||
vec_lut[1] = vec_lut[1] - vec_bs_0;
|
||||
vec_lut[2] = vdupq_n_s16(0);
|
||||
vec_lut[2] = vec_lut[2] - vec_bs_0;
|
||||
vec_lut[2] = vec_lut[2] + vec_bs_1;
|
||||
vec_lut[3] = vdupq_n_s16(0);
|
||||
vec_lut[3] = vec_lut[3] - vec_bs_1;
|
||||
vec_lut[4] = vdupq_n_s16(0);
|
||||
vec_lut[5] = vec_bs_1;
|
||||
vec_lut[6] = vec_bs_0;
|
||||
vec_lut[6] = vec_lut[6] - vec_bs_1;
|
||||
vec_lut[7] = vec_bs_0;
|
||||
vec_lut[8] = vec_bs_0;
|
||||
vec_lut[8] = vec_lut[8] + vec_bs_1;
|
||||
Transpose_8_8(&(vec_lut[0]), &(vec_lut[1]), &(vec_lut[2]), &(vec_lut[3]),
|
||||
&(vec_lut[4]), &(vec_lut[5]), &(vec_lut[6]), &(vec_lut[7]));
|
||||
Transpose_8_8(&(vec_lut[8]), &(vec_lut[9]), &(vec_lut[10]), &(vec_lut[11]),
|
||||
&(vec_lut[12]), &(vec_lut[13]), &(vec_lut[14]), &(vec_lut[15]));
|
||||
#pragma unroll
|
||||
for (int idx = 0; idx < 8; idx++) {{
|
||||
int8x16_t q0_s = vqtbl1q_s8(vreinterpretq_s8_s16(vec_lut[idx]), tbl_mask_q);
|
||||
int8x8_t q0_low = vget_low_s8(q0_s);
|
||||
int8x8_t q0_high = vget_high_s8(q0_s);
|
||||
int8x16_t q1_s = vqtbl1q_s8(vreinterpretq_s8_s16(vec_lut[idx + 8]), tbl_mask_q);
|
||||
int8x8_t q1_low = vget_low_s8(q1_s);
|
||||
int8x8_t q1_high = vget_high_s8(q1_s);
|
||||
vst1_s8(qlut + k * 16 * 8 * 2 + idx * 16 * 2, q0_high);
|
||||
vst1_s8(qlut + k * 16 * 8 * 2 + idx * 16 * 2 + 8, q1_high);
|
||||
vst1_s8(qlut + k * 16 * 8 * 2 + idx * 16 * 2 + 16, q0_low);
|
||||
vst1_s8(qlut + k * 16 * 8 * 2 + idx * 16 * 2 + 24, q1_low);
|
||||
}}
|
||||
}}
|
||||
#endif
|
||||
}}
|
||||
|
||||
static bool is_type_supported(enum ggml_type type) {{
|
||||
if (type == GGML_TYPE_Q4_0 ||
|
||||
type == GGML_TYPE_TL1) {{
|
||||
return true;
|
||||
}} else {{
|
||||
return false;
|
||||
}}
|
||||
}}
|
||||
#include <arm_neon.h>
|
||||
|
||||
#define BM3200_8640 160
|
||||
#define BBK3200_8640 64
|
||||
inline void tbl_impl_3200_8640(int32_t* c, int8_t* lut, uint8_t* a) {
|
||||
#ifdef __ARM_NEON
|
||||
const int KK = BBK3200_8640 / 2;
|
||||
const uint8x16_t vec_mask = vdupq_n_u8(0x0f);
|
||||
const int8x16_t vec_zero = vdupq_n_s16(0x0000);
|
||||
int8x16_t vec_lut[2 * KK];
|
||||
int16x8_t vec_c[4];
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 2 * KK; k++) {
|
||||
vec_lut[k] = vld1q_s8(lut + k * 16);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < BM3200_8640; i += 32) {
|
||||
#pragma unroll
|
||||
for (int i=0; i<4; i++) {
|
||||
vec_c[i] = vandq_s16(vec_c[i], vec_zero);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int k = 0; k < KK / 4; k++) {
|
||||
|
||||
uint8x16_t vec_a_0 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 0 * 16);
|
||||
uint8x16_t vec_a0_top = vshrq_n_u8(vec_a_0, 4);
|
||||
uint8x16_t vec_a0_bot = vandq_u8(vec_a_0, vec_mask);
|
||||
int8x16_t vec_v_0_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 0], vec_a0_top);
|
||||
int8x16_t vec_v_0_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 1], vec_a0_top);
|
||||
int8x16_t vec_v_0_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 2], vec_a0_bot);
|
||||
int8x16_t vec_v_0_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 3], vec_a0_bot);
|
||||
int8x16x2_t vec_v_left_0 = vzipq_s8(vec_v_0_left_tmp1, vec_v_0_left_tmp0);
|
||||
int8x16x2_t vec_v_right_0 = vzipq_s8(vec_v_0_right_tmp1, vec_v_0_right_tmp0);
|
||||
vec_c[0] += vec_v_left_0.val[0];
|
||||
vec_c[0] += vec_v_right_0.val[0];
|
||||
vec_c[1] += vec_v_left_0.val[1];
|
||||
vec_c[1] += vec_v_right_0.val[1];
|
||||
|
||||
uint8x16_t vec_a_1 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 1 * 16);
|
||||
uint8x16_t vec_a1_top = vshrq_n_u8(vec_a_1, 4);
|
||||
uint8x16_t vec_a1_bot = vandq_u8(vec_a_1, vec_mask);
|
||||
int8x16_t vec_v_1_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 4], vec_a1_top);
|
||||
int8x16_t vec_v_1_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 5], vec_a1_top);
|
||||
int8x16_t vec_v_1_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 6], vec_a1_bot);
|
||||
int8x16_t vec_v_1_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 7], vec_a1_bot);
|
||||
int8x16x2_t vec_v_left_1 = vzipq_s8(vec_v_1_left_tmp1, vec_v_1_left_tmp0);
|
||||
int8x16x2_t vec_v_right_1 = vzipq_s8(vec_v_1_right_tmp1, vec_v_1_right_tmp0);
|
||||
vec_c[0] += vec_v_left_1.val[0];
|
||||
vec_c[0] += vec_v_right_1.val[0];
|
||||
vec_c[1] += vec_v_left_1.val[1];
|
||||
vec_c[1] += vec_v_right_1.val[1];
|
||||
|
||||
uint8x16_t vec_a_2 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 2 * 16);
|
||||
uint8x16_t vec_a2_top = vshrq_n_u8(vec_a_2, 4);
|
||||
uint8x16_t vec_a2_bot = vandq_u8(vec_a_2, vec_mask);
|
||||
int8x16_t vec_v_2_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 0], vec_a2_top);
|
||||
int8x16_t vec_v_2_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 1], vec_a2_top);
|
||||
int8x16_t vec_v_2_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 2], vec_a2_bot);
|
||||
int8x16_t vec_v_2_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 3], vec_a2_bot);
|
||||
int8x16x2_t vec_v_left_2 = vzipq_s8(vec_v_2_left_tmp1, vec_v_2_left_tmp0);
|
||||
int8x16x2_t vec_v_right_2 = vzipq_s8(vec_v_2_right_tmp1, vec_v_2_right_tmp0);
|
||||
vec_c[2] += vec_v_left_2.val[0];
|
||||
vec_c[2] += vec_v_right_2.val[0];
|
||||
vec_c[3] += vec_v_left_2.val[1];
|
||||
vec_c[3] += vec_v_right_2.val[1];
|
||||
|
||||
uint8x16_t vec_a_3 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 3 * 16);
|
||||
uint8x16_t vec_a3_top = vshrq_n_u8(vec_a_3, 4);
|
||||
uint8x16_t vec_a3_bot = vandq_u8(vec_a_3, vec_mask);
|
||||
int8x16_t vec_v_3_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 4], vec_a3_top);
|
||||
int8x16_t vec_v_3_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 5], vec_a3_top);
|
||||
int8x16_t vec_v_3_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 6], vec_a3_bot);
|
||||
int8x16_t vec_v_3_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 7], vec_a3_bot);
|
||||
int8x16x2_t vec_v_left_3 = vzipq_s8(vec_v_3_left_tmp1, vec_v_3_left_tmp0);
|
||||
int8x16x2_t vec_v_right_3 = vzipq_s8(vec_v_3_right_tmp1, vec_v_3_right_tmp0);
|
||||
vec_c[2] += vec_v_left_3.val[0];
|
||||
vec_c[2] += vec_v_right_3.val[0];
|
||||
vec_c[3] += vec_v_left_3.val[1];
|
||||
vec_c[3] += vec_v_right_3.val[1];
|
||||
|
||||
}
|
||||
|
||||
int32x4_t vec_v_bot_low_low_0 = vmovl_s16(vget_low_s16(vec_c[0]));
|
||||
int32x4_t vec_v_bot_low_high_0 = vmovl_high_s16(vec_c[0]);
|
||||
vst1q_s32(c + i + 0, vld1q_s32(c + i + 0) + vec_v_bot_low_low_0);
|
||||
vst1q_s32(c + i + 4, vld1q_s32(c + i + 4) + vec_v_bot_low_high_0);
|
||||
int32x4_t vec_v_bot_low_low_1 = vmovl_s16(vget_low_s16(vec_c[1]));
|
||||
int32x4_t vec_v_bot_low_high_1 = vmovl_high_s16(vec_c[1]);
|
||||
vst1q_s32(c + i + 8, vld1q_s32(c + i + 8) + vec_v_bot_low_low_1);
|
||||
vst1q_s32(c + i + 12, vld1q_s32(c + i + 12) + vec_v_bot_low_high_1);
|
||||
int32x4_t vec_v_bot_low_low_2 = vmovl_s16(vget_low_s16(vec_c[2]));
|
||||
int32x4_t vec_v_bot_low_high_2 = vmovl_high_s16(vec_c[2]);
|
||||
vst1q_s32(c + i + 16, vld1q_s32(c + i + 16) + vec_v_bot_low_low_2);
|
||||
vst1q_s32(c + i + 20, vld1q_s32(c + i + 20) + vec_v_bot_low_high_2);
|
||||
int32x4_t vec_v_bot_low_low_3 = vmovl_s16(vget_low_s16(vec_c[3]));
|
||||
int32x4_t vec_v_bot_low_high_3 = vmovl_high_s16(vec_c[3]);
|
||||
vst1q_s32(c + i + 24, vld1q_s32(c + i + 24) + vec_v_bot_low_low_3);
|
||||
vst1q_s32(c + i + 28, vld1q_s32(c + i + 28) + vec_v_bot_low_high_3);
|
||||
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
int32_t qgemm_lut_3200_8640(void* A, void* LUT, void* Scales, void* LUT_Scales, void* C) {
|
||||
alignas(32) uint32_t CBits[BM3200_8640];
|
||||
memset(&(CBits[0]), 0, BM3200_8640 * sizeof(int32_t));
|
||||
#pragma unroll
|
||||
for (int32_t k_outer = 0; k_outer < 8640 / BBK3200_8640; ++k_outer) {
|
||||
tbl_impl_3200_8640((&(((int32_t*)CBits)[0])), (&(((int8_t*)LUT)[(k_outer * BBK3200_8640 / 2 * 32)])), (&(((uint8_t*)A)[(k_outer * BBK3200_8640 / 2 / 2 * BM3200_8640)])));
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i = 0; i < BM3200_8640; i++) {
|
||||
((bitnet_float_type*)C)[i] = (((int32_t*)CBits)[i]) / ((bitnet_float_type*)LUT_Scales)[0] * ((bitnet_float_type*)Scales)[0];
|
||||
}
|
||||
return 0;
|
||||
};
|
||||
#include <arm_neon.h>
|
||||
|
||||
#define BM3200_3200 320
|
||||
#define BBK3200_3200 128
|
||||
inline void tbl_impl_3200_3200(int32_t* c, int8_t* lut, uint8_t* a) {
|
||||
#ifdef __ARM_NEON
|
||||
const int KK = BBK3200_3200 / 2;
|
||||
const uint8x16_t vec_mask = vdupq_n_u8(0x0f);
|
||||
const int8x16_t vec_zero = vdupq_n_s16(0x0000);
|
||||
int8x16_t vec_lut[2 * KK];
|
||||
int16x8_t vec_c[8];
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 2 * KK; k++) {
|
||||
vec_lut[k] = vld1q_s8(lut + k * 16);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < BM3200_3200; i += 64) {
|
||||
#pragma unroll
|
||||
for (int i=0; i<8; i++) {
|
||||
vec_c[i] = vandq_s16(vec_c[i], vec_zero);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int k = 0; k < KK / 2; k++) {
|
||||
|
||||
uint8x16_t vec_a_0 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 0 * 16);
|
||||
uint8x16_t vec_a0_top = vshrq_n_u8(vec_a_0, 4);
|
||||
uint8x16_t vec_a0_bot = vandq_u8(vec_a_0, vec_mask);
|
||||
int8x16_t vec_v_0_left_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 0], vec_a0_top);
|
||||
int8x16_t vec_v_0_left_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 1], vec_a0_top);
|
||||
int8x16_t vec_v_0_right_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 2], vec_a0_bot);
|
||||
int8x16_t vec_v_0_right_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 3], vec_a0_bot);
|
||||
int8x16x2_t vec_v_left_0 = vzipq_s8(vec_v_0_left_tmp1, vec_v_0_left_tmp0);
|
||||
int8x16x2_t vec_v_right_0 = vzipq_s8(vec_v_0_right_tmp1, vec_v_0_right_tmp0);
|
||||
vec_c[0] += vec_v_left_0.val[0];
|
||||
vec_c[0] += vec_v_right_0.val[0];
|
||||
vec_c[1] += vec_v_left_0.val[1];
|
||||
vec_c[1] += vec_v_right_0.val[1];
|
||||
|
||||
uint8x16_t vec_a_1 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 1 * 16);
|
||||
uint8x16_t vec_a1_top = vshrq_n_u8(vec_a_1, 4);
|
||||
uint8x16_t vec_a1_bot = vandq_u8(vec_a_1, vec_mask);
|
||||
int8x16_t vec_v_1_left_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 0], vec_a1_top);
|
||||
int8x16_t vec_v_1_left_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 1], vec_a1_top);
|
||||
int8x16_t vec_v_1_right_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 2], vec_a1_bot);
|
||||
int8x16_t vec_v_1_right_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 3], vec_a1_bot);
|
||||
int8x16x2_t vec_v_left_1 = vzipq_s8(vec_v_1_left_tmp1, vec_v_1_left_tmp0);
|
||||
int8x16x2_t vec_v_right_1 = vzipq_s8(vec_v_1_right_tmp1, vec_v_1_right_tmp0);
|
||||
vec_c[2] += vec_v_left_1.val[0];
|
||||
vec_c[2] += vec_v_right_1.val[0];
|
||||
vec_c[3] += vec_v_left_1.val[1];
|
||||
vec_c[3] += vec_v_right_1.val[1];
|
||||
|
||||
uint8x16_t vec_a_2 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 2 * 16);
|
||||
uint8x16_t vec_a2_top = vshrq_n_u8(vec_a_2, 4);
|
||||
uint8x16_t vec_a2_bot = vandq_u8(vec_a_2, vec_mask);
|
||||
int8x16_t vec_v_2_left_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 0], vec_a2_top);
|
||||
int8x16_t vec_v_2_left_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 1], vec_a2_top);
|
||||
int8x16_t vec_v_2_right_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 2], vec_a2_bot);
|
||||
int8x16_t vec_v_2_right_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 3], vec_a2_bot);
|
||||
int8x16x2_t vec_v_left_2 = vzipq_s8(vec_v_2_left_tmp1, vec_v_2_left_tmp0);
|
||||
int8x16x2_t vec_v_right_2 = vzipq_s8(vec_v_2_right_tmp1, vec_v_2_right_tmp0);
|
||||
vec_c[4] += vec_v_left_2.val[0];
|
||||
vec_c[4] += vec_v_right_2.val[0];
|
||||
vec_c[5] += vec_v_left_2.val[1];
|
||||
vec_c[5] += vec_v_right_2.val[1];
|
||||
|
||||
uint8x16_t vec_a_3 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 3 * 16);
|
||||
uint8x16_t vec_a3_top = vshrq_n_u8(vec_a_3, 4);
|
||||
uint8x16_t vec_a3_bot = vandq_u8(vec_a_3, vec_mask);
|
||||
int8x16_t vec_v_3_left_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 0], vec_a3_top);
|
||||
int8x16_t vec_v_3_left_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 1], vec_a3_top);
|
||||
int8x16_t vec_v_3_right_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 2], vec_a3_bot);
|
||||
int8x16_t vec_v_3_right_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 3], vec_a3_bot);
|
||||
int8x16x2_t vec_v_left_3 = vzipq_s8(vec_v_3_left_tmp1, vec_v_3_left_tmp0);
|
||||
int8x16x2_t vec_v_right_3 = vzipq_s8(vec_v_3_right_tmp1, vec_v_3_right_tmp0);
|
||||
vec_c[6] += vec_v_left_3.val[0];
|
||||
vec_c[6] += vec_v_right_3.val[0];
|
||||
vec_c[7] += vec_v_left_3.val[1];
|
||||
vec_c[7] += vec_v_right_3.val[1];
|
||||
|
||||
}
|
||||
|
||||
int32x4_t vec_v_bot_low_low_0 = vmovl_s16(vget_low_s16(vec_c[0]));
|
||||
int32x4_t vec_v_bot_low_high_0 = vmovl_high_s16(vec_c[0]);
|
||||
vst1q_s32(c + i + 0, vld1q_s32(c + i + 0) + vec_v_bot_low_low_0);
|
||||
vst1q_s32(c + i + 4, vld1q_s32(c + i + 4) + vec_v_bot_low_high_0);
|
||||
int32x4_t vec_v_bot_low_low_1 = vmovl_s16(vget_low_s16(vec_c[1]));
|
||||
int32x4_t vec_v_bot_low_high_1 = vmovl_high_s16(vec_c[1]);
|
||||
vst1q_s32(c + i + 8, vld1q_s32(c + i + 8) + vec_v_bot_low_low_1);
|
||||
vst1q_s32(c + i + 12, vld1q_s32(c + i + 12) + vec_v_bot_low_high_1);
|
||||
int32x4_t vec_v_bot_low_low_2 = vmovl_s16(vget_low_s16(vec_c[2]));
|
||||
int32x4_t vec_v_bot_low_high_2 = vmovl_high_s16(vec_c[2]);
|
||||
vst1q_s32(c + i + 16, vld1q_s32(c + i + 16) + vec_v_bot_low_low_2);
|
||||
vst1q_s32(c + i + 20, vld1q_s32(c + i + 20) + vec_v_bot_low_high_2);
|
||||
int32x4_t vec_v_bot_low_low_3 = vmovl_s16(vget_low_s16(vec_c[3]));
|
||||
int32x4_t vec_v_bot_low_high_3 = vmovl_high_s16(vec_c[3]);
|
||||
vst1q_s32(c + i + 24, vld1q_s32(c + i + 24) + vec_v_bot_low_low_3);
|
||||
vst1q_s32(c + i + 28, vld1q_s32(c + i + 28) + vec_v_bot_low_high_3);
|
||||
int32x4_t vec_v_bot_low_low_4 = vmovl_s16(vget_low_s16(vec_c[4]));
|
||||
int32x4_t vec_v_bot_low_high_4 = vmovl_high_s16(vec_c[4]);
|
||||
vst1q_s32(c + i + 32, vld1q_s32(c + i + 32) + vec_v_bot_low_low_4);
|
||||
vst1q_s32(c + i + 36, vld1q_s32(c + i + 36) + vec_v_bot_low_high_4);
|
||||
int32x4_t vec_v_bot_low_low_5 = vmovl_s16(vget_low_s16(vec_c[5]));
|
||||
int32x4_t vec_v_bot_low_high_5 = vmovl_high_s16(vec_c[5]);
|
||||
vst1q_s32(c + i + 40, vld1q_s32(c + i + 40) + vec_v_bot_low_low_5);
|
||||
vst1q_s32(c + i + 44, vld1q_s32(c + i + 44) + vec_v_bot_low_high_5);
|
||||
int32x4_t vec_v_bot_low_low_6 = vmovl_s16(vget_low_s16(vec_c[6]));
|
||||
int32x4_t vec_v_bot_low_high_6 = vmovl_high_s16(vec_c[6]);
|
||||
vst1q_s32(c + i + 48, vld1q_s32(c + i + 48) + vec_v_bot_low_low_6);
|
||||
vst1q_s32(c + i + 52, vld1q_s32(c + i + 52) + vec_v_bot_low_high_6);
|
||||
int32x4_t vec_v_bot_low_low_7 = vmovl_s16(vget_low_s16(vec_c[7]));
|
||||
int32x4_t vec_v_bot_low_high_7 = vmovl_high_s16(vec_c[7]);
|
||||
vst1q_s32(c + i + 56, vld1q_s32(c + i + 56) + vec_v_bot_low_low_7);
|
||||
vst1q_s32(c + i + 60, vld1q_s32(c + i + 60) + vec_v_bot_low_high_7);
|
||||
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
int32_t qgemm_lut_3200_3200(void* A, void* LUT, void* Scales, void* LUT_Scales, void* C) {
|
||||
alignas(32) uint32_t CBits[BM3200_3200];
|
||||
memset(&(CBits[0]), 0, BM3200_3200 * sizeof(int32_t));
|
||||
#pragma unroll
|
||||
for (int32_t k_outer = 0; k_outer < 3200 / BBK3200_3200; ++k_outer) {
|
||||
tbl_impl_3200_3200((&(((int32_t*)CBits)[0])), (&(((int8_t*)LUT)[(k_outer * BBK3200_3200 / 2 * 32)])), (&(((uint8_t*)A)[(k_outer * BBK3200_3200 / 2 / 2 * BM3200_3200)])));
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i = 0; i < BM3200_3200; i++) {
|
||||
((bitnet_float_type*)C)[i] = (((int32_t*)CBits)[i]) / ((bitnet_float_type*)LUT_Scales)[0] * ((bitnet_float_type*)Scales)[0];
|
||||
}
|
||||
return 0;
|
||||
};
|
||||
#include <arm_neon.h>
|
||||
|
||||
#define BM8640_3200 320
|
||||
#define BBK8640_3200 64
|
||||
inline void tbl_impl_8640_3200(int32_t* c, int8_t* lut, uint8_t* a) {
|
||||
#ifdef __ARM_NEON
|
||||
const int KK = BBK8640_3200 / 2;
|
||||
const uint8x16_t vec_mask = vdupq_n_u8(0x0f);
|
||||
const int8x16_t vec_zero = vdupq_n_s16(0x0000);
|
||||
int8x16_t vec_lut[2 * KK];
|
||||
int16x8_t vec_c[4];
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 2 * KK; k++) {
|
||||
vec_lut[k] = vld1q_s8(lut + k * 16);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < BM8640_3200; i += 32) {
|
||||
#pragma unroll
|
||||
for (int i=0; i<4; i++) {
|
||||
vec_c[i] = vandq_s16(vec_c[i], vec_zero);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int k = 0; k < KK / 4; k++) {
|
||||
|
||||
uint8x16_t vec_a_0 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 0 * 16);
|
||||
uint8x16_t vec_a0_top = vshrq_n_u8(vec_a_0, 4);
|
||||
uint8x16_t vec_a0_bot = vandq_u8(vec_a_0, vec_mask);
|
||||
int8x16_t vec_v_0_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 0], vec_a0_top);
|
||||
int8x16_t vec_v_0_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 1], vec_a0_top);
|
||||
int8x16_t vec_v_0_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 2], vec_a0_bot);
|
||||
int8x16_t vec_v_0_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 3], vec_a0_bot);
|
||||
int8x16x2_t vec_v_left_0 = vzipq_s8(vec_v_0_left_tmp1, vec_v_0_left_tmp0);
|
||||
int8x16x2_t vec_v_right_0 = vzipq_s8(vec_v_0_right_tmp1, vec_v_0_right_tmp0);
|
||||
vec_c[0] += vec_v_left_0.val[0];
|
||||
vec_c[0] += vec_v_right_0.val[0];
|
||||
vec_c[1] += vec_v_left_0.val[1];
|
||||
vec_c[1] += vec_v_right_0.val[1];
|
||||
|
||||
uint8x16_t vec_a_1 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 1 * 16);
|
||||
uint8x16_t vec_a1_top = vshrq_n_u8(vec_a_1, 4);
|
||||
uint8x16_t vec_a1_bot = vandq_u8(vec_a_1, vec_mask);
|
||||
int8x16_t vec_v_1_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 4], vec_a1_top);
|
||||
int8x16_t vec_v_1_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 5], vec_a1_top);
|
||||
int8x16_t vec_v_1_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 6], vec_a1_bot);
|
||||
int8x16_t vec_v_1_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 7], vec_a1_bot);
|
||||
int8x16x2_t vec_v_left_1 = vzipq_s8(vec_v_1_left_tmp1, vec_v_1_left_tmp0);
|
||||
int8x16x2_t vec_v_right_1 = vzipq_s8(vec_v_1_right_tmp1, vec_v_1_right_tmp0);
|
||||
vec_c[0] += vec_v_left_1.val[0];
|
||||
vec_c[0] += vec_v_right_1.val[0];
|
||||
vec_c[1] += vec_v_left_1.val[1];
|
||||
vec_c[1] += vec_v_right_1.val[1];
|
||||
|
||||
uint8x16_t vec_a_2 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 2 * 16);
|
||||
uint8x16_t vec_a2_top = vshrq_n_u8(vec_a_2, 4);
|
||||
uint8x16_t vec_a2_bot = vandq_u8(vec_a_2, vec_mask);
|
||||
int8x16_t vec_v_2_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 0], vec_a2_top);
|
||||
int8x16_t vec_v_2_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 1], vec_a2_top);
|
||||
int8x16_t vec_v_2_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 2], vec_a2_bot);
|
||||
int8x16_t vec_v_2_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 3], vec_a2_bot);
|
||||
int8x16x2_t vec_v_left_2 = vzipq_s8(vec_v_2_left_tmp1, vec_v_2_left_tmp0);
|
||||
int8x16x2_t vec_v_right_2 = vzipq_s8(vec_v_2_right_tmp1, vec_v_2_right_tmp0);
|
||||
vec_c[2] += vec_v_left_2.val[0];
|
||||
vec_c[2] += vec_v_right_2.val[0];
|
||||
vec_c[3] += vec_v_left_2.val[1];
|
||||
vec_c[3] += vec_v_right_2.val[1];
|
||||
|
||||
uint8x16_t vec_a_3 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 3 * 16);
|
||||
uint8x16_t vec_a3_top = vshrq_n_u8(vec_a_3, 4);
|
||||
uint8x16_t vec_a3_bot = vandq_u8(vec_a_3, vec_mask);
|
||||
int8x16_t vec_v_3_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 4], vec_a3_top);
|
||||
int8x16_t vec_v_3_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 5], vec_a3_top);
|
||||
int8x16_t vec_v_3_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 6], vec_a3_bot);
|
||||
int8x16_t vec_v_3_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 7], vec_a3_bot);
|
||||
int8x16x2_t vec_v_left_3 = vzipq_s8(vec_v_3_left_tmp1, vec_v_3_left_tmp0);
|
||||
int8x16x2_t vec_v_right_3 = vzipq_s8(vec_v_3_right_tmp1, vec_v_3_right_tmp0);
|
||||
vec_c[2] += vec_v_left_3.val[0];
|
||||
vec_c[2] += vec_v_right_3.val[0];
|
||||
vec_c[3] += vec_v_left_3.val[1];
|
||||
vec_c[3] += vec_v_right_3.val[1];
|
||||
|
||||
}
|
||||
|
||||
int32x4_t vec_v_bot_low_low_0 = vmovl_s16(vget_low_s16(vec_c[0]));
|
||||
int32x4_t vec_v_bot_low_high_0 = vmovl_high_s16(vec_c[0]);
|
||||
vst1q_s32(c + i + 0, vld1q_s32(c + i + 0) + vec_v_bot_low_low_0);
|
||||
vst1q_s32(c + i + 4, vld1q_s32(c + i + 4) + vec_v_bot_low_high_0);
|
||||
int32x4_t vec_v_bot_low_low_1 = vmovl_s16(vget_low_s16(vec_c[1]));
|
||||
int32x4_t vec_v_bot_low_high_1 = vmovl_high_s16(vec_c[1]);
|
||||
vst1q_s32(c + i + 8, vld1q_s32(c + i + 8) + vec_v_bot_low_low_1);
|
||||
vst1q_s32(c + i + 12, vld1q_s32(c + i + 12) + vec_v_bot_low_high_1);
|
||||
int32x4_t vec_v_bot_low_low_2 = vmovl_s16(vget_low_s16(vec_c[2]));
|
||||
int32x4_t vec_v_bot_low_high_2 = vmovl_high_s16(vec_c[2]);
|
||||
vst1q_s32(c + i + 16, vld1q_s32(c + i + 16) + vec_v_bot_low_low_2);
|
||||
vst1q_s32(c + i + 20, vld1q_s32(c + i + 20) + vec_v_bot_low_high_2);
|
||||
int32x4_t vec_v_bot_low_low_3 = vmovl_s16(vget_low_s16(vec_c[3]));
|
||||
int32x4_t vec_v_bot_low_high_3 = vmovl_high_s16(vec_c[3]);
|
||||
vst1q_s32(c + i + 24, vld1q_s32(c + i + 24) + vec_v_bot_low_low_3);
|
||||
vst1q_s32(c + i + 28, vld1q_s32(c + i + 28) + vec_v_bot_low_high_3);
|
||||
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
int32_t qgemm_lut_8640_3200(void* A, void* LUT, void* Scales, void* LUT_Scales, void* C) {
|
||||
alignas(32) uint32_t CBits[BM8640_3200];
|
||||
memset(&(CBits[0]), 0, BM8640_3200 * sizeof(int32_t));
|
||||
#pragma unroll
|
||||
for (int32_t k_outer = 0; k_outer < 3200 / BBK8640_3200; ++k_outer) {
|
||||
tbl_impl_8640_3200((&(((int32_t*)CBits)[0])), (&(((int8_t*)LUT)[(k_outer * BBK8640_3200 / 2 * 32)])), (&(((uint8_t*)A)[(k_outer * BBK8640_3200 / 2 / 2 * BM8640_3200)])));
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i = 0; i < BM8640_3200; i++) {
|
||||
((bitnet_float_type*)C)[i] = (((int32_t*)CBits)[i]) / ((bitnet_float_type*)LUT_Scales)[0] * ((bitnet_float_type*)Scales)[0];
|
||||
}
|
||||
return 0;
|
||||
};
|
||||
|
||||
template<int K>
|
||||
void preprocessor_k(void* B, void* LUT_Scales, void* QLUT) {{
|
||||
partial_max_reset((&(((bitnet_float_type*)LUT_Scales)[0])));
|
||||
per_tensor_quant(K, (&(((bitnet_float_type*)LUT_Scales)[0])), (&(((bitnet_float_type*)B)[0])));
|
||||
|
||||
lut_ctor<K>((&(((int8_t*)QLUT)[0])), (&(((bitnet_float_type*)B)[0])), (&(((bitnet_float_type*)LUT_Scales)[0])));
|
||||
}}
|
||||
void ggml_preprocessor(int m, int k, void* B, void* LUT_Scales, void* QLUT) {
|
||||
if (m == 3200 && k == 8640) {
|
||||
preprocessor_k<8640>(B, LUT_Scales, QLUT);
|
||||
}
|
||||
else if (m == 3200 && k == 3200) {
|
||||
preprocessor_k<3200>(B, LUT_Scales, QLUT);
|
||||
}
|
||||
else if (m == 8640 && k == 3200) {
|
||||
preprocessor_k<3200>(B, LUT_Scales, QLUT);
|
||||
}
|
||||
}
|
||||
void ggml_qgemm_lut(int m, int k, void* A, void* LUT, void* Scales, void* LUT_Scales, void* C) {
|
||||
if (m == 3200 && k == 8640) {
|
||||
qgemm_lut_3200_8640(A, LUT, Scales, LUT_Scales, C);
|
||||
}
|
||||
else if (m == 3200 && k == 3200) {
|
||||
qgemm_lut_3200_3200(A, LUT, Scales, LUT_Scales, C);
|
||||
}
|
||||
else if (m == 8640 && k == 3200) {
|
||||
qgemm_lut_8640_3200(A, LUT, Scales, LUT_Scales, C);
|
||||
}
|
||||
}
|
||||
|
||||
void ggml_bitnet_transform_tensor(struct ggml_tensor * tensor) {
|
||||
if (!(is_type_supported(tensor->type) && tensor->backend == GGML_BACKEND_TYPE_CPU && tensor->extra == nullptr)) {
|
||||
return;
|
||||
}
|
||||
|
||||
int k = tensor->ne[0];
|
||||
int m = tensor->ne[1];
|
||||
const int lut_scales_size = 1;
|
||||
const int scales_size = 1;
|
||||
int bk = 0;
|
||||
int bm = 0;
|
||||
|
||||
if (m == 3200 && k == 8640) {
|
||||
bm = BM3200_8640;
|
||||
bk = BBK3200_8640;
|
||||
}
|
||||
else if (m == 3200 && k == 3200) {
|
||||
bm = BM3200_3200;
|
||||
bk = BBK3200_3200;
|
||||
}
|
||||
else if (m == 8640 && k == 3200) {
|
||||
bm = BM8640_3200;
|
||||
bk = BBK8640_3200;
|
||||
}
|
||||
|
||||
const int n_tile_num = m / bm;
|
||||
const int BK = bk;
|
||||
uint8_t * qweights;
|
||||
bitnet_float_type * scales;
|
||||
|
||||
scales = (bitnet_float_type *) aligned_malloc(sizeof(bitnet_float_type));
|
||||
qweights = (uint8_t *) tensor->data;
|
||||
float * i2_scales = (float * )(qweights + k * m / 4);
|
||||
scales[0] = (bitnet_float_type) i2_scales[0];
|
||||
|
||||
tensor->extra = bitnet_tensor_extras + bitnet_tensor_extras_index;
|
||||
bitnet_tensor_extras[bitnet_tensor_extras_index++] = {
|
||||
/* .lut_scales_size = */ lut_scales_size,
|
||||
/* .scales_size = */ scales_size,
|
||||
/* .n_tile_num = */ n_tile_num,
|
||||
/* .qweights = */ qweights,
|
||||
/* .scales = */ scales
|
||||
};
|
||||
}
|
||||
#endif
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,21 @@
|
||||
[Kernels_0]
|
||||
m = 3200
|
||||
k = 8640
|
||||
bm = 160
|
||||
bk = 64
|
||||
bmm = 32
|
||||
|
||||
[Kernels_1]
|
||||
m = 3200
|
||||
k = 3200
|
||||
bm = 320
|
||||
bk = 128
|
||||
bmm = 64
|
||||
|
||||
[Kernels_2]
|
||||
m = 8640
|
||||
k = 3200
|
||||
bm = 320
|
||||
bk = 64
|
||||
bmm = 32
|
||||
|
||||
@@ -0,0 +1,21 @@
|
||||
[Kernels_0]
|
||||
m = 3200
|
||||
k = 8640
|
||||
bm = 160
|
||||
bk = 96
|
||||
bmm = 32
|
||||
|
||||
[Kernels_1]
|
||||
m = 3200
|
||||
k = 3200
|
||||
bm = 320
|
||||
bk = 96
|
||||
bmm = 32
|
||||
|
||||
[Kernels_2]
|
||||
m = 8640
|
||||
k = 3200
|
||||
bm = 320
|
||||
bk = 96
|
||||
bmm = 32
|
||||
|
||||
@@ -0,0 +1,627 @@
|
||||
#if defined(GGML_BITNET_ARM_TL1)
|
||||
#include "ggml-bitnet.h"
|
||||
#define GGML_BITNET_MAX_NODES 8192
|
||||
static bool initialized = false;
|
||||
static bitnet_tensor_extra * bitnet_tensor_extras = nullptr;
|
||||
static size_t bitnet_tensor_extras_index = 0;
|
||||
static void * aligned_malloc(size_t size) {{
|
||||
#if defined(_WIN32)
|
||||
return _aligned_malloc(size, 64);
|
||||
#else
|
||||
void * ptr = nullptr;
|
||||
posix_memalign(&ptr, 64, size);
|
||||
return ptr;
|
||||
#endif
|
||||
}}
|
||||
static void aligned_free(void * ptr) {{
|
||||
#if defined(_WIN32)
|
||||
_aligned_free(ptr);
|
||||
#else
|
||||
free(ptr);
|
||||
#endif
|
||||
}}
|
||||
|
||||
void per_tensor_quant(int k, void* lut_scales_, void* b_) {{
|
||||
bitnet_float_type* lut_scales = (bitnet_float_type*)lut_scales_;
|
||||
bitnet_float_type* b = (bitnet_float_type*)b_;
|
||||
#ifdef __ARM_NEON
|
||||
float32x4_t temp_max = vdupq_n_f32(0);
|
||||
for (int i=0; i < k / 4; i++) {{
|
||||
float32x4_t vec_bs = vld1q_f32(b + 4 * i);
|
||||
float32x4_t abssum = vabsq_f32(vec_bs);
|
||||
temp_max = vmaxq_f32(abssum, temp_max);
|
||||
}}
|
||||
float32_t scales = 127 / vmaxvq_f32(temp_max);
|
||||
*lut_scales = scales;
|
||||
#elif defined __AVX2__
|
||||
__m256 max_vec = _mm256_set1_ps(0.f);
|
||||
const __m256 vec_sign = _mm256_set1_ps(-0.0f);
|
||||
// #pragma unroll
|
||||
for (int i = 0; i < k / 8; i++) {{
|
||||
__m256 vec_b = _mm256_loadu_ps(b + i * 8);
|
||||
__m256 vec_babs = _mm256_andnot_ps(vec_sign, vec_b);
|
||||
max_vec = _mm256_max_ps(vec_babs, max_vec);
|
||||
}}
|
||||
__m128 max1 = _mm_max_ps(_mm256_extractf128_ps(max_vec, 1), _mm256_castps256_ps128(max_vec));
|
||||
max1 = _mm_max_ps(max1, _mm_movehl_ps(max1, max1));
|
||||
max1 = _mm_max_ss(max1, _mm_movehdup_ps(max1));
|
||||
float scales = 127 / _mm_cvtss_f32(max1);
|
||||
*lut_scales = scales;
|
||||
#endif
|
||||
}}
|
||||
|
||||
void partial_max_reset(void* lut_scales_) {{
|
||||
bitnet_float_type* lut_scales = (bitnet_float_type*)lut_scales_;
|
||||
*lut_scales = 0.0;
|
||||
}}
|
||||
|
||||
#ifdef __ARM_NEON
|
||||
inline void Transpose_8_8(
|
||||
int16x8_t *v0,
|
||||
int16x8_t *v1,
|
||||
int16x8_t *v2,
|
||||
int16x8_t *v3,
|
||||
int16x8_t *v4,
|
||||
int16x8_t *v5,
|
||||
int16x8_t *v6,
|
||||
int16x8_t *v7)
|
||||
{{
|
||||
int16x8x2_t q04 = vzipq_s16(*v0, *v4);
|
||||
int16x8x2_t q15 = vzipq_s16(*v1, *v5);
|
||||
int16x8x2_t q26 = vzipq_s16(*v2, *v6);
|
||||
int16x8x2_t q37 = vzipq_s16(*v3, *v7);
|
||||
|
||||
int16x8x2_t q0246_0 = vzipq_s16(q04.val[0], q26.val[0]);
|
||||
int16x8x2_t q0246_1 = vzipq_s16(q04.val[1], q26.val[1]);
|
||||
int16x8x2_t q1357_0 = vzipq_s16(q15.val[0], q37.val[0]);
|
||||
int16x8x2_t q1357_1 = vzipq_s16(q15.val[1], q37.val[1]);
|
||||
|
||||
int16x8x2_t q_fin_0 = vzipq_s16(q0246_0.val[0], q1357_0.val[0]);
|
||||
int16x8x2_t q_fin_1 = vzipq_s16(q0246_0.val[1], q1357_0.val[1]);
|
||||
int16x8x2_t q_fin_2 = vzipq_s16(q0246_1.val[0], q1357_1.val[0]);
|
||||
int16x8x2_t q_fin_3 = vzipq_s16(q0246_1.val[1], q1357_1.val[1]);
|
||||
|
||||
*v0 = q_fin_0.val[0];
|
||||
*v1 = q_fin_0.val[1];
|
||||
*v2 = q_fin_1.val[0];
|
||||
*v3 = q_fin_1.val[1];
|
||||
*v4 = q_fin_2.val[0];
|
||||
*v5 = q_fin_2.val[1];
|
||||
*v6 = q_fin_3.val[0];
|
||||
*v7 = q_fin_3.val[1];
|
||||
}}
|
||||
#endif
|
||||
|
||||
template<int act_k>
|
||||
inline void lut_ctor(int8_t* qlut, bitnet_float_type* b, bitnet_float_type* lut_scales) {{
|
||||
#ifdef __ARM_NEON
|
||||
int16x8_t vec_lut[16];
|
||||
float32_t scales = *lut_scales;
|
||||
uint8_t tbl_mask[16];
|
||||
tbl_mask[0] = 0;
|
||||
tbl_mask[1] = 2;
|
||||
tbl_mask[2] = 4;
|
||||
tbl_mask[3] = 6;
|
||||
tbl_mask[4] = 8;
|
||||
tbl_mask[5] = 10;
|
||||
tbl_mask[6] = 12;
|
||||
tbl_mask[7] = 14;
|
||||
tbl_mask[8] = 1;
|
||||
tbl_mask[9] = 3;
|
||||
tbl_mask[10] = 5;
|
||||
tbl_mask[11] = 7;
|
||||
tbl_mask[12] = 9;
|
||||
tbl_mask[13] = 11;
|
||||
tbl_mask[14] = 13;
|
||||
tbl_mask[15] = 15;
|
||||
uint8x16_t tbl_mask_q = vld1q_u8(tbl_mask);
|
||||
#pragma unroll
|
||||
for (int k = 0; k < act_k / 16; ++k) {{
|
||||
float32x4x2_t vec_bs_x0 = vld2q_f32(b + k * 16);
|
||||
float32x4x2_t vec_bs_x1 = vld2q_f32(b + k * 16 + 8);
|
||||
float32x4_t vec_f_0 = vmulq_n_f32(vec_bs_x0.val[0], scales);
|
||||
float32x4_t vec_f_1 = vmulq_n_f32(vec_bs_x0.val[1], scales);
|
||||
float32x4_t vec_f_2 = vmulq_n_f32(vec_bs_x1.val[0], scales);
|
||||
float32x4_t vec_f_3 = vmulq_n_f32(vec_bs_x1.val[1], scales);
|
||||
int32x4_t vec_b_0 = vcvtnq_s32_f32(vec_f_0);
|
||||
int32x4_t vec_b_1 = vcvtnq_s32_f32(vec_f_1);
|
||||
int32x4_t vec_b_2 = vcvtnq_s32_f32(vec_f_2);
|
||||
int32x4_t vec_b_3 = vcvtnq_s32_f32(vec_f_3);
|
||||
int16x4_t vec_b16_0 = vmovn_s32(vec_b_0);
|
||||
int16x4_t vec_b16_1 = vmovn_s32(vec_b_1);
|
||||
int16x4_t vec_b16_2 = vmovn_s32(vec_b_2);
|
||||
int16x4_t vec_b16_3 = vmovn_s32(vec_b_3);
|
||||
int16x8_t vec_bs_0 = vcombine_s16(vec_b16_0, vec_b16_2);
|
||||
int16x8_t vec_bs_1 = vcombine_s16(vec_b16_1, vec_b16_3);
|
||||
vec_lut[0] = vdupq_n_s16(0);
|
||||
vec_lut[0] = vec_lut[0] - vec_bs_0;
|
||||
vec_lut[0] = vec_lut[0] - vec_bs_1;
|
||||
vec_lut[1] = vdupq_n_s16(0);
|
||||
vec_lut[1] = vec_lut[1] - vec_bs_0;
|
||||
vec_lut[2] = vdupq_n_s16(0);
|
||||
vec_lut[2] = vec_lut[2] - vec_bs_0;
|
||||
vec_lut[2] = vec_lut[2] + vec_bs_1;
|
||||
vec_lut[3] = vdupq_n_s16(0);
|
||||
vec_lut[3] = vec_lut[3] - vec_bs_1;
|
||||
vec_lut[4] = vdupq_n_s16(0);
|
||||
vec_lut[5] = vec_bs_1;
|
||||
vec_lut[6] = vec_bs_0;
|
||||
vec_lut[6] = vec_lut[6] - vec_bs_1;
|
||||
vec_lut[7] = vec_bs_0;
|
||||
vec_lut[8] = vec_bs_0;
|
||||
vec_lut[8] = vec_lut[8] + vec_bs_1;
|
||||
Transpose_8_8(&(vec_lut[0]), &(vec_lut[1]), &(vec_lut[2]), &(vec_lut[3]),
|
||||
&(vec_lut[4]), &(vec_lut[5]), &(vec_lut[6]), &(vec_lut[7]));
|
||||
Transpose_8_8(&(vec_lut[8]), &(vec_lut[9]), &(vec_lut[10]), &(vec_lut[11]),
|
||||
&(vec_lut[12]), &(vec_lut[13]), &(vec_lut[14]), &(vec_lut[15]));
|
||||
#pragma unroll
|
||||
for (int idx = 0; idx < 8; idx++) {{
|
||||
int8x16_t q0_s = vqtbl1q_s8(vreinterpretq_s8_s16(vec_lut[idx]), tbl_mask_q);
|
||||
int8x8_t q0_low = vget_low_s8(q0_s);
|
||||
int8x8_t q0_high = vget_high_s8(q0_s);
|
||||
int8x16_t q1_s = vqtbl1q_s8(vreinterpretq_s8_s16(vec_lut[idx + 8]), tbl_mask_q);
|
||||
int8x8_t q1_low = vget_low_s8(q1_s);
|
||||
int8x8_t q1_high = vget_high_s8(q1_s);
|
||||
vst1_s8(qlut + k * 16 * 8 * 2 + idx * 16 * 2, q0_high);
|
||||
vst1_s8(qlut + k * 16 * 8 * 2 + idx * 16 * 2 + 8, q1_high);
|
||||
vst1_s8(qlut + k * 16 * 8 * 2 + idx * 16 * 2 + 16, q0_low);
|
||||
vst1_s8(qlut + k * 16 * 8 * 2 + idx * 16 * 2 + 24, q1_low);
|
||||
}}
|
||||
}}
|
||||
#endif
|
||||
}}
|
||||
|
||||
static bool is_type_supported(enum ggml_type type) {{
|
||||
if (type == GGML_TYPE_Q4_0 ||
|
||||
type == GGML_TYPE_TL1) {{
|
||||
return true;
|
||||
}} else {{
|
||||
return false;
|
||||
}}
|
||||
}}
|
||||
#include <arm_neon.h>
|
||||
|
||||
#define BM1536_4096 256
|
||||
#define BBK1536_4096 128
|
||||
inline void tbl_impl_1536_4096(int32_t* c, int8_t* lut, uint8_t* a) {
|
||||
#ifdef __ARM_NEON
|
||||
const int KK = BBK1536_4096 / 2;
|
||||
const uint8x16_t vec_mask = vdupq_n_u8(0x0f);
|
||||
const int8x16_t vec_zero = vdupq_n_s16(0x0000);
|
||||
int8x16_t vec_lut[2 * KK];
|
||||
int16x8_t vec_c[4];
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 2 * KK; k++) {
|
||||
vec_lut[k] = vld1q_s8(lut + k * 16);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < BM1536_4096; i += 32) {
|
||||
#pragma unroll
|
||||
for (int i=0; i<4; i++) {
|
||||
vec_c[i] = vandq_s16(vec_c[i], vec_zero);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int k = 0; k < KK / 4; k++) {
|
||||
|
||||
uint8x16_t vec_a_0 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 0 * 16);
|
||||
uint8x16_t vec_a0_top = vshrq_n_u8(vec_a_0, 4);
|
||||
uint8x16_t vec_a0_bot = vandq_u8(vec_a_0, vec_mask);
|
||||
int8x16_t vec_v_0_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 0], vec_a0_top);
|
||||
int8x16_t vec_v_0_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 1], vec_a0_top);
|
||||
int8x16_t vec_v_0_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 2], vec_a0_bot);
|
||||
int8x16_t vec_v_0_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 3], vec_a0_bot);
|
||||
int8x16x2_t vec_v_left_0 = vzipq_s8(vec_v_0_left_tmp1, vec_v_0_left_tmp0);
|
||||
int8x16x2_t vec_v_right_0 = vzipq_s8(vec_v_0_right_tmp1, vec_v_0_right_tmp0);
|
||||
vec_c[0] += vec_v_left_0.val[0];
|
||||
vec_c[0] += vec_v_right_0.val[0];
|
||||
vec_c[1] += vec_v_left_0.val[1];
|
||||
vec_c[1] += vec_v_right_0.val[1];
|
||||
|
||||
uint8x16_t vec_a_1 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 1 * 16);
|
||||
uint8x16_t vec_a1_top = vshrq_n_u8(vec_a_1, 4);
|
||||
uint8x16_t vec_a1_bot = vandq_u8(vec_a_1, vec_mask);
|
||||
int8x16_t vec_v_1_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 4], vec_a1_top);
|
||||
int8x16_t vec_v_1_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 5], vec_a1_top);
|
||||
int8x16_t vec_v_1_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 6], vec_a1_bot);
|
||||
int8x16_t vec_v_1_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 7], vec_a1_bot);
|
||||
int8x16x2_t vec_v_left_1 = vzipq_s8(vec_v_1_left_tmp1, vec_v_1_left_tmp0);
|
||||
int8x16x2_t vec_v_right_1 = vzipq_s8(vec_v_1_right_tmp1, vec_v_1_right_tmp0);
|
||||
vec_c[0] += vec_v_left_1.val[0];
|
||||
vec_c[0] += vec_v_right_1.val[0];
|
||||
vec_c[1] += vec_v_left_1.val[1];
|
||||
vec_c[1] += vec_v_right_1.val[1];
|
||||
|
||||
uint8x16_t vec_a_2 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 2 * 16);
|
||||
uint8x16_t vec_a2_top = vshrq_n_u8(vec_a_2, 4);
|
||||
uint8x16_t vec_a2_bot = vandq_u8(vec_a_2, vec_mask);
|
||||
int8x16_t vec_v_2_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 0], vec_a2_top);
|
||||
int8x16_t vec_v_2_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 1], vec_a2_top);
|
||||
int8x16_t vec_v_2_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 2], vec_a2_bot);
|
||||
int8x16_t vec_v_2_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 3], vec_a2_bot);
|
||||
int8x16x2_t vec_v_left_2 = vzipq_s8(vec_v_2_left_tmp1, vec_v_2_left_tmp0);
|
||||
int8x16x2_t vec_v_right_2 = vzipq_s8(vec_v_2_right_tmp1, vec_v_2_right_tmp0);
|
||||
vec_c[2] += vec_v_left_2.val[0];
|
||||
vec_c[2] += vec_v_right_2.val[0];
|
||||
vec_c[3] += vec_v_left_2.val[1];
|
||||
vec_c[3] += vec_v_right_2.val[1];
|
||||
|
||||
uint8x16_t vec_a_3 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 3 * 16);
|
||||
uint8x16_t vec_a3_top = vshrq_n_u8(vec_a_3, 4);
|
||||
uint8x16_t vec_a3_bot = vandq_u8(vec_a_3, vec_mask);
|
||||
int8x16_t vec_v_3_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 4], vec_a3_top);
|
||||
int8x16_t vec_v_3_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 5], vec_a3_top);
|
||||
int8x16_t vec_v_3_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 6], vec_a3_bot);
|
||||
int8x16_t vec_v_3_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 7], vec_a3_bot);
|
||||
int8x16x2_t vec_v_left_3 = vzipq_s8(vec_v_3_left_tmp1, vec_v_3_left_tmp0);
|
||||
int8x16x2_t vec_v_right_3 = vzipq_s8(vec_v_3_right_tmp1, vec_v_3_right_tmp0);
|
||||
vec_c[2] += vec_v_left_3.val[0];
|
||||
vec_c[2] += vec_v_right_3.val[0];
|
||||
vec_c[3] += vec_v_left_3.val[1];
|
||||
vec_c[3] += vec_v_right_3.val[1];
|
||||
|
||||
}
|
||||
|
||||
int32x4_t vec_v_bot_low_low_0 = vmovl_s16(vget_low_s16(vec_c[0]));
|
||||
int32x4_t vec_v_bot_low_high_0 = vmovl_high_s16(vec_c[0]);
|
||||
vst1q_s32(c + i + 0, vld1q_s32(c + i + 0) + vec_v_bot_low_low_0);
|
||||
vst1q_s32(c + i + 4, vld1q_s32(c + i + 4) + vec_v_bot_low_high_0);
|
||||
int32x4_t vec_v_bot_low_low_1 = vmovl_s16(vget_low_s16(vec_c[1]));
|
||||
int32x4_t vec_v_bot_low_high_1 = vmovl_high_s16(vec_c[1]);
|
||||
vst1q_s32(c + i + 8, vld1q_s32(c + i + 8) + vec_v_bot_low_low_1);
|
||||
vst1q_s32(c + i + 12, vld1q_s32(c + i + 12) + vec_v_bot_low_high_1);
|
||||
int32x4_t vec_v_bot_low_low_2 = vmovl_s16(vget_low_s16(vec_c[2]));
|
||||
int32x4_t vec_v_bot_low_high_2 = vmovl_high_s16(vec_c[2]);
|
||||
vst1q_s32(c + i + 16, vld1q_s32(c + i + 16) + vec_v_bot_low_low_2);
|
||||
vst1q_s32(c + i + 20, vld1q_s32(c + i + 20) + vec_v_bot_low_high_2);
|
||||
int32x4_t vec_v_bot_low_low_3 = vmovl_s16(vget_low_s16(vec_c[3]));
|
||||
int32x4_t vec_v_bot_low_high_3 = vmovl_high_s16(vec_c[3]);
|
||||
vst1q_s32(c + i + 24, vld1q_s32(c + i + 24) + vec_v_bot_low_low_3);
|
||||
vst1q_s32(c + i + 28, vld1q_s32(c + i + 28) + vec_v_bot_low_high_3);
|
||||
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
int32_t qgemm_lut_1536_4096(void* A, void* LUT, void* Scales, void* LUT_Scales, void* C) {
|
||||
alignas(32) uint32_t CBits[BM1536_4096];
|
||||
memset(&(CBits[0]), 0, BM1536_4096 * sizeof(int32_t));
|
||||
#pragma unroll
|
||||
for (int32_t k_outer = 0; k_outer < 4096 / BBK1536_4096; ++k_outer) {
|
||||
tbl_impl_1536_4096((&(((int32_t*)CBits)[0])), (&(((int8_t*)LUT)[(k_outer * BBK1536_4096 / 2 * 32)])), (&(((uint8_t*)A)[(k_outer * BBK1536_4096 / 2 / 2 * BM1536_4096)])));
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i = 0; i < BM1536_4096; i++) {
|
||||
((bitnet_float_type*)C)[i] = (((int32_t*)CBits)[i]) / ((bitnet_float_type*)LUT_Scales)[0] * ((bitnet_float_type*)Scales)[0];
|
||||
}
|
||||
return 0;
|
||||
};
|
||||
#include <arm_neon.h>
|
||||
|
||||
#define BM1536_1536 128
|
||||
#define BBK1536_1536 64
|
||||
inline void tbl_impl_1536_1536(int32_t* c, int8_t* lut, uint8_t* a) {
|
||||
#ifdef __ARM_NEON
|
||||
const int KK = BBK1536_1536 / 2;
|
||||
const uint8x16_t vec_mask = vdupq_n_u8(0x0f);
|
||||
const int8x16_t vec_zero = vdupq_n_s16(0x0000);
|
||||
int8x16_t vec_lut[2 * KK];
|
||||
int16x8_t vec_c[8];
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 2 * KK; k++) {
|
||||
vec_lut[k] = vld1q_s8(lut + k * 16);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < BM1536_1536; i += 64) {
|
||||
#pragma unroll
|
||||
for (int i=0; i<8; i++) {
|
||||
vec_c[i] = vandq_s16(vec_c[i], vec_zero);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int k = 0; k < KK / 2; k++) {
|
||||
|
||||
uint8x16_t vec_a_0 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 0 * 16);
|
||||
uint8x16_t vec_a0_top = vshrq_n_u8(vec_a_0, 4);
|
||||
uint8x16_t vec_a0_bot = vandq_u8(vec_a_0, vec_mask);
|
||||
int8x16_t vec_v_0_left_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 0], vec_a0_top);
|
||||
int8x16_t vec_v_0_left_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 1], vec_a0_top);
|
||||
int8x16_t vec_v_0_right_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 2], vec_a0_bot);
|
||||
int8x16_t vec_v_0_right_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 3], vec_a0_bot);
|
||||
int8x16x2_t vec_v_left_0 = vzipq_s8(vec_v_0_left_tmp1, vec_v_0_left_tmp0);
|
||||
int8x16x2_t vec_v_right_0 = vzipq_s8(vec_v_0_right_tmp1, vec_v_0_right_tmp0);
|
||||
vec_c[0] += vec_v_left_0.val[0];
|
||||
vec_c[0] += vec_v_right_0.val[0];
|
||||
vec_c[1] += vec_v_left_0.val[1];
|
||||
vec_c[1] += vec_v_right_0.val[1];
|
||||
|
||||
uint8x16_t vec_a_1 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 1 * 16);
|
||||
uint8x16_t vec_a1_top = vshrq_n_u8(vec_a_1, 4);
|
||||
uint8x16_t vec_a1_bot = vandq_u8(vec_a_1, vec_mask);
|
||||
int8x16_t vec_v_1_left_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 0], vec_a1_top);
|
||||
int8x16_t vec_v_1_left_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 1], vec_a1_top);
|
||||
int8x16_t vec_v_1_right_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 2], vec_a1_bot);
|
||||
int8x16_t vec_v_1_right_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 3], vec_a1_bot);
|
||||
int8x16x2_t vec_v_left_1 = vzipq_s8(vec_v_1_left_tmp1, vec_v_1_left_tmp0);
|
||||
int8x16x2_t vec_v_right_1 = vzipq_s8(vec_v_1_right_tmp1, vec_v_1_right_tmp0);
|
||||
vec_c[2] += vec_v_left_1.val[0];
|
||||
vec_c[2] += vec_v_right_1.val[0];
|
||||
vec_c[3] += vec_v_left_1.val[1];
|
||||
vec_c[3] += vec_v_right_1.val[1];
|
||||
|
||||
uint8x16_t vec_a_2 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 2 * 16);
|
||||
uint8x16_t vec_a2_top = vshrq_n_u8(vec_a_2, 4);
|
||||
uint8x16_t vec_a2_bot = vandq_u8(vec_a_2, vec_mask);
|
||||
int8x16_t vec_v_2_left_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 0], vec_a2_top);
|
||||
int8x16_t vec_v_2_left_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 1], vec_a2_top);
|
||||
int8x16_t vec_v_2_right_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 2], vec_a2_bot);
|
||||
int8x16_t vec_v_2_right_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 3], vec_a2_bot);
|
||||
int8x16x2_t vec_v_left_2 = vzipq_s8(vec_v_2_left_tmp1, vec_v_2_left_tmp0);
|
||||
int8x16x2_t vec_v_right_2 = vzipq_s8(vec_v_2_right_tmp1, vec_v_2_right_tmp0);
|
||||
vec_c[4] += vec_v_left_2.val[0];
|
||||
vec_c[4] += vec_v_right_2.val[0];
|
||||
vec_c[5] += vec_v_left_2.val[1];
|
||||
vec_c[5] += vec_v_right_2.val[1];
|
||||
|
||||
uint8x16_t vec_a_3 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 3 * 16);
|
||||
uint8x16_t vec_a3_top = vshrq_n_u8(vec_a_3, 4);
|
||||
uint8x16_t vec_a3_bot = vandq_u8(vec_a_3, vec_mask);
|
||||
int8x16_t vec_v_3_left_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 0], vec_a3_top);
|
||||
int8x16_t vec_v_3_left_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 1], vec_a3_top);
|
||||
int8x16_t vec_v_3_right_tmp0 = vqtbl1q_s8(vec_lut[4 * k + 2], vec_a3_bot);
|
||||
int8x16_t vec_v_3_right_tmp1 = vqtbl1q_s8(vec_lut[4 * k + 3], vec_a3_bot);
|
||||
int8x16x2_t vec_v_left_3 = vzipq_s8(vec_v_3_left_tmp1, vec_v_3_left_tmp0);
|
||||
int8x16x2_t vec_v_right_3 = vzipq_s8(vec_v_3_right_tmp1, vec_v_3_right_tmp0);
|
||||
vec_c[6] += vec_v_left_3.val[0];
|
||||
vec_c[6] += vec_v_right_3.val[0];
|
||||
vec_c[7] += vec_v_left_3.val[1];
|
||||
vec_c[7] += vec_v_right_3.val[1];
|
||||
|
||||
}
|
||||
|
||||
int32x4_t vec_v_bot_low_low_0 = vmovl_s16(vget_low_s16(vec_c[0]));
|
||||
int32x4_t vec_v_bot_low_high_0 = vmovl_high_s16(vec_c[0]);
|
||||
vst1q_s32(c + i + 0, vld1q_s32(c + i + 0) + vec_v_bot_low_low_0);
|
||||
vst1q_s32(c + i + 4, vld1q_s32(c + i + 4) + vec_v_bot_low_high_0);
|
||||
int32x4_t vec_v_bot_low_low_1 = vmovl_s16(vget_low_s16(vec_c[1]));
|
||||
int32x4_t vec_v_bot_low_high_1 = vmovl_high_s16(vec_c[1]);
|
||||
vst1q_s32(c + i + 8, vld1q_s32(c + i + 8) + vec_v_bot_low_low_1);
|
||||
vst1q_s32(c + i + 12, vld1q_s32(c + i + 12) + vec_v_bot_low_high_1);
|
||||
int32x4_t vec_v_bot_low_low_2 = vmovl_s16(vget_low_s16(vec_c[2]));
|
||||
int32x4_t vec_v_bot_low_high_2 = vmovl_high_s16(vec_c[2]);
|
||||
vst1q_s32(c + i + 16, vld1q_s32(c + i + 16) + vec_v_bot_low_low_2);
|
||||
vst1q_s32(c + i + 20, vld1q_s32(c + i + 20) + vec_v_bot_low_high_2);
|
||||
int32x4_t vec_v_bot_low_low_3 = vmovl_s16(vget_low_s16(vec_c[3]));
|
||||
int32x4_t vec_v_bot_low_high_3 = vmovl_high_s16(vec_c[3]);
|
||||
vst1q_s32(c + i + 24, vld1q_s32(c + i + 24) + vec_v_bot_low_low_3);
|
||||
vst1q_s32(c + i + 28, vld1q_s32(c + i + 28) + vec_v_bot_low_high_3);
|
||||
int32x4_t vec_v_bot_low_low_4 = vmovl_s16(vget_low_s16(vec_c[4]));
|
||||
int32x4_t vec_v_bot_low_high_4 = vmovl_high_s16(vec_c[4]);
|
||||
vst1q_s32(c + i + 32, vld1q_s32(c + i + 32) + vec_v_bot_low_low_4);
|
||||
vst1q_s32(c + i + 36, vld1q_s32(c + i + 36) + vec_v_bot_low_high_4);
|
||||
int32x4_t vec_v_bot_low_low_5 = vmovl_s16(vget_low_s16(vec_c[5]));
|
||||
int32x4_t vec_v_bot_low_high_5 = vmovl_high_s16(vec_c[5]);
|
||||
vst1q_s32(c + i + 40, vld1q_s32(c + i + 40) + vec_v_bot_low_low_5);
|
||||
vst1q_s32(c + i + 44, vld1q_s32(c + i + 44) + vec_v_bot_low_high_5);
|
||||
int32x4_t vec_v_bot_low_low_6 = vmovl_s16(vget_low_s16(vec_c[6]));
|
||||
int32x4_t vec_v_bot_low_high_6 = vmovl_high_s16(vec_c[6]);
|
||||
vst1q_s32(c + i + 48, vld1q_s32(c + i + 48) + vec_v_bot_low_low_6);
|
||||
vst1q_s32(c + i + 52, vld1q_s32(c + i + 52) + vec_v_bot_low_high_6);
|
||||
int32x4_t vec_v_bot_low_low_7 = vmovl_s16(vget_low_s16(vec_c[7]));
|
||||
int32x4_t vec_v_bot_low_high_7 = vmovl_high_s16(vec_c[7]);
|
||||
vst1q_s32(c + i + 56, vld1q_s32(c + i + 56) + vec_v_bot_low_low_7);
|
||||
vst1q_s32(c + i + 60, vld1q_s32(c + i + 60) + vec_v_bot_low_high_7);
|
||||
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
int32_t qgemm_lut_1536_1536(void* A, void* LUT, void* Scales, void* LUT_Scales, void* C) {
|
||||
alignas(32) uint32_t CBits[BM1536_1536];
|
||||
memset(&(CBits[0]), 0, BM1536_1536 * sizeof(int32_t));
|
||||
#pragma unroll
|
||||
for (int32_t k_outer = 0; k_outer < 1536 / BBK1536_1536; ++k_outer) {
|
||||
tbl_impl_1536_1536((&(((int32_t*)CBits)[0])), (&(((int8_t*)LUT)[(k_outer * BBK1536_1536 / 2 * 32)])), (&(((uint8_t*)A)[(k_outer * BBK1536_1536 / 2 / 2 * BM1536_1536)])));
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i = 0; i < BM1536_1536; i++) {
|
||||
((bitnet_float_type*)C)[i] = (((int32_t*)CBits)[i]) / ((bitnet_float_type*)LUT_Scales)[0] * ((bitnet_float_type*)Scales)[0];
|
||||
}
|
||||
return 0;
|
||||
};
|
||||
#include <arm_neon.h>
|
||||
|
||||
#define BM4096_1536 256
|
||||
#define BBK4096_1536 128
|
||||
inline void tbl_impl_4096_1536(int32_t* c, int8_t* lut, uint8_t* a) {
|
||||
#ifdef __ARM_NEON
|
||||
const int KK = BBK4096_1536 / 2;
|
||||
const uint8x16_t vec_mask = vdupq_n_u8(0x0f);
|
||||
const int8x16_t vec_zero = vdupq_n_s16(0x0000);
|
||||
int8x16_t vec_lut[2 * KK];
|
||||
int16x8_t vec_c[4];
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 2 * KK; k++) {
|
||||
vec_lut[k] = vld1q_s8(lut + k * 16);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < BM4096_1536; i += 32) {
|
||||
#pragma unroll
|
||||
for (int i=0; i<4; i++) {
|
||||
vec_c[i] = vandq_s16(vec_c[i], vec_zero);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int k = 0; k < KK / 4; k++) {
|
||||
|
||||
uint8x16_t vec_a_0 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 0 * 16);
|
||||
uint8x16_t vec_a0_top = vshrq_n_u8(vec_a_0, 4);
|
||||
uint8x16_t vec_a0_bot = vandq_u8(vec_a_0, vec_mask);
|
||||
int8x16_t vec_v_0_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 0], vec_a0_top);
|
||||
int8x16_t vec_v_0_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 1], vec_a0_top);
|
||||
int8x16_t vec_v_0_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 2], vec_a0_bot);
|
||||
int8x16_t vec_v_0_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 3], vec_a0_bot);
|
||||
int8x16x2_t vec_v_left_0 = vzipq_s8(vec_v_0_left_tmp1, vec_v_0_left_tmp0);
|
||||
int8x16x2_t vec_v_right_0 = vzipq_s8(vec_v_0_right_tmp1, vec_v_0_right_tmp0);
|
||||
vec_c[0] += vec_v_left_0.val[0];
|
||||
vec_c[0] += vec_v_right_0.val[0];
|
||||
vec_c[1] += vec_v_left_0.val[1];
|
||||
vec_c[1] += vec_v_right_0.val[1];
|
||||
|
||||
uint8x16_t vec_a_1 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 1 * 16);
|
||||
uint8x16_t vec_a1_top = vshrq_n_u8(vec_a_1, 4);
|
||||
uint8x16_t vec_a1_bot = vandq_u8(vec_a_1, vec_mask);
|
||||
int8x16_t vec_v_1_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 4], vec_a1_top);
|
||||
int8x16_t vec_v_1_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 5], vec_a1_top);
|
||||
int8x16_t vec_v_1_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 6], vec_a1_bot);
|
||||
int8x16_t vec_v_1_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 7], vec_a1_bot);
|
||||
int8x16x2_t vec_v_left_1 = vzipq_s8(vec_v_1_left_tmp1, vec_v_1_left_tmp0);
|
||||
int8x16x2_t vec_v_right_1 = vzipq_s8(vec_v_1_right_tmp1, vec_v_1_right_tmp0);
|
||||
vec_c[0] += vec_v_left_1.val[0];
|
||||
vec_c[0] += vec_v_right_1.val[0];
|
||||
vec_c[1] += vec_v_left_1.val[1];
|
||||
vec_c[1] += vec_v_right_1.val[1];
|
||||
|
||||
uint8x16_t vec_a_2 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 2 * 16);
|
||||
uint8x16_t vec_a2_top = vshrq_n_u8(vec_a_2, 4);
|
||||
uint8x16_t vec_a2_bot = vandq_u8(vec_a_2, vec_mask);
|
||||
int8x16_t vec_v_2_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 0], vec_a2_top);
|
||||
int8x16_t vec_v_2_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 1], vec_a2_top);
|
||||
int8x16_t vec_v_2_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 2], vec_a2_bot);
|
||||
int8x16_t vec_v_2_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 3], vec_a2_bot);
|
||||
int8x16x2_t vec_v_left_2 = vzipq_s8(vec_v_2_left_tmp1, vec_v_2_left_tmp0);
|
||||
int8x16x2_t vec_v_right_2 = vzipq_s8(vec_v_2_right_tmp1, vec_v_2_right_tmp0);
|
||||
vec_c[2] += vec_v_left_2.val[0];
|
||||
vec_c[2] += vec_v_right_2.val[0];
|
||||
vec_c[3] += vec_v_left_2.val[1];
|
||||
vec_c[3] += vec_v_right_2.val[1];
|
||||
|
||||
uint8x16_t vec_a_3 = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + 3 * 16);
|
||||
uint8x16_t vec_a3_top = vshrq_n_u8(vec_a_3, 4);
|
||||
uint8x16_t vec_a3_bot = vandq_u8(vec_a_3, vec_mask);
|
||||
int8x16_t vec_v_3_left_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 4], vec_a3_top);
|
||||
int8x16_t vec_v_3_left_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 5], vec_a3_top);
|
||||
int8x16_t vec_v_3_right_tmp0 = vqtbl1q_s8(vec_lut[8 * k + 6], vec_a3_bot);
|
||||
int8x16_t vec_v_3_right_tmp1 = vqtbl1q_s8(vec_lut[8 * k + 7], vec_a3_bot);
|
||||
int8x16x2_t vec_v_left_3 = vzipq_s8(vec_v_3_left_tmp1, vec_v_3_left_tmp0);
|
||||
int8x16x2_t vec_v_right_3 = vzipq_s8(vec_v_3_right_tmp1, vec_v_3_right_tmp0);
|
||||
vec_c[2] += vec_v_left_3.val[0];
|
||||
vec_c[2] += vec_v_right_3.val[0];
|
||||
vec_c[3] += vec_v_left_3.val[1];
|
||||
vec_c[3] += vec_v_right_3.val[1];
|
||||
|
||||
}
|
||||
|
||||
int32x4_t vec_v_bot_low_low_0 = vmovl_s16(vget_low_s16(vec_c[0]));
|
||||
int32x4_t vec_v_bot_low_high_0 = vmovl_high_s16(vec_c[0]);
|
||||
vst1q_s32(c + i + 0, vld1q_s32(c + i + 0) + vec_v_bot_low_low_0);
|
||||
vst1q_s32(c + i + 4, vld1q_s32(c + i + 4) + vec_v_bot_low_high_0);
|
||||
int32x4_t vec_v_bot_low_low_1 = vmovl_s16(vget_low_s16(vec_c[1]));
|
||||
int32x4_t vec_v_bot_low_high_1 = vmovl_high_s16(vec_c[1]);
|
||||
vst1q_s32(c + i + 8, vld1q_s32(c + i + 8) + vec_v_bot_low_low_1);
|
||||
vst1q_s32(c + i + 12, vld1q_s32(c + i + 12) + vec_v_bot_low_high_1);
|
||||
int32x4_t vec_v_bot_low_low_2 = vmovl_s16(vget_low_s16(vec_c[2]));
|
||||
int32x4_t vec_v_bot_low_high_2 = vmovl_high_s16(vec_c[2]);
|
||||
vst1q_s32(c + i + 16, vld1q_s32(c + i + 16) + vec_v_bot_low_low_2);
|
||||
vst1q_s32(c + i + 20, vld1q_s32(c + i + 20) + vec_v_bot_low_high_2);
|
||||
int32x4_t vec_v_bot_low_low_3 = vmovl_s16(vget_low_s16(vec_c[3]));
|
||||
int32x4_t vec_v_bot_low_high_3 = vmovl_high_s16(vec_c[3]);
|
||||
vst1q_s32(c + i + 24, vld1q_s32(c + i + 24) + vec_v_bot_low_low_3);
|
||||
vst1q_s32(c + i + 28, vld1q_s32(c + i + 28) + vec_v_bot_low_high_3);
|
||||
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
int32_t qgemm_lut_4096_1536(void* A, void* LUT, void* Scales, void* LUT_Scales, void* C) {
|
||||
alignas(32) uint32_t CBits[BM4096_1536];
|
||||
memset(&(CBits[0]), 0, BM4096_1536 * sizeof(int32_t));
|
||||
#pragma unroll
|
||||
for (int32_t k_outer = 0; k_outer < 1536 / BBK4096_1536; ++k_outer) {
|
||||
tbl_impl_4096_1536((&(((int32_t*)CBits)[0])), (&(((int8_t*)LUT)[(k_outer * BBK4096_1536 / 2 * 32)])), (&(((uint8_t*)A)[(k_outer * BBK4096_1536 / 2 / 2 * BM4096_1536)])));
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i = 0; i < BM4096_1536; i++) {
|
||||
((bitnet_float_type*)C)[i] = (((int32_t*)CBits)[i]) / ((bitnet_float_type*)LUT_Scales)[0] * ((bitnet_float_type*)Scales)[0];
|
||||
}
|
||||
return 0;
|
||||
};
|
||||
|
||||
template<int K>
|
||||
void preprocessor_k(void* B, void* LUT_Scales, void* QLUT) {{
|
||||
partial_max_reset((&(((bitnet_float_type*)LUT_Scales)[0])));
|
||||
per_tensor_quant(K, (&(((bitnet_float_type*)LUT_Scales)[0])), (&(((bitnet_float_type*)B)[0])));
|
||||
|
||||
lut_ctor<K>((&(((int8_t*)QLUT)[0])), (&(((bitnet_float_type*)B)[0])), (&(((bitnet_float_type*)LUT_Scales)[0])));
|
||||
}}
|
||||
void ggml_preprocessor(int m, int k, void* B, void* LUT_Scales, void* QLUT) {
|
||||
if (m == 1536 && k == 4096) {
|
||||
preprocessor_k<4096>(B, LUT_Scales, QLUT);
|
||||
}
|
||||
else if (m == 1536 && k == 1536) {
|
||||
preprocessor_k<1536>(B, LUT_Scales, QLUT);
|
||||
}
|
||||
else if (m == 4096 && k == 1536) {
|
||||
preprocessor_k<1536>(B, LUT_Scales, QLUT);
|
||||
}
|
||||
}
|
||||
void ggml_qgemm_lut(int m, int k, void* A, void* LUT, void* Scales, void* LUT_Scales, void* C) {
|
||||
if (m == 1536 && k == 4096) {
|
||||
qgemm_lut_1536_4096(A, LUT, Scales, LUT_Scales, C);
|
||||
}
|
||||
else if (m == 1536 && k == 1536) {
|
||||
qgemm_lut_1536_1536(A, LUT, Scales, LUT_Scales, C);
|
||||
}
|
||||
else if (m == 4096 && k == 1536) {
|
||||
qgemm_lut_4096_1536(A, LUT, Scales, LUT_Scales, C);
|
||||
}
|
||||
}
|
||||
|
||||
void ggml_bitnet_transform_tensor(struct ggml_tensor * tensor) {
|
||||
if (!(is_type_supported(tensor->type) && tensor->backend == GGML_BACKEND_TYPE_CPU && tensor->extra == nullptr)) {
|
||||
return;
|
||||
}
|
||||
|
||||
int k = tensor->ne[0];
|
||||
int m = tensor->ne[1];
|
||||
const int lut_scales_size = 1;
|
||||
const int scales_size = 1;
|
||||
int bk = 0;
|
||||
int bm = 0;
|
||||
|
||||
if (m == 1536 && k == 4096) {
|
||||
bm = BM1536_4096;
|
||||
bk = BBK1536_4096;
|
||||
}
|
||||
else if (m == 1536 && k == 1536) {
|
||||
bm = BM1536_1536;
|
||||
bk = BBK1536_1536;
|
||||
}
|
||||
else if (m == 4096 && k == 1536) {
|
||||
bm = BM4096_1536;
|
||||
bk = BBK4096_1536;
|
||||
}
|
||||
|
||||
const int n_tile_num = m / bm;
|
||||
const int BK = bk;
|
||||
uint8_t * qweights;
|
||||
bitnet_float_type * scales;
|
||||
|
||||
scales = (bitnet_float_type *) aligned_malloc(sizeof(bitnet_float_type));
|
||||
qweights = (uint8_t *) tensor->data;
|
||||
float * i2_scales = (float * )(qweights + k * m / 4);
|
||||
scales[0] = (bitnet_float_type) i2_scales[0];
|
||||
|
||||
tensor->extra = bitnet_tensor_extras + bitnet_tensor_extras_index;
|
||||
bitnet_tensor_extras[bitnet_tensor_extras_index++] = {
|
||||
/* .lut_scales_size = */ lut_scales_size,
|
||||
/* .scales_size = */ scales_size,
|
||||
/* .n_tile_num = */ n_tile_num,
|
||||
/* .qweights = */ qweights,
|
||||
/* .scales = */ scales
|
||||
};
|
||||
}
|
||||
#endif
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,21 @@
|
||||
[Kernels_0]
|
||||
m = 1536
|
||||
k = 4096
|
||||
bm = 256
|
||||
bk = 128
|
||||
bmm = 32
|
||||
|
||||
[Kernels_1]
|
||||
m = 1536
|
||||
k = 1536
|
||||
bm = 128
|
||||
bk = 64
|
||||
bmm = 64
|
||||
|
||||
[Kernels_2]
|
||||
m = 4096
|
||||
k = 1536
|
||||
bm = 256
|
||||
bk = 128
|
||||
bmm = 32
|
||||
|
||||
@@ -0,0 +1,21 @@
|
||||
[Kernels_0]
|
||||
m = 1536
|
||||
k = 4096
|
||||
bm = 256
|
||||
bk = 96
|
||||
bmm = 32
|
||||
|
||||
[Kernels_1]
|
||||
m = 1536
|
||||
k = 1536
|
||||
bm = 128
|
||||
bk = 192
|
||||
bmm = 32
|
||||
|
||||
[Kernels_2]
|
||||
m = 4096
|
||||
k = 1536
|
||||
bm = 256
|
||||
bk = 96
|
||||
bmm = 64
|
||||
|
||||
@@ -0,0 +1,11 @@
|
||||
# These requirements include all dependencies for all top-level python scripts
|
||||
# for llama.cpp. Avoid adding packages here directly.
|
||||
#
|
||||
# Package versions must stay compatible across all top-level python scripts.
|
||||
#
|
||||
|
||||
-r 3rdparty/llama.cpp/requirements/requirements-convert_legacy_llama.txt
|
||||
-r 3rdparty/llama.cpp/requirements/requirements-convert_hf_to_gguf.txt
|
||||
-r 3rdparty/llama.cpp/requirements/requirements-convert_hf_to_gguf_update.txt
|
||||
-r 3rdparty/llama.cpp/requirements/requirements-convert_llama_ggml_to_gguf.txt
|
||||
-r 3rdparty/llama.cpp/requirements/requirements-convert_lora_to_gguf.txt
|
||||
@@ -0,0 +1,53 @@
|
||||
import os
|
||||
import sys
|
||||
import signal
|
||||
import platform
|
||||
import argparse
|
||||
import subprocess
|
||||
|
||||
def run_command(command, shell=False):
|
||||
"""Run a system command and ensure it succeeds."""
|
||||
try:
|
||||
subprocess.run(command, shell=shell, check=True)
|
||||
except subprocess.CalledProcessError as e:
|
||||
print(f"Error occurred while running command: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
def run_inference():
|
||||
build_dir = "build"
|
||||
if platform.system() == "Windows":
|
||||
main_path = os.path.join(build_dir, "bin", "Release", "llama-cli.exe")
|
||||
if not os.path.exists(main_path):
|
||||
main_path = os.path.join(build_dir, "bin", "llama-cli")
|
||||
else:
|
||||
main_path = os.path.join(build_dir, "bin", "llama-cli")
|
||||
command = [
|
||||
f'{main_path}',
|
||||
'-m', args.model,
|
||||
'-n', str(args.n_predict),
|
||||
'-t', str(args.threads),
|
||||
'-p', args.prompt,
|
||||
'-ngl', '0',
|
||||
'-c', str(args.ctx_size),
|
||||
'--temp', str(args.temperature),
|
||||
"-b", "1"
|
||||
]
|
||||
run_command(command)
|
||||
|
||||
def signal_handler(sig, frame):
|
||||
print("Ctrl+C pressed, exiting...")
|
||||
sys.exit(0)
|
||||
|
||||
if __name__ == "__main__":
|
||||
signal.signal(signal.SIGINT, signal_handler)
|
||||
# Usage: python run_inference.py -p "Microsoft Corporation is an American multinational corporation and technology company headquartered in Redmond, Washington."
|
||||
parser = argparse.ArgumentParser(description='Run inference')
|
||||
parser.add_argument("-m", "--model", type=str, help="Path to model file", required=False, default="models/bitnet_b1_58-3B/ggml-model-i2_s.gguf")
|
||||
parser.add_argument("-n", "--n-predict", type=int, help="Number of tokens to predict when generating text", required=False, default=128)
|
||||
parser.add_argument("-p", "--prompt", type=str, help="Prompt to generate text from", required=True)
|
||||
parser.add_argument("-t", "--threads", type=int, help="Number of threads to use", required=False, default=2)
|
||||
parser.add_argument("-c", "--ctx-size", type=int, help="Size of the prompt context", required=False, default=2048)
|
||||
parser.add_argument("-temp", "--temperature", type=float, help="Temperature, a hyperparameter that controls the randomness of the generated text", required=False, default=0.8)
|
||||
|
||||
args = parser.parse_args()
|
||||
run_inference()
|
||||
+202
@@ -0,0 +1,202 @@
|
||||
import subprocess
|
||||
import signal
|
||||
import sys
|
||||
import os
|
||||
import platform
|
||||
import argparse
|
||||
import logging
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
|
||||
logger = logging.getLogger("setup_env")
|
||||
|
||||
SUPPORTED_HF_MODELS = {
|
||||
"1bitLLM/bitnet_b1_58-large": {
|
||||
"model_name": "bitnet_b1_58-large",
|
||||
},
|
||||
"1bitLLM/bitnet_b1_58-3B": {
|
||||
"model_name": "bitnet_b1_58-3B",
|
||||
},
|
||||
"HF1BitLLM/Llama3-8B-1.58-100B-tokens": {
|
||||
"model_name": "Llama3-8B-1.58-100B-tokens",
|
||||
}
|
||||
}
|
||||
|
||||
SUPPORTED_QUANT_TYPES = {
|
||||
"arm64": ["i2_s", "tl1"],
|
||||
"x86_64": ["i2_s", "tl2"]
|
||||
}
|
||||
|
||||
COMPILER_EXTRA_ARGS = {
|
||||
"arm64": ["-DBITNET_ARM_TL1=ON"],
|
||||
"x86_64": ["-DBITNET_X86_TL2=ON"]
|
||||
}
|
||||
|
||||
OS_EXTRA_ARGS = {
|
||||
"Windows":["-T", "ClangCL"],
|
||||
"Linux": ["-DCMAKE_C_COMPILER=clang", "-DCMAKE_CXX_COMPILER=clang++"]
|
||||
}
|
||||
|
||||
ARCH_ALIAS = {
|
||||
"AMD64": "x86_64",
|
||||
"x86": "x86_64",
|
||||
"x86_64": "x86_64",
|
||||
"aarch64": "arm64",
|
||||
"arm64": "arm64",
|
||||
"ARM64": "arm64",
|
||||
}
|
||||
|
||||
def system_info():
|
||||
return platform.system(), ARCH_ALIAS[platform.machine()]
|
||||
|
||||
def get_model_name():
|
||||
if args.hf_repo:
|
||||
return SUPPORTED_HF_MODELS[args.hf_repo]["model_name"]
|
||||
return os.path.basename(os.path.normpath(args.model_dir))
|
||||
|
||||
def run_command(command, shell=False, log_step=None):
|
||||
"""Run a system command and ensure it succeeds."""
|
||||
if log_step:
|
||||
log_file = os.path.join(args.log_dir, log_step + ".log")
|
||||
with open(log_file, "w") as f:
|
||||
try:
|
||||
subprocess.run(command, shell=shell, check=True, stdout=f, stderr=f)
|
||||
except subprocess.CalledProcessError as e:
|
||||
logging.error(f"Error occurred while running command: {e}, check details in {log_file}")
|
||||
sys.exit(1)
|
||||
else:
|
||||
try:
|
||||
subprocess.run(command, shell=shell, check=True)
|
||||
except subprocess.CalledProcessError as e:
|
||||
logging.error(f"Error occurred while running command: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
def prepare_model():
|
||||
_, arch = system_info()
|
||||
hf_url = args.hf_repo
|
||||
model_dir = args.model_dir
|
||||
quant_type = args.quant_type
|
||||
quant_embd = args.quant_embd
|
||||
if hf_url is not None:
|
||||
# download the model
|
||||
model_dir = os.path.join(model_dir, SUPPORTED_HF_MODELS[hf_url]["model_name"])
|
||||
Path(model_dir).mkdir(parents=True, exist_ok=True)
|
||||
logging.info(f"Downloading model {hf_url} from HuggingFace to {model_dir}...")
|
||||
run_command(["huggingface-cli", "download", hf_url, "--local-dir", model_dir], log_step="download_model")
|
||||
elif not os.path.exists(model_dir):
|
||||
logging.error(f"Model directory {model_dir} does not exist.")
|
||||
sys.exit(1)
|
||||
else:
|
||||
logging.info(f"Loading model from directory {model_dir}.")
|
||||
gguf_path = os.path.join(model_dir, "ggml-model-" + quant_type + ".gguf")
|
||||
if not os.path.exists(gguf_path) or os.path.getsize(gguf_path) == 0:
|
||||
logging.info(f"Converting HF model to GGUF format...")
|
||||
if quant_type.startswith("tl"):
|
||||
run_command([sys.executable, "utils/convert-hf-to-gguf-bitnet.py", model_dir, "--outtype", quant_type, "--quant-embd"], log_step="convert_to_tl")
|
||||
else: # i2s
|
||||
# convert to f32
|
||||
run_command([sys.executable, "utils/convert-hf-to-gguf-bitnet.py", model_dir, "--outtype", "f32"], log_step="convert_to_f32_gguf")
|
||||
f32_model = os.path.join(model_dir, "ggml-model-f32.gguf")
|
||||
i2s_model = os.path.join(model_dir, "ggml-model-i2_s.gguf")
|
||||
# quantize to i2s
|
||||
if platform.system() != "Windows":
|
||||
if quant_embd:
|
||||
run_command(["./build/bin/llama-quantize", "--token-embedding-type", "f16", f32_model, i2s_model, "I2_S", "1", "1"], log_step="quantize_to_i2s")
|
||||
else:
|
||||
run_command(["./build/bin/llama-quantize", f32_model, i2s_model, "I2_S", "1"], log_step="quantize_to_i2s")
|
||||
else:
|
||||
if quant_embd:
|
||||
run_command(["./build/bin/Release/llama-quantize", "--token-embedding-type", "f16", f32_model, i2s_model, "I2_S", "1", "1"], log_step="quantize_to_i2s")
|
||||
else:
|
||||
run_command(["./build/bin/Release/llama-quantize", f32_model, i2s_model, "I2_S", "1"], log_step="quantize_to_i2s")
|
||||
|
||||
logging.info(f"GGUF model saved at {gguf_path}")
|
||||
else:
|
||||
logging.info(f"GGUF model already exists at {gguf_path}")
|
||||
|
||||
def setup_gguf():
|
||||
# Install the pip package
|
||||
run_command([sys.executable, "-m", "pip", "install", "3rdparty/llama.cpp/gguf-py"], log_step="install_gguf")
|
||||
|
||||
def gen_code():
|
||||
_, arch = system_info()
|
||||
if arch == "arm64":
|
||||
if args.use_pretuned:
|
||||
pretuned_kernels = os.path.join("preset_kernels", get_model_name())
|
||||
if not os.path.exists(pretuned_kernels):
|
||||
logging.error(f"Pretuned kernels not found for model {args.hf_repo}")
|
||||
sys.exit(1)
|
||||
if args.quant_type == "tl1":
|
||||
shutil.copyfile(os.path.join(pretuned_kernels, "bitnet-lut-kernels-tl1.h"), "include/bitnet-lut-kernels.h")
|
||||
shutil.copyfile(os.path.join(pretuned_kernels, "kernel_config_tl1.ini"), "include/kernel_config.ini")
|
||||
elif args.quant_type == "tl2":
|
||||
shutil.copyfile(os.path.join(pretuned_kernels, "bitnet-lut-kernels-tl2.h"), "include/bitnet-lut-kernels.h")
|
||||
shutil.copyfile(os.path.join(pretuned_kernels, "kernel_config_tl2.ini"), "include/kernel_config.ini")
|
||||
if get_model_name() == "bitnet_b1_58-large":
|
||||
run_command([sys.executable, "utils/codegen_tl1.py", "--model", "bitnet_b1_58-large", "--BM", "256,128,256", "--BK", "128,64,128", "--bm", "32,64,32"], log_step="codegen")
|
||||
elif get_model_name() == "Llama3-8B-1.58-100B-tokens":
|
||||
run_command([sys.executable, "utils/codegen_tl1.py", "--model", "Llama3-8B-1.58-100B-tokens", "--BM", "256,128,256,128", "--BK", "128,64,128,64", "--bm", "32,64,32,64"], log_step="codegen")
|
||||
elif get_model_name() == "bitnet_b1_58-3B":
|
||||
run_command([sys.executable, "utils/codegen_tl1.py", "--model", "bitnet_b1_58-3B", "--BM", "160,320,320", "--BK", "64,128,64", "--bm", "32,64,32"], log_step="codegen")
|
||||
else:
|
||||
raise NotImplementedError()
|
||||
else:
|
||||
if args.use_pretuned:
|
||||
# cp preset_kernels/model_name/bitnet-lut-kernels_tl1.h to include/bitnet-lut-kernels.h
|
||||
pretuned_kernels = os.path.join("preset_kernels", get_model_name())
|
||||
if not os.path.exists(pretuned_kernels):
|
||||
logging.error(f"Pretuned kernels not found for model {args.hf_repo}")
|
||||
sys.exit(1)
|
||||
shutil.copyfile(os.path.join(pretuned_kernels, "bitnet-lut-kernels-tl2.h"), "include/bitnet-lut-kernels.h")
|
||||
if get_model_name() == "bitnet_b1_58-large":
|
||||
run_command([sys.executable, "utils/codegen_tl2.py", "--model", "bitnet_b1_58-large", "--BM", "256,128,256", "--BK", "96,192,96", "--bm", "32,32,32"], log_step="codegen")
|
||||
elif get_model_name() == "Llama3-8B-1.58-100B-tokens":
|
||||
run_command([sys.executable, "utils/codegen_tl2.py", "--model", "Llama3-8B-1.58-100B-tokens", "--BM", "256,128,256,128", "--BK", "96,96,96,96", "--bm", "32,32,32,32"], log_step="codegen")
|
||||
elif get_model_name() == "bitnet_b1_58-3B":
|
||||
run_command([sys.executable, "utils/codegen_tl2.py", "--model", "bitnet_b1_58-3B", "--BM", "160,320,320", "--BK", "96,96,96", "--bm", "32,32,32"], log_step="codegen")
|
||||
else:
|
||||
raise NotImplementedError()
|
||||
|
||||
|
||||
def compile():
|
||||
# Check if cmake is installed
|
||||
cmake_exists = subprocess.run(["cmake", "--version"], capture_output=True)
|
||||
if cmake_exists.returncode != 0:
|
||||
logging.error("Cmake is not available. Please install CMake and try again.")
|
||||
sys.exit(1)
|
||||
_, arch = system_info()
|
||||
if arch not in COMPILER_EXTRA_ARGS.keys():
|
||||
logging.error(f"Arch {arch} is not supported yet")
|
||||
exit(0)
|
||||
logging.info("Compiling the code using CMake.")
|
||||
run_command(["cmake", "-B", "build", *COMPILER_EXTRA_ARGS[arch], *OS_EXTRA_ARGS.get(platform.system(), [])], log_step="generate_build_files")
|
||||
# run_command(["cmake", "--build", "build", "--target", "llama-cli", "--config", "Release"])
|
||||
run_command(["cmake", "--build", "build", "--config", "Release"], log_step="compile")
|
||||
|
||||
def main():
|
||||
setup_gguf()
|
||||
gen_code()
|
||||
compile()
|
||||
prepare_model()
|
||||
|
||||
def parse_args():
|
||||
_, arch = system_info()
|
||||
parser = argparse.ArgumentParser(description='Setup the environment for running the inference')
|
||||
parser.add_argument("--hf-repo", "-hr", type=str, help="Model used for inference", choices=SUPPORTED_HF_MODELS.keys())
|
||||
parser.add_argument("--model-dir", "-md", type=str, help="Directory to save/load the model", default="models")
|
||||
parser.add_argument("--log-dir", "-ld", type=str, help="Directory to save the logging info", default="logs")
|
||||
parser.add_argument("--quant-type", "-q", type=str, help="Quantization type", choices=SUPPORTED_QUANT_TYPES[arch], default="i2_s")
|
||||
parser.add_argument("--quant-embd", action="store_true", help="Quantize the embeddings to f16")
|
||||
parser.add_argument("--use-pretuned", "-p", action="store_true", help="Use the pretuned kernel parameters")
|
||||
return parser.parse_args()
|
||||
|
||||
def signal_handler(sig, frame):
|
||||
logging.info("Ctrl+C pressed, exiting...")
|
||||
sys.exit(0)
|
||||
|
||||
if __name__ == "__main__":
|
||||
signal.signal(signal.SIGINT, signal_handler)
|
||||
args = parse_args()
|
||||
Path(args.log_dir).mkdir(parents=True, exist_ok=True)
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
main()
|
||||
@@ -0,0 +1,10 @@
|
||||
set(GGML_HEADERS_BITNET ../include/ggml-bitnet.h)
|
||||
set(GGML_SOURCES_BITNET ggml-bitnet-mad.cpp)
|
||||
set(GGML_SOURCES_BITNET ggml-bitnet-lut.cpp)
|
||||
|
||||
include_directories(3rdparty/llama.cpp/ggml/include)
|
||||
|
||||
if ((NOT ${CMAKE_C_COMPILER_ID} MATCHES "Clang") OR
|
||||
(NOT ${CMAKE_CXX_COMPILER_ID} MATCHES "Clang"))
|
||||
message(FATAL_ERROR "Clang is required for Bitnet.cpp compilation")
|
||||
endif()
|
||||
@@ -0,0 +1,163 @@
|
||||
#include <vector>
|
||||
#include <type_traits>
|
||||
|
||||
#include "ggml-bitnet.h"
|
||||
#include "ggml-quants.h"
|
||||
#include "bitnet-lut-kernels.h"
|
||||
|
||||
#if defined(GGML_BITNET_ARM_TL1)
|
||||
|
||||
void ggml_bitnet_init(void) {
|
||||
// LOG(INFO) << "ggml_bitnet_init";
|
||||
|
||||
if (initialized) {
|
||||
return;
|
||||
}
|
||||
initialized = true;
|
||||
|
||||
// if (wrapper == nullptr) {
|
||||
// wrapper = new BITNET::BITNETGeMMWrapper<bitnet_bitnet_float_type>();
|
||||
// }
|
||||
if (bitnet_tensor_extras == nullptr) {
|
||||
bitnet_tensor_extras = new bitnet_tensor_extra[GGML_BITNET_MAX_NODES];
|
||||
}
|
||||
bitnet_tensor_extras_index = 0;
|
||||
}
|
||||
|
||||
void ggml_bitnet_free(void) {
|
||||
// LOG(INFO) << "ggml_bitnet_free";
|
||||
|
||||
if (!initialized) {
|
||||
return;
|
||||
}
|
||||
initialized = false;
|
||||
|
||||
// delete wrapper;
|
||||
// wrapper = nullptr;
|
||||
for (size_t i = 0; i < bitnet_tensor_extras_index; i++) {
|
||||
// aligned_free(bitnet_tensor_extras[i].qweights);
|
||||
// aligned_free(bitnet_tensor_extras[i].scales);
|
||||
}
|
||||
delete[] bitnet_tensor_extras;
|
||||
bitnet_tensor_extras = nullptr;
|
||||
}
|
||||
|
||||
static bool do_permutate(enum ggml_type type) {
|
||||
if (type == GGML_TYPE_TL1) {
|
||||
// Add additional args to decide if permuted I2 or naive I2
|
||||
return false;
|
||||
} else {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
|
||||
bool ggml_bitnet_can_mul_mat(const struct ggml_tensor * src0, const struct ggml_tensor * src1, const struct ggml_tensor * dst) {
|
||||
if ((is_type_supported(src0->type)) &&
|
||||
src1->type == GGML_TYPE_F32 &&
|
||||
dst->type == GGML_TYPE_F32 &&
|
||||
src0->backend == GGML_BACKEND_TYPE_CPU) {
|
||||
if (src1->ne[1] <= 1) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
size_t ggml_bitnet_mul_mat_get_wsize(const struct ggml_tensor * src0, const struct ggml_tensor * src1, const struct ggml_tensor * dst) {
|
||||
const size_t ne01 = src0->ne[1];
|
||||
const size_t ne10 = src1->ne[0];
|
||||
const size_t ne11 = src1->ne[1];
|
||||
const int bits = ggml_bitnet_get_type_bits(src0->type);
|
||||
|
||||
size_t wsize = ne10 * ne11 * 15 * sizeof(int8_t) + 1 * ne11 * 2 * sizeof(bitnet_float_type);
|
||||
if (sizeof(bitnet_float_type) == 2) {
|
||||
// Need fp32 to fp16 conversion
|
||||
wsize += std::max(ne10, ne01) * ne11 * sizeof(bitnet_float_type);
|
||||
}
|
||||
wsize = ((wsize - 1) / 64 + 1) * 64;
|
||||
return wsize;
|
||||
}
|
||||
|
||||
int ggml_bitnet_get_type_bits(enum ggml_type type) {
|
||||
switch (type) {
|
||||
case GGML_TYPE_TL1:
|
||||
return 2;
|
||||
case GGML_TYPE_Q4_0:
|
||||
return 4;
|
||||
default:
|
||||
return 0;
|
||||
}
|
||||
}
|
||||
|
||||
#endif
|
||||
#if defined(GGML_BITNET_X86_TL2)
|
||||
void ggml_bitnet_init(void) {
|
||||
// LOG(INFO) << "ggml_bitnet_init";
|
||||
|
||||
if (initialized) {
|
||||
return;
|
||||
}
|
||||
initialized = true;
|
||||
|
||||
// if (wrapper == nullptr) {
|
||||
// wrapper = new BITNET::BITNETGeMMWrapper<bitnet_bitnet_float_type>();
|
||||
// }
|
||||
if (bitnet_tensor_extras == nullptr) {
|
||||
bitnet_tensor_extras = new bitnet_tensor_extra[GGML_BITNET_MAX_NODES];
|
||||
}
|
||||
bitnet_tensor_extras_index = 0;
|
||||
}
|
||||
|
||||
void ggml_bitnet_free(void) {
|
||||
// LOG(INFO) << "ggml_bitnet_free";
|
||||
|
||||
if (!initialized) {
|
||||
return;
|
||||
}
|
||||
initialized = false;
|
||||
|
||||
// delete wrapper;
|
||||
// wrapper = nullptr;
|
||||
for (size_t i = 0; i < bitnet_tensor_extras_index; i++) {
|
||||
// aligned_free(bitnet_tensor_extras[i].qweights);
|
||||
// aligned_free(bitnet_tensor_extras[i].scales);
|
||||
}
|
||||
delete[] bitnet_tensor_extras;
|
||||
bitnet_tensor_extras = nullptr;
|
||||
}
|
||||
|
||||
bool ggml_bitnet_can_mul_mat(const struct ggml_tensor * src0, const struct ggml_tensor * src1, const struct ggml_tensor * dst) {
|
||||
if ((is_type_supported(src0->type)) &&
|
||||
src1->type == GGML_TYPE_F32 &&
|
||||
dst->type == GGML_TYPE_F32 &&
|
||||
src0->backend == GGML_BACKEND_TYPE_CPU) {
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
size_t ggml_bitnet_mul_mat_get_wsize(const struct ggml_tensor * src0, const struct ggml_tensor * src1, const struct ggml_tensor * dst) {
|
||||
const size_t ne01 = src0->ne[1];
|
||||
const size_t ne10 = src1->ne[0];
|
||||
const size_t ne11 = src1->ne[1];
|
||||
|
||||
size_t wsize = ne10 * ne11 * 11 * sizeof(int8_t) + 2 * ne11 * 2 * sizeof(bitnet_float_type);
|
||||
if (sizeof(bitnet_float_type) == 2) {
|
||||
// Need fp32 to fp16 conversion
|
||||
wsize += std::max(ne10, ne01) * ne11 * sizeof(bitnet_float_type);
|
||||
}
|
||||
wsize = ((wsize - 1) / 64 + 1) * 64;
|
||||
return wsize;
|
||||
}
|
||||
|
||||
int ggml_bitnet_get_type_bits(enum ggml_type type) {
|
||||
switch (type) {
|
||||
case GGML_TYPE_TL2:
|
||||
return 2;
|
||||
case GGML_TYPE_Q4_0:
|
||||
return 4;
|
||||
default:
|
||||
return 0;
|
||||
}
|
||||
}
|
||||
#endif
|
||||
@@ -0,0 +1,361 @@
|
||||
#include <vector>
|
||||
#include <type_traits>
|
||||
|
||||
#include "ggml-bitnet.h"
|
||||
#include "ggml-quants.h"
|
||||
#include <cmath>
|
||||
#include <cstring>
|
||||
|
||||
#define QK_I2_S 128
|
||||
#define QK_I2 128
|
||||
|
||||
#if defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__) || defined(__SSSE3__)
|
||||
#include <immintrin.h>
|
||||
// horizontally add 8 int32_t
|
||||
static inline int hsum_i32_8(const __m256i a) {
|
||||
const __m128i sum128 = _mm_add_epi32(_mm256_castsi256_si128(a), _mm256_extractf128_si256(a, 1));
|
||||
const __m128i hi64 = _mm_unpackhi_epi64(sum128, sum128);
|
||||
const __m128i sum64 = _mm_add_epi32(hi64, sum128);
|
||||
const __m128i hi32 = _mm_shuffle_epi32(sum64, _MM_SHUFFLE(2, 3, 0, 1));
|
||||
return _mm_cvtsi128_si32(_mm_add_epi32(sum64, hi32));
|
||||
}
|
||||
#elif defined(__loongarch_asx)
|
||||
// horizontally add 8 int32_t
|
||||
static inline int hsum_i32_8(const __m256i a) {
|
||||
|
||||
__m256i tmp1 = __lasx_xvpermi_q(a, a, 0x11);
|
||||
__m256i tmp2 = __lasx_xvpermi_q(a, a, 0x00);
|
||||
|
||||
__m128i tmp1_128 = lasx_extracti128_lo(tmp1);
|
||||
__m128i tmp2_128 = lasx_extracti128_lo(tmp2);
|
||||
|
||||
__m128i sum128 = __lsx_vadd_w(tmp1_128, tmp2_128);
|
||||
|
||||
__m128i ev = __lsx_vpickev_w(sum128, sum128);
|
||||
__m128i od = __lsx_vpickod_w(sum128, sum128);
|
||||
__m128i sum64 = __lsx_vadd_w(ev, od);
|
||||
|
||||
int sum64_1, sum64_2;
|
||||
sum64_1 = __lsx_vpickve2gr_w(sum64, 0);
|
||||
sum64_2 = __lsx_vpickve2gr_w(sum64, 1);
|
||||
|
||||
return sum64_1 + sum64_2;
|
||||
}
|
||||
#endif
|
||||
|
||||
size_t quantize_i2_s(const float * src, void * dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) {
|
||||
// 2 bits per weight
|
||||
|
||||
size_t row_size = ggml_row_size(GGML_TYPE_I2_S, n_per_row);
|
||||
|
||||
int n = nrow * n_per_row;
|
||||
|
||||
// f32 -> q8
|
||||
double max = 0;
|
||||
for (int i = 0; i < n; ++i) {
|
||||
max = fmax(max, (double)fabs((double)src[i]));
|
||||
}
|
||||
double i2_scale = max;
|
||||
|
||||
uint8_t* q8 = (uint8_t*)malloc(n * sizeof(uint8_t));
|
||||
for (int i=0; i<n; i++) {
|
||||
if (fabs((double)(src[i])) < 1e-6) {
|
||||
q8[i] = 1;
|
||||
continue;
|
||||
}
|
||||
q8[i] = (double)src[i] * i2_scale > 0 ? 2 : 0;
|
||||
}
|
||||
|
||||
memset(dst, 0, n * sizeof(uint8_t) / 4);
|
||||
|
||||
// q8 -> 0, 1, 2
|
||||
// | | |
|
||||
// -1, 0, 1
|
||||
|
||||
uint8_t* i2_weight = (uint8_t*)dst;
|
||||
for (int i = 0; i < n / QK_I2; i++) {
|
||||
for (int j = 0; j < QK_I2; j++) {
|
||||
int group_idx = j / 32;
|
||||
int group_pos = j % 32;
|
||||
uint8_t temp = (q8[i * QK_I2 + j] << (6 - 2 * group_idx));
|
||||
i2_weight[i * 32 + group_pos] |= temp;
|
||||
}
|
||||
}
|
||||
|
||||
float* scale_ptr = (float*)((char*)i2_weight + n / 4);
|
||||
scale_ptr[0] = i2_scale;
|
||||
|
||||
// 32B for alignment
|
||||
return nrow * row_size / 4 + 32;
|
||||
}
|
||||
|
||||
void ggml_vec_dot_i2_i8_s(int n, float * s, size_t bs, const void * vx, size_t bx, const void * vy, size_t by, int nrc) {
|
||||
const uint8_t * x = (uint8_t *)vx;
|
||||
const int8_t * y = (int8_t *)vy;
|
||||
|
||||
const int nb = n / QK_I2_S;
|
||||
const int group32_num = nb / 32;
|
||||
const int la_num = nb % 32;
|
||||
const int groupla_num = nb % 32 != 0 ? 1 : 0;
|
||||
|
||||
#if defined(__AVX2__)
|
||||
|
||||
__m256i mask = _mm256_set1_epi8(0x03);
|
||||
__m256i accu = _mm256_setzero_si256();
|
||||
|
||||
for (int i=0; i < group32_num; i++){
|
||||
__m256i accu32 = _mm256_setzero_si256();
|
||||
for (int j=0; j < 32; j++) {
|
||||
// 128 index
|
||||
__m256i xq8_3 = _mm256_loadu_si256((const __m256i*)(x + i * 32 * 32 + j * 32));
|
||||
__m256i xq8_2 = _mm256_srli_epi16(xq8_3, 2);
|
||||
__m256i xq8_1 = _mm256_srli_epi16(xq8_3, 4);
|
||||
__m256i xq8_0 = _mm256_srli_epi16(xq8_3, 6);
|
||||
|
||||
// each 32 index
|
||||
xq8_3 = _mm256_and_si256(xq8_3, mask);
|
||||
xq8_2 = _mm256_and_si256(xq8_2, mask);
|
||||
xq8_1 = _mm256_and_si256(xq8_1, mask);
|
||||
xq8_0 = _mm256_and_si256(xq8_0, mask);
|
||||
|
||||
// each 32 index
|
||||
__m256i yq8_0 = _mm256_loadu_si256((const __m256i*)(y + i * 128 * 32 + j * 128 + 0));
|
||||
__m256i yq8_1 = _mm256_loadu_si256((const __m256i*)(y + i * 128 * 32 + j * 128 + 32));
|
||||
__m256i yq8_2 = _mm256_loadu_si256((const __m256i*)(y + i * 128 * 32 + j * 128 + 64));
|
||||
__m256i yq8_3 = _mm256_loadu_si256((const __m256i*)(y + i * 128 * 32 + j * 128 + 96));
|
||||
|
||||
// 128 index accumulation add
|
||||
// split into 32 accumulation block
|
||||
// each block each 128 index accumulated 4index
|
||||
// each index maximum 256
|
||||
// each block maximum 4 * 256
|
||||
// each block accumulation maximum 127 * 256
|
||||
// each 32 group index (128 index in one group) needs cast to int32
|
||||
xq8_0 = _mm256_maddubs_epi16(xq8_0, yq8_0);
|
||||
xq8_1 = _mm256_maddubs_epi16(xq8_1, yq8_1);
|
||||
xq8_2 = _mm256_maddubs_epi16(xq8_2, yq8_2);
|
||||
xq8_3 = _mm256_maddubs_epi16(xq8_3, yq8_3);
|
||||
|
||||
accu32 = _mm256_add_epi16(accu32, _mm256_add_epi16(xq8_0, xq8_1));
|
||||
accu32 = _mm256_add_epi16(accu32, _mm256_add_epi16(xq8_2, xq8_3));
|
||||
}
|
||||
accu = _mm256_add_epi32(_mm256_madd_epi16(accu32, _mm256_set1_epi16(1)), accu);
|
||||
}
|
||||
|
||||
for (int i = 0; i < groupla_num; i++){
|
||||
__m256i accula = _mm256_setzero_si256();
|
||||
for (int j = 0; j < la_num; j++) {
|
||||
// 128 index
|
||||
__m256i xq8_3 = _mm256_loadu_si256((const __m256i*)(x + group32_num * 32 * 32 + j * 32));
|
||||
__m256i xq8_2 = _mm256_srli_epi16(xq8_3, 2);
|
||||
__m256i xq8_1 = _mm256_srli_epi16(xq8_3, 4);
|
||||
__m256i xq8_0 = _mm256_srli_epi16(xq8_3, 6);
|
||||
|
||||
// each 32 index
|
||||
xq8_3 = _mm256_and_si256(xq8_3, mask);
|
||||
xq8_2 = _mm256_and_si256(xq8_2, mask);
|
||||
xq8_1 = _mm256_and_si256(xq8_1, mask);
|
||||
xq8_0 = _mm256_and_si256(xq8_0, mask);
|
||||
|
||||
// each 32 index
|
||||
__m256i yq8_0 = _mm256_loadu_si256((const __m256i*)(y + group32_num * 128 * 32 + j * 128 + 0));
|
||||
__m256i yq8_1 = _mm256_loadu_si256((const __m256i*)(y + group32_num * 128 * 32 + j * 128 + 32));
|
||||
__m256i yq8_2 = _mm256_loadu_si256((const __m256i*)(y + group32_num * 128 * 32 + j * 128 + 64));
|
||||
__m256i yq8_3 = _mm256_loadu_si256((const __m256i*)(y + group32_num * 128 * 32 + j * 128 + 96));
|
||||
|
||||
// 128 index accumulation add
|
||||
// split into 32 accumulation block
|
||||
// each block each 128 index accumulated 4index
|
||||
// each index maximum 256
|
||||
// each block maximum 4 * 256
|
||||
// each block accumulation maximum 127 * 256
|
||||
// each 32 group index (128 index in one group) needs cast to int32
|
||||
xq8_0 = _mm256_maddubs_epi16(xq8_0, yq8_0);
|
||||
xq8_1 = _mm256_maddubs_epi16(xq8_1, yq8_1);
|
||||
xq8_2 = _mm256_maddubs_epi16(xq8_2, yq8_2);
|
||||
xq8_3 = _mm256_maddubs_epi16(xq8_3, yq8_3);
|
||||
|
||||
accula = _mm256_add_epi16(accula, _mm256_add_epi16(xq8_0, xq8_1));
|
||||
accula = _mm256_add_epi16(accula, _mm256_add_epi16(xq8_2, xq8_3));
|
||||
}
|
||||
accu = _mm256_add_epi32(accu, _mm256_madd_epi16(accula, _mm256_set1_epi16(1)));
|
||||
}
|
||||
int sumi = hsum_i32_8(accu);
|
||||
*s = (float)sumi;
|
||||
|
||||
#elif defined(__ARM_NEON)
|
||||
|
||||
int32x4_t accu_0 = vdupq_n_s32(0);
|
||||
int32x4_t accu_1 = vdupq_n_s32(0);
|
||||
int32x4_t accu_2 = vdupq_n_s32(0);
|
||||
int32x4_t accu_3 = vdupq_n_s32(0);
|
||||
const uint8x16_t mask = vdupq_n_u8(3);
|
||||
|
||||
for (int i=0; i < group32_num; i++) {
|
||||
|
||||
#if defined(__ARM_FEATURE_DOTPROD)
|
||||
|
||||
#else
|
||||
int16x8_t accu32_0 = vdupq_n_s16(0);
|
||||
int16x8_t accu32_1 = vdupq_n_s16(0);
|
||||
int16x8_t accu32_2 = vdupq_n_s16(0);
|
||||
int16x8_t accu32_3 = vdupq_n_s16(0);
|
||||
#endif
|
||||
|
||||
for (int j=0; j < 32; j++) {
|
||||
uint8x16_t xq8_6 = vld1q_u8(x + i * 32 * 32 + j * 32);
|
||||
uint8x16_t xq8_7 = vld1q_u8(x + i * 32 * 32 + j * 32 + 16);
|
||||
uint8x16_t xq8_4 = vshrq_n_u8(xq8_6, 2);
|
||||
uint8x16_t xq8_5 = vshrq_n_u8(xq8_7, 2);
|
||||
uint8x16_t xq8_2 = vshrq_n_u8(xq8_6, 4);
|
||||
uint8x16_t xq8_3 = vshrq_n_u8(xq8_7, 4);
|
||||
uint8x16_t xq8_0 = vshrq_n_u8(xq8_6, 6);
|
||||
uint8x16_t xq8_1 = vshrq_n_u8(xq8_7, 6);
|
||||
|
||||
int8x16_t q8_0 = vreinterpretq_s8_u8(vandq_u8(xq8_0, mask));
|
||||
int8x16_t q8_1 = vreinterpretq_s8_u8(vandq_u8(xq8_1, mask));
|
||||
int8x16_t q8_2 = vreinterpretq_s8_u8(vandq_u8(xq8_2, mask));
|
||||
int8x16_t q8_3 = vreinterpretq_s8_u8(vandq_u8(xq8_3, mask));
|
||||
int8x16_t q8_4 = vreinterpretq_s8_u8(vandq_u8(xq8_4, mask));
|
||||
int8x16_t q8_5 = vreinterpretq_s8_u8(vandq_u8(xq8_5, mask));
|
||||
int8x16_t q8_6 = vreinterpretq_s8_u8(vandq_u8(xq8_6, mask));
|
||||
int8x16_t q8_7 = vreinterpretq_s8_u8(vandq_u8(xq8_7, mask));
|
||||
|
||||
const int8x16_t yq8_0 = vld1q_s8(y + i * 128 * 32 + j * 128 + 0);
|
||||
const int8x16_t yq8_1 = vld1q_s8(y + i * 128 * 32 + j * 128 + 16);
|
||||
const int8x16_t yq8_2 = vld1q_s8(y + i * 128 * 32 + j * 128 + 32);
|
||||
const int8x16_t yq8_3 = vld1q_s8(y + i * 128 * 32 + j * 128 + 48);
|
||||
const int8x16_t yq8_4 = vld1q_s8(y + i * 128 * 32 + j * 128 + 64);
|
||||
const int8x16_t yq8_5 = vld1q_s8(y + i * 128 * 32 + j * 128 + 80);
|
||||
const int8x16_t yq8_6 = vld1q_s8(y + i * 128 * 32 + j * 128 + 96);
|
||||
const int8x16_t yq8_7 = vld1q_s8(y + i * 128 * 32 + j * 128 + 112);
|
||||
|
||||
#if defined(__ARM_FEATURE_DOTPROD)
|
||||
accu_0 = vdotq_s32(accu_0, q8_0, yq8_0);
|
||||
accu_1 = vdotq_s32(accu_1, q8_1, yq8_1);
|
||||
accu_2 = vdotq_s32(accu_2, q8_2, yq8_2);
|
||||
accu_3 = vdotq_s32(accu_3, q8_3, yq8_3);
|
||||
accu_0 = vdotq_s32(accu_0, q8_4, yq8_4);
|
||||
accu_1 = vdotq_s32(accu_1, q8_5, yq8_5);
|
||||
accu_2 = vdotq_s32(accu_2, q8_6, yq8_6);
|
||||
accu_3 = vdotq_s32(accu_3, q8_7, yq8_7);
|
||||
#else
|
||||
accu32_0 = vmlal_s8(accu32_0, vget_low_s8(q8_0), vget_low_s8(yq8_0));
|
||||
accu32_1 = vmlal_s8(accu32_1, vget_high_s8(q8_0), vget_high_s8(yq8_0));
|
||||
accu32_2 = vmlal_s8(accu32_2, vget_low_s8(q8_1), vget_low_s8(yq8_1));
|
||||
accu32_3 = vmlal_s8(accu32_3, vget_high_s8(q8_1), vget_high_s8(yq8_1));
|
||||
accu32_0 = vmlal_s8(accu32_0, vget_low_s8(q8_2), vget_low_s8(yq8_2));
|
||||
accu32_1 = vmlal_s8(accu32_1, vget_high_s8(q8_2), vget_high_s8(yq8_2));
|
||||
accu32_2 = vmlal_s8(accu32_2, vget_low_s8(q8_3), vget_low_s8(yq8_3));
|
||||
accu32_3 = vmlal_s8(accu32_3, vget_high_s8(q8_3), vget_high_s8(yq8_3));
|
||||
accu32_0 = vmlal_s8(accu32_0, vget_low_s8(q8_4), vget_low_s8(yq8_4));
|
||||
accu32_1 = vmlal_s8(accu32_1, vget_high_s8(q8_4), vget_high_s8(yq8_4));
|
||||
accu32_2 = vmlal_s8(accu32_2, vget_low_s8(q8_5), vget_low_s8(yq8_5));
|
||||
accu32_3 = vmlal_s8(accu32_3, vget_high_s8(q8_5), vget_high_s8(yq8_5));
|
||||
accu32_0 = vmlal_s8(accu32_0, vget_low_s8(q8_6), vget_low_s8(yq8_6));
|
||||
accu32_1 = vmlal_s8(accu32_1, vget_high_s8(q8_6), vget_high_s8(yq8_6));
|
||||
accu32_2 = vmlal_s8(accu32_2, vget_low_s8(q8_7), vget_low_s8(yq8_7));
|
||||
accu32_3 = vmlal_s8(accu32_3, vget_high_s8(q8_7), vget_high_s8(yq8_7));
|
||||
#endif
|
||||
}
|
||||
|
||||
#if defined(__ARM_FEATURE_DOTPROD)
|
||||
|
||||
#else
|
||||
accu_0 = vaddq_s32(accu_0, vmovl_s16(vget_low_s16(accu32_0)));
|
||||
accu_0 = vaddq_s32(accu_0, vmovl_high_s16(accu32_0));
|
||||
accu_1 = vaddq_s32(accu_1, vmovl_s16(vget_low_s16(accu32_1)));
|
||||
accu_1 = vaddq_s32(accu_1, vmovl_high_s16(accu32_1));
|
||||
accu_2 = vaddq_s32(accu_2, vmovl_s16(vget_low_s16(accu32_2)));
|
||||
accu_2 = vaddq_s32(accu_2, vmovl_high_s16(accu32_2));
|
||||
accu_3 = vaddq_s32(accu_3, vmovl_s16(vget_low_s16(accu32_3)));
|
||||
accu_3 = vaddq_s32(accu_3, vmovl_high_s16(accu32_3));
|
||||
#endif
|
||||
}
|
||||
|
||||
for (int i = 0; i < groupla_num; i++){
|
||||
#if defined(__ARM_FEATURE_DOTPROD)
|
||||
|
||||
#else
|
||||
int16x8_t accula_0 = vdupq_n_s16(0);
|
||||
int16x8_t accula_1 = vdupq_n_s16(0);
|
||||
int16x8_t accula_2 = vdupq_n_s16(0);
|
||||
int16x8_t accula_3 = vdupq_n_s16(0);
|
||||
#endif
|
||||
for (int j = 0; j < la_num; j++) {
|
||||
uint8x16_t xq8_6 = vld1q_u8(x + group32_num * 32 * 32 + j * 32);
|
||||
uint8x16_t xq8_7 = vld1q_u8(x + group32_num * 32 * 32 + j * 32 + 16);
|
||||
uint8x16_t xq8_4 = vshrq_n_u8(xq8_6, 2);
|
||||
uint8x16_t xq8_5 = vshrq_n_u8(xq8_7, 2);
|
||||
uint8x16_t xq8_2 = vshrq_n_u8(xq8_6, 4);
|
||||
uint8x16_t xq8_3 = vshrq_n_u8(xq8_7, 4);
|
||||
uint8x16_t xq8_0 = vshrq_n_u8(xq8_6, 6);
|
||||
uint8x16_t xq8_1 = vshrq_n_u8(xq8_7, 6);
|
||||
|
||||
int8x16_t q8_0 = vreinterpretq_s8_u8(vandq_u8(xq8_0, mask));
|
||||
int8x16_t q8_1 = vreinterpretq_s8_u8(vandq_u8(xq8_1, mask));
|
||||
int8x16_t q8_2 = vreinterpretq_s8_u8(vandq_u8(xq8_2, mask));
|
||||
int8x16_t q8_3 = vreinterpretq_s8_u8(vandq_u8(xq8_3, mask));
|
||||
int8x16_t q8_4 = vreinterpretq_s8_u8(vandq_u8(xq8_4, mask));
|
||||
int8x16_t q8_5 = vreinterpretq_s8_u8(vandq_u8(xq8_5, mask));
|
||||
int8x16_t q8_6 = vreinterpretq_s8_u8(vandq_u8(xq8_6, mask));
|
||||
int8x16_t q8_7 = vreinterpretq_s8_u8(vandq_u8(xq8_7, mask));
|
||||
|
||||
const int8x16_t yq8_0 = vld1q_s8(y + group32_num * 128 * 32 + j * 128 + 0);
|
||||
const int8x16_t yq8_1 = vld1q_s8(y + group32_num * 128 * 32 + j * 128 + 16);
|
||||
const int8x16_t yq8_2 = vld1q_s8(y + group32_num * 128 * 32 + j * 128 + 32);
|
||||
const int8x16_t yq8_3 = vld1q_s8(y + group32_num * 128 * 32 + j * 128 + 48);
|
||||
const int8x16_t yq8_4 = vld1q_s8(y + group32_num * 128 * 32 + j * 128 + 64);
|
||||
const int8x16_t yq8_5 = vld1q_s8(y + group32_num * 128 * 32 + j * 128 + 80);
|
||||
const int8x16_t yq8_6 = vld1q_s8(y + group32_num * 128 * 32 + j * 128 + 96);
|
||||
const int8x16_t yq8_7 = vld1q_s8(y + group32_num * 128 * 32 + j * 128 + 112);
|
||||
|
||||
#if defined(__ARM_FEATURE_DOTPROD)
|
||||
accu_0 = vdotq_s32(accu_0, q8_0, yq8_0);
|
||||
accu_1 = vdotq_s32(accu_1, q8_1, yq8_1);
|
||||
accu_2 = vdotq_s32(accu_2, q8_2, yq8_2);
|
||||
accu_3 = vdotq_s32(accu_3, q8_3, yq8_3);
|
||||
accu_0 = vdotq_s32(accu_0, q8_4, yq8_4);
|
||||
accu_1 = vdotq_s32(accu_1, q8_5, yq8_5);
|
||||
accu_2 = vdotq_s32(accu_2, q8_6, yq8_6);
|
||||
accu_3 = vdotq_s32(accu_3, q8_7, yq8_7);
|
||||
#else
|
||||
accula_0 = vmlal_s8(accula_0, vget_low_s8(q8_0), vget_low_s8(yq8_0));
|
||||
accula_1 = vmlal_s8(accula_1, vget_high_s8(q8_0), vget_high_s8(yq8_0));
|
||||
accula_2 = vmlal_s8(accula_2, vget_low_s8(q8_1), vget_low_s8(yq8_1));
|
||||
accula_3 = vmlal_s8(accula_3, vget_high_s8(q8_1), vget_high_s8(yq8_1));
|
||||
accula_0 = vmlal_s8(accula_0, vget_low_s8(q8_2), vget_low_s8(yq8_2));
|
||||
accula_1 = vmlal_s8(accula_1, vget_high_s8(q8_2), vget_high_s8(yq8_2));
|
||||
accula_2 = vmlal_s8(accula_2, vget_low_s8(q8_3), vget_low_s8(yq8_3));
|
||||
accula_3 = vmlal_s8(accula_3, vget_high_s8(q8_3), vget_high_s8(yq8_3));
|
||||
accula_0 = vmlal_s8(accula_0, vget_low_s8(q8_4), vget_low_s8(yq8_4));
|
||||
accula_1 = vmlal_s8(accula_1, vget_high_s8(q8_4), vget_high_s8(yq8_4));
|
||||
accula_2 = vmlal_s8(accula_2, vget_low_s8(q8_5), vget_low_s8(yq8_5));
|
||||
accula_3 = vmlal_s8(accula_3, vget_high_s8(q8_5), vget_high_s8(yq8_5));
|
||||
accula_0 = vmlal_s8(accula_0, vget_low_s8(q8_6), vget_low_s8(yq8_6));
|
||||
accula_1 = vmlal_s8(accula_1, vget_high_s8(q8_6), vget_high_s8(yq8_6));
|
||||
accula_2 = vmlal_s8(accula_2, vget_low_s8(q8_7), vget_low_s8(yq8_7));
|
||||
accula_3 = vmlal_s8(accula_3, vget_high_s8(q8_7), vget_high_s8(yq8_7));
|
||||
#endif
|
||||
}
|
||||
#if defined(__ARM_FEATURE_DOTPROD)
|
||||
|
||||
#else
|
||||
accu_0 = vaddq_s32(accu_0, vmovl_s16(vget_low_s16(accula_0)));
|
||||
accu_0 = vaddq_s32(accu_0, vmovl_high_s16(accula_0));
|
||||
accu_1 = vaddq_s32(accu_1, vmovl_s16(vget_low_s16(accula_1)));
|
||||
accu_1 = vaddq_s32(accu_1, vmovl_high_s16(accula_1));
|
||||
accu_2 = vaddq_s32(accu_2, vmovl_s16(vget_low_s16(accula_2)));
|
||||
accu_2 = vaddq_s32(accu_2, vmovl_high_s16(accula_2));
|
||||
accu_3 = vaddq_s32(accu_3, vmovl_s16(vget_low_s16(accula_3)));
|
||||
accu_3 = vaddq_s32(accu_3, vmovl_high_s16(accula_3));
|
||||
#endif
|
||||
}
|
||||
accu_0 = vaddq_s32(accu_0, accu_1);
|
||||
accu_2 = vaddq_s32(accu_2, accu_3);
|
||||
accu_0 = vaddq_s32(accu_0, accu_2);
|
||||
int sumi = vaddlvq_s32(accu_0);
|
||||
*s = (float)sumi;
|
||||
|
||||
#endif
|
||||
}
|
||||
@@ -0,0 +1,442 @@
|
||||
import argparse
|
||||
import os
|
||||
from configparser import ConfigParser
|
||||
|
||||
def gen_ctor_code():
|
||||
kernel_code = "\n\
|
||||
#include \"ggml-bitnet.h\"\n\
|
||||
#define GGML_BITNET_MAX_NODES 8192\n\
|
||||
static bool initialized = false;\n\
|
||||
static bitnet_tensor_extra * bitnet_tensor_extras = nullptr;\n\
|
||||
static size_t bitnet_tensor_extras_index = 0;\n\
|
||||
static void * aligned_malloc(size_t size) {{\n\
|
||||
#if defined(_WIN32)\n\
|
||||
return _aligned_malloc(size, 64);\n\
|
||||
#else\n\
|
||||
void * ptr = nullptr;\n\
|
||||
posix_memalign(&ptr, 64, size);\n\
|
||||
return ptr;\n\
|
||||
#endif\n\
|
||||
}}\n\
|
||||
static void aligned_free(void * ptr) {{\n\
|
||||
#if defined(_WIN32)\n\
|
||||
_aligned_free(ptr);\n\
|
||||
#else\n\
|
||||
free(ptr);\n\
|
||||
#endif\n\
|
||||
}}\n\
|
||||
\n\
|
||||
void per_tensor_quant(int k, void* lut_scales_, void* b_) {{\n\
|
||||
bitnet_float_type* lut_scales = (bitnet_float_type*)lut_scales_;\n\
|
||||
bitnet_float_type* b = (bitnet_float_type*)b_;\n\
|
||||
#ifdef __ARM_NEON\n\
|
||||
float32x4_t temp_max = vdupq_n_f32(0);\n\
|
||||
for (int i=0; i < k / 4; i++) {{\n\
|
||||
float32x4_t vec_bs = vld1q_f32(b + 4 * i);\n\
|
||||
float32x4_t abssum = vabsq_f32(vec_bs);\n\
|
||||
temp_max = vmaxq_f32(abssum, temp_max);\n\
|
||||
}}\n\
|
||||
float32_t scales = 127 / vmaxvq_f32(temp_max);\n\
|
||||
*lut_scales = scales;\n\
|
||||
#elif defined __AVX2__\n\
|
||||
__m256 max_vec = _mm256_set1_ps(0.f);\n\
|
||||
const __m256 vec_sign = _mm256_set1_ps(-0.0f);\n\
|
||||
// #pragma unroll\n\
|
||||
for (int i = 0; i < k / 8; i++) {{\n\
|
||||
__m256 vec_b = _mm256_loadu_ps(b + i * 8);\n\
|
||||
__m256 vec_babs = _mm256_andnot_ps(vec_sign, vec_b);\n\
|
||||
max_vec = _mm256_max_ps(vec_babs, max_vec);\n\
|
||||
}}\n\
|
||||
__m128 max1 = _mm_max_ps(_mm256_extractf128_ps(max_vec, 1), _mm256_castps256_ps128(max_vec));\n\
|
||||
max1 = _mm_max_ps(max1, _mm_movehl_ps(max1, max1));\n\
|
||||
max1 = _mm_max_ss(max1, _mm_movehdup_ps(max1));\n\
|
||||
float scales = 127 / _mm_cvtss_f32(max1);\n\
|
||||
*lut_scales = scales;\n\
|
||||
#endif\n\
|
||||
}}\n\
|
||||
\n\
|
||||
void partial_max_reset(void* lut_scales_) {{\n\
|
||||
bitnet_float_type* lut_scales = (bitnet_float_type*)lut_scales_;\n\
|
||||
*lut_scales = 0.0;\n\
|
||||
}}\n\
|
||||
\n\
|
||||
#ifdef __ARM_NEON\n\
|
||||
inline void Transpose_8_8(\n\
|
||||
int16x8_t *v0,\n\
|
||||
int16x8_t *v1,\n\
|
||||
int16x8_t *v2,\n\
|
||||
int16x8_t *v3,\n\
|
||||
int16x8_t *v4,\n\
|
||||
int16x8_t *v5,\n\
|
||||
int16x8_t *v6,\n\
|
||||
int16x8_t *v7)\n\
|
||||
{{\n\
|
||||
int16x8x2_t q04 = vzipq_s16(*v0, *v4);\n\
|
||||
int16x8x2_t q15 = vzipq_s16(*v1, *v5);\n\
|
||||
int16x8x2_t q26 = vzipq_s16(*v2, *v6);\n\
|
||||
int16x8x2_t q37 = vzipq_s16(*v3, *v7);\n\
|
||||
\n\
|
||||
int16x8x2_t q0246_0 = vzipq_s16(q04.val[0], q26.val[0]);\n\
|
||||
int16x8x2_t q0246_1 = vzipq_s16(q04.val[1], q26.val[1]);\n\
|
||||
int16x8x2_t q1357_0 = vzipq_s16(q15.val[0], q37.val[0]);\n\
|
||||
int16x8x2_t q1357_1 = vzipq_s16(q15.val[1], q37.val[1]);\n\
|
||||
\n\
|
||||
int16x8x2_t q_fin_0 = vzipq_s16(q0246_0.val[0], q1357_0.val[0]);\n\
|
||||
int16x8x2_t q_fin_1 = vzipq_s16(q0246_0.val[1], q1357_0.val[1]);\n\
|
||||
int16x8x2_t q_fin_2 = vzipq_s16(q0246_1.val[0], q1357_1.val[0]);\n\
|
||||
int16x8x2_t q_fin_3 = vzipq_s16(q0246_1.val[1], q1357_1.val[1]);\n\
|
||||
\n\
|
||||
*v0 = q_fin_0.val[0];\n\
|
||||
*v1 = q_fin_0.val[1];\n\
|
||||
*v2 = q_fin_1.val[0];\n\
|
||||
*v3 = q_fin_1.val[1];\n\
|
||||
*v4 = q_fin_2.val[0];\n\
|
||||
*v5 = q_fin_2.val[1];\n\
|
||||
*v6 = q_fin_3.val[0];\n\
|
||||
*v7 = q_fin_3.val[1];\n\
|
||||
}}\n\
|
||||
#endif\n\
|
||||
\n\
|
||||
template<int act_k>\n\
|
||||
inline void lut_ctor(int8_t* qlut, bitnet_float_type* b, bitnet_float_type* lut_scales) {{\n\
|
||||
#ifdef __ARM_NEON\n\
|
||||
int16x8_t vec_lut[16];\n\
|
||||
float32_t scales = *lut_scales;\n\
|
||||
uint8_t tbl_mask[16];\n\
|
||||
tbl_mask[0] = 0;\n\
|
||||
tbl_mask[1] = 2;\n\
|
||||
tbl_mask[2] = 4;\n\
|
||||
tbl_mask[3] = 6;\n\
|
||||
tbl_mask[4] = 8;\n\
|
||||
tbl_mask[5] = 10;\n\
|
||||
tbl_mask[6] = 12;\n\
|
||||
tbl_mask[7] = 14;\n\
|
||||
tbl_mask[8] = 1;\n\
|
||||
tbl_mask[9] = 3;\n\
|
||||
tbl_mask[10] = 5;\n\
|
||||
tbl_mask[11] = 7;\n\
|
||||
tbl_mask[12] = 9;\n\
|
||||
tbl_mask[13] = 11;\n\
|
||||
tbl_mask[14] = 13;\n\
|
||||
tbl_mask[15] = 15;\n\
|
||||
uint8x16_t tbl_mask_q = vld1q_u8(tbl_mask);\n\
|
||||
#pragma unroll\n\
|
||||
for (int k = 0; k < act_k / 16; ++k) {{\n\
|
||||
float32x4x2_t vec_bs_x0 = vld2q_f32(b + k * 16);\n\
|
||||
float32x4x2_t vec_bs_x1 = vld2q_f32(b + k * 16 + 8);\n\
|
||||
float32x4_t vec_f_0 = vmulq_n_f32(vec_bs_x0.val[0], scales);\n\
|
||||
float32x4_t vec_f_1 = vmulq_n_f32(vec_bs_x0.val[1], scales);\n\
|
||||
float32x4_t vec_f_2 = vmulq_n_f32(vec_bs_x1.val[0], scales);\n\
|
||||
float32x4_t vec_f_3 = vmulq_n_f32(vec_bs_x1.val[1], scales);\n\
|
||||
int32x4_t vec_b_0 = vcvtnq_s32_f32(vec_f_0);\n\
|
||||
int32x4_t vec_b_1 = vcvtnq_s32_f32(vec_f_1);\n\
|
||||
int32x4_t vec_b_2 = vcvtnq_s32_f32(vec_f_2);\n\
|
||||
int32x4_t vec_b_3 = vcvtnq_s32_f32(vec_f_3);\n\
|
||||
int16x4_t vec_b16_0 = vmovn_s32(vec_b_0);\n\
|
||||
int16x4_t vec_b16_1 = vmovn_s32(vec_b_1);\n\
|
||||
int16x4_t vec_b16_2 = vmovn_s32(vec_b_2);\n\
|
||||
int16x4_t vec_b16_3 = vmovn_s32(vec_b_3);\n\
|
||||
int16x8_t vec_bs_0 = vcombine_s16(vec_b16_0, vec_b16_2);\n\
|
||||
int16x8_t vec_bs_1 = vcombine_s16(vec_b16_1, vec_b16_3);\n\
|
||||
vec_lut[0] = vdupq_n_s16(0);\n\
|
||||
vec_lut[0] = vec_lut[0] - vec_bs_0;\n\
|
||||
vec_lut[0] = vec_lut[0] - vec_bs_1;\n\
|
||||
vec_lut[1] = vdupq_n_s16(0);\n\
|
||||
vec_lut[1] = vec_lut[1] - vec_bs_0;\n\
|
||||
vec_lut[2] = vdupq_n_s16(0);\n\
|
||||
vec_lut[2] = vec_lut[2] - vec_bs_0;\n\
|
||||
vec_lut[2] = vec_lut[2] + vec_bs_1;\n\
|
||||
vec_lut[3] = vdupq_n_s16(0);\n\
|
||||
vec_lut[3] = vec_lut[3] - vec_bs_1;\n\
|
||||
vec_lut[4] = vdupq_n_s16(0);\n\
|
||||
vec_lut[5] = vec_bs_1;\n\
|
||||
vec_lut[6] = vec_bs_0;\n\
|
||||
vec_lut[6] = vec_lut[6] - vec_bs_1;\n\
|
||||
vec_lut[7] = vec_bs_0;\n\
|
||||
vec_lut[8] = vec_bs_0;\n\
|
||||
vec_lut[8] = vec_lut[8] + vec_bs_1;\n\
|
||||
Transpose_8_8(&(vec_lut[0]), &(vec_lut[1]), &(vec_lut[2]), &(vec_lut[3]),\n\
|
||||
&(vec_lut[4]), &(vec_lut[5]), &(vec_lut[6]), &(vec_lut[7]));\n\
|
||||
Transpose_8_8(&(vec_lut[8]), &(vec_lut[9]), &(vec_lut[10]), &(vec_lut[11]),\n\
|
||||
&(vec_lut[12]), &(vec_lut[13]), &(vec_lut[14]), &(vec_lut[15]));\n\
|
||||
#pragma unroll\n\
|
||||
for (int idx = 0; idx < 8; idx++) {{\n\
|
||||
int8x16_t q0_s = vqtbl1q_s8(vreinterpretq_s8_s16(vec_lut[idx]), tbl_mask_q);\n\
|
||||
int8x8_t q0_low = vget_low_s8(q0_s);\n\
|
||||
int8x8_t q0_high = vget_high_s8(q0_s);\n\
|
||||
int8x16_t q1_s = vqtbl1q_s8(vreinterpretq_s8_s16(vec_lut[idx + 8]), tbl_mask_q);\n\
|
||||
int8x8_t q1_low = vget_low_s8(q1_s);\n\
|
||||
int8x8_t q1_high = vget_high_s8(q1_s);\n\
|
||||
vst1_s8(qlut + k * 16 * 8 * 2 + idx * 16 * 2, q0_high);\n\
|
||||
vst1_s8(qlut + k * 16 * 8 * 2 + idx * 16 * 2 + 8, q1_high);\n\
|
||||
vst1_s8(qlut + k * 16 * 8 * 2 + idx * 16 * 2 + 16, q0_low);\n\
|
||||
vst1_s8(qlut + k * 16 * 8 * 2 + idx * 16 * 2 + 24, q1_low);\n\
|
||||
}}\n\
|
||||
}}\n\
|
||||
#endif\n\
|
||||
}}\n\
|
||||
\n\
|
||||
static bool is_type_supported(enum ggml_type type) {{\n\
|
||||
if (type == GGML_TYPE_Q4_0 ||\n\
|
||||
type == GGML_TYPE_TL1) {{\n\
|
||||
return true;\n\
|
||||
}} else {{\n\
|
||||
return false;\n\
|
||||
}}\n\
|
||||
}}\n\
|
||||
"
|
||||
return kernel_code
|
||||
|
||||
def gen_body_core_code(bm, by):
|
||||
length = 4
|
||||
all_code = ""
|
||||
for i in range(length):
|
||||
core_code = "\n\
|
||||
uint8x16_t vec_a_{0} = vld1q_u8(a + i * KK / 2 + k * 32 * 2 + {0} * 16);\n\
|
||||
uint8x16_t vec_a{0}_top = vshrq_n_u8(vec_a_{0}, 4);\n\
|
||||
uint8x16_t vec_a{0}_bot = vandq_u8(vec_a_{0}, vec_mask);\n\
|
||||
int8x16_t vec_v_{0}_left_tmp0 = vqtbl1q_s8(vec_lut[{1} * k + {2}], vec_a{0}_top);\n\
|
||||
int8x16_t vec_v_{0}_left_tmp1 = vqtbl1q_s8(vec_lut[{1} * k + {3}], vec_a{0}_top);\n\
|
||||
int8x16_t vec_v_{0}_right_tmp0 = vqtbl1q_s8(vec_lut[{1} * k + {4}], vec_a{0}_bot);\n\
|
||||
int8x16_t vec_v_{0}_right_tmp1 = vqtbl1q_s8(vec_lut[{1} * k + {5}], vec_a{0}_bot);\n\
|
||||
int8x16x2_t vec_v_left_{0} = vzipq_s8(vec_v_{0}_left_tmp1, vec_v_{0}_left_tmp0);\n\
|
||||
int8x16x2_t vec_v_right_{0} = vzipq_s8(vec_v_{0}_right_tmp1, vec_v_{0}_right_tmp0);\n\
|
||||
vec_c[{6}] += vec_v_left_{0}.val[0];\n\
|
||||
vec_c[{6}] += vec_v_right_{0}.val[0];\n\
|
||||
vec_c[{7}] += vec_v_left_{0}.val[1];\n\
|
||||
vec_c[{7}] += vec_v_right_{0}.val[1];\n\
|
||||
".format(i, 2 * by // 2, (4 * i) % (2 * by // 2), (4 * i + 1) % (2 * by // 2), (4 * i + 2) % (2 * by // 2), (4 * i + 3) % (2 * by // 2), (i * 2) // (by // 2) * 2 + 0, (i * 2) // (by // 2) * 2 + 1)
|
||||
|
||||
all_code = "".join([all_code, core_code])
|
||||
|
||||
all_code = "".join([all_code, "\n }\n\n"])
|
||||
|
||||
for i in range(bm // 8):
|
||||
core_code = "\
|
||||
int32x4_t vec_v_bot_low_low_{0} = vmovl_s16(vget_low_s16(vec_c[{0}]));\n\
|
||||
int32x4_t vec_v_bot_low_high_{0} = vmovl_high_s16(vec_c[{0}]);\n\
|
||||
vst1q_s32(c + i + {1}, vld1q_s32(c + i + {1}) + vec_v_bot_low_low_{0});\n\
|
||||
vst1q_s32(c + i + {2}, vld1q_s32(c + i + {2}) + vec_v_bot_low_high_{0});\n".format(i, i * 8, i * 8 + 4)
|
||||
all_code = "".join([all_code, core_code])
|
||||
|
||||
return all_code
|
||||
|
||||
def gen_tbl_impl(pre, BM, BK, bm, k):
|
||||
|
||||
kernel_code = "\
|
||||
#include <arm_neon.h>\n\
|
||||
\n\
|
||||
#define BM{0} {1}\n\
|
||||
#define BBK{0} {2}\n\
|
||||
inline void tbl_impl_{0}(int32_t* c, int8_t* lut, uint8_t* a) {{\n\
|
||||
#ifdef __ARM_NEON\n\
|
||||
const int KK = BBK{0} / 2;\n\
|
||||
const uint8x16_t vec_mask = vdupq_n_u8(0x0f);\n\
|
||||
const int8x16_t vec_zero = vdupq_n_s16(0x0000);\n\
|
||||
int8x16_t vec_lut[2 * KK];\n\
|
||||
".format(pre, BM, BK)
|
||||
|
||||
kernel_code = "".join([kernel_code, " int16x8_t vec_c[{}];".format(bm // 8)])
|
||||
|
||||
kernel_code = "".join([kernel_code, "\n\
|
||||
#pragma unroll\n\
|
||||
for (int k = 0; k < 2 * KK; k++) {\n\
|
||||
vec_lut[k] = vld1q_s8(lut + k * 16);\n\
|
||||
}\n"])
|
||||
|
||||
pre_core_code = "\n\
|
||||
#pragma unroll\n\
|
||||
for (int i = 0; i < BM{}; i += {}) {{\n\
|
||||
#pragma unroll\n\
|
||||
for (int i=0; i<{}; i++) {{\n\
|
||||
vec_c[i] = vandq_s16(vec_c[i], vec_zero);\n\
|
||||
}}\n".format(pre, bm, bm // 8)
|
||||
|
||||
body_core_pre_code = "\n\
|
||||
#pragma unroll\n\
|
||||
for (int k = 0; k < KK / {}; k++) {{\n\
|
||||
".format(256 // bm // 2)
|
||||
|
||||
body_core_post_code = "\n\
|
||||
}\n\
|
||||
\
|
||||
#endif\n\
|
||||
}\n"
|
||||
|
||||
kernel_code = "".join([kernel_code, pre_core_code, body_core_pre_code, gen_body_core_code(bm, 256 // bm), body_core_post_code])
|
||||
|
||||
kernel_code = "".join([kernel_code, "\n\
|
||||
int32_t qgemm_lut_{0}(void* A, void* LUT, void* Scales, void* LUT_Scales, void* C) {{\n\
|
||||
alignas({1}) uint32_t CBits[BM{0}];\n\
|
||||
memset(&(CBits[0]), 0, BM{0} * sizeof(int32_t));\n\
|
||||
#pragma unroll\n\
|
||||
for (int32_t k_outer = 0; k_outer < {2} / BBK{0}; ++k_outer) {{\n\
|
||||
tbl_impl_{0}((&(((int32_t*)CBits)[0])), (&(((int8_t*)LUT)[(k_outer * BBK{0} / 2 * 32)])), (&(((uint8_t*)A)[(k_outer * BBK{0} / 2 / 2 * BM{0})])));\n\
|
||||
}}\n\
|
||||
#pragma unroll\n\
|
||||
for (int i = 0; i < BM{0}; i++) {{\n\
|
||||
((bitnet_float_type*)C)[i] = (((int32_t*)CBits)[i]) / ((bitnet_float_type*)LUT_Scales)[0] * ((bitnet_float_type*)Scales)[0];\n\
|
||||
}}\n\
|
||||
return 0;\n\
|
||||
}};\n".format(pre, min(32, BK), k)])
|
||||
|
||||
return kernel_code
|
||||
|
||||
def gen_top_api(kernel_shapes):
|
||||
|
||||
kernel_code = "void ggml_preprocessor(int m, int k, void* B, void* LUT_Scales, void* QLUT) {{\n\
|
||||
if (m == {0} && k == {1}) {{\n\
|
||||
preprocessor_k<{1}>(B, LUT_Scales, QLUT);\n\
|
||||
}}\n\
|
||||
".format(kernel_shapes[0][0], kernel_shapes[0][1])
|
||||
for i in range(1, len(kernel_shapes)):
|
||||
kernel_code = "".join([kernel_code, " else if (m == {0} && k == {1}) {{\n\
|
||||
preprocessor_k<{1}>(B, LUT_Scales, QLUT);\n\
|
||||
}}\n".format(kernel_shapes[i][0], kernel_shapes[i][1])])
|
||||
kernel_code = "".join([kernel_code, "}\n"])
|
||||
kernel_code = "".join([kernel_code, "void ggml_qgemm_lut(int m, int k, void* A, void* LUT, void* Scales, void* LUT_Scales, void* C) {{\n\
|
||||
if (m == {0} && k == {1}) {{\n\
|
||||
qgemm_lut_{0}_{1}(A, LUT, Scales, LUT_Scales, C);\n\
|
||||
}}\n\
|
||||
".format(kernel_shapes[0][0], kernel_shapes[0][1])])
|
||||
for i in range(1, len(kernel_shapes)):
|
||||
kernel_code = "".join([kernel_code, " else if (m == {0} && k == {1}) {{\n\
|
||||
qgemm_lut_{0}_{1}(A, LUT, Scales, LUT_Scales, C);\n\
|
||||
}}\n\
|
||||
".format(kernel_shapes[i][0], kernel_shapes[i][1])])
|
||||
kernel_code = "".join([kernel_code, "}\n"])
|
||||
return kernel_code
|
||||
|
||||
def gen_preprocess_code():
|
||||
kernel_code = "\n\
|
||||
template<int K>\n\
|
||||
void preprocessor_k(void* B, void* LUT_Scales, void* QLUT) {{\n\
|
||||
partial_max_reset((&(((bitnet_float_type*)LUT_Scales)[0])));\n\
|
||||
per_tensor_quant(K, (&(((bitnet_float_type*)LUT_Scales)[0])), (&(((bitnet_float_type*)B)[0])));\n\
|
||||
\n\
|
||||
lut_ctor<K>((&(((int8_t*)QLUT)[0])), (&(((bitnet_float_type*)B)[0])), (&(((bitnet_float_type*)LUT_Scales)[0])));\n\
|
||||
}}\n"
|
||||
return kernel_code
|
||||
|
||||
def gen_transform_code(kernel_shape):
|
||||
kernel_code = "\n\
|
||||
void ggml_bitnet_transform_tensor(struct ggml_tensor * tensor) {\n\
|
||||
if (!(is_type_supported(tensor->type) && tensor->backend == GGML_BACKEND_TYPE_CPU && tensor->extra == nullptr)) {\n\
|
||||
return;\n\
|
||||
}\n\
|
||||
\n\
|
||||
int k = tensor->ne[0];\n\
|
||||
int m = tensor->ne[1];\n\
|
||||
const int lut_scales_size = 1;\n\
|
||||
const int scales_size = 1;\n\
|
||||
int bk = 0;\n\
|
||||
int bm = 0;\n"
|
||||
|
||||
kernel_code = "".join([kernel_code, "\n\
|
||||
if (m == {0} && k == {1}) {{\n\
|
||||
bm = BM{0}_{1};\n\
|
||||
bk = BBK{0}_{1};\n\
|
||||
}}\n".format(kernel_shapes[0][0], kernel_shapes[0][1])])
|
||||
|
||||
for i in range(1, len(kernel_shapes)):
|
||||
kernel_code = "".join([kernel_code, "else if (m == {0} && k == {1}) {{\n\
|
||||
bm = BM{0}_{1};\n\
|
||||
bk = BBK{0}_{1};\n\
|
||||
}}\n".format(kernel_shapes[i][0], kernel_shapes[i][1])])
|
||||
|
||||
kernel_code = "".join([kernel_code, "\n\
|
||||
const int n_tile_num = m / bm;\n\
|
||||
const int BK = bk;\n\
|
||||
uint8_t * qweights;\n\
|
||||
bitnet_float_type * scales;\n\
|
||||
\n\
|
||||
scales = (bitnet_float_type *) aligned_malloc(sizeof(bitnet_float_type));\n\
|
||||
qweights = (uint8_t *) tensor->data;\n\
|
||||
float * i2_scales = (float * )(qweights + k * m / 4);\n\
|
||||
scales[0] = (bitnet_float_type) i2_scales[0];\n\
|
||||
\n\
|
||||
tensor->extra = bitnet_tensor_extras + bitnet_tensor_extras_index;\n\
|
||||
bitnet_tensor_extras[bitnet_tensor_extras_index++] = {\n\
|
||||
/* .lut_scales_size = */ lut_scales_size,\n\
|
||||
/* .BK = */ BK,\n\
|
||||
/* .n_tile_num = */ n_tile_num,\n\
|
||||
/* .qweights = */ qweights,\n\
|
||||
/* .scales = */ scales\n\
|
||||
};\n\
|
||||
}\n"])
|
||||
|
||||
return kernel_code
|
||||
|
||||
if __name__ == "__main__":
|
||||
ModelShapeDict = {
|
||||
"bitnet_b1_58-large" : [[1536, 4096],
|
||||
[1536, 1536],
|
||||
[4096, 1536]],
|
||||
"bitnet_b1_58-3B" : [[3200, 8640],
|
||||
[3200, 3200],
|
||||
[8640, 3200]],
|
||||
"Llama3-8B-1.58-100B-tokens" : [[14336, 4096],
|
||||
[4096, 14336],
|
||||
[1024, 4096],
|
||||
[4096, 4096]]
|
||||
}
|
||||
|
||||
parser = argparse.ArgumentParser(description='gen impl')
|
||||
parser.add_argument('--model',default="input", type=str, dest="model",
|
||||
help="choose from bitnet_b1_58-large/bitnet_b1_58-3B/Llama3-8B-1.58-100B-tokens.")
|
||||
parser.add_argument('--BM',default="input", type=str,
|
||||
help="block length when cutting one weight (M, K) into M / BM weights (BM, K).")
|
||||
parser.add_argument('--BK',default="input", type=str,
|
||||
help="block length when cutting one weight (M, K) into K / BK weights (M, BK).")
|
||||
parser.add_argument('--bm',default="input", type=str,
|
||||
help="using simd instructions to compute (bm, 256 / bm) in one block")
|
||||
args = parser.parse_args()
|
||||
|
||||
kernel_shapes = ModelShapeDict[args.model]
|
||||
|
||||
BM_list = [int(item) for item in args.BM.split(',')]
|
||||
BK_list = [int(item) for item in args.BK.split(',')]
|
||||
bm_list = [int(item) for item in args.bm.split(',')]
|
||||
|
||||
assert(len(BM_list) == len(BK_list) == len(bm_list) == len(kernel_shapes)), "number of BM / BK / bm shoud be {}".format(len(kernel_shapes))
|
||||
|
||||
for i in range(len(kernel_shapes)):
|
||||
assert kernel_shapes[i][0] % BM_list[i] == 0, "M %% BM should be 0"
|
||||
assert kernel_shapes[i][1] % BK_list[i] == 0, "K %% BK should be 0"
|
||||
assert bm_list[i] in [32, 64], "choose bm from [32, 64]"
|
||||
|
||||
tbl_impl_code = []
|
||||
|
||||
for i in range(len(kernel_shapes)):
|
||||
tbl_impl_code.append(
|
||||
gen_tbl_impl("{}_{}".format(kernel_shapes[i][0], kernel_shapes[i][1]), BM_list[i], BK_list[i], bm_list[i], kernel_shapes[i][1])
|
||||
)
|
||||
api_code = gen_top_api(kernel_shapes)
|
||||
pre_code = gen_preprocess_code()
|
||||
ctor_code = gen_ctor_code()
|
||||
trans_code = gen_transform_code(kernel_shapes)
|
||||
|
||||
output_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "include")
|
||||
|
||||
with open(''.join([output_dir, "/bitnet-lut-kernels.h"]), 'w') as f:
|
||||
f.write(''.join("#if defined(GGML_BITNET_ARM_TL1)"))
|
||||
f.write(''.join(ctor_code))
|
||||
for code in tbl_impl_code:
|
||||
f.write(''.join(code))
|
||||
f.write(''.join(pre_code))
|
||||
f.write(''.join(api_code))
|
||||
f.write(''.join(trans_code))
|
||||
f.write(''.join("#endif"))
|
||||
|
||||
config = ConfigParser()
|
||||
|
||||
for i in range(len(kernel_shapes)):
|
||||
config.add_section('Kernels_{}'.format(i))
|
||||
config.set('Kernels_{}'.format(i), 'M'.format(i), str(kernel_shapes[i][0]))
|
||||
config.set('Kernels_{}'.format(i), 'K'.format(i), str(kernel_shapes[i][1]))
|
||||
config.set('Kernels_{}'.format(i), 'BM'.format(i), str(BM_list[i]))
|
||||
config.set('Kernels_{}'.format(i), 'BK'.format(i), str(BK_list[i]))
|
||||
config.set('Kernels_{}'.format(i), 'bmm'.format(i), str(bm_list[i]))
|
||||
|
||||
with open(''.join([output_dir, "/kernel_config.ini"]), 'w') as configfile:
|
||||
config.write(configfile)
|
||||
@@ -0,0 +1,757 @@
|
||||
import argparse
|
||||
import os
|
||||
from configparser import ConfigParser
|
||||
|
||||
def gen_ctor_code():
|
||||
kernel_code = "\n\
|
||||
#include \"ggml-bitnet.h\"\n\
|
||||
#include <cstring>\n\
|
||||
#include <immintrin.h>\n\
|
||||
#define GGML_BITNET_MAX_NODES 8192\n\
|
||||
static bool initialized = false;\n\
|
||||
static bitnet_tensor_extra * bitnet_tensor_extras = nullptr;\n\
|
||||
static size_t bitnet_tensor_extras_index = 0;\n\
|
||||
static void * aligned_malloc(size_t size) {\n\
|
||||
#if defined(_WIN32)\n\
|
||||
return _aligned_malloc(size, 64);\n\
|
||||
#else\n\
|
||||
void * ptr = nullptr;\n\
|
||||
posix_memalign(&ptr, 64, size);\n\
|
||||
return ptr;\n\
|
||||
#endif\n\
|
||||
}\n\
|
||||
\n\
|
||||
static void aligned_free(void * ptr) {\n\
|
||||
#if defined(_WIN32)\n\
|
||||
_aligned_free(ptr);\n\
|
||||
#else\n\
|
||||
free(ptr);\n\
|
||||
#endif\n\
|
||||
}\n\
|
||||
#define BK2 32\n\
|
||||
#if defined __AVX2__\n\
|
||||
inline void _mm256_merge_epi32(const __m256i v0, const __m256i v1, __m256i *vl, __m256i *vh)\n\
|
||||
{\n\
|
||||
__m256i va = _mm256_permute4x64_epi64(v0, _MM_SHUFFLE(3, 1, 2, 0));\n\
|
||||
__m256i vb = _mm256_permute4x64_epi64(v1, _MM_SHUFFLE(3, 1, 2, 0));\n\
|
||||
*vl = _mm256_unpacklo_epi32(va, vb);\n\
|
||||
*vh = _mm256_unpackhi_epi32(va, vb);\n\
|
||||
}\n\
|
||||
inline void _mm256_merge_epi64(const __m256i v0, const __m256i v1, __m256i *vl, __m256i *vh)\n\
|
||||
{\n\
|
||||
__m256i va = _mm256_permute4x64_epi64(v0, _MM_SHUFFLE(3, 1, 2, 0));\n\
|
||||
__m256i vb = _mm256_permute4x64_epi64(v1, _MM_SHUFFLE(3, 1, 2, 0));\n\
|
||||
*vl = _mm256_unpacklo_epi64(va, vb);\n\
|
||||
*vh = _mm256_unpackhi_epi64(va, vb);\n\
|
||||
}\n\
|
||||
inline void _mm256_merge_si128(const __m256i v0, const __m256i v1, __m256i *vl, __m256i *vh)\n\
|
||||
{\n\
|
||||
*vl = _mm256_permute2x128_si256(v0, v1, _MM_SHUFFLE(0, 2, 0, 0));\n\
|
||||
*vh = _mm256_permute2x128_si256(v0, v1, _MM_SHUFFLE(0, 3, 0, 1));\n\
|
||||
}\n\
|
||||
inline void Transpose_8_8(\n\
|
||||
__m256i *v0,\n\
|
||||
__m256i *v1,\n\
|
||||
__m256i *v2,\n\
|
||||
__m256i *v3,\n\
|
||||
__m256i *v4,\n\
|
||||
__m256i *v5,\n\
|
||||
__m256i *v6,\n\
|
||||
__m256i *v7)\n\
|
||||
{\n\
|
||||
__m256i w0, w1, w2, w3, w4, w5, w6, w7;\n\
|
||||
__m256i x0, x1, x2, x3, x4, x5, x6, x7;\n\
|
||||
_mm256_merge_epi32(*v0, *v1, &w0, &w1);\n\
|
||||
_mm256_merge_epi32(*v2, *v3, &w2, &w3);\n\
|
||||
_mm256_merge_epi32(*v4, *v5, &w4, &w5);\n\
|
||||
_mm256_merge_epi32(*v6, *v7, &w6, &w7);\n\
|
||||
_mm256_merge_epi64(w0, w2, &x0, &x1);\n\
|
||||
_mm256_merge_epi64(w1, w3, &x2, &x3);\n\
|
||||
_mm256_merge_epi64(w4, w6, &x4, &x5);\n\
|
||||
_mm256_merge_epi64(w5, w7, &x6, &x7);\n\
|
||||
_mm256_merge_si128(x0, x4, v0, v1);\n\
|
||||
_mm256_merge_si128(x1, x5, v2, v3);\n\
|
||||
_mm256_merge_si128(x2, x6, v4, v5);\n\
|
||||
_mm256_merge_si128(x3, x7, v6, v7);\n\
|
||||
}\n\
|
||||
#endif\n\
|
||||
inline int32_t per_tensor_quant(int k, void* lut_scales_, void* b_) {\n\
|
||||
bitnet_float_type* lut_scales = (bitnet_float_type*)lut_scales_;\n\
|
||||
bitnet_float_type* b = (bitnet_float_type*)b_;\n\
|
||||
#if defined __AVX2__\n\
|
||||
__m256 max_vec = _mm256_set1_ps(0.f);\n\
|
||||
const __m256 vec_sign = _mm256_set1_ps(-0.0f);\n\
|
||||
for (int i = 0; i < k / 8; i++) {\n\
|
||||
__m256 vec_b = _mm256_loadu_ps(b + i * 8);\n\
|
||||
__m256 vec_babs = _mm256_andnot_ps(vec_sign, vec_b);\n\
|
||||
max_vec = _mm256_max_ps(vec_babs, max_vec);\n\
|
||||
}\n\
|
||||
__m128 max1 = _mm_max_ps(_mm256_extractf128_ps(max_vec, 1), _mm256_castps256_ps128(max_vec));\n\
|
||||
max1 = _mm_max_ps(max1, _mm_movehl_ps(max1, max1));\n\
|
||||
max1 = _mm_max_ss(max1, _mm_movehdup_ps(max1));\n\
|
||||
float scales = 127 / _mm_cvtss_f32(max1);\n\
|
||||
*lut_scales = scales;\n\
|
||||
#endif\n\
|
||||
return 0;\n\
|
||||
}\n\
|
||||
inline int32_t partial_max_reset(int32_t bs, void* lut_scales_) {\n\
|
||||
bitnet_float_type* lut_scales = (bitnet_float_type*)lut_scales_;\n\
|
||||
#pragma unroll\n\
|
||||
for (int i=0; i< bs; i++) {\n\
|
||||
lut_scales[i] = 0.0;\n\
|
||||
}\n\
|
||||
return 0;\n\
|
||||
}\n\
|
||||
template<int act_k>\n\
|
||||
inline int32_t three_lut_ctor(int8_t* qlut, bitnet_float_type* b, bitnet_float_type* lut_scales) {\n\
|
||||
#if defined __AVX2__\n\
|
||||
__m256 vec_lut[16];\n\
|
||||
const __m256i vec_bi = _mm256_set_epi32(84, 72, 60, 48, 36, 24, 12, 0);\n\
|
||||
float scales = *lut_scales;\n\
|
||||
__m256i shuffle_mask = _mm256_set_epi8(\n\
|
||||
0x0f, 0x0d, 0x0b, 0x09, 0x07, 0x05, 0x03, 0x01,\n\
|
||||
0x0e, 0x0c, 0x0a, 0x08, 0x06, 0x04, 0x02, 0x00,\n\
|
||||
0x0f, 0x0d, 0x0b, 0x09, 0x07, 0x05, 0x03, 0x01,\n\
|
||||
0x0e, 0x0c, 0x0a, 0x08, 0x06, 0x04, 0x02, 0x00\n\
|
||||
);\n\
|
||||
#pragma unroll\n\
|
||||
for (int k = 0; k < act_k / 24; ++k) {\n\
|
||||
__m256 vec_b0 = _mm256_i32gather_ps(b + k * 24 + 0, vec_bi, 1);\n\
|
||||
__m256 vec_b1 = _mm256_i32gather_ps(b + k * 24 + 1, vec_bi, 1);\n\
|
||||
__m256 vec_b2 = _mm256_i32gather_ps(b + k * 24 + 2, vec_bi, 1);\n\
|
||||
\n\
|
||||
__m256i vec_b0i = _mm256_cvtps_epi32(_mm256_round_ps(_mm256_mul_ps(vec_b0, _mm256_set1_ps(scales)), _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC));\n\
|
||||
__m256i vec_b1i = _mm256_cvtps_epi32(_mm256_round_ps(_mm256_mul_ps(vec_b1, _mm256_set1_ps(scales)), _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC));\n\
|
||||
__m256i vec_b2i = _mm256_cvtps_epi32(_mm256_round_ps(_mm256_mul_ps(vec_b2, _mm256_set1_ps(scales)), _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC));\n\
|
||||
\n\
|
||||
vec_lut[15] = _mm256_setzero_si256();\n\
|
||||
vec_lut[14] = _mm256_setzero_si256();\n\
|
||||
vec_lut[13] = vec_b0i;\n\
|
||||
vec_lut[13] = _mm256_add_epi32(vec_lut[13], vec_b1i);\n\
|
||||
vec_lut[13] = _mm256_add_epi32(vec_lut[13], vec_b2i);\n\
|
||||
vec_lut[12] = vec_b0i;\n\
|
||||
vec_lut[12] = _mm256_add_epi32(vec_lut[12], vec_b1i);\n\
|
||||
vec_lut[11] = vec_b0i;\n\
|
||||
vec_lut[11] = _mm256_add_epi32(vec_lut[11], vec_b1i);\n\
|
||||
vec_lut[11] = _mm256_sub_epi32(vec_lut[11], vec_b2i);\n\
|
||||
vec_lut[10] = vec_b0i;\n\
|
||||
vec_lut[10] = _mm256_add_epi32(vec_lut[10], vec_b2i);\n\
|
||||
vec_lut[9] = vec_b0i;\n\
|
||||
vec_lut[8] = vec_b0i;\n\
|
||||
vec_lut[8] = _mm256_sub_epi32(vec_lut[8], vec_b2i);\n\
|
||||
vec_lut[7] = vec_b0i;\n\
|
||||
vec_lut[7] = _mm256_sub_epi32(vec_lut[7], vec_b1i);\n\
|
||||
vec_lut[7] = _mm256_add_epi32(vec_lut[7], vec_b2i);\n\
|
||||
vec_lut[6] = vec_b0i;\n\
|
||||
vec_lut[6] = _mm256_sub_epi32(vec_lut[6], vec_b1i);\n\
|
||||
vec_lut[5] = vec_b0i;\n\
|
||||
vec_lut[5] = _mm256_sub_epi32(vec_lut[5], vec_b1i);\n\
|
||||
vec_lut[5] = _mm256_sub_epi32(vec_lut[5], vec_b2i);\n\
|
||||
vec_lut[4] = vec_b1i;\n\
|
||||
vec_lut[4] = _mm256_add_epi32(vec_lut[4], vec_b2i);\n\
|
||||
vec_lut[3] = vec_b1i;\n\
|
||||
vec_lut[2] = vec_b1i;\n\
|
||||
vec_lut[2] = _mm256_sub_epi32(vec_lut[2], vec_b2i);\n\
|
||||
vec_lut[1] = vec_b2i;\n\
|
||||
vec_lut[0] = _mm256_setzero_si256();\n\
|
||||
__m256i ix[16];\n\
|
||||
\n\
|
||||
#pragma unroll\n\
|
||||
for (int g = 0; g < 16; ++g) {\n\
|
||||
ix[g] = vec_lut[g];\n\
|
||||
}\n\
|
||||
\n\
|
||||
Transpose_8_8(&(ix[0]), &(ix[1]), &(ix[2]), &(ix[3]), &(ix[4]), &(ix[5]),&(ix[6]), &(ix[7]));\n\
|
||||
Transpose_8_8(&(ix[8]), &(ix[9]), &(ix[10]), &(ix[11]), &(ix[12]), &(ix[13]),&(ix[14]), &(ix[15]));\n\
|
||||
\n\
|
||||
#pragma unroll\n\
|
||||
for (int g = 0; g < 8; ++g) {\n\
|
||||
ix[g] = _mm256_packs_epi32(ix[g], ix[g + 8]);\n\
|
||||
ix[g] = _mm256_permute4x64_epi64(ix[g], _MM_SHUFFLE(3, 1, 2, 0));\n\
|
||||
ix[g] = _mm256_shuffle_epi8(ix[g], shuffle_mask);\n\
|
||||
ix[g] = _mm256_permute4x64_epi64(ix[g], _MM_SHUFFLE(3, 1, 2, 0));\n\
|
||||
}\n\
|
||||
int8_t* qlut_i8 = reinterpret_cast<int8_t*>(qlut);\n\
|
||||
_mm256_storeu_si256(reinterpret_cast<__m256i*>(qlut_i8 + k * 256 + 0 * 32 + 0), ix[0]);\n\
|
||||
_mm256_storeu_si256(reinterpret_cast<__m256i*>(qlut_i8 + k * 256 + 1 * 32 + 0), ix[1]);\n\
|
||||
_mm256_storeu_si256(reinterpret_cast<__m256i*>(qlut_i8 + k * 256 + 2 * 32 + 0), ix[2]);\n\
|
||||
_mm256_storeu_si256(reinterpret_cast<__m256i*>(qlut_i8 + k * 256 + 3 * 32 + 0), ix[3]);\n\
|
||||
_mm256_storeu_si256(reinterpret_cast<__m256i*>(qlut_i8 + k * 256 + 4 * 32 + 0), ix[4]);\n\
|
||||
_mm256_storeu_si256(reinterpret_cast<__m256i*>(qlut_i8 + k * 256 + 5 * 32 + 0), ix[5]);\n\
|
||||
_mm256_storeu_si256(reinterpret_cast<__m256i*>(qlut_i8 + k * 256 + 6 * 32 + 0), ix[6]);\n\
|
||||
_mm256_storeu_si256(reinterpret_cast<__m256i*>(qlut_i8 + k * 256 + 7 * 32 + 0), ix[7]);\n\
|
||||
\n\
|
||||
}\n\
|
||||
\n\
|
||||
*lut_scales = scales;\n\
|
||||
#endif\n\
|
||||
return 0;\n\
|
||||
}\n\
|
||||
\n\
|
||||
template<int act_k>\n\
|
||||
inline int32_t two_lut_ctor(int8_t* qlut, bitnet_float_type* b, bitnet_float_type* lut_scales) {\n\
|
||||
#if defined __AVX2__\n\
|
||||
__m256 vec_lut[16];\n\
|
||||
const __m256i vec_bi = _mm256_set_epi32(56, 48, 40, 32, 24, 16, 8, 0);\n\
|
||||
float scales = *lut_scales;\n\
|
||||
__m256i shuffle_mask = _mm256_set_epi8(\n\
|
||||
0x0f, 0x0d, 0x0b, 0x09, 0x07, 0x05, 0x03, 0x01,\n\
|
||||
0x0e, 0x0c, 0x0a, 0x08, 0x06, 0x04, 0x02, 0x00,\n\
|
||||
0x0f, 0x0d, 0x0b, 0x09, 0x07, 0x05, 0x03, 0x01,\n\
|
||||
0x0e, 0x0c, 0x0a, 0x08, 0x06, 0x04, 0x02, 0x00\n\
|
||||
);\n\
|
||||
#pragma unroll\n\
|
||||
for (int k = 0; k < act_k / 16; ++k) {\n\
|
||||
__m256 vec_b0f = _mm256_i32gather_ps(b + k * 16 + 0, vec_bi, 1);\n\
|
||||
__m256 vec_b1f = _mm256_i32gather_ps(b + k * 16 + 1, vec_bi, 1);\n\
|
||||
\n\
|
||||
__m256i vec_b0 = _mm256_cvtps_epi32(_mm256_round_ps(_mm256_mul_ps(vec_b0f, _mm256_set1_ps(scales)), _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC));\n\
|
||||
__m256i vec_b1 = _mm256_cvtps_epi32(_mm256_round_ps(_mm256_mul_ps(vec_b1f, _mm256_set1_ps(scales)), _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC));\n\
|
||||
vec_lut[15] = _mm256_setzero_si256();\n\
|
||||
vec_lut[14] = _mm256_setzero_si256();\n\
|
||||
vec_lut[13] = _mm256_setzero_si256();\n\
|
||||
vec_lut[12] = _mm256_setzero_si256();\n\
|
||||
vec_lut[11] = _mm256_setzero_si256();\n\
|
||||
vec_lut[10] = _mm256_setzero_si256();\n\
|
||||
vec_lut[9] = _mm256_setzero_si256();\n\
|
||||
vec_lut[8] = vec_b0;\n\
|
||||
vec_lut[8] = _mm256_add_epi32(vec_lut[8], vec_b1);\n\
|
||||
vec_lut[7] = vec_b0;\n\
|
||||
vec_lut[6] = vec_b0;\n\
|
||||
vec_lut[6] = _mm256_sub_epi32(vec_lut[6], vec_b1);\n\
|
||||
vec_lut[5] = vec_b1;\n\
|
||||
vec_lut[4] = _mm256_setzero_si256();\n\
|
||||
vec_lut[3] = _mm256_setzero_si256();\n\
|
||||
vec_lut[3] = _mm256_sub_epi32(vec_lut[3], vec_b1);\n\
|
||||
vec_lut[2] = _mm256_setzero_si256();\n\
|
||||
vec_lut[2] = _mm256_sub_epi32(vec_lut[2], vec_b0);\n\
|
||||
vec_lut[2] = _mm256_add_epi32(vec_lut[2], vec_b1);\n\
|
||||
vec_lut[1] = _mm256_setzero_si256();\n\
|
||||
vec_lut[1] = _mm256_sub_epi32(vec_lut[1], vec_b0);\n\
|
||||
vec_lut[0] = _mm256_setzero_si256();\n\
|
||||
vec_lut[0] = _mm256_sub_epi32(vec_lut[0], vec_b0);\n\
|
||||
vec_lut[0] = _mm256_sub_epi32(vec_lut[0], vec_b1);\n\
|
||||
\n\
|
||||
__m256i ix[16];\n\
|
||||
#pragma unroll\n\
|
||||
for (int g = 0; g < 16; ++g) {\n\
|
||||
ix[g] = vec_lut[g];\n\
|
||||
}\n\
|
||||
\n\
|
||||
Transpose_8_8(&(ix[0]), &(ix[1]), &(ix[2]), &(ix[3]), &(ix[4]), &(ix[5]),&(ix[6]), &(ix[7]));\n\
|
||||
Transpose_8_8(&(ix[8]), &(ix[9]), &(ix[10]), &(ix[11]), &(ix[12]), &(ix[13]),&(ix[14]), &(ix[15]));\n\
|
||||
\n\
|
||||
#pragma unroll\n\
|
||||
for (int g = 0; g < 8; ++g) {\n\
|
||||
ix[g] = _mm256_packs_epi32(ix[g], ix[g + 8]);\n\
|
||||
ix[g] = _mm256_permute4x64_epi64(ix[g], _MM_SHUFFLE(3, 1, 2, 0));\n\
|
||||
ix[g] = _mm256_shuffle_epi8(ix[g], shuffle_mask);\n\
|
||||
ix[g] = _mm256_permute4x64_epi64(ix[g], _MM_SHUFFLE(3, 1, 2, 0));\n\
|
||||
}\n\
|
||||
\n\
|
||||
int8_t* qlut_i8 = reinterpret_cast<int8_t*>(qlut);\n\
|
||||
\n\
|
||||
_mm256_storeu_si256(reinterpret_cast<__m256i*>(qlut_i8 + k * 256 + 0 * 32 + 0), ix[0]);\n\
|
||||
_mm256_storeu_si256(reinterpret_cast<__m256i*>(qlut_i8 + k * 256 + 1 * 32 + 0), ix[1]);\n\
|
||||
_mm256_storeu_si256(reinterpret_cast<__m256i*>(qlut_i8 + k * 256 + 2 * 32 + 0), ix[2]);\n\
|
||||
_mm256_storeu_si256(reinterpret_cast<__m256i*>(qlut_i8 + k * 256 + 3 * 32 + 0), ix[3]);\n\
|
||||
_mm256_storeu_si256(reinterpret_cast<__m256i*>(qlut_i8 + k * 256 + 4 * 32 + 0), ix[4]);\n\
|
||||
_mm256_storeu_si256(reinterpret_cast<__m256i*>(qlut_i8 + k * 256 + 5 * 32 + 0), ix[5]);\n\
|
||||
_mm256_storeu_si256(reinterpret_cast<__m256i*>(qlut_i8 + k * 256 + 6 * 32 + 0), ix[6]);\n\
|
||||
_mm256_storeu_si256(reinterpret_cast<__m256i*>(qlut_i8 + k * 256 + 7 * 32 + 0), ix[7]);\n\
|
||||
\n\
|
||||
}\n\
|
||||
*lut_scales = scales;\n\
|
||||
#endif\n\
|
||||
return 0;\n\
|
||||
}\n\
|
||||
static bool is_type_supported(enum ggml_type type) {\n\
|
||||
if (type == GGML_TYPE_Q4_0 ||\n\
|
||||
type == GGML_TYPE_TL2) {\n\
|
||||
return true;\n\
|
||||
} else {\n\
|
||||
return false;\n\
|
||||
}\n\
|
||||
}\n\
|
||||
"
|
||||
return kernel_code
|
||||
|
||||
def gen_tbl_impl(pre, BM, BK, bm, k_list):
|
||||
|
||||
kernel_code = "\
|
||||
#include <immintrin.h>\n\
|
||||
\n\
|
||||
#define BM{0} {1}\n\
|
||||
#define BBK{0} {2}\n\
|
||||
template<int batch_size, int K3>\n\
|
||||
inline void three_tbl_impl_{0}(int32_t* c, int8_t* lut, uint8_t* a, uint8_t* sign) {{\n\
|
||||
".format(pre, BM, BK)
|
||||
|
||||
kernel_code = "".join([kernel_code, "\
|
||||
#ifdef __AVX2__\n\
|
||||
const __m256i vec_mask = _mm256_set1_epi8(0x0f);\n\
|
||||
const __m256i vec_sign_mask = _mm256_set1_epi16(0x8000);\n\
|
||||
const __m256i vec_zero = _mm256_set1_epi8(0x00);\n\
|
||||
const __m256i vec_one = _mm256_set1_epi8(0xff);\n\
|
||||
const int KK = BBK{0} / 3;\n\
|
||||
#pragma unroll\n\
|
||||
for (int i = 0; i < BM{0}; i += 32) {{\n\
|
||||
__m256i vec_as[KK / 2];\n\
|
||||
__m256i vec_signs[KK / 8];\n\
|
||||
#pragma unroll\n\
|
||||
for (int ai = 0; ai < KK / 2; ai++) {{\n\
|
||||
vec_as[ai] = _mm256_loadu_si256(reinterpret_cast<__m256i*>(a + i * KK / 2 + ai * 32));\n\
|
||||
}}\n\
|
||||
#pragma unroll\n\
|
||||
for (int as = 0; as < KK / 8; as++) {{\n\
|
||||
vec_signs[as] = _mm256_loadu_si256(reinterpret_cast<__m256i*>(sign + i * KK / 8 + as * 32));\n\
|
||||
}}\n\
|
||||
#pragma unroll\n\
|
||||
for (int bs = 0; bs < batch_size; bs++) {{\n\
|
||||
__m256i vec_c0 = _mm256_setzero_si256();\n\
|
||||
__m256i vec_c1 = _mm256_setzero_si256();\n\
|
||||
#pragma unroll\n\
|
||||
for (int k = 0; k < KK / 8; k++) {{\n\
|
||||
__m256i vec_sign = vec_signs[k];\n\
|
||||
__m256i vec_a_0 = vec_as[k * 4 + 0];\n\
|
||||
__m128i vec_k1_0 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 0 * 64 + 0 + K3 / 3 * 32 * bs));\n\
|
||||
__m128i vec_k2_0 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 0 * 64 + 16 + K3 / 3 * 32 * bs));\n\
|
||||
__m128i vec_k3_0 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 0 * 64 + 32 + K3 / 3 * 32 * bs));\n\
|
||||
__m128i vec_k4_0 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 0 * 64 + 48 + K3 / 3 * 32 * bs));\n\
|
||||
__m256i vec_sign_left_hi_0 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 0)), 15);\n\
|
||||
__m256i vec_sign_left_lo_0 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 0 + 1)), 15);\n\
|
||||
__m256i vec_v_top_0 = _mm256_and_si256(_mm256_srli_epi16(vec_a_0, 4), vec_mask);\n\
|
||||
__m256i vec_v_top_fir_0 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k1_0, vec_k1_0), vec_v_top_0);\n\
|
||||
__m256i vec_v_top_sec_0 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k2_0, vec_k2_0), vec_v_top_0);\n\
|
||||
__m256i vec_sign_right_hi_0 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 0 + 2)), 15);\n\
|
||||
__m256i vec_sign_right_lo_0 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 0 + 3)), 15);\n\
|
||||
__m256i vec_v_bot_0 = _mm256_and_si256(vec_a_0, vec_mask);\n\
|
||||
__m256i vec_v_bot_fir_0 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k3_0, vec_k3_0), vec_v_bot_0);\n\
|
||||
__m256i vec_v_bot_sec_0 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k4_0, vec_k4_0), vec_v_bot_0);\n\
|
||||
__m256i vec_v_top_lo_0 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpackhi_epi8(vec_v_top_fir_0, vec_v_top_sec_0), vec_sign_left_lo_0), vec_sign_left_lo_0);\n\
|
||||
__m256i vec_v_top_hi_0 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpacklo_epi8(vec_v_top_fir_0, vec_v_top_sec_0), vec_sign_left_hi_0), vec_sign_left_hi_0);\n\
|
||||
__m256i vec_v_bot_lo_0 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpackhi_epi8(vec_v_bot_fir_0, vec_v_bot_sec_0), vec_sign_right_lo_0), vec_sign_right_lo_0);\n\
|
||||
__m256i vec_v_bot_hi_0 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpacklo_epi8(vec_v_bot_fir_0, vec_v_bot_sec_0), vec_sign_right_hi_0), vec_sign_right_hi_0);\n\
|
||||
vec_c0 = _mm256_add_epi16(vec_c0, vec_v_top_hi_0);\n\
|
||||
vec_c0 = _mm256_add_epi16(vec_c0, vec_v_bot_hi_0);\n\
|
||||
vec_c1 = _mm256_add_epi16(vec_c1, vec_v_top_lo_0);\n\
|
||||
vec_c1 = _mm256_add_epi16(vec_c1, vec_v_bot_lo_0);\n\
|
||||
__m256i vec_a_1 = vec_as[k * 4 + 1];\n\
|
||||
__m128i vec_k1_1 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 1 * 64 + 0 + K3 / 3 * 32 * bs));\n\
|
||||
__m128i vec_k2_1 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 1 * 64 + 16 + K3 / 3 * 32 * bs));\n\
|
||||
__m128i vec_k3_1 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 1 * 64 + 32 + K3 / 3 * 32 * bs));\n\
|
||||
__m128i vec_k4_1 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 1 * 64 + 48 + K3 / 3 * 32 * bs));\n\
|
||||
__m256i vec_sign_left_hi_1 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 1)), 15);\n\
|
||||
__m256i vec_sign_left_lo_1 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 1 + 1)), 15);\n\
|
||||
__m256i vec_v_top_1 = _mm256_and_si256(_mm256_srli_epi16(vec_a_1, 4), vec_mask);\n\
|
||||
__m256i vec_v_top_fir_1 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k1_1, vec_k1_1), vec_v_top_1);\n\
|
||||
__m256i vec_v_top_sec_1 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k2_1, vec_k2_1), vec_v_top_1);\n\
|
||||
__m256i vec_sign_right_hi_1 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 1 + 2)), 15);\n\
|
||||
__m256i vec_sign_right_lo_1 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 1 + 3)), 15);\n\
|
||||
__m256i vec_v_bot_1 = _mm256_and_si256(vec_a_1, vec_mask);\n\
|
||||
__m256i vec_v_bot_fir_1 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k3_1, vec_k3_1), vec_v_bot_1);\n\
|
||||
__m256i vec_v_bot_sec_1 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k4_1, vec_k4_1), vec_v_bot_1);\n\
|
||||
__m256i vec_v_top_lo_1 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpackhi_epi8(vec_v_top_fir_1, vec_v_top_sec_1), vec_sign_left_lo_1), vec_sign_left_lo_1);\n\
|
||||
__m256i vec_v_top_hi_1 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpacklo_epi8(vec_v_top_fir_1, vec_v_top_sec_1), vec_sign_left_hi_1), vec_sign_left_hi_1);\n\
|
||||
__m256i vec_v_bot_lo_1 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpackhi_epi8(vec_v_bot_fir_1, vec_v_bot_sec_1), vec_sign_right_lo_1), vec_sign_right_lo_1);\n\
|
||||
__m256i vec_v_bot_hi_1 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpacklo_epi8(vec_v_bot_fir_1, vec_v_bot_sec_1), vec_sign_right_hi_1), vec_sign_right_hi_1);\n\
|
||||
vec_c0 = _mm256_add_epi16(vec_c0, vec_v_top_hi_1);\n\
|
||||
vec_c0 = _mm256_add_epi16(vec_c0, vec_v_bot_hi_1);\n\
|
||||
vec_c1 = _mm256_add_epi16(vec_c1, vec_v_top_lo_1);\n\
|
||||
vec_c1 = _mm256_add_epi16(vec_c1, vec_v_bot_lo_1);\n\
|
||||
__m256i vec_a_2 = vec_as[k * 4 + 2];\n\
|
||||
__m128i vec_k1_2 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 2 * 64 + 0 + K3 / 3 * 32 * bs));\n\
|
||||
__m128i vec_k2_2 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 2 * 64 + 16 + K3 / 3 * 32 * bs));\n\
|
||||
__m128i vec_k3_2 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 2 * 64 + 32 + K3 / 3 * 32 * bs));\n\
|
||||
__m128i vec_k4_2 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 2 * 64 + 48 + K3 / 3 * 32 * bs));\n\
|
||||
__m256i vec_sign_left_hi_2 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 2)), 15);\n\
|
||||
__m256i vec_sign_left_lo_2 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 2 + 1)), 15);\n\
|
||||
__m256i vec_v_top_2 = _mm256_and_si256(_mm256_srli_epi16(vec_a_2, 4), vec_mask);\n\
|
||||
__m256i vec_v_top_fir_2 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k1_2, vec_k1_2), vec_v_top_2);\n\
|
||||
__m256i vec_v_top_sec_2 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k2_2, vec_k2_2), vec_v_top_2);\n\
|
||||
__m256i vec_sign_right_hi_2 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 2 + 2)), 15);\n\
|
||||
__m256i vec_sign_right_lo_2 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 2 + 3)), 15);\n\
|
||||
__m256i vec_v_bot_2 = _mm256_and_si256(vec_a_2, vec_mask);\n\
|
||||
__m256i vec_v_bot_fir_2 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k3_2, vec_k3_2), vec_v_bot_2);\n\
|
||||
__m256i vec_v_bot_sec_2 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k4_2, vec_k4_2), vec_v_bot_2);\n\
|
||||
__m256i vec_v_top_lo_2 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpackhi_epi8(vec_v_top_fir_2, vec_v_top_sec_2), vec_sign_left_lo_2), vec_sign_left_lo_2);\n\
|
||||
__m256i vec_v_top_hi_2 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpacklo_epi8(vec_v_top_fir_2, vec_v_top_sec_2), vec_sign_left_hi_2), vec_sign_left_hi_2);\n\
|
||||
__m256i vec_v_bot_lo_2 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpackhi_epi8(vec_v_bot_fir_2, vec_v_bot_sec_2), vec_sign_right_lo_2), vec_sign_right_lo_2);\n\
|
||||
__m256i vec_v_bot_hi_2 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpacklo_epi8(vec_v_bot_fir_2, vec_v_bot_sec_2), vec_sign_right_hi_2), vec_sign_right_hi_2);\n\
|
||||
vec_c0 = _mm256_add_epi16(vec_c0, vec_v_top_hi_2);\n\
|
||||
vec_c0 = _mm256_add_epi16(vec_c0, vec_v_bot_hi_2);\n\
|
||||
vec_c1 = _mm256_add_epi16(vec_c1, vec_v_top_lo_2);\n\
|
||||
vec_c1 = _mm256_add_epi16(vec_c1, vec_v_bot_lo_2);\n\
|
||||
__m256i vec_a_3 = vec_as[k * 4 + 3];\n\
|
||||
__m128i vec_k1_3 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 3 * 64 + 0 + K3 / 3 * 32 * bs));\n\
|
||||
__m128i vec_k2_3 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 3 * 64 + 16 + K3 / 3 * 32 * bs));\n\
|
||||
__m128i vec_k3_3 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 3 * 64 + 32 + K3 / 3 * 32 * bs));\n\
|
||||
__m128i vec_k4_3 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + 3 * 64 + 48 + K3 / 3 * 32 * bs));\n\
|
||||
__m256i vec_sign_left_hi_3 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 3)), 15);\n\
|
||||
__m256i vec_sign_left_lo_3 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 3 + 1)), 15);\n\
|
||||
__m256i vec_v_top_3 = _mm256_and_si256(_mm256_srli_epi16(vec_a_3, 4), vec_mask);\n\
|
||||
__m256i vec_v_top_fir_3 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k1_3, vec_k1_3), vec_v_top_3);\n\
|
||||
__m256i vec_v_top_sec_3 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k2_3, vec_k2_3), vec_v_top_3);\n\
|
||||
__m256i vec_sign_right_hi_3 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 3 + 2)), 15);\n\
|
||||
__m256i vec_sign_right_lo_3 = _mm256_srai_epi16(_mm256_slli_epi16(vec_sign, (4 * 3 + 3)), 15);\n\
|
||||
__m256i vec_v_bot_3 = _mm256_and_si256(vec_a_3, vec_mask);\n\
|
||||
__m256i vec_v_bot_fir_3 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k3_3, vec_k3_3), vec_v_bot_3);\n\
|
||||
__m256i vec_v_bot_sec_3 = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k4_3, vec_k4_3), vec_v_bot_3);\n\
|
||||
__m256i vec_v_top_lo_3 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpackhi_epi8(vec_v_top_fir_3, vec_v_top_sec_3), vec_sign_left_lo_3), vec_sign_left_lo_3);\n\
|
||||
__m256i vec_v_top_hi_3 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpacklo_epi8(vec_v_top_fir_3, vec_v_top_sec_3), vec_sign_left_hi_3), vec_sign_left_hi_3);\n\
|
||||
__m256i vec_v_bot_lo_3 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpackhi_epi8(vec_v_bot_fir_3, vec_v_bot_sec_3), vec_sign_right_lo_3), vec_sign_right_lo_3);\n\
|
||||
__m256i vec_v_bot_hi_3 = _mm256_xor_si256(_mm256_add_epi16(_mm256_unpacklo_epi8(vec_v_bot_fir_3, vec_v_bot_sec_3), vec_sign_right_hi_3), vec_sign_right_hi_3);\n\
|
||||
vec_c0 = _mm256_add_epi16(vec_c0, vec_v_top_hi_3);\n\
|
||||
vec_c0 = _mm256_add_epi16(vec_c0, vec_v_bot_hi_3);\n\
|
||||
vec_c1 = _mm256_add_epi16(vec_c1, vec_v_top_lo_3);\n\
|
||||
vec_c1 = _mm256_add_epi16(vec_c1, vec_v_bot_lo_3);\n\
|
||||
}}\n\
|
||||
__m256i vec_gc0 = _mm256_loadu_si256(reinterpret_cast<__m256i*>(c + i + BM{0} * bs));\n\
|
||||
__m256i vec_gc1 = _mm256_loadu_si256(reinterpret_cast<__m256i*>(c + i + 8 + BM{0} * bs));\n\
|
||||
__m256i vec_gc2 = _mm256_loadu_si256(reinterpret_cast<__m256i*>(c + i + 16 + BM{0} * bs));\n\
|
||||
__m256i vec_gc3 = _mm256_loadu_si256(reinterpret_cast<__m256i*>(c + i + 24 + BM{0} * bs));\n\
|
||||
vec_gc0 = _mm256_add_epi32(vec_gc0, _mm256_cvtepi16_epi32(_mm256_castsi256_si128(vec_c0)));\n\
|
||||
vec_gc1 = _mm256_add_epi32(vec_gc1, _mm256_cvtepi16_epi32(_mm256_extracti128_si256(vec_c0, 1)));\n\
|
||||
vec_gc2 = _mm256_add_epi32(vec_gc2, _mm256_cvtepi16_epi32(_mm256_castsi256_si128(vec_c1)));\n\
|
||||
vec_gc3 = _mm256_add_epi32(vec_gc3, _mm256_cvtepi16_epi32(_mm256_extracti128_si256(vec_c1, 1)));\n\
|
||||
_mm256_storeu_si256(reinterpret_cast<__m256i*>(c + i + BM{0} * bs), vec_gc0);\n\
|
||||
_mm256_storeu_si256(reinterpret_cast<__m256i*>(c + i + 8 + BM{0} * bs), vec_gc1);\n\
|
||||
_mm256_storeu_si256(reinterpret_cast<__m256i*>(c + i + 16 + BM{0} * bs), vec_gc2);\n\
|
||||
_mm256_storeu_si256(reinterpret_cast<__m256i*>(c + i + 24 + BM{0} * bs), vec_gc3);\n\
|
||||
}}\n\
|
||||
}}\n\
|
||||
#endif\n\
|
||||
}}\n\
|
||||
\n\
|
||||
template<int batch_size, int K2>\n\
|
||||
inline int32_t two_tbl_impl{0}(int32_t* c, int8_t* lut, uint8_t* a) {{\n\
|
||||
#ifdef __AVX2__\n\
|
||||
const __m256i vec_mask = _mm256_set1_epi8(0x0f);\n\
|
||||
const int KK = BK2 / 2;\n\
|
||||
#pragma unroll\n\
|
||||
for (int i = 0; i < BM{0}; i += 32) {{\n\
|
||||
__m256i vec_as[KK / 2];\n\
|
||||
#pragma unroll\n\
|
||||
for (int ai = 0; ai < KK / 2; ai++) {{\n\
|
||||
vec_as[ai] = _mm256_loadu_si256(reinterpret_cast<__m256i*>(a + i * KK / 2 + ai * 32));\n\
|
||||
}}\n\
|
||||
#pragma unroll\n\
|
||||
for (int bs = 0; bs < batch_size; bs++) {{\n\
|
||||
__m256i vec_c0 = _mm256_setzero_si256();\n\
|
||||
__m256i vec_c1 = _mm256_setzero_si256();\n\
|
||||
#pragma unroll\n\
|
||||
for (int k = 0; k < KK / 8; k++) {{\n\
|
||||
#pragma unroll\n\
|
||||
for (int j = 0; j < 4; j++) {{\n\
|
||||
__m256i vec_a = vec_as[k * 4 + j];\n\
|
||||
\n\
|
||||
__m128i vec_k1 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + j * 64 + 0 + K2 / 2 * 32 * bs));\n\
|
||||
__m128i vec_k2 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + j * 64 + 16 + K2 / 2 * 32 * bs));\n\
|
||||
__m128i vec_k3 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + j * 64 + 32 + K2 / 2 * 32 * bs));\n\
|
||||
__m128i vec_k4 = _mm_loadu_si128(reinterpret_cast<__m128i*>(lut + k * 32 * 8 + j * 64 + 48 + K2 / 2 * 32 * bs));\n\
|
||||
\n\
|
||||
__m256i vec_v_top = _mm256_and_si256(_mm256_srli_epi16(vec_a, 4), vec_mask);\n\
|
||||
__m256i vec_v_top_fir = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k1, vec_k1), vec_v_top);\n\
|
||||
__m256i vec_v_top_sec = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k2, vec_k2), vec_v_top);\n\
|
||||
\n\
|
||||
__m256i vec_v_bot = _mm256_and_si256(vec_a, vec_mask);\n\
|
||||
__m256i vec_v_bot_fir = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k3, vec_k3), vec_v_bot);\n\
|
||||
__m256i vec_v_bot_sec = _mm256_shuffle_epi8(_mm256_set_m128i(vec_k4, vec_k4), vec_v_bot);\n\
|
||||
\n\
|
||||
__m256i vec_v_top_lo = _mm256_unpackhi_epi8(vec_v_top_fir, vec_v_top_sec);\n\
|
||||
__m256i vec_v_top_hi = _mm256_unpacklo_epi8(vec_v_top_fir, vec_v_top_sec);\n\
|
||||
__m256i vec_v_bot_lo = _mm256_unpackhi_epi8(vec_v_bot_fir, vec_v_bot_sec);\n\
|
||||
__m256i vec_v_bot_hi = _mm256_unpacklo_epi8(vec_v_bot_fir, vec_v_bot_sec);\n\
|
||||
vec_c0 = _mm256_add_epi16(vec_c0, vec_v_top_hi);\n\
|
||||
vec_c0 = _mm256_add_epi16(vec_c0, vec_v_bot_hi);\n\
|
||||
vec_c1 = _mm256_add_epi16(vec_c1, vec_v_top_lo);\n\
|
||||
vec_c1 = _mm256_add_epi16(vec_c1, vec_v_bot_lo); \n\
|
||||
}}\n\
|
||||
}}\n\
|
||||
\n\
|
||||
__m256i vec_gc0 = _mm256_loadu_si256(reinterpret_cast<__m256i*>(c + i + BM{0} * bs));\n\
|
||||
__m256i vec_gc1 = _mm256_loadu_si256(reinterpret_cast<__m256i*>(c + i + 8 + BM{0} * bs));\n\
|
||||
__m256i vec_gc2 = _mm256_loadu_si256(reinterpret_cast<__m256i*>(c + i + 16 + BM{0} * bs));\n\
|
||||
__m256i vec_gc3 = _mm256_loadu_si256(reinterpret_cast<__m256i*>(c + i + 24 + BM{0} * bs));\n\
|
||||
\n\
|
||||
vec_gc0 = _mm256_add_epi32(vec_gc0, _mm256_cvtepi16_epi32(_mm256_castsi256_si128(vec_c0)));\n\
|
||||
vec_gc1 = _mm256_add_epi32(vec_gc1, _mm256_cvtepi16_epi32(_mm256_extracti128_si256(vec_c0, 1)));\n\
|
||||
vec_gc2 = _mm256_add_epi32(vec_gc2, _mm256_cvtepi16_epi32(_mm256_castsi256_si128(vec_c1)));\n\
|
||||
vec_gc3 = _mm256_add_epi32(vec_gc3, _mm256_cvtepi16_epi32(_mm256_extracti128_si256(vec_c1, 1)));\n\
|
||||
\n\
|
||||
_mm256_storeu_si256(reinterpret_cast<__m256i*>(c + i + BM{0} * bs), vec_gc0);\n\
|
||||
_mm256_storeu_si256(reinterpret_cast<__m256i*>(c + i + 8 + BM{0} * bs), vec_gc1);\n\
|
||||
_mm256_storeu_si256(reinterpret_cast<__m256i*>(c + i + 16 + BM{0} * bs), vec_gc2);\n\
|
||||
_mm256_storeu_si256(reinterpret_cast<__m256i*>(c + i + 24 + BM{0} * bs), vec_gc3);\n\
|
||||
}}\n\
|
||||
}}\n\
|
||||
#endif\n\
|
||||
return 0;\n\
|
||||
}}\n\
|
||||
\n\
|
||||
template<int BATCH_SIZE>\n\
|
||||
int32_t three_qgemm_lut_{0}(void* A, void* sign, void* LUT, void* Scales, void* LUT_Scales, void* C) {{\n\
|
||||
alignas(32) uint32_t CBits[BATCH_SIZE * BM{0}];\n\
|
||||
memset(&(CBits[0]), 0, BATCH_SIZE * BM{0} * sizeof(int32_t));\n\
|
||||
#pragma unroll\n\
|
||||
for (int32_t k_outer = 0; k_outer < {1} / BBK{0}; ++k_outer) {{\n\
|
||||
three_tbl_impl_{0}<BATCH_SIZE, {1}>((&(((int32_t*)CBits)[0])), (&(((int8_t*)LUT)[(k_outer * BBK{0} / 3 * 32)])), (&(((uint8_t*)A)[(k_outer * BBK{0} / 3 / 2 * BM{0})])), (&(((uint8_t*)sign)[(k_outer * BBK{0} / 3 / 8 * BM{0})])));\n\
|
||||
}}\n\
|
||||
#pragma unroll\n\
|
||||
for (int bs = 0; bs < BATCH_SIZE; bs++) {{\n\
|
||||
#pragma unroll\n\
|
||||
for (int i = 0; i < BM{0}; i++) {{\n\
|
||||
((int32_t*)C)[i] = (int32_t)(((int32_t*)CBits)[i + bs * BM{0}]);\n\
|
||||
}}\n\
|
||||
}}\n\
|
||||
return 0;\n\
|
||||
}}\n\
|
||||
\n\
|
||||
template<int BATCH_SIZE>\n\
|
||||
int32_t two_qgemm_lut_{0}(void* A, void* LUT, void* Scales, void* LUT_Scales, void* C) {{\n\
|
||||
alignas(32) uint32_t CBits[BATCH_SIZE * BM{0}];\n\
|
||||
memset(&(CBits[0]), 0, BATCH_SIZE * BM{0} * sizeof(int32_t));\n\
|
||||
#pragma unroll\n\
|
||||
for (int32_t k_outer = 0; k_outer < {2} / 32; ++k_outer) {{\n\
|
||||
two_tbl_impl{0}<BATCH_SIZE, {2}>((&(((int32_t*)CBits)[0])), (&(((int8_t*)LUT)[(k_outer * BK2 / 2 * 32)])), (&(((uint8_t*)A)[(k_outer * BK2 / 2 / 2 * BM{0})])));\n\
|
||||
}}\n\
|
||||
#pragma unroll\n\
|
||||
for (int bs = 0; bs < BATCH_SIZE; bs++) {{\n\
|
||||
#pragma unroll\n\
|
||||
for (int i = 0; i < BM{0}; i++) {{\n\
|
||||
((int32_t*)C)[i] += (int32_t)(((int32_t*)CBits)[i + bs * BM{0}]);\n\
|
||||
((float*)C)[i] = (float)(((int32_t*)C)[i]) / ((float*)LUT_Scales)[bs] * ((float*)Scales)[0];\n\
|
||||
}}\n\
|
||||
}}\n\
|
||||
return 0;\n\
|
||||
}}\n\
|
||||
\n\
|
||||
".format(pre, k_list[1], k_list[0])])
|
||||
return kernel_code
|
||||
|
||||
def gen_top_api(kernel_shapes, k_list):
|
||||
|
||||
kernel_code = "void ggml_preprocessor(int bs, int m, int three_k, int two_k, void* B, void* LUT_Scales, void* Three_QLUT, void* Two_QLUT) {{\n\
|
||||
partial_max_reset(bs, (&(((float*)LUT_Scales)[0])));\n\
|
||||
if (m == {0} && two_k == {1} && three_k == {2}) {{\n\
|
||||
for (int32_t b = 0; b < bs; b++) {{\n\
|
||||
per_tensor_quant(two_k + three_k, (&(((float*)LUT_Scales)[b])), (&(((float*)B)[b * (two_k + three_k)])));\n\
|
||||
three_lut_ctor<{2}>((&(((int8_t*)Three_QLUT)[b * three_k / 3 * 32])), (&(((float*)B)[b * (three_k + two_k)])), (&(((float*)LUT_Scales)[b])));\n\
|
||||
two_lut_ctor<{1}>((&(((int8_t*)Two_QLUT)[b * two_k / 2 * 32])), (&(((float*)B)[b * (three_k + two_k) + {2}])), (&(((float*)LUT_Scales)[b])));\n\
|
||||
}}\n\
|
||||
}}\n\
|
||||
".format(kernel_shapes[0][0], k_list[0][0], k_list[0][1])
|
||||
for i in range(1, len(kernel_shapes)):
|
||||
kernel_code = "".join([kernel_code, " else if (m == {0} && two_k == {1} && three_k == {2}) {{\n\
|
||||
for (int32_t b = 0; b < bs; b++) {{\n\
|
||||
per_tensor_quant(two_k + three_k, (&(((float*)LUT_Scales)[b])), (&(((float*)B)[b * (two_k + three_k)])));\n\
|
||||
three_lut_ctor<{2}>((&(((int8_t*)Three_QLUT)[b * three_k / 3 * 32])), (&(((float*)B)[b * (three_k + two_k)])), (&(((float*)LUT_Scales)[b])));\n\
|
||||
two_lut_ctor<{1}>((&(((int8_t*)Two_QLUT)[b * two_k / 2 * 32])), (&(((float*)B)[b * (three_k + two_k) + {2}])), (&(((float*)LUT_Scales)[b])));\n\
|
||||
}}\n\
|
||||
}}\n".format(kernel_shapes[i][0], k_list[i][0], k_list[i][1])])
|
||||
kernel_code = "".join([kernel_code, "}\n"])
|
||||
|
||||
|
||||
kernel_code = "".join([kernel_code, "void ggml_qgemm_lut(int bs, int m, int k, int BK, void* A, void* sign, void* LUT, void* Scales, void* LUT_Scales, void* C) {{\n\
|
||||
if (m == {0} && k == {1}) {{\n\
|
||||
if (BK == {2}) {{\n\
|
||||
if (bs == 1) {{\n\
|
||||
two_qgemm_lut_{4}<1>(A, LUT, Scales, LUT_Scales, C);\n\
|
||||
}} else if (bs == 8) {{\n\
|
||||
two_qgemm_lut_{4}<8>(A, LUT, Scales, LUT_Scales, C);\n\
|
||||
}} else if (bs == 32) {{\n\
|
||||
two_qgemm_lut_{4}<32>(A, LUT, Scales, LUT_Scales, C);\n\
|
||||
}} else if (bs == 128) {{\n\
|
||||
two_qgemm_lut_{4}<128>(A, LUT, Scales, LUT_Scales, C);\n\
|
||||
}} else if (bs == 256) {{\n\
|
||||
two_qgemm_lut_{4}<256>(A, LUT, Scales, LUT_Scales, C);\n\
|
||||
}} else if (bs == 512) {{\n\
|
||||
two_qgemm_lut_{4}<512>(A, LUT, Scales, LUT_Scales, C);\n\
|
||||
}}\n\
|
||||
}}\n\
|
||||
else if (BK == {3}) {{\n\
|
||||
if (bs == 1) {{\n\
|
||||
three_qgemm_lut_{4}<1>(A, sign, LUT, Scales, LUT_Scales, C);\n\
|
||||
}}else if (bs == 8) {{\n\
|
||||
three_qgemm_lut_{4}<8>(A, sign, LUT, Scales, LUT_Scales, C);\n\
|
||||
}}else if (bs == 32) {{\n\
|
||||
three_qgemm_lut_{4}<32>(A, sign, LUT, Scales, LUT_Scales, C);\n\
|
||||
}}else if (bs == 128) {{\n\
|
||||
three_qgemm_lut_{4}<128>(A, sign, LUT, Scales, LUT_Scales, C);\n\
|
||||
}}else if (bs == 256) {{\n\
|
||||
three_qgemm_lut_{4}<256>(A, sign, LUT, Scales, LUT_Scales, C);\n\
|
||||
}}else if (bs == 512) {{\n\
|
||||
three_qgemm_lut_{4}<512>(A, sign, LUT, Scales, LUT_Scales, C);\n\
|
||||
}}\n\
|
||||
}}\n\
|
||||
}}\n\
|
||||
".format(kernel_shapes[0][0], kernel_shapes[0][1], k_list[0][0], k_list[0][1], "{}_{}".format(kernel_shapes[0][0], kernel_shapes[0][1]))])
|
||||
for i in range(1, len(kernel_shapes)):
|
||||
kernel_code = "".join([kernel_code, " else if (m == {0} && k == {1}) {{\n\
|
||||
if (BK == {2}) {{\n\
|
||||
if (bs == 1) {{\n\
|
||||
two_qgemm_lut_{4}<1>(A, LUT, Scales, LUT_Scales, C);\n\
|
||||
}} else if (bs == 8) {{\n\
|
||||
two_qgemm_lut_{4}<8>(A, LUT, Scales, LUT_Scales, C);\n\
|
||||
}} else if (bs == 32) {{\n\
|
||||
two_qgemm_lut_{4}<32>(A, LUT, Scales, LUT_Scales, C);\n\
|
||||
}} else if (bs == 128) {{\n\
|
||||
two_qgemm_lut_{4}<128>(A, LUT, Scales, LUT_Scales, C);\n\
|
||||
}} else if (bs == 256) {{\n\
|
||||
two_qgemm_lut_{4}<256>(A, LUT, Scales, LUT_Scales, C);\n\
|
||||
}} else if (bs == 512) {{\n\
|
||||
two_qgemm_lut_{4}<512>(A, LUT, Scales, LUT_Scales, C);\n\
|
||||
}}\n\
|
||||
}}\n\
|
||||
else if (BK == {3}) {{\n\
|
||||
if (bs == 1) {{\n\
|
||||
three_qgemm_lut_{4}<1>(A, sign, LUT, Scales, LUT_Scales, C);\n\
|
||||
}}else if (bs == 8) {{\n\
|
||||
three_qgemm_lut_{4}<8>(A, sign, LUT, Scales, LUT_Scales, C);\n\
|
||||
}}else if (bs == 32) {{\n\
|
||||
three_qgemm_lut_{4}<32>(A, sign, LUT, Scales, LUT_Scales, C);\n\
|
||||
}}else if (bs == 128) {{\n\
|
||||
three_qgemm_lut_{4}<128>(A, sign, LUT, Scales, LUT_Scales, C);\n\
|
||||
}}else if (bs == 256) {{\n\
|
||||
three_qgemm_lut_{4}<256>(A, sign, LUT, Scales, LUT_Scales, C);\n\
|
||||
}}else if (bs == 512) {{\n\
|
||||
three_qgemm_lut_{4}<512>(A, sign, LUT, Scales, LUT_Scales, C);\n\
|
||||
}}\n\
|
||||
}}\n\
|
||||
}}\n\
|
||||
".format(kernel_shapes[i][0], kernel_shapes[i][1], k_list[i][0], k_list[i][1], "{}_{}".format(kernel_shapes[i][0], kernel_shapes[i][1]))])
|
||||
kernel_code = "".join([kernel_code, "}\n"])
|
||||
return kernel_code
|
||||
|
||||
def gen_transform_code(kernel_shapes):
|
||||
kernel_code = "\n\
|
||||
void ggml_bitnet_transform_tensor(struct ggml_tensor * tensor) {\n\
|
||||
if (!(is_type_supported(tensor->type) && tensor->backend == GGML_BACKEND_TYPE_CPU && tensor->extra == nullptr)) {\n\
|
||||
return;\n\
|
||||
}\n\
|
||||
\n\
|
||||
int k = tensor->ne[0];\n\
|
||||
int m = tensor->ne[1];\n\
|
||||
const int lut_scales_size = 1;\n\
|
||||
int bk = 0;\n\
|
||||
int bm = 0;\n"
|
||||
|
||||
kernel_code = "".join([kernel_code, "\n\
|
||||
if (m == {0} && k == {1}) {{\n\
|
||||
bm = BM{0}_{1};\n\
|
||||
bk = BBK{0}_{1};\n\
|
||||
}}\n".format(kernel_shapes[0][0], kernel_shapes[0][1])])
|
||||
|
||||
for i in range(1, len(kernel_shapes)):
|
||||
kernel_code = "".join([kernel_code, "else if (m == {0} && k == {1}) {{\n\
|
||||
bm = BM{0}_{1};\n\
|
||||
bk = BBK{0}_{1};\n\
|
||||
}}\n".format(kernel_shapes[i][0], kernel_shapes[i][1])])
|
||||
|
||||
kernel_code = "".join([kernel_code, "\n\
|
||||
const int n_tile_num = m / bm;\n\
|
||||
const int BK = bk;\n\
|
||||
uint8_t * qweights;\n\
|
||||
bitnet_float_type * scales;\n\
|
||||
\n\
|
||||
scales = (bitnet_float_type *) aligned_malloc(sizeof(bitnet_float_type));\n\
|
||||
qweights = (uint8_t *) tensor->data;\n\
|
||||
int nbytes = (k - 256) * m / 3 * 5 / 8 + 256 * m / 2 * 4 / 8;\n\
|
||||
if (nbytes % 32 != 0) nbytes = 32 - nbytes % 32 + nbytes;\n\
|
||||
float * i2_scales = (float * )(qweights + nbytes);\n\
|
||||
scales[0] = (bitnet_float_type) i2_scales[0];\n\
|
||||
\n\
|
||||
tensor->extra = bitnet_tensor_extras + bitnet_tensor_extras_index;\n\
|
||||
bitnet_tensor_extras[bitnet_tensor_extras_index++] = {\n\
|
||||
/* .lut_scales_size = */ lut_scales_size,\n\
|
||||
/* .BK = */ BK,\n\
|
||||
/* .n_tile_num = */ n_tile_num,\n\
|
||||
/* .qweights = */ qweights,\n\
|
||||
/* .scales = */ scales\n\
|
||||
};\n\
|
||||
}\n"])
|
||||
|
||||
return kernel_code
|
||||
|
||||
def get_three_k_two_k(K, bk):
|
||||
bk_num = K // bk
|
||||
three_k = bk_num * bk
|
||||
two_k = K - three_k
|
||||
return two_k, three_k
|
||||
|
||||
if __name__ == "__main__":
|
||||
ModelShapeDict = {
|
||||
"bitnet_b1_58-large" : [[1536, 4096],
|
||||
[1536, 1536],
|
||||
[4096, 1536]],
|
||||
"bitnet_b1_58-3B" : [[3200, 8640],
|
||||
[3200, 3200],
|
||||
[8640, 3200]],
|
||||
"Llama3-8B-1.58-100B-tokens" : [[14336, 4096],
|
||||
[4096, 14336],
|
||||
[1024, 4096],
|
||||
[4096, 4096]]
|
||||
}
|
||||
|
||||
parser = argparse.ArgumentParser(description='gen impl')
|
||||
parser.add_argument('--model',default="input", type=str, dest="model",
|
||||
help="choose from bitnet_b1_58-large/bitnet_b1_58-3B/Llama3-8B-1.58-100B-tokens.")
|
||||
parser.add_argument('--BM',default="input", type=str,
|
||||
help="block length when cutting one weight (M, K) into M / BM weights (BM, K).")
|
||||
parser.add_argument('--BK',default="input", type=str,
|
||||
help="block length when cutting one weight (M, K) into K / BK weights (M, BK).")
|
||||
parser.add_argument('--bm',default="input", type=str,
|
||||
help="using simd instructions to compute (bm, 192 / bm) in one block")
|
||||
args = parser.parse_args()
|
||||
|
||||
kernel_shapes = ModelShapeDict[args.model]
|
||||
|
||||
BM_list = [int(item) for item in args.BM.split(',')]
|
||||
BK_list = [int(item) for item in args.BK.split(',')]
|
||||
bm_list = [int(item) for item in args.bm.split(',')]
|
||||
|
||||
tbl_impl_code = []
|
||||
k_list = []
|
||||
|
||||
for i in range(len(kernel_shapes)):
|
||||
k_list.append(get_three_k_two_k(kernel_shapes[i][1], BK_list[i]))
|
||||
|
||||
for i in range(len(kernel_shapes)):
|
||||
tbl_impl_code.append(
|
||||
gen_tbl_impl("{}_{}".format(kernel_shapes[i][0], kernel_shapes[i][1]), BM_list[i], BK_list[i], bm_list[i], k_list[i])
|
||||
)
|
||||
|
||||
assert(len(BM_list) == len(BK_list) == len(bm_list) == len(kernel_shapes)), "number of BM / BK / bm shoud be {}".format(len(kernel_shapes))
|
||||
|
||||
for i in range(len(kernel_shapes)):
|
||||
assert kernel_shapes[i][0] % BM_list[i] == 0, "M %% BM should be 0"
|
||||
assert (kernel_shapes[i][1] % BK_list[i]) % 32 == 0, "K %% BK %% 32 should be 0"
|
||||
assert bm_list[i] in [32], "choose bm from [32]"
|
||||
|
||||
ctor_code = gen_ctor_code()
|
||||
api_code = gen_top_api(kernel_shapes, k_list)
|
||||
trans_code = gen_transform_code(kernel_shapes)
|
||||
|
||||
output_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "include")
|
||||
|
||||
with open(''.join([output_dir, "/bitnet-lut-kernels.h"]), 'w') as f:
|
||||
f.write(''.join("#if defined(GGML_BITNET_X86_TL2)"))
|
||||
f.write(''.join(ctor_code))
|
||||
for code in tbl_impl_code:
|
||||
f.write(''.join(code))
|
||||
f.write(''.join(api_code))
|
||||
f.write(''.join(trans_code))
|
||||
f.write(''.join("#endif"))
|
||||
|
||||
config = ConfigParser()
|
||||
|
||||
for i in range(len(kernel_shapes)):
|
||||
config.add_section('Kernels_{}'.format(i))
|
||||
config.set('Kernels_{}'.format(i), 'M'.format(i), str(kernel_shapes[i][0]))
|
||||
config.set('Kernels_{}'.format(i), 'K'.format(i), str(kernel_shapes[i][1]))
|
||||
config.set('Kernels_{}'.format(i), 'BM'.format(i), str(BM_list[i]))
|
||||
config.set('Kernels_{}'.format(i), 'BK'.format(i), str(BK_list[i]))
|
||||
config.set('Kernels_{}'.format(i), 'bmm'.format(i), str(bm_list[i]))
|
||||
|
||||
with open(''.join([output_dir, "/kernel_config.ini"]), 'w') as configfile:
|
||||
config.write(configfile)
|
||||
File diff suppressed because it is too large
Load Diff
+1711
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,52 @@
|
||||
import os
|
||||
import sys
|
||||
import logging
|
||||
import argparse
|
||||
import subprocess
|
||||
|
||||
def run_command(command, shell=False, log_step=None):
|
||||
"""Run a system command and ensure it succeeds."""
|
||||
if log_step:
|
||||
log_file = os.path.join(args.log_dir, log_step + ".log")
|
||||
with open(log_file, "w") as f:
|
||||
try:
|
||||
subprocess.run(command, shell=shell, check=True, stdout=f, stderr=f)
|
||||
except subprocess.CalledProcessError as e:
|
||||
logging.error(f"Error occurred while running command: {e}, check details in {log_file}")
|
||||
sys.exit(1)
|
||||
else:
|
||||
try:
|
||||
subprocess.run(command, shell=shell, check=True)
|
||||
except subprocess.CalledProcessError as e:
|
||||
logging.error(f"Error occurred while running command: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
def run_benchmark():
|
||||
bench_path = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "build/bin/llama-bench")
|
||||
if not os.path.exists(bench_path):
|
||||
logging.error(f"Benchmark binary not found, please build first.")
|
||||
sys.exit(1)
|
||||
command = [
|
||||
f'{bench_path}',
|
||||
'-m', args.model,
|
||||
'-n', str(args.n_token),
|
||||
'-ngl', '0',
|
||||
'-b', '1',
|
||||
'-t', str(args.threads),
|
||||
'-p', str(args.n_prompt),
|
||||
'-r', '5'
|
||||
]
|
||||
run_command(command)
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description='Setup the environment for running the inference')
|
||||
parser.add_argument("-m", "--model", type=str, help="Path to model file", required=True)
|
||||
parser.add_argument("-n", "--n-token", type=int, help="Number of generated tokens", required=False, default=128)
|
||||
parser.add_argument("-p", "--n-prompt", type=int, help="Prompt to generate text from", required=False, default=512)
|
||||
parser.add_argument("-t", "--threads", type=int, help="Number of threads to use", required=False, default=2)
|
||||
return parser.parse_args()
|
||||
|
||||
if __name__ == "__main__":
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
args = parse_args()
|
||||
run_benchmark()
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user