Почему -affine-scalrep не может удалить избыточную загрузку/сохранение во внутреннем цикле MLIR? ⇐ C++

[Anonymous]

Я пытаюсь оптимизировать ядро, подобное GEMM, используя аффинный диалект MLIR. Внутри моего самого внутреннего цикла (%arg5) у меня есть избыточные affine.load и affine.store на %alloc.
Несмотря на выполнение прохода -affine-scalrep (скалярная замена), эти операции с памятью не переводятся в регистры/значения SSA. Есть ли конкретная причина, по которой scalrep может игнорировать это, или есть другой проход, который я должен использовать, чтобы устранить эти накладные расходы?
module {
func.func @main(%arg0: memref, %arg1: memref, %arg2: memref) -> memref attributes {llvm.emit_c_interface} {
%cst = arith.constant 0.000000e+00 : f32
%alloc = memref.alloc() {alignment = 64 : i64} : memref
%alloc_0 = memref.alloc() {alignment = 64 : i64} : memref
affine.for %arg3 = 0 to 8 {
affine.for %arg4 = 0 to 64 {
%0 = affine.load %arg2[0, %arg4] : memref
affine.store %0, %alloc[%arg3, 0, %arg4] : memref
affine.for %arg5 = 0 to 16 {
%3 = affine.load %arg0[%arg3, %arg5] : memref
%4 = affine.load %arg1[%arg5, %arg4] : memref
%5 = affine.load %alloc[%arg3, 0, %arg4] : memref
%6 = arith.mulf %3, %4 : f32
%7 = arith.addf %5, %6 : f32
affine.store %7, %alloc[%arg3, 0, %arg4] : memref
}
%1 = affine.load %alloc[%arg3, 0, %arg4] : memref
%2 = arith.maximumf %1, %cst : f32
affine.store %2, %alloc_0[%arg3, %arg4] : memref
}
}
memref.dealloc %alloc : memref
return %alloc_0 : memref
}
}


Подробнее здесь: https://stackoverflow.com/questions/798 ... ir-inner-l