#pragma once

#include <torch/torch.h>
#include <c10/util/Optional.h>

void quantize_act(torch::Tensor &Aq,
                  torch::Tensor &scale,
                  torch::Tensor const &A);

void bitnet_gemv_fused(torch::Tensor &out,
                       torch::Tensor const &a,
                       torch::Tensor const &w_packed,
                       torch::Tensor const &scale_wt);

void bitnet_gemm(torch::Tensor &out,
                 torch::Tensor const &a_int8,
                 torch::Tensor const &w_packed,
                 torch::Tensor const &scale_act,
                 torch::Tensor const &scale_wt,
                 c10::optional<torch::Tensor> const &scratch);