Loop fusion does not work in my case

MLIR
1

I have the following IR. For some reason, LoopFusion does not work for it.
Do you have any ideas why?
Thanks!

  func @f(%arg0: !arrow.record_batch) -> !arrow.record_batch {
      %c5_i64 = constant 5 : i64
      %0 = "arrow.get_column"(%arg0) {columnName = "a"} : (!arrow.record_batch) -> !arrow.array<i64>
      %1 = "arrow.get_rows_count"(%arg0) : (!arrow.record_batch) -> index
      %2 = alloc(%1) : memref<?xi1>
      %3 = alloc(%1) : memref<?xi64>
      %4 = alloc(%1) : memref<?xi64>
      %5 = "arrow.get_data_buffer"(%0) : (!arrow.array<i64>) -> memref<?xi64>
      %6 = "arrow.get_null_bitmap"(%0) : (!arrow.array<i64>) -> memref<?xi1>
      affine.for %arg1 = 0 to %1 {
        %9 = affine.load %5[%arg1] : memref<?xi64>
        %10 = addi %9, %c5_i64 : i64
        affine.store %10, %4[%arg1] : memref<?xi64>
      }
      affine.for %arg1 = 0 to %1 {
        %9 = affine.load %4[%arg1] : memref<?xi64>
        %10 = muli %9, %c5_i64 : i64
        affine.store %10, %3[%arg1] : memref<?xi64>
      }
      affine.for %arg1 = 0 to %1 {
        %9 = affine.load %3[%arg1] : memref<?xi64>
        %10 = cmpi "sge", %9, %c5_i64 : i64
        affine.store %10, %2[%arg1] : memref<?xi1>
      }
      %7 = "arrow.make_array"(%6, %2) : (memref<?xi1>, memref<?xi1>) -> !arrow.array<i1>
      %8 = "arrow.make_recordbatch"(%7) : (!arrow.array<i1>) -> !arrow.record_batch
      dealloc %4 : memref<?xi64>
      dealloc %3 : memref<?xi64>
      return %8 : !arrow.record_batch
    }
2

I think it’s because your bounds aren’t constant. It’s easy to extend the pass to non-constant bounds but it hasn’t been done yet. Please try replacing %1 with a constant to confirm.

1 Like
3

I see, thanks! It works with constant bounds. I’ll try to extend the pass then.