Assert: registry.instances.find(key) == registry.instances.end()

halide_test.cpp

#include "Halide.h"

#include <cstdlib>
#include <vector>
#include <algorithm>
#include <iostream>
#include <memory>
#include <random>

using namespace Halide;

namespace {
    constexpr auto seed = 1234;
    std::default_random_engine engine(seed);
    std::uniform_real_distribution<float> uniform(-1.f, 1.f);
    auto generator = [] { return uniform(engine); };
}

template <typename ElementType>
struct ILayer {
    virtual void SetInput(Func) = 0;
    virtual Func OutputFunc() const = 0;
    virtual void Apply(int, int, std::vector<ElementType>&) const = 0;
    virtual void Apply() const = 0;
    virtual void Compile() = 0;
    virtual const std::vector<ElementType>& GetOutput() const = 0;
};

template <typename ElementType>
struct Group : ILayer<ElementType>
{
    Group()
    {
        _inputData = ImageParam(type_of<ElementType>(), 2, "Group_InputData");
        _input = Func("Group_Input");

        Var i("i"), j("j");
        _input(i, j) = _inputData(i, j);
    }

    void AddLayer(std::unique_ptr<ILayer<ElementType>> layer)
    {
        if (_layers.empty())
        {
            layer->SetInput(_input);
        }
        else
        {
            layer->SetInput(_layers.back()->OutputFunc());
        }
        _layers.emplace_back(std::move(layer));

        _output = _layers.back()->OutputFunc();
    }

    // ILayer
    void SetInput(Func f) override
    {
        Halide::Var i, j;
        _input = Halide::Func("Group_Input");

        _input(i, j) = f(i, j);

        Halide::Func newInput = _input;
        for (auto& layer : _layers)
        {
            layer->SetInput(newInput);
            newInput = layer->OutputFunc();
        }
    }

    Func OutputFunc() const override
    {
        return _output;
    }

    void Apply(int rows, int cols, std::vector<ElementType>& A) const override
    {
        Buffer<ElementType> buffer(A.data(), rows, cols);

        _inputData.set(buffer);

        Apply();
    }

    void Apply() const override
    {
        _layers.back()->Apply();
    }

    void Compile() override
    {
        _layers.back()->Compile();
    }

    const std::vector<ElementType>& GetOutput() const override
    {
        return _layers.back()->GetOutput();
    }

private:
    std::vector<std::unique_ptr<ILayer<ElementType>>> _layers;
    Func _input;
    Func _output;
    mutable ImageParam _inputData;
};

template <typename ElementType>
struct MatMulLayer : ILayer<ElementType> {
    bool Initialize(int m, int k, int n)
    {
        _B = Buffer<ElementType>(k, n);
        _B.fill(generator);

        _m = m; _k = k; _n = n;

        _C = std::vector<ElementType>(m * n);
        _outputBuffer = Buffer<ElementType>(_C.data(), _m, _n);
        _A = ImageParam(type_of<ElementType>(), 2, "MatMulLayer_InputBuffer");

        Func A("MatMulLayer_A");
        Var x("x"), y("y");
        A(x, y) = _A(x, y);

        CreateMainFunc(A);

        return true;
    }

    // ILayer
    void SetInput(Func f) override
	{
        CreateMainFunc(f);
	}

    Func OutputFunc() const override
	{
        return C;
	}

    void Apply(int rows, int cols, std::vector<ElementType>& A) const override
	{
        _A.set(Buffer<ElementType>(A.data(), rows, cols));

        Apply();
	}

    void Apply() const override
	{
        Realization realization(_outputBuffer);
        C.realize(realization);
	}

    void Compile() override
	{
        C.compile_jit();
	}

    const std::vector<ElementType>& GetOutput() const override
	{
        return _C;
	}

private:
    void CreateMainFunc(Func f)
    {
        Var x("x"), y("y");
        Func A;
        if (f.name() == "MatMulLayer_A")
        {
            A = f;
        }
        else
        {
            A = Func("MatMulLayer_A");
            A(x,y) = f(x, y);
        }

        RDom k(0, _k, "k");

        C = Func("MatMulLayer_C");
        C(x, y) = sum(A(x, k) * _B(k, y));

        C
            .bound(x, 0, _m)
            .bound(y, 0, _n)
            ;
    }

    std::vector<ElementType> _C;


    int _m, _k, _n;
    mutable ImageParam _A;
    Buffer<ElementType> _B;
    mutable Buffer<ElementType> _outputBuffer;
    mutable Func C;
};

template <typename LayerT>
void test(int m, int k, int n, int iterations)
{
    Group<float> group;

    auto layer = std::unique_ptr<LayerT>(new LayerT);
    layer->Initialize(m, k, n);
    group.AddLayer(std::move(layer));
    group.Compile();

    std::vector<std::vector<float>> vecA(iterations);
    for (auto& A : vecA)
    {
        A.resize(m * k);
        std::generate(A.begin(), A.end(), generator);
    }

    for (auto& A : vecA)
    {
        group.Apply(m, k, A);
        volatile const auto& output = group.GetOutput();
        (void)output;
    }
}

int main(int argc, char **argv) {
    while (true)
    {
        test<MatMulLayer<float>>(10, 12, 14, 1000);
        test<MatMulLayer<float>>(11, 12, 14, 1000);
        test<MatMulLayer<float>>(10, 11, 14, 1000);
        test<MatMulLayer<float>>(10, 12, 12, 1000);
        test<MatMulLayer<float>>(10, 10, 12, 1000);
        test<MatMulLayer<float>>(10, 10, 10, 1000);

        test<MatMulLayer<float>>(15, 12, 14, 1000);
        test<MatMulLayer<float>>(15, 12, 14, 1000);
        test<MatMulLayer<float>>(15, 15, 14, 1000);
        test<MatMulLayer<float>>(15, 12, 12, 1000);
        test<MatMulLayer<float>>(15, 15, 12, 1000);
        test<MatMulLayer<float>>(15, 15, 15, 1000);

        test<MatMulLayer<float>>(9, 12, 14, 1000);
        test<MatMulLayer<float>>(9, 12, 14, 1000);
        test<MatMulLayer<float>>(9, 9, 14, 1000);
        test<MatMulLayer<float>>(9, 12, 12, 1000);
        test<MatMulLayer<float>>(9, 9, 12, 1000);
        test<MatMulLayer<float>>(9, 9, 9, 1000);

        test<MatMulLayer<float>>(11, 12, 14, 1000);
        test<MatMulLayer<float>>(11, 12, 14, 1000);
        test<MatMulLayer<float>>(11, 11, 14, 1000);
        test<MatMulLayer<float>>(11, 12, 12, 1000);
        test<MatMulLayer<float>>(11, 11, 12, 1000);
        test<MatMulLayer<float>>(11, 11, 11, 1000);

        test<MatMulLayer<float>>(13, 12, 14, 1000);
        test<MatMulLayer<float>>(13, 12, 14, 1000);
        test<MatMulLayer<float>>(13, 13, 14, 1000);
        test<MatMulLayer<float>>(13, 12, 12, 1000);
        test<MatMulLayer<float>>(13, 13, 12, 1000);
        test<MatMulLayer<float>>(13, 13, 13, 1000);

        test<MatMulLayer<float>>(9, 12, 14, 1000);
        test<MatMulLayer<float>>(9, 12, 14, 1000);
        test<MatMulLayer<float>>(9, 9, 14, 1000);
        test<MatMulLayer<float>>(9, 12, 12, 1000);
        test<MatMulLayer<float>>(9, 9, 12, 1000);
        test<MatMulLayer<float>>(9, 9, 9, 1000);


        test<MatMulLayer<float>>(14, 12, 14, 1000);
        test<MatMulLayer<float>>(14, 12, 14, 1000);
        test<MatMulLayer<float>>(14, 14, 14, 1000);
        test<MatMulLayer<float>>(14, 12, 12, 1000);
        test<MatMulLayer<float>>(14, 14, 12, 1000);
        test<MatMulLayer<float>>(14, 14, 14, 1000);


        test<MatMulLayer<float>>(16, 12, 14, 1000);
        test<MatMulLayer<float>>(16, 12, 14, 1000);
        test<MatMulLayer<float>>(16, 16, 14, 1000);
        test<MatMulLayer<float>>(16, 12, 12, 1000);
        test<MatMulLayer<float>>(16, 16, 12, 1000);
        test<MatMulLayer<float>>(16, 16, 16, 1000);


        test<MatMulLayer<float>>(18, 12, 14, 1000);
        test<MatMulLayer<float>>(18, 12, 14, 1000);
        test<MatMulLayer<float>>(18, 18, 14, 1000);
        test<MatMulLayer<float>>(18, 12, 12, 1000);
        test<MatMulLayer<float>>(18, 18, 12, 1000);
        test<MatMulLayer<float>>(18, 18, 18, 1000);

    }

    return 0;
}