[mlir] Modify LinalgStructuredInterface to allow the computation block to have arguments only for a subset of operands.

[olegshyshkov]

Summary

Currently there is an expectations that there is a corresponsing block argument
for each operand. For some operation, it leads to unused arguments. For example,
in map, only input operands are used for the computation.

[pifon2a]

Thanks!

mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td

581

What about calling it getTiedOpOperands? Because it returns the list of getTiedOpOperand(i) for every block arg.

[mravishankar]

I think this is a big change. Having unused arguments is not a big deal. It removes some consistency that would be very involved to work around.

[nicolasvasilache]

I think this is a big change. Having unused arguments is not a big deal. It removes some consistency that would be very involved to work around.

I've actually been annoyed by the trailing unused bbarg hanging around for some time ..
In the linalg.generic case I agree that consistency trumps the rest.
It is innocuous to me to relax the interface so we can create ops with nicer ergonomics.

If anything, I think this even increases consistency because the bbArgs region verifier directly relates to the matching block operands without an implicit level of indirection.

mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td

581

The "tied" terminology would be a misnomer here I think as it relates to DPS

See e.g. https://stackoverflow.com/questions/57551135/how-does-llvm-translate-three-address-llvm-ir-add-into-x86-two-address-add ()

"Tied" means two operands are mapped to the same register. When we write Constraints = "$src1 = $dst" in tablegen, generated code would set $src1 and $dst as "tied" (GlobalISel/Utils.cpp):

I introduced it improperly a while back but would appreciate the cleanup.
e.g. getTiedOpOperand(bbArg) -> getMatchingOpOperand(bbArg) and other places (getTiedIndexingMap -> getMatchingIndexingMap etc)

Then this one could be called getOpOperandsMatchingBBargs ?

[olegshyshkov]

Renamed the method to getOpOperandsMatchingBBargs.

[olegshyshkov]

mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td

581

getOpOperandsMatchingBBargs sounds good. I was hesitant to add such a verbose name at first.

[mravishankar]

removing my blocker based on @nicolasvasilache comments.

[olegshyshkov]

mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td

581

I'll prepare a separate patch to cleanup getTiedOpOperand(bbArg) -> getMatchingOpOperand(bbArg) tomorrow.

[nicolasvasilache]

removing my blocker based on @nicolasvasilache comments.

Thanks for your flexibility @mravishankar!

I'll prepare a separate patch to cleanup getTiedOpOperand(bbArg) -> getMatchingOpOperand(bbArg) tomorrow.

Great, thank you!

mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td

InterfaceMethod<
          Return whether the op accesses the iteration indices.
       }],
       /*retTy=*/"bool",
       /*methodName=*/"hasIndexSemantics",
       /*args=*/(ins),
       /*methodBody=*/"",
       /*defaultImplementation=*/""
     >,
+    InterfaceMethod<
+      /*desc=*/[{
+        Return op operands that have a corresponding argument in the basic block.
+        By default, the block should have an argument for each operand, but there
+        are expection. For example, in `map` output operand isn't used in
+        the block.
+      }],
+      /*retTy=*/"OpOperandVector",
+      /*methodName=*/"getOpOperandsMatchingBBargs",
+      /*args=*/(ins),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        return $_op.getInputAndOutputOperands();
+      }]
+    >,
     //===------------------------------------------------------------------===//
     // Linalg generalization hooks.
     //===------------------------------------------------------------------===//
     InterfaceMethod<
       /*desc=*/[{
         Hook to provide a custom AffineMap used to compute all the operand
         subshapes given loop bounds. This is used to answer the question: "given
         an iteration space over the codomain, what are the subshapes of the

mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp

LogicalResult mlir::linalg::detail::verifyStructuredOpInterface(Operation *op) {
   // Simplifying assumption: bbargs match 1-1 with shape operands elemental
   // types.
   // TODO: once ranked shape types are plugged in, we may want to drop the
   // corresponding bbargs, that can never be read from. This will be subject to
   // consistency discussions (i.e. what to do with output tensors whose bbarg is
   // not used).
   Block &block = linalgOp->getRegion(0).front();
 
-  if (linalgOp.getNumInputsAndOutputs() != block.getNumArguments())
+  if (linalgOp.getOpOperandsMatchingBBargs().size() != block.getNumArguments())
     return op->emitOpError("expected as many non-induction variable region "
                            "arguments as the number of input/output operands");
 
-  for (OpOperand *opOperand : linalgOp.getInputAndOutputOperands()) {
+  for (OpOperand *opOperand : linalgOp.getOpOperandsMatchingBBargs()) {
     Type elementType = getElementTypeOrSelf(opOperand->get());
     Type argType = block.getArgument(opOperand->getOperandNumber()).getType();
     if (elementType != argType)
       return op->emitOpError("expected type of bb argument #")
              << opOperand->getOperandNumber() << " (" << argType << ")"
              << " to match element or self type of the corresponding operand ("
              << elementType << ")";
   }