Keyboard shortcuts

Press ← or → to navigate between chapters

Press S or / to search in the book

Press ? to show this help

Press Esc to hide this help

mlir-egglog

Affine 다이얼렉트와 OpenMP

· 태그: ai, python, llms, compilers, mlir · 원문 · 한국어 번역(원문에서 변경됨)

MLIR의 affine 다이얼렉트는 병렬 실행에 적합한 루프와 배열 접근을 표현하기 위한 추상화를 제공합니다. 앞서 살펴본 scf 다이얼렉트와 달리, affine 다이얼렉트는 고급 루프 변환과 자동 병렬화를 가능하게 합니다. 이는 본질적으로 “일급 루프“이며, 우리가 늘 LLVM 위에 덧붙여 구현해 왔던 것입니다.

affine 다이얼렉트에는 자동 병렬화 기회, 내장 벡터화 기능, 정교한 루프 최적화, 배열 연산의 자연스러운 표현, OpenMP와의 통합 등 많은 장점이 있습니다.

다면체 최적화 모델

다면체 모델은 루프 중첩과 그 연산을 다차원 공간의 기하학적 객체로 표현하는 수학적 프레임워크입니다.

핵심(Core)적으로, 다면체 모델은 루프 중첩 안에서 실행되는 문장을 다차원 공간의 점으로 표현합니다. 각 점은 실행 중인 문장의 특정 인스턴스에 대응하며, 해당 반복 좌표로 정의됩니다. 이 표현은 복잡한 루프 최적화 문제를 다면체에 대한 기하학적 연산으로 바꾸어 프로그램 분석 및 변환 기법을 가능하게 합니다.

다면체 모델의 핵심 구성 요소로는 반복 도메인, 접근 관계, 스케줄링 함수가 있습니다. 반복 도메인은 루프 중첩이 실행하는 모든 반복의 집합을 정의하며, 이는 affine 제약 조건으로 경계가 정해진 정수 점의 집합으로 표현됩니다. 접근 관계는 반복 점을 각 문장이 접근하는 메모리 위치에 대응시킵니다. 스케줄링 함수는 반복 점의 실행 순서를 결정하며, 프로그램의 의미를 보존하면서 성능 특성을 개선하는 변환을 가능하게 합니다.

MLIR은 다면체 개념을 일급 시민으로 통합한 affine을 제공합니다. 이 하이브리드 설계는 고성능 데이터 병렬 시스템을 위한 고수준 데이터 흐름 그래프와 타깃별 코드의 표현, 분석, 변환을 최적화합니다.

우리가 이해해야 할 핵심(Core) 개념은 세 가지입니다:

  1. Affine 맵: 차원과 심볼을 결과에 대응시키는 다차원 준선형 함수입니다. 예를 들어, (d0, d1, d2, s0) → (d0 + d1, s0 * d2)는 차원 인수 d0, d1, d2와 심볼 인수 s0을 두 개의 결과에 대응시키는 2차원 affine 맵을 나타냅니다.

  2. 정수 집합: affine 부등식과 등식으로 표현되는, 차원과 심볼에 대한 제약 조건입니다. 예를 들어, (i)[N, M] : (i >= 0, -i + N >= 0, N - 5 == 0, -i + M + 1 >= 0)는 0 <= i < N, N = 5, i <= M + 1을 만족하는 값 i의 집합을 나타냅니다.

  3. Affine 연산: affine 맵과 정수 집합을 활용하여 affine 제약 조건이 있는 루프 중첩과 조건문을 표현하는 affine.for, affine.if, affine.parallel 같은 연산입니다.

이론을 먼저 늘어놓기보다는 몇 가지 예제를 바로 살펴봅시다.

행렬 곱셈

두 행렬 와 의 고전적인 행렬 곱은 다음과 같습니다:

이를 MLIR에서 명시적인 루프 중첩으로 표현하면 다음과 같습니다:

func.func @matmul(%A: memref<?x?xf32>, %B: memref<?x?xf32>, %C: memref<?x?xf32>) {
  %c0 = arith.constant 0 : index
  %c1 = arith.constant 1 : index

  %M = memref.dim %A, %c0 : memref<?x?xf32>
  %N = memref.dim %B, %c1 : memref<?x?xf32>
  %K = memref.dim %A, %c1 : memref<?x?xf32>

  // Sequential implementation
  affine.for %i = 0 to %M {
    affine.for %j = 0 to %N {
      affine.for %k = 0 to %K {
        %a = affine.load %A[%i, %k] : memref<?x?xf32>
        %b = affine.load %B[%k, %j] : memref<?x?xf32>
        %c = affine.load %C[%i, %j] : memref<?x?xf32>
        %prod = arith.mulf %a, %b : f32
        %sum = arith.addf %c, %prod : f32
        affine.store %sum, %C[%i, %j] : memref<?x?xf32>
      }
    }
  }
  return
}

%i와 %j에 대한 두 루프는 서로 다른 반복 사이에 의존성이 없다는 점에 주목하십시오. -affine-parallelize 패스를 사용하면 이를 다음과 같이 변환할 수 있습니다:

func.func @matmul_parallel(%A: memref<?x?xf32>, %B: memref<?x?xf32>, %C: memref<?x?xf32>) {
  %c0 = arith.constant 0 : index
  %c1 = arith.constant 1 : index

  %M = memref.dim %A, %c0 : memref<?x?xf32>
  %N = memref.dim %B, %c1 : memref<?x?xf32>
  %K = memref.dim %A, %c1 : memref<?x?xf32>

  // Parallel implementation
  affine.parallel (%i, %j) = (0, 0) to (%M, %N) {
    affine.for %k = 0 to %K {
      %a = affine.load %A[%i, %k] : memref<?x?xf32>
      %b = affine.load %B[%k, %j] : memref<?x?xf32>
      %c = affine.load %C[%i, %j] : memref<?x?xf32>
      %prod = arith.mulf %a, %b : f32
      %sum = arith.addf %c, %prod : f32
      affine.store %sum, %C[%i, %j] : memref<?x?xf32>
    }
  }
  return
}

affine.parallel 연산은 루프의 여러 반복이 독립적으로, 그리고 잠재적으로 동시에 실행될 수 있음을 표현하는 방법입니다. 이는 “이 계산들은 어떤 순서로든, 또는 한꺼번에 수행될 수 있다“라고 말하는 것과 같다고 생각하십시오. 서로 소통할 필요가 없는 여러 작업자에게 일을 나누어 주는 방식과 비슷합니다.

더 정확히 말하면, affine.parallel은 다음과 같은 중첩 루프의 집합을 정의합니다:

  • 각 반복은 다른 모든 반복과 독립적입니다
  • 반복들은 어떤 순서로든 실행될 수 있으며 동시에 실행될 수도 있습니다
  • 이 연산은 결과를 결합하는 순서가 중요하지 않은 리덕션(값을 합산하는 것 등)을 통해 결과를 생성합니다

기술적으로는 “초직사각형 병렬 밴드(hyper-rectangular parallel band)“를 생성합니다. 즉, 중첩된 여러 병렬 루프를 표현할 수 있으며, 각 루프는 고유한 귀납 변수(루프 카운터), 하한과 상한, 스텝 크기(각 반복마다 얼마나 증가하는지를 나타냄), 그리고 리덕션 연산(서로 다른 반복의 결과를 어떻게 결합할지를 정의함)으로 특징지어집니다.

병렬 영역의 본문은 반드시 affine.yield로 끝나야 하며, 이는 덧셈이나 곱셈 같은 연산을 사용해 서로 다른 반복의 결과를 어떻게 결합할지를 지정합니다. 0번 실행되는 루프의 경우, 결과는 해당 연산의 항등원(덧셈은 0, 곱셈은 1)입니다.

아핀 맵(affine.apply)

아핀 맵은 Affine 다이얼렉트의 핵심 구성 요소로, 다차원 아핀 변환을 나타냅니다. 아핀 맵은 루프 인덱스가 메모리 접근 패턴에 어떻게 매핑되는지를 기술하는 방법을 제공하며, 이는 고성능 컴퓨팅에서 메모리 레이아웃과 접근 패턴을 최적화하는 데 필수적입니다.

affine.apply의 다양한 용도를 보여 주는 종합적인 예제를 살펴봅시다:

// Define some reusable affine maps
#tile_map = affine_map<(d0) -> (d0 floordiv 32)>
#offset_map = affine_map<(d0)[s0] -> (d0 + s0)>
#complex_map = affine_map<(d0, d1)[s0] -> (d0 * 2 + d1 floordiv 4 + s0)>

func.func @affine_apply_examples() {
  // Create some test indices
  %c0 = arith.constant 0 : index
  %c42 = arith.constant 42 : index
  %c128 = arith.constant 128 : index

  // Example 1: Simple tiling calculation
  // Calculate which tile a given index belongs to (tile size = 32)
  %tile_idx = affine.apply #tile_map(%c42)

  // Example 2: Inline affine map for offset calculation
  %offset = affine.apply affine_map<(i) -> (i + 10)>(%c42)

  // Example 3: Using a symbol parameter
  %shifted = affine.apply #offset_map(%c42)[%c128]

  // Example 4: Multiple dimensions and a symbol
  %complex = affine.apply #complex_map(%c42, %c128)[%c0]

  // Example 5: Composition of affine applies
  %temp = affine.apply affine_map<(i) -> (i * 2)>(%c42)
  %final = affine.apply affine_map<(i) -> (i + 5)>(%temp)

  return
}

// Example showing how affine.apply can be used in a practical loop context
func.func @tiled_loop(%arg0: memref<256xf32>) {
  %c0 = arith.constant 0 : index
  %c256 = arith.constant 256 : index
  %c32 = arith.constant 32 : index

  // Outer loop iterates over tiles
  affine.for %i = 0 to 256 step 32 {
    // Inner loop processes elements within the tile
    affine.for %j = 0 to 32 {
      %idx = affine.apply affine_map<(d0, d1) -> (d0 + d1)>(%i, %j)
      %val = memref.load %arg0[%idx] : memref<256xf32>
      // Process value...
      memref.store %val, %arg0[%idx] : memref<256xf32>
    }
  }
  return
}

tiled_loop 함수는 실제 시나리오에서 affine.apply를 사용해 배열의 타일링된 처리를 구현하는 방법을 보여 줍니다. 바깥쪽 루프는 타일 크기만큼씩 진행하며, affine.apply는 각 타일을 처리하는 안쪽 루프의 경계를 계산하는 데 사용됩니다.

아핀 맵에 관한 몇 가지 핵심 사항으로, floordiv는 내림하는 정수 나눗셈을 수행하며, (d0)[s0] 구문은 차원 매개변수 d0 하나와 심볼 매개변수 s0 하나를 나타냅니다. 또한 아핀 맵에는 덧셈, 뺄셈, 상수에 의한 곱셈, 상수에 의한 나눗셈, 상수에 의한 나머지 연산, 상수에 의한 floordiv/ceildiv만 포함될 수 있습니다. 결과 타입은 항상 index입니다.

아핀 최적화 패스

MLIR 최적화기는 Affine 다이얼렉트를 위해 특별히 설계된 풍부한 패스 집합을 제공합니다. 이 패스들은 다양한 최적화 기법을 통해 코드를 자동으로 변환하여 성능을 향상시킬 수 있습니다. 가장 중요한 몇 가지를 살펴봅시다:

  1. 루프 불변 코드 이동(Loop invariant code motion): 이 변환은 루프 불변 코드(루프 인덱스에 의존하지 않는 코드)를 루프 밖으로 옮깁니다. 이는 루프 인덱스에 의존하지 않는 코드를 바깥쪽 루프로 옮김으로써 병렬성을 드러내는 데 도움이 될 수 있습니다.
// Before
func.func @before(%A: memref<10xf32>, %B: memref<10xf32>) {
  %c0 = arith.constant 0 : index
  %c10 = arith.constant 10 : index
  %c1 = arith.constant 1 : index

  affine.for %i = 0 to 10 {
    %x = arith.constant 42.0 : f32  // This is loop invariant!
    %v = memref.load %A[%i] : memref<10xf32>
    %sum = arith.addf %v, %x : f32
    memref.store %sum, %B[%i] : memref<10xf32>
  }
  return
}

// After
func.func @after(%A: memref<10xf32>, %B: memref<10xf32>) {
  %c0 = arith.constant 0 : index
  %c10 = arith.constant 10 : index
  %c1 = arith.constant 1 : index

  %x = arith.constant 42.0 : f32  // Moved outside the loop
  affine.for %i = 0 to 10 {
    %v = memref.load %A[%i] : memref<10xf32>
    %sum = arith.addf %v, %x : f32
    memref.store %sum, %B[%i] : memref<10xf32>
  }
  return
}
  1. 루프 스큐잉(Loop skewing): 이 변환은 다른 루프 인덱스의 값에 따라 반복을 이동시켜 반복 공간을 변경합니다. 이는 의존성을 변환함으로써 병렬성을 드러내는 데 도움이 될 수 있습니다.
// Before - Dependencies prevent parallelization
func.func @before(%A: memref<10x10xf32>) {
  affine.for %i = 1 to 10 {
    affine.for %j = 1 to 10 {
      %v1 = affine.load %A[%i-1, %j] : memref<10x10xf32>
      %v2 = affine.load %A[%i, %j-1] : memref<10x10xf32>
      %sum = arith.addf %v1, %v2 : f32
      affine.store %sum, %A[%i, %j] : memref<10x10xf32>
    }
  }
}

// After skewing - Inner loop can now be parallelized
func.func @after(%A: memref<10x10xf32>) {
  affine.for %i = 1 to 10 {
    affine.for %j = %i to %i + 9 {
      %i_new = affine.apply affine_map<(d0, d1) -> (d0)>(%i, %j)
      %j_new = affine.apply affine_map<(d0, d1) -> (d1 - d0)>(%i, %j)
      %v1 = affine.load %A[%i_new-1, %j_new] : memref<10x10xf32>
      %v2 = affine.load %A[%i_new, %j_new-1] : memref<10x10xf32>
      %sum = arith.addf %v1, %v2 : f32
      affine.store %sum, %A[%i_new, %j_new] : memref<10x10xf32>
    }
  }
}
  1. 루프 교환(Loop interchange): 이 변환은 루프의 중첩 순서를 바꾸며, 이를 통해 캐시 지역성을 개선하고 병렬성을 드러낼 수 있습니다.
// Before - Column-major access pattern
func.func @before(%A: memref<16x16xf32>, %B: memref<16x16xf32>) {
  affine.for %i = 0 to 16 {
    affine.for %j = 0 to 16 {
      %v = affine.load %A[%j, %i] : memref<16x16xf32>
      affine.store %v, %B[%j, %i] : memref<16x16xf32>
    }
  }
}

// After - Row-major access pattern (better locality)
func.func @after(%A: memref<16x16xf32>, %B: memref<16x16xf32>) {
  affine.for %j = 0 to 16 {
    affine.for %i = 0 to 16 {
      %v = affine.load %A[%j, %i] : memref<16x16xf32>
      affine.store %v, %B[%j, %i] : memref<16x16xf32>
    }
  }
}
  1. 루프 융합(Loop fusion): 이 변환은 인접한 루프들을 결합하여 지역성을 개선하고 루프 오버헤드를 줄입니다.
// Before fusion
func.func @before(%A: memref<10xf32>, %B: memref<10xf32>, %C: memref<10xf32>) {
  affine.for %i = 0 to 10 {
    %v1 = affine.load %A[%i] : memref<10xf32>
    %v2 = arith.mulf %v1, %v1 : f32
    affine.store %v2, %B[%i] : memref<10xf32>
  }

  affine.for %i = 0 to 10 {
    %v3 = affine.load %B[%i] : memref<10xf32>
    %v4 = arith.addf %v3, %v3 : f32
    affine.store %v4, %C[%i] : memref<10xf32>
  }
}

// After fusion
func.func @after(%A: memref<10xf32>, %B: memref<10xf32>, %C: memref<10xf32>) {
  affine.for %i = 0 to 10 {
    %v1 = affine.load %A[%i] : memref<10xf32>
    %v2 = arith.mulf %v1, %v1 : f32
    affine.store %v2, %B[%i] : memref<10xf32>
    %v3 = arith.addf %v2, %v2 : f32
    affine.store %v3, %C[%i] : memref<10xf32>
  }
}
  1. 루프 타일링(Loop tiling): 이 변환은 캐시 지역성을 개선하기 위해 루프를 더 작은 덩어리로 나눕니다.
// Before tiling
func.func @before(%A: memref<32x32xf32>) {
  affine.for %i = 0 to 32 {
    affine.for %j = 0 to 32 {
      %v = affine.load %A[%i, %j] : memref<32x32xf32>
      // ... computation ...
      affine.store %v, %A[%i, %j] : memref<32x32xf32>
    }
  }
}

// After tiling (tile size 8x8)
func.func @after(%A: memref<32x32xf32>) {
  affine.for %ti = 0 to 32 step 8 {
    affine.for %tj = 0 to 32 step 8 {
      affine.for %i = #map(%ti) to #map(%ti + 8) {
        affine.for %j = #map(%tj) to #map(%tj + 8) {
          %v = affine.load %A[%i, %j] : memref<32x32xf32>
          // ... computation ...
          affine.store %v, %A[%i, %j] : memref<32x32xf32>
        }
      }
    }
  }
}

이러한 최적화는 mlir-opt 도구에서 사용할 수 있습니다:

  • -affine-loop-coalescing - 연속된 루프들을 하나의 루프로 병합합니다.
  • -affine-loop-fusion - 경계와 스텝이 같은 루프들을 병합합니다.
  • -affine-loop-invariant-code-motion - 루프 불변 코드를 루프 밖으로 옮깁니다.
  • -affine-loop-normalize - 루프 경계와 스텝을 정규화합니다.
  • -affine-loop-tile - 루프를 타일링합니다.
  • -affine-loop-unroll - 루프를 언롤링합니다.
  • -affine-loop-unroll-jam - 루프를 언롤하고 하나로 합칩니다(jam).
  • -affine-parallelize - 루프를 병렬화합니다.
  • -affine-pipeline-data-transfer - 메모리 계층 구조에서 명시적으로 관리되는 레벨 간의 논블로킹 데이터 전송을 파이프라인화합니다
  • -affine-scalrep - 저장(store)을 로드(load)로 전달하고 중복 로드를 제거하여 affine memref 접근을 스칼라로 대체합니다
  • -affine-simplify-structures - 맵/집합의 affine 표현식을 단순화하고 memref를 정규화합니다
  • -affine-super-vectorize - 타깃에 독립적인 n차원 벡터 추상화로 벡터화합니다

나중에 사용할 -convert-affine-for-to-gpu도 있습니다. 이 패스는 affine 루프를 여러 GPU 아키텍처를 타깃으로 할 수 있는 GPU 커널로 변환합니다.

합성곱 커널

이제 affine 다이얼렉트에서 사용할 수 있는 기본 최적화를 이해했으니, 더 실용적인 예제인 2D 합성곱 구현을 살펴봅시다. 이 연산은 이미지 처리와 딥러닝에서 근본적인 연산으로, 특징 추출과 패턴 인식에 사용됩니다.

2D 합성곱의 수학 공식은 다음과 같습니다:

module {
  func.func @conv_2d(%input: memref<128x128xf32>, %filter: memref<16x16xf32>, %output: memref<113x113xf32>) {
    // Loop over the output matrix dimensions (113x113)
    affine.for %i = 0 to 113 {
      affine.for %j = 0 to 113 {
        // Use affine.parallel to accumulate values into %acc using iter_args
        %zero = arith.constant 0.0 : f32
        %acc = affine.for %fi = 0 to 16 iter_args(%acc = %zero) -> (f32) {
          %acc_inner = affine.for %fj = 0 to 16 iter_args(%acc_inner = %acc) -> (f32) {
            // Load filter value
            %filter_val = affine.load %filter[%fi, %fj] : memref<16x16xf32>

            // Load corresponding input value from the input matrix
            %input_val = affine.load %input[%i + %fi, %j + %fj] : memref<128x128xf32>

            // Multiply input value with filter value
            %prod = arith.mulf %input_val, %filter_val : f32

            // Add product to the accumulator
            %new_acc = arith.addf %acc_inner, %prod : f32
            affine.yield %new_acc : f32
          }
          affine.yield %acc_inner : f32
        }

        // Store the accumulated result in the output matrix
        affine.store %acc, %output[%i, %j] : memref<113x113xf32>
      }
    }
    return
  }
}

affine 연산을 제대로 최적화하고 병렬화하려면 특정한 순서의 패스가 필요합니다:

mlir-opt conv_2d.mlir \
  --affine-loop-normalize \
  --affine-parallelize \
  -lower-affine \
  --convert-scf-to-cf \
  --convert-cf-to-llvm \
  --convert-math-to-llvm \
  --convert-arith-to-llvm \
  --finalize-memref-to-llvm \
  --convert-func-to-llvm \
  --reconcile-unrealized-casts \
  -o conv_2d_opt.mlir

# Translate to LLVM IR
mlir-translate conv_2d_opt.mlir -mlir-to-llvmir -o conv_2d.ll

# Compile to object file
llc -filetype=obj --relocation-model=pic conv_2d.ll -o conv_2d.o

# Compile to shared library
clang -shared -fopenmp -o libconv2d.dylib conv_2d.o

여기서 -lower-affine 패스는 affine 다이얼렉트를 scf 다이얼렉트와 arith 다이얼렉트의 조합으로 낮춥니다.

병렬화된 코드를 Python에서 사용하기 위해, NumPy와 MLIR MemRef 디스크립터 간의 2D 배열 변환을 처리하는 헬퍼 모듈을 만들었습니다. 이 모듈은 NumPy 배열을 변환하고 컴파일된 MLIR 코드를 실행하는 유틸리티를 제공합니다:

import ctypes
import numpy as np
from ctypes import c_void_p, c_longlong, Structure


class MemRef2DDescriptor(Structure):
    """Structure matching MLIR's 2D MemRef descriptor"""

    _fields_ = [
        ("allocated", c_void_p),  # Allocated pointer
        ("aligned", c_void_p),  # Aligned pointer
        ("offset", c_longlong),  # Offset in elements
        ("shape", c_longlong * 2),  # Array shape (2D)
        ("stride", c_longlong * 2),  # Strides in elements
    ]


def numpy_to_memref2d(arr):
    """Convert a 2D NumPy array to a MemRef descriptor"""
    if not arr.flags["C_CONTIGUOUS"]:
        arr = np.ascontiguousarray(arr)

    desc = MemRef2DDescriptor()
    desc.allocated = arr.ctypes.data_as(c_void_p)
    desc.aligned = desc.allocated
    desc.offset = 0
    desc.shape[0] = arr.shape[0]
    desc.shape[1] = arr.shape[1]
    desc.stride[0] = arr.strides[0] // arr.itemsize
    desc.stride[1] = arr.strides[1] // arr.itemsize

    return desc


def run_conv2d():
    """Run 2D convolution using MLIR compiled module"""
    # Create input arrays
    input_matrix = np.ones((10, 10), dtype=np.float32)
    conv_filter = np.arange(9, dtype=np.float32).reshape(3, 3)
    result = np.zeros((8, 8), dtype=np.float32)

    # Load compiled module
    module = ctypes.CDLL("./libconv2d.dylib")

    # Prepare MemRef descriptors
    input_memref = numpy_to_memref2d(input_matrix)
    filter_memref = numpy_to_memref2d(conv_filter)
    result_memref = numpy_to_memref2d(result)

    # Set function argument types
    module.conv_2d.argtypes = [
        *[type(x) for x in input_memref._fields_],
        *[type(x) for x in filter_memref._fields_],
        *[type(x) for x in result_memref._fields_],
    ]

    # Call the function
    module.conv_2d(
        *[getattr(input_memref, field[0]) for field in input_memref._fields_],
        *[getattr(filter_memref, field[0]) for field in filter_memref._fields_],
        *[getattr(result_memref, field[0]) for field in result_memref._fields_]
    )

    return result


if __name__ == "__main__":
    result = run_conv2d()
    print(result)

OpenMP

affine 다이얼렉트는 병렬 실행을 위한 추상화를 제공하지만, 때로는 병렬화를 더 직접적으로 제어하고 싶을 때가 있습니다. 이때 OpenMP가 등장합니다. OpenMP는 병렬 실행에 대한 명시적인 제어를 제공하는 지시문 기반 모델로, 고성능 컴퓨팅에서 널리 사용됩니다.

OpenMP 병렬화를 보여 주는 간단한 C 예제를 살펴봅시다:

#include <stdlib.h>
#include <stdio.h>
#include <omp.h>

int kernel(float *input, float *output) {
    #pragma omp parallel
    {
        #pragma omp for
        for (int i = 0; i < 10; i++) {
            output[i] = input[i] * 2.0f;
        }
    }
}

int main() {
    float input[10] = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10};
    float output[10];
    kernel(input, output);
    for (int i = 0; i < 10; i++) {
        printf("%f ", output[i]);
    }
    return 0;
}

OpenMP 지원(clang -fopenmp)을 켜고 컴파일하면, 이 프로그램은 루프를 여러 스레드에 걸쳐 병렬화합니다. 같은 계산을 MLIR에서는 OpenMP 다이얼렉트로 표현할 수 있습니다:

병렬 영역 (omp.parallel)

OpenMP는 병렬 영역이라는 개념을 통해 병렬 실행을 구성합니다. MLIR에서 병렬 영역은 omp.parallel 구문으로 표현되며, 이 구문은 코드를 동시에 실행할 수 있는 스레드 팀을 생성합니다. 이는 affine 다이얼렉트에서 본 자동 병렬화에 비해 병렬화를 더 명시적으로 제어하는 방법을 제공합니다.

영역은 omp.terminator 연산으로 종료되며, 이 연산은 모든 OpenMP 영역에 필수입니다.

작업 공유 루프 (omp.wsloop)

병렬 루프를 위해 OpenMP 다이얼렉트는 omp.wsloop 연산을 제공하며, 이 연산은 루프 반복이 현재 컨텍스트의 스레드들에 의해 병렬로 실행됨을 지정합니다. 작업 공유 루프는 일반적으로 병렬 영역 안에서 사용되어 루프 반복을 사용 가능한 스레드들에 분배합니다:

omp.wsloop {
  omp.loop_nest (%i1, %i2) : index = (%c0, %c0) to (%c10, %c10) step (%c1, %c1) {
    %a = load %arrA[%i1, %i2] : memref<?x?xf32>
    %b = load %arrB[%i1, %i2] : memref<?x?xf32>
    %sum = arith.addf %a, %b : f32
    store %sum, %arrC[%i1, %i2] : memref<?x?xf32>
    omp.yield
  }
}

이 예제는 두 행렬을 더하기 위해 중첩 루프를 병렬화하는 모습을 보여 줍니다. omp.wsloop 안의 omp.loop_nest 연산은 반복 변수, 범위, 스텝을 갖춘 중첩 루프 구조를 나타냅니다.

배리어를 이용한 동기화 (omp.barrier)

스레드 동기화를 위해 OpenMP 다이얼렉트는 omp.barrier 연산을 제공하며, 이 연산은 자신이 나타나는 지점에 명시적 배리어를 지정합니다. 이를 통해 팀 내의 모든 스레드가 배리어에 도달한 후에야 어떤 스레드든 진행할 수 있도록 보장합니다:

// Code executed by all threads
omp.barrier
// Code executed after all threads reach the barrier

omp.barrier 연산은 입력이나 출력이 없는 단순한 연산이며, 어셈블리 형식에는 이름과 속성만 포함됩니다.

omp 다이얼렉트

위의 병렬 루프 C 코드를 MLIR로 변환해 봅시다:

func.func private @kernel(%input: memref<10xf32>, %output: memref<10xf32>) {
  %loop_ub = llvm.mlir.constant(9 : i32) : i32
  %loop_lb = llvm.mlir.constant(0 : i32) : i32
  %loop_step = llvm.mlir.constant(1 : i32) : i32

  omp.parallel {
      omp.wsloop {
        omp.loop_nest (%i) : i32 = (%loop_lb) to (%loop_ub) inclusive step (%loop_step) {
          %ix = arith.index_cast %i : i32 to index
          %input_val = memref.load %input[%ix] : memref<10xf32>
          %two = arith.constant 2.0 : f32
          %result = arith.mulf %input_val, %two : f32
          memref.store %result, %output[%ix] : memref<10xf32>
          omp.yield
        }
      }
    omp.barrier
    omp.terminator
  }

  return
}

그런 다음 정적으로 정의된 입력에 대해 커널을 호출하는 main 함수를 작성할 수 있습니다:

memref.global constant @input : memref<10xf32> = dense<[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0]>

// Define external references to printf and kernel
llvm.func @printf(!llvm.ptr, ...) -> i32
func.func private @kernel(%input: memref<10xf32>, %output: memref<10xf32>) -> ()

func.func @main() {
  %input = memref.get_global @input : memref<10xf32>
  %output = memref.alloc() : memref<10xf32>

  func.call @kernel(%input, %output) : (memref<10xf32>, memref<10xf32>) -> ()

  %lb = index.constant 0
  %ub = index.constant 10
  %step = index.constant 1

  %fs = llvm.mlir.addressof @fmt : !llvm.ptr

  scf.for %iv = %lb to %ub step %step {
    %el = memref.load %output[%iv] : memref<10xf32>
    llvm.call @printf(%fs, %el) vararg(!llvm.func<i32 (ptr, ...)>) : (!llvm.ptr, f32) -> i32
  }

  return
}

// Define a constant string for the format specifier
llvm.mlir.global private constant @fmt("%f\0A\00") {addr_space = 0 : i32}

덧붙여 말하자면, MLIR의 가변 인자 함수는 C와 마찬가지로 유연한 인자 전달을 가능하게 합니다. 우리가 사용할 일반적인 패턴은 가변 개수의 인자를 받을 수 있는 printf와 같은 함수를 호출하는 것입니다. main 함수에서는 결과를 출력하기 위해 (가변 인자 함수인) printf 함수를 호출합니다. 다음 코드는 C의 printf("%f\n", el)과 동일하며, MLIR에서 printf를 사용하는 방법을 보여 줍니다.

// Define external printf function from stdio.h
llvm.func @printf(!llvm.ptr, ...) -> i32

// Define a constant string for the format specifier
llvm.mlir.global private constant @fmt("%f\0A\00") {addr_space = 0 : i32}

// Call the variadic printf function with the format specifier and type specialized arguments
%fs = llvm.mlir.addressof @fmt : !llvm.ptr
%el = llvm.mlir.constant(1.0 : f32) : f32
llvm.call @printf(%fs, %el) vararg(!llvm.func<i32 (ptr, ...)>) : (!llvm.ptr, f32) -> i32

다음은 병렬 루프와 main 함수를 포함한 전체 예제입니다:

memref.global constant @input : memref<10xf32> = dense<[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0]>

llvm.func @printf(!llvm.ptr, ...) -> i32
llvm.mlir.global private constant @fmt("%f\0A\00") {addr_space = 0 : i32}

func.func private @kernel(%input: memref<10xf32>, %output: memref<10xf32>) {
  %loop_ub = llvm.mlir.constant(9 : i32) : i32
  %loop_lb = llvm.mlir.constant(0 : i32) : i32
  %loop_step = llvm.mlir.constant(1 : i32) : i32

  omp.parallel {
      omp.wsloop {
        omp.loop_nest (%i) : i32 = (%loop_lb) to (%loop_ub) inclusive step (%loop_step) {
          %ix = arith.index_cast %i : i32 to index
          %input_val = memref.load %input[%ix] : memref<10xf32>
          %two = arith.constant 2.0 : f32
          %result = arith.mulf %input_val, %two : f32
          memref.store %result, %output[%ix] : memref<10xf32>
          omp.yield
        }
      }
    omp.barrier
    omp.terminator
  }

  return
}

func.func private @main() {
  %input = memref.get_global @input : memref<10xf32>
  %output = memref.alloc() : memref<10xf32>

  call @kernel(%input, %output) : (memref<10xf32>, memref<10xf32>) -> ()

  %lb = index.constant 0
  %ub = index.constant 10
  %step = index.constant 1

  %fs = llvm.mlir.addressof @fmt : !llvm.ptr

  scf.for %iv = %lb to %ub step %step {
    %el = memref.load %output[%iv] : memref<10xf32>
    llvm.call @printf(%fs, %el) vararg(!llvm.func<i32 (ptr, ...)>) : (!llvm.ptr, f32) -> i32
  }

  return
}

SCF에서 OpenMP로의 변환 패스

-convert-scf-to-openmp 패스는 SCF 다이얼렉트의 병렬 루프를 동등한 OpenMP 연산으로 자동 변환합니다. 이 패스는 OpenMP 다이얼렉트 코드를 직접 작성하지 않고 OpenMP 실행을 대상으로 삼고자 할 때 특히 유용합니다.

예를 들어, 다음과 같은 SCF 병렬 루프는:

scf.parallel (%i) = (%c0) to (%c100) step (%c1) {
  %val = memref.load %A[%i] : memref<100xf32>
  %result = arith.addf %val, %val : f32
  memref.store %result, %B[%i] : memref<100xf32>
  scf.yield
}

OpenMP 연산으로 변환됩니다:

omp.parallel {
  omp.wsloop for (%i) : i32 = (%c0) to (%c100) step (%c1) {
    %val = memref.load %A[%i] : memref<100xf32>
    %result = arith.addf %val, %val : f32
    memref.store %result, %B[%i] : memref<100xf32>
    omp.yield
  }
  omp.terminator
}

이 패스는 루프 경계와 스텝 크기를 보존하면서 scf.parallel 연산을 omp.wsloop을 포함하는 omp.parallel 영역으로 변환하는 것을 비롯해 여러 중요한 변환을 처리합니다. 또한 SCF의 리덕션 연산은 OpenMP 리덕션으로 매핑되며, 중첩된 병렬 루프도 적절히 처리됩니다. 이 변환은 일반적으로 더 큰 로어링 파이프라인의 일부로 사용되며, 그 뒤에는 OpenMP 런타임 지원으로 실행할 수 있는 LLVM IR을 생성하기 위해 -convert-openmp-to-llvm이 이어지는 경우가 많습니다.

외부 자료

이전: 2부 - MLIR의 메모리 · 시리즈 목차 · 다음: 4부 - MLIR의 선형대수

원문 © Stephen Diehl, stephendiehl.com · 글 CC BY-SA 4.0, 코드 MIT · Stephen Diehl은 런던의 소프트웨어 엔지니어이자 작가입니다(저자가 정한 소개 문구). 이 글은 저자가 개인 자격으로 쓴 것이며 어떤 조직과의 연관도 밝히거나 암시해서는 안 됩니다 (라이선스와 고지).
이 페이지는 원문을 한국어로 옮기며 변경한 비공식 번역이며 같은 CC BY-SA 4.0으로 제공합니다 · yeokja 한국어 번역 · Anthropic 사의 claude-opus-5-5 모델을 활용하여 번역되었으며 학습을 모두 비허용한 상태로 작업하였습니다.