Input Fusion

[yiqian1]

Background

Regularization

An important pre-fusion stage is the regularization pass which rewrites the code into fusion-friendly form.

For this code, when recomputing the coordinates for indexing %out_view, it's hard (impossible) to get the coordinates of %write_view and %alloc. Regularization rewrites

// pull the reader's view
%alloc = alloc(): T
%write_view = rock.transform %alloc: T->TW
gemm() -> %write_view
%read_view = rock.transform %alloc : T->TR
%out_view = rock.transform %out: TO->TR
linalg.generic(%read_view) -> %out_view

as

// push the writer's view
%read_view = alloc(): TR
%alloc_view = rock.tranform %read_view: TR->T
%write_view = rock.transform %alloc_view : T-> TW
gemm() ->%write_view
%out_view = rock.transform %out: TO->TR
linalg.generic(%read_view) -> %out_view

and then:

%read_view = alloc(): TR
%alloc_view = rock.tranform %read_view: TR->T
%write_view = rock.transform %alloc_view : T-> TW
gemm() ->%write_view
%out_view = rock.transform %out: TO->TR
%out_view_T = rock.transform %out_view: TR->T
%out_view_transform = rock.transform %out_view_T: T->TW
linalg.generic(%write_view) -> %out_view_transform

Input fusion:

%alloc = alloc(): T
linalg.generic(%argA) -> %alloc  // type cast: f16->f32
%A_view = rock.transform %alloc
%B_view = rock.transform %argB
gemm(%A_view, %B_view) -> %gemm_out

After blockwise-to-threadwise:

%alloc = alloc(): T
linalg.generic(%argA) -> %alloc
%A_view = rock.transform %alloc
%B_view = rock.transform %argB
%generic_out = rock.transform %A_view
…
rock.threadwise_read_into (%generic_out, …): -> %gemm_argA_tiled (memref<2xf32>) 

After fusion:

%argA_tiled = alloc(): (memref<2xf16, #gpu.address_space<private>>)
%argA_tr = rock.transform %argA: 
%argA_tr_tr = rock.transform %argA_tr
threadwise_read_into(%argA_tr_tr) -> %argA_tiled (memref<2xf16>)
linalg.generic (%argA_tiled) -> %generic_tiled (memref<2xf32>)

Problems and potential solutions with input fusion

Vectorization

When linalg.generic has multiple inputs, the major issue is how to determine which input needs to be vectorized during gridwise-gemm-to-blockwise.

Potential solutions:
1.At MiGraphX or Tosa level pick the first largest function argument as the primary tensor that will be vectorized
2. annotate it and pass through to the regularization which will mark which linalg.generic input leads to the primary tensor.

Update regularization

Secondly we need to rewrite the regularization to push linalg.generic output view to gemm's input.

%alloc = alloc() :T
%out_view = rock.transform %alloc: T->TO
Linalg.generic(%arg1_view, %arg2_view) -> %out_view
%gemm_A = rock.transform %alloc : T->TA
gemm(%gemm_A,  %gemm_B) -> %gemm_out 

Instead of

%A = alloc() : TA
%out_view_inv = rock.transform %A: TA->T
%out_view = rock.transform %out_view_inv: T->TO
linalg.generic(%arg1_view, %arg2_view) -> %out_view
gemm(%A, %gemm_B)->%gemm_out

do this transformation:

%out = alloc(): TO
Linalg.generic(%arg1_view, %arg2_view) -> %out
%gemm_A_inv = rock.transform %out: TO->T
%gemm_A = rock.transform %gemm_A_inv : T->TA
gemm(%gemm_A,  %gemm_B) -> %gemm_out