Início / coding / Perguntas / 79558025
Accepted
Olibarer
Asked: 2025-04-06 17:31:27 +0800 CST

Maneira eficiente e legível de obter uma matriz de índice N-dimensional em ordem C usando NumPy

  • 772

Quando preciso gerar uma matriz de índices N-dimensional em ordem C, tento algumas abordagens diferentes do NumPy.

O mais rápido para matrizes maiores, mas menos legível:

np.stack(np.meshgrid(*[np.arange(i, dtype=dtype) for i in sizes], indexing="ij"), axis=-1).reshape(-1, len(sizes))

Mais legível com bom desempenho:

np.ascontiguousarray(np.indices(sizes, dtype=dtype).reshape(len(sizes), -1).T)

Aqui, não tenho certeza se a cópia ascontiguousarray é realmente necessária ou se há uma maneira melhor de garantir que o resultado esteja na ordem C sem forçar uma cópia.

Mais legível, mas de longe o mais lento:

np.vstack([*np.ndindex(sizes)], dtype=dtype)

A conversão do iterador é bastante lenta para matrizes maiores.

Existe uma maneira integrada mais direta e legível de fazer isso que corresponda ao desempenho de np.meshgrid ou np.indices usando NumPy? Caso contrário, as abordagens meshgrid ou indices podem ser otimizadas para evitar cópias de memória desnecessárias (como ascontiguousarray), garantindo ainda que o array seja C-contíguo?

Exemplo:

sizes = (3, 1, 2)
idx = np.ascontiguousarray(np.indices(sizes).reshape(len(sizes), -1).T)
print(idx)
print(f"C_CONTIGUOUS: {idx.flags['C_CONTIGUOUS']}")
# [[0 0 0]
#  [0 0 1]
#  [1 0 0]
#  [1 0 1]
#  [2 0 0]
#  [2 0 1]]
# C_CONTIGUOUS: True
python

1 respostas

  1. Best Answer

    Aqui está uma solução (bastante ingênua) no Numba usando múltiplos threads:

    import numba as nb

@nb.njit(
    [
        # Eagerly compiled for common types
        # Please add your type if it is missing
        '(int32[:,:], int32[:])',
        '(int64[:,:], int32[:])',
        '(float32[:,:], int32[:])',
        '(float64[:,:], int32[:])',
    ],
    parallel=True,
    cache=True
)
def nb_kernel(res, sizes):
    n = np.prod(sizes)
    m = sizes.size
    chunk_size = 1024
    assert n > 0 and m > 0
    for i in range(m):
        assert sizes[i] > 0
    # Compute blocks of 256 rows.
    # Multiple threads compute separate blocks.
    for block in nb.prange((n + chunk_size - 1) // chunk_size):
        start = block * chunk_size
        end = min(start + chunk_size, n)
        # Compute the first row of the block
        jump = 1
        for j in range(m-1, -1, -1):
            res[start, j] = (start // jump) % sizes[j]
            jump *= sizes[j]
        # The next rows of the block incrementally
        for i in range(start+1, end):
            inc = 1
            for j in range(m-1, -1, -1):
                val = res[i-1, j] + inc
                if val >= sizes[j]:
                    val = 0
                    inc = 1
                else:
                    inc = 0
                res[i, j] = val

def nb_compute(sizes, dtype):
    res = np.empty((np.prod(sizes), len(sizes)), dtype=dtype)
    nb_kernel(res, np.array(sizes, dtype=np.int32))
    return res

    Referência

    Na minha máquina (CPU i5-9600KF), no Windows, aqui estão os resultados com sizes=(101,53,71)e dtype=np.int32:

    np.vstack(...):              626.1 ms
np.ascontiguousarray(...):     3.5 ms
np.stack(...):                 2.6 ms

nb_compute(...):               1.1 ms   <----
nb_kernel(...):                0.5 ms   <----

With a fixed length of `sizes` known at compile time:
    nb_compute(...):           0.8 ms   <----
    nb_kernel(...):            0.2 ms   <----

    Análises e otimizações

    Podemos ver que chamar nb_kerneldiretamente um array pré-alocado é significativamente mais rápido. De fato, quando preenchemos um array pela primeira vez, isso causa muitas falhas de página de memória , que são inerentemente lentas. Fazer isso em paralelo é melhor (mas não é escalável no Windows).

    Se você já fizer isso em cada thread de um código paralelo, nb_kernelisso não tornará o processo significativamente mais rápido. De fato, a maior parte da velocidade do Numba vem do uso de múltiplas threads. Consequentemente, neste caso, precisamos otimizar o kernel do Numba. Uma otimização importante é especializar a função para um comprimento específico desizes (então conhecido em tempo de compilação). De fato, o código é mais que o dobro se substituirmos mpor 3 (portanto, suportamos apenas len(sizes)3). Espero que a maioria dos casos tenha um tamanho muito pequeno, len(sizes)então você pode especializar a função para os casos 2, 3, 4 e 5 e escrever uma função Python pura chamando a boa especialização. Essa otimização também torna o código paralelo mais rápido.

    Para melhor desempenho, evite preencher arrays grandes devido à lentidão da DRAM. Isso é especialmente verdadeiro para arrays temporários (arrays que são preenchidos uma vez e nunca mais reutilizados) devido a falhas de página. Acredito que o código acima seja ideal para arrays de saída que não cabem no cache de último nível (LLC) da sua CPU.

    Para matrizes de saída que se encaixam no LLC, há implementações mais rápidas do que a acima, por exemplo, usando uma linguagem nativa compatível com SIMD (mas é bem complexa de implementar).

    • 3

relate perguntas