Iterators: Add boilerplate code for new 'States to LLVM' pass. (#605)

Author: ingomueller-net

experimental/iterators/include/iterators/Conversion/Passes.h

@@ -10,6 +10,7 @@
 #define ITERATORS_CONVERSION_PASSES_H
 
 #include "iterators/Conversion/IteratorsToLLVM/IteratorsToLLVM.h"
+#include "iterators/Conversion/StatesToLLVM/StatesToLLVM.h"
 #include "iterators/Conversion/TabularToLLVM/TabularToLLVM.h"
 #include "mlir/Pass/Pass.h"
 

experimental/iterators/include/iterators/Conversion/Passes.td

@@ -29,8 +29,8 @@ def ConvertIteratorsToLLVM : Pass<"convert-iterators-to-llvm", "ModuleOp"> {
     currently only works if the use-def chains of `Stream`s form a tree, i.e.,
     every `Stream` is used as an operand by exactly one subsequent iterator.
 
-    More precisely, for each iterator, the lowering produces a state in the form
-    of an `!llvm.struct<...>`, including any local state that the iterator might
+    More precisely, for each iterator, the lowering produces a state with a
+    number of typed fields, including any local state that the iterator might
     require **plus the states of all iterators in the transitive use-def chain**
     of its operands. The computations are expressed as three functions, `Open`,
     `Next`, and `Close`, which operate on that state and which continuously pass
@@ -58,6 +58,26 @@ def ConvertIteratorsToLLVM : Pass<"convert-iterators-to-llvm", "ModuleOp"> {
   ];
 }
 
+//===----------------------------------------------------------------------===//
+// IteratorsToLLVM
+//===----------------------------------------------------------------------===//
+
+def ConvertStatesToLLVM : Pass<"convert-states-to-llvm", "ModuleOp"> {
+  let summary = "Convert the operations on iterator states into the LLVM "
+                "dialect";
+  let description = [{
+    This lowering pass converts operations on iterator states into equivalent
+    operations of the LLVM dialect. Currently, the ops on iterator states are
+    essentially equivalent to the LLVM ops dealing with structs (but allow
+    arbitrary types), so the lowering only consists of straightforward,
+    one-to-one patterns.
+  }];
+  let constructor = "mlir::createConvertStatesToLLVMPass()";
+  let dependentDialects = [
+    "LLVM::LLVMDialect"
+  ];
+}
+
 //===----------------------------------------------------------------------===//
 // TabularToLLVM
 //===----------------------------------------------------------------------===//

experimental/iterators/include/iterators/Conversion/StatesToLLVM/StatesToLLVM.h

@@ -0,0 +1,35 @@
+//===-- StatesToLLVM.h - Utils to convert from Iterators states -*- C++ -*-===//
+//
+// Licensed 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 ITERATORS_CONVERSION_STATESTOLLVM_STATESTOLLVM_H
+#define ITERATORS_CONVERSION_STATESTOLLVM_STATESTOLLVM_H
+
+#include <memory>
+
+namespace mlir {
+class ModuleOp;
+template <typename T>
+class OperationPass;
+class RewritePatternSet;
+class TypeConverter;
+
+namespace iterators {
+
+/// Populate the given list with patterns that convert from Iterator states to
+/// LLVM.
+void populateStatesToLLVMConversionPatterns(RewritePatternSet &patterns,
+                                            TypeConverter &typeConverter);
+
+} // namespace iterators
+
+/// Create a pass to convert operations on Iterator states to the LLVM dialect.
+std::unique_ptr<OperationPass<ModuleOp>> createConvertStatesToLLVMPass();
+
+} // namespace mlir
+
+#endif // ITERATORS_CONVERSION_STATESTOLLVM_STATESTOLLVM_H

experimental/iterators/lib/CAPI/CMakeLists.txt

@@ -7,4 +7,5 @@ add_mlir_public_c_api_library(IteratorsCAPI
     MLIRTabular
     MLIRTabularToLLVM
     MLIRPass
+    MLIRStatesToLLVM
 )

experimental/iterators/lib/Conversion/CMakeLists.txt

@@ -1,2 +1,3 @@
 add_subdirectory(IteratorsToLLVM)
+add_subdirectory(StatesToLLVM)
 add_subdirectory(TabularToLLVM)

experimental/iterators/lib/Conversion/StatesToLLVM/CMakeLists.txt

@@ -0,0 +1,14 @@
+add_mlir_conversion_library(MLIRStatesToLLVM
+  StatesToLLVM.cpp
+
+  DEPENDS
+  MLIRIteratorsConversionIncGen
+
+  LINK_LIBS PUBLIC
+  IteratorsUtils
+  MLIRFuncDialect
+  MLIRFuncTransforms
+  MLIRIterators
+  MLIRLLVMDialect
+  MLIRPass
+)

experimental/iterators/lib/Conversion/StatesToLLVM/StatesToLLVM.cpp

@@ -0,0 +1,99 @@
+//===-- StatesToLLVM.h - Conversion from Iterator states to LLVM-*- C++ -*-===//
+//
+// Licensed 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "iterators/Conversion/StatesToLLVM/StatesToLLVM.h"
+
+#include "../PassDetail.h"
+#include "iterators/Dialect/Iterators/IR/Iterators.h"
+#include "iterators/Utils/MLIRSupport.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Func/Transforms/FuncConversions.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/IR/ImplicitLocOpBuilder.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+namespace mlir {
+class MLIRContext;
+} // namespace mlir
+
+using namespace mlir;
+using namespace mlir::func;
+using namespace mlir::iterators;
+using namespace mlir::LLVM;
+
+namespace {
+struct ConvertStatesToLLVMPass
+    : public ConvertStatesToLLVMBase<ConvertStatesToLLVMPass> {
+  void runOnOperation() override;
+};
+} // namespace
+
+/// Maps state types from the Iterators dialect to corresponding types in LLVM.
+class StateTypeConverter : public TypeConverter {
+public:
+  StateTypeConverter() {
+    addConversion([](Type type) { return type; });
+  }
+
+private:
+};
+
+void mlir::iterators::populateStatesToLLVMConversionPatterns(
+    RewritePatternSet &patterns, TypeConverter &typeConverter) {
+  // patterns.add<
+  //     // clang-format off
+  //     // clang-format on
+  //     >(typeConverter, patterns.getContext());
+}
+
+void ConvertStatesToLLVMPass::runOnOperation() {
+  auto module = getOperation();
+  StateTypeConverter typeConverter;
+
+  // Convert the remaining ops of this dialect using dialect conversion.
+  ConversionTarget target(getContext());
+  target.addLegalDialect<LLVMDialect>();
+  target.addLegalOp<ModuleOp>();
+  RewritePatternSet patterns(&getContext());
+
+  populateStatesToLLVMConversionPatterns(patterns, typeConverter);
+
+  // Add patterns that converts function signature and calls.
+  populateFunctionOpInterfaceTypeConversionPattern<FuncOp>(patterns,
+                                                           typeConverter);
+  populateCallOpTypeConversionPattern(patterns, typeConverter);
+  populateReturnOpTypeConversionPattern(patterns, typeConverter);
+
+  // Force application of that pattern if signature is not legal yet.
+  target.addDynamicallyLegalOp<func::FuncOp>([&](func::FuncOp op) {
+    return typeConverter.isSignatureLegal(op.getFunctionType());
+  });
+  target.addDynamicallyLegalOp<func::ReturnOp>([&](func::ReturnOp op) {
+    return typeConverter.isLegal(op.getOperandTypes());
+  });
+  target.addDynamicallyLegalOp<func::CallOp>([&](func::CallOp op) {
+    return typeConverter.isSignatureLegal(op.getCalleeType());
+  });
+
+  // Use UnrealizedConversionCast as materializations, which have to be cleaned
+  // up by later passes.
+  auto addUnrealizedCast = [](OpBuilder &builder, Type type, ValueRange inputs,
+                              Location loc) {
+    auto cast = builder.create<UnrealizedConversionCastOp>(loc, type, inputs);
+    return Optional<Value>(cast.getResult(0));
+  };
+  typeConverter.addSourceMaterialization(addUnrealizedCast);
+  typeConverter.addTargetMaterialization(addUnrealizedCast);
+
+  if (failed(applyPartialConversion(module, target, std::move(patterns))))
+    signalPassFailure();
+}
+
+std::unique_ptr<OperationPass<ModuleOp>> mlir::createConvertStatesToLLVMPass() {
+  return std::make_unique<ConvertStatesToLLVMPass>();
+}

tools/mlir-proto-lsp-server/CMakeLists.txt

@@ -16,6 +16,7 @@ if(SANDBOX_ENABLE_ITERATORS)
   list(APPEND dialect_libs MLIRTabular)
   list(APPEND conversion_libs MLIRIteratorsToLLVM)
   list(APPEND conversion_libs MLIRTabularToLLVM)
+  list(APPEND conversion_libs MLIRStatesToLLVM)
 endif()
 
 target_link_libraries(mlir-proto-lsp-server

tools/mlir-proto-opt/CMakeLists.txt

@@ -16,6 +16,7 @@ if(SANDBOX_ENABLE_ITERATORS)
   list(APPEND dialect_libs MLIRTabular)
   list(APPEND conversion_libs MLIRIteratorsToLLVM)
   list(APPEND conversion_libs MLIRTabularToLLVM)
+  list(APPEND conversion_libs MLIRStatesToLLVM)
 endif()
 
 target_link_libraries(mlir-proto-opt