test/spmm-microkernel-tester.h - platform/external/XNNPACK - Git at Google

 // Copyright 2019 Google LLC
 //
 // This source code is licensed under the BSD-style license found in the
 // LICENSE file in the root directory of this source tree.

 #pragma once

 #include <gtest/gtest.h>

 #include <algorithm>
 #include <cassert>
 #include <cmath>
 #include <cstddef>
 #include <cstdlib>
 #include <functional>
 #include <random>
 #include <vector>

 #include <xnnpack.h>
 #include <xnnpack/AlignedAllocator.h>
 #include <xnnpack/params-init.h>
 #include <xnnpack/params.h>


 static inline bool is_fp16_zero(uint16_t x) {
   const uint32_t ext_x = x;
   const uint32_t two_x = ext_x + ext_x;
   return (uint16_t) two_x == 0;
 }

 class SpMMMicrokernelTester {
  public:
   enum class Variant {
     Native,
     Scalar,
   };

   inline SpMMMicrokernelTester& mr(size_t mr) {
     this->mr_ = mr;
     return *this;
   }

   inline size_t mr() const {
     return this->mr_;
   }

   inline SpMMMicrokernelTester& nr(size_t nr) {
     this->nr_ = nr;
     return *this;
   }

   inline size_t nr() const {
     return this->nr_;
   }

   inline SpMMMicrokernelTester& m(size_t m) {
     this->m_ = m;
     return *this;
   }

   inline size_t m() const {
     return this->m_;
   }

   inline SpMMMicrokernelTester& n(size_t n) {
     this->n_ = n;
     return *this;
   }

   inline size_t n() const {
     return this->n_;
   }

   inline SpMMMicrokernelTester& k(size_t k) {
     this->k_ = k;
     return *this;
   }

   inline size_t k() const {
     return this->k_;
   }

   inline SpMMMicrokernelTester& sparsity(float sparsity) {
     this->sparsity_ = sparsity;
     return *this;
   }

   inline float sparsity() const {
     return this->sparsity_;
   }

   inline SpMMMicrokernelTester& qmin(uint8_t qmin) {
     this->qmin_ = qmin;
     return *this;
   }

   inline uint8_t qmin() const {
     return this->qmin_;
   }

   inline SpMMMicrokernelTester& qmax(uint8_t qmax) {
     this->qmax_ = qmax;
     return *this;
   }

   inline uint8_t qmax() const {
     return this->qmax_;
   }

   inline SpMMMicrokernelTester& iterations(size_t iterations) {
     this->iterations_ = iterations;
     return *this;
   }

   inline size_t iterations() const {
     return this->iterations_;
   }

   void Test(xnn_f32_spmm_ukernel_function spmm, Variant variant = Variant::Native) const {
     ASSERT_GE(m(), 1);
     ASSERT_GE(n(), 1);
     ASSERT_GE(k(), 1);

     std::random_device random_device;
     auto rng = std::mt19937(random_device());
     auto f32rng = std::bind(std::uniform_real_distribution<float>(), rng);
     auto prng = std::bind(std::uniform_real_distribution<float>(), rng);

     std::vector<float, AlignedAllocator<float, 64>> a(k() * m());
     // Think of b as (n/nr + n % nr) x k, expansion happens later.
     const size_t ncols = n() / nr() + n() % nr();
     std::vector<float> b(ncols * k());
     std::vector<float> bias(n());
     // Number of non-zero weights per N (output channel).
     std::vector<uint32_t> nmap(n());
     // Mapping from index of non-zero weight to increment of K (input channel) following this index.
     std::vector<int32_t> dmap(n() * k());
     std::vector<float> w(n() * k() + n());
     std::vector<float> c(n() * m());
     std::vector<float> c_ref(n() * m());

     for (size_t iteration = 0; iteration < iterations(); iteration++) {
       std::generate(a.begin(), a.end(), std::ref(f32rng));
       std::generate(b.begin(), b.end(), std::ref(f32rng));
       std::generate(bias.begin(), bias.end(), std::ref(f32rng));
       std::fill(c.begin(), c.end(), nanf(""));
       std::fill(c_ref.begin(), c_ref.end(), 0.0f);
       std::fill(nmap.begin(), nmap.end(), 0);
       std::fill(dmap.begin(), dmap.end(), 0);
       std::fill(w.begin(), w.end(), 0.0f);

       for (float& b_value : b) {
         if (prng() <= sparsity()) {
           b_value = 0.0f;
         }
       }

       uint32_t nnz = 0;
       uint32_t wcnt = 0;
       size_t last_kk = 0;
       bool first_nzz = true;
       size_t first_kk = 0;
       for (size_t nn = 0; nn < n() / nr(); nn++) {
         for (size_t i = 0; i < nr(); ++i)
           w[wcnt++] = bias[nr() * nn + i];
         for (size_t kk = 0; kk < k(); kk++) {
           if (b[nn * k() + kk] != 0.0f) {
             // Every non-zero actually corresponds to nr adjacent non-zeros.
             for (size_t i = 0; i < nr(); ++i)
               w[wcnt++] = b[nn * k() + kk] + static_cast<float>(i);
             // Skip the very first non-zero weight as we record only the difference.
             if (first_nzz) {
               first_kk = kk;
             } else {
               const int32_t increment = int32_t(kk - last_kk) * int32_t(m() * sizeof(float));
               dmap[nnz++] = increment;
             }
             last_kk = kk;
             first_nzz = false;
             nmap[nn] += 1;
           }
         }
       }

       // now we've constructed the matrix for the blocked part and switch to the
       // leftovers, which we do as nr=1 always.
       for (size_t nn = n() / nr(); nn < ncols; nn++) {
         w[wcnt++] = bias[(n() / nr()) * nr() + (nn - n() / nr())];
         for (size_t kk = 0; kk < k(); kk++) {
           if (b[nn * k() + kk] != 0.0f) {
             // Every non-zero actually corresponds to nr adjacent non-zeros.
             w[wcnt++] = b[nn * k() + kk];
             // Skip the very first non-zero weight as we record only the difference.
             if (first_nzz) {
               first_kk = kk;
             } else {
               const int32_t increment = int32_t(kk - last_kk) * int32_t(m() * sizeof(float));
               dmap[nnz++] = increment;
             }
             last_kk = kk;
             first_nzz = false;
             nmap[nn] += 1;
           }
         }
       }
       // In the end, we must return input pointer to the initial value.
       const int64_t increment = int32_t(first_kk - last_kk) * int32_t(m() * sizeof(float));
       dmap[nnz++] = increment;

       // Generate expanded b which will be used in reference calculation.
       // Everywhere there is a non-zero in the original we copy it and add an
       // adjacent non-zero with incremented weight value.
       std::vector<float> b_full(n() * k());
       if (nr() == 1) {
          b_full = b;
       }
       else {
         for (size_t nn = 0; nn < n() / nr(); nn++) {
           for (size_t kk = 0; kk < k(); kk++) {
             if (b[nn * k() + kk] != 0.0f) {
               for (size_t i = 0; i < nr(); ++i)
                 b_full[nr() * nn * k() + i * k() + kk] = b[nn * k() + kk] + static_cast<float>(i);
             }
           }
         }
         for (size_t nn = n() / nr(); nn < ncols; nn++) {
           for (size_t kk = 0; kk < k(); kk++) {
             if (b[nn * k() + kk] != 0.0f) {
               b_full[nr() * (n() / nr()) * k() + (nn - n() / nr()) * k() + kk] = b[nn * k() + kk];
             }
           }
         }
       }

       for (size_t oc = 0; oc < n(); oc++) {
         for (size_t pxb = 0; pxb < m(); pxb++) {
           c_ref[oc * m() + pxb] = bias[oc];
           for (size_t ic = 0; ic < k(); ic++) {
             c_ref[oc * m() + pxb] += a[ic * m() + pxb] * b_full[oc * k() + ic];
           }
         }
       }

       // Micro-kernel can access one element beyond w and dmap for software pipelining.
       w.resize(wcnt + 1);
       dmap.resize(nnz + 1);

       // Compute clamping parameters.
       const float accumulated_min = *std::min_element(c_ref.cbegin(), c_ref.cend());
       const float accumulated_max = *std::max_element(c_ref.cbegin(), c_ref.cend());
       const float c_min = accumulated_min + (accumulated_max - accumulated_min) / 255.0f * float(qmin());
       const float c_max = accumulated_max - (accumulated_max - accumulated_min) / 255.0f * float(255 - qmax());

       // Clamp reference results.
       for (float& c_value : c_ref) {
         c_value = std::min(std::max(c_value, c_min), c_max);
       }

       // Prepare output parameters.
       xnn_f32_output_params output_params = { };
       switch (variant) {
         case Variant::Native:
           output_params = xnn_init_f32_output_params(c_min, c_max);
           break;
         case Variant::Scalar:
           output_params = xnn_init_scalar_f32_output_params(c_min, c_max);
           break;
       }

       spmm(m(), n(),
         a.data() + first_kk * m(), w.data(), dmap.data(), nmap.data(), c.data(),
         &output_params);

       // Validate micro-kernel outputs.
       for (size_t pxb = 0; pxb < n(); pxb++) {
         for (size_t oc = 0; oc < m(); oc++) {
           ASSERT_NEAR(
             c[pxb * m() + oc],
             c_ref[pxb * m() + oc],
             std::abs(c_ref[pxb * m() + oc]) * 1.0e-6f)
             << "at " << pxb << ", " << oc
             << ": Mr x Nr x Kr = " << mr() << " x " << nr()
             << ", M x N x K = " << m() << " x " << n() << " x " << k();
         }
       }
     }
   }

   void Test(xnn_f16_spmm_ukernel_function spmm) const {
     ASSERT_GE(m(), 1);
     ASSERT_GE(n(), 1);
     ASSERT_GE(k(), 1);

     std::random_device random_device;
     auto rng = std::mt19937(random_device());
     auto f32rng = std::bind(std::uniform_real_distribution<float>(), rng);
     auto f16rng = std::bind(fp16_ieee_from_fp32_value, f32rng);
     auto prng = std::bind(std::uniform_real_distribution<float>(), rng);

     std::vector<uint16_t, AlignedAllocator<uint16_t, 64>> a(k() * m());
     // Think of b as (n/nr + n % nr) x k, expansion happens later.
     const size_t ncols = n() / nr() + n() % nr();
     std::vector<uint16_t> b(ncols * k());
     std::vector<uint16_t> bias(n());
     // Number of non-zero weights per N (output channel).
     std::vector<uint32_t> nmap(n());
     // Mapping from index of non-zero weight to increment of K (input channel) following this index.
     std::vector<int32_t> dmap(n() * k());
     std::vector<uint16_t> w(n() * k() + n());
     std::vector<uint16_t> c(n() * m());
     std::vector<float> c_ref(n() * m());

     for (size_t iteration = 0; iteration < iterations(); iteration++) {
       std::generate(a.begin(), a.end(), std::ref(f16rng));
       std::generate(b.begin(), b.end(), std::ref(f16rng));
       std::generate(bias.begin(), bias.end(), std::ref(f16rng));
       std::fill(c.begin(), c.end(), 0xC000);
       std::fill(c_ref.begin(), c_ref.end(), 0.0f);
       std::fill(nmap.begin(), nmap.end(), 0);
       std::fill(dmap.begin(), dmap.end(), 0);
       std::fill(w.begin(), w.end(), 0);

       for (uint16_t& b_value : b) {
         if (prng() <= sparsity()) {
           b_value = 0;
         }
       }

       uint32_t nnz = 0;
       uint32_t wcnt = 0;
       size_t last_kk = 0;
       bool first_nzz = true;
       size_t first_kk = 0;
       for (size_t nn = 0; nn < n() / nr(); nn++) {
         for (size_t i = 0; i < nr(); ++i)
           w[wcnt++] = bias[nr() * nn + i];
         for (size_t kk = 0; kk < k(); kk++) {
           if (!is_fp16_zero(b[nn * k() + kk])) {
             // Every non-zero actually corresponds to nr adjacent non-zeros.
             for (size_t i = 0; i < nr(); ++i)
               w[wcnt++] = fp16_ieee_from_fp32_value(fp16_ieee_to_fp32_value(b[nn * k() + kk]) + static_cast<float>(i));
             // Skip the very first non-zero weight as we record only the difference.
             if (first_nzz) {
               first_kk = kk;
             } else {
               const int32_t increment = int32_t(kk - last_kk) * int32_t(m() * sizeof(uint16_t));
               dmap[nnz++] = increment;
             }
             last_kk = kk;
             first_nzz = false;
             nmap[nn] += 1;
           }
         }
       }

       // now we've constructed the matrix for the blocked part and switch to the
       // leftovers, which we do as nr=1 always.
       for (size_t nn = n() / nr(); nn < ncols; nn++) {
         w[wcnt++] = bias[(n() / nr()) * nr() + (nn - n() / nr())];
         for (size_t kk = 0; kk < k(); kk++) {
           if (!is_fp16_zero(b[nn * k() + kk])) {
             // Every non-zero actually corresponds to nr adjacent non-zeros.
             w[wcnt++] = b[nn * k() + kk];
             // Skip the very first non-zero weight as we record only the difference.
             if (first_nzz) {
               first_kk = kk;
             } else {
               const int32_t increment = int32_t(kk - last_kk) * int32_t(m() * sizeof(uint16_t));
               dmap[nnz++] = increment;
             }
             last_kk = kk;
             first_nzz = false;
             nmap[nn] += 1;
           }
         }
       }
       // In the end, we must return input pointer to the initial value.
       const int64_t increment = int32_t(first_kk - last_kk) * int32_t(m() * sizeof(uint16_t));
       dmap[nnz++] = increment;

       // Generate expanded b which will be used in reference calculation.
       // Everywhere there is a non-zero in the original we copy it and add an
       // adjacent non-zero with incremented weight value.
       std::vector<uint16_t> b_full(n() * k());
       if (nr() == 1) {
          b_full = b;
       }
       else {
         for (size_t nn = 0; nn < n() / nr(); nn++) {
           for (size_t kk = 0; kk < k(); kk++) {
             if (b[nn * k() + kk] != 0.0f) {
               for (size_t i = 0; i < nr(); ++i)
                 b_full[nr() * nn * k() + i * k() + kk] = fp16_ieee_from_fp32_value(
                   fp16_ieee_to_fp32_value(b[nn * k() + kk]) + static_cast<float>(i));
             }
           }
         }
         for (size_t nn = n() / nr(); nn < ncols; nn++) {
           for (size_t kk = 0; kk < k(); kk++) {
             if (b[nn * k() + kk] != 0.0f) {
               b_full[nr() * (n() / nr()) * k() + (nn - n() / nr()) * k() + kk] = b[nn * k() + kk];
             }
           }
         }
       }

       for (size_t oc = 0; oc < n(); oc++) {
         for (size_t pxb = 0; pxb < m(); pxb++) {
           c_ref[oc * m() + pxb] = fp16_ieee_to_fp32_value(bias[oc]);
           for (size_t ic = 0; ic < k(); ic++) {
             c_ref[oc * m() + pxb] += fp16_ieee_to_fp32_value(a[ic * m() + pxb]) * fp16_ieee_to_fp32_value(b_full[oc * k() + ic]);
           }
         }
       }

       // Micro-kernel can access one element beyond w and dmap for software pipelining.
       w.resize(wcnt + 1);
       dmap.resize(nnz + 1);

       // Compute clamping parameters.
       const float accumulated_min = *std::min_element(c_ref.cbegin(), c_ref.cend());
       const float accumulated_max = *std::max_element(c_ref.cbegin(), c_ref.cend());
       const float c_min = accumulated_min + (accumulated_max - accumulated_min) / 255.0f * float(qmin());
       const float c_max = accumulated_max - (accumulated_max - accumulated_min) / 255.0f * float(255 - qmax());

       // Clamp reference results.
       for (float& c_value : c_ref) {
         c_value = std::min(std::max(c_value, c_min), c_max);
       }

       // Prepare output parameters.
       xnn_f16_output_params output_params;
       output_params.scale = UINT16_C(0x3C00) /* 1.0 */;
       output_params.max = fp16_ieee_from_fp32_value(c_max);
       output_params.min = fp16_ieee_from_fp32_value(c_min);

       spmm(m(), n(),
         a.data() + first_kk * m(), w.data(), dmap.data(), nmap.data(), c.data(),
         &output_params);

       // Validate micro-kernel outputs.
       for (size_t pxb = 0; pxb < n(); pxb++) {
         for (size_t oc = 0; oc < m(); oc++) {
           ASSERT_NEAR(
             fp16_ieee_to_fp32_value(c[pxb * m() + oc]),
             c_ref[pxb * m() + oc],
             std::abs(c_ref[pxb * m() + oc]) * 1.0e-2f)
             << "at " << pxb << ", " << oc
             << ": Mr x Nr x Kr = " << mr() << " x " << nr()
             << ", M x N x K = " << m() << " x " << n() << " x " << k();
         }
       }
     }
   }

  private:
   size_t mr_{1};
   size_t nr_{1};
   size_t m_{1};
   size_t n_{1};
   size_t k_{1};
   float sparsity_{0.5f};
   uint8_t qmin_{0};
   uint8_t qmax_{255};
   size_t iterations_{1};
 };
	// Copyright 2019 Google LLC
	//
	// This source code is licensed under the BSD-style license found in the
	// LICENSE file in the root directory of this source tree.

	#pragma once

	#include <gtest/gtest.h>

	#include <algorithm>
	#include <cassert>
	#include <cmath>
	#include <cstddef>
	#include <cstdlib>
	#include <functional>
	#include <random>
	#include <vector>

	#include <xnnpack.h>
	#include <xnnpack/AlignedAllocator.h>
	#include <xnnpack/params-init.h>
	#include <xnnpack/params.h>


	static inline bool is_fp16_zero(uint16_t x) {
	const uint32_t ext_x = x;
	const uint32_t two_x = ext_x + ext_x;
	return (uint16_t) two_x == 0;
	}

	class SpMMMicrokernelTester {
	public:
	enum class Variant {
	Native,
	Scalar,
	};

	inline SpMMMicrokernelTester& mr(size_t mr) {
	this->mr_ = mr;
	return *this;
	}

	inline size_t mr() const {
	return this->mr_;
	}

	inline SpMMMicrokernelTester& nr(size_t nr) {
	this->nr_ = nr;
	return *this;
	}

	inline size_t nr() const {
	return this->nr_;
	}

	inline SpMMMicrokernelTester& m(size_t m) {
	this->m_ = m;
	return *this;
	}

	inline size_t m() const {
	return this->m_;
	}

	inline SpMMMicrokernelTester& n(size_t n) {
	this->n_ = n;
	return *this;
	}

	inline size_t n() const {
	return this->n_;
	}

	inline SpMMMicrokernelTester& k(size_t k) {
	this->k_ = k;
	return *this;
	}

	inline size_t k() const {
	return this->k_;
	}

	inline SpMMMicrokernelTester& sparsity(float sparsity) {
	this->sparsity_ = sparsity;
	return *this;
	}

	inline float sparsity() const {
	return this->sparsity_;
	}

	inline SpMMMicrokernelTester& qmin(uint8_t qmin) {
	this->qmin_ = qmin;
	return *this;
	}

	inline uint8_t qmin() const {
	return this->qmin_;
	}

	inline SpMMMicrokernelTester& qmax(uint8_t qmax) {
	this->qmax_ = qmax;
	return *this;
	}

	inline uint8_t qmax() const {
	return this->qmax_;
	}

	inline SpMMMicrokernelTester& iterations(size_t iterations) {
	this->iterations_ = iterations;
	return *this;
	}

	inline size_t iterations() const {
	return this->iterations_;
	}

	void Test(xnn_f32_spmm_ukernel_function spmm, Variant variant = Variant::Native) const {
	ASSERT_GE(m(), 1);
	ASSERT_GE(n(), 1);
	ASSERT_GE(k(), 1);

	std::random_device random_device;
	auto rng = std::mt19937(random_device());
	auto f32rng = std::bind(std::uniform_real_distribution<float>(), rng);
	auto prng = std::bind(std::uniform_real_distribution<float>(), rng);

	std::vector<float, AlignedAllocator<float, 64>> a(k() * m());
	// Think of b as (n/nr + n % nr) x k, expansion happens later.
	const size_t ncols = n() / nr() + n() % nr();
	std::vector<float> b(ncols * k());
	std::vector<float> bias(n());
	// Number of non-zero weights per N (output channel).
	std::vector<uint32_t> nmap(n());
	// Mapping from index of non-zero weight to increment of K (input channel) following this index.
	std::vector<int32_t> dmap(n() * k());
	std::vector<float> w(n() * k() + n());
	std::vector<float> c(n() * m());
	std::vector<float> c_ref(n() * m());

	for (size_t iteration = 0; iteration < iterations(); iteration++) {
	std::generate(a.begin(), a.end(), std::ref(f32rng));
	std::generate(b.begin(), b.end(), std::ref(f32rng));
	std::generate(bias.begin(), bias.end(), std::ref(f32rng));
	std::fill(c.begin(), c.end(), nanf(""));
	std::fill(c_ref.begin(), c_ref.end(), 0.0f);
	std::fill(nmap.begin(), nmap.end(), 0);
	std::fill(dmap.begin(), dmap.end(), 0);
	std::fill(w.begin(), w.end(), 0.0f);

	for (float& b_value : b) {
	if (prng() <= sparsity()) {
	b_value = 0.0f;
	}
	}

	uint32_t nnz = 0;
	uint32_t wcnt = 0;
	size_t last_kk = 0;
	bool first_nzz = true;
	size_t first_kk = 0;
	for (size_t nn = 0; nn < n() / nr(); nn++) {
	for (size_t i = 0; i < nr(); ++i)
	w[wcnt++] = bias[nr() * nn + i];
	for (size_t kk = 0; kk < k(); kk++) {
	if (b[nn * k() + kk] != 0.0f) {
	// Every non-zero actually corresponds to nr adjacent non-zeros.
	for (size_t i = 0; i < nr(); ++i)
	w[wcnt++] = b[nn * k() + kk] + static_cast<float>(i);
	// Skip the very first non-zero weight as we record only the difference.
	if (first_nzz) {
	first_kk = kk;
	} else {
	const int32_t increment = int32_t(kk - last_kk) * int32_t(m() * sizeof(float));
	dmap[nnz++] = increment;
	}
	last_kk = kk;
	first_nzz = false;
	nmap[nn] += 1;
	}
	}
	}

	// now we've constructed the matrix for the blocked part and switch to the
	// leftovers, which we do as nr=1 always.
	for (size_t nn = n() / nr(); nn < ncols; nn++) {
	w[wcnt++] = bias[(n() / nr()) * nr() + (nn - n() / nr())];
	for (size_t kk = 0; kk < k(); kk++) {
	if (b[nn * k() + kk] != 0.0f) {
	// Every non-zero actually corresponds to nr adjacent non-zeros.
	w[wcnt++] = b[nn * k() + kk];
	// Skip the very first non-zero weight as we record only the difference.
	if (first_nzz) {
	first_kk = kk;
	} else {
	const int32_t increment = int32_t(kk - last_kk) * int32_t(m() * sizeof(float));
	dmap[nnz++] = increment;
	}
	last_kk = kk;
	first_nzz = false;
	nmap[nn] += 1;
	}
	}
	}
	// In the end, we must return input pointer to the initial value.
	const int64_t increment = int32_t(first_kk - last_kk) * int32_t(m() * sizeof(float));
	dmap[nnz++] = increment;

	// Generate expanded b which will be used in reference calculation.
	// Everywhere there is a non-zero in the original we copy it and add an
	// adjacent non-zero with incremented weight value.
	std::vector<float> b_full(n() * k());
	if (nr() == 1) {
	b_full = b;
	}
	else {
	for (size_t nn = 0; nn < n() / nr(); nn++) {
	for (size_t kk = 0; kk < k(); kk++) {
	if (b[nn * k() + kk] != 0.0f) {
	for (size_t i = 0; i < nr(); ++i)
	b_full[nr() * nn * k() + i * k() + kk] = b[nn * k() + kk] + static_cast<float>(i);
	}
	}
	}
	for (size_t nn = n() / nr(); nn < ncols; nn++) {
	for (size_t kk = 0; kk < k(); kk++) {
	if (b[nn * k() + kk] != 0.0f) {
	b_full[nr() * (n() / nr()) * k() + (nn - n() / nr()) * k() + kk] = b[nn * k() + kk];
	}
	}
	}
	}

	for (size_t oc = 0; oc < n(); oc++) {
	for (size_t pxb = 0; pxb < m(); pxb++) {
	c_ref[oc * m() + pxb] = bias[oc];
	for (size_t ic = 0; ic < k(); ic++) {
	c_ref[oc * m() + pxb] += a[ic * m() + pxb] * b_full[oc * k() + ic];
	}
	}
	}

	// Micro-kernel can access one element beyond w and dmap for software pipelining.
	w.resize(wcnt + 1);
	dmap.resize(nnz + 1);

	// Compute clamping parameters.
	const float accumulated_min = *std::min_element(c_ref.cbegin(), c_ref.cend());
	const float accumulated_max = *std::max_element(c_ref.cbegin(), c_ref.cend());
	const float c_min = accumulated_min + (accumulated_max - accumulated_min) / 255.0f * float(qmin());
	const float c_max = accumulated_max - (accumulated_max - accumulated_min) / 255.0f * float(255 - qmax());

	// Clamp reference results.
	for (float& c_value : c_ref) {
	c_value = std::min(std::max(c_value, c_min), c_max);
	}

	// Prepare output parameters.
	xnn_f32_output_params output_params = { };
	switch (variant) {
	case Variant::Native:
	output_params = xnn_init_f32_output_params(c_min, c_max);
	break;
	case Variant::Scalar:
	output_params = xnn_init_scalar_f32_output_params(c_min, c_max);
	break;
	}

	spmm(m(), n(),
	a.data() + first_kk * m(), w.data(), dmap.data(), nmap.data(), c.data(),
	&output_params);

	// Validate micro-kernel outputs.
	for (size_t pxb = 0; pxb < n(); pxb++) {
	for (size_t oc = 0; oc < m(); oc++) {
	ASSERT_NEAR(
	c[pxb * m() + oc],
	c_ref[pxb * m() + oc],
	std::abs(c_ref[pxb * m() + oc]) * 1.0e-6f)
	<< "at " << pxb << ", " << oc
	<< ": Mr x Nr x Kr = " << mr() << " x " << nr()
	<< ", M x N x K = " << m() << " x " << n() << " x " << k();
	}
	}
	}
	}

	void Test(xnn_f16_spmm_ukernel_function spmm) const {
	ASSERT_GE(m(), 1);
	ASSERT_GE(n(), 1);
	ASSERT_GE(k(), 1);

	std::random_device random_device;
	auto rng = std::mt19937(random_device());
	auto f32rng = std::bind(std::uniform_real_distribution<float>(), rng);
	auto f16rng = std::bind(fp16_ieee_from_fp32_value, f32rng);
	auto prng = std::bind(std::uniform_real_distribution<float>(), rng);

	std::vector<uint16_t, AlignedAllocator<uint16_t, 64>> a(k() * m());
	// Think of b as (n/nr + n % nr) x k, expansion happens later.
	const size_t ncols = n() / nr() + n() % nr();
	std::vector<uint16_t> b(ncols * k());
	std::vector<uint16_t> bias(n());
	// Number of non-zero weights per N (output channel).
	std::vector<uint32_t> nmap(n());
	// Mapping from index of non-zero weight to increment of K (input channel) following this index.
	std::vector<int32_t> dmap(n() * k());
	std::vector<uint16_t> w(n() * k() + n());
	std::vector<uint16_t> c(n() * m());
	std::vector<float> c_ref(n() * m());

	for (size_t iteration = 0; iteration < iterations(); iteration++) {
	std::generate(a.begin(), a.end(), std::ref(f16rng));
	std::generate(b.begin(), b.end(), std::ref(f16rng));
	std::generate(bias.begin(), bias.end(), std::ref(f16rng));
	std::fill(c.begin(), c.end(), 0xC000);
	std::fill(c_ref.begin(), c_ref.end(), 0.0f);
	std::fill(nmap.begin(), nmap.end(), 0);
	std::fill(dmap.begin(), dmap.end(), 0);
	std::fill(w.begin(), w.end(), 0);

	for (uint16_t& b_value : b) {
	if (prng() <= sparsity()) {
	b_value = 0;
	}
	}

	uint32_t nnz = 0;
	uint32_t wcnt = 0;
	size_t last_kk = 0;
	bool first_nzz = true;
	size_t first_kk = 0;
	for (size_t nn = 0; nn < n() / nr(); nn++) {
	for (size_t i = 0; i < nr(); ++i)
	w[wcnt++] = bias[nr() * nn + i];
	for (size_t kk = 0; kk < k(); kk++) {
	if (!is_fp16_zero(b[nn * k() + kk])) {
	// Every non-zero actually corresponds to nr adjacent non-zeros.
	for (size_t i = 0; i < nr(); ++i)
	w[wcnt++] = fp16_ieee_from_fp32_value(fp16_ieee_to_fp32_value(b[nn * k() + kk]) + static_cast<float>(i));
	// Skip the very first non-zero weight as we record only the difference.
	if (first_nzz) {
	first_kk = kk;
	} else {
	const int32_t increment = int32_t(kk - last_kk) * int32_t(m() * sizeof(uint16_t));
	dmap[nnz++] = increment;
	}
	last_kk = kk;
	first_nzz = false;
	nmap[nn] += 1;
	}
	}
	}

	// now we've constructed the matrix for the blocked part and switch to the
	// leftovers, which we do as nr=1 always.
	for (size_t nn = n() / nr(); nn < ncols; nn++) {
	w[wcnt++] = bias[(n() / nr()) * nr() + (nn - n() / nr())];
	for (size_t kk = 0; kk < k(); kk++) {
	if (!is_fp16_zero(b[nn * k() + kk])) {
	// Every non-zero actually corresponds to nr adjacent non-zeros.
	w[wcnt++] = b[nn * k() + kk];
	// Skip the very first non-zero weight as we record only the difference.
	if (first_nzz) {
	first_kk = kk;
	} else {
	const int32_t increment = int32_t(kk - last_kk) * int32_t(m() * sizeof(uint16_t));
	dmap[nnz++] = increment;
	}
	last_kk = kk;
	first_nzz = false;
	nmap[nn] += 1;
	}
	}
	}
	// In the end, we must return input pointer to the initial value.
	const int64_t increment = int32_t(first_kk - last_kk) * int32_t(m() * sizeof(uint16_t));
	dmap[nnz++] = increment;

	// Generate expanded b which will be used in reference calculation.
	// Everywhere there is a non-zero in the original we copy it and add an
	// adjacent non-zero with incremented weight value.
	std::vector<uint16_t> b_full(n() * k());
	if (nr() == 1) {
	b_full = b;
	}
	else {
	for (size_t nn = 0; nn < n() / nr(); nn++) {
	for (size_t kk = 0; kk < k(); kk++) {
	if (b[nn * k() + kk] != 0.0f) {
	for (size_t i = 0; i < nr(); ++i)
	b_full[nr() * nn * k() + i * k() + kk] = fp16_ieee_from_fp32_value(
	fp16_ieee_to_fp32_value(b[nn * k() + kk]) + static_cast<float>(i));
	}
	}
	}
	for (size_t nn = n() / nr(); nn < ncols; nn++) {
	for (size_t kk = 0; kk < k(); kk++) {
	if (b[nn * k() + kk] != 0.0f) {
	b_full[nr() * (n() / nr()) * k() + (nn - n() / nr()) * k() + kk] = b[nn * k() + kk];
	}
	}
	}
	}

	for (size_t oc = 0; oc < n(); oc++) {
	for (size_t pxb = 0; pxb < m(); pxb++) {
	c_ref[oc * m() + pxb] = fp16_ieee_to_fp32_value(bias[oc]);
	for (size_t ic = 0; ic < k(); ic++) {
	c_ref[oc * m() + pxb] += fp16_ieee_to_fp32_value(a[ic * m() + pxb]) * fp16_ieee_to_fp32_value(b_full[oc * k() + ic]);
	}
	}
	}

	// Micro-kernel can access one element beyond w and dmap for software pipelining.
	w.resize(wcnt + 1);
	dmap.resize(nnz + 1);

	// Compute clamping parameters.
	const float accumulated_min = *std::min_element(c_ref.cbegin(), c_ref.cend());
	const float accumulated_max = *std::max_element(c_ref.cbegin(), c_ref.cend());
	const float c_min = accumulated_min + (accumulated_max - accumulated_min) / 255.0f * float(qmin());
	const float c_max = accumulated_max - (accumulated_max - accumulated_min) / 255.0f * float(255 - qmax());

	// Clamp reference results.
	for (float& c_value : c_ref) {
	c_value = std::min(std::max(c_value, c_min), c_max);
	}

	// Prepare output parameters.
	xnn_f16_output_params output_params;
	output_params.scale = UINT16_C(0x3C00) /* 1.0 */;
	output_params.max = fp16_ieee_from_fp32_value(c_max);
	output_params.min = fp16_ieee_from_fp32_value(c_min);

	spmm(m(), n(),
	a.data() + first_kk * m(), w.data(), dmap.data(), nmap.data(), c.data(),
	&output_params);

	// Validate micro-kernel outputs.
	for (size_t pxb = 0; pxb < n(); pxb++) {
	for (size_t oc = 0; oc < m(); oc++) {
	ASSERT_NEAR(
	fp16_ieee_to_fp32_value(c[pxb * m() + oc]),
	c_ref[pxb * m() + oc],
	std::abs(c_ref[pxb * m() + oc]) * 1.0e-2f)
	<< "at " << pxb << ", " << oc
	<< ": Mr x Nr x Kr = " << mr() << " x " << nr()
	<< ", M x N x K = " << m() << " x " << n() << " x " << k();
	}
	}
	}
	}

	private:
	size_t mr_{1};
	size_t nr_{1};
	size_t m_{1};
	size_t n_{1};
	size_t k_{1};
	float sparsity_{0.5f};
	uint8_t qmin_{0};
	uint8_t qmax_{255};
	size_t iterations_{1};
	};