WWDC의 자동 메모이제이션 코드 분석

시간이 오래 걸릴 수 있는 비싼 연산 작업이 있고, 이러한 연산을 반복해서 수행해야 하는 상황이 있을 때, 가장 손쉬운 성능 개선 방법은 연산 결과를 캐싱하는 것이다. 연산을 수행하는 함수가 입력이 같을 때 출력이 같은 것이 보장되는 순수한 함수라면, 재계산을 수행하는 것보다 캐시된 내용을 읽어와서 사용하는 것이 훨씬 빠른 것은 자명한 사실이다.

메모이제이션(Memoization)

실제로 이런 종류의 문제는 제법 흔하고 실제로도 계산 결과를 캐시해서 사용하는 것은 좋은 해결책 중 하나이다. 어떤 문제가 더 작은 부분문제들로 나뉠 수 있고, 전체 문제는 작은 부분문제들의 반복 계산에 의해 산출되는 것이라면 이 패턴을 적용하여 성능을 높일 수 있다. 이러한 패턴을 메모이제이션(memoization)이라 한다.

메모이제이션은 캐시를 위해 약간의 메모리 자원을 소모하는 대신에 수행 속도를 그만큼 높일 수 있는 테크닉이다. 실제로도 기존 함수를 간단하게 수정하는 것으로 메모이제이션을 적용할 수 있는데, 만약 기존 함수가 제 3자가 작성한 라이브러리에 있는 것이라 수정이 불가능하다면 메모이제이션이 적용된 형태로 변경하는 함수를 만들어서 사용할 수 있을 것이다. (이러한 패턴은 파이썬에서는 데코레이터 등을 사용해서 많이 쓰인다.)

이와 관련해서 14년 WWDC에서는 자동 메모이제이션(auto memoization)이라는 테크닉을 소개하였는데, 가히 그 코드가 mind-blowing 하다고 할 만큼 언뜻 봐서는 이해가 불가능해 보일 지경으로 신묘하다.

해당 세션에서 소개된 코드의 최종형태는 아래와 같다.1

func memoize<T: Hashable, U> (body: @escaping ((T) -> U, T) -> U) -> (T) -> U
{
  var memo = Dictionary<T, U>()
  var result: ((T)->U)!
  result = { x in 
    if let q = memo[x] { return q }
    let r = body(result, x)
    memo[x] = r
    return r
  return result
}

// 적용

let fib: (Int) -> Int = memoize{ fib, n in
    guard n > 1 else { return 0 }
    return fib(n-1) + fib(n-2)
}

기본적인 메모이제이션 함수

기본적으로 메모이제이션은 한 번 계산한 결과를 캐싱해두고 다시 똑같은 입력이 들어왔을 때는 계산을 수행하지 않고 그 값을 활용하는 것이다. 따라서 메모이제이션이 적용된 함수는 다음과 같은 로직으로 동작하면 된다.

  1. 캐시로 사용할 사전(Dictionary)을 하나 준비한다.[^2]
  2. 입력된 인자값 x 가 캐시의 키로 존재하는 경우, 연산 과정 없이 해당 키의 값을 꺼내어 리턴한다.
  3. 입력된 인자에 대해 계산된 캐시가 없는 경우, 실제 연산을 수행한 결과를 얻고
  4. 입력값과 계산 결과 값의 쌍을 캐시에 추가한 후
  5. 결과값을 리턴한다.

여기서 실제 연산을 처리할 함수를 일반화하면 다음과 같이 메모이제이션 함수를 만들 수 있다.

func simpleMemoization <T: Hashable, U> (_ f: @escaping (T) -> U) -> (T) -> U
{
  var memo = Dictionary<T, U>()
  return { (x: T) -> U in
    if let cachedResult = memo[x] { return x }
    let newResult = f(x)
    memo[x] = newResult
    return newResult
  }
}

이 함수는 (T) -> U 타입의 함수를 받아서 동일한 타입의 함수(클로저)를 리턴한다. 이 때 리턴되는 함수는 사전 객체 하나를 캡쳐하고 있고, 이 사전에 한 번 계산한 값들을 저장한다.

확인

그렇다면, 실제로 이 함수가 효과가 있는지를 살펴보자. 어떤 함수가 좋을까?  시간이 오래걸리면서 작성하기도 쉬운 (그리고 만만한) 소수 판별 함수가 좋겠다.  소수 판별 함수를 반복 적용하면서 계산에 드는 시간을 간이로 재어보자.

func isPrime(_ n: Int) -> Int
{
  guard n > 1 else { return false }
  for k in 2..<n {
    if n % k == 0 { return false }
  }
  return true
}

위는 정말 무책임한(?) 소수 판별함수이고, 다음은 반복과 시간 측정을 위한 함수들을 간단히 작성해본다.

import Foundation

func timeit(_ f: () -> ()){
  let begin = Date
  f()
  let elapsed = Date().timeIntervalSince(begin)
  print("time: \(elapsed * 1000)ms")
}

func repeater(_ count: Int, f: () -> ()) {
  guard count > 0 else { return }
  for _ in 0..<count { f() }
}

테스트 : 메모이제이션 없을 때

그리고 적당한 구간의 소수의 개수를 구하는 동작을 처리하도록 해보자. isPrime 이 메모아이즈 되지 않았다면 매 시행마다 거의 비슷한 시간이 소요될 것이다.

repeater(3){
  timeit{
    print(Array(1...30000).filter(isPrime).count)
  }
}

/*
3245
time: 1016.9049501419ms
3245
time: 999.333988692ms
3245
time: 1017.28498395699ms
*/

테스트: 메모이제이션 적용시

다음은 메모이제이션을 적용한 상태로 테스트하는 것이다.

let f = memoize(isPrime)
repeater(3){ timeit {
  print(Array(1...30000).filter(f).count)
}}

/*
3245
time: 1064.60899114609ms
3245
time: 42.5939559936523ms
3245
time: 44.7009801864624ms
*/

실행결과, 첫 번째 시행에서는 앞의 테스트와 거의 같은 시간이 걸리는 것이 확인된다. 그런데 2차와 3차의 테스트에서는 1차의 테스트에 비해서 1/200 정도의 시간만이 소요되었다. 왜냐면 1차의 테스트에서 30000개의 시행 결과가 모두 캐시되있고, 그 이후 시행에서는 모두 캐시된 결과를 그대로 사용하기 때문이다.

재귀 호출을 포함하는 함수의 메모이제이션

문제는 재귀호출을 하는 함수를 메모이제이션하게 되었을 때이다. 임의의 함수 f가 재귀호출을 포함한다고 가정하자. 이 함수 f는 수행중에 다시 f를 호출하는 형태의 코드를 포함하게 된다. 그리고 이 함수 f를 메모이제이션 한 함수 g를 생각해보자.

  • 함수 g가 인자 a를 가지고 호출된다.
  • 인자 a에 대한 리턴 값은 캐시된적이 없다고 하면 이 부분은 실제로 계산이 필요하다. 따라서 f(a)를 호출한다.
  • f는 가정한대로 재귀호출을 포함하며 이 때 a가 아닌 b를 인자로 사용한다.
  • 만약 b가 g에 의해서 이미 계산된 적이 있다 하더라도 한 번 f가 호출된 이후부터는 항상 f만 호출되며 g가 호출되지 않는 이상 이전에 계산된 파라미터 값이 중간부터 등장하더라도 캐시의 혜택을 받지 못한다.

이와 유사한 케이스를 만들어보자. 주어진 x 이하의 소수의 합을 구하는 함수를 재귀 형태로 만들어보겠다.

func sumOfPrimesUpto(_ n: Int) -> Int {
  guard n > 1 else { return 0 }
  return (isPrime(n) ? n : 0) + sumOfPrimeUpto(n - 1)
}

함수의 특성상 만약 여기에 3000을 대입했다고 하면, sumOfPrimeUpto(2999)가 분명히 호출되므로 그 다음 번에 다시 2999 이하의 소수의 합을 구하는 호출에는 수행시간이 매우 짧아져야 할 것이다.  하지만 여기서는 성공적으로 캐싱이 되지 않으리라는 것을 충분히 예측할 수 있고, 결과는 아래와 같다.

45675864
time: 979.550004005432ms
45675864
time: 985.236048698425ms

재귀호출에서도 메모아이징이 되도록 하기 위한 준비

그렇다면 어떻게 재귀호출하는 패턴의 함수가 메모아이징을 할 수 있을까? 바로 “캐시된 값이 없을 떄에 수행하는 계산” 자체에서 메모이제이션된 함수를 호출해야 하는 아이러니한 상황이 된다.  즉 이미 정의해버린 재귀호출 함수를, 메모이제이션 함수를 통해서 수정할 수는 없는 상황이된다. 따라서 자동 메모이제이션 함수는 클로저를 기반으로 작성되어야 한다는 결론에 다다를 수 있다. 그런데 문제가 있다.

다시 아래 코드를 보자. 만약 우리가 모종의  memoize 함수를 작성하는데 성공했다고 가정하자. 그렇다면 위에서 작성한 재귀꼴의 sumOfPrimes 를 다음과 같이 정의할 수 있게 된다.

let sumOfPrimes:(Int) -> Int = memoize{ n in 
    guard n > 1 else { return 0 }
    return (isPrime(n) ? n : 0) + sumOfPrimes(n-1)
}

그런데 이 구문은 실질적으로는 에러가 된다. 왜냐하면 세 번째 행에서 참조하고 있는 sumOfPrimes 클로저는 아직 선언되지 않은 (이 식은 현재 sumOfPrimes를 정의하기 위한 식의 우변이므로) 상태에서 참조되기 때문이다.

이 문제를 해결하기 위해서는 이 함수가 호출해야 할 실체가 사실은 메모아이즈를 적용받은 함수여야 한다는 조건이 필요하다. 그런데 지금 그 함수를 만드는 중이니 자신을 정의하는 중에 자신을 참조해야 한다는 아이러니가 발생했다….  어떻게 하면 좋을까?

그 제약 자체를 무시하자. 그러니까, 자기 자신을 정의하는 중에 호출하는 것이 아니라, 어떤 “메모아이즈된 함수를 인자로 받고, 이를 호출하는 것”이다.  뜬금없지만 -그리고 당연하지만-문법상으로는 아무런 문제가 되지 않는다. 따라서, 다음과 같이 메모아이즈된 함수를 인자로 받아서 호출하는 것으로 쓸 수는 있을 것이다.

let sumOfPrimes:(Int) -> Int = memoize{ memoizedFunc, n in 
  guard n > 1 else { return 0 }
  return (isPrime(n) ? n : 0) + memoizedFunc(n-1)
}

그리고 이 방법이 나쁘지 않은 것으로 판단된다면 그에 맞게 memoize 함수를 다시 고쳐쓰면 될 노릇이다.  memoizedFunc  라는 인자를 누가 주는가?에 대한 의문이 남아있기는 하지만(당연히 memoize 함수가 정해줘야 한다),  memoize 함수의 인자가 ((T)->U) -> (T) -> U 에서 함수 타입이 추가되어 (((T) -> U, T) -> U) -> (T) -> U가 되었다.  그럼 메모이제이션 함수는 아래와 같이 수정할 수 있을 것이다. 그리고 이 모양이 바로 WWDC의 바로 그 코드가 된다.

fun memoize <T: Hashable, U> (_ body: @escaping (((T) -> U, T) -> U)) -> (T) -> U
{
  var memo = [T:U]()
  var result: ((T) -> U)!
  result = { n in 
    if let cachedResult = memo[n] { return cachedResult }
    let newResult = body(result, n)
    memo[n] = newResult
    return newResult
  }
  return result
}

이제 마술의 안개를 걷어낼 차례이다.

  1.  memoize 함수는 어떤 클로저를 받은 후, 캐시에서 답을 찾고 답이 없는 경우만 해당 클로저를 호출하여 계산을 수행하는 함수를 리턴한다.
  2. 그런데 재귀등에 의해 중첩된 계산 내에서도 캐시를 이용하려면 원래 클로저가 캐시를 이용하도록 변형된 함수를 호출해야 한다.
  3. 그런데 원래의 클로저가 아직 만들어지지 않은 자기 자신의 메모이제이션 버전을 미리 아는 것은 말이 안된다.
  4. 따라서 원래의 재귀호출을 하던 클로저는 자신의 캐시화 버전을 인자로 받고, 자신을 재귀호출 하는 대신 인자로 받은 메모아이즈된 함수를 호출하는 식으로 변경할 수 있고,
  5. memoize 함수내에서 이 캐시되는 함수는 정의하기 전에 선언하여 (옵셔널 타입이면 가능하니까) 원래 함수의 인자로 전달될 수 있게 한다.

자 이제 최종적으로 테스트해보자. memoize 함수가 위와 같다면 올바른 메모이제이션 테스트는 아래와 같이, 이미 계산됐을 것으로 보이는 인자를 전달해주면 즉시 답을 낼 것이다.

let sumOfPrimesUpto: (Int) -> Int = memoize{ m, n in
  guard n > 1 else { return 0 }
  return (isPrime(n) ? n : 0) +  m(n-1)
}

timeit {
  print(sumOfPrimesUpto(10000)) // 이 계산 중에 sumOfPrimesUpto(9999)가 캐시될 것이다.
}
timeit {
  print(sumOfPrimesUpto(9999) // 위 계산에서 9999에 대한 계산을 마쳐서 캐시하고 있으므로, 거의 시간이 걸리지 않는 수준으로 찍힌다. 
}

 


  1. 당시 Swift 1에서는 인자로 받는 클로저는 기본적으로 탈출 가능했으나, Swift3부터는 클로저 인자의 기본 속성이 탈출 불가능하게끔 바뀌었다. 따라서 @escaping을 추가해주어야 한다.