[PoC][TabgleGen] Add new bang operator !apply

[wangpc]

Summary

There is no way to pass function as a value in TableGen, but we can
treat DAGs as S-expression and evaluate them.

In this patch, a new bang operator is added:

type value = !apply<type>(dag[, param]);

We can treat dag as function and evaluate it on params.

To not increase complexity, we map DAGs to subset of existed bang
operators. But we may support more bang operator or even user-definded
subroutines.

The need of !apply has been discussed in D146198.

[tra]

It would be great to update tablegen documentation, too.

On the surface, I think we're reinventing a lambda here.

llvm/lib/TableGen/Record.cpp

1898

How about mul, div, and mod ?

1898

Does it mean that we make all these ops a reserved keyword which would no longer be available to users for other purposes?

E.g. here: https://github.com/llvm/llvm-project/blob/main/llvm/lib/Target/AMDGPU/SIRegisterInfo.td#L123

[• pcwang-thead]

In D147131#4231378, @tra wrote:

It would be great to update tablegen documentation, too.

Thanks. This patch is still in the early stage, doc and more operations will be added later.

On the surface, I think we're reinventing a lambda here.

Yes, so I post this patch here to gather feedbacks to see if I am on the right track. :-)

llvm/lib/TableGen/Record.cpp

1898

mul, div, mod and others will be added later.

1898

No, I think they won't be reserved keywords. These ops are valid only if they are in the context of !apply. And they are just mapped to existed bang operators and the semantic meanings are target-independent.
For the example you pointed out, they are some set operations and implemented in SetTheory.h/.cpp and the semantic meanings are target-dependent.

[nhaehnle]

This is an interesting proposal, but I think it really needs a bit more thought on the overall design. If we're going to add something like this, it should be nicely orthogonal and general.

Some examples of problems in what's there right now:

• Parallel hard-coding of add, sub etc. in the parsing of IDs. Not great, since it affects parsing everywhere -- we could do that, but then what's the overall direction here with regards to !apply vs. the existing ops?
• Hard-coding of argument names $aN.

I agree with @tra that in a way this is reinventing lambdas, so perhaps we should instead just literally take lambdas? That means a syntax for defining lambdas and a syntax for applying them. Strawman riffing off apply.td, something like:

defvar add_lambda = (lambda (lhs, rhs) !add(lhs, rhs));

class A<int a, int b> {
  int add_result = !apply(add_lambda, a, b);
}

What I'm thinking here is that the expression (lambda (args...) expr) is not a DAG but a new expression/init kind, with lambda being the keyword to detect that. The expression is evaluated with a scope where the arguments are substituted.

We could consider alternative syntaxes, two possibilities that would probably work:

defvar add_lambda = ($lhs, $rhs) -> !add(lhs, rhs);
defvar add_lambda = |lhs, rhs| !add(lhs, rhs);

And perhaps we can shortcut the evaluation syntax to just:

class A<int a, int b> {
  int add_result = !(add_lambda a b);
}

Or, perhaps:

class A<int a, int b> {
  int add_result = !add_lambda(a, b);
};

In other words, the syntax is ! expression ( args... ) where the expression must evaluate to a lambda. We could turn most (but not all) bang operators into builtin variables and turn a few exceptional cases into reserved keywords.

I haven't thought this through fully, just throwing some ideas out there that hopefully help suggest a cleaner and more robust design. I think this is large enough that it would also merit having a discussion on Discourse for visibility.

[• pcwang-thead]

In D147131#4232652, @nhaehnle wrote:

This is an interesting proposal, but I think it really needs a bit more thought on the overall design. If we're going to add something like this, it should be nicely orthogonal and general.

Some examples of problems in what's there right now:

• Parallel hard-coding of add, sub etc. in the parsing of IDs. Not great, since it affects parsing everywhere -- we could do that, but then what's the overall direction here with regards to !apply vs. the existing ops?
• Hard-coding of argument names $aN.

I agree with @tra that in a way this is reinventing lambdas, so perhaps we should instead just literally take lambdas? That means a syntax for defining lambdas and a syntax for applying them. Strawman riffing off apply.td, something like:

defvar add_lambda = (lambda (lhs, rhs) !add(lhs, rhs));

class A<int a, int b> {
  int add_result = !apply(add_lambda, a, b);
}

What I'm thinking here is that the expression (lambda (args...) expr) is not a DAG but a new expression/init kind, with lambda being the keyword to detect that. The expression is evaluated with a scope where the arguments are substituted.

We could consider alternative syntaxes, two possibilities that would probably work:

defvar add_lambda = ($lhs, $rhs) -> !add(lhs, rhs);
defvar add_lambda = |lhs, rhs| !add(lhs, rhs);

And perhaps we can shortcut the evaluation syntax to just:

class A<int a, int b> {
  int add_result = !(add_lambda a b);
}

Or, perhaps:

class A<int a, int b> {
  int add_result = !add_lambda(a, b);
};

In other words, the syntax is ! expression ( args... ) where the expression must evaluate to a lambda. We could turn most (but not all) bang operators into builtin variables and turn a few exceptional cases into reserved keywords.

I haven't thought this through fully, just throwing some ideas out there that hopefully help suggest a cleaner and more robust design. I think this is large enough that it would also merit having a discussion on Discourse for visibility.

Thanks! It's really helpful and inspiring!
I will redesign this and make it much more complete before discussing on Discourse.

[tra]

While we're on the subject of lambdas, the pattern we currently use to deal with the lack of lambdas or functions is to instantiate a record with a field evaluated to be a particular value. A lot of those case could be better served by lambdas, but they would need to be able to deal with all tablegen types, not just integers.

I think the current patch's functionality could already be done using this approach, even if it's a bit more verbose.
E.g: https://github.com/llvm/llvm-project/blob/main/llvm/lib/Target/AMDGPU/SIInstrInfo.td#L305
https://github.com/llvm/llvm-project/blob/main/llvm/lib/Target/AMDGPU/SIInstrInfo.td#L2469

We may be just looking for a better syntax, not a new functionality.

[wangpc]

Close it since I will continue finishing similar feature in D148915.

llvm/include/llvm/TableGen/Record.h

enum InitKind : uint8_t {
     IK_BinOpInit,
     IK_TernOpInit,
     IK_UnOpInit,
     IK_LastOpInit,
     IK_CondOpInit,
     IK_FoldOpInit,
     IK_IsAOpInit,
     IK_ExistsOpInit,
+    IK_ApplyOpInit,
     IK_AnonymousNameInit,
     IK_StringInit,
     IK_VarInit,
     IK_VarListElementInit,
     IK_VarBitInit,
     IK_VarDefInit,
     IK_LastTypedInit,
     IK_UnsetInit
public:
   inline size_t              name_size () const { return NumArgNames; }
   inline bool                name_empty() const { return NumArgNames == 0; }
 
   Init *getBit(unsigned Bit) const override {
     llvm_unreachable("Illegal bit reference off dag");
   }
 };
 
+/// !apply<type>(dag, params) - Apply params to function represented in dag.
+class ApplyOpInit : public TypedInit, public FoldingSetNode {
+private:
+  RecTy *ReturnType;
+  Init *Func;
+  SmallVector<Init *, 2> Params;
+
+  ApplyOpInit(RecTy *ReturnType, Init *Func, SmallVector<Init *, 2> Params)
+      : TypedInit(IK_ApplyOpInit, ReturnType), ReturnType(ReturnType),
+        Func(Func), Params(Params) {}
+
+  Init *parseDAG(Record *CurRec, DagInit *Dag) const;
+
+public:
+  ApplyOpInit(const ApplyOpInit &) = delete;
+  ApplyOpInit &operator=(const ApplyOpInit &) = delete;
+
+  static bool classof(const Init *I) { return I->getKind() == IK_ApplyOpInit; }
+
+  static ApplyOpInit *get(RecTy *ReturnType, Init *Func,
+                          SmallVector<Init *, 2> Params);
+
+  void Profile(FoldingSetNodeID &ID) const;
+
+  // Fold - If possible, fold this to a simpler init.  Return this if not
+  // possible to fold.
+  Init *Fold(Record *CurRec, bool IsFinal = false) const;
+
+  Init *resolveReferences(Resolver &R) const override;
+
+  Init *getBit(unsigned Bit) const override;
+
+  std::string getAsString() const override;
+};
+
 //===----------------------------------------------------------------------===//
 //  High-Level Classes
 //===----------------------------------------------------------------------===//
 
 /// This class represents a field in a record, including its name, type,
 /// value, and source location.
 class RecordVal {
   friend class Record;

llvm/lib/TableGen/Record.cpp

 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/MathExtras.h"
 #include "llvm/Support/SMLoc.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/TableGen/Error.h"
 #include <cassert>
 #include <cstdint>
+#include <cstdlib>
 #include <map>
 #include <memory>
 #include <string>
 #include <utility>
 #include <vector>
 
 using namespace llvm;
 
struct RecordKeeperImpl {
   StringMap<StringInit *, BumpPtrAllocator &> StringInitCodePool;
   FoldingSet<ListInit> TheListInitPool;
   FoldingSet<UnOpInit> TheUnOpInitPool;
   FoldingSet<BinOpInit> TheBinOpInitPool;
   FoldingSet<TernOpInit> TheTernOpInitPool;
   FoldingSet<FoldOpInit> TheFoldOpInitPool;
   FoldingSet<IsAOpInit> TheIsAOpInitPool;
   FoldingSet<ExistsOpInit> TheExistsOpInitPool;
+  FoldingSet<ApplyOpInit> TheApplyOpInitPool;
   DenseMap<std::pair<RecTy *, Init *>, VarInit *> TheVarInitPool;
   DenseMap<std::pair<TypedInit *, unsigned>, VarBitInit *> TheVarBitInitPool;
   DenseMap<std::pair<TypedInit *, unsigned>, VarListElementInit *>
       TheVarListElementInitPool;
   FoldingSet<VarDefInit> TheVarDefInitPool;
   DenseMap<std::pair<Init *, StringInit *>, FieldInit *> TheFieldInitPool;
   FoldingSet<CondOpInit> TheCondOpInitPool;
   FoldingSet<DagInit> TheDagInitPool;

 }
 
 std::string ExistsOpInit::getAsString() const {
   return (Twine("!exists<") + CheckType->getAsString() + ">(" +
           Expr->getAsString() + ")")
       .str();
 }
 
+static void ProfileApplyOpInit(FoldingSetNodeID &ID, RecTy *ReturnType,
+                               Init *Dag, SmallVector<Init *, 2> Params) {
+  ID.AddPointer(ReturnType);
+  ID.AddPointer(Dag);
+  for (auto *Param : Params)
+    ID.AddPointer(Param);
+}
+
+ApplyOpInit *ApplyOpInit::get(RecTy *ReturnType, Init *Dag,
+                              SmallVector<Init *, 2> Params) {
+  FoldingSetNodeID ID;
+  ProfileApplyOpInit(ID, ReturnType, Dag, Params);
+
+  detail::RecordKeeperImpl &RK = ReturnType->getRecordKeeper().getImpl();
+  void *IP = nullptr;
+  if (ApplyOpInit *I = RK.TheApplyOpInitPool.FindNodeOrInsertPos(ID, IP))
+    return I;
+
+  ApplyOpInit *I = new (RK.Allocator) ApplyOpInit(ReturnType, Dag, Params);
+  RK.TheApplyOpInitPool.InsertNode(I, IP);
+  return I;
+}
+
+void ApplyOpInit::Profile(FoldingSetNodeID &ID) const {
+  ProfileApplyOpInit(ID, ReturnType, Func, Params);
+}
+
+Init *ApplyOpInit::parseDAG(Record *CurRec, DagInit *Dag) const {
+  SmallVector<Init *, 2> Args;
+  for (unsigned I = 0; I < Dag->getNumArgs(); I++) {
+    Init *Value = Dag->getArg(I);
+    StringRef Name = Dag->getArgNameStr(I);
+    UnsetInit *ValueUnset = dyn_cast<UnsetInit>(Value);
+    if (!ValueUnset && !Name.empty())
+      PrintError(CurRec->getLoc(), Twine("The value and name of DAG in !apply "
+                                         "can't exist at the same time: ") +
+                                       Value->getAsUnquotedString() + ": " +
+                                       Name);
+    Init *Arg = Value;
+    // If the argument is refered to inputs. We replace it with input argument.
+    if (!Name.empty()) {
+      if (!Name.starts_with("a"))
+        PrintError(CurRec->getLoc(),
+                   Twine("The argument should be prefixed with 'a': ") + Name);
+      int ArgIndex = atoi(Name.drop_front().str().c_str());
+      Arg = Params[ArgIndex];
+    }
+    // If the value is a DAG, we parse it recursively.
+    if (DagInit *ValueDag = dyn_cast<DagInit>(Value))
+      Arg = parseDAG(CurRec, ValueDag);
+    Args.push_back(Arg);
+  }
+
+  Init *Op = Dag->getOperator();
+  std::string OpStr = Op->getAsUnquotedString();
+  if (OpStr == "add")
+    return BinOpInit::get(BinOpInit::ADD, Args[0], Args[1], ReturnType)
+        ->Fold(CurRec);
+  if (OpStr == "sub")
+    return BinOpInit::get(BinOpInit::SUB, Args[0], Args[1], ReturnType)
+        ->Fold(CurRec);
+  if (OpStr == "and")
+    return BinOpInit::get(BinOpInit::AND, Args[0], Args[1], ReturnType)
+        ->Fold(CurRec);
+  if (OpStr == "or")
+    return BinOpInit::get(BinOpInit::OR, Args[0], Args[1], ReturnType)
+        ->Fold(CurRec);
+  if (OpStr == "xor")
+    return BinOpInit::get(BinOpInit::XOR, Args[0], Args[1], ReturnType)
+        ->Fold(CurRec);
+  if (OpStr == "srl")
+    return BinOpInit::get(BinOpInit::SRL, Args[0], Args[1], ReturnType)
+        ->Fold(CurRec);
+  if (OpStr == "sra")
+    return BinOpInit::get(BinOpInit::SRA, Args[0], Args[1], ReturnType)
+        ->Fold(CurRec);
+  if (OpStr == "shl")
+    return BinOpInit::get(BinOpInit::SHL, Args[0], Args[1], ReturnType)
+        ->Fold(CurRec);
+
+  PrintError(CurRec->getLoc(),
+             Twine("Unsupported operator in !apply: ") + OpStr);
+  return UnsetInit::get(CurRec->getRecords());
+}
+
+Init *ApplyOpInit::Fold(Record *CurRec, bool IsFinal) const {
+  if (DagInit *Dag = dyn_cast<DagInit>(Func))
+    return parseDAG(CurRec, Dag);
+  return const_cast<ApplyOpInit *>(this);
+}
+
+Init *ApplyOpInit::resolveReferences(Resolver &R) const {
+  Init *NewDag = Func->resolveReferences(R);
+  SmallVector<Init *, 2> NewParams;
+  for (auto *Param : Params) {
+    Init *NewParam = Param->resolveReferences(R);
+    NewParams.push_back(NewParam);
+  }
+
+  ApplyOpInit *Ret = get(ReturnType, NewDag, NewParams);
+  if (isa<DagInit>(NewDag))
+    return Ret->Fold(R.getCurrentRecord());
+
+  return const_cast<ApplyOpInit *>(Ret);
+}
+
+Init *ApplyOpInit::getBit(unsigned Bit) const {
+  return VarBitInit::get(const_cast<ApplyOpInit *>(this), Bit);
+}
+
+std::string ApplyOpInit::getAsString() const {
+  std::string Result =
+      "!apply<" + ReturnType->getAsString() + ">(" + Func->getAsString();
+  for (auto *Param : Params) {
+    Result += ", ";
+    Result += Param->getAsString();
+  }
+  return Result + ")";
+}
+
 RecTy *TypedInit::getFieldType(StringInit *FieldName) const {
   if (RecordRecTy *RecordType = dyn_cast<RecordRecTy>(getType())) {
     for (Record *Rec : RecordType->getClasses()) {
       if (RecordVal *Field = Rec->getValue(FieldName))
         return Field->getType();
     }
   }
   return nullptr;

llvm/lib/TableGen/TGLexer.h

// existing 'Else' that means the preprocessor #else.)
     FalseKW, Field, Foreach, If, In, Include, Int, Let, List, MultiClass,
     String, Then, TrueKW,
 
   // Bang operators.
     XConcat, XADD, XSUB, XMUL, XDIV, XNOT, XLOG2, XAND, XOR, XXOR, XSRA, XSRL,
     XSHL, XListConcat, XListSplat, XStrConcat, XInterleave, XSubstr, XFind,
     XCast, XSubst, XForEach, XFilter, XFoldl, XHead, XTail, XSize, XEmpty, XIf,
     XCond, XEq, XIsA, XDag, XNe, XLe, XLt, XGe, XGt, XSetDagOp, XGetDagOp,
-    XExists, XListRemove,  XToLower, XToUpper,
+    XExists, XListRemove,  XToLower, XToUpper, XApply,
 
   // Boolean literals.
     TrueVal, FalseVal,
 
   // Integer value.
   IntVal,
 
   // Binary constant.  Note that these are sized according to the number of

llvm/lib/TableGen/TGLexer.cpp

tgtok::TokKind Kind =
           .Case("interleave", tgtok::XInterleave)
           .Case("substr", tgtok::XSubstr)
           .Case("find", tgtok::XFind)
           .Cases("setdagop", "setop", tgtok::XSetDagOp) // !setop is deprecated.
           .Cases("getdagop", "getop", tgtok::XGetDagOp) // !getop is deprecated.
           .Case("exists", tgtok::XExists)
           .Case("tolower", tgtok::XToLower)
           .Case("toupper", tgtok::XToUpper)
+          .Case("apply", tgtok::XApply)
           .Default(tgtok::Error);
 
   return Kind != tgtok::Error ? Kind : ReturnError(Start-1, "Unknown operator");
 }
 
 bool TGLexer::prepExitInclude(bool IncludeStackMustBeEmpty) {
   // Report an error, if preprocessor control stack for the current
   // file is not empty.

llvm/lib/TableGen/TGParser.cpp

 //===- TGParser.cpp - Parser for TableGen Files ---------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // Implement the Parser for TableGen.
 //
 //===----------------------------------------------------------------------===//
 
 #include "TGParser.h"
 #include "llvm/ADT/DenseMapInfo.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/Config/llvm-config.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/raw_ostream.h"
+#include "llvm/TableGen/Record.h"
 #include <algorithm>
 #include <cassert>
 #include <cstdint>
 #include <limits>
 
 using namespace llvm;
 
 //===----------------------------------------------------------------------===//
Init *TGParser::ParseIDValue(Record *CurRec, StringInit *Name, SMRange NameLoc,
   }
 
   // Allow self-references of concrete defs, but delay the lookup so that we
   // get the correct type.
   if (CurRec && !CurRec->isClass() && !CurMultiClass &&
       CurRec->getNameInit() == Name)
     return UnOpInit::get(UnOpInit::CAST, Name, CurRec->getType());
 
+  static std::string ApplySubroutines[] = {"add", "sub", "and", "or",
+                                           "xor", "srl", "sra", "shl"};
+  if (find_if(ApplySubroutines,
+              [&](std::string &a) { return a == Name->getValue(); })) {
+    return Name;
+  }
+
   Error(NameLoc.Start, "Variable not defined: '" + Name->getValue() + "'");
   return nullptr;
 }
 
 /// ParseOperation - Parse an operator.  This returns null on error.
 ///
 /// Operation ::= XOperator ['<' Type '>'] '(' Args ')'
 ///
case tgtok::XExists: {
     if (!consume(tgtok::r_paren)) {
       TokError("expected ')' in !exists");
       return nullptr;
     }
 
     return (ExistsOpInit::get(Type, Expr))->Fold(CurRec);
   }
 
+  case tgtok::XApply: {
+    // Value ::= !apply '<' Type '>' '(' dag [',' Value]  ')'
+    Lex.Lex(); // eat the operation
+
+    RecTy *Type = ParseOperatorType();
+    if (!Type)
+      return nullptr;
+
+    if (!consume(tgtok::l_paren)) {
+      TokError("expected '(' after type of !apply");
+      return nullptr;
+    }
+    SMLoc DagLoc = Lex.getLoc();
+    Init *Dag = ParseValue(CurRec);
+    if (!Dag)
+      return nullptr;
+
+    TypedInit *TypedValue = dyn_cast<TypedInit>(Dag);
+    if (!TypedValue) {
+      Error(DagLoc, "expected dag type argument in !apply operator");
+      return nullptr;
+    }
+
+    DagRecTy *DagType = dyn_cast<DagRecTy>(TypedValue->getType());
+    if (!DagType) {
+      Error(DagLoc, "expected dag type argument in !apply operator");
+      return nullptr;
+    }
+
+    SmallVector<Init *, 2> Params;
+    if (consume(tgtok::comma)) {
+      for (;;) {
+        Params.push_back(ParseValue(CurRec));
+        if (!consume(tgtok::comma))
+          break;
+      }
+    }
+
+    if (!consume(tgtok::r_paren)) {
+      TokError("expected ')' in operator");
+      return nullptr;
+    }
+
+    return (ApplyOpInit::get(Type, Dag, Params))->Fold(CurRec);
+  }
+
   case tgtok::XConcat:
   case tgtok::XADD:
   case tgtok::XSUB:
   case tgtok::XMUL:
   case tgtok::XDIV:
   case tgtok::XAND:
   case tgtok::XOR:
   case tgtok::XXOR:
Init *TGParser::ParseSimpleValue(Record *CurRec, RecTy *ItemType,
   case tgtok::XSetDagOp: // Value ::= !binop '(' Value ',' Value ')'
   case tgtok::XIf:
   case tgtok::XCond:
   case tgtok::XFoldl:
   case tgtok::XForEach:
   case tgtok::XFilter:
   case tgtok::XSubst:
   case tgtok::XSubstr:
+  case tgtok::XApply:
   case tgtok::XFind: { // Value ::= !ternop '(' Value ',' Value ',' Value ')'
     return ParseOperation(CurRec, ItemType);
   }
   }
 
   return R;
 }
 

llvm/test/TableGen/apply.td

+// RUN: llvm-tblgen %s | FileCheck %s
+// XFAIL: vg_leak
+
+class Dag<dag v>{
+  dag value = v;
+}
+
+// CHECK: --- Defs ---
+// CHECK: def A0 {
+// CHECK:   int add_result = 24;
+// CHECK:   int sub_result = 18;
+// CHECK:   int and_result = 1;
+// CHECK:   int or_result = 23;
+// CHECK:   int xor_result = 22;
+// CHECK:   int srl_result = 2;
+// CHECK:   int sra_result = 2;
+// CHECK:   int shl_result = 168;
+
+defvar add_dag = Dag<(add $a0, $a1)>.value;
+defvar sub_dag = Dag<(sub $a0, $a1)>.value;
+defvar xor_dag = !dag(xor, [?, ?]<int>, ["a0", "a1"]);
+defvar srl_dag = !dag(srl, [?, ?]<int>, ["a0", "a1"]);
+
+class A<int a, int b> {
+  int add_result = !apply<int>(add_dag, a, b);
+  int sub_result = !apply<int>(sub_dag, a, b);
+  int and_result = !apply<int>(Dag<(and $a0, $a1)>.value, a, b);
+  int or_result  = !apply<int>(Dag<(or  $a0, $a1)>.value,  a, b);
+  int xor_result = !apply<int>(xor_dag, a, b);
+  int srl_result = !apply<int>(srl_dag, a, b);
+  int sra_result = !apply<int>((sra $a0, $a1), a, b);
+  int shl_result = !apply<int>((shl $a0, $a1), a, b);
+}
+
+def A0 : A<21, 3>;
+
+// CHECK: def B0 {
+// CHECK:   int result = 254;
+
+class B<Dag op, int a, int b> {
+  int result = !apply<int>(op.value, a, b);
+}
+
+defvar c = 233;
+// Capture variables in lexical scope.
+def B0 : B<Dag<(add $a0, c)>, 21, 3>;