[mlir][SCF] populateSCFStructuralTypeConversionsAndLegality WhileOp support

[ftynse]

mlir/lib/Dialect/SCF/Transforms/StructuralTypeConversions.cpp

144

Nit: expand auto here please (should be a reference IIRC).

[Hardcode84]

Review comment and warning

mlir/lib/Dialect/SCF/Transforms/StructuralTypeConversions.cpp

public:
   matchAndRewrite(scf::YieldOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     rewriter.replaceOpWithNewOp<scf::YieldOp>(op, operands);
     return success();
   }
 };
 } // namespace
 
+namespace {
+class ConvertWhileOpTypes : public OpConversionPattern<WhileOp> {
+public:
+  using OpConversionPattern<WhileOp>::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(WhileOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto converter = getTypeConverter();
+    assert(converter);
+    SmallVector<Type> newResultTypes;
+    if (failed(converter->convertTypes(op.getResultTypes(), newResultTypes)))
+      return failure();
+
+    WhileOp::Adaptor adaptor(operands);
+    auto newOp = rewriter.create<WhileOp>(op.getLoc(), newResultTypes,
+                                          adaptor.getOperands());
+    for (auto i : {0u, 1u}) {
+      auto &dstRegion = newOp.getRegion(i);
+      rewriter.inlineRegionBefore(op.getRegion(i), dstRegion, dstRegion.end());
+      if (failed(rewriter.convertRegionTypes(&dstRegion, *converter)))
+        return rewriter.notifyMatchFailure(op, "could not convert body types");
+    }
+    rewriter.replaceOp(op, newOp.getResults());
+    return success();
+  }
+};
+} // namespace
+
+namespace {
+class ConvertConditionOpTypes : public OpConversionPattern<ConditionOp> {
+public:
+  using OpConversionPattern<ConditionOp>::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(ConditionOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    rewriter.updateRootInPlace(op, [&]() { op->setOperands(operands); });
+    return success();
+  }
+};
+} // namespace
+
 void mlir::scf::populateSCFStructuralTypeConversionsAndLegality(
     TypeConverter &typeConverter, RewritePatternSet &patterns,
     ConversionTarget &target) {
-  patterns.add<ConvertForOpTypes, ConvertIfOpTypes, ConvertYieldOpTypes>(
+  patterns.add<ConvertForOpTypes, ConvertIfOpTypes, ConvertYieldOpTypes,
+               ConvertWhileOpTypes, ConvertConditionOpTypes>(
       typeConverter, patterns.getContext());
   target.addDynamicallyLegalOp<ForOp, IfOp>([&](Operation *op) {
     return typeConverter.isLegal(op->getResultTypes());
   });
   target.addDynamicallyLegalOp<scf::YieldOp>([&](scf::YieldOp op) {
     // We only have conversions for a subset of ops that use scf.yield
     // terminators.
-    if (!isa<ForOp, IfOp>(op->getParentOp()))
+    if (!isa<ForOp, IfOp, WhileOp>(op->getParentOp()))
       return true;
     return typeConverter.isLegal(op.getOperandTypes());
   });
+  target.addDynamicallyLegalOp<WhileOp, ConditionOp>(
+      [&](Operation *op) { return typeConverter.isLegal(op); });
 }

mlir/test/Dialect/SCF/bufferize.mlir

 // CHECK:         }
 func @for_correct_recursive_legalization_behavior(%arg0: tensor<f32>, %index: index) -> tensor<f32> {
   %ret = scf.for %iv = %index to %index step %index iter_args(%iter = %arg0) -> tensor<f32> {
     %0 = "test.munge_tensor"(%iter) : (tensor<f32>) -> (tensor<f32>)
     scf.yield %0 : tensor<f32>
   }
   return %ret : tensor<f32>
 }
+
+// CHECK-LABEL:   func @bufferize_while(
+// CHECK-SAME: %[[ARG0:.*]]: i64, %[[ARG1:.*]]: i64, %[[ARG2:.*]]: tensor<f32>
+// CHECK: %[[M:.*]] = memref.buffer_cast %[[ARG2]] : memref<f32>
+// CHECK: %[[RES1:.*]]:3 = scf.while (%{{.*}} = %[[ARG0]], %{{.*}} = %[[M]]) : (i64, memref<f32>) -> (i64, i64, memref<f32>)
+// CHECK: scf.condition(%{{.*}}) %{{.*}}, %{{.*}}, %{{.*}} : i64, i64, memref<f32>
+// CHECK: ^bb0(%{{.*}}: i64, %{{.*}}: i64, %{{.*}}: memref<f32>):
+// CHECK: scf.yield %{{.*}}, %{{.*}} : i64, memref<f32>
+// CHECK:  %[[RES2:.*]] = memref.tensor_load %[[RES1]]#2 : memref<f32>
+// CHECK:  return %[[RES1]]#1, %[[RES2]] : i64, tensor<f32>
+func @bufferize_while(%arg0: i64, %arg1: i64, %arg2: tensor<f32>) -> (i64, tensor<f32>) {
+  %c2_i64 = constant 2 : i64
+  %0:3 = scf.while (%arg3 = %arg0, %arg4 = %arg2) : (i64, tensor<f32>) -> (i64, i64, tensor<f32>) {
+    %1 = cmpi slt, %arg3, %arg1 : i64
+    scf.condition(%1) %arg3, %arg3, %arg4 : i64, i64, tensor<f32>
+  } do {
+  ^bb0(%arg5: i64, %arg6: i64, %arg7: tensor<f32>):
+    %1 = muli %arg6, %c2_i64 : i64
+    scf.yield %1, %arg7 : i64, tensor<f32>
+  }
+  return %0#1, %0#2 : i64, tensor<f32>
+}