Handle optional node arguments more gracefully

Optional arguments to nodes were treated exactly the same as required ones. This makes op_specification and the nodes_header.tmpl treat them specially and create <operand>_operand_number functions to get the actual operand number instead of having (potentially incorrect) kOperandIndex constexprs. PiperOrigin-RevId: 577359837
kammoh · Oct 28, 2023 · 523e5a9 · 523e5a9
1 parent 24c0e73
commit 523e5a9
Show file tree

Hide file tree

Showing 5 changed files with 220 additions and 85 deletions.
diff --git a/xls/ir/BUILD b/xls/ir/BUILD
@@ -1019,8 +1019,9 @@ build_test(
 
 genrule(
     name = "op_header",
+    srcs = ["op_header.tmpl"],
     outs = ["op.h"],
-    cmd = "$(location :render_specification_against_template) ./xls/ir/op_header.tmpl" +
+    cmd = "$(location :render_specification_against_template) $(location :op_header.tmpl)" +
           " | $(location @llvm_toolchain_llvm//:bin/clang-format)" +
           " > $(OUTS)",
     tools = [
@@ -1031,8 +1032,9 @@ genrule(
 
 genrule(
     name = "gen_op_proto",
+    srcs = ["op_proto.tmpl"],
     outs = ["op.generated.proto"],
-    cmd = "$(location :render_specification_against_template) ./xls/ir/op_proto.tmpl > $(OUTS)",
+    cmd = "$(location :render_specification_against_template) $(location :op_proto.tmpl) > $(OUTS)",
     tools = [
         ":render_specification_against_template",
     ],
@@ -1055,8 +1057,9 @@ sh_test(
 
 genrule(
     name = "op_source",
+    srcs = ["op_source.tmpl"],
     outs = ["op.cc"],
-    cmd = "$(location :render_specification_against_template) ./xls/ir/op_source.tmpl" +
+    cmd = "$(location :render_specification_against_template) $(location :op_source.tmpl)" +
           " | $(location @llvm_toolchain_llvm//:bin/clang-format)" +
           " > $(OUTS)",
     tools = [
@@ -1067,8 +1070,9 @@ genrule(
 
 genrule(
     name = "nodes_header",
+    srcs = ["nodes_header.tmpl"],
     outs = ["nodes.h"],
-    cmd = "$(location :render_specification_against_template) ./xls/ir/nodes_header.tmpl" +
+    cmd = "$(location :render_specification_against_template) $(location :nodes_header.tmpl)" +
           " | $(location @llvm_toolchain_llvm//:bin/clang-format)" +
           " > $(OUTS)",
     tools = [
@@ -1079,8 +1083,9 @@ genrule(
 
 genrule(
     name = "nodes_source",
+    srcs = ["nodes_source.tmpl"],
     outs = ["nodes.cc"],
-    cmd = "$(location :render_specification_against_template) ./xls/ir/nodes_source.tmpl" +
+    cmd = "$(location :render_specification_against_template) $(location :nodes_source.tmpl)" +
           " | $(location @llvm_toolchain_llvm//:bin/clang-format)" +
           " > $(OUTS)",
     tools = [

diff --git a/xls/ir/nodes_header.tmpl b/xls/ir/nodes_header.tmpl
@@ -28,8 +28,8 @@ struct SliceData {
 {% for op_class in spec.OpClass.kinds.values() -%}
 class {{ op_class.name }} final : public Node {
  public:
-{% for op in op_class.operands -%}
-  static constexpr int64_t k{{op.camel_case_name()}}Operand = {{loop.index0}};
+{% for op_set in op_class.fixed_operands() -%}
+  static constexpr int64_t k{{op_set.operand.camel_case_name()}}Operand = {{op_set.index}};
 {% endfor %}
 
   {{ op_class.name }}({{ op_class.constructor_args_str() }});
@@ -42,7 +42,28 @@ class {{ op_class.name }} final : public Node {
   {% if method.expression %}  {% if method.expression_is_body %} { {{ method.expression }} } {% else %} { return {{ method.expression }}; }  {% endif %}
   {% else %}; {% endif %}
 {% endfor -%}
-{%- if op_class.data_members() %}
+{% for op_set in op_class.optional_operands() -%}
+{% if not op_set.operand.manual_optional_implementation %}
+  absl::StatusOr<int64_t> {{op_set.operand.name}}_operand_number() const {
+    if (!has_{{op_set.operand.name}}_) {
+      return absl::InternalError("{{op_set.operand.name}} is not present");
+    }
+    int64_t ret = {{op_set.index}};
+
+    {% for other_optional in op_class.optional_operands()  %}
+      {% if other_optional.index < op_set.index %}
+    if (!has_{{other_optional.operand.name}}_) {
+      ret--;
+    }
+      {% endif %}
+    {% endfor %}
+    return ret;
+  }
+
+{% endif %}
+{% endfor -%}
+
+{%- if op_class.has_data_members() %}
   bool IsDefinitelyEqualTo(const Node* other) const final;
 
  private:

diff --git a/xls/ir/nodes_source.tmpl b/xls/ir/nodes_source.tmpl
@@ -95,7 +95,7 @@ absl::StatusOr<Node*> {{ op_class.name }}::CloneInNewFunction(
 }
 {% endif %}
 
-{% if op_class.data_members() %}
+{% if op_class.has_data_members() %}
 bool {{ op_class.name }}::IsDefinitelyEqualTo(const Node* other) const {
   if (this == other) {
     return true;