[mlir] Add structural conversion to async dialect lowering.

[csigg]

Summary

Lowering of async dialect uses a fixed type converter and therefore does not support lowering non-standard types.

This revision adds a structural conversion so that non-standard types in !async.values can be lowered to LLVM before lowering the async dialect itself.

[ezhulenev]

mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp

1171

Not needed?

[csigg]

Remove commented code left-over.

[csigg]

mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp

1171

Yeah, that's just left-over from experimenting. Removed.

mlir/include/mlir/Conversion/AsyncToLLVM/AsyncToLLVM.h

 //===- AsyncToLLVM.h - Convert Async to LLVM dialect ------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_CONVERSION_ASYNCTOLLVM_ASYNCTOLLVM_H
 #define MLIR_CONVERSION_ASYNCTOLLVM_ASYNCTOLLVM_H
 
 #include <memory>
 
 namespace mlir {
 
+class ConversionTarget;
 class ModuleOp;
 template <typename T>
 class OperationPass;
+class MLIRContext;
+class OwningRewritePatternList;
+class TypeConverter;
 
 /// Create a pass to convert Async operations to the LLVM dialect.
 std::unique_ptr<OperationPass<ModuleOp>> createConvertAsyncToLLVMPass();
 
+/// Populates patterns for async structural type conversions.
+///
+/// A "structural" type conversion is one where the underlying ops are
+/// completely agnostic to the actual types involved and simply need to update
+/// their types. An example of this is async.execute -- the async.execute op and
+/// the corresponding async.yield ops need to update their types accordingly to
+/// the TypeConverter, but otherwise don't care what type conversions are
+/// happening.
+void populateAsyncStructuralTypeConversionsAndLegality(
+    MLIRContext *context, TypeConverter &typeConverter,
+    OwningRewritePatternList &patterns, ConversionTarget &target);
+
 } // namespace mlir
 
 #endif // MLIR_CONVERSION_ASYNCTOLLVM_ASYNCTOLLVM_H

mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp

target.addDynamicallyLegalOp<CallOp>([&](CallOp op) {
     return converter.isSignatureLegal(op.getCalleeType());
   });
 
   if (failed(applyPartialConversion(module, target, std::move(patterns))))
     signalPassFailure();
 }
 } // namespace
 
+namespace {
+class ConvertExecuteOpTypes : public OpConversionPattern<ExecuteOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(ExecuteOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    ExecuteOp newOp =
+        cast<ExecuteOp>(rewriter.cloneWithoutRegions(*op.getOperation()));
+    rewriter.inlineRegionBefore(op.getRegion(), newOp.getRegion(),
+                                newOp.getRegion().end());
+
+    // Set operands and update block argument and result types.
+    newOp->setOperands(operands);
+    if (failed(rewriter.convertRegionTypes(&newOp.getRegion(), *typeConverter)))
+      return failure();
+    for (auto result : newOp.getResults())
+      result.setType(typeConverter->convertType(result.getType()));
+
+    rewriter.replaceOp(op, newOp.getResults());
+    return success();
+  }
+};
+
+// Dummy pattern to trigger the appropriate type conversion / materialization.
+class ConvertAwaitOpTypes : public OpConversionPattern<AwaitOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(AwaitOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    rewriter.replaceOpWithNewOp<AwaitOp>(op, operands.front());
+    // rewriter.updateRootInPlace(op, [] {});
+    return success();
+  }
+};
+
+// Dummy pattern to trigger the appropriate type conversion / materialization.
+class ConvertYieldOpTypes : public OpConversionPattern<async::YieldOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(async::YieldOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    rewriter.replaceOpWithNewOp<async::YieldOp>(op, operands);
+    return success();
+  }
+};
+} // namespace
+
 std::unique_ptr<OperationPass<ModuleOp>> mlir::createConvertAsyncToLLVMPass() {
   return std::make_unique<ConvertAsyncToLLVMPass>();
 }
+
+void mlir::populateAsyncStructuralTypeConversionsAndLegality(
+    MLIRContext *context, TypeConverter &typeConverter,
+    OwningRewritePatternList &patterns, ConversionTarget &target) {
+  typeConverter.addConversion([&](TokenType type) { return type; });
+  typeConverter.addConversion([&](ValueType type) {
+    return ValueType::get(typeConverter.convertType(type.getValueType()));
+  });
+
+  patterns
+      .insert<ConvertExecuteOpTypes, ConvertAwaitOpTypes, ConvertYieldOpTypes>(
+          typeConverter, context);
+
+  target.addDynamicallyLegalOp<AwaitOp, ExecuteOp, async::YieldOp>(
+      [&](Operation *op) { return typeConverter.isLegal(op); });
+}