본문으로 건너뛰기

[triton] Triton GPU 최적화: 스레드 지역성 향상을 위한 Reduce 연산 개선

PR 링크: triton-lang/triton#11503 상태: Merged | 변경: +176 / -35

들어가며

GPU 컴퓨팅에서 Reduce 연산은 데이터 집계에 필수적인 연산입니다. 하지만 Reduce 연산의 구현 방식은 스레드 간의 데이터 접근 패턴과 메모리 대역폭 활용에 큰 영향을 미칠 수 있습니다. Triton은 이러한 Reduce 연산을 효율적으로 처리하기 위해 다양한 최적화를 적용합니다. 이번 글에서는 triton-lang/triton 레포지토리의 PR #11370 후속 작업인 "[OptimizeThreadLocality] Support rank-one, non-innermost-axis, and cross CTA reduce" PR을 분석하여, Reduce 연산의 스레드 지역성(Thread Locality)을 개선하고 성능을 향상시키는 코드 변경 사항들을 자세히 살펴보겠습니다.

이 PR은 기존 Triton 컴파일러가 Reduce 연산을 최적화할 때 적용했던 몇 가지 제약 조건을 완화하여, 더 넓은 범위의 Reduce 패턴에 대해 효율적인 코드를 생성하도록 개선합니다. 특히, rank-1 Reduce, 가장 안쪽 축이 아닌 축에서의 Reduce, 그리고 CTA(Cooperative Thread Array)를 넘나드는 Reduce 연산을 지원하는 데 중점을 둡니다.

코드 분석

이번 PR의 핵심 변경 사항은 lib/Dialect/TritonGPU/Transforms/OptimizeThreadLocality.cpp 파일에 집중되어 있습니다. 변경 사항을 파일별로 나누어 살펴보겠습니다.

OptimizeThreadLocality.cpp

이 파일은 Triton 커널의 스레드 지역성을 최적화하는 로직을 담고 있습니다. 이번 PR에서는 OptimizeThreadLocalityPass 내의 여러 함수와 관련 로직이 수정되었습니다.

1. matchAccumulatedReduce 함수의 제약 조건 완화

matchAccumulatedReduce 함수는 scf.for 루프 내에서 누적되는 Reduce 패턴을 식별하는 역할을 합니다. 기존에는 Reduce 연산이 가장 안쪽 축(rank - 1)에서만 수행되고, 텐서의 인코딩이 BlockedEncodingAttr이며 rank가 1보다 커야 한다는 제약 조건이 있었습니다. 이번 PR에서는 이러한 제약 조건들이 완화되었습니다.

Before:

-  if (!(isa<triton::gpu::BlockedEncodingAttr>(srcType.getEncoding()) &&
-        rank > 1))
-    return std::nullopt;
-  // The code currently assumes that the reduction is happening on the most
-  // inner dim.
-  if (axis != rank - 1)
+  if (!isa<triton::gpu::BlockedEncodingAttr>(srcType.getEncoding()))
     return std::nullopt;

After:

-  if (!(isa<triton::gpu::BlockedEncodingAttr>(srcType.getEncoding()) &&
-        rank > 1))
-    return std::nullopt;
-  // The code currently assumes that the reduction is happening on the most
-  // inner dim.
-  if (axis != rank - 1)
+  if (!isa<triton::gpu::BlockedEncodingAttr>(srcType.getEncoding()))
     return std::nullopt;

기존 코드에서는 rank > 1 조건과 axis != rank - 1 조건을 통해 Reduce 연산이 rank-1 축에서만 수행되어야 한다는 가정을 강제했습니다. 하지만 이 PR에서는 이러한 제약이 제거되었습니다. 이제 BlockedEncodingAttr만 만족하면 rank-1 축이 아니더라도, 즉 가장 안쪽 축이 아닌 다른 축에서도 Reduce 연산이 최적화 대상이 될 수 있습니다. 이는 rank-1 Reduce 뿐만 아니라, 더 일반적인 경우의 Reduce 연산에 대한 최적화를 가능하게 합니다.

2. incorporateOriginalAccumulatorValue 함수의 스칼라 처리 개선

incorporateOriginalAccumulatorValue 함수는 루프 내에서 계산된 Reduce 결과를 기존 누적 값과 결합하는 역할을 합니다. Rank-1 Reduce의 경우, Reduce 연산의 결과가 스칼라가 되는데, 기존에는 cloneWithInferType 유틸리티가 스칼라 타입을 제대로 처리하지 못하는 문제가 있었습니다.

Before:

-        Type destType = newLoop.getResult(argIdx).getType();
-        auto cvtLayout = createConvertLayout(builder, destType, newReduce2);
-        // incorporate the original accumulator value into the final result
-        auto finalOp = incorporateOriginalAccumulatorValue(builder, oldUpdate,
-                                                           cvtLayout, oldAccum);
+        // add convert_layout to get back to original layout, the result layout
+        // should now match the layout of the old accumulator (%init); a
+        // rank-one reduce yields a scalar, which needs no conversion
+        Type destType = newLoop.getResult(argIdx).getType();
+        Operation *newResult = newReduce2;
+        if (isa<RankedTensorType>(destType))
+          newResult = createConvertLayout(builder, destType, newReduce2);
+        // incorporate the original accumulator value into the final result
+        auto finalOp = incorporateOriginalAccumulatorValue(builder, oldUpdate,
+                                                           newResult, oldAccum);

After:

-    builder.setInsertionPointAfter(cvtLayout);
+    builder.setInsertionPointAfter(newResult);
     IRMapping mapping;
     mapping.map(oldUpdate->getOperand(0), oldAccum);
-    mapping.map(oldUpdate->getOperand(1), cvtLayout->getResult(0));
+    mapping.map(oldUpdate->getOperand(1), newResult->getResult(0));
     auto finalOp = cloneWithInferType(builder, &(*oldUpdate), mapping);
     return finalOp;

기존 코드에서는 createConvertLayout을 항상 호출하여 레이아웃 변환을 시도했습니다. 하지만 rank-1 Reduce의 경우 결과가 스칼라이므로 convert_layout이 불필요하며, 오히려 스칼라 타입을 제대로 처리하지 못하는 문제가 있었습니다. PR에서는 isa<RankedTensorType>(destType) 조건을 추가하여, 결과 타입이 텐서일 경우에만 createConvertLayout을 호출하도록 수정했습니다. 이를 통해 rank-1 Reduce 결과가 스칼라인 경우에도 올바르게 처리될 수 있습니다. 리뷰어 Jokeren은 이 변경에 대해 "why new update cannot be a scalar?"라고 질문했지만, he-weiwen은 rank-1 Reduce의 경우 부분 누적 값이 스칼라가 아닌 rank-1 텐서가 되어야 비-스레드 로컬 차원을 유지할 수 있다고 설명했습니다. 또한 cloneWithInferType이 스칼라 타입에서 실패하는 것을 해결하기 위해 직접 newUpdate를 빌드하는 방식이 더 간단하다고 언급했습니다.

3. createReduce 함수의 축 재정렬 로직 개선

createReduce 함수는 실제 Reduce 연산을 수행하기 위한 tt.reshape, tt.trans, tt.reshape 연산을 생성합니다. 이 과정에서 스레드 지역성을 최적화하기 위해 축(axis)의 순서를 재정렬합니다. 이번 PR에서는 CTA split, register wraparound, warp/lane, size per thread 간의 관계를 고려하여 축 재정렬 로직을 더욱 일반화했습니다.

Before (주요 변경 부분):

-    SmallVector<int64_t> factorShape(srcType.getShape().begin(),
-                                     srcType.getShape().end());
-    factorShape.back() = elemsPerThread / sizePerThread;
-    factorShape.push_back(dstShape[rank - 1]);
-    factorShape.push_back(sizePerThread);
-    SmallVector<int32_t> transposeOrder(rank + 2);
+    // [.., ctaSplit, R, H, S, ..] -> [.., ctaSplit, H, .., R, S]
+    //   -> [.., ctaSplit * H, .., R * S].
+    unsigned axis = reduce.getAxis();
+    int64_t ctaSplit =
+        std::min<int64_t>(getCTASplitNum(blocked)[axis], dstShape[axis]);
+    int64_t R = elemsPerThread / sizePerThread; // register wraparounds
+    int64_t H = dstShape[axis] / ctaSplit;      // warps & threads
+    SmallVector<int64_t> factorShape(srcType.getShape().begin(),
+                                     srcType.getShape().end());
+    // [.., ctaSplit, R, H, S, ..]
+    factorShape[axis] = ctaSplit;
+    factorShape.insert(factorShape.begin() + axis + 1, {R, H, sizePerThread});
+    SmallVector<int32_t> transposeOrder(rank + 3);
     std::iota(transposeOrder.begin(), transposeOrder.end(), 0);
-    std::swap(transposeOrder[rank - 1], transposeOrder[rank]);
+    // [.., ctaSplit, (R, H), S, ..] -> [.., ctaSplit, (H, R), S, ..]
+    std::swap(transposeOrder[axis + 1], transposeOrder[axis + 2]);
+    // [.., ctaSplit, H, (R, S), ..] -> [.., ctaSplit, H, .., (R, S)]
+    std::rotate(transposeOrder.begin() + axis + 2,
+                transposeOrder.begin() + axis + 4, transposeOrder.end());

After (주요 변경 부분):

-    SmallVector<int64_t> factorShape(srcType.getShape().begin(),
-                                     srcType.getShape().end());
-    factorShape.back() = elemsPerThread / sizePerThread;
-    factorShape.push_back(dstShape[rank - 1]);
-    factorShape.push_back(sizePerThread);
-    SmallVector<int32_t> transposeOrder(rank + 2);
+    // [.., ctaSplit, R, H, S, ..] -> [.., ctaSplit, H, .., R, S]
+    //   -> [.., ctaSplit * H, .., R * S].
+    unsigned axis = reduce.getAxis();
+    int64_t ctaSplit =
+        std::min<int64_t>(getCTASplitNum(blocked)[axis], dstShape[axis]);
+    int64_t R = elemsPerThread / sizePerThread; // register wraparounds
+    int64_t H = dstShape[axis] / ctaSplit;      // warps & threads
+    SmallVector<int64_t> factorShape(srcType.getShape().begin(),
+                                     srcType.getShape().end());
+    // [.., ctaSplit, R, H, S, ..]
+    factorShape[axis] = ctaSplit;
+    factorShape.insert(factorShape.begin() + axis + 1, {R, H, sizePerThread});
+    SmallVector<int32_t> transposeOrder(rank + 3);
     std::iota(transposeOrder.begin(), transposeOrder.end(), 0);
-    std::swap(transposeOrder[rank - 1], transposeOrder[rank]);
+    // [.., ctaSplit, (R, H), S, ..] -> [.., ctaSplit, (H, R), S, ..]
+    std::swap(transposeOrder[axis + 1], transposeOrder[axis + 2]);
+    // [.., ctaSplit, H, (R, S), ..] -> [.., ctaSplit, H, .., (R, S)]
+    std::rotate(transposeOrder.begin() + axis + 2,
+                transposeOrder.begin() + axis + 4, transposeOrder.end());

기존에는 reshape -> transpose -> reshape 과정에서 축의 순서를 고정적으로 처리했습니다. 하지만 이 PR에서는 ctaSplit, R (register wraparound), H (warps & lanes), S (size per thread)를 명시적으로 계산하고, 이를 기반으로 factorShape와 transposeOrder를 동적으로 생성합니다. 특히, transposeOrder 계산 시 std::swap과 std::rotate를 사용하여 ctaSplit, H, R, S 축을 원하는 순서로 재배열합니다. 리뷰어 he-weiwen은 "in the 2nd reshape, ctaSplit & H need to be combined into the same axis, or the output shape wouldn't work."라고 지적했으며, 이는 ctaSplit과 H를 하나의 축으로 결합해야 최종적인 efficient_layout으로 변환될 수 있음을 의미합니다. 이 변경은 CTA split과 register wraparound이 모두 복잡할 때 발생하던 cvtReordersRegisters assert 오류를 해결하고, 더 일반적인 형태의 Reduce 연산에 대한 최적화를 가능하게 합니다.

4. createPostLoopReduce 함수의 스칼라 처리

createPostLoopReduce 함수는 루프 이후에 수행될 Reduce 연산을 생성합니다. Rank-1 Reduce의 경우, 루프 내 Reduce 결과가 스칼라이므로 후속 convert_layout이 불필요합니다.

Before:

-        // add convert_layout to get back to original layout, the result layout
-        // should now match the layout of the old accumulator (%init)
-        Type destType = newLoop.getResult(argIdx).getType();
-        auto cvtLayout = createConvertLayout(builder, destType, newReduce2);
-        // incorporate the original accumulator value into the final result
-        auto finalOp = incorporateOriginalAccumulatorValue(builder, oldUpdate,
-                                                           cvtLayout, oldAccum);
+        // add convert_layout to get back to original layout, the result layout
+        // should now match the layout of the old accumulator (%init); a
+        // rank-one reduce yields a scalar, which needs no conversion
+        Type destType = newLoop.getResult(argIdx).getType();
+        Operation *newResult = newReduce2;
+        if (isa<RankedTensorType>(destType))
+          newResult = createConvertLayout(builder, destType, newReduce2);
+        // incorporate the original accumulator value into the final result
+        auto finalOp = incorporateOriginalAccumulatorValue(builder, oldUpdate,
+                                                           newResult, oldAccum);

After:

-        auto cvtLayout = createConvertLayout(builder, destType, newReduce2);
-        // incorporate the original accumulator value into the final result
-        auto finalOp = incorporateOriginalAccumulatorValue(builder, oldUpdate,
-                                                           cvtLayout, oldAccum);
+        Operation *newResult = newReduce2;
+        if (isa<RankedTensorType>(destType))
+          newResult = createConvertLayout(builder, destType, newReduce2);
+        // incorporate the original accumulator value into the final result
+        auto finalOp = incorporateOriginalAccumulatorValue(builder, oldUpdate,
+                                                           newResult, oldAccum);

이 변경은 위에서 설명한 incorporateOriginalAccumulatorValue 함수의 스칼라 처리 개선과 밀접하게 연관되어 있습니다. Rank-1 Reduce의 경우, newReduce2의 결과가 스칼라이므로 createConvertLayout을 건너뛰고 newResult로 직접 전달하여 incorporateOriginalAccumulatorValue 함수가 올바르게 작동하도록 합니다.

테스트 케이스 (test/TritonGPU/optimize-locality.mlir)

PR에는 다양한 시나리오에 대한 테스트 케이스가 포함되어 있습니다. 변경된 코드의 동작을 검증하기 위해 CHECK 지시어를 사용하여 예상되는 MLIR 출력을 명시하고 있습니다. 특히, rank-one reduce, reduce_axis_zero (non-innermost axis reduce) 등의 테스트 케이스에서 tt.reshape과 tt.trans 연산의 order 속성이 변경된 것을 확인할 수 있습니다. 예를 들어, 기존에는 order = array<i32: 0, 2, 1, 3> 였던 것이 order = array<i32: 0, 1, 3, 2, 4> 와 같이 변경되어, 축 재정렬 로직이 더 복잡한 경우를 지원함을 보여줍니다.

왜 이게 좋은가?

이번 PR은 Triton GPU 컴파일러의 Reduce 연산 최적화 기능을 크게 향상시킵니다. 주요 이점은 다음과 같습니다.

  1. 지원 범위 확대: Rank-1 Reduce, 가장 안쪽 축이 아닌 축에서의 Reduce, 그리고 CTA를 넘나드는 Reduce 연산까지 최적화 대상에 포함시킴으로써, 더 다양한 종류의 Triton 커널에서 성능 향상을 기대할 수 있습니다.
  2. 스레드 지역성 향상: Reduce 연산 시 데이터를 스레드 로컬 메모리 또는 레지스터에 최대한 가깝게 배치하고, CTA 간 통신을 최소화하도록 축을 재정렬합니다. 이는 메모리 접근 지연 시간을 줄이고 GPU의 연산 장치 활용률을 높여 전반적인 성능을 향상시킵니다.
  3. 컴파일러 견고성 증대: 기존에는 특정 조건에서만 작동하거나 assert 오류를 발생시키던 최적화 로직을 일반화하고 견고하게 만들어, 컴파일러의 안정성과 신뢰도를 높였습니다.

구체적인 성능 수치는 PR 자체에 명시되어 있지 않지만, 이러한 최적화는 일반적으로 Reduce 연산이 많은 커널에서 상당한 속도 향상(수 %에서 수십 %까지)을 가져올 수 있습니다. 일반적인 교훈은 다음과 같습니다:

  • 컴파일러 최적화의 제약 조건 완화: 실제 정확성이나 수익성에 문제가 없는 한, 컴파일러의 최적화 제약 조건을 점진적으로 완화하는 것은 더 많은 코드에 최적화를 적용할 수 있게 합니다.
  • 축 재정렬의 중요성: GPU에서 데이터 레이아웃과 축의 순서는 성능에 결정적인 영향을 미칩니다. reshape, transpose 연산을 통해 데이터를 효율적으로 재배열하는 것은 스레드 지역성을 극대화하는 핵심 기법입니다.
  • 엣지 케이스 처리: Rank-1 Reduce와 같이 특수한 경우나, CTA split과 register wraparound이 동시에 복잡하게 얽히는 경우와 같은 엣지 케이스를 올바르게 처리하는 것이 컴파일러의 완성도를 높이는 데 중요합니다.

결론

이번 PR은 Triton GPU 컴파일러가 Reduce 연산을 최적화하는 방식을 더욱 정교하고 강력하게 만들었습니다. Rank-1 Reduce, 비-내부 축 Reduce, 그리고 CTA 간 Reduce 연산에 대한 지원 확대를 통해, 개발자들은 더 넓은 범위의 커널에서 성능 향상을 누릴 수 있게 되었습니다. 이러한 최적화는 GPU 하드웨어의 특성을 깊이 이해하고, 컴파일러 기술을 통해 이를 효과적으로 활용하는 좋은 사례입니다.

References

참고 자료

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.

댓글

관련 포스트

PR Analysis 의 다른글