Various enhancements and refactorings on caten/ajit #110

hikettei · 2024-09-20T08:50:53Z

Fix: Optimize gemm (no extra copy by !contiguous)
Fix: pass all tests
Scheduler: Fuse nested loops
- Merge inner (:LOCAL) Loop to get an optimal solution of Embedding
- 同一のタスクにScheduleする条件をゆるくした方がいいかもしれない (MultiExpr or buffer-intersect-p)
Fix: 4D Tensor ScaledDotProductAttention
Add: ConvND Testing
say goodbye to double corruption
Scheduler: ISLのASTでなるべくFuseしておく
IR: merge viewを考える
segv/randn w/ 2~3d inputs are failing
INDEX-COMPONENT is always scalar
stあるからone-shotでカーネル求めれないかのぅ
よくわかんなくなってきた
すぐ治したい: transpose matmulのtransposeがzero_costじゃない (Broadcastも...)
Improve the debugger visualization

hikettei · 2024-09-20T12:48:18Z

Embedding in a single kernel!!! (If I propagate :INDEX_COMPONENTS, it should be an optimal solution)
(Plus, fuse _gid2 and _gid3)

void main111178_e24_k0(float* val_35, float* val_54, boolean* val_48, float* val_31, float* val_37, float* val_41) {
  for(int _gid0=0;(_gid0<=9);_gid0+=1) {
    val_41[10*0+10*0+_gid0+0] = (10*0+10*0+1*_gid0+1*0);
    for(int _gid1=0;(_gid1<=25);_gid1+=4) {
      for(int _gid2=0;(_gid2<=9);_gid2+=1) {
        float val_45_0 = val_41[0+0+_gid2+0];
        float val_45_1 = val_41[0+0+_gid2+0];
        float val_45_2 = val_41[0+0+_gid2+0];
        float val_45_3 = val_41[0+0+_gid2+0];
        boolean val_47_0 = (val_45_0!=val_37[10*(_gid1+0)+_gid0+0+0]);
        boolean val_47_1 = (val_45_1!=val_37[10*(_gid1+1)+_gid0+0+0]);
        boolean val_47_2 = (val_45_2!=val_37[10*(_gid1+2)+_gid0+0+0]);
        boolean val_47_3 = (val_45_3!=val_37[10*(_gid1+3)+_gid0+0+0]);
        val_48[100*(_gid1+0)+10*_gid0+_gid2+0] = !val_47_0;
        val_48[100*(_gid1+1)+10*_gid0+_gid2+0] = !val_47_1;
        val_48[100*(_gid1+2)+10*_gid0+_gid2+0] = !val_47_2;
        val_48[100*(_gid1+3)+10*_gid0+_gid2+0] = !val_47_3;
      }
      for(int _gid2=0;(_gid2<=9);_gid2+=1) {
        val_54[100*(_gid1+0)+10*_gid0+10*0+_gid2] = 0.0;
        val_54[100*(_gid1+1)+10*_gid0+10*0+_gid2] = 0.0;
        val_54[100*(_gid1+2)+10*_gid0+10*0+_gid2] = 0.0;
        val_54[100*(_gid1+3)+10*_gid0+10*0+_gid2] = 0.0;
        for(int _gid3=0;(_gid3<=9);_gid3+=1) {
          val_35[1000*(_gid1+0)+100*_gid0+10*_gid3+_gid2] = (val_48[100*(_gid1+0)+10*_gid0+_gid3+0] ? val_31[0+0+10*_gid3+_gid2] : 0.0);
          val_35[1000*(_gid1+1)+100*_gid0+10*_gid3+_gid2] = (val_48[100*(_gid1+1)+10*_gid0+_gid3+0] ? val_31[0+0+10*_gid3+_gid2] : 0.0);
          val_35[1000*(_gid1+2)+100*_gid0+10*_gid3+_gid2] = (val_48[100*(_gid1+2)+10*_gid0+_gid3+0] ? val_31[0+0+10*_gid3+_gid2] : 0.0);
          val_35[1000*(_gid1+3)+100*_gid0+10*_gid3+_gid2] = (val_48[100*(_gid1+3)+10*_gid0+_gid3+0] ? val_31[0+0+10*_gid3+_gid2] : 0.0);
        }
        for(int _gid3=0;(_gid3<=9);_gid3+=1) {
          val_54[100*(_gid1+0)+10*_gid0+0+_gid2] = (val_54[100*(_gid1+0)+10*_gid0+0+_gid2]+val_35[1000*(_gid1+0)+100*_gid0+10*_gid3+_gid2]);
          val_54[100*(_gid1+1)+10*_gid0+0+_gid2] = (val_54[100*(_gid1+1)+10*_gid0+0+_gid2]+val_35[1000*(_gid1+1)+100*_gid0+10*_gid3+_gid2]);
          val_54[100*(_gid1+2)+10*_gid0+0+_gid2] = (val_54[100*(_gid1+2)+10*_gid0+0+_gid2]+val_35[1000*(_gid1+2)+100*_gid0+10*_gid3+_gid2]);
          val_54[100*(_gid1+3)+10*_gid0+0+_gid2] = (val_54[100*(_gid1+3)+10*_gid0+0+_gid2]+val_35[1000*(_gid1+3)+100*_gid0+10*_gid3+_gid2]);
        }
      }
    }
    for(int _gid1=28;(_gid1<=29);_gid1+=1) {
      for(int _gid2=0;(_gid2<=9);_gid2+=1) {
        float val_45 = val_41[0+0+_gid2+0];
        boolean val_47 = (val_45!=val_37[10*_gid1+_gid0+0+0]);
        val_48[100*_gid1+10*_gid0+_gid2+0] = !val_47;
      }
      for(int _gid2=0;(_gid2<=9);_gid2+=1) {
        val_54[100*_gid1+10*_gid0+10*0+_gid2] = 0.0;
        for(int _gid3=0;(_gid3<=9);_gid3+=1) {
          val_35[1000*_gid1+100*_gid0+10*_gid3+_gid2] = (val_48[100*_gid1+10*_gid0+_gid3+0] ? val_31[0+0+10*_gid3+_gid2] : 0.0);
        }
        for(int _gid3=0;(_gid3<=9);_gid3+=1) {
          val_54[100*_gid1+10*_gid0+0+_gid2] = (val_54[100*_gid1+10*_gid0+0+_gid2]+val_35[1000*_gid1+100*_gid0+10*_gid3+_gid2]);
        }
      }
    }
  }
}

hikettei · 2024-09-21T07:11:29Z

we have to revise the semantic of !reshape ...

(caten (!sin (!reshape (!sin (!reshape (make-tensor `(3 3)) `(9))) `(3 3))))

hikettei · 2024-09-21T09:40:02Z

JIT in Caten

(文章化してちゃんと考える ...)

リファクタしたい，一つのPolyhedral IRにつき一つCLOS Classを用意する

Polyhedral IR in Caten of Embedding, (gained by SERIALIZE=1).

(with-no-grad
    (caten (call (Embedding 10 10) (make-tensor `(10 30)))))
Compiled[e23]:
Compiled[e24]:

/*
Arrays:
  - val_54[float32]: (10 30 1 10) // OUTPUT, TMP
*/
void main2228060_e24_k0(float* val_54);
void main2228060_e24_k0(float* val_54) {
  for(int _gid0=0;(_gid0<=29);_gid0+=1) {
    for(int _gid1=0;(_gid1<=9);_gid1+=1) {
      for(int _gid2=0;(_gid2<=9);_gid2+=1) {
        val_54[300*_gid1+10*_gid0+10*0+_gid2] = 0.0;
      }
    }
  }
}

/*
Arrays:
  - val_41[float32]: (1 1 10 1) // IO, TMP
*/
void main2228060_e24_k1(float* val_41);
void main2228060_e24_k1(float* val_41) {
  for(int _gid0=0;(_gid0<=9);_gid0+=1) {
    val_41[10*0+10*0+_gid0+0] = (10*0+10*0+1*_gid0+1*0);
  }
}

/*
Arrays:
  - val_45[float32]: (10 30 10 1) // OUTPUT, TMP
  - val_41[float32]: (1 1 10 1) // INPUT, TMP
*/
void main2228060_e24_k2(float* val_45, float* val_41);
void main2228060_e24_k2(float* val_45, float* val_41) {
  for(int _gid0=0;(_gid0<=29);_gid0+=1) {
    for(int _gid1=0;(_gid1<=9);_gid1+=1) {
      for(int _gid2=0;(_gid2<=9);_gid2+=1) {
        val_45[300*_gid2+10*_gid0+_gid1+0] = val_41[0+0+_gid1+0];
      }
    }
  }
}

/*
Arrays:
  - val_47[bool]: (10 30 10 1) // OUTPUT, TMP
  - val_45[float32]: (10 30 10 1) // INPUT, TMP
  - val_37[float32]: (10 30 1 1) // INPUT, TMP
*/
void main2228060_e24_k3(boolean* val_47, float* val_45, float* val_37);
void main2228060_e24_k3(boolean* val_47, float* val_45, float* val_37) {
  for(int _gid0=0;(_gid0<=29);_gid0+=1) {
    for(int _gid1=0;(_gid1<=9);_gid1+=1) {
      for(int _gid2=0;(_gid2<=9);_gid2+=1) {
        val_47[300*_gid2+10*_gid0+_gid1+0] = (val_45[300*_gid2+10*_gid0+_gid1+0]!=val_37[30*_gid2+_gid0+0+0]);
      }
    }
  }
}

/*
Arrays:
  - val_48[bool]: (10 30 10 1) // OUTPUT, TMP
  - val_47[bool]: (10 30 10 1) // INPUT, TMP
*/
void main2228060_e24_k4(boolean* val_48, boolean* val_47);
void main2228060_e24_k4(boolean* val_48, boolean* val_47) {
  for(int _gid0=0;(_gid0<=29);_gid0+=1) {
    for(int _gid1=0;(_gid1<=9);_gid1+=1) {
      for(int _gid2=0;(_gid2<=9);_gid2+=1) {
        val_48[300*_gid2+10*_gid0+_gid1+0] = !val_47[300*_gid2+10*_gid0+_gid1+0];
      }
    }
  }
}

/*
Arrays:
  - val_35[float32]: (10 30 10 10) // OUTPUT, TMP
  - val_48[bool]: (10 30 10 1) // INPUT, TMP
  - val_31[float32]: (1 1 10 10) // INPUT, TMP
*/
void main2228060_e24_k5(float* val_35, boolean* val_48, float* val_31);
void main2228060_e24_k5(float* val_35, boolean* val_48, float* val_31) {
  for(int _gid0=0;(_gid0<=9);_gid0+=1) {
    for(int _gid1=0;(_gid1<=29);_gid1+=1) {
      for(int _gid2=0;(_gid2<=9);_gid2+=1) {
        for(int _gid3=0;(_gid3<=9);_gid3+=1) {
          val_35[3000*_gid0+100*_gid1+10*_gid2+_gid3] = (val_48[300*_gid0+10*_gid1+_gid2+0] ? val_31[0+0+10*_gid2+_gid3] : 0.0);
        }
      }
    }
  }
}

/*
Arrays:
  - val_54[float32]: (10 30 1 10) // IO, TMP
  - val_35[float32]: (10 30 10 10) // INPUT, TMP
*/
void main2228060_e24_k6(float* val_54, float* val_35);
void main2228060_e24_k6(float* val_54, float* val_35) {
  for(int _gid0=0;(_gid0<=9);_gid0+=1) {
    for(int _gid1=0;(_gid1<=29);_gid1+=1) {
      for(int _gid2=0;(_gid2<=9);_gid2+=1) {
        for(int _gid3=0;(_gid3<=9);_gid3+=1) {
          val_54[300*_gid0+10*_gid1+0+_gid2] = (val_54[300*_gid0+10*_gid1+0+_gid2]+val_35[3000*_gid0+100*_gid1+10*_gid3+_gid2]);
        }
      }
    }
  }
}

hikettei · 2024-09-24T12:22:29Z

!randnは修正するべき (OK)
- Dynamicだと通る，Staticだと通らない (OK)
mainとコンパイル結果が一致するように，特に: (OK)
- !softmax unrolling (packed-funcall, Loopが2つ以上でもUnroll)
- !log-softmax in-place (Post-MultiExpr) !matmul packing (:LOCAL :GLOBALの付与がおかしい？)
- !randn (Initial Scheduleが悪い？）(nth=199, is scheduled to T0 -> T2 -> T1)
  - TODO: Fuse val_14
everything is an initial schedule issue.
Complete Embedding Folding to pass randn test

TODO: Remaining Task

fix randn compiler
The way determining coincident (to unroll !softmax, !matmul properly)
Implement Post-MultiExpr to fuse !log-softmax and Embedding
Implement Index-Component Fusion
Delete Unused Node After PostMultiExpr

hikettei · 2024-09-25T05:53:01Z

Softmax (old vs new)

/*
Arrays:
  - B[uint32]: NIL // INPUT, SHAPE
  - A[uint32]: NIL // INPUT, SHAPE
  - val_34[float32]: (A B) // IO, TMP
*/
void main862852_e30_k0(uint32_t b, uint32_t a, float* val_34);
void main862852_e30_k0(uint32_t b, uint32_t a, float* val_34) {
  for(int _gid0=0;(_gid0<a);_gid0+=1) {
    float val_25 = 0.0;
    for(int _gid1=0;((_gid1+4)<=b);_gid1+=4) {
      val_25 = max(val_25, val_34[b*_gid0+(_gid1+0)]);
      val_25 = max(val_25, val_34[b*_gid0+(_gid1+1)]);
      val_25 = max(val_25, val_34[b*_gid0+(_gid1+2)]);
      val_25 = max(val_25, val_34[b*_gid0+(_gid1+3)]);
    }
    for(int _gid1=(b-(b%4));(_gid1<b);_gid1+=1) {
      val_25 = max(val_25, val_34[b*_gid0+_gid1]);
    }
    float val_26 = -(val_25);
    val_25 = 0.0;
    for(int _gid1=0;((_gid1+4)<=b);_gid1+=4) {
      val_34[b*_gid0+(_gid1+0)] = exp2(((val_34[b*_gid0+(_gid1+0)]+val_26)*1.442695));
      val_34[b*_gid0+(_gid1+1)] = exp2(((val_34[b*_gid0+(_gid1+1)]+val_26)*1.442695));
      val_34[b*_gid0+(_gid1+2)] = exp2(((val_34[b*_gid0+(_gid1+2)]+val_26)*1.442695));
      val_34[b*_gid0+(_gid1+3)] = exp2(((val_34[b*_gid0+(_gid1+3)]+val_26)*1.442695));
      val_25 = (val_25+val_34[b*_gid0+(_gid1+0)]);
      val_25 = (val_25+val_34[b*_gid0+(_gid1+1)]);
      val_25 = (val_25+val_34[b*_gid0+(_gid1+2)]);
      val_25 = (val_25+val_34[b*_gid0+(_gid1+3)]);
    }
    for(int _gid1=(b-(b%4));(_gid1<b);_gid1+=1) {
      val_34[b*_gid0+_gid1] = exp2(((val_34[b*_gid0+_gid1]+val_26)*1.442695));
      val_25 = (val_25+val_34[b*_gid0+_gid1]);
    }
    float val_11 = 1/(val_25);
    for(int _gid1=0;((_gid1+4)<=b);_gid1+=4) {
      val_34[b*_gid0+(_gid1+0)] = (val_34[b*_gid0+(_gid1+0)]*val_11);
      val_34[b*_gid0+(_gid1+1)] = (val_34[b*_gid0+(_gid1+1)]*val_11);
      val_34[b*_gid0+(_gid1+2)] = (val_34[b*_gid0+(_gid1+2)]*val_11);
      val_34[b*_gid0+(_gid1+3)] = (val_34[b*_gid0+(_gid1+3)]*val_11);
    }
    for(int _gid1=(b-(b%4));(_gid1<b);_gid1+=1) {
      val_34[b*_gid0+_gid1] = (val_34[b*_gid0+_gid1]*val_11);
    }
  }
}

/*
Arrays:
  - B[uint32]: NIL // INPUT, SHAPE
  - A[uint32]: NIL // INPUT, SHAPE
  - val_34[float32]: (A B) // IO, TMP
*/
void main736348_e30_k0(uint32_t b, uint32_t a, float* val_34);
void main736348_e30_k0(uint32_t b, uint32_t a, float* val_34) {
  for(int _gid0=0;(_gid0<a);_gid0+=1) {
    float val_11 = 0.0;
    float val_26 = 0.0;
    for(int _gid1=0;((_gid1+4)<=b);_gid1+=4) {
      val_26 = max(val_26, val_34[((b*_gid0)+(_gid1+0))]);
      val_26 = max(val_26, val_34[((b*_gid0)+(_gid1+1))]);
      val_26 = max(val_26, val_34[((b*_gid0)+(_gid1+2))]);
      val_26 = max(val_26, val_34[((b*_gid0)+(_gid1+3))]);
    }
    for(int _gid1=(b-(b%4));(_gid1<b);_gid1+=1) {
      val_26 = max(val_26, val_34[((b*_gid0)+_gid1)]);
    }
    val_26 = -(val_26);
    for(int _gid1=0;((_gid1+4)<=b);_gid1+=4) {
      val_34[((b*_gid0)+(_gid1+0))] = exp2(((val_34[((b*_gid0)+(_gid1+0))]+val_26)*1.442695));
      val_34[((b*_gid0)+(_gid1+1))] = exp2(((val_34[((b*_gid0)+(_gid1+1))]+val_26)*1.442695));
      val_34[((b*_gid0)+(_gid1+2))] = exp2(((val_34[((b*_gid0)+(_gid1+2))]+val_26)*1.442695));
      val_34[((b*_gid0)+(_gid1+3))] = exp2(((val_34[((b*_gid0)+(_gid1+3))]+val_26)*1.442695));
    }
    for(int _gid1=(b-(b%4));(_gid1<b);_gid1+=1) {
      val_34[((b*_gid0)+_gid1)] = exp2(((val_34[((b*_gid0)+_gid1)]+val_26)*1.442695));
    }
    for(int _gid1=0;((_gid1+4)<=b);_gid1+=4) {
      val_11 = (val_11+val_34[((b*_gid0)+(_gid1+0))]);
      val_11 = (val_11+val_34[((b*_gid0)+(_gid1+1))]);
      val_11 = (val_11+val_34[((b*_gid0)+(_gid1+2))]);
      val_11 = (val_11+val_34[((b*_gid0)+(_gid1+3))]);
    }
    for(int _gid1=(b-(b%4));(_gid1<b);_gid1+=1) {
      val_11 = (val_11+val_34[((b*_gid0)+_gid1)]);
    }
    val_11 = 1/(val_11);
    for(int _gid1=0;((_gid1+4)<=b);_gid1+=4) {
      val_34[((b*_gid0)+(_gid1+0))] = (val_34[((b*_gid0)+(_gid1+0))]*val_11);
      val_34[((b*_gid0)+(_gid1+1))] = (val_34[((b*_gid0)+(_gid1+1))]*val_11);
      val_34[((b*_gid0)+(_gid1+2))] = (val_34[((b*_gid0)+(_gid1+2))]*val_11);
      val_34[((b*_gid0)+(_gid1+3))] = (val_34[((b*_gid0)+(_gid1+3))]*val_11);
    }
    for(int _gid1=(b-(b%4));(_gid1<b);_gid1+=1) {
      val_34[((b*_gid0)+_gid1)] = (val_34[((b*_gid0)+_gid1)]*val_11);
    }
  }
}

This reverts commit d3f9fdf.

hikettei added 9 commits September 20, 2024 17:50

wip

8a88278

needs refactor for isl ...

999c36e

still misunderstanding waw/war/raw?

443f6a8

wip

49ae1b1

refactor: get 4d matmul working

658e328

fixed?

c64f8eb

serial t

5c7688d

a little tweak for double free?

7aa5044

Embedding in a single kernel

349b96e

fix syntax error

4501a5c

hikettei changed the title ~~[WIP] Fix JIT~~ Fix Scheduler Sep 20, 2024

hikettei added 6 commits September 20, 2024 22:14

update

f8f5e41

a lil tweak (normal, randn still wont work)

8ce740f

cannot permute assertion

8103234

spacing

85faa26

fix around scalar handling

4eb1060

.

495ca35

hikettei added 3 commits September 21, 2024 16:14

hotfix: fix typo

61333e7

fuse w/ prev-rank

d86a183

disable manual loop fusion for now

d1c6300

hikettei added 7 commits September 22, 2024 14:47

progresses on refactoring

74be6d3

wip

5297644

hmmmmm

7433c11

hmm

f8f2d97

rem

b62935d

Enhancement: EXPR and EXPR Simplifier

57a898b

MultiExpr Simplify

2e655de

hikettei added 8 commits September 24, 2024 20:17

passes on jit

36abc34

18 and 1 3 are not mergeable

5920a5b

skip ConvND

e9ebab7

Skip ConvND

b500d41

skip?

7af3f71

padding schedule idx

c338ad2

logsoftmax is not a inplace?

f09747e

2d+1 padiding

9fcafc5

hikettei added 10 commits September 24, 2024 21:44

unroll anywhere

e74e87d

i think we should abandon this branch ...

f546d7e

i think we should abandon this branch ...

ccd5f02

wip

ff0d5ac

Fix: memory-planner (different views are not merged)

0580760

fix memory corruption

943bbcc

wip: no transpose?

6562656

del comment

e5a9cb3

Fix: Unrolling Softmax

014c605

remove context-handle

f821c15

hikettei added 4 commits September 25, 2024 15:15

Fix: Unrollment especially when transposed

190fea4

regression tests

f05f29c

giving up index-component folding for a now

813c6c5

clean up codes

d05bcfa

hikettei marked this pull request as ready for review September 25, 2024 08:59

hikettei added 3 commits September 25, 2024 18:13

update test

1113308

no padding?

af18f26

need pads

4674966

hikettei merged commit d3f9fdf into main Sep 25, 2024
2 of 4 checks passed

hikettei added a commit that referenced this pull request Sep 25, 2024

Revert "Various enhancements and refactorings on caten/ajit (#110)"

91641a1

This reverts commit d3f9fdf.

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Various enhancements and refactorings on caten/ajit #110

Various enhancements and refactorings on caten/ajit #110

hikettei commented Sep 20, 2024 •

edited

Loading

hikettei commented Sep 20, 2024 •

edited

Loading

hikettei commented Sep 21, 2024

hikettei commented Sep 21, 2024 •

edited

Loading

hikettei commented Sep 24, 2024 •

edited

Loading

hikettei commented Sep 25, 2024

Various enhancements and refactorings on caten/ajit #110

Various enhancements and refactorings on caten/ajit #110

Conversation

hikettei commented Sep 20, 2024 • edited Loading

hikettei commented Sep 20, 2024 • edited Loading

hikettei commented Sep 21, 2024

hikettei commented Sep 21, 2024 • edited Loading

JIT in Caten

hikettei commented Sep 24, 2024 • edited Loading

TODO: Remaining Task

hikettei commented Sep 25, 2024

hikettei commented Sep 20, 2024 •

edited

Loading

hikettei commented Sep 20, 2024 •

edited

Loading

hikettei commented Sep 21, 2024 •

edited

Loading

hikettei commented Sep 24, 2024 •

edited

Loading