[Heavy] Add ApplyActionOp

Author: ricejasonf

heavy/include/nbdl_gen/Nbdl.td

@@ -49,6 +49,7 @@ class Nbdl_TypeBase<string name, string type_mnemonic,
   let assemblyFormat = "`<` $cpp_typename `>`";
 }
 
+// DEPRECATED (I think)
 def Nbdl_Opaque : TypeDef<Nbdl_Dialect, "Opaque", []> {
   let mnemonic = "opaque";
   let description = [{
@@ -70,29 +71,12 @@ def Nbdl_Store : TypeDef<Nbdl_Dialect, "Store", []> {
   }];
 }
 
-/* FIXME Types should only be based on corresponding concepts in nbdl.
-def Nbdl_Struct : Nbdl_TypeBase<"Struct", "struct"> {
-  let summary = "C++ semiregular aggregate type";
-  let description = [{
-    Represent C++ struct types whose members
-    are accessed by their name.
-  }];
-}
-
-def Nbdl_State : TypeDef<Nbdl_Dialect, "State", []> {
-  let mnemonic = "state";
-  let description = [{
-    A type satisfying the nbdl::State concept
-  }];
-}
-
 def Nbdl_Variant : TypeDef<Nbdl_Dialect, "Variant", []> {
   let mnemonic = "variant";
   let description = [{
     A type satisfying the nbdl::Variant concept
   }];
 }
-*/
 
 def Nbdl_Tag : TypeDef<Nbdl_Dialect, "Tag", []> {
   let mnemonic = "tag";
@@ -112,10 +96,10 @@ def Nbdl_Symbol : TypeDef<Nbdl_Dialect, "Symbol"> {
 
 // Keys are optional, but we are using variadic arguments
 // so we resort to using this sum type.
-def Nbdl_KeyType : AnyTypeOf<[Nbdl_Opaque, Nbdl_Tag,
+def Nbdl_KeyType : AnyTypeOf<[Nbdl_Store, Nbdl_Tag,
                               Nbdl_Symbol, Nbdl_Unit]>;
 
-def Nbdl_Type : AnyTypeOf<[Nbdl_Unit, Nbdl_Empty, Nbdl_Opaque,
+def Nbdl_Type : AnyTypeOf<[Nbdl_Unit, Nbdl_Empty,
                            Nbdl_Store, Nbdl_Tag]>;
 
 class Nbdl_Op<string mnemonic, list<Trait> traits = []> :
@@ -141,9 +125,12 @@ def Nbdl_GetOp : Nbdl_Op<"get", []> {
   let results   = (outs Nbdl_Type:$result);
 }
 
-def Nbdl_VisitOp : Nbdl_Op<"visit", [Terminator]> {
+def Nbdl_ResolveOp : Nbdl_Op<"resolve", [Terminator]> {
   let description = [{
-    Visit matched results with a function object.
+    Finalize resolving stores by invoking the callback
+    provided in the definition of a visit function.
+    This is similar to nbdl.visit except it does not return
+    a value and terminates the block.
   }];
 
   let arguments = (ins Nbdl_Type:$fn, Variadic<Nbdl_Type>:$args);
@@ -202,7 +189,7 @@ def Nbdl_MatchIfOp : Nbdl_Op<"match_if", [Terminator]> {
     like an if/else statement.
   }];
 
-  let arguments = (ins Nbdl_Type:$input, Nbdl_Opaque:$pred);
+  let arguments = (ins Nbdl_Type:$input, Nbdl_Type:$pred);
   let regions   = (region SizedRegion<1>:$thenRegion,
                           SizedRegion<1>:$elseRegion);
 }
@@ -322,13 +309,24 @@ def Nbdl_UnitOp : Nbdl_Op<"unit", []> {
   let results = (outs Nbdl_Unit:$result);
 }
 
-def Nbdl_ApplyOp : Nbdl_Op<"apply", []> {
+def Nbdl_VisitOp : Nbdl_Op<"visit", []> {
   let description = [{
-    Apply a function.
+    Call a function with resolved stores as arguments
+    returning a result. This is analogous to std::visit
+    except only one function is generated.
   }];
-  let arguments = (ins Nbdl_Opaque:$fn,
+  let arguments = (ins Nbdl_Type:$fn,
                        Variadic<Nbdl_Type>:$args);
-  let results = (outs Nbdl_Opaque:$result);
+  let results = (outs Nbdl_Type:$result);
+}
+
+def Nbdl_ApplyActionOp : Nbdl_Op<"apply_action", [Terminator]> {
+  let description = [{
+    Create a call to nbdl::apply_action which is dispatched
+    by the type of the LHS in C++.
+  }];
+  let arguments = (ins Nbdl_Type:$lhs, Variadic<Nbdl_Type>:$args);
+  let results = (outs);
 }
 
 def Nbdl_ConsumerOp : Nbdl_Op<"consumer", []> {
@@ -345,4 +343,38 @@ def Nbdl_ConsumerOp : Nbdl_Op<"consumer", []> {
   let results = (outs Nbdl_Type:$result);
 }
 
+def Nbdl_ConstOp : Nbdl_Op<"const", []> {
+  let description = [{
+    Cast an object to const likely using std::as_const.
+
+    Notes:
+    We don't modify the type in IR land but rely on C++
+    to manage underlying types. This is analogous to an
+    immutable borrow I think.
+    We generally only allow mutation to occur to the operands
+    of nbdl.apply_action or nbdl.resolve.
+  }];
+
+  let arguments = (ins Nbdl_Type:$arg);
+  let results = (outs Nbdl_Type:$result);
+}
+
+def Nbdl_ScopeOp : Nbdl_Op<"scope", []> {
+  let description = [{
+    Terminating operations in a function may mutate a store
+    thus invalidating any matched store with that object as
+    a node in its root path.  We can nest intermediate operations
+    in a nbdl.scope region to prevent subsequent uses of such
+    objects.
+  }];
+
+  let arguments = (ins);
+  let results = (outs);
+  let regions = (region SizedRegion<1>:$body);
+
+  let builders = [
+    OpBuilder<(ins "std::unique_ptr<::mlir::Region>":$body)>
+  ];
+}
+
 #endif

heavy/lib/Nbdl.cpp

@@ -29,6 +29,7 @@ translate_cpp(heavy::LexerWriterFnRef FnRef, mlir::Operation* Op);
 namespace heavy::nbdl_bind_var {
 heavy::ContextLocal current_nbdl_module;
 heavy::ExternFunction translate_cpp;
+heavy::ExternFunction close_previous_scope;
 heavy::ExternFunction build_match_params_impl;
 heavy::ExternFunction build_overload_impl;
 heavy::ExternFunction build_match_if_impl;
@@ -95,7 +96,7 @@ void build_match_params_impl(Context& C, ValueRefs Args) {
     InputTypes.push_back(Builder.getType<nbdl_gen::StoreType>());
 
   // Push the visitor fn argument.
-  InputTypes.push_back(Builder.getType<nbdl_gen::OpaqueType>());
+  InputTypes.push_back(Builder.getType<nbdl_gen::StoreType>());
 
   mlir::FunctionType FT = Builder.getFunctionType(InputTypes,
                                                   /*ResultTypes*/{});
@@ -113,6 +114,9 @@ void build_match_params_impl(Context& C, ValueRefs Args) {
       BlockArgs.push_back(V);
     }
 
+    C.PushCont([FuncOp](Context& C, ValueRefs) mutable {
+      C.Cont(FuncOp.getOperation());
+    }, CaptureList{});
     C.Apply(Callback, BlockArgs);
   }, CaptureList{Callback});
 
@@ -156,7 +160,7 @@ void build_overload_impl(Context& C, ValueRefs Args) {
       mlir::Region& Body = OverloadOp.getBody();
       mlir::Location MLoc = OverloadOp.getLoc();
       mlir::Type Type = mlir::OpBuilder(OverloadOp)
-        .getType<nbdl_gen::OpaqueType>();
+        .getType<nbdl_gen::StoreType>();
       mlir::BlockArgument BlockArg = Body.addArgument(Type, MLoc);
       heavy::Value V = C.CreateAny(mlir::Value(BlockArg));
       C.Apply(Callback, V);
@@ -203,6 +207,8 @@ void build_match_if_impl(Context& C, ValueRefs Args) {
 
 // Translate a nbdl dialect operation to C++.
 // (translate-cpp op port)
+// The parameter `op` may be an mlir::Operation* or a StringLike
+// which will be used to look up the name in the module.
 // 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)
@@ -249,6 +255,35 @@ void translate_cpp(Context& C, ValueRefs Args) {
   C.Cont();
 }
 
+// If the current block has a terminator, wrap the
+// entire block in a nbdl.scope. This supports the
+// convention that only terminators may perform an
+// operation that may invalidate child stores.
+void close_previous_scope(Context& C, ValueRefs Args) {
+  if (Args.size() != 0)
+    return C.RaiseError("invalid arity");
+  mlir::OpBuilder* Builder = mlir_helper::getCurrentBuilder(C);
+  if (!Builder)
+    return;  // error is already raised by getCurrentBuilder
+  mlir::Block* Block = Builder->getBlock();
+  assert(Block && isa<mlir::func::FuncOp>(Block->getParentOp())
+      && "expecting func insertion point");
+  if (Block->empty() || !Block->back().hasTrait<mlir::OpTrait::IsTerminator>())
+    return C.Cont();
+
+  mlir::Location Loc = Block->back().getLoc();
+
+  // Create new Region for ScopeOp.
+  auto ScopeBody = std::make_unique<mlir::Region>();
+  mlir::Block& NewBlock = ScopeBody->emplaceBlock();
+  while (!Block->empty())
+    Block->front().moveBefore(&NewBlock, NewBlock.end());
+  mlir::Operation* ScopeOp = Builder->create<nbdl_gen::ScopeOp>(Loc, std::move(ScopeBody));
+  Builder->setInsertionPointAfter(ScopeOp);
+
+  C.Cont();
+}
+
 void build_context_impl(Context& C, ValueRefs Args) {
   if (Args.size() != 3)
     return C.RaiseError("invalid arity");
@@ -282,9 +317,9 @@ void build_context_impl(Context& C, ValueRefs Args) {
   mlir::Block& EntryBlock = ContextOp.getBody().emplaceBlock();
 
   // Create the arguments.
-  auto OpaqueTy = Builder.getType<nbdl_gen::OpaqueType>();
+  auto StoreTy = Builder.getType<nbdl_gen::StoreType>();
   for (int32_t i = 0; i < NumParams; i++)
-    EntryBlock.addArgument(OpaqueTy, MLoc);
+    EntryBlock.addArgument(StoreTy, MLoc);
 
   heavy::Value Thunk = C.CreateLambda([ContextOp](Context& C,
                                                   ValueRefs) mutable {
@@ -324,6 +359,8 @@ void HEAVY_NBDL_INIT(heavy::Context& C) {
 
   heavy::nbdl_bind_var::current_nbdl_module.init(C, ModuleOp.getOperation());
   heavy::nbdl_bind_var::translate_cpp = heavy::nbdl_bind::translate_cpp;
+  heavy::nbdl_bind_var::close_previous_scope
+    = heavy::nbdl_bind::close_previous_scope;
   heavy::nbdl_bind_var::build_match_params_impl
     = heavy::nbdl_bind::build_match_params_impl;
   heavy::nbdl_bind_var::build_overload_impl
@@ -341,6 +378,8 @@ void HEAVY_NBDL_LOAD_MODULE(heavy::Context& C) {
   heavy::initModuleNames(C, HEAVY_NBDL_LIB_STR, {
     {"current-nbdl-module", heavy::nbdl_bind_var::current_nbdl_module.get(C)},
     {"translate-cpp", heavy::nbdl_bind_var::translate_cpp},
+    {"close-previous-scope",
+                  heavy::nbdl_bind_var::close_previous_scope},
     {"%build-match-params", heavy::nbdl_bind_var::build_match_params_impl},
     {"%build-overload", heavy::nbdl_bind_var::build_overload_impl},
     {"%build-match-if", heavy::nbdl_bind_var::build_match_if_impl},

heavy/lib/Nbdl/NbdlDialect.cpp

@@ -41,3 +41,9 @@ void nbdl_gen::NbdlDialect::initialize() {
       >();
 
 }
+
+void nbdl_gen::ScopeOp::build(::mlir::OpBuilder&,
+                              ::mlir::OperationState& odsState,
+                              std::unique_ptr<::mlir::Region> body) {
+  odsState.addRegion(std::move(body));
+}