MaheshRavishankar
commented
1 year ago I think the issue is the input linalg.generic
%3 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d1)>], iterator_types = ["reduction", "parallel"]} ins(%arg0 : tensor<4x32xi32>) outs(%cst : tensor<32xi32>) {
^bb0(%in: i32, %out: i32):
%4 = arith.addi %in, %out : i32
linalg.yield %4 : i32
} -> tensor<32xi32>is ill-formed. This should be split out as
%0 = tensor.empty()
%1 = linalg.fill (%cst : f32) outs(%0)
%2 = linalg.matmul ins(...) outs(%1)I would say its a front-end problem, but we can try to massage the IR within IREE to the right form...
@julianwa @jpienaar anyone who could pick this up?
benvanik
jpienaar
allieculp
I've noticed that constants get captured as
!flow.dispatch.tensor<readwrite...>. This has a lot of implications and is currently only barely working due to a combination of other things we do (like very early constant inlining) and may be hiding performance issues (we can't mark bindings as read-only when lowering to LLVM/SPIR-V/etc). This seems to have been introduced by the changes around dispatch region formation to first go through the region op and then go to the workgroups op.Root issue is that linalg ops specify constant values as outputs:
When forming the dispatch region the code looks at the linalg op, sees that %cst is an output, and marks it as readwrite:
The stream dialect correctly handles this by introducing a copy-on-write but this is the least efficient way of doing this kind of work. We should not be doing this in-place - doing so is a pessimization: it results in a read of the input + a write to the temporary + a full barrier + a read again on the inside of the dispatch + a write again of the output. So we end up needing to reserve the new memory for the output (what we need to do anyway) and then issue a copy from the constant memory to the temporary output that must complete prior to launching the dispatch. We don't see this in benchmarks much because the stream dialect schedules the copies concurrently with other dispatches but it increases memory consumption (something we don't track), code size (which we do but no one is looking at it), and adds memory bandwidth pressure that only shows as hurting dispatches around the dispatches doing this instead of being attributed to the dispatches that introduce the issue.
The IR above comes when disabling inlining of small constants during FormDispatchRegions, but if either the maximum size to inline is smaller or the tensor is larger than 256 bytes it'll look like the above today. This means that if you have a 100MB constant we'll reserve 100MB of memory, memcpy the whole constant over, and then do all those extra reads/writes to it in a very unlikely-to-be-cache-friendly manner.
Maybe there's a pass we could run to decouple linalg ops writing into constants? The above could be transformed into a tensor.empty as output and then just read the constant instead. Maybe this is related to the linalg decomposition stuff?
Since we barely manage to squeak by today for small constants with the current FormDispatchRegions constant inlining most people don't notice this in unit tests but it's blocking improvements to deduplication we need to reduce compile time and code size. It's also confounding other optimization work deeper down as with readwrite we can't add no-alias attributes during code generation and we're generating more code than we need to (all the additional loads/etc).