[mlir-tblgen] Emit named operand indices

[Meinersbur]

An operation's operands are defined by the arguments field in the tablegen definition. mlir-tblgen generates accessors for them: getXYZ() and setXYZ(...) to set an operation's operands without knowing the operand's index, but it does not expose the operand index itself. Yet some use cases requires knowing the operand index that is now covered by just getters and setters. For instance:

• Given an mlir::OpOperand, find out whether it is a specific argument (from the arguments field in the .td file)
• For operation with variable number of operands (variadic, AttrSizedOperandSegments), get the value to pass to getODSOperands or getODSOperandIndexAndLength.

Emitting a named constant avoids hardcoding the index in the source as magic constant, such as 2 in

llvm-project/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp

Line 244 in d56c06e

getODSOperands(2).getTypes().front(), /*isAccumulator=*/true);

Extracted out of #144785

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-core

Author: Michael Kruse (Meinersbur)

Changes

An operation's operands are defined by the arguments field in the tablegen definition. mlir-tblgen generates accessors for them: getXYZ() and setXYZ(...) to set an operation's operands without knowing the operand's index, but it does not expose the operand index itself. Yet some use cases requires knowing the operand index that is now covered by just getters and setters. For instance:

• Given an mlir::OpOperand, find out whether it is a specific argument (from the arguments field in the .td file)
• For operation with variable number of operands (variadic, AttrSizedOperandSegments), get the value to pass to getODSOperands or getODSOperandIndexAndLength.

Emitting a named constant avoids hardcoding the index in the source as magic constant, such as 2 in

llvm-project/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp

Line 244 in d56c06e

getODSOperands(2).getTypes().front(), /*isAccumulator=*/true);

Extracted out of #144785

---

Full diff: https://github.com/llvm/llvm-project/pull/146839.diff

2 Files Affected:

• (modified) mlir/test/mlir-tblgen/op-operand.td (+11)
• (modified) mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp (+11)

diff --git a/mlir/test/mlir-tblgen/op-operand.td b/mlir/test/mlir-tblgen/op-operand.td
index a2fa1f7046a97..ab8d721ed5427 100644
--- a/mlir/test/mlir-tblgen/op-operand.td
+++ b/mlir/test/mlir-tblgen/op-operand.td
@@ -13,6 +13,9 @@ def OpA : NS_Op<"one_normal_operand_op", []> {
   let arguments = (ins I32:$input);
 }
 
+// DECL-LABEL: class OpA : {{.*}} {
+// DECL: static constexpr int odsIndex_input = 0;
+
 // CHECK-LABEL: OpA definitions
 
 // CHECK:      void OpA::build
@@ -28,6 +31,9 @@ def OpB : NS_Op<"one_variadic_operand_op", []> {
   let arguments = (ins Variadic<I32>:$input);
 }
 
+// DECL-LABEL: class OpB : {{.*}} {
+// DECL: static constexpr int odsIndex_input = 0;
+
 // CHECK-LABEL: OpB::build
 // CHECK:         ::mlir::ValueRange input
 // CHECK-NOT:     assert
@@ -37,6 +43,11 @@ def OpD : NS_Op<"mix_variadic_and_normal_inputs_op", [SameVariadicOperandSize]>
   let arguments = (ins Variadic<AnyTensor>:$input1, AnyTensor:$input2, Variadic<AnyTensor>:$input3);
 }
 
+// DECL-LABEL: class OpD : {{.*}} {
+// DECL: static constexpr int odsIndex_input1 = 0;
+// DECL: static constexpr int odsIndex_input2 = 1;
+// DECL: static constexpr int odsIndex_input3 = 2;
+
 // DECL-LABEL: ::mlir::Operation::operand_range getInput1
 // DECL-NEXT: return getODSOperands(0);
 
diff --git a/mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp b/mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp
index 6008ed4673d1b..cbb4030f3adb4 100644
--- a/mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp
+++ b/mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp
@@ -2223,6 +2223,17 @@ generateNamedOperandGetters(const Operator &op, Class &opClass,
                     "'SameVariadicOperandSize' traits");
   }
 
+  // Print the ods names so they don't need to be hardcoded in the source.
+  for (int i = 0; i != numOperands; ++i) {
+    const auto &operand = op.getOperand(i);
+    if (operand.name.empty())
+      continue;
+
+    opClass.declare<Field>("static constexpr int", Twine("odsIndex_") +
+                                                       operand.name + " = " +
+                                                       Twine(i));
+  }
+
   // First emit a few "sink" getter methods upon which we layer all nicer named
   // getter methods.
   // If generating for an adaptor, the method is put into the non-templated

[joker-eph]

Given an mlir::OpOperand, find out whether it is a specific argument (from the arguments field in the .td file)

The method to do this is non-trivial because of variadics. It is possible to emit a method computing this from TableGen (the opposite of getODSOperands() ?), but I'm not sure how this patch solves the problem right now?

[Meinersbur]

It is possible to emit a method computing this from TableGen (the opposite of getODSOperands() ?

Is this a question or a statement? Even if we had an oppsite of getODSOperands() (Input: OpOperand, output: an index of the arguments list in the .td file), it one would still need compare that return index to an hardcoded magic number, instead of by the name of that argument.

but I'm not sure how this patch solves the problem right now?

It is done in #144785 which uses it for this pupose. In principle: Check whether a given operand index falls into the range of getODSOperandIndexAndLength(argnum), where argnum is one of the named constants such as odsIndex_input. Equivalently, but less efficient, llvm::contains(opoperand, getODSOperand(OpB::odsIndex_input)).

[joker-eph]

Is this a question or a statement?

This was a statement (such method is possible I believe), with a question in the parenthesis (whether it would be the opposite of ...).

It is done in #144785 which uses it for this pupose. In principle: Check whether a given operand index falls into the range of getODSOperandIndexAndLength(argnum), where argnum is one of the named constants such as odsIndex_input. Equivalently, but less efficient, llvm::contains(opoperand, getODSOperand(OpB::odsIndex_input)).

Gotcha, thanks for explaining!

Anyway I agree that being able to use getODSOperandIndexAndLength() without a magic constant is a nice improvement.

[Meinersbur]

Thanks for the review