Add conversions from vm list operations to emitc (#5350)

This adds basic support for all list operations except for the ones dealing with refs. Things missing include lists as function arguments/results and basic block arguments. Co-authored-by: Marius Brehler <[email protected]>
iree-org · Apr 21, 2021 · 5eae15e · 5eae15e
1 parent 7deb25f
commit 5eae15e
Show file tree

Hide file tree

Showing 13 changed files with 642 additions and 24 deletions.
diff --git a/iree/compiler/Dialect/VM/Conversion/VMToEmitC/CMakeLists.txt b/iree/compiler/Dialect/VM/Conversion/VMToEmitC/CMakeLists.txt
@@ -27,6 +27,7 @@ if(${IREE_ENABLE_EMITC})
       MLIRPass
       MLIREmitC
       MLIRTransforms
+      iree::compiler::Dialect::VM::Analysis
       iree::compiler::Dialect::VM::IR
     INCLUDES
       "${PROJECT_SOURCE_DIR}/third_party/mlir-emitc/include"

diff --git a/iree/compiler/Dialect/VM/Conversion/VMToEmitC/ConvertVMToEmitC.cpp b/iree/compiler/Dialect/VM/Conversion/VMToEmitC/ConvertVMToEmitC.cpp
diff --git a/iree/compiler/Dialect/VM/Target/C/CMakeLists.txt b/iree/compiler/Dialect/VM/Target/C/CMakeLists.txt
@@ -32,6 +32,7 @@ if(${IREE_ENABLE_EMITC})
       MLIRIR
       MLIRPass
       MLIRSupport
+      iree::compiler::Dialect::VM::Analysis
       iree::compiler::Dialect::VM::IR
       iree::compiler::Dialect::VM::Conversion::VMToEmitC
       iree::compiler::Dialect::VM::Target::CallingConventionUtils

diff --git a/iree/compiler/Dialect/VM/Target/C/CModuleTarget.cpp b/iree/compiler/Dialect/VM/Target/C/CModuleTarget.cpp
@@ -17,6 +17,7 @@
 #include "emitc/Target/Cpp.h"
 #include "iree/compiler/Dialect/IREE/IR/IREEOps.h"
 #include "iree/compiler/Dialect/IREE/Transforms/Passes.h"
+#include "iree/compiler/Dialect/VM/Analysis/RegisterAllocation.h"
 #include "iree/compiler/Dialect/VM/Conversion/VMToEmitC/ConvertVMToEmitC.h"
 #include "iree/compiler/Dialect/VM/Target/CallingConventionUtils.h"
 #include "iree/compiler/Dialect/VM/Transforms/Passes.h"
@@ -219,45 +220,57 @@ static LogicalResult translateCondBranchOp(IREE::VM::CondBranchOp condBranchOp,
 }
 
 static LogicalResult translateFailOp(IREE::VM::FailOp failOp,
-                                     mlir::emitc::CppEmitter &emitter) {
+                                     mlir::emitc::CppEmitter &emitter,
+                                     bool hasRefs) {
   llvm::raw_ostream &output = emitter.ostream();
 
   auto status = failOp.status();
 
+  if (hasRefs) {
+    output << "VM_REF_ARRAY_RELEASE(local_refs);\n";
+  }
+
   output << "return vm_fail_or_ok(" << emitter.getOrCreateName(status)
          << ", iree_make_cstring_view(\"" << failOp.message() << "\"));\n";
   return success();
 }
 
 static LogicalResult translateReturnOpToC(
     IREE::VM::ReturnOp returnOp, mlir::emitc::CppEmitter &emitter,
-    SmallVector<std::string, 4> resultNames) {
+    SmallVector<std::string, 4> resultNames, bool hasRefs) {
+  llvm::raw_ostream &output = emitter.ostream();
+
   for (std::tuple<Value, std::string> tuple :
        llvm::zip(returnOp.getOperands(), resultNames)) {
     Value operand = std::get<0>(tuple);
     std::string resultName = std::get<1>(tuple);
-    emitter.ostream() << "*" << resultName << " = "
-                      << emitter.getOrCreateName(operand) << ";\n";
+    output << "*" << resultName << " = " << emitter.getOrCreateName(operand)
+           << ";\n";
   }
 
-  emitter.ostream() << "return iree_ok_status();\n";
+  if (hasRefs) {
+    output << "VM_REF_ARRAY_RELEASE(local_refs);\n";
+  }
+
+  output << "return iree_ok_status();\n";
 
   return success();
 }
 
 static LogicalResult translateOpToC(Operation &op,
                                     mlir::emitc::CppEmitter &emitter,
-                                    SmallVector<std::string, 4> resultNames) {
+                                    SmallVector<std::string, 4> resultNames,
+                                    bool hasRefs) {
   if (auto branchOp = dyn_cast<IREE::VM::BranchOp>(op))
     return translateBranchOp(branchOp, emitter);
   if (auto callOp = dyn_cast<IREE::VM::CallOp>(op))
     return translateCallOpToC(callOp, emitter);
   if (auto condBranchOp = dyn_cast<IREE::VM::CondBranchOp>(op))
     return translateCondBranchOp(condBranchOp, emitter);
   if (auto failOp = dyn_cast<IREE::VM::FailOp>(op))
-    return translateFailOp(failOp, emitter);
+    return translateFailOp(failOp, emitter, hasRefs);
   if (auto returnOp = dyn_cast<IREE::VM::ReturnOp>(op))
-    return translateReturnOpToC(returnOp, emitter, resultNames);
+    return translateReturnOpToC(returnOp, emitter, resultNames, hasRefs);
   // Fall back to generic emitc printer
   if (succeeded(emitter.emitOperation(op, /*trailingSemicolon=*/true))) {
     return success();
@@ -305,34 +318,62 @@ static LogicalResult translateFunctionToC(IREE::VM::ModuleOp &moduleOp,
   // struct argument name here must not be changed.
   output << moduleName << "_state_t* state) {\n";
 
-  // We forward declare all result variables.
+  // We forward declare all result variables except for the ones with RefType.
+  output << "// VARIABLE DECLARATIONS\n";
+  output << "// RESULTS\n";
   for (auto &op : funcOp.getOps()) {
     for (auto result : op.getResults()) {
+      if (result.getType().isa<IREE::VM::RefType>()) {
+        continue;
+      }
       if (failed(emitter.emitVariableDeclaration(result,
                                                  /*trailingSemicolon=*/true))) {
         return op.emitError() << "Unable to declare result variable for op";
       }
     }
   }
+  output << "// BASIC BLOCK ARGUMENTS\n";
 
   auto &blocks = funcOp.getBlocks();
   // Create label names for basic blocks.
   for (auto &block : blocks) {
     emitter.getOrCreateName(block);
   }
 
-  // Emit variables for basic block arguments.
+  // Emit variables for basic block arguments (omitting the first).
   for (auto it = std::next(blocks.begin()); it != blocks.end(); ++it) {
     Block &block = *it;
     for (auto &arg : block.getArguments()) {
-      if (emitter.hasValueInScope(arg)) return failure();
+      if (emitter.hasValueInScope(arg)) {
+        // This shouldn't happen
+        return failure();
+      }
       if (failed(emitter.emitType(arg.getType()))) {
         return failure();
       }
       output << " " << emitter.getOrCreateName(arg) << ";\n";
     }
   }
 
+  output << "// END VARIABLE DECLARATIONS\n";
+
+  // We reuse the register allocation pass and emit an array for all Values with
+  // ref type instead of generating one variable per Value. This makes the
+  // deallocation process easier for us.
+  RegisterAllocation registerAllocation;
+  if (failed(registerAllocation.recalculate(funcOp))) {
+    return funcOp.emitOpError() << "unable to perform register allocation";
+  }
+
+  const size_t numRefs = registerAllocation.getMaxRefRegisterOrdinal() + 1;
+  const bool hasRefs = numRefs > 0;
+
+  if (hasRefs) {
+    auto ref_initializers = SmallVector<StringRef, 4>{numRefs, "{0}"};
+    output << "iree_vm_ref_t local_refs[" << numRefs << "] = {"
+           << llvm::join(ref_initializers, ", ") << "};\n";
+  }
+
   for (auto &block : blocks) {
     // Only print a label if there is more than one block.
     if (blocks.size() > 1) {
@@ -341,7 +382,8 @@ static LogicalResult translateFunctionToC(IREE::VM::ModuleOp &moduleOp,
       }
     }
     for (Operation &op : block.getOperations()) {
-      if (failed(translateOpToC(op, emitter, resultNames))) {
+      if (failed(
+              translateOpToC(op, emitter, resultNames, /*hasRefs=*/hasRefs))) {
         return failure();
       }
     }
@@ -655,6 +697,7 @@ LogicalResult translateModuleToC(IREE::VM::ModuleOp moduleOp,
   printInclude("iree/vm/api.h");
   printInclude("iree/vm/ops.h");
   printInclude("iree/vm/shims_emitc.h");
+  printInclude("iree/vm/value.h");
   output << "\n";
 
   printModuleComment(moduleOp, output);

diff --git a/iree/compiler/Dialect/VM/Target/C/test/add.mlir b/iree/compiler/Dialect/VM/Target/C/test/add.mlir
@@ -4,8 +4,12 @@
 vm.module @add_module {
   // CHECK: iree_status_t add_module_add_1_impl(int32_t v1, int32_t v2, int32_t *out0, int32_t *out1, add_module_state_t* state) {
   vm.func @add_1(%arg0 : i32, %arg1 : i32) -> (i32, i32) {
+    // CHECK-NEXT: VARIABLE DECLARATIONS
+    // CHECK-NEXT: RESULTS
     // CHECK-NEXT: int32_t v3;
     // CHECK-NEXT: int32_t v4;
+    // CHECK-NEXT: BASIC BLOCK ARGUMENTS
+    // CHECK-NEXT: END VARIABLE DECLARATIONS
     // CHECK-NEXT: v3 = vm_add_i32(v1, v2);
     %0 = vm.add.i32 %arg0, %arg1 : i32
     // CHECK-NEXT: v4 = vm_add_i32(v3, v3);

diff --git a/iree/compiler/Dialect/VM/Target/C/test/calling_convention.mlir b/iree/compiler/Dialect/VM/Target/C/test/calling_convention.mlir
@@ -4,19 +4,31 @@
 vm.module @calling_convention_test {
   // CHECK: iree_status_t calling_convention_test_no_in_no_return_impl(calling_convention_test_state_t* state) {
   vm.func @no_in_no_return() -> () {
+    // CHECK-NEXT: VARIABLE DECLARATIONS
+    // CHECK-NEXT: RESULTS
+    // CHECK-NEXT: BASIC BLOCK ARGUMENTS
+    // CHECK-NEXT: END VARIABLE DECLARATIONS
     // CHECK-NEXT: return iree_ok_status();
     vm.return
   }
 
   // CHECK: iree_status_t calling_convention_test_i32_in_no_return_impl(int32_t v1, calling_convention_test_state_t* state) {
   vm.func @i32_in_no_return(%arg0 : i32) -> () {
+    // CHECK-NEXT: VARIABLE DECLARATIONS
+    // CHECK-NEXT: RESULTS
+    // CHECK-NEXT: BASIC BLOCK ARGUMENTS
+    // CHECK-NEXT: END VARIABLE DECLARATIONS
     // CHECK-NEXT: return iree_ok_status();
     vm.return
   }
 
   // CHECK: iree_status_t calling_convention_test_no_in_i32_return_impl(int32_t *out0, calling_convention_test_state_t* state) {
   vm.func @no_in_i32_return() -> (i32) {
+    // CHECK-NEXT: VARIABLE DECLARATIONS
+    // CHECK-NEXT: RESULTS
     // CHECK-NEXT: int32_t v1;
+    // CHECK-NEXT: BASIC BLOCK ARGUMENTS
+    // CHECK-NEXT: END VARIABLE DECLARATIONS
     // CHECK-NEXT: v1 = 32;
     %0 = vm.const.i32 32 : i32
     // CHECK-NEXT: *out0 = v1;
@@ -26,7 +38,11 @@ vm.module @calling_convention_test {
 
   // CHECK: iree_status_t calling_convention_test_i32_in_i32_return_impl(int32_t v1, int32_t *out0, calling_convention_test_state_t* state) {
   vm.func @i32_in_i32_return(%arg0 : i32) -> (i32) {
+    // CHECK-NEXT: VARIABLE DECLARATIONS
+    // CHECK-NEXT: RESULTS
     // CHECK-NEXT: int32_t v2;
+    // CHECK-NEXT: BASIC BLOCK ARGUMENTS
+    // CHECK-NEXT: END VARIABLE DECLARATIONS
     // CHECK-NEXT: v2 = 32;
     %0 = vm.const.i32 32 : i32
     // CHECK-NEXT: *out0 = v2;

diff --git a/iree/compiler/Dialect/VM/Target/C/test/control_flow.mlir b/iree/compiler/Dialect/VM/Target/C/test/control_flow.mlir
@@ -16,11 +16,15 @@ vm.module @control_flow_module {
   }
 }
 // CHECK: iree_status_t control_flow_module_control_flow_test_impl(int32_t [[A:[^ ]*]], int32_t [[COND:[^ ]*]], int32_t *[[RESULT:[^ ]*]], control_flow_module_state_t* [[STATE:[^ ]*]]) {
+  // CHECK-NEXT: VARIABLE DECLARATIONS
+  // CHECK-NEXT: RESULTS
   // CHECK-NEXT: int32_t [[B:[^ ]*]];
   // CHECK-NEXT: int32_t [[V0:[^ ]*]];
+  // CHECK-NEXT: BASIC BLOCK ARGUMENTS
   // CHECK-NEXT: int32_t [[C:[^ ]*]];
   // CHECK-NEXT: int32_t [[D:[^ ]*]];
   // CHECK-NEXT: int32_t [[E:[^ ]*]];
+  // CHECK-NEXT: END VARIABLE DECLARATIONS
   // CHECK-NEXT: [[BB0:[^ ]*]]:
   // CHECK-NEXT: if ([[COND]]) {
   // CHECK-NEXT: goto [[BB1:[^ ]*]];

diff --git a/iree/compiler/Dialect/VM/Target/C/test/global_ops.mlir b/iree/compiler/Dialect/VM/Target/C/test/global_ops.mlir
@@ -14,7 +14,11 @@ vm.module @global_ops {
   vm.export @test_global_load_i32
   // CHECK-LABEL: iree_status_t global_ops_test_global_load_i32_impl(
   vm.func @test_global_load_i32() -> i32 {
+    // CHECK-NEXT: VARIABLE DECLARATIONS
+    // CHECK-NEXT: RESULTS
     // CHECK-NEXT: int32_t v1;
+    // CHECK-NEXT: BASIC BLOCK ARGUMENTS
+    // CHECK-NEXT: END VARIABLE DECLARATIONS
     // CHECK-NEXT: v1 = vm_global_load_i32(state->rwdata, 0);
     %value = vm.global.load.i32 @c42 : i32
     vm.return %value : i32
@@ -23,8 +27,12 @@ vm.module @global_ops {
   vm.export @test_global_store_i32
   // CHECK-LABEL: iree_status_t global_ops_test_global_store_i32_impl(
   vm.func @test_global_store_i32() -> i32 {
+    // CHECK-NEXT: VARIABLE DECLARATIONS
+    // CHECK-NEXT: RESULTS
     // CHECK-NEXT: int32_t v1;
     // CHECK-NEXT: int32_t v2;
+    // CHECK-NEXT: BASIC BLOCK ARGUMENTS
+    // CHECK-NEXT: END VARIABLE DECLARATIONS
     // CHECK-NEXT: v1 = 17;
     %c17 = vm.const.i32 17 : i32
     // CHECK-NEXT: vm_global_store_i32(state->rwdata, 4, v1);

diff --git a/iree/vm/ops.h b/iree/vm/ops.h
@@ -18,6 +18,7 @@
 #include <stdint.h>
 
 #include "iree/base/api.h"
+#include "iree/vm/value.h"
 
 //===------------------------------------------------------------------===//
 // Globals
@@ -222,4 +223,32 @@ static inline int32_t vm_cmp_nz_i64(int64_t operand) {
   return (operand != 0) ? 1 : 0;
 }
 
+//===------------------------------------------------------------------===//
+// Utility macros (Used for things that EmitC can't hadnle)
+//===------------------------------------------------------------------===//
+
+// Get the address of an array element
+#define VM_ARRAY_ELEMENT_ADDRESS(array, index) &array[index]
+
+// Release all refs from the given array
+#define VM_REF_ARRAY_RELEASE(array)                          \
+  for (int i = 0; i < IREE_ARRAYSIZE(array); i++) {          \
+    iree_vm_ref_release(VM_ARRAY_ELEMENT_ADDRESS(array, i)); \
+  }
+
+// TODO(simon-camp): This macro should resemble the error handling part of the
+// IREE_RETURN_IF_ERROR macro. There are two different definitions in
+// iree/base/api.h depending on a feature flag.
+#define VM_RETURN_IF_ERROR(status, array) \
+  if (status) {                           \
+    VM_REF_ARRAY_RELEASE(array);          \
+    return status;                        \
+  }
+
+#define VM_RETURN_IF_LIST_NULL(list, array)                \
+  if (!list) {                                             \
+    VM_REF_ARRAY_RELEASE(array);                           \
+    return iree_make_status(IREE_STATUS_INVALID_ARGUMENT); \
+  }
+
 #endif  // IREE_VM_OPS_H_
diff --git a/iree/vm/test/emitc/CMakeLists.txt b/iree/vm/test/emitc/CMakeLists.txt
@@ -40,6 +40,7 @@ iree_cc_test(
     ::conversion_ops
     ::conversion_ops_i64
     ::global_ops
+    ::list_ops
     ::shift_ops
     ::shift_ops_i64
 )
@@ -134,6 +135,15 @@ iree_c_module(
     "global_ops.h"
 )
 
+iree_c_module(
+  NAME
+    list_ops
+  SRC
+    "../list_ops.mlir"
+  H_FILE_OUTPUT
+    "list_ops.h"
+)
+
 iree_c_module(
   NAME
     shift_ops

diff --git a/iree/vm/test/emitc/module_test.cc b/iree/vm/test/emitc/module_test.cc
@@ -28,6 +28,7 @@
 #include "iree/vm/test/emitc/conversion_ops.h"
 #include "iree/vm/test/emitc/conversion_ops_i64.h"
 #include "iree/vm/test/emitc/global_ops.h"
+#include "iree/vm/test/emitc/list_ops.h"
 #include "iree/vm/test/emitc/shift_ops.h"
 #include "iree/vm/test/emitc/shift_ops_i64.h"
 
@@ -48,7 +49,8 @@ struct ModuleDescription {
 };
 
 std::ostream& operator<<(std::ostream& os, const TestParams& params) {
-  return os << absl::StrReplaceAll(params.local_name, {{":", "_"}, {".", "_"}});
+  std::string qualified_name = params.module_name + "." + params.local_name;
+  return os << absl::StrReplaceAll(qualified_name, {{":", "_"}, {".", "_"}});
 }
 
 std::vector<TestParams> GetModuleTestParams() {
@@ -66,6 +68,7 @@ std::vector<TestParams> GetModuleTestParams() {
       {conversion_ops_descriptor_, conversion_ops_create},
       {conversion_ops_i64_descriptor_, conversion_ops_i64_create},
       {global_ops_descriptor_, global_ops_create},
+      {list_ops_descriptor_, list_ops_create},
       {shift_ops_descriptor_, shift_ops_create},
       {shift_ops_i64_descriptor_, shift_ops_i64_create}};
 
@@ -130,7 +133,6 @@ class VMCModuleTest : public ::testing::Test,
 
   iree_vm_instance_t* instance_ = nullptr;
   iree_vm_context_t* context_ = nullptr;
-  iree_vm_module_t* bytecode_module_ = nullptr;
 };
 
 TEST_P(VMCModuleTest, Check) {