wehub-resource-sync
112 lines
2.4 KiB
Python
112 lines
2.4 KiB
Python
from pathlib import Path
|
|
|
|
import numpy as np
|
|
|
|
# had_16 = """
|
|
# ++++++++++++++++
|
|
# +-+-+-+-+-+-+-+-
|
|
# ++--++--++--++--
|
|
# +--++--++--++--+
|
|
# ++++----++++----
|
|
# +-+--+-++-+--+-+
|
|
# ++----++++----++
|
|
# +--+-++-+--+-++-
|
|
# ++++++++--------
|
|
# +-+-+-+--+-+-+-+
|
|
# ++--++----++--++
|
|
# +--++--+-++--++-
|
|
# ++++--------++++
|
|
# +-+--+-+-+-++-+-
|
|
# ++----++--++++--
|
|
# +--+-++--++-+--+
|
|
# """
|
|
|
|
# transformerEngine style rht matrix
|
|
had_16 = """
|
|
+++-+------+-+--
|
|
+-++++-+-+-----+
|
|
++-++-++--+--+++
|
|
+---+++--+++--+-
|
|
+++--+++---++-++
|
|
+-++--+--+--+++-
|
|
++-+-+----+-+---
|
|
+------+-+++++-+
|
|
+++-+---+++-+-++
|
|
+-++++-++-+++++-
|
|
++-++-++++-++---
|
|
+---+++-+---++-+
|
|
+++--++++++--+--
|
|
+-++--+-+-++---+
|
|
++-+-+--++-+-+++
|
|
+------++-----+-
|
|
"""
|
|
|
|
|
|
header = """
|
|
/******************************************************************************
|
|
* Copyright (c) 2023, Tri Dao.
|
|
* Adapted by Junxian Guo from https://github.com/Dao-AILab/fast-hadamard-transform/blob/master/csrc/code_gen.py
|
|
* Copyright (c) 2025, FourOverSix Team.
|
|
******************************************************************************/
|
|
|
|
// This file is auto-generated. See "hadamard_code_gen.py"\n
|
|
|
|
#pragma once
|
|
|
|
"""
|
|
|
|
template = """
|
|
namespace fouroversix {{
|
|
|
|
__device__ __forceinline__ void hadamard_mult_thread_{N}(float x[{N}]) {
|
|
float out[{N}];
|
|
{code}
|
|
#pragma unroll
|
|
for (int i = 0; i < {N}; i++) { x[i] = out[i]; }
|
|
}
|
|
|
|
}} // namespace fouroversix
|
|
|
|
"""
|
|
|
|
|
|
def string_to_array(string: str) -> np.ndarray:
|
|
# Convert strings of + and - to bool arrays
|
|
string = string.strip().replace("+", "1").replace("-", "-1").split()
|
|
return np.stack(
|
|
[
|
|
np.fromstring(" ".join(string[i]), dtype=np.int32, sep=" ")
|
|
for i in range(len(string))
|
|
],
|
|
)
|
|
|
|
|
|
def array_code_gen(arr: np.ndarray) -> str:
|
|
n = arr.shape[0]
|
|
if arr.shape[0] != arr.shape[1]:
|
|
msg = f"Hadamard matrix is not square: {arr.shape}"
|
|
raise ValueError(msg)
|
|
out = [
|
|
f"out[{i}] = "
|
|
+ " ".join([f"{'+' if arr[i, j] == 1 else '-'} x[{j}]" for j in range(n)])
|
|
+ ";"
|
|
for i in range(n)
|
|
]
|
|
return template.replace("{N}", str(n)).replace("{code}", "\n ".join(out))
|
|
|
|
|
|
def main() -> None:
|
|
output_dir = (
|
|
Path(__file__).parent.parent
|
|
/ "src"
|
|
/ "fouroversix"
|
|
/ "csrc"
|
|
/ "include"
|
|
/ "hadamard_transform_te.h"
|
|
)
|
|
output_dir.write_text(header + array_code_gen(string_to_array(had_16)))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|