Convert type of scf::ForOp initArgs and results

MLIR
1

I have the following Op:

    %vec_sum = scf.for %i = %c0 to %input_size step %c64
            iter_args(%sum_iter = %sum_init) -> (vector<64xf32>) {
         %dif = arith.subi %input_size, %i : index
         %mask = vector.create_mask %dif : vector<64xi1>
         %pass_thru = arith.constant dense<0.0> : vector<64xf32>

         // Compute squared values of the input vector
         %t = vector.maskedload %input_ub[%i], %mask, %pass_thru
             : memref<?xf32>, vector<64xi1>, vector<64xf32> into vector<64xf32>
         %squared = arith.mulf %t, %t : vector<64xf32>
         %sum_next = arith.addf %sum_iter, %squared : vector<64xf32>

         scf.yield %sum_next : vector<64xf32>
    }

and we defined our custom vector type to represent our machine feature, said !my.vector<64xf32>, when lowering, we need to convert all vector type values into our machine vector types, for example

    %vec_sum = scf.for %i = %c0 to %input_size step %c64
            iter_args(%sum_iter = %sum_init) -> (!my.vector<64xf32>) {
         // .....

         scf.yield %sum_next : !my.vector<64xf32>
    }

I found the builder of scf::ForOp is :

  static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, 
     Value lowerBound, Value upperBound, Value step, 
     ValueRange initArgs = std::nullopt,
     function_ref<void(OpBuilder &, Location, Value, ValueRange)> odsArg4 = nullptr);

I can’t find an method to map the original Op’s region into new Op’s region, which only modify the last yieldOp’s operands.

I tried make yieldOp with vectorType is illegal, and convert yieldOp’s operands’ types. I hope argumentMaterialization can do this magic. But it didn’t work.

    target.addDynamicallyLegalOp<scf::YieldOp>([&](Operation *op) -> bool {
      for (auto result : dyn_cast<scf::YieldOp>(op).getResults()) {
        if (result.getType().dyn_cast<VectorType>()) {
          return false;
        }
      }
      return true;
    });

Thanks for your attention!

2

I tried method mentioned in The correct way to convert scf::ForOp's signature? - #2 by ThomasRaoux

However, the post was writed at 2020, I have to adapt code into current LLVM project version. I changed code as follow:

struct SCFForOpPattern : public OpConversionPattern<scf::ForOp> {
  using OpConversionPattern<scf::ForOp>::OpConversionPattern;
  LogicalResult
  matchAndRewrite(scf::ForOp op, scf::ForOpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    // rewriter.replaceOpWithNewOp<scf::ForOp>();
    auto loc = op.getLoc();
    auto converter = getTypeConverter();
    TypeConverter::SignatureConversion signatureConverter(
        adaptor.getInitArgs().size() + 1);
    scf::ForOp newFor = rewriter.create<scf::ForOp>(
        op.getLoc(), adaptor.getLowerBound(), adaptor.getUpperBound(),
        adaptor.getStep(), adaptor.getInitArgs());
    signatureConverter.addInputs(0, newFor.getInductionVar().getType());
    for (auto it : llvm::enumerate(newFor.getRegionIterArgs()))
      signatureConverter.addInputs(it.index() + 1, it.value().getType());
    rewriter.applySignatureConversion(op.getBody(), signatureConverter,
                                      converter);

    rewriter.inlineRegionBefore(op.getBodyRegion(), newFor.getBodyRegion(),
                                newFor.getBodyRegion().begin());
    newFor.getBodyRegion().back().erase();
    newFor.emitRemark();

    rewriter.replaceOp(op, newFor.getResults());
    return success();
  }
};

However, this doesn’t work. I got a compiler crash, and the dump result is:

    %6 = "scf.for"(%4, %arg3, %5, %0) ({
^bb0(%arg5: index, %arg6: !my.vec<64 x f32, xxx, xxx>):
  %22 = "builtin.unrealized_conversion_cast"(%arg6) : (!my.vec<64 x f32, xxx, xxx>) -> vector<64xf32>
  %23 = "builtin.unrealized_conversion_cast"(%22) : (vector<64xf32>) -> !my.vec<64 x f32, xxx, xxx>
  %24 = "arith.subi"(%arg3, <<UNKNOWN SSA VALUE>>) <{overflowFlags = #arith.overflow<none>}> : (index, index) -> index
  %25 = "vector.create_mask"(%24) : (index) -> vector<64xi1>
  %26 = "arith.constant"() <{value = dense<0.000000e+00> : vector<64xf32>}> : () -> vector<64xf32>
  %27 = "my.load"(%arg0, <<UNKNOWN SSA VALUE>>) : (memref<?xf32>, index) -> vector<64xf32>
  %28 = "my.select"(%25, %27, %26) : (vector<64xi1>, vector<64xf32>, vector<64xf32>) -> vector<64xf32>
  %29 = "arith.mulf"(%28, %28) <{denormal = #arith.denormal<ieee>, fastmath = #arith.fastmath<none>}> : (vector<64xf32>, vector<64xf32>) -> vector<64xf32>
  %30 = "arith.addf"(<<UNKNOWN SSA VALUE>>, %29) <{denormal = #arith.denormal<ieee>, fastmath = #arith.fastmath<none>}> : (vector<64xf32>, vector<64xf32>) -> vector<64xf32>
  "scf.yield"(%30) : (vector<64xf32>) -> ()
}) : (index, index, index, !my.vec<64 x f32, xxx, xxx>

seens like ops that use %arg5 refers to an deleted or (maybe) obsoleted value. What’s the problem, and how can I get this work in 2025?

Thanks in advance!

3

I find out the correct way. Here is the code:

struct SCFForOpPattern : public OpConversionPattern<scf::ForOp> {
  using OpConversionPattern<scf::ForOp>::OpConversionPattern;
  LogicalResult
  matchAndRewrite(scf::ForOp op, scf::ForOpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    // rewriter.replaceOpWithNewOp<scf::ForOp>();
    auto loc = op.getLoc();
    auto converter = getTypeConverter();
    scf::ForOp newFor = rewriter.create<scf::ForOp>(
        op.getLoc(), adaptor.getLowerBound(), adaptor.getUpperBound(),
        adaptor.getStep(), adaptor.getInitArgs());

    rewriter.inlineRegionBefore(op.getBodyRegion(), newFor.getBodyRegion(),
                                newFor.getBodyRegion().end());
    newFor.getBodyRegion().front().erase();

    auto signatureConverter =
        converter->convertBlockSignature(newFor.getBody()).value();
    rewriter.applySignatureConversion(newFor.getBody(), signatureConverter,
                                      converter);

    rewriter.replaceOp(op, newFor.getResults());
    return success();
  }
};

Hope it helpful