diff --git a/src/iluvatar/linear/kernel.h b/src/iluvatar/linear/kernel.h new file mode 100644 index 00000000..d146d0b2 --- /dev/null +++ b/src/iluvatar/linear/kernel.h @@ -0,0 +1,71 @@ +#ifndef INFINI_OPS_ILUVATAR_LINEAR_KERNEL_H_ +#define INFINI_OPS_ILUVATAR_LINEAR_KERNEL_H_ + +#include + +#include "base/linear.h" +#include "iluvatar/add/kernel.h" +#include "iluvatar/gemm/cublas.h" + +namespace infini::ops { + +template <> +class Operator : public Linear { + public: + Operator(const Tensor a, const Tensor b, std::optional bias, + bool trans_a, bool trans_b, Tensor out) + : Linear(a, b, bias, trans_a, trans_b, out), + gemm_(a, b, std::optional{1.0f}, std::optional{0.0f}, + std::optional{static_cast(trans_a)}, + std::optional{static_cast(trans_b)}, out) { + if (has_bias_) { + add_.emplace(out, BroadcastBias(*bias, out), out); + } + } + + void operator()(const Tensor a, const Tensor b, std::optional bias, + bool trans_a, bool trans_b, Tensor out) const override { + assert(has_bias_ == bias.has_value()); + + ConfigureSubOperator(gemm_); + gemm_(a, b, std::optional{1.0f}, std::optional{0.0f}, + std::optional{static_cast(trans_a)}, + std::optional{static_cast(trans_b)}, out); + + if (has_bias_) { + auto bias_view = BroadcastBias(*bias, out); + ConfigureSubOperator(*add_); + (*add_)(out, bias_view, out); + } + } + + private: + static Tensor BroadcastBias(const Tensor& bias, const Tensor& out) { + assert(bias.ndim() == 1); + assert(bias.size(0) == out.size(-1)); + + auto shape = out.shape(); + Tensor::Strides strides(out.ndim(), 0); + strides.back() = bias.stride(0); + + return Tensor{const_cast(bias.data()), shape, bias.dtype(), + bias.device(), strides}; + } + + template + void ConfigureSubOperator(Op& op) const { + op.set_handle(handle_); + op.set_config(config_); + op.set_stream(stream_); + op.set_workspace(workspace_); + op.set_workspace_size_in_bytes(workspace_size_in_bytes_); + } + + mutable Operator gemm_; + + mutable std::optional> add_; +}; + +} // namespace infini::ops + +#endif diff --git a/src/metax/linear/kernel.h b/src/metax/linear/kernel.h new file mode 100644 index 00000000..3418fc30 --- /dev/null +++ b/src/metax/linear/kernel.h @@ -0,0 +1,71 @@ +#ifndef INFINI_OPS_METAX_LINEAR_KERNEL_H_ +#define INFINI_OPS_METAX_LINEAR_KERNEL_H_ + +#include + +#include "base/linear.h" +#include "metax/add/kernel.h" +#include "metax/gemm/mcblas.h" + +namespace infini::ops { + +template <> +class Operator : public Linear { + public: + Operator(const Tensor a, const Tensor b, std::optional bias, + bool trans_a, bool trans_b, Tensor out) + : Linear(a, b, bias, trans_a, trans_b, out), + gemm_(a, b, std::optional{1.0f}, std::optional{0.0f}, + std::optional{static_cast(trans_a)}, + std::optional{static_cast(trans_b)}, out) { + if (has_bias_) { + add_.emplace(out, BroadcastBias(*bias, out), out); + } + } + + void operator()(const Tensor a, const Tensor b, std::optional bias, + bool trans_a, bool trans_b, Tensor out) const override { + assert(has_bias_ == bias.has_value()); + + ConfigureSubOperator(gemm_); + gemm_(a, b, std::optional{1.0f}, std::optional{0.0f}, + std::optional{static_cast(trans_a)}, + std::optional{static_cast(trans_b)}, out); + + if (has_bias_) { + auto bias_view = BroadcastBias(*bias, out); + ConfigureSubOperator(*add_); + (*add_)(out, bias_view, out); + } + } + + private: + static Tensor BroadcastBias(const Tensor& bias, const Tensor& out) { + assert(bias.ndim() == 1); + assert(bias.size(0) == out.size(-1)); + + auto shape = out.shape(); + Tensor::Strides strides(out.ndim(), 0); + strides.back() = bias.stride(0); + + return Tensor{const_cast(bias.data()), shape, bias.dtype(), + bias.device(), strides}; + } + + template + void ConfigureSubOperator(Op& op) const { + op.set_handle(handle_); + op.set_config(config_); + op.set_stream(stream_); + op.set_workspace(workspace_); + op.set_workspace_size_in_bytes(workspace_size_in_bytes_); + } + + mutable Operator gemm_; + + mutable std::optional> add_; +}; + +} // namespace infini::ops + +#endif diff --git a/src/moore/linear/kernel.h b/src/moore/linear/kernel.h new file mode 100644 index 00000000..38b85782 --- /dev/null +++ b/src/moore/linear/kernel.h @@ -0,0 +1,71 @@ +#ifndef INFINI_OPS_MOORE_LINEAR_KERNEL_H_ +#define INFINI_OPS_MOORE_LINEAR_KERNEL_H_ + +#include + +#include "base/linear.h" +#include "moore/add/kernel.h" +#include "moore/gemm/mublas.h" + +namespace infini::ops { + +template <> +class Operator : public Linear { + public: + Operator(const Tensor a, const Tensor b, std::optional bias, + bool trans_a, bool trans_b, Tensor out) + : Linear(a, b, bias, trans_a, trans_b, out), + gemm_(a, b, std::optional{1.0f}, std::optional{0.0f}, + std::optional{static_cast(trans_a)}, + std::optional{static_cast(trans_b)}, out) { + if (has_bias_) { + add_.emplace(out, BroadcastBias(*bias, out), out); + } + } + + void operator()(const Tensor a, const Tensor b, std::optional bias, + bool trans_a, bool trans_b, Tensor out) const override { + assert(has_bias_ == bias.has_value()); + + ConfigureSubOperator(gemm_); + gemm_(a, b, std::optional{1.0f}, std::optional{0.0f}, + std::optional{static_cast(trans_a)}, + std::optional{static_cast(trans_b)}, out); + + if (has_bias_) { + auto bias_view = BroadcastBias(*bias, out); + ConfigureSubOperator(*add_); + (*add_)(out, bias_view, out); + } + } + + private: + static Tensor BroadcastBias(const Tensor& bias, const Tensor& out) { + assert(bias.ndim() == 1); + assert(bias.size(0) == out.size(-1)); + + auto shape = out.shape(); + Tensor::Strides strides(out.ndim(), 0); + strides.back() = bias.stride(0); + + return Tensor{const_cast(bias.data()), shape, bias.dtype(), + bias.device(), strides}; + } + + template + void ConfigureSubOperator(Op& op) const { + op.set_handle(handle_); + op.set_config(config_); + op.set_stream(stream_); + op.set_workspace(workspace_); + op.set_workspace_size_in_bytes(workspace_size_in_bytes_); + } + + mutable Operator gemm_; + + mutable std::optional> add_; +}; + +} // namespace infini::ops + +#endif diff --git a/src/nvidia/linear/kernel.h b/src/nvidia/linear/kernel.h new file mode 100644 index 00000000..67c3baf4 --- /dev/null +++ b/src/nvidia/linear/kernel.h @@ -0,0 +1,71 @@ +#ifndef INFINI_OPS_NVIDIA_LINEAR_KERNEL_H_ +#define INFINI_OPS_NVIDIA_LINEAR_KERNEL_H_ + +#include + +#include "base/linear.h" +#include "nvidia/add/kernel.h" +#include "nvidia/gemm/cublas.h" + +namespace infini::ops { + +template <> +class Operator : public Linear { + public: + Operator(const Tensor a, const Tensor b, std::optional bias, + bool trans_a, bool trans_b, Tensor out) + : Linear(a, b, bias, trans_a, trans_b, out), + gemm_(a, b, std::optional{1.0f}, std::optional{0.0f}, + std::optional{static_cast(trans_a)}, + std::optional{static_cast(trans_b)}, out) { + if (has_bias_) { + add_.emplace(out, BroadcastBias(*bias, out), out); + } + } + + void operator()(const Tensor a, const Tensor b, std::optional bias, + bool trans_a, bool trans_b, Tensor out) const override { + assert(has_bias_ == bias.has_value()); + + ConfigureSubOperator(gemm_); + gemm_(a, b, std::optional{1.0f}, std::optional{0.0f}, + std::optional{static_cast(trans_a)}, + std::optional{static_cast(trans_b)}, out); + + if (has_bias_) { + auto bias_view = BroadcastBias(*bias, out); + ConfigureSubOperator(*add_); + (*add_)(out, bias_view, out); + } + } + + private: + static Tensor BroadcastBias(const Tensor& bias, const Tensor& out) { + assert(bias.ndim() == 1); + assert(bias.size(0) == out.size(-1)); + + auto shape = out.shape(); + Tensor::Strides strides(out.ndim(), 0); + strides.back() = bias.stride(0); + + return Tensor{const_cast(bias.data()), shape, bias.dtype(), + bias.device(), strides}; + } + + template + void ConfigureSubOperator(Op& op) const { + op.set_handle(handle_); + op.set_config(config_); + op.set_stream(stream_); + op.set_workspace(workspace_); + op.set_workspace_size_in_bytes(workspace_size_in_bytes_); + } + + mutable Operator gemm_; + + mutable std::optional> add_; +}; + +} // namespace infini::ops + +#endif