Affine buildloopNest indices and non-constant values

Beginners
,
1

This is a beginner level question. So, I have to implement toy.delay operation with 2 input arguments, a vector and a single tensor. Here, this operation is just a shift in data ie, for a vector of size 3 – [10, 20, 30 ] , and unit = 2, output will be [0 , 0, 10].

Now, when I am lowering it to affine loops – what should be the right way to buildAffineLoopNest when the upper index ex- %5 below is coming from an operation and not an attribute? I am assuming this has to be converted to arith::ConstantIndexOp & for this, we need an integer constant value but this is not a constant but rather a SSA value which can’t be casted to IntegerAttr –

%3 = toy.add %1, %2 : tensor<f64>
%4 = arith.fptoui %3 : ui
%5 = arith.ConstanIndex %4 : index
affine.for %arg0 = 0 to %5 {
  affine.store %cst,%alloc[%arg0]:memref<3xf64>  
}

Also, for the affine-map , can we use SSA value and not a constant value ie,
#map2 = affine_map<(d0) -> (d0 + %4)>

Is this allowed in MLIR to use SSA values for affine-loop indices and affine-map? If yes, what is the correct sequence from say, tensor<f64> to affine-loop index and affine-map? I have tried tensor<f64> → arith.ui → arith.index but I am getting the error : error: failed to materialize conversion for result #0 of operation ‘toy.constant’ that remained live after conversion . Any particular example code would be really helpful.

2

There are multiple parts in this question:

  1. Yes, it is certainly possible to use SSA values as dimensions or symbols in affine expressions that appear in affine.for bounds, but there are documented restrictions as to which values are accepted. The verifier will fail for the IR that does not respect these restrictions.

  2. buildAffineLoopNest is a family of utility functions that simplify the common cases when the bound is either a constant or a single value that is fed into the identity map as loop bounds. To use a more complex bound, you can resort to using OpBuilder object and calling builder.create<affine::AffineForOp> with the appropriate operands (the list of operands can be found by looking at affine::AffineForOp::build, unfortunately c++ tooling is currently unable to see through the layers of templating we use).

  3. Tensor values cannot be used in affine maps.

  4. Floating-point values cannot be used in affine maps either.

  5. The likely correct way to pass a value extracted from a tensor to an affine loop is tensor.extract → arith.fptosi → arith.index_cast. Note that the restrictions must still be preserved. And beware of overflows.