ircd::math: Add softmax functions.

include/ircd/math/math.h

@@ -19,3 +19,4 @@
 #include "exp.h"
 #include "fmma.h"
 #include "norm.h"
+#include "smax.h"

include/ircd/math/smax.h

@@ -0,0 +1,166 @@
+// Matrix Construct
+//
+// Copyright (C) Matrix Construct Developers, Authors & Contributors
+// Copyright (C) 2016-2021 Jason Volk <jason@zemos.net>
+//
+// Permission to use, copy, modify, and/or distribute this software for any
+// purpose with or without fee is hereby granted, provided that the above
+// copyright notice and this permission notice is present in all copies. The
+// full license for this software is available in the LICENSE file.
+
+#pragma once
+#define HAVE_IRCD_MATH_SMAX_H
+
+namespace ircd::math
+{
+	template<class T,
+	         class D = T>
+	typename std::enable_if<!simd::is<T>(), void>::type
+	smax(const vector_view<T> &out,
+	     const vector_view<const T> &in,
+	     const vector_view<D> &tmp0,
+	     const vector_view<D> &tmp1);
+
+	template<class T,
+	         class D = T>
+	typename std::enable_if<!simd::is<T>(), void>::type
+	smax(const vector_view<T> &out,
+	     const vector_view<const T> &in,
+	     const vector_view<D> &tmp);
+
+	template<class T,
+	         class D = T>
+	typename std::enable_if<simd::is<T>(), void>::type
+	smax(const vector_view<T> &out,
+	     const vector_view<const T> &in,
+	     const vector_view<D> &tmp,
+	     const vector_view<D> &tmp1);
+
+	template<class T,
+	         class D = T>
+	typename std::enable_if<simd::is<T>(), void>::type
+	smax(const vector_view<T> &out,
+	     const vector_view<const T> &in,
+	     const vector_view<D> &tmp);
+}
+
+/// Softmax
+template<class T,
+         class D>
+inline typename std::enable_if<ircd::simd::is<T>(), void>::type
+ircd::math::smax(const vector_view<T> &__restrict__ out,
+                 const vector_view<const T> &__restrict__ in,
+                 const vector_view<D> &__restrict__ acc)
+{
+	assert(out.size() >= in.size());
+	assert(acc.size() >= in.size());
+
+	for(uint i(0); i < size(in); ++i)
+		acc[i] = {0};
+
+	for(uint i(0); i < size(in); ++i)
+	{
+		const auto precise
+		{
+			exp(lane_cast<D>(in[i]))
+		};
+
+		for(uint j(0); j < size(in); ++j)
+			for(uint k(0); k < simd::lanes<D>(); ++k)
+				acc[j] += precise[k];
+	}
+
+	for(uint i(0); i < size(in); ++i)
+	{
+		const auto precise
+		{
+			exp(lane_cast<D>(in[i]))
+		};
+
+		out[i] = lane_cast<T>(precise / acc[i]);
+	}
+}
+
+/// Softmax
+template<class T,
+         class D>
+inline typename std::enable_if<ircd::simd::is<T>(), void>::type
+ircd::math::smax(const vector_view<T> &__restrict__ out,
+                 const vector_view<const T> &__restrict__ in,
+                 const vector_view<D> &__restrict__ acc,
+                 const vector_view<D> &__restrict__ exps)
+{
+	assert(out.size() >= in.size());
+	assert(acc.size() >= in.size());
+	assert(exps.size() >= in.size());
+
+	for(uint i(0); i < size(in); ++i)
+		acc[i] = {0};
+
+	for(uint i(0); i < size(in); ++i)
+		exps[i] = exp(lane_cast<D>(in[i]));
+
+	for(uint i(0); i < size(in); ++i)
+		for(uint j(0); j < size(in); ++j)
+			for(uint k(0); k < simd::lanes<D>(); ++k)
+				acc[j] += exps[i][k];
+
+	for(uint i(0); i < size(in); ++i)
+		out[i] = lane_cast<T>(exps[i] / acc[i]);
+}
+
+/// Softmax
+template<class T,
+         class D>
+inline typename std::enable_if<!ircd::simd::is<T>(), void>::type
+ircd::math::smax(const vector_view<T> &__restrict__ out,
+                 const vector_view<const T> &__restrict__ in,
+                 const vector_view<D> &__restrict__ acc)
+{
+	using ircd::exp;
+
+	assert(out.size() >= in.size());
+	assert(acc.size() >= in.size());
+
+	for(uint i(0); i < size(in); ++i)
+		acc[i] = {0};
+
+	for(uint i(0); i < size(in); ++i)
+		out[i] = exp(in[i]);
+
+	for(uint i(0); i < size(in); ++i)
+		for(uint j(0); j < size(in); ++j)
+			acc[j] += out[i];
+
+	for(uint i(0); i < size(in); ++i)
+		out[i] = out[i] / acc[i];
+}
+
+/// Softmax
+template<class T,
+         class D>
+inline typename std::enable_if<!ircd::simd::is<T>(), void>::type
+ircd::math::smax(const vector_view<T> &__restrict__ out,
+                 const vector_view<const T> &__restrict__ in,
+                 const vector_view<D> &__restrict__ acc,
+                 const vector_view<D> &__restrict__ exps)
+{
+	using ircd::exp;
+
+	assert(out.size() >= in.size());
+	assert(acc.size() >= in.size());
+	assert(exps.size() >= in.size());
+
+	for(uint i(0); i < size(in); ++i)
+		acc[i] = {0};
+
+	for(uint i(0); i < size(in); ++i)
+		exps[i] = exp(in[i]);
+
+	for(uint i(0); i < size(in); ++i)
+		for(uint j(0); j < size(in); ++j)
+			acc[j] += exps[i];
+
+	for(uint i(0); i < size(in); ++i)
+		out[i] = exps[i] / acc[i];
+}