【C#】多次元配列と Span

多次元配列（int[,]のような配列）は、標準でAsSpan()が用意されていない。そのため通常では、配列全体を連続した領域のデータとみなしたSpanとして扱えない。

列挙以外に全体を連続として扱える方法は C# ではない……？

MemoryMarshal.CreateSpanを使う方法がある。
参考：多次元配列の AsSpan · GitHub

public static Span<T> AsSpan<T>(this T[,] array) => MemoryMarshal.CreateSpan(ref array[0, 0], array.Length);

C# は行優先なので、[0, 0], [0, 1], ..., [1, 0], [1, 1], ... と整列したSpanを得ることができる。

（※ なお少し調べた限り、必ず連続した整列なのか確証がないのは念のため触れておく。こういう環境でダメ、という具体的な情報はなかったので触れる程度で。）

今回、比較用の例として、以下の記事を参考とさせていただいた。
[C#]Span<T>構造体で高速かつ安全に配列へアクセスして、行列積を素早く計算する。

ジャグ配列はオミットし、二次元配列の ① 普通のループ、② 事前に転置、③ インデックスアクセスを減、そして ④ スパン利用を較べてみる。

①・②については BenchmarkDotNet で測る用に調整したものの、基本的に上記記事内容のとおりなので以下折りたたみ。

折りたたみ

public int N = 2022;
public int D = 2022;
public int M = 2022;

public double[,] left2D = default!;
public double[,] right2D = default!;

[GlobalSetup]
public void Setup()
{
    left2D = new double[N, D];
    right2D = new double[D, M];

    for (int i = 0; i < N; i++)
    {
        for (int k = 0; k < D; k++)
        {
            left2D[i, k] = i + k;
        }
    }
    for (int k = 0; k < D; k++)
    {
        for (int j = 0; j < M; j++)
        {
            right2D[k, j] = k + j + 1;
        }
    }
}

[Benchmark]
public double[,] ByRegular()
{
    var result = new double[N, M];
    for (int i = 0; i < N; i++)
    {
        for (int j = 0; j < M; j++)
        {
            for (int k = 0; k < D; k++)
            {
                result[i, j] += left2D[i, k] * right2D[k, j];
            }
        }
    }

    return result;
}

[Benchmark]
public double[,] ByTransposed()
{
    var transposed = new double[M, D];
    for (int j = 0; j < M; j++)
    {
        for (int k = 0; k < D; k++)
        {
            transposed[j, k] = right2D[k, j];
        }
    }

    var result = new double[N, M];
    for (int i = 0; i < N; i++)
    {
        for (int j = 0; j < M; j++)
        {
            for (int k = 0; k < D; k++)
            {
                result[i, j] += left2D[i, k] * transposed[j, k];
            }
        }
    }

    return result;
}

③については、result[i, j]へ k 回アクセスしていたので変数を置いて1回のみとしたもの。
影響を見たいもののSpanと関係ないので別段階の枠として、折りたたみ。

折りたたみ

[Benchmark]
public double[,] BySum()
{
    var transposed = new double[M, D];
    for (int j = 0; j < M; j++)
    {
        for (int k = 0; k < D; k++)
        {
            transposed[j, k] = right2D[k, j];
        }
    }

    var result = new double[N, M];
    for (int i = 0; i < N; i++)
    {
        for (int j = 0; j < M; j++)
        {
            var sum = 0d;
            for (int k = 0; k < D; k++)
            {
                // result[i, j] += left2D[i, k] * transposed[j, k];
                sum += left2D[i, k] * transposed[j, k];
            }
            result[i, j] = sum;
        }
    }

    return result;
}

④、Spanを用いたもの。
上述の記事だとジャグ配列の配列のAsSpan()を行っているが、ここでは違ったかたちで扱う。

private static Span<T> AsSpan<T>(T[,] span) => MemoryMarshal.CreateSpan(ref span[0, 0], span.Length);

[Benchmark]
public double[,] By1DSpan()
{
    var transposed = ArrayPool<double>.Shared.Rent(M * D);
    Span<double> tSpan = transposed.AsSpan();
    try
    {
        for (int k = 0; k < M; k++)
        {
            for (int j = 0; j < D; j++)
            {
                tSpan[(j * D) + k] = right2D[k, j];
            }
        }

        var result = new double[N, M];
        var lSpan = AsSpan(left2D);
        for (int i = 0; i < N; i++)
        {
            for (int j = 0; j < M; j++)
            {
                var sum = 0d;
                var lPart = lSpan.Slice(i * N, D);
                var tPart = tSpan.Slice(j * M, D);
                for(var k = 0; k < D; k++)
                {
                    sum += lPart[k] * tPart[k];
                }
                // result は二次元配列のインデックスアクセスで妥協
                result[i, j] = sum;
            }
        }

        return result;
    }
    finally
    {
        ArrayPool<double>.Shared.Return(transposed);
    }
}

       Method |      Mean |    Error |   StdDev | Allocated |
------------- |----------:|---------:|---------:|----------:|
    ByRegular | 117.189 s |  4.910 s | 0.2691 s |  31.19 MB |
 ByTransposed |  31.871 s | 10.289 s | 0.5640 s |  62.39 MB |
        BySum |  19.351 s | 12.898 s | 0.7070 s |  62.39 MB |
     By1DSpan |   7.772 s |  3.306 s | 0.1812 s |  31.19 MB |

正直かなり違って驚いた。
（一応SequenceEqualsで結果の整合は取ったので変な出力にはなってない、はず）

実施前の見込みでは、裏でよしなに SIMD してくれたりするAPIを使っていなかったり、Spanを通したとはいえ結局は列単位のアクセスと同等だからメリットが薄そうに感じたりと、実はそこまで違いはないんじゃないかと思っていた。
とすると、Spanの最適化やら計算機科学的なメリットやらが今回の例では大きかったということになる。

Span強し。

多次元配列は基本Spanで扱いたい、と思ったけれど、前述のとおりちょっと確かでない部分があるので、とりあえず引き出し。

てくメモ

trivial な notes

【C#】多次元配列と Span