Unable to reconstruct for loop in lambda function

[PietroGhg]

Hello, I noticed that the CanonicalizeFor pass is unable to reconstruct the for loop from this code:

struct A{
  int end;
  int step;
  A(int end, int step) : end(end), step(step) {}
};
void use(int);

void f(int step, int end){
  A a(end, step);
  auto l = [a](){
    for(int i = 0; i < a.end; i += a.step){
      use(i);
    }
  };
  l();
}

Running with cgeist -O3 -function='*' loop.cpp -S -o loop.mlir, it generates this MLIR:

  func.func @_Z1fii(%arg0: i32, %arg1: i32) attributes {llvm.linkage = #llvm.linkage<external>} {
    %c0_i32 = arith.constant 0 : i32
    %c1_i64 = arith.constant 1 : i64
    %0 = llvm.alloca %c1_i64 x !llvm.struct<(struct<(i32, i32)>)> : (i64) -> !llvm.ptr<struct<(struct<(i32, i32)>)>>
    %1 = llvm.alloca %c1_i64 x !llvm.struct<(struct<(i32, i32)>)> : (i64) -> !llvm.ptr<struct<(struct<(i32, i32)>)>>
    %2 = llvm.alloca %c1_i64 x !llvm.struct<(struct<(i32, i32)>)> : (i64) -> !llvm.ptr<struct<(struct<(i32, i32)>)>>
    %3 = llvm.alloca %c1_i64 x !llvm.struct<(struct<(i32, i32)>)> : (i64) -> !llvm.ptr<struct<(struct<(i32, i32)>)>>
    %4 = llvm.getelementptr %2[%c0_i32, 0] : (!llvm.ptr<struct<(struct<(i32, i32)>)>>, i32) -> !llvm.ptr<struct<(i32, i32)>>
    %5 = llvm.getelementptr %4[%c0_i32, 0] : (!llvm.ptr<struct<(i32, i32)>>, i32) -> !llvm.ptr<i32>
    llvm.store %arg1, %5 : !llvm.ptr<i32>
    %6 = llvm.getelementptr %4[%c0_i32, 1] : (!llvm.ptr<struct<(i32, i32)>>, i32) -> !llvm.ptr<i32>
    llvm.store %arg0, %6 : !llvm.ptr<i32>
    %7 = llvm.load %2 : !llvm.ptr<struct<(struct<(i32, i32)>)>>
    llvm.store %7, %1 : !llvm.ptr<struct<(struct<(i32, i32)>)>>
    call @_ZZ1fiiEN3$_0C1EOS_(%3, %1) : (!llvm.ptr<struct<(struct<(i32, i32)>)>>, !llvm.ptr<struct<(struct<(i32, i32)>)>>) -> ()
    %8 = llvm.load %3 : !llvm.ptr<struct<(struct<(i32, i32)>)>>
    llvm.store %8, %0 : !llvm.ptr<struct<(struct<(i32, i32)>)>>
    call @_ZZ1fiiENK3$_0clEv(%0) : (!llvm.ptr<struct<(struct<(i32, i32)>)>>) -> ()
    return
  }

...

  func.func private @_ZZ1fiiENK3$_0clEv(%arg0: !llvm.ptr<struct<(struct<(i32, i32)>)>>) attributes {llvm.linkage = #llvm.linkage<internal>} {
    %c0_i32 = arith.constant 0 : i32
    %0 = llvm.getelementptr %arg0[%c0_i32, 0] : (!llvm.ptr<struct<(struct<(i32, i32)>)>>, i32) -> !llvm.ptr<struct<(i32, i32)>>
    %1 = llvm.getelementptr %0[%c0_i32, 0] : (!llvm.ptr<struct<(i32, i32)>>, i32) -> !llvm.ptr<i32>
    %2 = llvm.getelementptr %0[%c0_i32, 1] : (!llvm.ptr<struct<(i32, i32)>>, i32) -> !llvm.ptr<i32>
    %3 = scf.while (%arg1 = %c0_i32) : (i32) -> i32 {
      %4 = llvm.load %1 : !llvm.ptr<i32>
      %5 = arith.cmpi slt, %arg1, %4 : i32
      scf.condition(%5) %arg1 : i32
    } do {
    ^bb0(%arg1: i32):
      func.call @_Z3usei(%arg1) : (i32) -> ()
      %4 = llvm.load %2 : !llvm.ptr<i32>
      %5 = arith.addi %arg1, %4 : i32
      scf.yield %5 : i32
    }
    return
  }

The scf.while can't be raised to a scf.for in this case because the step and end are loaded from the struct pointer, and so they are not loop invariants. This could (I think) be solved by inlining the function call, but the LLVMIR MLIR dialect doesn't seem to implement the DialectInlinerInterface, and so the MLIR inliner pass bails out. Why do Polygeist emits so many LLVMIR operations in this sample? Is there any flag that I can set to prevent that? What are your thought about this? Thanks, Pietro

[chelini]

Can you not load step and end in local variables?

[chelini]

struct A{
  int end;
  int step;
  A(int end, int step) : end(end), step(step) {}
};
void use(int);

void f(int step, int end){
  A a(end, step);
  auto l = [a](){
    int cst_end = a.end;
    int cst_step = a.step;
    for(int i = 0; i < cst_end; i += cst_step){
      use(i);
    }
  };
  l();
}

[PietroGhg]

Hi, thanks for your reply. Yes of course I can do that, and it works nicely on this simple code snippet, but this snippet is the narrowed-down version of more complex code that I'm trying to analyse, and in general "changing the input code" is not the ideal solution :). But again, using local variables is an alright temporary solution for me. I don't know if there's a plan for the LLVMIR dialect to implement the inliner interface, but this will probably be solved for free if/when it happens. I opened this issue just to have a brief discussion on this behaviour, but feel free to close it if you feel that no further actions are needed, thanks :)