ThomasRaoux
commented
1 year ago Interestingly enough the original IR looks simpler than the one we get at region formation time:
%11 = linalg.batch_matmul ins(%3, %10 : tensor<1x128x768xf32>, tensor<1x768x768xf32>) outs(%6 : tensor<1x128x768xf32>) -> tensor<1x128x768xf32>
%12 = linalg.generic {indexing_maps = [#map2, #map5, #map3], iterator_types = ["parallel", "parallel", "parallel"]} ins(%11, %cst_3 : tensor<1x128x768xf32>, tensor<768xf32>) outs(%2 : tensor<1x128x768xf32>) {
^bb0(%in: f32, %in_16: f32, %out: f32):
%47 = arith.addf %in, %in_16 : f32
linalg.yield %47 : f32
} -> tensor<1x128x768xf32>
%expanded = tensor.expand_shape %12 [[0], [1], [2, 3]] : tensor<1x128x768xf32> into tensor<1x128x12x64xf32>
%13 = tensor.empty() : tensor<1x12x128x64xf32>
%14 = linalg.generic {indexing_maps = [#map6, #map7], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded : tensor<1x128x12x64xf32>) outs(%13 : tensor<1x12x128x64xf32>) {
^bb0(%in: f32, %out: f32):
linalg.yield %in : f32
} -> tensor<1x12x128x64xf32>
%15 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel"]} ins(%cst_0 : tensor<768x768xf32>) outs(%0 : tensor<768x768xf32>) {
^bb0(%in: f32, %out: f32):
linalg.yield %in : f32
} -> tensor<768x768xf32>
%16 = linalg.generic {indexing_maps = [#map4, #map3], iterator_types = ["parallel", "parallel", "parallel"]} ins(%15 : tensor<768x768xf32>) outs(%4 : tensor<1x768x768xf32>) {
^bb0(%in: f32, %out: f32):
linalg.yield %in : f32
} -> tensor<1x768x768xf32>
%17 = linalg.batch_matmul ins(%3, %16 : tensor<1x128x768xf32>, tensor<1x768x768xf32>) outs(%6 : tensor<1x128x768xf32>) -> tensor<1x128x768xf32>Here it is very clear that we can fuse the batchmatmul (%11) with the bias (%12). Then the transpose could be fused into the following batchmatmul to create a batchmatmul_transpose_b. Right in FusionOfTensorOpsPass we will re-order the expandShape and linalg generic because we think we can't fuse the transpose with the following matmul. Then once we re-ordered the expandShape and generic there are no more opportunity to fuse.
One simple solution would be opportunistically fuse the transpose of rhs operand in the matmul, for the GPU this tends to be a slightly better layout anyway and I suspect it is the case on all platforms.
I had filed an issue about this when we were looking at Bert performance. We got stuck at the time as we weren't sure this was the right direction but I think we should revisit it so that for those cases we fuse the transpose(b) into the matmul, I added a named op for matmul_tranpose_b to support exactly this kind of cases.
MaheshRavishankar
stellaraccident
silvasean
nicolasvasilache
benvanik
mfuntowicz
allieculp
pjannaty
Young768
What happened?
Here is the IR for a typical attention layer of bert-base-uncased HuggingFace model with 1 sentence, sequence length 128: gist -- see bottom for the version with full weights. (here is the program that generates it) if that is of interest.
There are some missing fusion opportunities marked in the graph below (generated with --iree-flow-dump-dispatch-graph for the IR). For this particular workload, getting these fusions should be about 33% speedup of this workload (the unfused code occupies 25% of the execution time). The main issue is that fusions are not happening into the outputs of the batch_matmul ops. There are missing biasadd fusions and an output transpose/reshape too.
Steps to reproduce your issue
What component(s) does this issue relate to?
Compiler
Version information
Additional context
No response