Preserve iree.reflection attribute when converting from std.func -> vm.func.
PiperOrigin-RevId: 286500207
diff --git a/docs/function_abi.md b/docs/function_abi.md
index 0296a6a..bfd2002 100644
--- a/docs/function_abi.md
+++ b/docs/function_abi.md
@@ -52,7 +52,7 @@
'R' length-prefixed(type-sequence)
type-sequence ::= (arg-result-type)*
-arg-result-type ::= buffer-type | ref-object-type
+arg-result-type ::= buffer-type | ref-object-type | unrecognized-type
buffer-type ::= 'B' length-prefixed(scalar-type? dim*)
scalar-type ::= 't' (
'0' # IEEE float32 (default if not specified)
@@ -70,7 +70,7 @@
)
dim :: = 'd' integer # -1 indicates a dynamic dim
ref-object-type ::= 'O' length-prefixed() # Details TBD
-
+unrecognized-type ::= 'U' length-prefixed()
# Lexical primitives
integer ::= -?[0-9]+
diff --git a/iree/base/signature_mangle.cc b/iree/base/signature_mangle.cc
index 2009d6b..c156ceb 100644
--- a/iree/base/signature_mangle.cc
+++ b/iree/base/signature_mangle.cc
@@ -149,6 +149,10 @@
return func_builder;
}
+void RawSignatureMangler::AddUnrecognized() {
+ builder_.Span(absl::string_view(), 'U');
+}
+
void RawSignatureMangler::AddAnyReference() {
// A more constrained ref object would have a non empty span.
builder_.Span(absl::string_view(), 'O');
diff --git a/iree/base/signature_mangle.h b/iree/base/signature_mangle.h
index c398066..e1971e0 100644
--- a/iree/base/signature_mangle.h
+++ b/iree/base/signature_mangle.h
@@ -168,6 +168,10 @@
return ToFunctionSignature(inputs.builder(), results.builder());
}
+ // Adds an unrecognized type. By default, this is an empty span, but in the
+ // future, it may contain some further description.
+ void AddUnrecognized();
+
// Adds an unconstrained reference-type object.
void AddAnyReference();
diff --git a/iree/compiler/Dialect/Flow/Transforms/MaterializeExportedReflection.cpp b/iree/compiler/Dialect/Flow/Transforms/MaterializeExportedReflection.cpp
index ff806e4..f3969dd 100644
--- a/iree/compiler/Dialect/Flow/Transforms/MaterializeExportedReflection.cpp
+++ b/iree/compiler/Dialect/Flow/Transforms/MaterializeExportedReflection.cpp
@@ -94,8 +94,17 @@
auto mangledTensor = mangleTensorType(tensorType);
if (!mangledTensor) return nullptr;
mangledTensor->builder().AppendTo(fBuilder, tag);
+ return builder.getStringAttr(fBuilder.encoded());
}
+ return nullptr;
+}
+
+StringAttr unrecognizedTypeAttr(Builder builder, char tag) {
+ SignatureBuilder fBuilder;
+ RawSignatureMangler mangler;
+ mangler.AddUnrecognized();
+ mangler.builder().AppendTo(fBuilder, tag);
return builder.getStringAttr(fBuilder.encoded());
}
@@ -114,7 +123,13 @@
// Arguments.
for (int i = 0, e = funcType.getNumInputs(); i < e; ++i) {
auto mangled = mangleType(builder, funcType.getInput(i), 'I');
- if (!mangled) continue;
+ if (!mangled) {
+ func.emitWarning()
+ << "Argument #" << i << " of function " << func.getName()
+ << " is not a recognized public ABI type and the function"
+ << " may not be invokable by standard tools";
+ mangled = unrecognizedTypeAttr(builder, 'I');
+ }
NamedAttributeList l(
func.getArgAttrOfType<DictionaryAttr>(i, "iree.reflection"));
l.set(builder.getIdentifier("f_partial"), mangled);
@@ -124,7 +139,13 @@
// Results.
for (int i = 0, e = funcType.getNumResults(); i < e; ++i) {
auto mangled = mangleType(builder, funcType.getResult(i), 'R');
- if (!mangled) continue;
+ if (!mangled) {
+ func.emitWarning()
+ << "Result #" << i << " of function " << func.getName()
+ << " is not a recognized public ABI type and the function"
+ << " may not be invokable by standard tools";
+ mangled = unrecognizedTypeAttr(builder, 'R');
+ }
NamedAttributeList l(
func.getResultAttrOfType<DictionaryAttr>(i, "iree.reflection"));
l.set(builder.getIdentifier("f_partial"), mangled);
diff --git a/iree/compiler/Dialect/Flow/Transforms/Passes.cpp b/iree/compiler/Dialect/Flow/Transforms/Passes.cpp
index 58839e9..ecd5e26 100644
--- a/iree/compiler/Dialect/Flow/Transforms/Passes.cpp
+++ b/iree/compiler/Dialect/Flow/Transforms/Passes.cpp
@@ -93,16 +93,14 @@
// TODO(benvanik): run symbol DCE pass.
// Materialize default arg/result reflection metadata.
- // TODO(laurenzo): Enable this.
- // passManager.addPass(IREE::Flow::createMaterializeExportedReflection());
+ passManager.addPass(IREE::Flow::createMaterializeExportedReflection());
// Merge arg/result reflection metadata.
// NOTE(laurenzo): This will eventually not be the right place for this as
// it should happen after the HAL has further annotated the exported
// functions (which will be needed for dynamic shapes and synthetic barrier
// arguments).
- // TODO(laurenzo): Enable this.
- // passManager.addPass(IREE::Flow::createMergeExportedReflection());
+ passManager.addPass(IREE::Flow::createMergeExportedReflection());
}
static PassPipelineRegistration<> transformPassPipeline(
diff --git a/iree/compiler/Dialect/Flow/Transforms/test/materialize_exported_reflection.mlir b/iree/compiler/Dialect/Flow/Transforms/test/materialize_exported_reflection.mlir
index 2921d57..eaba6ab 100644
--- a/iree/compiler/Dialect/Flow/Transforms/test/materialize_exported_reflection.mlir
+++ b/iree/compiler/Dialect/Flow/Transforms/test/materialize_exported_reflection.mlir
@@ -19,6 +19,27 @@
}
// ----
+// CHECK-LABEL: func @unrecognizedArgument
+// CHECK-SAME: iree.reflection = {f_partial = "I4!U1!"}
+// expected-warning @+1 {{Argument #0 of function unrecognizedArgument is not a recognized public ABI type and the function may not be invokable by standard tools}}
+func @unrecognizedArgument(%arg0 : i1) -> ()
+ attributes {iree.module.export}
+{
+ return
+}
+
+// ----
+// CHECK-LABEL: func @unrecognizedResult
+// CHECK-SAME: iree.reflection = {f_partial = "R4!U1!"}
+// expected-warning @+1 {{Result #0 of function unrecognizedResult is not a recognized public ABI type and the function may not be invokable by standard tools}}
+func @unrecognizedResult() -> (i1)
+ attributes {iree.module.export}
+{
+ %0 = constant 0 : i1
+ return %0 : i1
+}
+
+// ----
// CHECK-LABEL: func @dynamicDim
// CHECK-SAME: iree.reflection = {f_partial = "I11!B8!t7d-1d4"}
func @dynamicDim(%arg0 : tensor<?x4xi64>) -> () attributes {iree.module.export}
diff --git a/iree/compiler/Dialect/VM/Conversion/StandardToVM/ConvertStandardToVM.cpp b/iree/compiler/Dialect/VM/Conversion/StandardToVM/ConvertStandardToVM.cpp
index e86465a..6ccb736 100644
--- a/iree/compiler/Dialect/VM/Conversion/StandardToVM/ConvertStandardToVM.cpp
+++ b/iree/compiler/Dialect/VM/Conversion/StandardToVM/ConvertStandardToVM.cpp
@@ -61,6 +61,11 @@
class FuncOpConversion : public OpConversionPattern<FuncOp> {
using OpConversionPattern::OpConversionPattern;
+ // Whitelist of function attributes to retain when converting to vm.func.
+ static constexpr std::array<const char *, 1> kRetainedAttributes = {
+ "iree.reflection",
+ };
+
PatternMatchResult matchAndRewrite(
FuncOp srcOp, ArrayRef<Value *> operands,
ConversionPatternRewriter &rewriter) const override {
@@ -92,6 +97,15 @@
auto newFuncOp = rewriter.create<IREE::VM::FuncOp>(
srcOp.getLoc(), srcOp.getName(), newFuncType);
+ // Retain function attributes in the whitelist.
+ for (auto retainAttrName : kRetainedAttributes) {
+ StringRef attrName(retainAttrName);
+ Attribute attr = srcOp.getAttr(attrName);
+ if (attr) {
+ newFuncOp.setAttr(attrName, attr);
+ }
+ }
+
// Move the body region from src -> new.
auto &srcRegion = srcOp.getOperation()->getRegion(0);
auto &newRegion = newFuncOp.getOperation()->getRegion(0);
diff --git a/iree/compiler/Dialect/VM/Conversion/StandardToVM/test/func_attrs.mlir b/iree/compiler/Dialect/VM/Conversion/StandardToVM/test/func_attrs.mlir
new file mode 100644
index 0000000..e6bc0f9
--- /dev/null
+++ b/iree/compiler/Dialect/VM/Conversion/StandardToVM/test/func_attrs.mlir
@@ -0,0 +1,16 @@
+// RUN: iree-opt -split-input-file -pass-pipeline='iree-convert-std-to-vm' %s | IreeFileCheck %s
+
+// -----
+// CHECK-LABEL: @t001_iree_reflection
+module @t001_iree_reflection {
+
+module {
+ // CHECK: func @t001_iree_reflection
+ // CHECK-SAME: iree.reflection = {f = "FOOBAR"}
+ func @t001_iree_reflection(%arg0: i32) -> (i32) attributes {iree.reflection = {f = "FOOBAR"}}
+ {
+ return %arg0 : i32
+ }
+}
+
+}
