Javascript - 1억 계산 코드 개선하기

자바스크립트 코드 중 속도가 느려 불만을 일으키는 코드가 있다.

norm(src, dest) {
    return src.map((v, i) => v - dest[i])
};

sqrt(m) {
    let sum = 0;
    m.forEach((row) => {
        sum += row * row;
    })
    return sum * 0.5;
}

computeDiagramPoint(mtx) {
    let n = mtx.length;
    console.time();
    const digonal = this.makeDiagonal(n);
    for (let i = 0; i < n; i++) {
        for (let j = 0; j < n; j++) {
            const m = this.norm(mtx[i], mtx[j]);
            digonal[i][j] = this.sqrt(m);
        }
    }
    console.timeEnd();

    return digonal;
}

이 코드는 interpolation 보강 무게를 구하는 코드로 입력 곡선 데이터가 n개의 연속된 배열이 있다면 무게를 구하기 전에 n*n 매트릭스 형태로 크기를 확장하고 무게를 구하는 코드인데, 문제는 속도가 최소 O^2 급이므로 입력되는 곡선이 클수록 연산량이 굉장히 많다

예를 들어 곡선 데이터 특성상 만개의 배열이 담긴 데이터가 오는 경우가 많은데 이럴 경우 10000*10000 매트릭스 확장되어 1억이 넘는 연산이 필요하다.

먼저 위 코드를 실행하여 속도가 얼마나 되는지 살펴보자.

속도 측정

default: 45031.274169921875 ms

1초의 밀리세컨드는 1000ms 이므로 45초가 걸린 셈이다.

코드를 개선해서 속도를 확보해보자.


함수 오버헤드 제거

코드를 살펴보면 sqrt, norm 함수가 있다. 코드 수행에는 함수로 이동해서 연산에 대한 오버헤드가 존재하므로 제거해보자. 제거 시 코드는 SOLID 원칙에 위배되지만 이 코드는 속도가 생명이므로 원칙을 어길 수 밖에 없는 상황이다.

함수 안에 있던 계산 로직들은 내부로 다시 적분 하듯 통합한다.

computeDiagramPoint(mtx) {
    let n = mtx.length;
    const m = Array(3).fill(0);
    console.time();
    const digonal = this.makeDiagonal(n);
    for (let i = 0; i < n; i++) {
        for (let j = 0; j < n; j++) {
            m[0] = mtx[i][0] - mtx[j][0];
            m[1] = mtx[i][1] - mtx[j][1];
            m[2] = mtx[i][2] - mtx[j][2];
            digonal[i][j] = (m[0] * m[0] + m[1] * m[1] + m[2] * m[2]) * 0.5;
        }
    }
    console.timeEnd();

    return digonal;
}

코드 수행 결과는 다음과 같이 발생했다.

default: 54738.47216796875 ms

되려 10초 가까이 더 늘어났다.
원인을 분석하니 생각없이 m 배열로 3개 공간을 잡아 마구잡이로 연산하다 포인터 연산으로 인해 더 늘어난 형태가 되었다. JIT 컴파일러가 헷갈려하면 최적화가 안될 수 있으니 잊지 말고 잘 다뤄보자.

배열 접근은 아래와 같이 거치고 있는 과정이 많으므로 속도가 느려지는 것으로 보인다.
m[i] // 배열 → 포인터 → 배열 → 숫자

다시 한번 코드를 수정해보자.

computeDiagramPoint(mtx) {
    let n = mtx.length;
    let x = 0;
    let y = 0;
    let z = 0;
    console.time();
    const digonal = this.makeDiagonal(n);
    for (let i = 0; i < n; i++) {
        for (let j = 0; j < n; j++) {
            x = mtx[i][0] - mtx[j][0];
            y = mtx[i][1] - mtx[j][1];
            z = mtx[i][2] - mtx[j][2];
            digonal[i][j] = (x * x + y * y + z * z) * 0.5;
        }
    }
    console.timeEnd();

    return digonal;
}
default: 41523.679931640625 ms

!!! 3초나 줄어들었다. 별로 체감되는 수치는 아니지만 무려 8퍼 가깝게 줄어든 것이다. 희망이 보이니 좀더 개선해보자.


대칭 행렬 연산 로직 제거

이번엔 계산 로직을 건들어볼 생각이다. 이중 for 문 i 다음 j 연산에서 i와 함께 j를 연산하는 대칭구조인데, i연산은 구지 j for문 안에서 돌 필요가 없어 보인다. 연산 로직 자체를 변경해보자.

computeDiagramPoint(mtx) {
    let n = mtx.length;
    let x = 0;
    let y = 0;
    let z = 0;
    console.time();
    const digonal = this.makeDiagonal(n);
    for (let i = 0; i < n; i++) {
        const xi = mtx[i][0];
        const yi = mtx[i][1];
        const zi = mtx[i][2];
        for (let j = 0; j < n; j++) {
            x = xi - mtx[j][0];
            y = yi - mtx[j][1];
            z = zi - mtx[j][2];
            digonal[i][j] = (x * x + y * y + z * z) * 0.5;
        }
    }
    console.timeEnd();

    return digonal;
}
default: 28771.144775390625 ms

계산량이 반절로 줄어들었고 i for문에서 미리 캐싱 연산 하였으므로 j for문에서도 부담감이 줄어들었다.

실제 체감 느낌이 좋고 피로감이 없어지고 있다..! 긍정적인 속도 체감이므로 좀 더 노력해보자.


SoA 구조 변경

코드 아키텍처 구조적으로 AoS에서 SoA으로 변경하여 개선하는 과정이다.

먼저 앞에서 설명했듯 배열로 이용한 계산은 포인터로 인한 오버헤드가 존재한다.

mtx[i][0] // 배열 → 포인터 → 배열 → 숫자

메모리 접근 비용이 크다.

배열이 존재한 채로 계산하면 메모리 접근 비용이 존재하므로 없애주기로 한다.

computeDiagramPoint(mtx) {
    let n = mtx.length;
    let x = new Float32Array(n);
    let y = new Float32Array(n);
    let z = new Float32Array(n);
    console.time();

    for (let i = 0; i < n; i++) {
        x[i] = mtx[i][0];
        y[i] = mtx[i][1];
        z[i] = mtx[i][2];
    }
    const digonal = this.makeDiagonal(n);
    for (let i = 0; i < n; i++) {
        const xi = x[i], yi = y[i], zi = z[i];
        for (let j = 0; j < n; j++) {
            const dx = xi - x[j];
            const dy = yi - y[j];
            const dz = zi - z[j];
            digonal[i][j] = (dx * dx + dy * dy + dz * dz) * 0.5;
        }
    }
    console.timeEnd();

    return digonal;
}
default: 23819.849853515625 ms

앞에서 for 문으로 미리 mtx 이중 포인터 계산 전 앞부분을 캐싱하여 연산하였더니 체감 속도가 느껴진다. 배열의 값들은 불필요한 공간 확장하지 않도록 Float32Array 고정 배열로 사용하였다.


대칭 로직 개선해보기

코드를 자세히 살펴본 사람이라면 눈치 채겠지만, 이 n배열을 n*n 확장하고 값들이 대칭으로 펼쳐진 것을 볼 수 있다. 그렇다면 대칭적으로 값들을 할당하면 O(N^2) 속도에서 O(N log N) 속도를 기대할 수 있다.

computeDiagramPoint(mtx) {
    console.time();
    let n = mtx.length;
    let x = new Float32Array(n);
    let y = new Float32Array(n);
    let z = new Float32Array(n);
    let syn = [];
    let dx = 0;
    let dy = 0;
    let dz = 0;

    for (let i = 0; i < n; i++) {
        x[i] = mtx[i][0];
        y[i] = mtx[i][1];
        z[i] = mtx[i][2];
        syn[i] = new Float32Array(n).fill(0);
    }

    const digonal = this.makeDiagonal(n);
    for (let i = 0; i < n; i++) {
        const xi = x[i], yi = y[i], zi = z[i];
        const di = digonal[i];
        for (let j = i + 1; j < n; j++) {
            dx = xi - x[j];
            dy = yi - y[j];
            dz = zi - z[j];
            di[i][j] = (dx * dx + dy * dy + dz * dz) * 0.5;
        }
    }

    for (let i = 0; i < n; i++) {
        for (let j = i + 1; j < n; j++) {
            di[j][i] = di[i][j];
        }
    }

    console.timeEnd();

    return digonal;
}

가독성은 나락으로 갔지만, 자바스크립트 문제인지 배열을 고정으로 만들고 값 대입하는 곳에서 알 수 없는 메모리가 생성되어 Out-of-Memory 가 발생한다. 그래서 불필요한 syn 배열과 새 for문을 만들고 대칭으로 할당하는 for문이 추가되 예상치 못한 오버헤드가 있다.

default: 15938.491943359375 ms

그래도 가독성과 동적 확장 그리고 SOLID 원칙을 희생했지만, 속도는 확실히 개선되었다.

마지막으로 범용으로 사용하고 코드이다.

computeDiagramPoint(mtx) {
    console.time()
    let cols = mtx.length;
    let rows = mtx[0].length;
    if (cols < rows) {
        throw new Error(`[ERROR] computeDiagramPoint Parameter.`);
    }

    const diagonal = Array.from({length: cols}, _ =>
        new Float32Array(cols).fill(1)
    );
    for (let i = 0; i < cols; i++) {
        const xi = mtx[i];
        const di = diagonal[i];
        for (let j = i + 1; j < cols; j++) {
            const dav = this.norm2(xi, mtx[j]);
            di[j] = this.sqrt(dav);
        }
    }

    for (let i = 0; i < cols; i++) {
        const di = diagonal[i];
        for (let j = i + 1; j < cols; j++) {
            diagonal[j][i] = di[j];
        }
    }
    console.timeEnd();

    return diagonal;
}
default: 20788.2119140625 ms


GPU 사용하기

브라우져 GPU 관련 API 기능이 있다. 제약 조건이 250MB 그리고 스레드 256 까지 할당 제한이 있어서 chunk 함수가 별도로 필요한데, 구현이 어려웠다. (이터레이터도 제한되어 있었데, 2^16 만큼만 지정할 수 있다.)

GPU 데이터 처리 코드를 모두 구현하였다면 50MB chunk와 256 스레드로 할당하여 병렬 처리에 네이티브보다 20배 빠른 속도로 처리할 수 있었다.

default: 1843.1290578092 ms

자바스크립트 코드

/**
 *
 * @param mtx Import Data
 * @param mb Chunk is Slice Mib
 */
async makeDiagonalChunk(mtx ,mb) {
    console.time();

    const code = `
        @group(0) @binding(0) var<storage, read_write> GPU : array<f32>;
        @group(0) @binding(1) var<storage, read_write> RGB : array<f32>;
        @group(0) @binding(2) var<storage, read_write> CHUNK : array<f32>;
        @group(0) @binding(3) var<storage, read_write> COLORS : array<f32>;

        @compute @workgroup_size(256)
        fn main(@builtin(global_invocation_id) gid : vec3<u32>) {
            let o = gid.x;
            let i = u32(CHUNK[0]) + o;
            let j = 0;
            let len = arrayLength(&RGB);
            let n = u32(COLORS[0]);
            
            // find index
            let cols = len / n;
            let rows = i % cols;
            let ix = i / cols;
            let jx = rows;
            
            var v = f32(0);
            
            // Extended Support Calculation
            for (var j : u32 = 0u; j < n; j++) {
                let dd = RGB[ix + cols * j] - RGB[jx + cols * j];
                v += dd * dd;               
            }
            v *= 0.5;
            
            // Diagonal Skip
            if ( v == 0 ) {
                return;
            }
            
            // Variable Array Out
            GPU[o] = v;    
        }
    `;

    const chunk = 250000 * (mb ?? 1);
    const cols = mtx.length;
    const rows = mtx[0].length;
    const max = rows * rows;
    const cake = Math.ceil(max / chunk);
    const promises = [];
    const result = [];
    let setIterGPU = 0;
    const gpu = new Gpu()
        .setCode(code)
    await gpu.init();

    const sendData = new Float32Array(cols * rows).map(_ => 0);
    const sendChunk = new Float32Array(1).fill(0);
    const sendColors = new Float32Array(1).fill(cols);
    
    for (let i = 0; i < cols; i++) {
        sendData.set(mtx[i], i * rows);
    }

    for (let s = 0; s < cake; s++) {
        let sub = chunk * s;
        let reserve = s !== cake - 1 ? chunk : max - chunk * (cake -1);
        const poke = [];
        sendChunk.fill(sub);
        const response = new Float32Array(reserve).fill(1);

        poke.push(response);
        poke.push(sendData);
        poke.push(sendChunk);
        poke.push(sendColors);

        setIterGPU = reserve / 256;
        gpu.setCode(code)
            .setData(poke)
            .setByte(4)
            .setIter(setIterGPU)
        promises.push(gpu.run());
        if (cake === s + 1 || s % 10 === 0) {
            const res = await Promise.all(promises);
            res.map(v => {
                result.push(v);
                promises.pop();
            });
        }
    }

    gpu.free();
    console.log(result);
    console.timeEnd();
    // return load;
}

gpu 계산하려면 코드를 병렬 처리가 가능한 코드로 작성이 필요하다. 코딩 실력이 부족해 최적화가 더 진행할 수 없지만, 25^25 공간 계산에 50초 걸렸던 작업이 초단위로 처리한 것을 보면 만족스럽다.

다만, 전용 GPU 메모리가 부족하면 버퍼 크래쉬가 발생하니 앞부분에 계산이 가능한지 예외 처리가 필요하다.