[Heavy] Add bindings for creating function types and module lookup

Author: ricejasonf

heavy/include/nbdl_gen/Nbdl.td

@@ -40,6 +40,7 @@ def Nbdl_Void : TypeDef<Nbdl_Dialect, "Void", []> {
 
 // We keep cpp_typename around for when
 // we want to go back to c++.
+// DEPRECATED
 class Nbdl_TypeBase<string name, string type_mnemonic,
                     list<Trait> traits = []>
   : TypeDef<Nbdl_Dialect, name, traits> {
@@ -48,7 +49,8 @@ class Nbdl_TypeBase<string name, string type_mnemonic,
   let assemblyFormat = "`<` $cpp_typename `>`";
 }
 
-def Nbdl_Opaque : Nbdl_TypeBase<"Opaque", "opaque"> {
+def Nbdl_Opaque : TypeDef<Nbdl_Dialect, "Opaque", []> {
+  let mnemonic = "opaque";
   let description = [{
     Represent an unknown C++ type.
   }];
@@ -64,46 +66,41 @@ def Nbdl_Struct : Nbdl_TypeBase<"Struct", "struct"> {
 }
 */
 
-/* FIXME This is not useful without the key.
-def Nbdl_State : Nbdl_TypeBase<"State", "state"> {
-  let summary = "Nbdl State type";
+def Nbdl_State : TypeDef<Nbdl_Dialect, "State", []> {
+  let mnemonic = "state";
   let description = [{
-    In Nbdl, a State object uses a static key to access
-    elements where paths are resolved at compile-time.
+    A type satisfying the nbdl::State concept
   }];
 }
-*/
 
-def Nbdl_Store : Nbdl_TypeBase<"Store", "store"> {
-  let summary = "Nbdl Store type";
+def Nbdl_Store : TypeDef<Nbdl_Dialect, "Store", []> {
+  let mnemonic = "store";
   let description = [{
-    In Nbdl, a Store object has state that may have multiple
-    alternative paths that are determined at run-time.
-    A State object is also a Store.
+    A type satisfying the nbdl::Store concept
   }];
-  // TODO Stores should have a list of keys that we know they match.
 }
 
-def Nbdl_Variant : Nbdl_TypeBase<"Variant", "variant"> {
-  let summary = "Nbdl Variant type";
+def Nbdl_Variant : TypeDef<Nbdl_Dialect, "Variant", []> {
+  let mnemonic = "variant";
   let description = [{
-    Variant is a store that is a sum type with a type index.
+    A type satisfying the nbdl::Variant concept
   }];
 }
 
-def Nbdl_Tag : Nbdl_TypeBase<"Tag", "tag_type"> {
-  let summary = "Nbdl tag type";
+def Nbdl_Tag : TypeDef<Nbdl_Dialect, "Tag", []> {
+  let mnemonic = "tag";
   let description = [{
-    In Nbdl, an empty type that only carries information
+    An empty type that only carries information
     at compile-time is called a tag type.
   }];
 }
 
-def Nbdl_Symbol : Nbdl_TypeBase<"Symbol", "symbol_type"> {
-  let summary = "Nbdl symbol type";
+def Nbdl_Symbol : Nbdl_TypeBase<"Symbol", "symbol"> {
   let description = [{
     Represent a string that is a valid C++ identifier.
     The intended primary use case is to specify a member of a Struct.
+
+    FIXME using this?
   }];
 }
 
@@ -113,7 +110,7 @@ def Nbdl_KeyType : AnyTypeOf<[Nbdl_Opaque, Nbdl_Tag,
                               Nbdl_Symbol, NoneType]>;
 
 def Nbdl_Type : AnyTypeOf<[Nbdl_Void, Nbdl_Empty, Nbdl_Opaque,
-                           Nbdl_Store, Nbdl_Tag]>;
+                           Nbdl_Store, Nbdl_State, Nbdl_Tag]>;
 
 def Nbdl_TagAttr : AttrDef<Nbdl_Dialect, "tag_attr"> {
   let attrName = "nbdl.tag_attr";

heavy/lib/Mlir.cpp

@@ -48,7 +48,9 @@ heavy::ExternFunction block_op;
 heavy::ExternFunction set_insertion_point;
 heavy::ExternFunction set_insertion_after;
 heavy::ExternFunction type;
+heavy::ExternFunction function_type_impl;
 heavy::ExternFunction attr;
+heavy::ExternFunction type_attr;
 heavy::ExternFunction value_attr;
 template <typename AttrTy>
 heavy::ExternFunction string_attr;
@@ -58,6 +60,7 @@ heavy::ExternFunction load_dialect;
 heavy::ExternFunction parent_op;
 heavy::ExternFunction op_next;
 heavy::ExternFunction verify;
+heavy::ExternFunction module_lookup;
 }
 
 namespace {
@@ -577,6 +580,45 @@ void type(Context& C, ValueRefs Args) {
   C.Cont(CreateTagged(C, kind::mlir_type, Type.getImpl()));
 }
 
+// Create a function type (using vector literals.
+// (%function-type #(<arg-types>...) #(<result-types>...))
+void function_type_impl(Context& C, ValueRefs Args) {
+  mlir::MLIRContext* MLIRContext = getCurrentContext(C);
+  if (Args.size() != 2)
+    return C.RaiseError("invalid arity");
+  auto* ArgTypeVals = heavy::dyn_cast<heavy::Vector>(Args[0]);
+  auto* ResultTypeVals = heavy::dyn_cast<heavy::Vector>(Args[1]);
+  if (!ArgTypeVals || !ResultTypeVals)
+    return C.RaiseError("expecting vectors");
+
+  // Arg types
+  llvm::SmallVector<mlir::Type, 8> ArgTypes;
+  Args = Args.drop_front();
+  for (heavy::Value Arg : ArgTypeVals->getElements()) {
+    mlir::Type ArgType = GetTagged<mlir::Type>(C, kind::mlir_type, Arg);
+    if (!ArgType)
+      return C.RaiseError("expecting mlir.type");
+    ArgTypes.push_back(ArgType);
+  }
+
+  // Result types
+  llvm::SmallVector<mlir::Type, 8> ResultTypes;
+  for (heavy::Value Result : ResultTypeVals->getElements()) {
+    mlir::Type ResultType = GetTagged<mlir::Type>(C, kind::mlir_type, Result);
+    if (!ResultType)
+      return C.RaiseError("expecting mlir.type");
+    ResultTypes.push_back(ResultType);
+  }
+
+  mlir::Type Type = mlir::FunctionType::get(MLIRContext, ArgTypes,
+                                            ResultTypes);
+
+  if (!Type)
+    return C.RaiseError("mlir build function type failed");
+
+  C.Cont(CreateTagged(C, kind::mlir_type, Type.getImpl()));
+}
+
 // Get an attribute by parsing a string.
 //  Usage: (attr _attr_str [_type_])
 //    attr_str - the string to be parsed
@@ -621,6 +663,17 @@ void attr(Context& C, ValueRefs Args) {
   C.Cont(CreateTagged(C, kind::mlir_attr, Attr.getImpl()));
 }
 
+// (type-attr (type "!type-goes-here"))
+void type_attr(Context& C, ValueRefs Args) {
+  if (Args.size() != 1)
+    return C.RaiseError("invalid arity");
+  mlir::Type Type = GetTagged<mlir::Type>(C, kind::mlir_type, Args[0]);
+  if (!Type)
+    return C.RaiseError("expecting a mlir.type");
+  mlir::Attribute Attr = mlir::TypeAttr::get(Type);
+  C.Cont(CreateTagged(C, kind::mlir_attr, Attr.getImpl()));
+}
+
 // Create a heavy scheme value attribute of type !heavy.value
 void value_attr(Context& C, ValueRefs Args) {
   if (Args.size() != 1)
@@ -769,6 +822,24 @@ void verify(Context& C, heavy::ValueRefs Args) {
   }
 }
 
+// (symbol-table-lookup _module-op_ _"symbolname"_)
+void module_lookup(Context& C, heavy::ValueRefs Args) {
+  if (Args.size() != 2)
+    return C.RaiseError("invalid arity");
+
+  mlir::ModuleOp ModuleOp = dyn_cast_or_null<mlir::ModuleOp>(
+      dyn_cast<mlir::Operation>(Args[0]));
+  llvm::StringRef SymbolName = Args[1].getStringRef();
+  if (!ModuleOp)
+    return C.RaiseError("expecting mlir::ModuleOp");
+  if (SymbolName.empty())
+    return C.RaiseError("expecting nonempty string-like object");
+
+  mlir::Operation* Op = ModuleOp.lookupSymbol(SymbolName);
+  heavy::Value Result = Op != nullptr ? heavy::Value(Op) : heavy::Empty();
+  C.Cont(Result);
+}
+
 }  // namespace heavy::mlir_bind
 
 extern "C" {
@@ -795,7 +866,9 @@ void HEAVY_MLIR_INIT(heavy::Context& C) {
   HEAVY_MLIR_VAR(set_insertion_point) = heavy::mlir_bind::set_insertion_point;
   HEAVY_MLIR_VAR(set_insertion_after) = heavy::mlir_bind::set_insertion_after;
   HEAVY_MLIR_VAR(type) = heavy::mlir_bind::type;
+  HEAVY_MLIR_VAR(function_type_impl) = heavy::mlir_bind::function_type_impl;
   HEAVY_MLIR_VAR(attr) = heavy::mlir_bind::attr;
+  HEAVY_MLIR_VAR(type_attr) = heavy::mlir_bind::type_attr;
   HEAVY_MLIR_VAR(value_attr) = heavy::mlir_bind::value_attr;
   HEAVY_MLIR_VAR(string_attr<mlir::StringAttr>)
     = heavy::mlir_bind::string_attr<mlir::StringAttr>;
@@ -805,6 +878,7 @@ void HEAVY_MLIR_INIT(heavy::Context& C) {
   HEAVY_MLIR_VAR(with_builder) = heavy::mlir_bind::with_builder;
   HEAVY_MLIR_VAR(with_new_context) = heavy::mlir_bind::with_new_context;
   HEAVY_MLIR_VAR(verify) = heavy::mlir_bind::verify;
+  HEAVY_MLIR_VAR(module_lookup) = heavy::mlir_bind::module_lookup;
 }
 
 void HEAVY_MLIR_LOAD_MODULE(heavy::Context& C) {
@@ -825,7 +899,9 @@ void HEAVY_MLIR_LOAD_MODULE(heavy::Context& C) {
     {"set-insertion-point", HEAVY_MLIR_VAR(set_insertion_point)},
     {"set-insertion-after", HEAVY_MLIR_VAR(set_insertion_after)},
     {"type", HEAVY_MLIR_VAR(type)},
+    {"%function-type", HEAVY_MLIR_VAR(function_type_impl)},
     {"attr", HEAVY_MLIR_VAR(attr)},
+    {"type-attr", HEAVY_MLIR_VAR(type_attr)},
     {"value-attr", HEAVY_MLIR_VAR(value_attr)},
     {"string-attr", HEAVY_MLIR_VAR(string_attr<mlir::StringAttr>)},
     {"flat-symbolref-attr",
@@ -834,6 +910,7 @@ void HEAVY_MLIR_LOAD_MODULE(heavy::Context& C) {
     {"with-builder", HEAVY_MLIR_VAR(with_builder)},
     {"load-dialect", HEAVY_MLIR_VAR(load_dialect)},
     {"verify", HEAVY_MLIR_VAR(verify)},
+    {"module-lookup", HEAVY_MLIR_VAR(module_lookup)},
   });
 }
 }

heavy/lib/Nbdl.cpp

@@ -19,34 +19,47 @@
 #include <tuple>
 
 namespace nbdl_gen {
-std::tuple<std::string, mlir::Location, mlir::Operation*>
+std::tuple<std::string, heavy::SourceLocationEncoding*, mlir::Operation*>
 translate_cpp(llvm::raw_ostream& OS, mlir::Operation* Op);
 }
 
 namespace heavy::nbdl_bind_var {
-heavy::ExternFunction write_cpp;
+heavy::ExternFunction translate_cpp;
 }
 
 namespace heavy::nbdl_bind {
 // Translate a nbdl dialect operation to C++.
-void write_cpp(Context& C, ValueRefs Args) {
-  if (Args.size() != 2)
+// (translate-cpp op port)
+// Currently the "port" has to be a tagged llvm::raw_ostream.
+void translate_cpp(Context& C, ValueRefs Args) {
+  if (Args.size() != 2 && Args.size() != 1)
     return C.RaiseError("invalid arity");
   auto* Op = dyn_cast<mlir::Operation>(Args[0]);
-  // Do not captured the emphemeral Tagged object.
-  auto* Tagged = dyn_cast<heavy::Tagged>(Args[1]);
-  heavy::Symbol* KindSym = C.CreateSymbol("::llvm::raw_ostream");
   if (!Op)
     return C.RaiseError("expecting mlir.operation");
-  if (!Tagged || !Tagged->isa(KindSym))
-    return C.RaiseError("expecting ::llvm::raw_ostream");
 
-  auto& OS = Tagged->cast<llvm::raw_ostream>();
-  auto&& [ErrMsg, ErrLoc, Irritant] = nbdl_gen::translate_cpp(OS, Op);
+  llvm::raw_ostream* OS = nullptr;
+
+  // Do not capture the emphemeral Tagged object.
+  if (Args.size() == 2) {
+    auto* Tagged = dyn_cast<heavy::Tagged>(Args[1]);
+    heavy::Symbol* KindSym = C.CreateSymbol("::llvm::raw_ostream");
+    if (!Tagged || !Tagged->isa(KindSym))
+      return C.RaiseError("expecting ::llvm::raw_ostream");
+    OS = &(Tagged->cast<llvm::raw_ostream>());
+  } else {
+    OS = &llvm::outs();
+  }
+
+  auto&& [ErrMsg, ErrLoc, Irritant] = nbdl_gen::translate_cpp(*OS, Op);
   if (!ErrMsg.empty()) {
-    auto Loc = heavy::SourceLocation(mlir::OpaqueLoc
-      ::getUnderlyingLocationOrNull<heavy::SourceLocationEncoding*>(
-        mlir::dyn_cast<mlir::OpaqueLoc>(ErrLoc)));
+    heavy::SourceLocation Loc(ErrLoc);
+#if 0
+    if (mlir::isa<mlir::OpaqueLoc>(ErrLoc))
+      Loc = heavy::SourceLocation(mlir::OpaqueLoc
+        ::getUnderlyingLocationOrNull<heavy::SourceLocationEncoding*>(
+          mlir::dyn_cast<mlir::OpaqueLoc>(ErrLoc)));
+#endif
     heavy::Error* Err = C.CreateError(Loc, ErrMsg, 
         Irritant ? heavy::Value(Irritant) : Undefined());
     return C.Raise(Err);
@@ -60,13 +73,13 @@ extern "C" {
 void HEAVY_NBDL_INIT(heavy::Context& C) {
   C.DialectRegistry->insert<nbdl_gen::NbdlDialect>();
 
-  heavy::nbdl_bind_var::write_cpp = heavy::nbdl_bind::write_cpp;
+  heavy::nbdl_bind_var::translate_cpp = heavy::nbdl_bind::translate_cpp;
 }
 
 void HEAVY_NBDL_LOAD_MODULE(heavy::Context& C) {
   HEAVY_NBDL_INIT(C);
   heavy::initModuleNames(C, HEAVY_NBDL_LIB_STR, {
-    {"write-cpp", heavy::nbdl_bind_var::write_cpp},
+    {"translate-cpp", heavy::nbdl_bind_var::translate_cpp},
   });
 }
 }

heavy/lib/NbdlWriter.cpp

@@ -264,8 +264,7 @@ class NbdlWriter {
   }
 
   void VisitType(StoreOp Op) {
-    mlir::Value Result = Op.getResult();
-    VisitType(Op.getLoc(), Result);
+    OS << Op.getName();
   }
 
   void VisitType(VariantOp Op) {
@@ -307,16 +306,7 @@ class NbdlWriter {
   }
 
   void VisitType(mlir::Location Loc, mlir::Type Type) {
-    if (auto OpaqueType = dyn_cast<nbdl_gen::OpaqueType>(Type))
-      OS << OpaqueType.getCppTypename();
-    else if (auto StoreType = dyn_cast<nbdl_gen::StoreType>(Type))
-      OS << StoreType.getCppTypename();
-    else if (auto VariantType = dyn_cast<nbdl_gen::VariantType>(Type))
-      OS << VariantType.getCppTypename();
-    else if (auto TagType = dyn_cast<nbdl_gen::TagType>(Type))
-      OS << TagType.getCppTypename();
-    else
-      SetError(Loc, "unprintable type");
+    SetError(Loc, "unprintable type");
   }
 };
 }