commit c78eee264dd3f20d123997e6881ed974af5842f3
author Stella Laurenzo <laurenzo@google.com> Sun Mar 22 16:59:41 2020 -0700
committer Copybara-Service <copybara-worker@google.com> Sun Mar 22 17:03:31 2020 -0700
tree 10b4bf0287ee54658bd5f336e99b590e47d61993
parent eec44822a658c01066d0322ae4f14f108cb59b81

Update handling of iree.module.export attribute to support specifying an alias.

PiperOrigin-RevId: 302335050

iree/compiler/Dialect/VM/Conversion/StandardToVM/ConvertStandardToVM.cpp

diff --git a/iree/compiler/Dialect/VM/Conversion/StandardToVM/ConvertStandardToVM.cpp b/iree/compiler/Dialect/VM/Conversion/StandardToVM/ConvertStandardToVM.cpp
index 84f9691..35bd5a1 100644
--- a/iree/compiler/Dialect/VM/Conversion/StandardToVM/ConvertStandardToVM.cpp
+++ b/iree/compiler/Dialect/VM/Conversion/StandardToVM/ConvertStandardToVM.cpp
@@ -126,10 +126,19 @@
     rewriter.applySignatureConversion(&newFuncOp.getBody(),
                                       signatureConversion);
 
-    // Also add an export if the function is tagged. Ideally this would be
-    // replaced with river sym_vis magic.
-    if (srcOp.getAttr("iree.module.export")) {
-      rewriter.create<IREE::VM::ExportOp>(srcOp.getLoc(), newFuncOp);
+    // Also add an export for the "raw" form of this function, which operates
+    // on low level VM types and does no verification. A later pass will
+    // materialize high level API-friendly wrappers.
+    if (auto exportAttr = srcOp.getAttr("iree.module.export")) {
+      StringRef exportName = newFuncOp.getName();
+      if (auto exportStrAttr = exportAttr.dyn_cast<StringAttr>()) {
+        exportName = exportStrAttr.getValue();
+      } else {
+        assert(exportAttr.isa<UnitAttr>());
+      }
+
+      rewriter.create<IREE::VM::ExportOp>(srcOp.getLoc(), newFuncOp,
+                                          exportName);
     }
 
     rewriter.replaceOp(srcOp, llvm::None);

iree/compiler/Dialect/VM/Conversion/StandardToVM/test/func_attrs.mlir

diff --git a/iree/compiler/Dialect/VM/Conversion/StandardToVM/test/func_attrs.mlir b/iree/compiler/Dialect/VM/Conversion/StandardToVM/test/func_attrs.mlir
index e6bc0f9..e17e5af 100644
--- a/iree/compiler/Dialect/VM/Conversion/StandardToVM/test/func_attrs.mlir
+++ b/iree/compiler/Dialect/VM/Conversion/StandardToVM/test/func_attrs.mlir
@@ -14,3 +14,36 @@
 }
 
 }
+
+// -----
+// CHECK-LABEL: @t002_iree_module_export_default
+module @t002_iree_module_export_default {
+
+module {
+  // CHECK: func @internal_function_name
+  // CHECK-NOT: iree.module.export
+  // CHECK: vm.export @internal_function_name
+  func @internal_function_name(%arg0: i32) -> (i32) attributes { iree.module.export }
+  {
+    return %arg0 : i32
+  }
+}
+
+}
+
+// -----
+// CHECK-LABEL: @t003_iree_module_export_alias
+module @t003_iree_module_export_alias {
+
+module {
+  // CHECK: func @internal_function_name
+  // CHECK-NOT: iree.module.export
+  // CHECK: vm.export @internal_function_name as("external_function_name")
+  func @internal_function_name(%arg0: i32) -> (i32)
+      attributes { iree.module.export = "external_function_name" }
+  {
+    return %arg0 : i32
+  }
+}
+
+}