본문 바로가기

SASS_Probe

1차 연구 노트 : CUDA Source에서 SASS Primitive 로 보는 Operator Materialization

1. 연구 목적

이 실험의 목적은 CUDA source code 가 SASS 에서 어떤 primitive 로 내려가는지 관찰하고, 단순 instruction mapping 을 넘어 operator-level 최적화 구조를 읽는 것이다.

분석 프레임

  • CUDA source expression
    • compiler graph rewrite
    • SASS primitive
    • intermediate materialization 여부

초기 실험은 단순 memory/arithmetic/control kernel 에서 시작

이후 reduction 과 softmax 로 확장

 

2. 실험 환경

빌드 옵션으로 -O3, --use_fast_math, --generate-line-info 를 사용

 

3. 기본 memory skeleton

copy_global 에서 기본 1D elementwise kernel skeleton 을 확인

int i = blockIdx.x * blockDim.x + threadIdx.x;

if (i < n) {
    y[i] = x[i];
}

SASS role pattern:

S2R blockIdx
S2R threadIdx
IMAD index
ISETP.GE bounds
@P EXIT

IMAD.WIDE x address
LDG x value
IMAD.WIDE y address
STG y value

핵심 관찰

CUDA 의 if(i<n) 은 SASS 에서 i >= n 이면 EXIT 하는 조기 종료 형태로 나타남

 

4. arithmetic primitive

add_f32, mul_f32, fma_f32 를 통해 floating-point 산술 primitive 를 확인

특히 fma_f32 는 multiply 와 add 의 두 연ㅅ난이 FFMA.FTZ 하나로 fusion 된다.

즉 source-level 연산 개수와 SASS-level instruction 개수는 일치하지 않을 수 있다.

 

5. control expression 의 branchless lowering

relu_f32 와 clamp_f32 에서 조건식이 실제 branch 로 내려가지 않는 사례 확인

source-level if 가 항상 SASS-level branch 가 되는 것은 아니다.

짧은 값 선택 로직은 branchless select 또는 min/max primitive 로 내려갈 수 있다.

 

6. shared-memory reduction materialization

reduce_sum_f32 에서 naive shared-memory tree reduction 을 확인했다.

sdata[tid] = v;
__syncthreads();

for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
    if (tid < stride) {
        sdata[tid] += sdata[tid + stride];
    }

    __syncthreads();
}

SASS role pattern:

STS
BAR.SYNC

loop:
    LDS
    LDS
    FADD
    STS
    BAR.SYNC
    BRA

partial sum 은 register 에 장기 유지되지 않는다.

각 stride 마다 shared memory 에 materialize 된다.

즉 naive reduction 은 다음 구조로 볼 수 있다.

global memory

  • register
  • shared memory
  • register
  • shared memory
  • ...
  • global memory

 

7. softmax_small_f32 : materialized softmax baseline

shared memory 를 재사용한 baseline softmax

max_v = reduce_max(x);
e = exp(x[i] - max_v);
sum_v = reduce_sum(e);
y[i] = e / sum_v;

SASS role pattern:

initial:
    LDG or -FLT_MAX
    STS
    BAR.SYNC

max reduction:
    LDS
    LDS
    FSETP
    FSEL
    STS
    BAR.SYNC
    BRA

exp:
    FADD x - max
    FMUL * log2(e)
    MUFU.EX2
    STS
    BAR.SYNC

sum reduction:
    LDS
    LDS
    FADD
    STS
    BAR.SYNC
    BRA

normalize:
    LDS sum
    MUFU.RCP
    FMUL
    STG

핵심 관찰

max partials, exp values, sum partials 가 shared memory 에 materialize 된다.

특히 expf 는 다음으로 내려간다

  • expf(a)
  • a * log2(e)
  • MUFU.EX2

division 은 다음으로 내려간다

  • e / sum
  • e * rcp(sum)
  • MUFU.RCP + FMUL

 

8. online_softmax_f32 : register accumulator baseline

online_softmax_f32 는 shared memory softmax 와 다른 구조를 가진다.

new_m = max(m, v);

old_scale = expf(m - new_m);
new_term  = expf(v - new_m);

s = s * old_scale + new_term;
m = new_m;

SASS role pattern:

new_m:
    FSETP
    FSEL

old_scale:
    FADD
    FMUL * log2(e)
    MUFU.EX2

new_term:
    FADD
    FMUL * log2(e)
    MUFU.EX2

sum update:
    FFMA

핵심 관찰

running max m 과 running sum s 가 register accumulator 로 유지된다.

shared memory materialization 이 없ㄷ.

 

특히 다음 source expression

  • s = s * old_scale + new_term

은 SASS 에서

FFMA.FTZ

로 나타난다. 

따라서 online softmax update 는 다음과 같이 해석할 수 있다.

running max : FSETP + FSEL

running sum : MUFU.EX2 + FFMA

amterialization : register accumulator

 

9. softmax_small vs online_softmax

softmax_small_f32

x

  • shared memory
  • max reduction
  • shared memory
  • exp(x - max)
  • shared memory
  • sum reduction
  • shared memory
  • normalize
  • y

특징

STS / LDS / BAR.SYNC 중심

중간값이 shared memroy 에 저장됨

 

online_softmax_f32

x stream

  • register running max
  • register running sum
  • final normalize
  • y

특징

STS / LDS / BAR.SYNC ㅈ우심

중간값이 shared memory 에 저장됨

 

online_softmax_f32

x stream

  • register running max
  • register running sum
  • final normalize
  • y

특징

FSETP / FSEL / MUFU.EX2 / FFMA 중심

중간값이 register accumulator 에 유지됨

핵심 차이

softmax_small  

  • exp 전체 벡터를 shared memory 에 저장하고 sum reduction 수행

online_softmax

  • exp 전체 벡터를 저장하지 않고 running sum 을 update

 

10. 현재까지의 결론

1. 단순 elementwise 연산은 instruction primitive 로 깔끔하게 드러난다

copy : LDG - STG

add : LDG - LDG - FADD - STG

mul : LDG - LDG - FMUL - STG

fma : LDG - LDG - LDG - FFMA - STG

 

2. 조건식은 branch 가 아닐 수 있다.

relu : FMNMX

clamp : FSETP + FSEL

즉 source-level contol flow 가 SASS-level branch 로 유지된다고 가정하면 안된다.

 

3. reduction 은 materialization boundary 를 만든다.

reduce_sum : STS / LDS / BAR.SYNC / FADD / STS

shared memory 에 partial sum이 반복적으로 저장된다.

 

4. softmax 는 materialization 방식에 따라 SASS 구조가 크게 달라진다

softmax_small : shared-memory materialized softmax

online_softmax : register-accumulator streaming softmax

 

5. FlashAttention 식 online update 의 축소판은 SASS 에서 읽을 수 있다.

핵심 update

s = s * exp(old_m - new_m) + exp(v - new_m);

SASS primitive

MUFU.EX2

FFMA

shared memory 없음, register accumulator 유지

 

11. 연구적 의미

이번 실험의 가장 중요한 결과는 SASS 를 단순 명령어 사전으로 보는 것이 아니라. operator / layer 최적화 구조를 읽는 도구로 사용할 수 있다는 점이다.

특히 다음 구분이 가능해졌다.

  • 단순 arithmetic primitive
  • branchless control lowering
  • shared-memory materialized reduction
  • register-accumulator online update

이 구분은 이후 attention, layernorm, softmax fusion, FlashAttention 류 kernel 을 읽을 때 중요한 기준이 된다.

 

12. 다음 연구 방향

online softmax 를 attention toy kernel 로 확장

 

13. 1차 결론 문장

1차 SASS probing 실험의 결론은 다음과 같다

CUDA source expression 은 SASS 에서 단순히 1 :1 로 번역되지 않는다

컴파일러는 expression 을 hardware primitive 로 rewrite하고

그 결과는 instruction 종류뿐 아니라 intermediate materialization 

 

특히 softmax 계열에서 이 차이는 명확

baselien softmax : shared memory materialization pattern

online softmax : register accumulator update pattern

따라서 SASS 분석은 단순히 단순히 어떤 명령어가 나왔는가를 보는 수준을 넘어, operator 가 memory hieirarchy 위에서 어떻게 실행 구조로 변환되었는지 읽는 방향으로 확장될 수 있다.