LoopVariable manipulation in affine::buildAffineLoopNest

I want to lower a custom operation in toy dialect example – toy.delay operation to affine dialect. Here, the operation is just a shift operation – Ex- if Input is [ 10,20,30] and 2 , then output will be [0 ,0, 10] ie, shift by 2 units.
To perform this – the steps I am performing is building 2 loops – 1 for constant 0 and other for loading from the input array/buffer –
> //Input is

%0 = toy.constant dense<[1.000000e+01, 2.000000e+01, 3.000000e+01]> : tensor<3xf64>
%1 = toy.constant dense<2.000000e+00> : tensor<f64>
%2 = "toy.delay"(%0, %1) : (tensor<6xf64>, tensor<f64>) -> tensor<3xf64>
//from 0 to delayOp 2nd argumentt 
affine.for %arg0 = 0 to %cst ie 2 {
    %0 = affine.load %alloc[] : memref<6xf64>
    affine.store %0, %out_alloc_6[%arg0] : memref<6x1xf64>
}

//another loop for loading from input buffer
affine.for %arg0 = 0 to 1 {
%0 = affine.load %alloc_5[%arg0] : memref<3xf64>
%1 = arith.const 2 : i64
%2 = arith.add %arg0 , %1 : i64
%3 = arith.index %2: index
affine.store %0, %alloc[%2] : memref<3xf64>
}

Can I get some help on how can I perform this particular operation – affine.store %0, %alloc*[%arg0 + 2 ]* : memref<3xf64> something like this as shown below ?

affine.for %arg0 = 0 to 1 {
  %0 = affine.load %alloc_6[%arg0 ] : memref<3xf64>
  affine.store %0, %alloc[%arg0 + 2 ] : memref<3xf64>
}  

Basically, I need some help with affine::buildAffineLoopNest(rewriter, loc, lowerBounds, upperBounds, steps, [&](OpBuilder &nestedBuilder, Location loc, ValueRange ivs) – how can I work with ivs here in a proper way?

The ivs argument is your list of IVs, which would simply be the %arg0 in your example. You can use that IV along with a map to pass the map + operands for the affine load and store in your loop body, i.e., use rewriter.create<...>(...) for each of those two ops.

1 Like