monorimet
commented
9 months ago The attached out.txt shows that TileAndDecomposeAttention was the last successful pass, and EliminateEmptyTensors fails first.
The IR after TileAndDecomposeAttention shows an address being reused many times in the innermost decomposed block of computation, which might be triggering a failure of OneShotAnalysis:
// -----// IR Dump After TileAndDecomposeWinogradTransform (iree-linalg-ext-tile-and-decompose-winograd) //----- //
func.func @main_dispatch_0_attention_20x4096x64xf16() {
%c4096 = arith.constant 4096 : index
%c20 = arith.constant 20 : index
%c0 = arith.constant 0 : index
%0 = hal.interface.binding.subspan set(0) binding(0) type(storage_buffer) alignment(64) offset(%c0) flags(ReadOnly) : !flow.dispatch.tensor<readonly:tensor<20x4096x64xf16>>
%1 = hal.interface.binding.subspan set(0) binding(1) type(storage_buffer) alignment(64) offset(%c0) flags(ReadOnly) : !flow.dispatch.tensor<readonly:tensor<20x4096x64xf16>>
%2 = hal.interface.binding.subspan set(0) binding(2) type(storage_buffer) alignment(64) offset(%c0) flags(ReadOnly) : !flow.dispatch.tensor<readonly:tensor<20x4096x64xf16>>
%3 = hal.interface.binding.subspan set(0) binding(3) type(storage_buffer) alignment(64) offset(%c0) : !flow.dispatch.tensor<writeonly:tensor<20x4096x64xf16>>
%workgroup_id_x = hal.interface.workgroup.id[0] : index
%workgroup_count_x = hal.interface.workgroup.count[0] : index
%workgroup_id_y = hal.interface.workgroup.id[1] : index
%workgroup_count_y = hal.interface.workgroup.count[1] : index
%4 = affine.apply affine_map<()[s0] -> (s0 * 20)>()[%workgroup_id_y]
%5 = affine.apply affine_map<()[s0] -> (s0 * 20)>()[%workgroup_count_y]
scf.for %arg0 = %4 to %c20 step %5 {
%6 = affine.apply affine_map<()[s0] -> (s0 * 64)>()[%workgroup_id_x]
%7 = affine.apply affine_map<()[s0] -> (s0 * 64)>()[%workgroup_count_x]
scf.for %arg1 = %6 to %c4096 step %7 {
%8 = flow.dispatch.tensor.load %3, offsets = [%arg0, %arg1, 0], sizes = [20, 64, 64], strides = [1, 1, 1] : !flow.dispatch.tensor<writeonly:tensor<20x4096x64xf16>> ->
tensor<20x64x64xf16>
%9 = flow.dispatch.tensor.load %0, offsets = [%arg0, %arg1, 0], sizes = [20, 64, 64], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<20x4096x64xf16>> ->
tensor<20x64x64xf16>
%10 = flow.dispatch.tensor.load %1, offsets = [%arg0, 0, 0], sizes = [20, 4096, 64], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<20x4096x64xf16>> ->
tensor<20x4096x64xf16>
%11 = flow.dispatch.tensor.load %2, offsets = [%arg0, 0, 0], sizes = [20, 4096, 64], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<20x4096x64xf16>> ->
tensor<20x4096x64xf16>
%12 = tensor.empty() : tensor<64x64xf32>
%extracted_slice = tensor.extract_slice %8[0, 0, 0] [1, 64, 64] [1, 1, 1] : tensor<20x64x64xf16> to tensor<64x64xf16>
%cst = arith.constant 0.000000e+00 : f32
%13 = linalg.fill ins(%cst : f32) outs(%12 : tensor<64x64xf32>) -> tensor<64x64xf32>
%cst_0 = arith.constant -1.000000e+30 : f32
%14 = tensor.empty() : tensor<64xf32>
%15 = linalg.fill ins(%cst_0 : f32) outs(%14 : tensor<64xf32>) -> tensor<64xf32>
%16 = tensor.empty() : tensor<64xf32>
%17 = linalg.fill ins(%cst : f32) outs(%16 : tensor<64xf32>) -> tensor<64xf32>
%c0_1 = arith.constant 0 : index
%c4096_2 = arith.constant 4096 : index
%c64 = arith.constant 64 : index
%18:3 = scf.for %arg2 = %c0_1 to %c4096_2 step %c64 iter_args(%arg3 = %13, %arg4 = %15, %arg5 = %17) -> (tensor<64x64xf32>, tensor<64xf32>, tensor<64xf32>) {
%extracted_slice_3 = tensor.extract_slice %10[0, %arg2, 0] [1, 64, 64] [1, 1, 1] : tensor<20x4096x64xf16> to tensor<64x64xf16>
%extracted_slice_4 = tensor.extract_slice %11[0, %arg2, 0] [1, 64, 64] [1, 1, 1] : tensor<20x4096x64xf16> to tensor<64x64xf16>
%extracted_slice_5 = tensor.extract_slice %9[0, 0, 0] [1, 64, 64] [1, 1, 1] : tensor<20x64x64xf16> to tensor<64x64xf16>
%cst_6 = arith.constant 0.000000e+00 : f32
%21 = tensor.empty() : tensor<64x64xf32>
%22 = linalg.fill ins(%cst_6 : f32) outs(%21 : tensor<64x64xf32>) -> tensor<64x64xf32>
%23 = linalg.matmul_transpose_b ins(%extracted_slice_5, %extracted_slice_3 : tensor<64x64xf16>, tensor<64x64xf16>) outs(%22 : tensor<64x64xf32>) -> tensor<64x64xf32>
%24 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0)>], iterator_types = ["parallel", "reduction"]} ins(%23 :
tensor<64x64xf32>) outs(%arg4 : tensor<64xf32>) {
^bb0(%in: f32, %out: f32):
%33 = arith.maximumf %in, %out : f32
linalg.yield %33 : f32
} -> tensor<64xf32>
%25 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0)>, affine_map<(d0, d1) -> (d0, d1)>], iterator_types = ["parallel", "parallel"]} ins(%24 : tensor<64xf32>)
outs(%23 : tensor<64x64xf32>) {
^bb0(%in: f32, %out: f32):
%33 = arith.subf %out, %in : f32
%34 = math.exp2 %33 : f32
linalg.yield %34 : f32
} -> tensor<64x64xf32>
%26 = linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>], iterator_types = ["parallel"]} ins(%24 : tensor<64xf32>) outs(%arg4 :
tensor<64xf32>) {
^bb0(%in: f32, %out: f32):
%33 = arith.subf %out, %in : f32
%34 = math.exp2 %33 : f32
linalg.yield %34 : f32
} -> tensor<64xf32>
%27 = linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>], iterator_types = ["parallel"]} ins(%26 : tensor<64xf32>) outs(%arg5 :
tensor<64xf32>) {
^bb0(%in: f32, %out: f32):
%33 = arith.mulf %in, %out : f32
linalg.yield %33 : f32
} -> tensor<64xf32>
%28 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0)>], iterator_types = ["parallel", "reduction"]} ins(%25 :
tensor<64x64xf32>) outs(%27 : tensor<64xf32>) {
^bb0(%in: f32, %out: f32):
%33 = arith.addf %in, %out : f32
linalg.yield %33 : f32
} -> tensor<64xf32>
%29 = tensor.empty() : tensor<64x64xf16>
%30 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>], iterator_types = ["parallel", "parallel"]} ins(%25 :
tensor<64x64xf32>) outs(%29 : tensor<64x64xf16>) {
^bb0(%in: f32, %out: f16):
%33 = arith.truncf %in : f32 to f16
linalg.yield %33 : f16
} -> tensor<64x64xf16>
%31 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0)>, affine_map<(d0, d1) -> (d0, d1)>], iterator_types = ["parallel", "parallel"]} ins(%26 : tensor<64xf32>)
outs(%arg3 : tensor<64x64xf32>) {
^bb0(%in: f32, %out: f32):
%33 = arith.mulf %in, %out : f32
linalg.yield %33 : f32
} -> tensor<64x64xf32>
%32 = linalg.matmul ins(%30, %extracted_slice_4 : tensor<64x64xf16>, tensor<64x64xf16>) outs(%31 : tensor<64x64xf32>) -> tensor<64x64xf32>
scf.yield %32, %24, %28 : tensor<64x64xf32>, tensor<64xf32>, tensor<64xf32>
}
%19 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0)>, affine_map<(d0, d1) -> (d0, d1)>], iterator_types = ["parallel", "parallel"]} ins(%18#2 : tensor<64xf32>)
outs(%18#0 : tensor<64x64xf32>) {
^bb0(%in: f32, %out: f32):
%cst_3 = arith.constant 1.000000e+00 : f32
%21 = arith.divf %cst_3, %in : f32
%22 = arith.mulf %21, %out : f32
linalg.yield %22 : f32
} -> tensor<64x64xf32>
%20 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>], iterator_types = ["parallel", "parallel"]} ins(%19 :
tensor<64x64xf32>) outs(%extracted_slice : tensor<64x64xf16>) {
^bb0(%in: f32, %out: f16):
%21 = arith.truncf %in : f32 to f16
linalg.yield %21 : f16
} -> tensor<64x64xf16>
%inserted_slice = tensor.insert_slice %20 into %8[0, 0, 0] [1, 64, 64] [1, 1, 1] : tensor<64x64xf16> into tensor<20x64x64xf16>
flow.dispatch.tensor.store %inserted_slice, %3, offsets = [%arg0, %arg1, 0], sizes = [20, 64, 64], strides = [1, 1, 1] : tensor<20x64x64xf16> ->
!flow.dispatch.tensor<writeonly:tensor<20x4096x64xf16>>
}
}
return
}
MaheshRavishankar
harsh-nod
hanhanW
erman-gurses
pashu123
What happened?
We have IR for SDXL that preserves the torch.aten._scaled_dot_product_attention op as tm_tensor.attention -> iree_linalg_ext.attention, but I'm seeing issues trying to lower the resulting IR through IREE-compiler.
@rsuderman managed to narrow the issue down to comprehensive bufferization of the tiled and decomposed attention op.
It seems that a buffer/address is being reused many times in the inner loop, causing bufferization/analysis to fail, though the error message given only mentions a mismatched yield / iterArg pair.
The minimized IR (minimal_attn.mlir):
The CLI input:
The error message:
The full log with --mlir-print-ir-after-all:
out.txt
What component(s) does this issue relate to?
iree-compiler, one-shot bufferization/analysis
Version information
Reproduced on latest IREE (0c61f77) and on a source build from contents of this PR: https://github.com/openxla/iree/pull/16416
Additional context
No response