[kernel slimming] Move fast_hadamard_transform to jit_kernel (#18475)
This commit is contained in:
Xiaoyu Zhang
@@ -0,0 +1,197 @@
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
|
||||
# From https://en.wikipedia.org/wiki/Paley_construction (construction II for q = 5)
|
||||
|
||||
had_12_paley = """
|
||||
+-++++++++++
|
||||
--+-+-+-+-+-
|
||||
+++-++----++
|
||||
+---+--+-++-
|
||||
+++++-++----
|
||||
+-+---+--+-+
|
||||
++--+++-++--
|
||||
+--++---+--+
|
||||
++----+++-++
|
||||
+--+-++---+-
|
||||
++++----+++-
|
||||
+-+--+-++---
|
||||
"""
|
||||
|
||||
# From http://neilsloane.com/hadamard/
|
||||
|
||||
had_12 = """
|
||||
+-----------
|
||||
++-+---+++-+
|
||||
+++-+---+++-
|
||||
+-++-+---+++
|
||||
++-++-+---++
|
||||
+++-++-+---+
|
||||
++++-++-+---
|
||||
+-+++-++-+--
|
||||
+--+++-++-+-
|
||||
+---+++-++-+
|
||||
++---+++-++-
|
||||
+-+---+++-++
|
||||
"""
|
||||
|
||||
had_20_will = """
|
||||
+----+----++--++-++-
|
||||
-+----+---+++---+-++
|
||||
--+----+---+++-+-+-+
|
||||
---+----+---+++++-+-
|
||||
----+----++--++-++-+
|
||||
-+++++-----+--+++--+
|
||||
+-+++-+---+-+--+++--
|
||||
++-++--+---+-+--+++-
|
||||
+++-+---+---+-+--+++
|
||||
++++-----++--+-+--++
|
||||
--++-+-++-+-----++++
|
||||
---++-+-++-+---+-+++
|
||||
+---++-+-+--+--++-++
|
||||
++---++-+----+-+++-+
|
||||
-++---++-+----+++++-
|
||||
-+--+--++-+----+----
|
||||
+-+-----++-+----+---
|
||||
-+-+-+---+--+----+--
|
||||
--+-+++------+----+-
|
||||
+--+--++------+----+
|
||||
"""
|
||||
|
||||
|
||||
had_28_will = """
|
||||
+------++----++-+--+-+--++--
|
||||
-+-----+++-----+-+--+-+--++-
|
||||
--+-----+++---+-+-+----+--++
|
||||
---+-----+++---+-+-+-+--+--+
|
||||
----+-----+++---+-+-+++--+--
|
||||
-----+-----++++--+-+--++--+-
|
||||
------++----++-+--+-+--++--+
|
||||
--++++-+-------++--+++-+--+-
|
||||
---++++-+-----+-++--+-+-+--+
|
||||
+---+++--+----++-++--+-+-+--
|
||||
++---++---+----++-++--+-+-+-
|
||||
+++---+----+----++-++--+-+-+
|
||||
++++--------+-+--++-++--+-+-
|
||||
-++++--------+++--++--+--+-+
|
||||
-+-++-++--++--+--------++++-
|
||||
+-+-++--+--++--+--------++++
|
||||
-+-+-++--+--++--+----+---+++
|
||||
+-+-+-++--+--+---+---++---++
|
||||
++-+-+-++--+------+--+++---+
|
||||
-++-+-+-++--+------+-++++---
|
||||
+-++-+---++--+------+-++++--
|
||||
-++--++-+-++-+++----++------
|
||||
+-++--++-+-++-+++-----+-----
|
||||
++-++---+-+-++-+++-----+----
|
||||
-++-++-+-+-+-+--+++-----+---
|
||||
--++-++++-+-+----+++-----+--
|
||||
+--++-+-++-+-+----+++-----+-
|
||||
++--++-+-++-+-+----++------+
|
||||
"""
|
||||
|
||||
|
||||
had_40_tpal = """
|
||||
+-------------------+-------------------
|
||||
++-++----+-+-++++--+++-++----+-+-++++--+
|
||||
+++-++----+-+-++++--+++-++----+-+-++++--
|
||||
+-++-++----+-+-++++-+-++-++----+-+-++++-
|
||||
+--++-++----+-+-+++++--++-++----+-+-++++
|
||||
++--++-++----+-+-+++++--++-++----+-+-+++
|
||||
+++--++-++----+-+-+++++--++-++----+-+-++
|
||||
++++--++-++----+-+-+++++--++-++----+-+-+
|
||||
+++++--++-++----+-+-+++++--++-++----+-+-
|
||||
+-++++--++-++----+-++-++++--++-++----+-+
|
||||
++-++++--++-++----+-++-++++--++-++----+-
|
||||
+-+-++++--++-++----++-+-++++--++-++----+
|
||||
++-+-++++--++-++----++-+-++++--++-++----
|
||||
+-+-+-++++--++-++---+-+-+-++++--++-++---
|
||||
+--+-+-++++--++-++--+--+-+-++++--++-++--
|
||||
+---+-+-++++--++-++-+---+-+-++++--++-++-
|
||||
+----+-+-++++--++-+++----+-+-++++--++-++
|
||||
++----+-+-++++--++-+++----+-+-++++--++-+
|
||||
+++----+-+-++++--++-+++----+-+-++++--++-
|
||||
+-++----+-+-++++--+++-++----+-+-++++--++
|
||||
+--------------------+++++++++++++++++++
|
||||
++-++----+-+-++++--+--+--++++-+-+----++-
|
||||
+++-++----+-+-++++-----+--++++-+-+----++
|
||||
+-++-++----+-+-++++--+--+--++++-+-+----+
|
||||
+--++-++----+-+-++++-++--+--++++-+-+----
|
||||
++--++-++----+-+-+++--++--+--++++-+-+---
|
||||
+++--++-++----+-+-++---++--+--++++-+-+--
|
||||
++++--++-++----+-+-+----++--+--++++-+-+-
|
||||
+++++--++-++----+-+------++--+--++++-+-+
|
||||
+-++++--++-++----+-+-+----++--+--++++-+-
|
||||
++-++++--++-++----+---+----++--+--++++-+
|
||||
+-+-++++--++-++----+-+-+----++--+--++++-
|
||||
++-+-++++--++-++------+-+----++--+--++++
|
||||
+-+-+-++++--++-++----+-+-+----++--+--+++
|
||||
+--+-+-++++--++-++---++-+-+----++--+--++
|
||||
+---+-+-++++--++-++--+++-+-+----++--+--+
|
||||
+----+-+-++++--++-++-++++-+-+----++--+--
|
||||
++----+-+-++++--++-+--++++-+-+----++--+-
|
||||
+++----+-+-++++--++----++++-+-+----++--+
|
||||
+-++----+-+-++++--++-+--++++-+-+----++--
|
||||
"""
|
||||
|
||||
|
||||
header = """
|
||||
/******************************************************************************
|
||||
* Copyright (c) 2023, Tri Dao.
|
||||
******************************************************************************/
|
||||
|
||||
// This file is auto-generated. See "code_gen.py"\n
|
||||
|
||||
#pragma once
|
||||
|
||||
"""
|
||||
|
||||
template = """
|
||||
__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]; }
|
||||
}
|
||||
|
||||
"""
|
||||
|
||||
|
||||
def string_to_array(string):
|
||||
# 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):
|
||||
N = arr.shape[0]
|
||||
assert arr.shape[0] == arr.shape[1]
|
||||
out = []
|
||||
for i in range(N):
|
||||
out.append(
|
||||
f"out[{i}] = "
|
||||
+ " ".join([f"{'+' if arr[i, j] == 1 else '-'} x[{j}]" for j in range(N)])
|
||||
+ ";"
|
||||
)
|
||||
return template.replace("{N}", str(N)).replace("{code}", "\n ".join(out))
|
||||
|
||||
|
||||
def main():
|
||||
output_dir = Path(__file__).parent / "fast_hadamard_transform_special.h"
|
||||
output_dir.write_text(
|
||||
header
|
||||
+ array_code_gen(string_to_array(had_12_paley))
|
||||
+ array_code_gen(string_to_array(had_20_will))
|
||||
+ array_code_gen(string_to_array(had_28_will))
|
||||
+ array_code_gen(string_to_array(had_40_tpal))
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,321 @@
|
||||
/******************************************************************************
|
||||
* Copyright (c) 2023, Tri Dao.
|
||||
******************************************************************************/
|
||||
|
||||
// Copied from https://github.com/sgl-project/fast-hadamard-transform
|
||||
|
||||
#include "fast_hadamard_transform.h"
|
||||
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#include <torch/extension.h>
|
||||
|
||||
#include <vector>
|
||||
|
||||
#define CHECK_SHAPE(x, ...) \
|
||||
TORCH_CHECK(x.sizes() == torch::IntArrayRef({__VA_ARGS__}), #x " must have shape (" #__VA_ARGS__ ")")
|
||||
|
||||
#define DISPATCH_ITYPE_FLOAT_AND_HALF_AND_BF16(ITYPE, NAME, ...) \
|
||||
if (ITYPE == at::ScalarType::Half) { \
|
||||
using input_t = at::Half; \
|
||||
__VA_ARGS__(); \
|
||||
} else if (ITYPE == at::ScalarType::BFloat16) { \
|
||||
using input_t = at::BFloat16; \
|
||||
__VA_ARGS__(); \
|
||||
} else if (ITYPE == at::ScalarType::Float) { \
|
||||
using input_t = float; \
|
||||
__VA_ARGS__(); \
|
||||
} else { \
|
||||
AT_ERROR(#NAME, " not implemented for input type '", toString(ITYPE), "'"); \
|
||||
}
|
||||
|
||||
template <typename input_t>
|
||||
void fast_hadamard_transform_cuda(HadamardParamsBase& params, cudaStream_t stream);
|
||||
|
||||
template <typename input_t>
|
||||
void fast_hadamard_transform_12N_cuda(HadamardParamsBase& params, cudaStream_t stream);
|
||||
|
||||
template <typename input_t>
|
||||
void fast_hadamard_transform_20N_cuda(HadamardParamsBase& params, cudaStream_t stream);
|
||||
|
||||
template <typename input_t>
|
||||
void fast_hadamard_transform_28N_cuda(HadamardParamsBase& params, cudaStream_t stream);
|
||||
|
||||
template <typename input_t>
|
||||
void fast_hadamard_transform_40N_cuda(HadamardParamsBase& params, cudaStream_t stream);
|
||||
|
||||
void set_hadamard_params(
|
||||
HadamardParamsBase& params,
|
||||
// sizes
|
||||
const size_t batch,
|
||||
const size_t dim,
|
||||
const size_t multiple,
|
||||
// device pointers
|
||||
const at::Tensor x,
|
||||
const at::Tensor out,
|
||||
float scale) {
|
||||
// Reset the parameters
|
||||
memset(¶ms, 0, sizeof(params));
|
||||
|
||||
params.batch = batch;
|
||||
params.dim = dim;
|
||||
params.log_N = int(ceil(std::log2(dim / multiple)));
|
||||
|
||||
// Set the pointers and strides.
|
||||
params.x_ptr = x.data_ptr();
|
||||
params.out_ptr = out.data_ptr();
|
||||
// All stride are in elements, not bytes.
|
||||
params.x_batch_stride = x.stride(0);
|
||||
params.out_batch_stride = out.stride(0);
|
||||
|
||||
params.scale = scale;
|
||||
}
|
||||
|
||||
at::Tensor fast_hadamard_transform(at::Tensor& x, float scale) {
|
||||
auto input_type = x.scalar_type();
|
||||
TORCH_CHECK(
|
||||
input_type == at::ScalarType::Float || input_type == at::ScalarType::Half ||
|
||||
input_type == at::ScalarType::BFloat16);
|
||||
|
||||
TORCH_CHECK(x.is_cuda());
|
||||
|
||||
const auto shapes_og = x.sizes();
|
||||
const int dim_og = x.size(-1);
|
||||
x = x.reshape({-1, dim_og});
|
||||
if (x.stride(-1) != 1) {
|
||||
x = x.contiguous();
|
||||
}
|
||||
const auto sizes = x.sizes();
|
||||
const int batch_size = sizes[0];
|
||||
|
||||
CHECK_SHAPE(x, batch_size, dim_og);
|
||||
TORCH_CHECK(x.stride(1) == 1);
|
||||
|
||||
if (dim_og % 8 != 0) {
|
||||
x = torch::nn::functional::pad(x, torch::nn::functional::PadFuncOptions({0, 8 - dim_og % 8}));
|
||||
}
|
||||
const int dim = x.size(1);
|
||||
|
||||
TORCH_CHECK(dim % 8 == 0, "fast_hadamard_transform only supports hidden dimension divisible by 8 for now");
|
||||
TORCH_CHECK(dim <= 32768, "fast_hadamard_transform only supports hidden dimension at most 32768 for now");
|
||||
|
||||
at::Tensor out = torch::empty_like(x);
|
||||
|
||||
HadamardParamsBase params;
|
||||
set_hadamard_params(params, batch_size, dim, 1, x, out, scale);
|
||||
|
||||
// Otherwise the kernel will be launched from cuda:0 device
|
||||
// Cast to char to avoid compiler warning about narrowing
|
||||
at::cuda::CUDAGuard device_guard{(char)x.get_device()};
|
||||
auto stream = at::cuda::getCurrentCUDAStream().stream();
|
||||
DISPATCH_ITYPE_FLOAT_AND_HALF_AND_BF16(
|
||||
x.scalar_type(), "fast_hadamard_transform", [&] { fast_hadamard_transform_cuda<input_t>(params, stream); });
|
||||
if (dim_og % 8 != 0) {
|
||||
out = out.index({torch::indexing::Slice(), torch::indexing::Slice(torch::indexing::None, dim_og)});
|
||||
}
|
||||
return out.reshape(shapes_og);
|
||||
}
|
||||
|
||||
at::Tensor fast_hadamard_transform_12N(at::Tensor& x, float scale) {
|
||||
auto input_type = x.scalar_type();
|
||||
TORCH_CHECK(
|
||||
input_type == at::ScalarType::Float || input_type == at::ScalarType::Half ||
|
||||
input_type == at::ScalarType::BFloat16);
|
||||
|
||||
TORCH_CHECK(x.is_cuda());
|
||||
|
||||
const auto shapes_og = x.sizes();
|
||||
const int dim_og = x.size(-1);
|
||||
x = x.reshape({-1, dim_og});
|
||||
if (x.stride(-1) != 1) {
|
||||
x = x.contiguous();
|
||||
}
|
||||
const auto sizes = x.sizes();
|
||||
const int batch_size = sizes[0];
|
||||
|
||||
CHECK_SHAPE(x, batch_size, dim_og);
|
||||
TORCH_CHECK(x.stride(1) == 1);
|
||||
|
||||
if (dim_og % (4 * 12) != 0) {
|
||||
x = torch::nn::functional::pad(x, torch::nn::functional::PadFuncOptions({0, (4 * 12) - dim_og % (4 * 12)}));
|
||||
}
|
||||
const int dim = x.size(1);
|
||||
|
||||
TORCH_CHECK(
|
||||
dim % (4 * 12) == 0, "fast_hadamard_transform_12N only supports hidden dimension divisible by 48 for now");
|
||||
TORCH_CHECK(dim <= 12 * 1024, "fast_hadamard_transform_12N only supports hidden dimension at most 12288 for now");
|
||||
|
||||
at::Tensor out = torch::empty_like(x);
|
||||
|
||||
HadamardParamsBase params;
|
||||
set_hadamard_params(params, batch_size, dim, 12, x, out, scale);
|
||||
|
||||
// Otherwise the kernel will be launched from cuda:0 device
|
||||
// Cast to char to avoid compiler warning about narrowing
|
||||
at::cuda::CUDAGuard device_guard{(char)x.get_device()};
|
||||
auto stream = at::cuda::getCurrentCUDAStream().stream();
|
||||
DISPATCH_ITYPE_FLOAT_AND_HALF_AND_BF16(
|
||||
x.scalar_type(), "fast_hadamard_transform", [&] { fast_hadamard_transform_12N_cuda<input_t>(params, stream); });
|
||||
if (dim_og % (4 * 12) != 0) {
|
||||
out = out.index({torch::indexing::Slice(), torch::indexing::Slice(torch::indexing::None, dim_og)});
|
||||
}
|
||||
return out.reshape(shapes_og);
|
||||
}
|
||||
|
||||
at::Tensor fast_hadamard_transform_20N(at::Tensor& x, float scale) {
|
||||
auto input_type = x.scalar_type();
|
||||
TORCH_CHECK(
|
||||
input_type == at::ScalarType::Float || input_type == at::ScalarType::Half ||
|
||||
input_type == at::ScalarType::BFloat16);
|
||||
|
||||
TORCH_CHECK(x.is_cuda());
|
||||
|
||||
const auto shapes_og = x.sizes();
|
||||
const int dim_og = x.size(-1);
|
||||
x = x.reshape({-1, dim_og});
|
||||
if (x.stride(-1) != 1) {
|
||||
x = x.contiguous();
|
||||
}
|
||||
const auto sizes = x.sizes();
|
||||
const int batch_size = sizes[0];
|
||||
|
||||
CHECK_SHAPE(x, batch_size, dim_og);
|
||||
TORCH_CHECK(x.stride(1) == 1);
|
||||
|
||||
if (dim_og % (4 * 20) != 0) {
|
||||
x = torch::nn::functional::pad(x, torch::nn::functional::PadFuncOptions({0, (4 * 20) - dim_og % (4 * 20)}));
|
||||
}
|
||||
const int dim = x.size(1);
|
||||
|
||||
TORCH_CHECK(
|
||||
dim % (4 * 20) == 0, "fast_hadamard_transform_20N only supports hidden dimension divisible by 80 for now");
|
||||
TORCH_CHECK(dim <= 20 * 1024, "fast_hadamard_transform_20N only supports hidden dimension at most 20480 for now");
|
||||
|
||||
at::Tensor out = torch::empty_like(x);
|
||||
|
||||
HadamardParamsBase params;
|
||||
set_hadamard_params(params, batch_size, dim, 20, x, out, scale);
|
||||
|
||||
// Otherwise the kernel will be launched from cuda:0 device
|
||||
// Cast to char to avoid compiler warning about narrowing
|
||||
at::cuda::CUDAGuard device_guard{(char)x.get_device()};
|
||||
auto stream = at::cuda::getCurrentCUDAStream().stream();
|
||||
DISPATCH_ITYPE_FLOAT_AND_HALF_AND_BF16(
|
||||
x.scalar_type(), "fast_hadamard_transform", [&] { fast_hadamard_transform_20N_cuda<input_t>(params, stream); });
|
||||
if (dim_og % (4 * 20) != 0) {
|
||||
out = out.index({torch::indexing::Slice(), torch::indexing::Slice(torch::indexing::None, dim_og)});
|
||||
}
|
||||
return out.reshape(shapes_og);
|
||||
}
|
||||
|
||||
at::Tensor fast_hadamard_transform_28N(at::Tensor& x, float scale) {
|
||||
auto input_type = x.scalar_type();
|
||||
TORCH_CHECK(
|
||||
input_type == at::ScalarType::Float || input_type == at::ScalarType::Half ||
|
||||
input_type == at::ScalarType::BFloat16);
|
||||
|
||||
TORCH_CHECK(x.is_cuda());
|
||||
|
||||
const auto shapes_og = x.sizes();
|
||||
const int dim_og = x.size(-1);
|
||||
x = x.reshape({-1, dim_og});
|
||||
if (x.stride(-1) != 1) {
|
||||
x = x.contiguous();
|
||||
}
|
||||
const auto sizes = x.sizes();
|
||||
const int batch_size = sizes[0];
|
||||
|
||||
CHECK_SHAPE(x, batch_size, dim_og);
|
||||
TORCH_CHECK(x.stride(1) == 1);
|
||||
|
||||
if (dim_og % (4 * 28) != 0) {
|
||||
x = torch::nn::functional::pad(x, torch::nn::functional::PadFuncOptions({0, (4 * 28) - dim_og % (4 * 28)}));
|
||||
}
|
||||
const int dim = x.size(1);
|
||||
|
||||
TORCH_CHECK(
|
||||
dim % (4 * 28) == 0, "fast_hadamard_transform_28N only supports hidden dimension divisible by 112 for now");
|
||||
TORCH_CHECK(dim <= 28 * 1024, "fast_hadamard_transform_28N only supports hidden dimension at most 28672 for now");
|
||||
|
||||
at::Tensor out = torch::empty_like(x);
|
||||
|
||||
HadamardParamsBase params;
|
||||
set_hadamard_params(params, batch_size, dim, 28, x, out, scale);
|
||||
|
||||
// Otherwise the kernel will be launched from cuda:0 device
|
||||
// Cast to char to avoid compiler warning about narrowing
|
||||
at::cuda::CUDAGuard device_guard{(char)x.get_device()};
|
||||
auto stream = at::cuda::getCurrentCUDAStream().stream();
|
||||
DISPATCH_ITYPE_FLOAT_AND_HALF_AND_BF16(
|
||||
x.scalar_type(), "fast_hadamard_transform", [&] { fast_hadamard_transform_28N_cuda<input_t>(params, stream); });
|
||||
if (dim_og % (8 * 28) != 0) {
|
||||
out = out.index({torch::indexing::Slice(), torch::indexing::Slice(torch::indexing::None, dim_og)});
|
||||
}
|
||||
return out.reshape(shapes_og);
|
||||
}
|
||||
|
||||
at::Tensor fast_hadamard_transform_40N(at::Tensor& x, float scale) {
|
||||
auto input_type = x.scalar_type();
|
||||
TORCH_CHECK(
|
||||
input_type == at::ScalarType::Float || input_type == at::ScalarType::Half ||
|
||||
input_type == at::ScalarType::BFloat16);
|
||||
|
||||
TORCH_CHECK(x.is_cuda());
|
||||
|
||||
const auto shapes_og = x.sizes();
|
||||
const int dim_og = x.size(-1);
|
||||
x = x.reshape({-1, dim_og});
|
||||
if (x.stride(-1) != 1) {
|
||||
x = x.contiguous();
|
||||
}
|
||||
const auto sizes = x.sizes();
|
||||
const int batch_size = sizes[0];
|
||||
|
||||
CHECK_SHAPE(x, batch_size, dim_og);
|
||||
TORCH_CHECK(x.stride(1) == 1);
|
||||
|
||||
if (dim_og % (4 * 40) != 0) {
|
||||
x = torch::nn::functional::pad(x, torch::nn::functional::PadFuncOptions({0, (4 * 40) - dim_og % (4 * 40)}));
|
||||
}
|
||||
const int dim = x.size(1);
|
||||
|
||||
TORCH_CHECK(
|
||||
dim % (4 * 40) == 0, "fast_hadamard_transform_40N only supports hidden dimension divisible by 160 for now");
|
||||
TORCH_CHECK(dim <= 40 * 1024, "fast_hadamard_transform_40N only supports hidden dimension at most 40960 for now");
|
||||
|
||||
at::Tensor out = torch::empty_like(x);
|
||||
|
||||
HadamardParamsBase params;
|
||||
set_hadamard_params(params, batch_size, dim, 40, x, out, scale);
|
||||
|
||||
// Otherwise the kernel will be launched from cuda:0 device
|
||||
// Cast to char to avoid compiler warning about narrowing
|
||||
at::cuda::CUDAGuard device_guard{(char)x.get_device()};
|
||||
auto stream = at::cuda::getCurrentCUDAStream().stream();
|
||||
DISPATCH_ITYPE_FLOAT_AND_HALF_AND_BF16(
|
||||
x.scalar_type(), "fast_hadamard_transform", [&] { fast_hadamard_transform_40N_cuda<input_t>(params, stream); });
|
||||
if (dim_og % (8 * 40) != 0) {
|
||||
out = out.index({torch::indexing::Slice(), torch::indexing::Slice(torch::indexing::None, dim_og)});
|
||||
}
|
||||
return out.reshape(shapes_og);
|
||||
}
|
||||
|
||||
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
||||
m.def("fast_hadamard_transform", &fast_hadamard_transform, "Fast Hadamard transform");
|
||||
m.def(
|
||||
"fast_hadamard_transform_12N",
|
||||
&fast_hadamard_transform_12N,
|
||||
"Fast Hadamard transform with dimension = 12 * power of 2");
|
||||
m.def(
|
||||
"fast_hadamard_transform_20N",
|
||||
&fast_hadamard_transform_20N,
|
||||
"Fast Hadamard transform with dimension = 20 * power of 2");
|
||||
m.def(
|
||||
"fast_hadamard_transform_28N",
|
||||
&fast_hadamard_transform_28N,
|
||||
"Fast Hadamard transform with dimension = 28 * power of 2");
|
||||
m.def(
|
||||
"fast_hadamard_transform_40N",
|
||||
&fast_hadamard_transform_40N,
|
||||
"Fast Hadamard transform with dimension = 40 * power of 2");
|
||||
}
|
||||
@@ -0,0 +1,24 @@
|
||||
/******************************************************************************
|
||||
* Copyright (c) 2023, Tri Dao.
|
||||
******************************************************************************/
|
||||
|
||||
// Copied from https://github.com/sgl-project/fast-hadamard-transform
|
||||
|
||||
#pragma once
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
struct HadamardParamsBase {
|
||||
using index_t = int64_t;
|
||||
|
||||
int batch, dim, log_N;
|
||||
|
||||
index_t x_batch_stride;
|
||||
index_t out_batch_stride;
|
||||
|
||||
float scale;
|
||||
|
||||
// Common data pointers.
|
||||
void* __restrict__ x_ptr;
|
||||
void* __restrict__ out_ptr;
|
||||
};
|
||||
@@ -0,0 +1,214 @@
|
||||
/******************************************************************************
|
||||
* Copyright (c) 2023, Tri Dao.
|
||||
******************************************************************************/
|
||||
|
||||
// Copied from https://github.com/sgl-project/fast-hadamard-transform
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <cuda_bf16.h>
|
||||
#include <cuda_fp16.h>
|
||||
|
||||
#define FULL_MASK 0xffffffff
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
struct uint8 {
|
||||
uint4 u;
|
||||
uint4 v;
|
||||
};
|
||||
|
||||
template <int BYTES>
|
||||
struct BytesToType {};
|
||||
|
||||
template <>
|
||||
struct BytesToType<32> {
|
||||
using Type = uint8;
|
||||
static_assert(sizeof(Type) == 32);
|
||||
};
|
||||
|
||||
template <>
|
||||
struct BytesToType<16> {
|
||||
using Type = uint4;
|
||||
static_assert(sizeof(Type) == 16);
|
||||
};
|
||||
|
||||
template <>
|
||||
struct BytesToType<8> {
|
||||
using Type = uint64_t;
|
||||
static_assert(sizeof(Type) == 8);
|
||||
};
|
||||
|
||||
template <>
|
||||
struct BytesToType<4> {
|
||||
using Type = uint32_t;
|
||||
static_assert(sizeof(Type) == 4);
|
||||
};
|
||||
|
||||
template <>
|
||||
struct BytesToType<2> {
|
||||
using Type = uint16_t;
|
||||
static_assert(sizeof(Type) == 2);
|
||||
};
|
||||
|
||||
template <>
|
||||
struct BytesToType<1> {
|
||||
using Type = uint8_t;
|
||||
static_assert(sizeof(Type) == 1);
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <typename T>
|
||||
struct SumOp {
|
||||
__device__ inline T operator()(T const& x, T const& y) {
|
||||
return x + y;
|
||||
}
|
||||
};
|
||||
|
||||
template <int THREADS>
|
||||
struct Allreduce {
|
||||
static_assert(THREADS == 32 || THREADS == 16 || THREADS == 8 || THREADS == 4);
|
||||
template <typename T, typename Operator>
|
||||
static __device__ inline T run(T x, Operator& op) {
|
||||
constexpr int OFFSET = THREADS / 2;
|
||||
x = op(x, __shfl_xor_sync(uint32_t(-1), x, OFFSET));
|
||||
return Allreduce<OFFSET>::run(x, op);
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct Allreduce<2> {
|
||||
template <typename T, typename Operator>
|
||||
static __device__ inline T run(T x, Operator& op) {
|
||||
x = op(x, __shfl_xor_sync(uint32_t(-1), x, 1));
|
||||
return x;
|
||||
}
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
// https://stackoverflow.com/questions/35311711/whats-the-right-way-to-compute-integral-base-2-logarithms-at-compile-time
|
||||
constexpr int cilog2(int val) {
|
||||
return val > 0 ? 1 + cilog2(val >> 1) : -1;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <int kLogN, int kNChunks>
|
||||
__device__ __forceinline__ void hadamard_mult_thread(float x[kNChunks][1 << kLogN]) {
|
||||
constexpr int N = 1 << kLogN;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kLogN; ++i) {
|
||||
const int stride = 1 << i;
|
||||
#pragma unroll
|
||||
for (int j = 0; j < N / 2; ++j) {
|
||||
const int lo = j & (stride - 1);
|
||||
const int idx = (j - lo) * 2 + lo;
|
||||
#pragma unroll
|
||||
for (int c = 0; c < kNChunks; ++c) {
|
||||
const float a = x[c][idx];
|
||||
const float b = x[c][idx + stride];
|
||||
x[c][idx] = a + b;
|
||||
x[c][idx + stride] = a - b;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <int kLogWarpSize, int kStepStart, int kNChunks, int kNItems>
|
||||
__device__ __forceinline__ void hadamard_mult_warp(float x[kNChunks][kNItems]) {
|
||||
constexpr int N = 1 << kLogWarpSize;
|
||||
int lane_id = threadIdx.x % N;
|
||||
#pragma unroll
|
||||
for (int step = kStepStart; step < kLogWarpSize; ++step) {
|
||||
const int lane_mask = 1 << step;
|
||||
const float sign = (lane_id & lane_mask) ? -1.f : 1.f;
|
||||
#pragma unroll
|
||||
for (int c = 0; c < kNChunks; ++c) {
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kNItems; ++i) {
|
||||
float x_val_other = __shfl_xor_sync(FULL_MASK, x[c][i], lane_mask);
|
||||
x[c][i] = sign * x[c][i] + x_val_other;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <int kNChunks, int kNElts, typename input_t>
|
||||
inline __device__ void load_input(input_t* x, float x_vals[kNChunks][kNElts], int dim) {
|
||||
using vec_t = typename BytesToType<sizeof(input_t) * kNElts>::Type;
|
||||
input_t x_vals_load[kNChunks][kNElts] = {0};
|
||||
#pragma unroll
|
||||
for (int c = 0; c < kNChunks; ++c) {
|
||||
if ((c * blockDim.x + threadIdx.x) * kNElts < dim) {
|
||||
reinterpret_cast<vec_t*>(x_vals_load)[c] = reinterpret_cast<const vec_t*>(x)[c * blockDim.x + threadIdx.x];
|
||||
}
|
||||
}
|
||||
#pragma unroll
|
||||
for (int c = 0; c < kNChunks; ++c) {
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kNElts; ++i) {
|
||||
x_vals[c][i] = float(x_vals_load[c][i]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <int kNChunks, int kNElts, typename output_t>
|
||||
inline __device__ void store_output(output_t* out, float out_vals[kNChunks][kNElts], int dim, float scale = 1.f) {
|
||||
using vec_t = typename BytesToType<sizeof(output_t) * kNElts>::Type;
|
||||
output_t out_vals_store[kNChunks][kNElts];
|
||||
#pragma unroll
|
||||
for (int c = 0; c < kNChunks; ++c) {
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kNElts; ++i) {
|
||||
out_vals_store[c][i] = out_vals[c][i] * scale;
|
||||
}
|
||||
}
|
||||
#pragma unroll
|
||||
for (int c = 0; c < kNChunks; ++c) {
|
||||
if ((c * blockDim.x + threadIdx.x) * kNElts < dim) {
|
||||
reinterpret_cast<vec_t*>(out)[c * blockDim.x + threadIdx.x] = reinterpret_cast<const vec_t*>(out_vals_store)[c];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
// Pre=true means the exchange before the hadamard_mult_warp, Pre=false means after.
|
||||
template <int kNChunks, int kChunksPerExchange, int kNElts, int kWarpSize, int kNWarps, bool Pre, typename vec_t>
|
||||
inline __device__ void exchange_smem_pre(float x_vals[kNChunks][kNElts], vec_t* smem) {
|
||||
constexpr int kNThreads = kWarpSize * kNWarps;
|
||||
constexpr int kNExchangePerVec = kNElts / (sizeof(vec_t) / sizeof(float));
|
||||
const int warp_id = threadIdx.x / kWarpSize;
|
||||
const int lane_id = threadIdx.x % kWarpSize;
|
||||
const int row_t = threadIdx.x % kNWarps;
|
||||
const int col_t = threadIdx.x / kNWarps;
|
||||
// We use the XOR swizzle trick (new_col = col ^ row) to avoid / reduce smem bank conflicts.
|
||||
#pragma unroll
|
||||
for (int c0 = 0; c0 < kNChunks / kChunksPerExchange; ++c0) {
|
||||
__syncthreads();
|
||||
#pragma unroll
|
||||
for (int c1 = 0; c1 < kChunksPerExchange; ++c1) {
|
||||
#pragma unroll
|
||||
for (int r = 0; r < kNExchangePerVec; ++r) {
|
||||
smem
|
||||
[(c1 * kNExchangePerVec + r) * kNThreads +
|
||||
(Pre ? warp_id * kWarpSize + lane_id ^ warp_id : row_t * kWarpSize + col_t ^ row_t)] =
|
||||
reinterpret_cast<vec_t*>(x_vals[c0 * kChunksPerExchange + c1])[r];
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
#pragma unroll
|
||||
for (int c1 = 0; c1 < kChunksPerExchange; ++c1) {
|
||||
#pragma unroll
|
||||
for (int r = 0; r < kNExchangePerVec; ++r) {
|
||||
reinterpret_cast<vec_t*>(x_vals[c0 * kChunksPerExchange + c1])[r] = smem
|
||||
[(c1 * kNExchangePerVec + r) * kNThreads +
|
||||
(Pre ? row_t * kWarpSize + col_t ^ row_t : warp_id * kWarpSize + lane_id ^ warp_id)];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,452 @@
|
||||
/******************************************************************************
|
||||
* Copyright (c) 2023, Tri Dao.
|
||||
******************************************************************************/
|
||||
|
||||
// Copied from https://github.com/sgl-project/fast-hadamard-transform
|
||||
|
||||
// #pragma once
|
||||
|
||||
#include <c10/cuda/CUDAException.h> // For C10_CUDA_CHECK and C10_CUDA_KERNEL_LAUNCH_CHECK
|
||||
#include <c10/util/BFloat16.h>
|
||||
#include <c10/util/Half.h>
|
||||
|
||||
#include "fast_hadamard_transform.h"
|
||||
#include "fast_hadamard_transform_common.h"
|
||||
#include "fast_hadamard_transform_special.h"
|
||||
#include "static_switch.h"
|
||||
|
||||
template <int kNThreads_, int kLogN_, typename input_t_>
|
||||
struct fast_hadamard_transform_kernel_traits {
|
||||
using input_t = input_t_;
|
||||
static constexpr int kNThreads = kNThreads_;
|
||||
static constexpr int kLogN = kLogN_;
|
||||
static constexpr int N = 1 << kLogN;
|
||||
static constexpr int kNBytes = sizeof(input_t);
|
||||
static_assert(kNBytes == 2 || kNBytes == 4);
|
||||
static constexpr int kNElts = kNBytes == 4 ? 4 : 8;
|
||||
// It's possible that we need to do 2 rounds of exchange if input_t is 16 bits
|
||||
// (since then we'd have 8 values of float, and each round we can exchange 4 floats).
|
||||
static constexpr int kNExchangePerVec = sizeof(float) / sizeof(input_t);
|
||||
using vec_t = typename BytesToType<kNBytes * kNElts>::Type;
|
||||
static constexpr int kNChunks = N / (kNElts * kNThreads);
|
||||
// We don't want to use more than 32 KB of shared memory.
|
||||
static constexpr int kSmemExchangeSize = std::min(N * 4, 32 * 1024);
|
||||
static constexpr int kNExchangeRounds = N * 4 / kSmemExchangeSize;
|
||||
static_assert(kNExchangeRounds * kSmemExchangeSize == N * 4);
|
||||
static constexpr int kSmemSize = kSmemExchangeSize;
|
||||
};
|
||||
|
||||
template <int kNThreads_, int kLogN_, typename input_t_>
|
||||
struct fast_hadamard_transform_12N_kernel_traits {
|
||||
using input_t = input_t_;
|
||||
static constexpr int kNThreads = kNThreads_;
|
||||
static constexpr int kLogN = kLogN_;
|
||||
static constexpr int N = (1 << kLogN) * 12;
|
||||
static_assert(N <= 12 * 1024, "fast_hadamard_transform_12 only supports dim <= 12288");
|
||||
static constexpr int kNBytes = sizeof(input_t);
|
||||
static_assert(kNBytes == 2 || kNBytes == 4);
|
||||
static constexpr int kNElts = 4;
|
||||
// It's possible that we need to do 2 rounds of exchange if input_t is 16 bits
|
||||
// (since then we'd have 8 values of float, and each round we can exchange 4 floats).
|
||||
static constexpr int kNExchangePerVec = sizeof(float) / sizeof(input_t);
|
||||
using vec_t = typename BytesToType<kNBytes * kNElts>::Type;
|
||||
static constexpr int kNChunks = N / (kNElts * kNThreads);
|
||||
static_assert(kNChunks == 12);
|
||||
// We don't want to use more than 24 KB of shared memory.
|
||||
static constexpr int kSmemExchangeSize = std::min(N * 4, 24 * 1024);
|
||||
static constexpr int kNExchangeRounds = N * 4 / kSmemExchangeSize;
|
||||
static_assert(kNExchangeRounds * kSmemExchangeSize == N * 4);
|
||||
static constexpr int kSmemSize = kSmemExchangeSize;
|
||||
};
|
||||
|
||||
template <int kNThreads_, int kLogN_, typename input_t_>
|
||||
struct fast_hadamard_transform_20N_kernel_traits {
|
||||
using input_t = input_t_;
|
||||
static constexpr int kNThreads = kNThreads_;
|
||||
static constexpr int kLogN = kLogN_;
|
||||
static constexpr int N = (1 << kLogN) * 20;
|
||||
static_assert(N <= 20 * 1024, "fast_hadamard_transform_20 only supports dim <= 20480");
|
||||
static constexpr int kNBytes = sizeof(input_t);
|
||||
static_assert(kNBytes == 2 || kNBytes == 4);
|
||||
static constexpr int kNElts = 4;
|
||||
// It's possible that we need to do 2 rounds of exchange if input_t is 16 bits
|
||||
// (since then we'd have 8 values of float, and each round we can exchange 4 floats).
|
||||
static constexpr int kNExchangePerVec = sizeof(float) / sizeof(input_t);
|
||||
using vec_t = typename BytesToType<kNBytes * kNElts>::Type;
|
||||
static constexpr int kNChunks = N / (kNElts * kNThreads);
|
||||
static_assert(kNChunks == 20);
|
||||
// We don't want to use more than 40 KB of shared memory.
|
||||
static constexpr int kSmemExchangeSize = std::min(N * 4, 40 * 1024);
|
||||
static constexpr int kNExchangeRounds = N * 4 / kSmemExchangeSize;
|
||||
static_assert(kNExchangeRounds * kSmemExchangeSize == N * 4);
|
||||
static constexpr int kSmemSize = kSmemExchangeSize;
|
||||
};
|
||||
|
||||
template <int kNThreads_, int kLogN_, typename input_t_>
|
||||
struct fast_hadamard_transform_28N_kernel_traits {
|
||||
using input_t = input_t_;
|
||||
static constexpr int kNThreads = kNThreads_;
|
||||
static constexpr int kLogN = kLogN_;
|
||||
static constexpr int N = (1 << kLogN) * 28;
|
||||
static_assert(N <= 28 * 1024, "fast_hadamard_transform_28 only supports dim <= 28672");
|
||||
static constexpr int kNBytes = sizeof(input_t);
|
||||
static_assert(kNBytes == 2 || kNBytes == 4);
|
||||
static constexpr int kNElts = 4;
|
||||
// It's possible that we need to do 2 rounds of exchange if input_t is 16 bits
|
||||
// (since then we'd have 8 values of float, and each round we can exchange 4 floats).
|
||||
static constexpr int kNExchangePerVec = sizeof(float) / sizeof(input_t);
|
||||
using vec_t = typename BytesToType<kNBytes * kNElts>::Type;
|
||||
static constexpr int kNChunks = N / (kNElts * kNThreads);
|
||||
static_assert(kNChunks == 28);
|
||||
// We don't want to use more than 28 KB of shared memory.
|
||||
static constexpr int kSmemExchangeSize = std::min(N * 4, 28 * 1024);
|
||||
static constexpr int kNExchangeRounds = N * 4 / kSmemExchangeSize;
|
||||
static_assert(kNExchangeRounds * kSmemExchangeSize == N * 4);
|
||||
static constexpr int kSmemSize = kSmemExchangeSize;
|
||||
};
|
||||
|
||||
template <int kNThreads_, int kLogN_, typename input_t_>
|
||||
struct fast_hadamard_transform_40N_kernel_traits {
|
||||
using input_t = input_t_;
|
||||
static constexpr int kNThreads = kNThreads_;
|
||||
static constexpr int kLogN = kLogN_;
|
||||
static constexpr int N = (1 << kLogN) * 40;
|
||||
static_assert(N <= 40 * 1024, "fast_hadamard_transform_40 only supports dim <= 40960");
|
||||
static constexpr int kNBytes = sizeof(input_t);
|
||||
static_assert(kNBytes == 2 || kNBytes == 4);
|
||||
static constexpr int kNElts = 4;
|
||||
// It's possible that we need to do 2 rounds of exchange if input_t is 16 bits
|
||||
// (since then we'd have 8 values of float, and each round we can exchange 4 floats).
|
||||
static constexpr int kNExchangePerVec = sizeof(float) / sizeof(input_t);
|
||||
using vec_t = typename BytesToType<kNBytes * kNElts>::Type;
|
||||
static constexpr int kNChunks = N / (kNElts * kNThreads);
|
||||
static_assert(kNChunks == 40);
|
||||
// We don't want to use more than 40 KB of shared memory.
|
||||
static constexpr int kSmemExchangeSize = std::min(N * 4, 40 * 1024);
|
||||
static constexpr int kNExchangeRounds = N * 4 / kSmemExchangeSize;
|
||||
static_assert(kNExchangeRounds * kSmemExchangeSize == N * 4);
|
||||
static constexpr int kSmemSize = kSmemExchangeSize;
|
||||
};
|
||||
|
||||
template <int kNChunks>
|
||||
__device__ __forceinline__ void hadamard_mult_thread_chunk_12(float x[kNChunks][12]) {
|
||||
#pragma unroll
|
||||
for (int c = 0; c < kNChunks; ++c) {
|
||||
hadamard_mult_thread_12(x[c]);
|
||||
}
|
||||
}
|
||||
|
||||
template <int kNChunks>
|
||||
__device__ __forceinline__ void hadamard_mult_thread_chunk_20(float x[kNChunks][20]) {
|
||||
#pragma unroll
|
||||
for (int c = 0; c < kNChunks; ++c) {
|
||||
hadamard_mult_thread_20(x[c]);
|
||||
}
|
||||
}
|
||||
|
||||
template <int kNChunks>
|
||||
__device__ __forceinline__ void hadamard_mult_thread_chunk_28(float x[kNChunks][28]) {
|
||||
#pragma unroll
|
||||
for (int c = 0; c < kNChunks; ++c) {
|
||||
hadamard_mult_thread_28(x[c]);
|
||||
}
|
||||
}
|
||||
|
||||
template <int kNChunks>
|
||||
__device__ __forceinline__ void hadamard_mult_thread_chunk_40(float x[kNChunks][40]) {
|
||||
#pragma unroll
|
||||
for (int c = 0; c < kNChunks; ++c) {
|
||||
hadamard_mult_thread_40(x[c]);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename Ktraits>
|
||||
__global__ __launch_bounds__(Ktraits::kNThreads) void fast_hadamard_transform_kernel(HadamardParamsBase params) {
|
||||
constexpr int kNThreads = Ktraits::kNThreads;
|
||||
constexpr int kNElts = Ktraits::kNElts;
|
||||
constexpr int kNExchangePerVec = Ktraits::kNExchangePerVec;
|
||||
constexpr int kNExchangeRounds = Ktraits::kNExchangeRounds;
|
||||
constexpr int kNChunks = Ktraits::kNChunks;
|
||||
using input_t = typename Ktraits::input_t;
|
||||
using vec_t = typename Ktraits::vec_t;
|
||||
|
||||
constexpr int kLogNElts = cilog2(Ktraits::kNElts);
|
||||
static_assert(1 << kLogNElts == kNElts, "kNElts must be a power of 2");
|
||||
constexpr int kWarpSize = std::min(kNThreads, 32);
|
||||
constexpr int kLogWarpSize = cilog2(kWarpSize);
|
||||
static_assert(1 << kLogWarpSize == kWarpSize, "Warp size must be a power of 2");
|
||||
constexpr int kNWarps = kNThreads / kWarpSize;
|
||||
constexpr int kLogNWarps = cilog2(kNWarps);
|
||||
static_assert(1 << kLogNWarps == kNWarps, "kNWarps must be a power of 2");
|
||||
constexpr int kLoadsPerExchange = Ktraits::kSmemExchangeSize / (sizeof(vec_t) * kNThreads);
|
||||
static_assert(
|
||||
kLoadsPerExchange * sizeof(vec_t) * kNThreads == Ktraits::kSmemExchangeSize,
|
||||
"kSmemExchangeSize should be a power of 2");
|
||||
static_assert(kNExchangeRounds * kLoadsPerExchange * sizeof(vec_t) == kNChunks * kNElts * sizeof(float));
|
||||
|
||||
constexpr int kChunksPerExchange = Ktraits::kSmemExchangeSize / (sizeof(vec_t) * kNExchangePerVec * kNThreads);
|
||||
static_assert(kChunksPerExchange * sizeof(vec_t) * kNExchangePerVec * kNThreads == Ktraits::kSmemExchangeSize);
|
||||
constexpr int kNExchanges = kNChunks / kChunksPerExchange;
|
||||
static_assert(kNExchanges * kChunksPerExchange == kNChunks);
|
||||
|
||||
// Shared memory.
|
||||
extern __shared__ char smem_[];
|
||||
vec_t* smem_exchange = reinterpret_cast<vec_t*>(smem_);
|
||||
|
||||
const int batch_id = blockIdx.x;
|
||||
input_t* x = reinterpret_cast<input_t*>(params.x_ptr) + batch_id * params.x_batch_stride;
|
||||
input_t* out = reinterpret_cast<input_t*>(params.out_ptr) + batch_id * params.out_batch_stride;
|
||||
|
||||
float x_vals[kNChunks][kNElts];
|
||||
load_input<kNChunks, kNElts, input_t>(x, x_vals, params.dim);
|
||||
|
||||
hadamard_mult_thread<kLogNElts, kNChunks>(x_vals);
|
||||
hadamard_mult_warp<kLogWarpSize, 0, kNChunks, kNElts>(x_vals);
|
||||
|
||||
if constexpr (kNWarps > 1) {
|
||||
exchange_smem_pre<kNChunks, kChunksPerExchange, kNElts, kWarpSize, kNWarps, true, vec_t>(x_vals, smem_exchange);
|
||||
hadamard_mult_warp<kLogNWarps, 0, kNChunks, kNElts>(x_vals);
|
||||
exchange_smem_pre<kNChunks, kChunksPerExchange, kNElts, kWarpSize, kNWarps, false, vec_t>(x_vals, smem_exchange);
|
||||
}
|
||||
|
||||
if constexpr (kNChunks > 1) {
|
||||
float x_vals_transposed[kNElts][kNChunks];
|
||||
#pragma unroll
|
||||
for (int c = 0; c < kNChunks; ++c) {
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kNElts; ++i) {
|
||||
x_vals_transposed[i][c] = x_vals[c][i];
|
||||
}
|
||||
}
|
||||
if constexpr (kNChunks == 12) {
|
||||
hadamard_mult_thread_chunk_12<kNElts>(x_vals_transposed);
|
||||
} else if constexpr (kNChunks == 20) {
|
||||
hadamard_mult_thread_chunk_20<kNElts>(x_vals_transposed);
|
||||
} else if constexpr (kNChunks == 28) {
|
||||
hadamard_mult_thread_chunk_28<kNElts>(x_vals_transposed);
|
||||
} else if constexpr (kNChunks == 40) {
|
||||
hadamard_mult_thread_chunk_40<kNElts>(x_vals_transposed);
|
||||
} else {
|
||||
constexpr int kLogNChunks = cilog2(kNChunks);
|
||||
static_assert(1 << kLogNChunks == kNChunks, "kNChunks must be a power of 2");
|
||||
hadamard_mult_thread<kLogNChunks, kNElts>(x_vals_transposed);
|
||||
}
|
||||
#pragma unroll
|
||||
for (int c = 0; c < kNChunks; ++c) {
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kNElts; ++i) {
|
||||
x_vals[c][i] = x_vals_transposed[i][c];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
store_output<kNChunks, kNElts, input_t>(out, x_vals, params.dim, params.scale);
|
||||
}
|
||||
|
||||
template <int kNThreads, int kLogN, typename input_t>
|
||||
void fast_hadamard_transform_launch(HadamardParamsBase& params, cudaStream_t stream) {
|
||||
using Ktraits = fast_hadamard_transform_kernel_traits<kNThreads, kLogN, input_t>;
|
||||
constexpr int kSmemSize = Ktraits::kSmemSize;
|
||||
dim3 grid(params.batch);
|
||||
auto kernel = &fast_hadamard_transform_kernel<Ktraits>;
|
||||
if (kSmemSize >= 48 * 1024) {
|
||||
C10_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemSize));
|
||||
}
|
||||
kernel<<<grid, Ktraits::kNThreads, kSmemSize, stream>>>(params);
|
||||
C10_CUDA_KERNEL_LAUNCH_CHECK();
|
||||
}
|
||||
|
||||
template <typename input_t>
|
||||
void fast_hadamard_transform_cuda(HadamardParamsBase& params, cudaStream_t stream) {
|
||||
if (params.log_N == 3) {
|
||||
fast_hadamard_transform_launch<1, 3, input_t>(params, stream);
|
||||
} else if (params.log_N == 4) {
|
||||
fast_hadamard_transform_launch<2, 4, input_t>(params, stream);
|
||||
} else if (params.log_N == 5) {
|
||||
fast_hadamard_transform_launch<4, 5, input_t>(params, stream);
|
||||
} else if (params.log_N == 6) {
|
||||
fast_hadamard_transform_launch<8, 6, input_t>(params, stream);
|
||||
} else if (params.log_N == 7) {
|
||||
fast_hadamard_transform_launch<16, 7, input_t>(params, stream);
|
||||
} else if (params.log_N == 8) {
|
||||
fast_hadamard_transform_launch<32, 8, input_t>(params, stream);
|
||||
} else if (params.log_N == 9) {
|
||||
fast_hadamard_transform_launch<32, 9, input_t>(params, stream);
|
||||
} else if (params.log_N == 10) {
|
||||
fast_hadamard_transform_launch<128, 10, input_t>(params, stream);
|
||||
} else if (params.log_N == 11) {
|
||||
fast_hadamard_transform_launch<256, 11, input_t>(params, stream);
|
||||
} else if (params.log_N == 12) {
|
||||
fast_hadamard_transform_launch<256, 12, input_t>(params, stream);
|
||||
} else if (params.log_N == 13) {
|
||||
fast_hadamard_transform_launch<256, 13, input_t>(params, stream);
|
||||
} else if (params.log_N == 14) {
|
||||
fast_hadamard_transform_launch<256, 14, input_t>(params, stream);
|
||||
} else if (params.log_N == 15) {
|
||||
fast_hadamard_transform_launch<256, 15, input_t>(params, stream);
|
||||
}
|
||||
}
|
||||
|
||||
template <int kNThreads, int kLogN, typename input_t>
|
||||
void fast_hadamard_transform_12N_launch(HadamardParamsBase& params, cudaStream_t stream) {
|
||||
using Ktraits = fast_hadamard_transform_12N_kernel_traits<kNThreads, kLogN, input_t>;
|
||||
constexpr int kSmemSize = Ktraits::kSmemSize;
|
||||
dim3 grid(params.batch);
|
||||
auto kernel = &fast_hadamard_transform_kernel<Ktraits>;
|
||||
if (kSmemSize >= 48 * 1024) {
|
||||
C10_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemSize));
|
||||
}
|
||||
kernel<<<grid, Ktraits::kNThreads, kSmemSize, stream>>>(params);
|
||||
C10_CUDA_KERNEL_LAUNCH_CHECK();
|
||||
}
|
||||
|
||||
template <typename input_t>
|
||||
void fast_hadamard_transform_12N_cuda(HadamardParamsBase& params, cudaStream_t stream) {
|
||||
if (params.log_N == 2) {
|
||||
fast_hadamard_transform_12N_launch<1, 2, input_t>(params, stream);
|
||||
} else if (params.log_N == 3) {
|
||||
fast_hadamard_transform_12N_launch<2, 3, input_t>(params, stream);
|
||||
} else if (params.log_N == 4) {
|
||||
fast_hadamard_transform_12N_launch<4, 4, input_t>(params, stream);
|
||||
} else if (params.log_N == 5) {
|
||||
fast_hadamard_transform_12N_launch<8, 5, input_t>(params, stream);
|
||||
} else if (params.log_N == 6) {
|
||||
fast_hadamard_transform_12N_launch<16, 6, input_t>(params, stream);
|
||||
} else if (params.log_N == 7) {
|
||||
fast_hadamard_transform_12N_launch<32, 7, input_t>(params, stream);
|
||||
} else if (params.log_N == 8) {
|
||||
fast_hadamard_transform_12N_launch<64, 8, input_t>(params, stream);
|
||||
} else if (params.log_N == 9) {
|
||||
fast_hadamard_transform_12N_launch<128, 9, input_t>(params, stream);
|
||||
} else if (params.log_N == 10) {
|
||||
fast_hadamard_transform_12N_launch<256, 10, input_t>(params, stream);
|
||||
}
|
||||
}
|
||||
|
||||
template <int kNThreads, int kLogN, typename input_t>
|
||||
void fast_hadamard_transform_20N_launch(HadamardParamsBase& params, cudaStream_t stream) {
|
||||
using Ktraits = fast_hadamard_transform_20N_kernel_traits<kNThreads, kLogN, input_t>;
|
||||
constexpr int kSmemSize = Ktraits::kSmemSize;
|
||||
dim3 grid(params.batch);
|
||||
auto kernel = &fast_hadamard_transform_kernel<Ktraits>;
|
||||
if (kSmemSize >= 48 * 1024) {
|
||||
C10_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemSize));
|
||||
}
|
||||
kernel<<<grid, Ktraits::kNThreads, kSmemSize, stream>>>(params);
|
||||
C10_CUDA_KERNEL_LAUNCH_CHECK();
|
||||
}
|
||||
|
||||
template <typename input_t>
|
||||
void fast_hadamard_transform_20N_cuda(HadamardParamsBase& params, cudaStream_t stream) {
|
||||
if (params.log_N == 2) {
|
||||
fast_hadamard_transform_20N_launch<1, 2, input_t>(params, stream);
|
||||
} else if (params.log_N == 3) {
|
||||
fast_hadamard_transform_20N_launch<2, 3, input_t>(params, stream);
|
||||
} else if (params.log_N == 4) {
|
||||
fast_hadamard_transform_20N_launch<4, 4, input_t>(params, stream);
|
||||
} else if (params.log_N == 5) {
|
||||
fast_hadamard_transform_20N_launch<8, 5, input_t>(params, stream);
|
||||
} else if (params.log_N == 6) {
|
||||
fast_hadamard_transform_20N_launch<16, 6, input_t>(params, stream);
|
||||
} else if (params.log_N == 7) {
|
||||
fast_hadamard_transform_20N_launch<32, 7, input_t>(params, stream);
|
||||
} else if (params.log_N == 8) {
|
||||
fast_hadamard_transform_20N_launch<64, 8, input_t>(params, stream);
|
||||
} else if (params.log_N == 9) {
|
||||
fast_hadamard_transform_20N_launch<128, 9, input_t>(params, stream);
|
||||
} else if (params.log_N == 10) {
|
||||
fast_hadamard_transform_20N_launch<256, 10, input_t>(params, stream);
|
||||
}
|
||||
}
|
||||
|
||||
template <int kNThreads, int kLogN, typename input_t>
|
||||
void fast_hadamard_transform_28N_launch(HadamardParamsBase& params, cudaStream_t stream) {
|
||||
using Ktraits = fast_hadamard_transform_28N_kernel_traits<kNThreads, kLogN, input_t>;
|
||||
constexpr int kSmemSize = Ktraits::kSmemSize;
|
||||
dim3 grid(params.batch);
|
||||
auto kernel = &fast_hadamard_transform_kernel<Ktraits>;
|
||||
if (kSmemSize >= 48 * 1024) {
|
||||
C10_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemSize));
|
||||
}
|
||||
kernel<<<grid, Ktraits::kNThreads, kSmemSize, stream>>>(params);
|
||||
C10_CUDA_KERNEL_LAUNCH_CHECK();
|
||||
}
|
||||
|
||||
template <typename input_t>
|
||||
void fast_hadamard_transform_28N_cuda(HadamardParamsBase& params, cudaStream_t stream) {
|
||||
if (params.log_N == 2) {
|
||||
fast_hadamard_transform_28N_launch<1, 2, input_t>(params, stream);
|
||||
} else if (params.log_N == 3) {
|
||||
fast_hadamard_transform_28N_launch<2, 3, input_t>(params, stream);
|
||||
} else if (params.log_N == 4) {
|
||||
fast_hadamard_transform_28N_launch<4, 4, input_t>(params, stream);
|
||||
} else if (params.log_N == 5) {
|
||||
fast_hadamard_transform_28N_launch<8, 5, input_t>(params, stream);
|
||||
} else if (params.log_N == 6) {
|
||||
fast_hadamard_transform_28N_launch<16, 6, input_t>(params, stream);
|
||||
} else if (params.log_N == 7) {
|
||||
fast_hadamard_transform_28N_launch<32, 7, input_t>(params, stream);
|
||||
} else if (params.log_N == 8) {
|
||||
fast_hadamard_transform_28N_launch<64, 8, input_t>(params, stream);
|
||||
} else if (params.log_N == 9) {
|
||||
fast_hadamard_transform_28N_launch<128, 9, input_t>(params, stream);
|
||||
} else if (params.log_N == 10) {
|
||||
fast_hadamard_transform_28N_launch<256, 10, input_t>(params, stream);
|
||||
}
|
||||
}
|
||||
|
||||
template <int kNThreads, int kLogN, typename input_t>
|
||||
void fast_hadamard_transform_40N_launch(HadamardParamsBase& params, cudaStream_t stream) {
|
||||
using Ktraits = fast_hadamard_transform_40N_kernel_traits<kNThreads, kLogN, input_t>;
|
||||
constexpr int kSmemSize = Ktraits::kSmemSize;
|
||||
dim3 grid(params.batch);
|
||||
auto kernel = &fast_hadamard_transform_kernel<Ktraits>;
|
||||
if (kSmemSize >= 48 * 1024) {
|
||||
C10_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemSize));
|
||||
}
|
||||
kernel<<<grid, Ktraits::kNThreads, kSmemSize, stream>>>(params);
|
||||
C10_CUDA_KERNEL_LAUNCH_CHECK();
|
||||
}
|
||||
|
||||
template <typename input_t>
|
||||
void fast_hadamard_transform_40N_cuda(HadamardParamsBase& params, cudaStream_t stream) {
|
||||
if (params.log_N == 2) {
|
||||
fast_hadamard_transform_40N_launch<1, 2, input_t>(params, stream);
|
||||
} else if (params.log_N == 3) {
|
||||
fast_hadamard_transform_40N_launch<2, 3, input_t>(params, stream);
|
||||
} else if (params.log_N == 4) {
|
||||
fast_hadamard_transform_40N_launch<4, 4, input_t>(params, stream);
|
||||
} else if (params.log_N == 5) {
|
||||
fast_hadamard_transform_40N_launch<8, 5, input_t>(params, stream);
|
||||
} else if (params.log_N == 6) {
|
||||
fast_hadamard_transform_40N_launch<16, 6, input_t>(params, stream);
|
||||
} else if (params.log_N == 7) {
|
||||
fast_hadamard_transform_40N_launch<32, 7, input_t>(params, stream);
|
||||
} else if (params.log_N == 8) {
|
||||
fast_hadamard_transform_40N_launch<64, 8, input_t>(params, stream);
|
||||
} else if (params.log_N == 9) {
|
||||
fast_hadamard_transform_40N_launch<128, 9, input_t>(params, stream);
|
||||
} else if (params.log_N == 10) {
|
||||
fast_hadamard_transform_40N_launch<256, 10, input_t>(params, stream);
|
||||
}
|
||||
}
|
||||
|
||||
template void fast_hadamard_transform_cuda<float>(HadamardParamsBase& params, cudaStream_t stream);
|
||||
template void fast_hadamard_transform_cuda<at::Half>(HadamardParamsBase& params, cudaStream_t stream);
|
||||
template void fast_hadamard_transform_cuda<at::BFloat16>(HadamardParamsBase& params, cudaStream_t stream);
|
||||
|
||||
template void fast_hadamard_transform_12N_cuda<float>(HadamardParamsBase& params, cudaStream_t stream);
|
||||
template void fast_hadamard_transform_12N_cuda<at::Half>(HadamardParamsBase& params, cudaStream_t stream);
|
||||
template void fast_hadamard_transform_12N_cuda<at::BFloat16>(HadamardParamsBase& params, cudaStream_t stream);
|
||||
|
||||
template void fast_hadamard_transform_20N_cuda<float>(HadamardParamsBase& params, cudaStream_t stream);
|
||||
template void fast_hadamard_transform_20N_cuda<at::Half>(HadamardParamsBase& params, cudaStream_t stream);
|
||||
template void fast_hadamard_transform_20N_cuda<at::BFloat16>(HadamardParamsBase& params, cudaStream_t stream);
|
||||
|
||||
template void fast_hadamard_transform_28N_cuda<float>(HadamardParamsBase& params, cudaStream_t stream);
|
||||
template void fast_hadamard_transform_28N_cuda<at::Half>(HadamardParamsBase& params, cudaStream_t stream);
|
||||
template void fast_hadamard_transform_28N_cuda<at::BFloat16>(HadamardParamsBase& params, cudaStream_t stream);
|
||||
|
||||
template void fast_hadamard_transform_40N_cuda<float>(HadamardParamsBase& params, cudaStream_t stream);
|
||||
template void fast_hadamard_transform_40N_cuda<at::Half>(HadamardParamsBase& params, cudaStream_t stream);
|
||||
template void fast_hadamard_transform_40N_cuda<at::BFloat16>(HadamardParamsBase& params, cudaStream_t stream);
|
||||
@@ -0,0 +1,298 @@
|
||||
|
||||
/******************************************************************************
|
||||
* Copyright (c) 2023, Tri Dao.
|
||||
******************************************************************************/
|
||||
|
||||
// Copied from https://github.com/sgl-project/fast-hadamard-transform
|
||||
|
||||
// This file is auto-generated. See "code_gen.py"
|
||||
|
||||
#pragma once
|
||||
|
||||
__device__ __forceinline__ void hadamard_mult_thread_12(float x[12]) {
|
||||
float out[12];
|
||||
out[0] = +x[0] - x[1] + x[2] + x[3] + x[4] + x[5] + x[6] + x[7] + x[8] + x[9] + x[10] + x[11];
|
||||
out[1] = -x[0] - x[1] + x[2] - x[3] + x[4] - x[5] + x[6] - x[7] + x[8] - x[9] + x[10] - x[11];
|
||||
out[2] = +x[0] + x[1] + x[2] - x[3] + x[4] + x[5] - x[6] - x[7] - x[8] - x[9] + x[10] + x[11];
|
||||
out[3] = +x[0] - x[1] - x[2] - x[3] + x[4] - x[5] - x[6] + x[7] - x[8] + x[9] + x[10] - x[11];
|
||||
out[4] = +x[0] + x[1] + x[2] + x[3] + x[4] - x[5] + x[6] + x[7] - x[8] - x[9] - x[10] - x[11];
|
||||
out[5] = +x[0] - x[1] + x[2] - x[3] - x[4] - x[5] + x[6] - x[7] - x[8] + x[9] - x[10] + x[11];
|
||||
out[6] = +x[0] + x[1] - x[2] - x[3] + x[4] + x[5] + x[6] - x[7] + x[8] + x[9] - x[10] - x[11];
|
||||
out[7] = +x[0] - x[1] - x[2] + x[3] + x[4] - x[5] - x[6] - x[7] + x[8] - x[9] - x[10] + x[11];
|
||||
out[8] = +x[0] + x[1] - x[2] - x[3] - x[4] - x[5] + x[6] + x[7] + x[8] - x[9] + x[10] + x[11];
|
||||
out[9] = +x[0] - x[1] - x[2] + x[3] - x[4] + x[5] + x[6] - x[7] - x[8] - x[9] + x[10] - x[11];
|
||||
out[10] = +x[0] + x[1] + x[2] + x[3] - x[4] - x[5] - x[6] - x[7] + x[8] + x[9] + x[10] - x[11];
|
||||
out[11] = +x[0] - x[1] + x[2] - x[3] - x[4] + x[5] - x[6] + x[7] + x[8] - x[9] - x[10] - x[11];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 12; i++) {
|
||||
x[i] = out[i];
|
||||
}
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void hadamard_mult_thread_20(float x[20]) {
|
||||
float out[20];
|
||||
out[0] = +x[0] - x[1] - x[2] - x[3] - x[4] + x[5] - x[6] - x[7] - x[8] - x[9] + x[10] + x[11] - x[12] - x[13] +
|
||||
x[14] + x[15] - x[16] + x[17] + x[18] - x[19];
|
||||
out[1] = -x[0] + x[1] - x[2] - x[3] - x[4] - x[5] + x[6] - x[7] - x[8] - x[9] + x[10] + x[11] + x[12] - x[13] -
|
||||
x[14] - x[15] + x[16] - x[17] + x[18] + x[19];
|
||||
out[2] = -x[0] - x[1] + x[2] - x[3] - x[4] - x[5] - x[6] + x[7] - x[8] - x[9] - x[10] + x[11] + x[12] + x[13] -
|
||||
x[14] + x[15] - x[16] + x[17] - x[18] + x[19];
|
||||
out[3] = -x[0] - x[1] - x[2] + x[3] - x[4] - x[5] - x[6] - x[7] + x[8] - x[9] - x[10] - x[11] + x[12] + x[13] +
|
||||
x[14] + x[15] + x[16] - x[17] + x[18] - x[19];
|
||||
out[4] = -x[0] - x[1] - x[2] - x[3] + x[4] - x[5] - x[6] - x[7] - x[8] + x[9] + x[10] - x[11] - x[12] + x[13] +
|
||||
x[14] - x[15] + x[16] + x[17] - x[18] + x[19];
|
||||
out[5] = -x[0] + x[1] + x[2] + x[3] + x[4] + x[5] - x[6] - x[7] - x[8] - x[9] - x[10] + x[11] - x[12] - x[13] +
|
||||
x[14] + x[15] + x[16] - x[17] - x[18] + x[19];
|
||||
out[6] = +x[0] - x[1] + x[2] + x[3] + x[4] - x[5] + x[6] - x[7] - x[8] - x[9] + x[10] - x[11] + x[12] - x[13] -
|
||||
x[14] + x[15] + x[16] + x[17] - x[18] - x[19];
|
||||
out[7] = +x[0] + x[1] - x[2] + x[3] + x[4] - x[5] - x[6] + x[7] - x[8] - x[9] - x[10] + x[11] - x[12] + x[13] -
|
||||
x[14] - x[15] + x[16] + x[17] + x[18] - x[19];
|
||||
out[8] = +x[0] + x[1] + x[2] - x[3] + x[4] - x[5] - x[6] - x[7] + x[8] - x[9] - x[10] - x[11] + x[12] - x[13] +
|
||||
x[14] - x[15] - x[16] + x[17] + x[18] + x[19];
|
||||
out[9] = +x[0] + x[1] + x[2] + x[3] - x[4] - x[5] - x[6] - x[7] - x[8] + x[9] + x[10] - x[11] - x[12] + x[13] -
|
||||
x[14] + x[15] - x[16] - x[17] + x[18] + x[19];
|
||||
out[10] = -x[0] - x[1] + x[2] + x[3] - x[4] + x[5] - x[6] + x[7] + x[8] - x[9] + x[10] - x[11] - x[12] - x[13] -
|
||||
x[14] - x[15] + x[16] + x[17] + x[18] + x[19];
|
||||
out[11] = -x[0] - x[1] - x[2] + x[3] + x[4] - x[5] + x[6] - x[7] + x[8] + x[9] - x[10] + x[11] - x[12] - x[13] -
|
||||
x[14] + x[15] - x[16] + x[17] + x[18] + x[19];
|
||||
out[12] = +x[0] - x[1] - x[2] - x[3] + x[4] + x[5] - x[6] + x[7] - x[8] + x[9] - x[10] - x[11] + x[12] - x[13] -
|
||||
x[14] + x[15] + x[16] - x[17] + x[18] + x[19];
|
||||
out[13] = +x[0] + x[1] - x[2] - x[3] - x[4] + x[5] + x[6] - x[7] + x[8] - x[9] - x[10] - x[11] - x[12] + x[13] -
|
||||
x[14] + x[15] + x[16] + x[17] - x[18] + x[19];
|
||||
out[14] = -x[0] + x[1] + x[2] - x[3] - x[4] - x[5] + x[6] + x[7] - x[8] + x[9] - x[10] - x[11] - x[12] - x[13] +
|
||||
x[14] + x[15] + x[16] + x[17] + x[18] - x[19];
|
||||
out[15] = -x[0] + x[1] - x[2] - x[3] + x[4] - x[5] - x[6] + x[7] + x[8] - x[9] + x[10] - x[11] - x[12] - x[13] -
|
||||
x[14] + x[15] - x[16] - x[17] - x[18] - x[19];
|
||||
out[16] = +x[0] - x[1] + x[2] - x[3] - x[4] - x[5] - x[6] - x[7] + x[8] + x[9] - x[10] + x[11] - x[12] - x[13] -
|
||||
x[14] - x[15] + x[16] - x[17] - x[18] - x[19];
|
||||
out[17] = -x[0] + x[1] - x[2] + x[3] - x[4] + x[5] - x[6] - x[7] - x[8] + x[9] - x[10] - x[11] + x[12] - x[13] -
|
||||
x[14] - x[15] - x[16] + x[17] - x[18] - x[19];
|
||||
out[18] = -x[0] - x[1] + x[2] - x[3] + x[4] + x[5] + x[6] - x[7] - x[8] - x[9] - x[10] - x[11] - x[12] + x[13] -
|
||||
x[14] - x[15] - x[16] - x[17] + x[18] - x[19];
|
||||
out[19] = +x[0] - x[1] - x[2] + x[3] - x[4] - x[5] + x[6] + x[7] - x[8] - x[9] - x[10] - x[11] - x[12] - x[13] +
|
||||
x[14] - x[15] - x[16] - x[17] - x[18] + x[19];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 20; i++) {
|
||||
x[i] = out[i];
|
||||
}
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void hadamard_mult_thread_28(float x[28]) {
|
||||
float out[28];
|
||||
out[0] = +x[0] - x[1] - x[2] - x[3] - x[4] - x[5] - x[6] + x[7] + x[8] - x[9] - x[10] - x[11] - x[12] + x[13] +
|
||||
x[14] - x[15] + x[16] - x[17] - x[18] + x[19] - x[20] + x[21] - x[22] - x[23] + x[24] + x[25] - x[26] -
|
||||
x[27];
|
||||
out[1] = -x[0] + x[1] - x[2] - x[3] - x[4] - x[5] - x[6] + x[7] + x[8] + x[9] - x[10] - x[11] - x[12] - x[13] -
|
||||
x[14] + x[15] - x[16] + x[17] - x[18] - x[19] + x[20] - x[21] + x[22] - x[23] - x[24] + x[25] + x[26] -
|
||||
x[27];
|
||||
out[2] = -x[0] - x[1] + x[2] - x[3] - x[4] - x[5] - x[6] - x[7] + x[8] + x[9] + x[10] - x[11] - x[12] - x[13] +
|
||||
x[14] - x[15] + x[16] - x[17] + x[18] - x[19] - x[20] - x[21] - x[22] + x[23] - x[24] - x[25] + x[26] +
|
||||
x[27];
|
||||
out[3] = -x[0] - x[1] - x[2] + x[3] - x[4] - x[5] - x[6] - x[7] - x[8] + x[9] + x[10] + x[11] - x[12] - x[13] -
|
||||
x[14] + x[15] - x[16] + x[17] - x[18] + x[19] - x[20] + x[21] - x[22] - x[23] + x[24] - x[25] - x[26] +
|
||||
x[27];
|
||||
out[4] = -x[0] - x[1] - x[2] - x[3] + x[4] - x[5] - x[6] - x[7] - x[8] - x[9] + x[10] + x[11] + x[12] - x[13] -
|
||||
x[14] - x[15] + x[16] - x[17] + x[18] - x[19] + x[20] + x[21] + x[22] - x[23] - x[24] + x[25] - x[26] -
|
||||
x[27];
|
||||
out[5] = -x[0] - x[1] - x[2] - x[3] - x[4] + x[5] - x[6] - x[7] - x[8] - x[9] - x[10] + x[11] + x[12] + x[13] +
|
||||
x[14] - x[15] - x[16] + x[17] - x[18] + x[19] - x[20] - x[21] + x[22] + x[23] - x[24] - x[25] + x[26] -
|
||||
x[27];
|
||||
out[6] = -x[0] - x[1] - x[2] - x[3] - x[4] - x[5] + x[6] + x[7] - x[8] - x[9] - x[10] - x[11] + x[12] + x[13] -
|
||||
x[14] + x[15] - x[16] - x[17] + x[18] - x[19] + x[20] - x[21] - x[22] + x[23] + x[24] - x[25] - x[26] +
|
||||
x[27];
|
||||
out[7] = -x[0] - x[1] + x[2] + x[3] + x[4] + x[5] - x[6] + x[7] - x[8] - x[9] - x[10] - x[11] - x[12] - x[13] -
|
||||
x[14] + x[15] + x[16] - x[17] - x[18] + x[19] + x[20] + x[21] - x[22] + x[23] - x[24] - x[25] + x[26] -
|
||||
x[27];
|
||||
out[8] = -x[0] - x[1] - x[2] + x[3] + x[4] + x[5] + x[6] - x[7] + x[8] - x[9] - x[10] - x[11] - x[12] - x[13] +
|
||||
x[14] - x[15] + x[16] + x[17] - x[18] - x[19] + x[20] - x[21] + x[22] - x[23] + x[24] - x[25] - x[26] +
|
||||
x[27];
|
||||
out[9] = +x[0] - x[1] - x[2] - x[3] + x[4] + x[5] + x[6] - x[7] - x[8] + x[9] - x[10] - x[11] - x[12] - x[13] +
|
||||
x[14] + x[15] - x[16] + x[17] + x[18] - x[19] - x[20] + x[21] - x[22] + x[23] - x[24] + x[25] - x[26] -
|
||||
x[27];
|
||||
out[10] = +x[0] + x[1] - x[2] - x[3] - x[4] + x[5] + x[6] - x[7] - x[8] - x[9] + x[10] - x[11] - x[12] - x[13] -
|
||||
x[14] + x[15] + x[16] - x[17] + x[18] + x[19] - x[20] - x[21] + x[22] - x[23] + x[24] - x[25] + x[26] -
|
||||
x[27];
|
||||
out[11] = +x[0] + x[1] + x[2] - x[3] - x[4] - x[5] + x[6] - x[7] - x[8] - x[9] - x[10] + x[11] - x[12] - x[13] -
|
||||
x[14] - x[15] + x[16] + x[17] - x[18] + x[19] + x[20] - x[21] - x[22] + x[23] - x[24] + x[25] - x[26] +
|
||||
x[27];
|
||||
out[12] = +x[0] + x[1] + x[2] + x[3] - x[4] - x[5] - x[6] - x[7] - x[8] - x[9] - x[10] - x[11] + x[12] - x[13] +
|
||||
x[14] - x[15] - x[16] + x[17] + x[18] - x[19] + x[20] + x[21] - x[22] - x[23] + x[24] - x[25] + x[26] -
|
||||
x[27];
|
||||
out[13] = -x[0] + x[1] + x[2] + x[3] + x[4] - x[5] - x[6] - x[7] - x[8] - x[9] - x[10] - x[11] - x[12] + x[13] +
|
||||
x[14] + x[15] - x[16] - x[17] + x[18] + x[19] - x[20] - x[21] + x[22] - x[23] - x[24] + x[25] - x[26] +
|
||||
x[27];
|
||||
out[14] = -x[0] + x[1] - x[2] + x[3] + x[4] - x[5] + x[6] + x[7] - x[8] - x[9] + x[10] + x[11] - x[12] - x[13] +
|
||||
x[14] - x[15] - x[16] - x[17] - x[18] - x[19] - x[20] - x[21] - x[22] + x[23] + x[24] + x[25] + x[26] -
|
||||
x[27];
|
||||
out[15] = +x[0] - x[1] + x[2] - x[3] + x[4] + x[5] - x[6] - x[7] + x[8] - x[9] - x[10] + x[11] + x[12] - x[13] -
|
||||
x[14] + x[15] - x[16] - x[17] - x[18] - x[19] - x[20] - x[21] - x[22] - x[23] + x[24] + x[25] + x[26] +
|
||||
x[27];
|
||||
out[16] = -x[0] + x[1] - x[2] + x[3] - x[4] + x[5] + x[6] - x[7] - x[8] + x[9] - x[10] - x[11] + x[12] + x[13] -
|
||||
x[14] - x[15] + x[16] - x[17] - x[18] - x[19] - x[20] + x[21] - x[22] - x[23] - x[24] + x[25] + x[26] +
|
||||
x[27];
|
||||
out[17] = +x[0] - x[1] + x[2] - x[3] + x[4] - x[5] + x[6] + x[7] - x[8] - x[9] + x[10] - x[11] - x[12] + x[13] -
|
||||
x[14] - x[15] - x[16] + x[17] - x[18] - x[19] - x[20] + x[21] + x[22] - x[23] - x[24] - x[25] + x[26] +
|
||||
x[27];
|
||||
out[18] = +x[0] + x[1] - x[2] + x[3] - x[4] + x[5] - x[6] + x[7] + x[8] - x[9] - x[10] + x[11] - x[12] - x[13] -
|
||||
x[14] - x[15] - x[16] - x[17] + x[18] - x[19] - x[20] + x[21] + x[22] + x[23] - x[24] - x[25] - x[26] +
|
||||
x[27];
|
||||
out[19] = -x[0] + x[1] + x[2] - x[3] + x[4] - x[5] + x[6] - x[7] + x[8] + x[9] - x[10] - x[11] + x[12] - x[13] -
|
||||
x[14] - x[15] - x[16] - x[17] - x[18] + x[19] - x[20] + x[21] + x[22] + x[23] + x[24] - x[25] - x[26] -
|
||||
x[27];
|
||||
out[20] = +x[0] - x[1] + x[2] + x[3] - x[4] + x[5] - x[6] - x[7] - x[8] + x[9] + x[10] - x[11] - x[12] + x[13] -
|
||||
x[14] - x[15] - x[16] - x[17] - x[18] - x[19] + x[20] - x[21] + x[22] + x[23] + x[24] + x[25] - x[26] -
|
||||
x[27];
|
||||
out[21] = -x[0] + x[1] + x[2] - x[3] - x[4] + x[5] + x[6] - x[7] + x[8] - x[9] + x[10] + x[11] - x[12] + x[13] +
|
||||
x[14] + x[15] - x[16] - x[17] - x[18] - x[19] + x[20] + x[21] - x[22] - x[23] - x[24] - x[25] - x[26] -
|
||||
x[27];
|
||||
out[22] = +x[0] - x[1] + x[2] + x[3] - x[4] - x[5] + x[6] + x[7] - x[8] + x[9] - x[10] + x[11] + x[12] - x[13] +
|
||||
x[14] + x[15] + x[16] - x[17] - x[18] - x[19] - x[20] - x[21] + x[22] - x[23] - x[24] - x[25] - x[26] -
|
||||
x[27];
|
||||
out[23] = +x[0] + x[1] - x[2] + x[3] + x[4] - x[5] - x[6] - x[7] + x[8] - x[9] + x[10] - x[11] + x[12] + x[13] -
|
||||
x[14] + x[15] + x[16] + x[17] - x[18] - x[19] - x[20] - x[21] - x[22] + x[23] - x[24] - x[25] - x[26] -
|
||||
x[27];
|
||||
out[24] = -x[0] + x[1] + x[2] - x[3] + x[4] + x[5] - x[6] + x[7] - x[8] + x[9] - x[10] + x[11] - x[12] + x[13] -
|
||||
x[14] - x[15] + x[16] + x[17] + x[18] - x[19] - x[20] - x[21] - x[22] - x[23] + x[24] - x[25] - x[26] -
|
||||
x[27];
|
||||
out[25] = -x[0] - x[1] + x[2] + x[3] - x[4] + x[5] + x[6] + x[7] + x[8] - x[9] + x[10] - x[11] + x[12] - x[13] -
|
||||
x[14] - x[15] - x[16] + x[17] + x[18] + x[19] - x[20] - x[21] - x[22] - x[23] - x[24] + x[25] - x[26] -
|
||||
x[27];
|
||||
out[26] = +x[0] - x[1] - x[2] + x[3] + x[4] - x[5] + x[6] - x[7] + x[8] + x[9] - x[10] + x[11] - x[12] + x[13] -
|
||||
x[14] - x[15] - x[16] - x[17] + x[18] + x[19] + x[20] - x[21] - x[22] - x[23] - x[24] - x[25] + x[26] -
|
||||
x[27];
|
||||
out[27] = +x[0] + x[1] - x[2] - x[3] + x[4] + x[5] - x[6] + x[7] - x[8] + x[9] + x[10] - x[11] + x[12] - x[13] +
|
||||
x[14] - x[15] - x[16] - x[17] - x[18] + x[19] + x[20] - x[21] - x[22] - x[23] - x[24] - x[25] - x[26] +
|
||||
x[27];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 28; i++) {
|
||||
x[i] = out[i];
|
||||
}
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void hadamard_mult_thread_40(float x[40]) {
|
||||
float out[40];
|
||||
out[0] = +x[0] - x[1] - x[2] - x[3] - x[4] - x[5] - x[6] - x[7] - x[8] - x[9] - x[10] - x[11] - x[12] - x[13] -
|
||||
x[14] - x[15] - x[16] - x[17] - x[18] - x[19] + x[20] - x[21] - x[22] - x[23] - x[24] - x[25] - x[26] -
|
||||
x[27] - x[28] - x[29] - x[30] - x[31] - x[32] - x[33] - x[34] - x[35] - x[36] - x[37] - x[38] - x[39];
|
||||
out[1] = +x[0] + x[1] - x[2] + x[3] + x[4] - x[5] - x[6] - x[7] - x[8] + x[9] - x[10] + x[11] - x[12] + x[13] +
|
||||
x[14] + x[15] + x[16] - x[17] - x[18] + x[19] + x[20] + x[21] - x[22] + x[23] + x[24] - x[25] - x[26] -
|
||||
x[27] - x[28] + x[29] - x[30] + x[31] - x[32] + x[33] + x[34] + x[35] + x[36] - x[37] - x[38] + x[39];
|
||||
out[2] = +x[0] + x[1] + x[2] - x[3] + x[4] + x[5] - x[6] - x[7] - x[8] - x[9] + x[10] - x[11] + x[12] - x[13] +
|
||||
x[14] + x[15] + x[16] + x[17] - x[18] - x[19] + x[20] + x[21] + x[22] - x[23] + x[24] + x[25] - x[26] -
|
||||
x[27] - x[28] - x[29] + x[30] - x[31] + x[32] - x[33] + x[34] + x[35] + x[36] + x[37] - x[38] - x[39];
|
||||
out[3] = +x[0] - x[1] + x[2] + x[3] - x[4] + x[5] + x[6] - x[7] - x[8] - x[9] - x[10] + x[11] - x[12] + x[13] -
|
||||
x[14] + x[15] + x[16] + x[17] + x[18] - x[19] + x[20] - x[21] + x[22] + x[23] - x[24] + x[25] + x[26] -
|
||||
x[27] - x[28] - x[29] - x[30] + x[31] - x[32] + x[33] - x[34] + x[35] + x[36] + x[37] + x[38] - x[39];
|
||||
out[4] = +x[0] - x[1] - x[2] + x[3] + x[4] - x[5] + x[6] + x[7] - x[8] - x[9] - x[10] - x[11] + x[12] - x[13] +
|
||||
x[14] - x[15] + x[16] + x[17] + x[18] + x[19] + x[20] - x[21] - x[22] + x[23] + x[24] - x[25] + x[26] +
|
||||
x[27] - x[28] - x[29] - x[30] - x[31] + x[32] - x[33] + x[34] - x[35] + x[36] + x[37] + x[38] + x[39];
|
||||
out[5] = +x[0] + x[1] - x[2] - x[3] + x[4] + x[5] - x[6] + x[7] + x[8] - x[9] - x[10] - x[11] - x[12] + x[13] -
|
||||
x[14] + x[15] - x[16] + x[17] + x[18] + x[19] + x[20] + x[21] - x[22] - x[23] + x[24] + x[25] - x[26] +
|
||||
x[27] + x[28] - x[29] - x[30] - x[31] - x[32] + x[33] - x[34] + x[35] - x[36] + x[37] + x[38] + x[39];
|
||||
out[6] = +x[0] + x[1] + x[2] - x[3] - x[4] + x[5] + x[6] - x[7] + x[8] + x[9] - x[10] - x[11] - x[12] - x[13] +
|
||||
x[14] - x[15] + x[16] - x[17] + x[18] + x[19] + x[20] + x[21] + x[22] - x[23] - x[24] + x[25] + x[26] -
|
||||
x[27] + x[28] + x[29] - x[30] - x[31] - x[32] - x[33] + x[34] - x[35] + x[36] - x[37] + x[38] + x[39];
|
||||
out[7] = +x[0] + x[1] + x[2] + x[3] - x[4] - x[5] + x[6] + x[7] - x[8] + x[9] + x[10] - x[11] - x[12] - x[13] -
|
||||
x[14] + x[15] - x[16] + x[17] - x[18] + x[19] + x[20] + x[21] + x[22] + x[23] - x[24] - x[25] + x[26] +
|
||||
x[27] - x[28] + x[29] + x[30] - x[31] - x[32] - x[33] - x[34] + x[35] - x[36] + x[37] - x[38] + x[39];
|
||||
out[8] = +x[0] + x[1] + x[2] + x[3] + x[4] - x[5] - x[6] + x[7] + x[8] - x[9] + x[10] + x[11] - x[12] - x[13] -
|
||||
x[14] - x[15] + x[16] - x[17] + x[18] - x[19] + x[20] + x[21] + x[22] + x[23] + x[24] - x[25] - x[26] +
|
||||
x[27] + x[28] - x[29] + x[30] + x[31] - x[32] - x[33] - x[34] - x[35] + x[36] - x[37] + x[38] - x[39];
|
||||
out[9] = +x[0] - x[1] + x[2] + x[3] + x[4] + x[5] - x[6] - x[7] + x[8] + x[9] - x[10] + x[11] + x[12] - x[13] -
|
||||
x[14] - x[15] - x[16] + x[17] - x[18] + x[19] + x[20] - x[21] + x[22] + x[23] + x[24] + x[25] - x[26] -
|
||||
x[27] + x[28] + x[29] - x[30] + x[31] + x[32] - x[33] - x[34] - x[35] - x[36] + x[37] - x[38] + x[39];
|
||||
out[10] = +x[0] + x[1] - x[2] + x[3] + x[4] + x[5] + x[6] - x[7] - x[8] + x[9] + x[10] - x[11] + x[12] + x[13] -
|
||||
x[14] - x[15] - x[16] - x[17] + x[18] - x[19] + x[20] + x[21] - x[22] + x[23] + x[24] + x[25] + x[26] -
|
||||
x[27] - x[28] + x[29] + x[30] - x[31] + x[32] + x[33] - x[34] - x[35] - x[36] - x[37] + x[38] - x[39];
|
||||
out[11] = +x[0] - x[1] + x[2] - x[3] + x[4] + x[5] + x[6] + x[7] - x[8] - x[9] + x[10] + x[11] - x[12] + x[13] +
|
||||
x[14] - x[15] - x[16] - x[17] - x[18] + x[19] + x[20] - x[21] + x[22] - x[23] + x[24] + x[25] + x[26] +
|
||||
x[27] - x[28] - x[29] + x[30] + x[31] - x[32] + x[33] + x[34] - x[35] - x[36] - x[37] - x[38] + x[39];
|
||||
out[12] = +x[0] + x[1] - x[2] + x[3] - x[4] + x[5] + x[6] + x[7] + x[8] - x[9] - x[10] + x[11] + x[12] - x[13] +
|
||||
x[14] + x[15] - x[16] - x[17] - x[18] - x[19] + x[20] + x[21] - x[22] + x[23] - x[24] + x[25] + x[26] +
|
||||
x[27] + x[28] - x[29] - x[30] + x[31] + x[32] - x[33] + x[34] + x[35] - x[36] - x[37] - x[38] - x[39];
|
||||
out[13] = +x[0] - x[1] + x[2] - x[3] + x[4] - x[5] + x[6] + x[7] + x[8] + x[9] - x[10] - x[11] + x[12] + x[13] -
|
||||
x[14] + x[15] + x[16] - x[17] - x[18] - x[19] + x[20] - x[21] + x[22] - x[23] + x[24] - x[25] + x[26] +
|
||||
x[27] + x[28] + x[29] - x[30] - x[31] + x[32] + x[33] - x[34] + x[35] + x[36] - x[37] - x[38] - x[39];
|
||||
out[14] = +x[0] - x[1] - x[2] + x[3] - x[4] + x[5] - x[6] + x[7] + x[8] + x[9] + x[10] - x[11] - x[12] + x[13] +
|
||||
x[14] - x[15] + x[16] + x[17] - x[18] - x[19] + x[20] - x[21] - x[22] + x[23] - x[24] + x[25] - x[26] +
|
||||
x[27] + x[28] + x[29] + x[30] - x[31] - x[32] + x[33] + x[34] - x[35] + x[36] + x[37] - x[38] - x[39];
|
||||
out[15] = +x[0] - x[1] - x[2] - x[3] + x[4] - x[5] + x[6] - x[7] + x[8] + x[9] + x[10] + x[11] - x[12] - x[13] +
|
||||
x[14] + x[15] - x[16] + x[17] + x[18] - x[19] + x[20] - x[21] - x[22] - x[23] + x[24] - x[25] + x[26] -
|
||||
x[27] + x[28] + x[29] + x[30] + x[31] - x[32] - x[33] + x[34] + x[35] - x[36] + x[37] + x[38] - x[39];
|
||||
out[16] = +x[0] - x[1] - x[2] - x[3] - x[4] + x[5] - x[6] + x[7] - x[8] + x[9] + x[10] + x[11] + x[12] - x[13] -
|
||||
x[14] + x[15] + x[16] - x[17] + x[18] + x[19] + x[20] - x[21] - x[22] - x[23] - x[24] + x[25] - x[26] +
|
||||
x[27] - x[28] + x[29] + x[30] + x[31] + x[32] - x[33] - x[34] + x[35] + x[36] - x[37] + x[38] + x[39];
|
||||
out[17] = +x[0] + x[1] - x[2] - x[3] - x[4] - x[5] + x[6] - x[7] + x[8] - x[9] + x[10] + x[11] + x[12] + x[13] -
|
||||
x[14] - x[15] + x[16] + x[17] - x[18] + x[19] + x[20] + x[21] - x[22] - x[23] - x[24] - x[25] + x[26] -
|
||||
x[27] + x[28] - x[29] + x[30] + x[31] + x[32] + x[33] - x[34] - x[35] + x[36] + x[37] - x[38] + x[39];
|
||||
out[18] = +x[0] + x[1] + x[2] - x[3] - x[4] - x[5] - x[6] + x[7] - x[8] + x[9] - x[10] + x[11] + x[12] + x[13] +
|
||||
x[14] - x[15] - x[16] + x[17] + x[18] - x[19] + x[20] + x[21] + x[22] - x[23] - x[24] - x[25] - x[26] +
|
||||
x[27] - x[28] + x[29] - x[30] + x[31] + x[32] + x[33] + x[34] - x[35] - x[36] + x[37] + x[38] - x[39];
|
||||
out[19] = +x[0] - x[1] + x[2] + x[3] - x[4] - x[5] - x[6] - x[7] + x[8] - x[9] + x[10] - x[11] + x[12] + x[13] +
|
||||
x[14] + x[15] - x[16] - x[17] + x[18] + x[19] + x[20] - x[21] + x[22] + x[23] - x[24] - x[25] - x[26] -
|
||||
x[27] + x[28] - x[29] + x[30] - x[31] + x[32] + x[33] + x[34] + x[35] - x[36] - x[37] + x[38] + x[39];
|
||||
out[20] = +x[0] - x[1] - x[2] - x[3] - x[4] - x[5] - x[6] - x[7] - x[8] - x[9] - x[10] - x[11] - x[12] - x[13] -
|
||||
x[14] - x[15] - x[16] - x[17] - x[18] - x[19] - x[20] + x[21] + x[22] + x[23] + x[24] + x[25] + x[26] +
|
||||
x[27] + x[28] + x[29] + x[30] + x[31] + x[32] + x[33] + x[34] + x[35] + x[36] + x[37] + x[38] + x[39];
|
||||
out[21] = +x[0] + x[1] - x[2] + x[3] + x[4] - x[5] - x[6] - x[7] - x[8] + x[9] - x[10] + x[11] - x[12] + x[13] +
|
||||
x[14] + x[15] + x[16] - x[17] - x[18] + x[19] - x[20] - x[21] + x[22] - x[23] - x[24] + x[25] + x[26] +
|
||||
x[27] + x[28] - x[29] + x[30] - x[31] + x[32] - x[33] - x[34] - x[35] - x[36] + x[37] + x[38] - x[39];
|
||||
out[22] = +x[0] + x[1] + x[2] - x[3] + x[4] + x[5] - x[6] - x[7] - x[8] - x[9] + x[10] - x[11] + x[12] - x[13] +
|
||||
x[14] + x[15] + x[16] + x[17] - x[18] - x[19] - x[20] - x[21] - x[22] + x[23] - x[24] - x[25] + x[26] +
|
||||
x[27] + x[28] + x[29] - x[30] + x[31] - x[32] + x[33] - x[34] - x[35] - x[36] - x[37] + x[38] + x[39];
|
||||
out[23] = +x[0] - x[1] + x[2] + x[3] - x[4] + x[5] + x[6] - x[7] - x[8] - x[9] - x[10] + x[11] - x[12] + x[13] -
|
||||
x[14] + x[15] + x[16] + x[17] + x[18] - x[19] - x[20] + x[21] - x[22] - x[23] + x[24] - x[25] - x[26] +
|
||||
x[27] + x[28] + x[29] + x[30] - x[31] + x[32] - x[33] + x[34] - x[35] - x[36] - x[37] - x[38] + x[39];
|
||||
out[24] = +x[0] - x[1] - x[2] + x[3] + x[4] - x[5] + x[6] + x[7] - x[8] - x[9] - x[10] - x[11] + x[12] - x[13] +
|
||||
x[14] - x[15] + x[16] + x[17] + x[18] + x[19] - x[20] + x[21] + x[22] - x[23] - x[24] + x[25] - x[26] -
|
||||
x[27] + x[28] + x[29] + x[30] + x[31] - x[32] + x[33] - x[34] + x[35] - x[36] - x[37] - x[38] - x[39];
|
||||
out[25] = +x[0] + x[1] - x[2] - x[3] + x[4] + x[5] - x[6] + x[7] + x[8] - x[9] - x[10] - x[11] - x[12] + x[13] -
|
||||
x[14] + x[15] - x[16] + x[17] + x[18] + x[19] - x[20] - x[21] + x[22] + x[23] - x[24] - x[25] + x[26] -
|
||||
x[27] - x[28] + x[29] + x[30] + x[31] + x[32] - x[33] + x[34] - x[35] + x[36] - x[37] - x[38] - x[39];
|
||||
out[26] = +x[0] + x[1] + x[2] - x[3] - x[4] + x[5] + x[6] - x[7] + x[8] + x[9] - x[10] - x[11] - x[12] - x[13] +
|
||||
x[14] - x[15] + x[16] - x[17] + x[18] + x[19] - x[20] - x[21] - x[22] + x[23] + x[24] - x[25] - x[26] +
|
||||
x[27] - x[28] - x[29] + x[30] + x[31] + x[32] + x[33] - x[34] + x[35] - x[36] + x[37] - x[38] - x[39];
|
||||
out[27] = +x[0] + x[1] + x[2] + x[3] - x[4] - x[5] + x[6] + x[7] - x[8] + x[9] + x[10] - x[11] - x[12] - x[13] -
|
||||
x[14] + x[15] - x[16] + x[17] - x[18] + x[19] - x[20] - x[21] - x[22] - x[23] + x[24] + x[25] - x[26] -
|
||||
x[27] + x[28] - x[29] - x[30] + x[31] + x[32] + x[33] + x[34] - x[35] + x[36] - x[37] + x[38] - x[39];
|
||||
out[28] = +x[0] + x[1] + x[2] + x[3] + x[4] - x[5] - x[6] + x[7] + x[8] - x[9] + x[10] + x[11] - x[12] - x[13] -
|
||||
x[14] - x[15] + x[16] - x[17] + x[18] - x[19] - x[20] - x[21] - x[22] - x[23] - x[24] + x[25] + x[26] -
|
||||
x[27] - x[28] + x[29] - x[30] - x[31] + x[32] + x[33] + x[34] + x[35] - x[36] + x[37] - x[38] + x[39];
|
||||
out[29] = +x[0] - x[1] + x[2] + x[3] + x[4] + x[5] - x[6] - x[7] + x[8] + x[9] - x[10] + x[11] + x[12] - x[13] -
|
||||
x[14] - x[15] - x[16] + x[17] - x[18] + x[19] - x[20] + x[21] - x[22] - x[23] - x[24] - x[25] + x[26] +
|
||||
x[27] - x[28] - x[29] + x[30] - x[31] - x[32] + x[33] + x[34] + x[35] + x[36] - x[37] + x[38] - x[39];
|
||||
out[30] = +x[0] + x[1] - x[2] + x[3] + x[4] + x[5] + x[6] - x[7] - x[8] + x[9] + x[10] - x[11] + x[12] + x[13] -
|
||||
x[14] - x[15] - x[16] - x[17] + x[18] - x[19] - x[20] - x[21] + x[22] - x[23] - x[24] - x[25] - x[26] +
|
||||
x[27] + x[28] - x[29] - x[30] + x[31] - x[32] - x[33] + x[34] + x[35] + x[36] + x[37] - x[38] + x[39];
|
||||
out[31] = +x[0] - x[1] + x[2] - x[3] + x[4] + x[5] + x[6] + x[7] - x[8] - x[9] + x[10] + x[11] - x[12] + x[13] +
|
||||
x[14] - x[15] - x[16] - x[17] - x[18] + x[19] - x[20] + x[21] - x[22] + x[23] - x[24] - x[25] - x[26] -
|
||||
x[27] + x[28] + x[29] - x[30] - x[31] + x[32] - x[33] - x[34] + x[35] + x[36] + x[37] + x[38] - x[39];
|
||||
out[32] = +x[0] + x[1] - x[2] + x[3] - x[4] + x[5] + x[6] + x[7] + x[8] - x[9] - x[10] + x[11] + x[12] - x[13] +
|
||||
x[14] + x[15] - x[16] - x[17] - x[18] - x[19] - x[20] - x[21] + x[22] - x[23] + x[24] - x[25] - x[26] -
|
||||
x[27] - x[28] + x[29] + x[30] - x[31] - x[32] + x[33] - x[34] - x[35] + x[36] + x[37] + x[38] + x[39];
|
||||
out[33] = +x[0] - x[1] + x[2] - x[3] + x[4] - x[5] + x[6] + x[7] + x[8] + x[9] - x[10] - x[11] + x[12] + x[13] -
|
||||
x[14] + x[15] + x[16] - x[17] - x[18] - x[19] - x[20] + x[21] - x[22] + x[23] - x[24] + x[25] - x[26] -
|
||||
x[27] - x[28] - x[29] + x[30] + x[31] - x[32] - x[33] + x[34] - x[35] - x[36] + x[37] + x[38] + x[39];
|
||||
out[34] = +x[0] - x[1] - x[2] + x[3] - x[4] + x[5] - x[6] + x[7] + x[8] + x[9] + x[10] - x[11] - x[12] + x[13] +
|
||||
x[14] - x[15] + x[16] + x[17] - x[18] - x[19] - x[20] + x[21] + x[22] - x[23] + x[24] - x[25] + x[26] -
|
||||
x[27] - x[28] - x[29] - x[30] + x[31] + x[32] - x[33] - x[34] + x[35] - x[36] - x[37] + x[38] + x[39];
|
||||
out[35] = +x[0] - x[1] - x[2] - x[3] + x[4] - x[5] + x[6] - x[7] + x[8] + x[9] + x[10] + x[11] - x[12] - x[13] +
|
||||
x[14] + x[15] - x[16] + x[17] + x[18] - x[19] - x[20] + x[21] + x[22] + x[23] - x[24] + x[25] - x[26] +
|
||||
x[27] - x[28] - x[29] - x[30] - x[31] + x[32] + x[33] - x[34] - x[35] + x[36] - x[37] - x[38] + x[39];
|
||||
out[36] = +x[0] - x[1] - x[2] - x[3] - x[4] + x[5] - x[6] + x[7] - x[8] + x[9] + x[10] + x[11] + x[12] - x[13] -
|
||||
x[14] + x[15] + x[16] - x[17] + x[18] + x[19] - x[20] + x[21] + x[22] + x[23] + x[24] - x[25] + x[26] -
|
||||
x[27] + x[28] - x[29] - x[30] - x[31] - x[32] + x[33] + x[34] - x[35] - x[36] + x[37] - x[38] - x[39];
|
||||
out[37] = +x[0] + x[1] - x[2] - x[3] - x[4] - x[5] + x[6] - x[7] + x[8] - x[9] + x[10] + x[11] + x[12] + x[13] -
|
||||
x[14] - x[15] + x[16] + x[17] - x[18] + x[19] - x[20] - x[21] + x[22] + x[23] + x[24] + x[25] - x[26] +
|
||||
x[27] - x[28] + x[29] - x[30] - x[31] - x[32] - x[33] + x[34] + x[35] - x[36] - x[37] + x[38] - x[39];
|
||||
out[38] = +x[0] + x[1] + x[2] - x[3] - x[4] - x[5] - x[6] + x[7] - x[8] + x[9] - x[10] + x[11] + x[12] + x[13] +
|
||||
x[14] - x[15] - x[16] + x[17] + x[18] - x[19] - x[20] - x[21] - x[22] + x[23] + x[24] + x[25] + x[26] -
|
||||
x[27] + x[28] - x[29] + x[30] - x[31] - x[32] - x[33] - x[34] + x[35] + x[36] - x[37] - x[38] + x[39];
|
||||
out[39] = +x[0] - x[1] + x[2] + x[3] - x[4] - x[5] - x[6] - x[7] + x[8] - x[9] + x[10] - x[11] + x[12] + x[13] +
|
||||
x[14] + x[15] - x[16] - x[17] + x[18] + x[19] - x[20] + x[21] - x[22] - x[23] + x[24] + x[25] + x[26] +
|
||||
x[27] - x[28] + x[29] - x[30] + x[31] - x[32] - x[33] - x[34] - x[35] + x[36] + x[37] - x[38] - x[39];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 40; i++) {
|
||||
x[i] = out[i];
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,482 @@
|
||||
/******************************************************************************
|
||||
* Copyright (c) 2023, Tri Dao.
|
||||
******************************************************************************/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <sgl_kernel/tensor.h>
|
||||
#include <sgl_kernel/utils.h>
|
||||
|
||||
#include <sgl_kernel/utils.cuh>
|
||||
|
||||
#include <tvm/ffi/container/tensor.h>
|
||||
|
||||
#include "fast_hadamard_transform.h"
|
||||
#include "fast_hadamard_transform_common.h"
|
||||
#include "fast_hadamard_transform_special.h"
|
||||
#include "static_switch.h"
|
||||
#include <algorithm>
|
||||
#include <cstdint>
|
||||
#include <cstring>
|
||||
|
||||
namespace {
|
||||
|
||||
using ::bf16_t;
|
||||
using ::fp16_t;
|
||||
using ::HadamardParamsBase;
|
||||
|
||||
constexpr inline int ceil_log2(int val) {
|
||||
int log = 0;
|
||||
int p = 1;
|
||||
while (p < val) {
|
||||
p <<= 1;
|
||||
++log;
|
||||
}
|
||||
return log;
|
||||
}
|
||||
|
||||
template <int kNThreads_, int kLogN_, typename input_t_>
|
||||
struct FastHadamardKernelTraits {
|
||||
using input_t = input_t_;
|
||||
static constexpr int kNThreads = kNThreads_;
|
||||
static constexpr int kLogN = kLogN_;
|
||||
static constexpr int N = 1 << kLogN;
|
||||
static constexpr int kNBytes = sizeof(input_t);
|
||||
static_assert(kNBytes == 2 || kNBytes == 4);
|
||||
static constexpr int kNElts = kNBytes == 4 ? 4 : 8;
|
||||
static constexpr int kNExchangePerVec = sizeof(float) / sizeof(input_t);
|
||||
using vec_t = typename BytesToType<kNBytes * kNElts>::Type;
|
||||
static constexpr int kNChunks = N / (kNElts * kNThreads);
|
||||
static constexpr int kSmemExchangeSize = (N * 4) < (32 * 1024) ? (N * 4) : (32 * 1024);
|
||||
static constexpr int kNExchangeRounds = N * 4 / kSmemExchangeSize;
|
||||
static_assert(kNExchangeRounds * kSmemExchangeSize == N * 4);
|
||||
static constexpr int kSmemSize = kSmemExchangeSize;
|
||||
};
|
||||
|
||||
template <int kNThreads_, int kLogN_, int kMultiple, int kMaxDim, int kMaxSmem, typename input_t_>
|
||||
struct FastHadamardMNKernelTraits {
|
||||
using input_t = input_t_;
|
||||
static constexpr int kNThreads = kNThreads_;
|
||||
static constexpr int kLogN = kLogN_;
|
||||
static constexpr int N = (1 << kLogN) * kMultiple;
|
||||
static_assert(N <= kMaxDim);
|
||||
static constexpr int kNBytes = sizeof(input_t);
|
||||
static_assert(kNBytes == 2 || kNBytes == 4);
|
||||
static constexpr int kNElts = 4;
|
||||
static constexpr int kNExchangePerVec = sizeof(float) / sizeof(input_t);
|
||||
using vec_t = typename BytesToType<kNBytes * kNElts>::Type;
|
||||
static constexpr int kNChunks = N / (kNElts * kNThreads);
|
||||
static_assert(kNChunks == kMultiple);
|
||||
static constexpr int kSmemExchangeSize = (N * 4) < kMaxSmem ? (N * 4) : kMaxSmem;
|
||||
static constexpr int kNExchangeRounds = N * 4 / kSmemExchangeSize;
|
||||
static_assert(kNExchangeRounds * kSmemExchangeSize == N * 4);
|
||||
static constexpr int kSmemSize = kSmemExchangeSize;
|
||||
};
|
||||
|
||||
template <int kNThreads_, int kLogN_, typename input_t_>
|
||||
using FastHadamard12NTraits = FastHadamardMNKernelTraits<kNThreads_, kLogN_, 12, 12 * 1024, 24 * 1024, input_t_>;
|
||||
|
||||
template <int kNThreads_, int kLogN_, typename input_t_>
|
||||
using FastHadamard20NTraits = FastHadamardMNKernelTraits<kNThreads_, kLogN_, 20, 20 * 1024, 40 * 1024, input_t_>;
|
||||
|
||||
template <int kNThreads_, int kLogN_, typename input_t_>
|
||||
using FastHadamard28NTraits = FastHadamardMNKernelTraits<kNThreads_, kLogN_, 28, 28 * 1024, 28 * 1024, input_t_>;
|
||||
|
||||
template <int kNThreads_, int kLogN_, typename input_t_>
|
||||
using FastHadamard40NTraits = FastHadamardMNKernelTraits<kNThreads_, kLogN_, 40, 40 * 1024, 40 * 1024, input_t_>;
|
||||
|
||||
template <int kNChunks>
|
||||
SGL_DEVICE void hadamard_mult_thread_chunk_12(float x[kNChunks][12]) {
|
||||
#pragma unroll
|
||||
for (int c = 0; c < kNChunks; ++c) {
|
||||
hadamard_mult_thread_12(x[c]);
|
||||
}
|
||||
}
|
||||
|
||||
template <int kNChunks>
|
||||
SGL_DEVICE void hadamard_mult_thread_chunk_20(float x[kNChunks][20]) {
|
||||
#pragma unroll
|
||||
for (int c = 0; c < kNChunks; ++c) {
|
||||
hadamard_mult_thread_20(x[c]);
|
||||
}
|
||||
}
|
||||
|
||||
template <int kNChunks>
|
||||
SGL_DEVICE void hadamard_mult_thread_chunk_28(float x[kNChunks][28]) {
|
||||
#pragma unroll
|
||||
for (int c = 0; c < kNChunks; ++c) {
|
||||
hadamard_mult_thread_28(x[c]);
|
||||
}
|
||||
}
|
||||
|
||||
template <int kNChunks>
|
||||
SGL_DEVICE void hadamard_mult_thread_chunk_40(float x[kNChunks][40]) {
|
||||
#pragma unroll
|
||||
for (int c = 0; c < kNChunks; ++c) {
|
||||
hadamard_mult_thread_40(x[c]);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename Ktraits>
|
||||
__global__ __launch_bounds__(Ktraits::kNThreads) void fast_hadamard_transform_kernel(HadamardParamsBase params) {
|
||||
constexpr int kNThreads = Ktraits::kNThreads;
|
||||
constexpr int kNElts = Ktraits::kNElts;
|
||||
constexpr int kNExchangePerVec = Ktraits::kNExchangePerVec;
|
||||
constexpr int kNChunks = Ktraits::kNChunks;
|
||||
using input_t = typename Ktraits::input_t;
|
||||
using vec_t = typename Ktraits::vec_t;
|
||||
|
||||
constexpr int kLogNElts = cilog2(Ktraits::kNElts);
|
||||
static_assert(1 << kLogNElts == kNElts, "kNElts must be a power of 2");
|
||||
|
||||
constexpr int kWarpSize = kNThreads < 32 ? kNThreads : 32;
|
||||
constexpr int kLogWarpSize = cilog2(kWarpSize);
|
||||
static_assert(1 << kLogWarpSize == kWarpSize, "Warp size must be a power of 2");
|
||||
|
||||
constexpr int kNWarps = kNThreads / kWarpSize;
|
||||
constexpr int kLogNWarps = cilog2(kNWarps);
|
||||
static_assert(1 << kLogNWarps == kNWarps, "kNWarps must be a power of 2");
|
||||
|
||||
constexpr int kChunksPerExchange = Ktraits::kSmemExchangeSize / (sizeof(vec_t) * kNExchangePerVec * kNThreads);
|
||||
static_assert(kChunksPerExchange * sizeof(vec_t) * kNExchangePerVec * kNThreads == Ktraits::kSmemExchangeSize);
|
||||
constexpr int kNExchanges = kNChunks / kChunksPerExchange;
|
||||
static_assert(kNExchanges * kChunksPerExchange == kNChunks);
|
||||
|
||||
extern __shared__ char smem_[];
|
||||
vec_t* smem_exchange = reinterpret_cast<vec_t*>(smem_);
|
||||
|
||||
const int batch_id = static_cast<int>(blockIdx.x);
|
||||
input_t* x = reinterpret_cast<input_t*>(params.x_ptr) + batch_id * params.x_batch_stride;
|
||||
input_t* out = reinterpret_cast<input_t*>(params.out_ptr) + batch_id * params.out_batch_stride;
|
||||
|
||||
float x_vals[kNChunks][kNElts];
|
||||
load_input<kNChunks, kNElts, input_t>(x, x_vals, params.dim);
|
||||
|
||||
hadamard_mult_thread<kLogNElts, kNChunks>(x_vals);
|
||||
hadamard_mult_warp<kLogWarpSize, 0, kNChunks, kNElts>(x_vals);
|
||||
|
||||
if constexpr (kNWarps > 1) {
|
||||
exchange_smem_pre<kNChunks, kChunksPerExchange, kNElts, kWarpSize, kNWarps, true, vec_t>(x_vals, smem_exchange);
|
||||
hadamard_mult_warp<kLogNWarps, 0, kNChunks, kNElts>(x_vals);
|
||||
exchange_smem_pre<kNChunks, kChunksPerExchange, kNElts, kWarpSize, kNWarps, false, vec_t>(x_vals, smem_exchange);
|
||||
}
|
||||
|
||||
if constexpr (kNChunks > 1) {
|
||||
float x_vals_transposed[kNElts][kNChunks];
|
||||
#pragma unroll
|
||||
for (int c = 0; c < kNChunks; ++c) {
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kNElts; ++i) {
|
||||
x_vals_transposed[i][c] = x_vals[c][i];
|
||||
}
|
||||
}
|
||||
|
||||
if constexpr (kNChunks == 12) {
|
||||
hadamard_mult_thread_chunk_12<kNElts>(x_vals_transposed);
|
||||
} else if constexpr (kNChunks == 20) {
|
||||
hadamard_mult_thread_chunk_20<kNElts>(x_vals_transposed);
|
||||
} else if constexpr (kNChunks == 28) {
|
||||
hadamard_mult_thread_chunk_28<kNElts>(x_vals_transposed);
|
||||
} else if constexpr (kNChunks == 40) {
|
||||
hadamard_mult_thread_chunk_40<kNElts>(x_vals_transposed);
|
||||
} else {
|
||||
constexpr int kLogNChunks = cilog2(kNChunks);
|
||||
static_assert(1 << kLogNChunks == kNChunks, "kNChunks must be a power of 2");
|
||||
hadamard_mult_thread<kLogNChunks, kNElts>(x_vals_transposed);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int c = 0; c < kNChunks; ++c) {
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kNElts; ++i) {
|
||||
x_vals[c][i] = x_vals_transposed[i][c];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
store_output<kNChunks, kNElts, input_t>(out, x_vals, params.dim, params.scale);
|
||||
}
|
||||
|
||||
template <typename Ktraits>
|
||||
inline void set_max_dynamic_smem() {
|
||||
constexpr int kSmemSize = Ktraits::kSmemSize;
|
||||
if constexpr (kSmemSize >= 48 * 1024) {
|
||||
auto kernel = &fast_hadamard_transform_kernel<Ktraits>;
|
||||
host::RuntimeDeviceCheck(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemSize));
|
||||
}
|
||||
}
|
||||
|
||||
template <typename Ktraits>
|
||||
inline void launch_kernel(HadamardParamsBase& params, DLDevice device) {
|
||||
constexpr int kSmemSize = Ktraits::kSmemSize;
|
||||
set_max_dynamic_smem<Ktraits>();
|
||||
auto kernel = &fast_hadamard_transform_kernel<Ktraits>;
|
||||
host::LaunchKernel(dim3(params.batch), dim3(Ktraits::kNThreads), device, kSmemSize)(kernel, params);
|
||||
host::RuntimeDeviceCheck();
|
||||
}
|
||||
|
||||
template <int kNThreads, int kLogN, typename input_t>
|
||||
inline void fast_hadamard_transform_launch(HadamardParamsBase& params, DLDevice device) {
|
||||
using Ktraits = FastHadamardKernelTraits<kNThreads, kLogN, input_t>;
|
||||
launch_kernel<Ktraits>(params, device);
|
||||
}
|
||||
|
||||
template <typename input_t>
|
||||
inline void fast_hadamard_transform_cuda(HadamardParamsBase& params, DLDevice device) {
|
||||
if (params.log_N == 3) {
|
||||
fast_hadamard_transform_launch<1, 3, input_t>(params, device);
|
||||
} else if (params.log_N == 4) {
|
||||
fast_hadamard_transform_launch<2, 4, input_t>(params, device);
|
||||
} else if (params.log_N == 5) {
|
||||
fast_hadamard_transform_launch<4, 5, input_t>(params, device);
|
||||
} else if (params.log_N == 6) {
|
||||
fast_hadamard_transform_launch<8, 6, input_t>(params, device);
|
||||
} else if (params.log_N == 7) {
|
||||
fast_hadamard_transform_launch<16, 7, input_t>(params, device);
|
||||
} else if (params.log_N == 8) {
|
||||
fast_hadamard_transform_launch<32, 8, input_t>(params, device);
|
||||
} else if (params.log_N == 9) {
|
||||
fast_hadamard_transform_launch<32, 9, input_t>(params, device);
|
||||
} else if (params.log_N == 10) {
|
||||
fast_hadamard_transform_launch<128, 10, input_t>(params, device);
|
||||
} else if (params.log_N == 11) {
|
||||
fast_hadamard_transform_launch<256, 11, input_t>(params, device);
|
||||
} else if (params.log_N == 12) {
|
||||
fast_hadamard_transform_launch<256, 12, input_t>(params, device);
|
||||
} else if (params.log_N == 13) {
|
||||
fast_hadamard_transform_launch<256, 13, input_t>(params, device);
|
||||
} else if (params.log_N == 14) {
|
||||
fast_hadamard_transform_launch<256, 14, input_t>(params, device);
|
||||
} else if (params.log_N == 15) {
|
||||
fast_hadamard_transform_launch<256, 15, input_t>(params, device);
|
||||
} else {
|
||||
host::Panic("fast_hadamard_transform: unsupported log_N=", params.log_N);
|
||||
}
|
||||
}
|
||||
|
||||
template <int kNThreads, int kLogN, typename input_t>
|
||||
inline void fast_hadamard_transform_12N_launch(HadamardParamsBase& params, DLDevice device) {
|
||||
using Ktraits = FastHadamard12NTraits<kNThreads, kLogN, input_t>;
|
||||
launch_kernel<Ktraits>(params, device);
|
||||
}
|
||||
|
||||
template <typename input_t>
|
||||
inline void fast_hadamard_transform_12N_cuda(HadamardParamsBase& params, DLDevice device) {
|
||||
if (params.log_N == 2) {
|
||||
fast_hadamard_transform_12N_launch<1, 2, input_t>(params, device);
|
||||
} else if (params.log_N == 3) {
|
||||
fast_hadamard_transform_12N_launch<2, 3, input_t>(params, device);
|
||||
} else if (params.log_N == 4) {
|
||||
fast_hadamard_transform_12N_launch<4, 4, input_t>(params, device);
|
||||
} else if (params.log_N == 5) {
|
||||
fast_hadamard_transform_12N_launch<8, 5, input_t>(params, device);
|
||||
} else if (params.log_N == 6) {
|
||||
fast_hadamard_transform_12N_launch<16, 6, input_t>(params, device);
|
||||
} else if (params.log_N == 7) {
|
||||
fast_hadamard_transform_12N_launch<32, 7, input_t>(params, device);
|
||||
} else if (params.log_N == 8) {
|
||||
fast_hadamard_transform_12N_launch<64, 8, input_t>(params, device);
|
||||
} else if (params.log_N == 9) {
|
||||
fast_hadamard_transform_12N_launch<128, 9, input_t>(params, device);
|
||||
} else if (params.log_N == 10) {
|
||||
fast_hadamard_transform_12N_launch<256, 10, input_t>(params, device);
|
||||
} else {
|
||||
host::Panic("fast_hadamard_transform_12N: unsupported log_N=", params.log_N);
|
||||
}
|
||||
}
|
||||
|
||||
template <int kNThreads, int kLogN, typename input_t>
|
||||
inline void fast_hadamard_transform_20N_launch(HadamardParamsBase& params, DLDevice device) {
|
||||
using Ktraits = FastHadamard20NTraits<kNThreads, kLogN, input_t>;
|
||||
launch_kernel<Ktraits>(params, device);
|
||||
}
|
||||
|
||||
template <typename input_t>
|
||||
inline void fast_hadamard_transform_20N_cuda(HadamardParamsBase& params, DLDevice device) {
|
||||
if (params.log_N == 2) {
|
||||
fast_hadamard_transform_20N_launch<1, 2, input_t>(params, device);
|
||||
} else if (params.log_N == 3) {
|
||||
fast_hadamard_transform_20N_launch<2, 3, input_t>(params, device);
|
||||
} else if (params.log_N == 4) {
|
||||
fast_hadamard_transform_20N_launch<4, 4, input_t>(params, device);
|
||||
} else if (params.log_N == 5) {
|
||||
fast_hadamard_transform_20N_launch<8, 5, input_t>(params, device);
|
||||
} else if (params.log_N == 6) {
|
||||
fast_hadamard_transform_20N_launch<16, 6, input_t>(params, device);
|
||||
} else if (params.log_N == 7) {
|
||||
fast_hadamard_transform_20N_launch<32, 7, input_t>(params, device);
|
||||
} else if (params.log_N == 8) {
|
||||
fast_hadamard_transform_20N_launch<64, 8, input_t>(params, device);
|
||||
} else if (params.log_N == 9) {
|
||||
fast_hadamard_transform_20N_launch<128, 9, input_t>(params, device);
|
||||
} else if (params.log_N == 10) {
|
||||
fast_hadamard_transform_20N_launch<256, 10, input_t>(params, device);
|
||||
} else {
|
||||
host::Panic("fast_hadamard_transform_20N: unsupported log_N=", params.log_N);
|
||||
}
|
||||
}
|
||||
|
||||
template <int kNThreads, int kLogN, typename input_t>
|
||||
inline void fast_hadamard_transform_28N_launch(HadamardParamsBase& params, DLDevice device) {
|
||||
using Ktraits = FastHadamard28NTraits<kNThreads, kLogN, input_t>;
|
||||
launch_kernel<Ktraits>(params, device);
|
||||
}
|
||||
|
||||
template <typename input_t>
|
||||
inline void fast_hadamard_transform_28N_cuda(HadamardParamsBase& params, DLDevice device) {
|
||||
if (params.log_N == 2) {
|
||||
fast_hadamard_transform_28N_launch<1, 2, input_t>(params, device);
|
||||
} else if (params.log_N == 3) {
|
||||
fast_hadamard_transform_28N_launch<2, 3, input_t>(params, device);
|
||||
} else if (params.log_N == 4) {
|
||||
fast_hadamard_transform_28N_launch<4, 4, input_t>(params, device);
|
||||
} else if (params.log_N == 5) {
|
||||
fast_hadamard_transform_28N_launch<8, 5, input_t>(params, device);
|
||||
} else if (params.log_N == 6) {
|
||||
fast_hadamard_transform_28N_launch<16, 6, input_t>(params, device);
|
||||
} else if (params.log_N == 7) {
|
||||
fast_hadamard_transform_28N_launch<32, 7, input_t>(params, device);
|
||||
} else if (params.log_N == 8) {
|
||||
fast_hadamard_transform_28N_launch<64, 8, input_t>(params, device);
|
||||
} else if (params.log_N == 9) {
|
||||
fast_hadamard_transform_28N_launch<128, 9, input_t>(params, device);
|
||||
} else if (params.log_N == 10) {
|
||||
fast_hadamard_transform_28N_launch<256, 10, input_t>(params, device);
|
||||
} else {
|
||||
host::Panic("fast_hadamard_transform_28N: unsupported log_N=", params.log_N);
|
||||
}
|
||||
}
|
||||
|
||||
template <int kNThreads, int kLogN, typename input_t>
|
||||
inline void fast_hadamard_transform_40N_launch(HadamardParamsBase& params, DLDevice device) {
|
||||
using Ktraits = FastHadamard40NTraits<kNThreads, kLogN, input_t>;
|
||||
launch_kernel<Ktraits>(params, device);
|
||||
}
|
||||
|
||||
template <typename input_t>
|
||||
inline void fast_hadamard_transform_40N_cuda(HadamardParamsBase& params, DLDevice device) {
|
||||
if (params.log_N == 2) {
|
||||
fast_hadamard_transform_40N_launch<1, 2, input_t>(params, device);
|
||||
} else if (params.log_N == 3) {
|
||||
fast_hadamard_transform_40N_launch<2, 3, input_t>(params, device);
|
||||
} else if (params.log_N == 4) {
|
||||
fast_hadamard_transform_40N_launch<4, 4, input_t>(params, device);
|
||||
} else if (params.log_N == 5) {
|
||||
fast_hadamard_transform_40N_launch<8, 5, input_t>(params, device);
|
||||
} else if (params.log_N == 6) {
|
||||
fast_hadamard_transform_40N_launch<16, 6, input_t>(params, device);
|
||||
} else if (params.log_N == 7) {
|
||||
fast_hadamard_transform_40N_launch<32, 7, input_t>(params, device);
|
||||
} else if (params.log_N == 8) {
|
||||
fast_hadamard_transform_40N_launch<64, 8, input_t>(params, device);
|
||||
} else if (params.log_N == 9) {
|
||||
fast_hadamard_transform_40N_launch<128, 9, input_t>(params, device);
|
||||
} else if (params.log_N == 10) {
|
||||
fast_hadamard_transform_40N_launch<256, 10, input_t>(params, device);
|
||||
} else {
|
||||
host::Panic("fast_hadamard_transform_40N: unsupported log_N=", params.log_N);
|
||||
}
|
||||
}
|
||||
|
||||
inline void set_hadamard_params(
|
||||
HadamardParamsBase& params,
|
||||
int64_t batch,
|
||||
int64_t dim,
|
||||
int64_t multiple,
|
||||
const tvm::ffi::TensorView x,
|
||||
const tvm::ffi::TensorView out,
|
||||
float scale) {
|
||||
std::memset(¶ms, 0, sizeof(params));
|
||||
params.batch = static_cast<int>(batch);
|
||||
params.dim = static_cast<int>(dim);
|
||||
params.log_N = ceil_log2(static_cast<int>(dim / multiple));
|
||||
params.x_ptr = const_cast<void*>(x.data_ptr());
|
||||
params.out_ptr = const_cast<void*>(out.data_ptr());
|
||||
params.x_batch_stride = x.stride(0);
|
||||
params.out_batch_stride = out.stride(0);
|
||||
params.scale = scale;
|
||||
}
|
||||
|
||||
template <int kMultiple, typename DType>
|
||||
inline void run_hadamard(const tvm::ffi::TensorView x, const tvm::ffi::TensorView out, float scale) {
|
||||
using namespace host;
|
||||
|
||||
auto N = SymbolicSize{"batch"};
|
||||
auto D = SymbolicSize{"dim"};
|
||||
auto SX = SymbolicSize{"x_batch_stride"};
|
||||
auto SO = SymbolicSize{"out_batch_stride"};
|
||||
auto device = SymbolicDevice{};
|
||||
device.set_options<kDLCUDA>();
|
||||
|
||||
TensorMatcher({N, D}).with_strides({SX, 1}).with_dtype<DType>().with_device(device).verify(x);
|
||||
TensorMatcher({N, D}).with_strides({SO, 1}).with_dtype<DType>().with_device(device).verify(out);
|
||||
|
||||
const int64_t batch = N.unwrap();
|
||||
const int64_t dim = D.unwrap();
|
||||
|
||||
RuntimeCheck(dim % kMultiple == 0, "hadamard: dim must be divisible by ", kMultiple);
|
||||
|
||||
HadamardParamsBase params;
|
||||
set_hadamard_params(params, batch, dim, kMultiple, x, out, scale);
|
||||
|
||||
if constexpr (kMultiple == 1) {
|
||||
RuntimeCheck(dim % 8 == 0, "fast_hadamard_transform only supports hidden dim divisible by 8");
|
||||
RuntimeCheck(dim <= 32768, "fast_hadamard_transform only supports hidden dim <= 32768");
|
||||
fast_hadamard_transform_cuda<DType>(params, device.unwrap());
|
||||
} else if constexpr (kMultiple == 12) {
|
||||
RuntimeCheck(dim % (4 * 12) == 0, "fast_hadamard_transform_12N only supports hidden dim divisible by 48");
|
||||
RuntimeCheck(dim <= 12 * 1024, "fast_hadamard_transform_12N only supports hidden dim <= 12288");
|
||||
fast_hadamard_transform_12N_cuda<DType>(params, device.unwrap());
|
||||
} else if constexpr (kMultiple == 20) {
|
||||
RuntimeCheck(dim % (4 * 20) == 0, "fast_hadamard_transform_20N only supports hidden dim divisible by 80");
|
||||
RuntimeCheck(dim <= 20 * 1024, "fast_hadamard_transform_20N only supports hidden dim <= 20480");
|
||||
fast_hadamard_transform_20N_cuda<DType>(params, device.unwrap());
|
||||
} else if constexpr (kMultiple == 28) {
|
||||
RuntimeCheck(dim % (4 * 28) == 0, "fast_hadamard_transform_28N only supports hidden dim divisible by 112");
|
||||
RuntimeCheck(dim <= 28 * 1024, "fast_hadamard_transform_28N only supports hidden dim <= 28672");
|
||||
fast_hadamard_transform_28N_cuda<DType>(params, device.unwrap());
|
||||
} else if constexpr (kMultiple == 40) {
|
||||
RuntimeCheck(dim % (4 * 40) == 0, "fast_hadamard_transform_40N only supports hidden dim divisible by 160");
|
||||
RuntimeCheck(dim <= 40 * 1024, "fast_hadamard_transform_40N only supports hidden dim <= 40960");
|
||||
fast_hadamard_transform_40N_cuda<DType>(params, device.unwrap());
|
||||
} else {
|
||||
Panic("Unsupported multiple");
|
||||
}
|
||||
}
|
||||
|
||||
template <typename DType>
|
||||
struct HadamardKernel {
|
||||
static void run(const tvm::ffi::TensorView x, const tvm::ffi::TensorView out, float scale) {
|
||||
run_hadamard<1, DType>(x, out, scale);
|
||||
}
|
||||
};
|
||||
|
||||
template <typename DType>
|
||||
struct Hadamard12NKernel {
|
||||
static void run(const tvm::ffi::TensorView x, const tvm::ffi::TensorView out, float scale) {
|
||||
run_hadamard<12, DType>(x, out, scale);
|
||||
}
|
||||
};
|
||||
|
||||
template <typename DType>
|
||||
struct Hadamard20NKernel {
|
||||
static void run(const tvm::ffi::TensorView x, const tvm::ffi::TensorView out, float scale) {
|
||||
run_hadamard<20, DType>(x, out, scale);
|
||||
}
|
||||
};
|
||||
|
||||
template <typename DType>
|
||||
struct Hadamard28NKernel {
|
||||
static void run(const tvm::ffi::TensorView x, const tvm::ffi::TensorView out, float scale) {
|
||||
run_hadamard<28, DType>(x, out, scale);
|
||||
}
|
||||
};
|
||||
|
||||
template <typename DType>
|
||||
struct Hadamard40NKernel {
|
||||
static void run(const tvm::ffi::TensorView x, const tvm::ffi::TensorView out, float scale) {
|
||||
run_hadamard<40, DType>(x, out, scale);
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace
|
||||
@@ -0,0 +1,27 @@
|
||||
// Copied from https://github.com/sgl-project/fast-hadamard-transform
|
||||
|
||||
// Inspired by https://github.com/NVIDIA/DALI/blob/main/include/dali/core/static_switch.h
|
||||
// and https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/Dispatch.h
|
||||
|
||||
#pragma once
|
||||
|
||||
/// @param COND - a boolean expression to switch by
|
||||
/// @param CONST_NAME - a name given for the constexpr bool variable.
|
||||
/// @param ... - code to execute for true and false
|
||||
///
|
||||
/// Usage:
|
||||
/// ```
|
||||
/// BOOL_SWITCH(flag, BoolConst, [&] {
|
||||
/// some_function<BoolConst>(...);
|
||||
/// });
|
||||
/// ```
|
||||
#define BOOL_SWITCH(COND, CONST_NAME, ...) \
|
||||
[&] { \
|
||||
if (COND) { \
|
||||
static constexpr bool CONST_NAME = true; \
|
||||
return __VA_ARGS__(); \
|
||||
} else { \
|
||||
static constexpr bool CONST_NAME = false; \
|
||||
return __VA_ARGS__(); \
|
||||
} \
|
||||
}()
|
||||
144
python/sglang/jit_kernel/hadamard.py
Normal file
144
python/sglang/jit_kernel/hadamard.py
Normal file
@@ -0,0 +1,144 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.utils import KERNEL_PATH, cache_once, load_jit, make_cpp_args
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from tvm_ffi.module import Module
|
||||
|
||||
|
||||
@cache_once
|
||||
def _jit_hadamard_module(dtype: torch.dtype) -> Module:
|
||||
args = make_cpp_args(dtype)
|
||||
hadamard_include_dir = (KERNEL_PATH / "csrc" / "fast-hadamard-transform").resolve()
|
||||
return load_jit(
|
||||
"hadamard",
|
||||
*args,
|
||||
cuda_files=["fast-hadamard-transform/hadamard_jit.cuh"],
|
||||
cuda_wrappers=[
|
||||
("hadamard_transform", f"HadamardKernel<{args}>::run"),
|
||||
("hadamard_transform_12n", f"Hadamard12NKernel<{args}>::run"),
|
||||
("hadamard_transform_20n", f"Hadamard20NKernel<{args}>::run"),
|
||||
("hadamard_transform_28n", f"Hadamard28NKernel<{args}>::run"),
|
||||
("hadamard_transform_40n", f"Hadamard40NKernel<{args}>::run"),
|
||||
],
|
||||
extra_include_paths=[str(hadamard_include_dir)],
|
||||
)
|
||||
|
||||
|
||||
def hadamard_transform(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor:
|
||||
if not x.is_cuda:
|
||||
raise RuntimeError("hadamard_transform only supports CUDA tensors")
|
||||
|
||||
shapes_og = x.size()
|
||||
dim_og = x.size(-1)
|
||||
x = x.reshape(-1, dim_og)
|
||||
if x.stride(-1) != 1:
|
||||
x = x.contiguous()
|
||||
|
||||
if dim_og % 8 != 0:
|
||||
x = torch.nn.functional.pad(x, (0, 8 - dim_og % 8))
|
||||
dim = x.size(1)
|
||||
|
||||
out = torch.empty_like(x)
|
||||
module = _jit_hadamard_module(x.dtype)
|
||||
module.hadamard_transform(x, out, scale)
|
||||
|
||||
if dim_og % 8 != 0:
|
||||
out = out[:, :dim_og]
|
||||
return out.reshape(shapes_og)
|
||||
|
||||
|
||||
def hadamard_transform_12n(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor:
|
||||
if not x.is_cuda:
|
||||
raise RuntimeError("hadamard_transform_12n only supports CUDA tensors")
|
||||
|
||||
shapes_og = x.size()
|
||||
dim_og = x.size(-1)
|
||||
x = x.reshape(-1, dim_og)
|
||||
if x.stride(-1) != 1:
|
||||
x = x.contiguous()
|
||||
|
||||
pad_multiple = 4 * 12
|
||||
if dim_og % pad_multiple != 0:
|
||||
x = torch.nn.functional.pad(x, (0, pad_multiple - dim_og % pad_multiple))
|
||||
|
||||
out = torch.empty_like(x)
|
||||
module = _jit_hadamard_module(x.dtype)
|
||||
module.hadamard_transform_12n(x, out, scale)
|
||||
|
||||
if dim_og % pad_multiple != 0:
|
||||
out = out[:, :dim_og]
|
||||
return out.reshape(shapes_og)
|
||||
|
||||
|
||||
def hadamard_transform_20n(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor:
|
||||
if not x.is_cuda:
|
||||
raise RuntimeError("hadamard_transform_20n only supports CUDA tensors")
|
||||
|
||||
shapes_og = x.size()
|
||||
dim_og = x.size(-1)
|
||||
x = x.reshape(-1, dim_og)
|
||||
if x.stride(-1) != 1:
|
||||
x = x.contiguous()
|
||||
|
||||
pad_multiple = 4 * 20
|
||||
if dim_og % pad_multiple != 0:
|
||||
x = torch.nn.functional.pad(x, (0, pad_multiple - dim_og % pad_multiple))
|
||||
|
||||
out = torch.empty_like(x)
|
||||
module = _jit_hadamard_module(x.dtype)
|
||||
module.hadamard_transform_20n(x, out, scale)
|
||||
|
||||
if dim_og % pad_multiple != 0:
|
||||
out = out[:, :dim_og]
|
||||
return out.reshape(shapes_og)
|
||||
|
||||
|
||||
def hadamard_transform_28n(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor:
|
||||
if not x.is_cuda:
|
||||
raise RuntimeError("hadamard_transform_28n only supports CUDA tensors")
|
||||
|
||||
shapes_og = x.size()
|
||||
dim_og = x.size(-1)
|
||||
x = x.reshape(-1, dim_og)
|
||||
if x.stride(-1) != 1:
|
||||
x = x.contiguous()
|
||||
|
||||
pad_multiple = 4 * 28
|
||||
if dim_og % pad_multiple != 0:
|
||||
x = torch.nn.functional.pad(x, (0, pad_multiple - dim_og % pad_multiple))
|
||||
|
||||
out = torch.empty_like(x)
|
||||
module = _jit_hadamard_module(x.dtype)
|
||||
module.hadamard_transform_28n(x, out, scale)
|
||||
|
||||
if dim_og % pad_multiple != 0:
|
||||
out = out[:, :dim_og]
|
||||
return out.reshape(shapes_og)
|
||||
|
||||
|
||||
def hadamard_transform_40n(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor:
|
||||
if not x.is_cuda:
|
||||
raise RuntimeError("hadamard_transform_40n only supports CUDA tensors")
|
||||
|
||||
shapes_og = x.size()
|
||||
dim_og = x.size(-1)
|
||||
x = x.reshape(-1, dim_og)
|
||||
if x.stride(-1) != 1:
|
||||
x = x.contiguous()
|
||||
|
||||
pad_multiple = 4 * 40
|
||||
if dim_og % pad_multiple != 0:
|
||||
x = torch.nn.functional.pad(x, (0, pad_multiple - dim_og % pad_multiple))
|
||||
|
||||
out = torch.empty_like(x)
|
||||
module = _jit_hadamard_module(x.dtype)
|
||||
module.hadamard_transform_40n(x, out, scale)
|
||||
|
||||
if dim_og % pad_multiple != 0:
|
||||
out = out[:, :dim_og]
|
||||
return out.reshape(shapes_og)
|
||||
@@ -125,7 +125,7 @@ def rotate_activation(x: torch.Tensor) -> torch.Tensor:
|
||||
if _is_hip:
|
||||
from fast_hadamard_transform import hadamard_transform
|
||||
else:
|
||||
from sgl_kernel import hadamard_transform
|
||||
from sglang.jit_kernel.hadamard import hadamard_transform
|
||||
|
||||
hidden_size = x.size(-1)
|
||||
assert (
|
||||
|
||||
@@ -106,15 +106,6 @@ FetchContent_Declare(
|
||||
)
|
||||
FetchContent_Populate(repo-mscclpp)
|
||||
|
||||
# fast-hadamard-transform
|
||||
FetchContent_Declare(
|
||||
repo-fast-hadamard-transform
|
||||
GIT_REPOSITORY https://github.com/sgl-project/fast-hadamard-transform.git
|
||||
GIT_TAG 48f3c13764dc2ec662ade842a4696a90a137f1bc
|
||||
GIT_SHALLOW OFF
|
||||
)
|
||||
FetchContent_Populate(repo-fast-hadamard-transform)
|
||||
|
||||
# ccache option
|
||||
option(ENABLE_CCACHE "Whether to use ccache" ON)
|
||||
find_program(CCACHE_FOUND ccache)
|
||||
@@ -343,9 +334,6 @@ set(SOURCES
|
||||
"${repo-flashinfer_SOURCE_DIR}/csrc/renorm.cu"
|
||||
"${repo-flashinfer_SOURCE_DIR}/csrc/sampling.cu"
|
||||
|
||||
"${repo-fast-hadamard-transform_SOURCE_DIR}/csrc/fast_hadamard_transform_cuda.cu"
|
||||
"${repo-fast-hadamard-transform_SOURCE_DIR}/csrc/fast_hadamard_transform.cpp"
|
||||
|
||||
"${repo-flash-attention_SOURCE_DIR}/csrc/flash_attn/src/flash_fwd_sparse_hdim128_bf16_causal_sm80.cu"
|
||||
"${repo-flash-attention_SOURCE_DIR}/csrc/flash_attn/src/flash_fwd_sparse_hdim128_bf16_sm80.cu"
|
||||
"${repo-flash-attention_SOURCE_DIR}/csrc/flash_attn/src/flash_fwd_sparse_hdim128_fp16_causal_sm80.cu"
|
||||
|
||||
@@ -591,24 +591,6 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
|
||||
"es_sm100_mxfp8_blockscaled_grouped_quant(Tensor input, Tensor problem_sizes, Tensor expert_offsets, Tensor "
|
||||
"blockscale_offsets, Tensor quant_output, Tensor scale_factor) -> () ");
|
||||
m.impl("es_sm100_mxfp8_blockscaled_grouped_quant", &es_sm100_mxfp8_blockscaled_grouped_quant);
|
||||
|
||||
/*
|
||||
* From fast-hadamard-transform
|
||||
*/
|
||||
m.def("fast_hadamard_transform(Tensor x, float scale) -> Tensor");
|
||||
m.impl("fast_hadamard_transform", torch::kCUDA, &fast_hadamard_transform);
|
||||
|
||||
m.def("fast_hadamard_transform_12N(Tensor x, float scale) -> Tensor");
|
||||
m.impl("fast_hadamard_transform_12N", torch::kCUDA, &fast_hadamard_transform_12N);
|
||||
|
||||
m.def("fast_hadamard_transform_20N(Tensor x, float scale) -> Tensor");
|
||||
m.impl("fast_hadamard_transform_20N", torch::kCUDA, &fast_hadamard_transform_20N);
|
||||
|
||||
m.def("fast_hadamard_transform_28N(Tensor x, float scale) -> Tensor");
|
||||
m.impl("fast_hadamard_transform_28N", torch::kCUDA, &fast_hadamard_transform_28N);
|
||||
|
||||
m.def("fast_hadamard_transform_40N(Tensor x, float scale) -> Tensor");
|
||||
m.impl("fast_hadamard_transform_40N", torch::kCUDA, &fast_hadamard_transform_40N);
|
||||
}
|
||||
|
||||
REGISTER_EXTENSION(common_ops)
|
||||
|
||||
@@ -936,15 +936,6 @@ void es_sm100_mxfp8_blockscaled_grouped_quant(
|
||||
torch::Tensor& quant_output,
|
||||
torch::Tensor& scale_factor);
|
||||
|
||||
/*
|
||||
* From fast-hadamard-transform
|
||||
*/
|
||||
torch::Tensor fast_hadamard_transform(torch::Tensor& x, double scale);
|
||||
torch::Tensor fast_hadamard_transform_12N(torch::Tensor& x, double scale);
|
||||
torch::Tensor fast_hadamard_transform_20N(torch::Tensor& x, double scale);
|
||||
torch::Tensor fast_hadamard_transform_28N(torch::Tensor& x, double scale);
|
||||
torch::Tensor fast_hadamard_transform_40N(torch::Tensor& x, double scale);
|
||||
|
||||
/*
|
||||
* From flashmla
|
||||
*/
|
||||
|
||||
@@ -65,13 +65,6 @@ from sgl_kernel.gemm import (
|
||||
silu_and_mul_scaled_fp4_grouped_quant,
|
||||
)
|
||||
from sgl_kernel.grammar import apply_token_bitmask_inplace_cuda
|
||||
from sgl_kernel.hadamard import (
|
||||
hadamard_transform,
|
||||
hadamard_transform_12n,
|
||||
hadamard_transform_20n,
|
||||
hadamard_transform_28n,
|
||||
hadamard_transform_40n,
|
||||
)
|
||||
from sgl_kernel.kvcacheio import (
|
||||
transfer_kv_all_layer,
|
||||
transfer_kv_all_layer_mla,
|
||||
|
||||
@@ -1,21 +0,0 @@
|
||||
import torch
|
||||
|
||||
|
||||
def hadamard_transform(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor:
|
||||
return torch.ops.sgl_kernel.fast_hadamard_transform.default(x, scale)
|
||||
|
||||
|
||||
def hadamard_transform_12n(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor:
|
||||
return torch.ops.sgl_kernel.fast_hadamard_transform_12N.default(x, scale)
|
||||
|
||||
|
||||
def hadamard_transform_20n(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor:
|
||||
return torch.ops.sgl_kernel.fast_hadamard_transform_20N.default(x, scale)
|
||||
|
||||
|
||||
def hadamard_transform_28n(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor:
|
||||
return torch.ops.sgl_kernel.fast_hadamard_transform_28N.default(x, scale)
|
||||
|
||||
|
||||
def hadamard_transform_40n(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor:
|
||||
return torch.ops.sgl_kernel.fast_hadamard_transform_40N.default(x, scale)
|
||||
@@ -5,7 +5,14 @@ import torch
|
||||
import torch.nn.functional as F
|
||||
from einops import rearrange, repeat
|
||||
from scipy.linalg import hadamard
|
||||
from sgl_kernel import hadamard_transform
|
||||
|
||||
try:
|
||||
from sgl_kernel import hadamard_transform
|
||||
except Exception:
|
||||
pytest.skip(
|
||||
"sgl-kernel hadamard interface was removed (migrated to jit_kernel)",
|
||||
allow_module_level=True,
|
||||
)
|
||||
|
||||
|
||||
def hadamard_transform_ref(x, scale=1.0):
|
||||
|
||||
@@ -584,8 +584,8 @@ class TestNSAIndexer(CustomTestCase):
|
||||
output = rotate_activation(x)
|
||||
self.assertEqual(output.shape, x.shape)
|
||||
self.assertEqual(output.dtype, torch.bfloat16)
|
||||
except ImportError:
|
||||
self.skipTest("sgl_kernel not available for hadamard_transform")
|
||||
except Exception:
|
||||
self.skipTest("hadamard JIT kernel not available")
|
||||
|
||||
def test_rotate_activation_invalid_size(self):
|
||||
"""Test that rotate_activation fails with non-power-of-2 size."""
|
||||
|
||||
Reference in New Issue
Block a user