Browse Source

Add MLIR support (#1044)

Lutz Roeder 2 months ago
parent
commit
004ed028e1
4 changed files with 668 additions and 301 deletions
  1. 530 2
      source/mlir-metadata.json
  2. 133 296
      source/mlir.js
  3. 3 3
      tools/mlir
  4. 2 0
      tools/mlir-script.js

+ 530 - 2
source/mlir-metadata.json

@@ -42754,8 +42754,7 @@
       { "name": "signA", "type": "DefaultValuedAttr<I1Attr, 0>" },
       { "name": "signB", "type": "DefaultValuedAttr<I1Attr, 0>" },
       { "name": "reuseA", "type": "DefaultValuedAttr<I1Attr, 0>" },
-      { "name": "reuseB", "type": "DefaultValuedAttr<I1Attr, 0>" },
-      { "name": "clamp", "type": "DefaultValuedAttr<I1Attr, 0>" }
+      { "name": "reuseB", "type": "DefaultValuedAttr<I1Attr, 0>" }
     ],
     "assemblyFormat": "$a `,` $b `,` $c attr-dict `:` functional-type(operands, $res)"
   },
@@ -74602,6 +74601,535 @@
       { "name": "result", "type": "TFRT_ChainType" }
     ]
   },
+  {
+    "name": "tfg.Case",
+    "summary": "A case operation, in functional form",
+    "operands": [
+      { "name": "branch_index", "type": "TFGraph_OpaqueTensorOr<I32Tensor>" },
+      { "name": "args", "type": "Variadic<TFGraph_TensorOrControlType>" }
+    ],
+    "results": [
+      { "name": "outs", "type": "Variadic<TFGraph_Tensor>" },
+      { "name": "ctl", "type": "ControlType" }
+    ],
+    "attributes": [
+      { "name": "branches", "type": "TypedArrayAttrBase<TFType_FuncAttr, an array of functions>" },
+      { "name": "Tin", "type": "OptionalAttr<TypedArrayAttrBase<TypeAttr>>" },
+      { "name": "Tout", "type": "OptionalAttr<TypedArrayAttrBase<TypeAttr>>" },
+      { "name": "output_shapes", "type": "OptionalAttr<TypedArrayAttrBase<TFType_ShapeAttrDef>>" }
+    ]
+  },
+  {
+    "name": "tfg.CaseRegion",
+    "summary": "A case op with graph regions.",
+    "description": "A `Case` operation in region form. `Case` takes one data operand, an index,\n    and executes the region corresponding to that index. The execution of any\n    of its regions is blocked on the operation's control operands.\n\n    The regions have no block arguments; everything that escapes the scope is an\n    implicit capture. All regions return values of the same types which are\n    compatible with the result types of the `Case` operation.\n\n    For each region, the operation contains optional \"call\" attributes and\n    region attributes in arrays. There are never region argument attributes, and\n    the number of result attributes must match the number of data results.",
+    "operands": [
+      { "name": "branch_index", "type": "I32Tensor" },
+      { "name": "ctls", "type": "Variadic<ControlType>" }
+    ],
+    "results": [
+      { "name": "outs", "type": "Variadic<TFGraph_Tensor>" },
+      { "name": "ctl", "type": "ControlType" }
+    ],
+    "attributes": [
+      { "name": "branch_attrs", "type": "OptionalAttr<TypedArrayAttrBase<DictionaryAttr>>" },
+      { "name": "region_attrs", "type": "OptionalAttr<TypedArrayAttrBase<TFGraph_RegionAttrs, case region preserved attributes>>" }
+    ],
+    "regions": [
+      { "name": "branches", "type": "VariadicRegion<SizedRegion<1>>" }
+    ],
+    "defaultDialect": "tfg",
+    "assemblyFormat": "$branch_index (` ` `[` $ctls^ `]`)?\n      $branches\n    attr-dict `:` functional-type($branch_index, $outs)"
+  },
+  {
+    "name": "tfg.condition",
+    "summary": "TF graph terminator op for while-loop condition regions.",
+    "description": "The `condition` operation is a special terminator for the condition region\n    of while loops. Its only purpose is to support reasoning about region\n    control-flow by forwarding the arguments to the condition region to the body\n    region or back to the parent operation's results.\n\n    The `condition` operation accepts one `cond` data input, which if true, will\n    cause the execution to branch to the loop body, and if false, will cause the\n    execution to return to the parent operation.\n\n    There must be the same number of arguments `args` as there are operands and\n    results in the parent `While` loop.",
+    "operands": [
+      { "name": "cond", "type": "I1Tensor" },
+      { "name": "args", "type": "Variadic<TFGraph_Tensor>" },
+      { "name": "ctls", "type": "Variadic<ControlType>" }
+    ],
+    "traits": [
+      { "type": "AttrSizedOperandSegments" }
+    ],
+    "assemblyFormat": "$cond `:` type($cond) (` ` `(` $args^ `)`)?\n                          (`[` $ctls^ `]`)? attr-dict (`:` type($args)^)?"
+  },
+  {
+    "name": "tfg.For",
+    "summary": "A functional for-loop operation.",
+    "operands": [
+      { "name": "start", "type": "TFGraph_OpaqueTensorOr<I32Tensor>" },
+      { "name": "limit", "type": "TFGraph_OpaqueTensorOr<I32Tensor>" },
+      { "name": "delta", "type": "TFGraph_OpaqueTensorOr<I32Tensor>" },
+      { "name": "args", "type": "Variadic<TFGraph_TensorOrControlType>" }
+    ],
+    "results": [
+      { "name": "outs", "type": "Variadic<TFGraph_Tensor>" },
+      { "name": "ctl", "type": "ControlType" }
+    ],
+    "attributes": [
+      { "name": "body", "type": "TFType_FuncAttr" },
+      { "name": "T", "type": "OptionalAttr<TypedArrayAttrBase<TypeAttr>>" }
+    ]
+  },
+  {
+    "name": "tfg.ForRegion",
+    "description": "A `For` operation in region form. The operation takes three scalar `i32`\n    operands: a lower bound, an upper bound, and a step. As well, it takes N\n    operands and produces N results of the same types. The loop body region has\n    a scalar `i32` block argument and its corresponding control token and N data\n    block arguments of the same types as the N data operands and N results. The\n    scalar `i32` block argument is the current loop index.\n\n    The operation is blocked on its control operands. The execution enters the\n    loop body, which must return N data results. The values are either fed into\n    the next iteration of the body or taken as the operation's results.\n\n    For the body region, the operation contains optional \"call\" attributes and\n    region attributes. The number of argument attributes must match the number\n    of data operands and the number of result attributes must match the number\n    of data results.\n\n    Values can be propagated through the loop bodies as iteration variables,\n    but the regions allow implicitly captured values as well.",
+    "operands": [
+      { "name": "start", "type": "I32Tensor" },
+      { "name": "limit", "type": "I32Tensor" },
+      { "name": "delta", "type": "I32Tensor" },
+      { "name": "init", "type": "Variadic<TFGraph_Tensor>" },
+      { "name": "ctls", "type": "Variadic<ControlType>" }
+    ],
+    "results": [
+      { "name": "outs", "type": "Variadic<TFGraph_Tensor>" },
+      { "name": "ctl", "type": "ControlType" }
+    ],
+    "attributes": [
+      { "name": "body_attrs", "type": "OptionalAttr<DictionaryAttr>" },
+      { "name": "region_attrs", "type": "OptionalAttr<TFGraph_RegionAttrs>" }
+    ],
+    "regions": [
+      { "name": "body_region", "type": "SizedRegion<1>" }
+    ],
+    "defaultDialect": "tfg",
+    "traits": [
+      { "type": "AttrSizedOperandSegments" }
+    ],
+    "assemblyFormat": "(`(` $init^ `)`)? (`[` $ctls^ `]`)?\n    `from` $start `to` $limit `by` $delta\n      $body_region\n    attr-dict `:` `(` type($start) `,` type($limit) `,` type($delta)\n    (`,` type($init)^)? `)` (`->` `(` type($outs)^ `)`)?"
+  },
+  {
+    "name": "tfg.func",
+    "summary": "Define a function as modeled by FunctionDef.",
+    "description": "This operation is modeling a function definition, corresponding to the\n    FunctionDef proto. It defines a region `body` which contains an unordered\n    list of nodes forming a TensorFlow graph.\n    The body still has a return operation as terminator which must always be\n    the last operation.\n\n    A function can be generic if it specifies the `generic` attribute. Such\n    function is akin to a C++ template: it need \"instantiation attributes\" to be\n    specified at each call site in order to make it into a concrete function,\n    see `tensorflow::InstantiateFunction()` for more info.",
+    "attributes": [
+      { "name": "generic", "type": "UnitAttr" },
+      { "name": "sym_name", "type": "SymbolNameAttr" },
+      { "name": "function_type", "type": "TypeAttrOf<FunctionType>" },
+      { "name": "arg_attrs", "type": "OptionalAttr<ArrayAttr>" },
+      { "name": "res_attrs", "type": "OptionalAttr<ArrayAttr>" },
+      { "name": "description", "type": "OptionalAttr<StrAttr>" },
+      { "name": "is_stateful", "type": "UnitAttr" },
+      { "name": "gradient", "type": "OptionalAttr<FlatSymbolRefAttr>" },
+      { "name": "resource_arg_unique_ids_keys", "type": "OptionalAttr<I32ElementsAttr>" },
+      { "name": "resource_arg_unique_ids_values", "type": "OptionalAttr<I32ElementsAttr>" },
+      { "name": "control_output", "type": "OptionalAttr<TypedArrayAttrBase<StrAttr>>" }
+    ],
+    "regions": [
+      { "name": "body", "type": "AnyRegion" }
+    ],
+    "defaultDialect": "tfg",
+    "hasCustomAssemblyFormat": true
+  },
+  {
+    "name": "tfg.get_result",
+    "summary": "Extract a value out of an opaque tensor in a generic Function.",
+    "description": "The `get_result` operation exists in generic functions before instantiation,\n    when the number of output per node isn't known yet. In generic function\n    nodes define a single SSA value which is then \"unpacked\" by name using this\n    operation.\n    Example:\n\n    ```mlir\n    %Switch, %ctl_1 = tfg.Switch(%G_z, %G_z) name = \"cond/Switch\"\n    %s_true = tfg.get_result %Switch \"output_true” : 0\n    %s_false = tfg.get_result %Switch \"output_false” : 0\n    ```",
+    "operands": [
+      { "name": "value", "type": "OpaqueTensorType" }
+    ],
+    "results": [
+      { "name": "result", "type": "OpaqueTensorType" }
+    ],
+    "attributes": [
+      { "name": "name", "type": "StrAttr" },
+      { "name": "number", "type": "I32Attr" }
+    ],
+    "assemblyFormat": "`(` $value `)` $name `:` $number attr-dict-with-keyword"
+  },
+  {
+    "name": "tfg.graph",
+    "summary": "Define a flat graph of operations",
+    "description": "This operation is modeling a unordered list of nodes forming a TensorFlow\n    graph. This corresponds to the GraphDef proto definition minus the\n    FunctionDefLibrary.",
+    "attributes": [
+      { "name": "version", "type": "TFType_VersionAttr" }
+    ],
+    "regions": [
+      { "name": "nodes", "type": "SizedRegion<1>" }
+    ],
+    "defaultDialect": "tfg",
+    "assemblyFormat": "qualified($version) attr-dict-with-keyword $nodes"
+  },
+  {
+    "name": "tfg.If",
+    "summary": "An if operation, in functional form.",
+    "operands": [
+      { "name": "cond", "type": "TFGraph_OpaqueTensorOr<I1Tensor>" },
+      { "name": "args", "type": "Variadic<TFGraph_TensorOrControlType>" }
+    ],
+    "results": [
+      { "name": "outs", "type": "Variadic<TFGraph_Tensor>" },
+      { "name": "ctl", "type": "ControlType" }
+    ],
+    "attributes": [
+      { "name": "then_branch", "type": "TFType_FuncAttr" },
+      { "name": "else_branch", "type": "TFType_FuncAttr" },
+      { "name": "Tcond", "type": "OptionalAttr<TypeAttr>" },
+      { "name": "Tin", "type": "OptionalAttr<TypedArrayAttrBase<TypeAttr>>" },
+      { "name": "Tout", "type": "OptionalAttr<TypedArrayAttrBase<TypeAttr>>" },
+      { "name": "output_shapes", "type": "OptionalAttr<TypedArrayAttrBase<TFType_ShapeAttrDef>>" }
+    ]
+  },
+  {
+    "name": "tfg.IfRegion",
+    "summary": "An if-then-else op with graph regions.",
+    "description": "An `If` operation in region form. `If` takes on data operand, which if true,\n    executes the `then` region, otherwise the `else` region. The execution of\n    either region is blocked on the operation's control operands.\n\n    The regions have no block arguments; everything that escapes the scope is an\n    implicit capture. Both regions return values of the same types which are\n    compatible with the result types of the `If` operation.\n\n    For each region, the operation contains optional \"call\" attributes and\n    region attributes. There are never region argument attributes, and the\n    number of result attributes must match the number of data results.",
+    "operands": [
+      { "name": "cond", "type": "I1Tensor" },
+      { "name": "ctls", "type": "Variadic<ControlType>" }
+    ],
+    "results": [
+      { "name": "outs", "type": "Variadic<TFGraph_Tensor>" },
+      { "name": "ctl", "type": "ControlType" }
+    ],
+    "attributes": [
+      { "name": "then_attrs", "type": "OptionalAttr<DictionaryAttr>" },
+      { "name": "else_attrs", "type": "OptionalAttr<DictionaryAttr>" },
+      { "name": "then_region_attrs", "type": "OptionalAttr<TFGraph_RegionAttrs>" },
+      { "name": "else_region_attrs", "type": "OptionalAttr<TFGraph_RegionAttrs>" }
+    ],
+    "regions": [
+      { "name": "then_region", "type": "SizedRegion<1>" },
+      { "name": "else_region", "type": "SizedRegion<1>" }
+    ],
+    "defaultDialect": "tfg",
+    "assemblyFormat": "$cond (` ` `[` $ctls^ `]`)? `then`\n      $then_region\n    `else`\n      $else_region\n    attr-dict `:` functional-type($cond, $outs)"
+  },
+  {
+    "name": "tfg.return",
+    "summary": "Return values from a Function.",
+    "description": "The `return` operation represents a return operation within a function.\n    The operation takes variable number of operands and produces no results.\n    The operand number must match the signature of the function that contains\n    the operation. The types of the input to the call must be compatible with\n    the function signature in terms of shape and matching the element type, and\n    similarly for the function results,",
+    "attributes": [
+      { "name": "control_ret_attrs", "type": "TypedArrayAttrBase<DictionaryAttr>" }
+    ],
+    "hasCustomAssemblyFormat": true
+  },
+  {
+    "name": "tfg.StatefulCase",
+    "summary": "A stateful case operation, in functional form",
+    "operands": [
+      { "name": "branch_index", "type": "TFGraph_OpaqueTensorOr<I32Tensor>" },
+      { "name": "args", "type": "Variadic<TFGraph_TensorOrControlType>" }
+    ],
+    "results": [
+      { "name": "outs", "type": "Variadic<TFGraph_Tensor>" },
+      { "name": "ctl", "type": "ControlType" }
+    ],
+    "attributes": [
+      { "name": "branches", "type": "TypedArrayAttrBase<TFType_FuncAttr, an array of functions>" },
+      { "name": "Tin", "type": "OptionalAttr<TypedArrayAttrBase<TypeAttr>>" },
+      { "name": "Tout", "type": "OptionalAttr<TypedArrayAttrBase<TypeAttr>>" },
+      { "name": "output_shapes", "type": "OptionalAttr<TypedArrayAttrBase<TFType_ShapeAttrDef>>" }
+    ]
+  },
+  {
+    "name": "tfg.StatefulCaseRegion",
+    "summary": "A stateful case op with graph regions.",
+    "description": "A `Case` operation in region form. `Case` takes one data operand, an index,\n    and executes the region corresponding to that index. The execution of any\n    of its regions is blocked on the operation's control operands.\n\n    The regions have no block arguments; everything that escapes the scope is an\n    implicit capture. All regions return values of the same types which are\n    compatible with the result types of the `Case` operation.\n\n    For each region, the operation contains optional \"call\" attributes and\n    region attributes in arrays. There are never region argument attributes, and\n    the number of result attributes must match the number of data results.",
+    "operands": [
+      { "name": "branch_index", "type": "I32Tensor" },
+      { "name": "ctls", "type": "Variadic<ControlType>" }
+    ],
+    "results": [
+      { "name": "outs", "type": "Variadic<TFGraph_Tensor>" },
+      { "name": "ctl", "type": "ControlType" }
+    ],
+    "attributes": [
+      { "name": "branch_attrs", "type": "OptionalAttr<TypedArrayAttrBase<DictionaryAttr>>" },
+      { "name": "region_attrs", "type": "OptionalAttr<TypedArrayAttrBase<TFGraph_RegionAttrs, case region preserved attributes>>" }
+    ],
+    "regions": [
+      { "name": "branches", "type": "VariadicRegion<SizedRegion<1>>" }
+    ],
+    "defaultDialect": "tfg",
+    "assemblyFormat": "$branch_index (` ` `[` $ctls^ `]`)?\n      $branches\n    attr-dict `:` functional-type($branch_index, $outs)"
+  },
+  {
+    "name": "tfg.StatefulIf",
+    "summary": "A stateful if operation, in functional form.",
+    "operands": [
+      { "name": "cond", "type": "TFGraph_OpaqueTensorOr<I1Tensor>" },
+      { "name": "args", "type": "Variadic<TFGraph_TensorOrControlType>" }
+    ],
+    "results": [
+      { "name": "outs", "type": "Variadic<TFGraph_Tensor>" },
+      { "name": "ctl", "type": "ControlType" }
+    ],
+    "attributes": [
+      { "name": "then_branch", "type": "TFType_FuncAttr" },
+      { "name": "else_branch", "type": "TFType_FuncAttr" },
+      { "name": "Tcond", "type": "OptionalAttr<TypeAttr>" },
+      { "name": "Tin", "type": "OptionalAttr<TypedArrayAttrBase<TypeAttr>>" },
+      { "name": "Tout", "type": "OptionalAttr<TypedArrayAttrBase<TypeAttr>>" },
+      { "name": "output_shapes", "type": "OptionalAttr<TypedArrayAttrBase<TFType_ShapeAttrDef>>" }
+    ]
+  },
+  {
+    "name": "tfg.StatefulIfRegion",
+    "summary": "A stateful if-then-else op with graph regions.",
+    "description": "An `If` operation in region form. `If` takes on data operand, which if true,\n    executes the `then` region, otherwise the `else` region. The execution of\n    either region is blocked on the operation's control operands.\n\n    The regions have no block arguments; everything that escapes the scope is an\n    implicit capture. Both regions return values of the same types which are\n    compatible with the result types of the `If` operation.\n\n    For each region, the operation contains optional \"call\" attributes and\n    region attributes. There are never region argument attributes, and the\n    number of result attributes must match the number of data results.",
+    "operands": [
+      { "name": "cond", "type": "I1Tensor" },
+      { "name": "ctls", "type": "Variadic<ControlType>" }
+    ],
+    "results": [
+      { "name": "outs", "type": "Variadic<TFGraph_Tensor>" },
+      { "name": "ctl", "type": "ControlType" }
+    ],
+    "attributes": [
+      { "name": "then_attrs", "type": "OptionalAttr<DictionaryAttr>" },
+      { "name": "else_attrs", "type": "OptionalAttr<DictionaryAttr>" },
+      { "name": "then_region_attrs", "type": "OptionalAttr<TFGraph_RegionAttrs>" },
+      { "name": "else_region_attrs", "type": "OptionalAttr<TFGraph_RegionAttrs>" }
+    ],
+    "regions": [
+      { "name": "then_region", "type": "SizedRegion<1>" },
+      { "name": "else_region", "type": "SizedRegion<1>" }
+    ],
+    "defaultDialect": "tfg",
+    "assemblyFormat": "$cond (` ` `[` $ctls^ `]`)? `then`\n      $then_region\n    `else`\n      $else_region\n    attr-dict `:` functional-type($cond, $outs)"
+  },
+  {
+    "name": "tfg.StatefulWhile",
+    "summary": "A functional stateful while loop operation.",
+    "operands": [
+      { "name": "args", "type": "Variadic<TFGraph_TensorOrControlType>" }
+    ],
+    "results": [
+      { "name": "outs", "type": "Variadic<TFGraph_Tensor>" },
+      { "name": "ctl", "type": "ControlType" }
+    ],
+    "attributes": [
+      { "name": "cond", "type": "TFType_FuncAttr" },
+      { "name": "body", "type": "TFType_FuncAttr" },
+      { "name": "parallel_iterations", "type": "I64Attr" },
+      { "name": "T", "type": "OptionalAttr<TypedArrayAttrBase<TypeAttr>>" },
+      { "name": "output_shapes", "type": "OptionalAttr<TypedArrayAttrBase<TFType_ShapeAttrDef>>" }
+    ]
+  },
+  {
+    "name": "tfg.StatefulWhileRegion",
+    "summary": "A stateful region-based while loop operation.",
+    "description": "A `While` operation in region form. The operation takes N operands and\n    produces N results of the same types. Each region has N data block arguments\n    and N corresponding control tokens.\n\n    The entry into the condition region is blocked on the operation's control\n    operands. The operation always executes the condition region at least once\n    with the arguments as the operation's data operands. If the condition region\n    returns `true`, the body region is executed with the arguments provided to\n    the condition region. The body region will branch back to the condition\n    region with N results. If `false`, the condition region will branch back to\n    the parent `While` operation with the argument values as results.\n\n    For each region, the operation contains optional \"call\" attributes and\n    region attributes. The number of argument attributes must match the number\n    of data operands and the number of result attributes must match the number\n    of data results.\n\n    Values can be propagated through the loop bodies as iteration variables,\n    but the regions allow implicitly captured values as well.",
+    "operands": [
+      { "name": "init", "type": "Variadic<TFGraph_Tensor>" },
+      { "name": "ctls", "type": "Variadic<ControlType>" }
+    ],
+    "results": [
+      { "name": "outs", "type": "Variadic<TFGraph_Tensor>" },
+      { "name": "ctl", "type": "ControlType" }
+    ],
+    "attributes": [
+      { "name": "parallel_iterations", "type": "I64Attr" },
+      { "name": "cond_attrs", "type": "OptionalAttr<DictionaryAttr>" },
+      { "name": "body_attrs", "type": "OptionalAttr<DictionaryAttr>" },
+      { "name": "cond_region_attrs", "type": "OptionalAttr<TFGraph_RegionAttrs>" },
+      { "name": "body_region_attrs", "type": "OptionalAttr<TFGraph_RegionAttrs>" }
+    ],
+    "regions": [
+      { "name": "cond_region", "type": "SizedRegion<1>" },
+      { "name": "body_region", "type": "SizedRegion<1>" }
+    ],
+    "defaultDialect": "tfg",
+    "traits": [
+      { "type": "AttrSizedOperandSegments" }
+    ],
+    "assemblyFormat": "(`(` $init^ `)`)? (`[` $ctls^ `]`)?\n      $cond_region\n    `do`\n      $body_region\n    attr-dict (`:` functional-type($init, $outs)^)?"
+  },
+  {
+    "name": "tfg.StatelessCase",
+    "summary": "A stateless case operation, in functional form",
+    "operands": [
+      { "name": "branch_index", "type": "TFGraph_OpaqueTensorOr<I32Tensor>" },
+      { "name": "args", "type": "Variadic<TFGraph_TensorOrControlType>" }
+    ],
+    "results": [
+      { "name": "outs", "type": "Variadic<TFGraph_Tensor>" },
+      { "name": "ctl", "type": "ControlType" }
+    ],
+    "attributes": [
+      { "name": "branches", "type": "TypedArrayAttrBase<TFType_FuncAttr, an array of functions>" },
+      { "name": "Tin", "type": "OptionalAttr<TypedArrayAttrBase<TypeAttr>>" },
+      { "name": "Tout", "type": "OptionalAttr<TypedArrayAttrBase<TypeAttr>>" },
+      { "name": "output_shapes", "type": "OptionalAttr<TypedArrayAttrBase<TFType_ShapeAttrDef>>" }
+    ]
+  },
+  {
+    "name": "tfg.StatelessCaseRegion",
+    "summary": "A stateless case op with graph regions.",
+    "description": "A `Case` operation in region form. `Case` takes one data operand, an index,\n    and executes the region corresponding to that index. The execution of any\n    of its regions is blocked on the operation's control operands.\n\n    The regions have no block arguments; everything that escapes the scope is an\n    implicit capture. All regions return values of the same types which are\n    compatible with the result types of the `Case` operation.\n\n    For each region, the operation contains optional \"call\" attributes and\n    region attributes in arrays. There are never region argument attributes, and\n    the number of result attributes must match the number of data results.",
+    "operands": [
+      { "name": "branch_index", "type": "I32Tensor" },
+      { "name": "ctls", "type": "Variadic<ControlType>" }
+    ],
+    "results": [
+      { "name": "outs", "type": "Variadic<TFGraph_Tensor>" },
+      { "name": "ctl", "type": "ControlType" }
+    ],
+    "attributes": [
+      { "name": "branch_attrs", "type": "OptionalAttr<TypedArrayAttrBase<DictionaryAttr>>" },
+      { "name": "region_attrs", "type": "OptionalAttr<TypedArrayAttrBase<TFGraph_RegionAttrs, case region preserved attributes>>" }
+    ],
+    "regions": [
+      { "name": "branches", "type": "VariadicRegion<SizedRegion<1>>" }
+    ],
+    "defaultDialect": "tfg",
+    "assemblyFormat": "$branch_index (` ` `[` $ctls^ `]`)?\n      $branches\n    attr-dict `:` functional-type($branch_index, $outs)"
+  },
+  {
+    "name": "tfg.StatelessIf",
+    "summary": "A stateless if operation, in functional form.",
+    "operands": [
+      { "name": "cond", "type": "TFGraph_OpaqueTensorOr<I1Tensor>" },
+      { "name": "args", "type": "Variadic<TFGraph_TensorOrControlType>" }
+    ],
+    "results": [
+      { "name": "outs", "type": "Variadic<TFGraph_Tensor>" },
+      { "name": "ctl", "type": "ControlType" }
+    ],
+    "attributes": [
+      { "name": "then_branch", "type": "TFType_FuncAttr" },
+      { "name": "else_branch", "type": "TFType_FuncAttr" },
+      { "name": "Tcond", "type": "OptionalAttr<TypeAttr>" },
+      { "name": "Tin", "type": "OptionalAttr<TypedArrayAttrBase<TypeAttr>>" },
+      { "name": "Tout", "type": "OptionalAttr<TypedArrayAttrBase<TypeAttr>>" },
+      { "name": "output_shapes", "type": "OptionalAttr<TypedArrayAttrBase<TFType_ShapeAttrDef>>" }
+    ]
+  },
+  {
+    "name": "tfg.StatelessIfRegion",
+    "summary": "A stateless if-then-else op with graph regions.",
+    "description": "An `If` operation in region form. `If` takes on data operand, which if true,\n    executes the `then` region, otherwise the `else` region. The execution of\n    either region is blocked on the operation's control operands.\n\n    The regions have no block arguments; everything that escapes the scope is an\n    implicit capture. Both regions return values of the same types which are\n    compatible with the result types of the `If` operation.\n\n    For each region, the operation contains optional \"call\" attributes and\n    region attributes. There are never region argument attributes, and the\n    number of result attributes must match the number of data results.",
+    "operands": [
+      { "name": "cond", "type": "I1Tensor" },
+      { "name": "ctls", "type": "Variadic<ControlType>" }
+    ],
+    "results": [
+      { "name": "outs", "type": "Variadic<TFGraph_Tensor>" },
+      { "name": "ctl", "type": "ControlType" }
+    ],
+    "attributes": [
+      { "name": "then_attrs", "type": "OptionalAttr<DictionaryAttr>" },
+      { "name": "else_attrs", "type": "OptionalAttr<DictionaryAttr>" },
+      { "name": "then_region_attrs", "type": "OptionalAttr<TFGraph_RegionAttrs>" },
+      { "name": "else_region_attrs", "type": "OptionalAttr<TFGraph_RegionAttrs>" }
+    ],
+    "regions": [
+      { "name": "then_region", "type": "SizedRegion<1>" },
+      { "name": "else_region", "type": "SizedRegion<1>" }
+    ],
+    "defaultDialect": "tfg",
+    "assemblyFormat": "$cond (` ` `[` $ctls^ `]`)? `then`\n      $then_region\n    `else`\n      $else_region\n    attr-dict `:` functional-type($cond, $outs)"
+  },
+  {
+    "name": "tfg.StatelessWhile",
+    "summary": "A functional stateless while loop operation.",
+    "operands": [
+      { "name": "args", "type": "Variadic<TFGraph_TensorOrControlType>" }
+    ],
+    "results": [
+      { "name": "outs", "type": "Variadic<TFGraph_Tensor>" },
+      { "name": "ctl", "type": "ControlType" }
+    ],
+    "attributes": [
+      { "name": "cond", "type": "TFType_FuncAttr" },
+      { "name": "body", "type": "TFType_FuncAttr" },
+      { "name": "parallel_iterations", "type": "I64Attr" },
+      { "name": "T", "type": "OptionalAttr<TypedArrayAttrBase<TypeAttr>>" },
+      { "name": "output_shapes", "type": "OptionalAttr<TypedArrayAttrBase<TFType_ShapeAttrDef>>" }
+    ]
+  },
+  {
+    "name": "tfg.StatelessWhileRegion",
+    "summary": "A stateless region-based while loop operation.",
+    "description": "A `While` operation in region form. The operation takes N operands and\n    produces N results of the same types. Each region has N data block arguments\n    and N corresponding control tokens.\n\n    The entry into the condition region is blocked on the operation's control\n    operands. The operation always executes the condition region at least once\n    with the arguments as the operation's data operands. If the condition region\n    returns `true`, the body region is executed with the arguments provided to\n    the condition region. The body region will branch back to the condition\n    region with N results. If `false`, the condition region will branch back to\n    the parent `While` operation with the argument values as results.\n\n    For each region, the operation contains optional \"call\" attributes and\n    region attributes. The number of argument attributes must match the number\n    of data operands and the number of result attributes must match the number\n    of data results.\n\n    Values can be propagated through the loop bodies as iteration variables,\n    but the regions allow implicitly captured values as well.",
+    "operands": [
+      { "name": "init", "type": "Variadic<TFGraph_Tensor>" },
+      { "name": "ctls", "type": "Variadic<ControlType>" }
+    ],
+    "results": [
+      { "name": "outs", "type": "Variadic<TFGraph_Tensor>" },
+      { "name": "ctl", "type": "ControlType" }
+    ],
+    "attributes": [
+      { "name": "parallel_iterations", "type": "I64Attr" },
+      { "name": "cond_attrs", "type": "OptionalAttr<DictionaryAttr>" },
+      { "name": "body_attrs", "type": "OptionalAttr<DictionaryAttr>" },
+      { "name": "cond_region_attrs", "type": "OptionalAttr<TFGraph_RegionAttrs>" },
+      { "name": "body_region_attrs", "type": "OptionalAttr<TFGraph_RegionAttrs>" }
+    ],
+    "regions": [
+      { "name": "cond_region", "type": "SizedRegion<1>" },
+      { "name": "body_region", "type": "SizedRegion<1>" }
+    ],
+    "defaultDialect": "tfg",
+    "traits": [
+      { "type": "AttrSizedOperandSegments" }
+    ],
+    "assemblyFormat": "(`(` $init^ `)`)? (`[` $ctls^ `]`)?\n      $cond_region\n    `do`\n      $body_region\n    attr-dict (`:` functional-type($init, $outs)^)?"
+  },
+  {
+    "name": "tfg.While",
+    "summary": "A functional while loop operation.",
+    "operands": [
+      { "name": "args", "type": "Variadic<TFGraph_TensorOrControlType>" }
+    ],
+    "results": [
+      { "name": "outs", "type": "Variadic<TFGraph_Tensor>" },
+      { "name": "ctl", "type": "ControlType" }
+    ],
+    "attributes": [
+      { "name": "cond", "type": "TFType_FuncAttr" },
+      { "name": "body", "type": "TFType_FuncAttr" },
+      { "name": "parallel_iterations", "type": "I64Attr" },
+      { "name": "T", "type": "OptionalAttr<TypedArrayAttrBase<TypeAttr>>" },
+      { "name": "output_shapes", "type": "OptionalAttr<TypedArrayAttrBase<TFType_ShapeAttrDef>>" }
+    ]
+  },
+  {
+    "name": "tfg.WhileRegion",
+    "summary": "A region-based while loop operation.",
+    "description": "A `While` operation in region form. The operation takes N operands and\n    produces N results of the same types. Each region has N data block arguments\n    and N corresponding control tokens.\n\n    The entry into the condition region is blocked on the operation's control\n    operands. The operation always executes the condition region at least once\n    with the arguments as the operation's data operands. If the condition region\n    returns `true`, the body region is executed with the arguments provided to\n    the condition region. The body region will branch back to the condition\n    region with N results. If `false`, the condition region will branch back to\n    the parent `While` operation with the argument values as results.\n\n    For each region, the operation contains optional \"call\" attributes and\n    region attributes. The number of argument attributes must match the number\n    of data operands and the number of result attributes must match the number\n    of data results.\n\n    Values can be propagated through the loop bodies as iteration variables,\n    but the regions allow implicitly captured values as well.",
+    "operands": [
+      { "name": "init", "type": "Variadic<TFGraph_Tensor>" },
+      { "name": "ctls", "type": "Variadic<ControlType>" }
+    ],
+    "results": [
+      { "name": "outs", "type": "Variadic<TFGraph_Tensor>" },
+      { "name": "ctl", "type": "ControlType" }
+    ],
+    "attributes": [
+      { "name": "parallel_iterations", "type": "I64Attr" },
+      { "name": "cond_attrs", "type": "OptionalAttr<DictionaryAttr>" },
+      { "name": "body_attrs", "type": "OptionalAttr<DictionaryAttr>" },
+      { "name": "cond_region_attrs", "type": "OptionalAttr<TFGraph_RegionAttrs>" },
+      { "name": "body_region_attrs", "type": "OptionalAttr<TFGraph_RegionAttrs>" }
+    ],
+    "regions": [
+      { "name": "cond_region", "type": "SizedRegion<1>" },
+      { "name": "body_region", "type": "SizedRegion<1>" }
+    ],
+    "defaultDialect": "tfg",
+    "traits": [
+      { "type": "AttrSizedOperandSegments" }
+    ],
+    "assemblyFormat": "(`(` $init^ `)`)? (`[` $ctls^ `]`)?\n      $cond_region\n    `do`\n      $body_region\n    attr-dict (`:` functional-type($init, $outs)^)?"
+  },
+  {
+    "name": "tfg.yield",
+    "summary": "TF graph terminator for control-flow regions.",
+    "description": "The `yield` operation is the terminator for control-flow regions. The\n    operation takes data operands and control results. All data operands are\n    inputs for successor regions. Typically, this means that the number of\n    data operands matches the number of results of the parent operation.\n\n    The semantics of function calls in TensorFlow mean that the region \"returns\"\n    data operands as they become available so that users can start executing\n    immediately. All functional control-flow operations behave as multi-device\n    function calls, which means that the control result of the parent operation\n    depends on the control operands of the `yield`.",
+    "operands": [
+      { "name": "args", "type": "Variadic<TFGraph_Tensor>" },
+      { "name": "ctls", "type": "Variadic<ControlType>" }
+    ],
+    "traits": [
+      { "type": "AttrSizedOperandSegments" }
+    ],
+    "assemblyFormat": "(`(` $args^ `)`)? (`[` $ctls^ `]`)? attr-dict\n                          (`:` type($args)^)?"
+  },
   {
     "name": "tfl.abs",
     "summary": "Absolute value operator",

File diff suppressed because it is too large
+ 133 - 296
source/mlir.js


+ 3 - 3
tools/mlir

@@ -13,7 +13,7 @@ entries=(
     "https://github.com/onnx/onnx-mlir.git|main|${src_dir}/onnx-mlir|"
     "https://github.com/llvm/torch-mlir.git|main|${src_dir}/torch-mlir|"
     "https://github.com/triton-lang/triton|main|${src_dir}/triton|"
-    "https://github.com/tensorflow/tensorflow.git|master|${src_dir}/tensorflow|tensorflow/compiler/mlir"
+    "https://github.com/tensorflow/tensorflow.git|master|${src_dir}/tensorflow|tensorflow/compiler/mlir tensorflow/core/ir"
     "https://github.com/tensorflow/runtime.git|master|${src_dir}/runtime|"
     "https://github.com/NVIDIA/TensorRT-Incubator.git|main|${src_dir}/TensorRT-Incubator|mlir-tensorrt"
     "https://github.com/iree-org/iree.git|main|${src_dir}/iree|"
@@ -42,13 +42,13 @@ sync() {
         if [ ! -d "${dir}" ]; then
             if [ -n "${root}" ]; then
                 git clone --quiet --depth=1 --branch "${branch}" --single-branch --filter=blob:none --sparse "${url}" "${dir}" >/dev/null 2>&1
-                git -C "${dir}" sparse-checkout set "${root}" >/dev/null 2>&1
+                git -C "${dir}" sparse-checkout set ${root} >/dev/null 2>&1
             else
                 git clone --quiet --depth=1 --branch "${branch}" --single-branch "${url}" "${dir}"
             fi
         else
             if [ -n "${root}" ]; then
-                git -C "${dir}" sparse-checkout set "${root}" >/dev/null 2>&1
+                git -C "${dir}" sparse-checkout set ${root} >/dev/null 2>&1
             fi
             git -C "${dir}" fetch --quiet --depth=1 origin "${branch}"
             git -C "${dir}" reset --quiet --hard FETCH_HEAD >/dev/null 2>&1

+ 2 - 0
tools/mlir-script.js

@@ -81,6 +81,7 @@ const schema = async () => {
         path.join(source, 'FlashTensor', 'include'),
         path.join(source, 'tpu-mlir', 'include'),
         path.join(source, 'tensorflow'),
+        path.join(source, 'tensorflow', 'tensorflow', 'core', 'ir'),
         path.join(source, 'tensorflow', 'tensorflow', 'compiler', 'mlir', 'tfrt', 'ir'),
         path.join(source, 'xla', 'xla', 'backends', 'gpu', 'codegen', 'triton', 'ir'),
         path.join(source, 'runtime', 'include'),
@@ -236,6 +237,7 @@ const schema = async () => {
         'torch-mlir/Dialect/Torch/IR/TorchOps.td',
         'torch-mlir/Dialect/TorchConversion/IR/TorchConversionOps.td',
         'torch-mlir-dialects/Dialect/TMTensor/IR/TMTensorOps.td',
+        'tensorflow/core/ir/ops.td',
         'tensorflow/compiler/mlir/lite/ir/tfl_ops.td',
         'tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td',
         'tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model_ops.td',

Some files were not shown because too many files changed in this diff