[mlir][Linalg] Refactor Linalg op initTensors support - NFC

[nicolasvasilache]

Summary

Manually-defined named ops do not currently support init_tensors or return values and may never support them. Add extra interface to the StructuredOpInterface so that we can still write op-agnostic transformations based on StructuredOpInterface.

This is an NFC extension in preparation for tiling on tensors.

[ftynse]

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

30

Typo: "and usage"

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

388

I think I already asked this question elsewhere, but don't remember the answer. Do you expect non-default implementation of this method? If not, I don't see the point of paying the cost of having it as an interface method instead of having it in extraClassDeclarations.

mlir/include/mlir/Dialect/Linalg/IR/LinalgTraits.h

41

I don't think you need ::Impl for non-parametric traits.

45–46

Nit: doubly public

[pifon2a]

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

67

will all (or most) these ops with ZeroInitTensors have always one output? If yes, then it might make sense to just have a base class PolyadicLinalgStructured_Op which is parameterized by the number of inputs and already sets number of outputs to 1 and number of inits to 0? Just a suggestion, no need to do anything.

[nicolasvasilache]

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

388

The extraClassDeclarations of what ?
How do you write portable transformations without having the method as part of the interface ?

[nicolasvasilache]

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

67

Ideally these would just go away and be replaced by full auto-generated ops.
For now this is allows implementing tiling on tensors, even though these ops are not concerned with it.

[nicolasvasilache]

Address review

[ftynse]

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

388

Of the interface - https://github.com/llvm/llvm-project/blob/master/mlir/include/mlir/IR/OpBase.td#L1859.

[nicolasvasilache]

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

388

ok, I'll do a standalone CL to move all methods that don't depend on $_op to extraClassDeclaration.

[nicolasvasilache]

Rebase.

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

 include "mlir/Dialect/Affine/IR/AffineOpsBase.td"
 include "mlir/Dialect/Linalg/IR/LinalgBase.td"
 include "mlir/Dialect/Linalg/IR/LinalgStructuredOpsInterface.td"
 include "mlir/Interfaces/CopyOpInterface.td"
 
 // The Linalg `NInputs` trait provides the API for ops that are known
 // to have a specified number of inputs, all passed as operands.
 // See Linalg/LinalgTraits.h for implementation details and usage.
-class NInputs<int args_in> :
-  NativeOpTrait<"linalg::NInputs<" # !cast<string>(args_in) # ">::Impl"> {}
+class NInputs<int n> :
+  NativeOpTrait<"linalg::NInputs<" # !cast<string>(n) # ">::Impl"> {}
+
+// The Linalg `ZeroInitTensors` trait provides the API for ops that are known
+// to not have input tensor operands.
+// See Linalg/LinalgTraits.h for implementation details and usage.
+def ZeroInitTensors : NativeOpTrait<"linalg::ZeroInitTensors"> {}
 
 // The Linalg `NOutputs` trait provides the API for ops that are known
 // to have a specified number of outputs, all passed as operands.
 // See Linalg/LinalgTraits.h for implementation details and usage.
-class NOutputs<int args_out> :
-  NativeOpTrait<"linalg::NOutputs<" # !cast<string>(args_out) # ">::Impl"> {}
+class NOutputs<int n> :
+  NativeOpTrait<"linalg::NOutputs<" # !cast<string>(n) # ">::Impl"> {}
 
 def StructuredOpTraits : NativeOpTrait<"linalg::StructuredOpTraits">;
 def NamedStructuredOpTrait : NativeOpTrait<"linalg::NamedStructuredOpTrait">;
 
 // Base Tablegen class for Linalg ops.
 // Linalg ops that correspond to library calls operate on linalg::View as their
 // first operands. These may be optionally followed by non-view operands
 // depending on the specific Linalg op.
class LinalgStructured_Op<string mnemonic, list<OpTrait> props>
   let assemblyFormat = "`(` operands `)` attr-dict `:` type(operands)";
 }
 
 //===----------------------------------------------------------------------===//
 // Named Linalg ops, implemented as special configurations of generic ops.
 //===----------------------------------------------------------------------===//
 // At the moment these are not declarative and require a bunch of C++ code.
 // In the future, these should be migrated to a declarative specification.
 def CopyOp : LinalgStructured_Op<"copy", [
     CopyOpInterface,
     NInputs<1>,
+    ZeroInitTensors,
     NOutputs<1>
   ]> {
   let description = [{
     Copies the data in the input view into the output view.
 
     Usage:
 
     ```mlir
let extraClassDeclaration = libraryCallName # [{
     Value getTarget() { return output(); }
   }];
   let verifier = [{ return ::verify(*this); }];
 
   let hasFolder = 1;
   let hasCanonicalizer = 1;
 }
 
-def FillOp : LinalgStructured_Op<"fill", [NInputs<0>, NOutputs<1>]> {
+def FillOp : LinalgStructured_Op<"fill", [
+    NInputs<0>,
+    ZeroInitTensors,
+    NOutputs<1>]> {
 
   let arguments = (ins AnyStridedMemRef:$output,
                    AnyTypeOf<[AnyFloat, AnySignlessInteger, AnyVector]>:$value);
   let extraClassDeclaration = libraryCallName # [{
     // Rank-polymorphic.
     //   filling_value -> O(ivs) with parallel iterators.
     ArrayAttr iterator_types() {
       unsigned nPar = getOutputShapedType(0).getRank();
code commonUtils = libraryCallName # [{
     int64_t getHighPad(unsigned i) {
       assert(i < getNumWindowLoops());
       if (!padding().hasValue()) return 0;
       return padding().getValue().getValue<int64_t>({i, 1});
     }
   }];
 }
 
-def ConvOp : PoolingBase_Op<"conv", [NInputs<2>, NOutputs<1>]> {
+def ConvOp : PoolingBase_Op<"conv", [
+    NInputs<2>,
+    // Despite having reductions, this manually defined ConvOp may only take
+    // memref operands and can never have init tensors.
+    ZeroInitTensors,
+    NOutputs<1>]> {
 
   let description = [{
     Generic n-D convolution as described in the TF documentation:
     https://www.tensorflow.org/versions/r2.0/api_docs/python/tf/nn/convolution
 
     ```
       output[b, x[0], ..., x[N-1], k] =
       sum_{z[0], ..., z[N-1], q}
def ConvOp : PoolingBase_Op<"conv", [
 
   let verifier = [{ return ::verify(*this); }];
 
   let hasFolder = 1;
   let hasCanonicalizer = 1;
 }
 
 class SingleInputPoolingBase_Op<string mnemonic>
-  : PoolingBase_Op<mnemonic, [NInputs<2>, NOutputs<1>]> {
+  : PoolingBase_Op<mnemonic, [
+    NInputs<2>,
+    // Despite having reductions, this manually defined ConvOp may only take
+    // memref operands and can never have init tensors.
+    ZeroInitTensors,
+    NOutputs<1>]> {
   let description = [{
     A base class for single input pooling function.
 
     TODO: Figure out a better way to handle window dimensions, i.e., eliminate
     the fake memref.
     The window dimensions are specified by argument `windowDims`. The i-th
     dimension in the shape of `windowDims` denotes the size of the window along
     dimension i. For example, if the window size is 2x3, then a memref<2x3>

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

InterfaceMethod<
       /*args=*/(ins),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
         auto iters = $_op.iterator_types();
         return iters.size() == 1 &&
                getNumIterators(getReductionIteratorTypeName(), iters) == 1;
       }]>,
     //===------------------------------------------------------------------===//
-    // Num input/output arguments handling.
+    // Num input/output/initTensors arguments handling.
     //===------------------------------------------------------------------===//
     // These special methods must be defined by each op that wants to implement
     // the LinalgStructuredInterface. For now, this is either:
-    // - inherited statically by using the NInputs<unsigned> or
-    //   NOutputs<unsigned> traits.
-    // - derived from args_in/args_out attributes (for linalg.generic and
+    // - Explicitly specified in the op definition.
+    // - Derived from variadic attributes (for "named" ops, linalg.generic and
     //   linalg.indexed_generic ops).
     InterfaceMethod<
       /*desc=*/[{
         Return the number of inputs.
       }],
       /*retTy=*/"unsigned",
       /*methodName=*/"getNumInputs"
     >,
     InterfaceMethod<
       /*desc=*/[{
+        Return the number of init tensors.
+      }],
+      /*retTy=*/"unsigned",
+      /*methodName=*/"getNumInitTensors"
+    >,
+    InterfaceMethod<
+      /*desc=*/[{
         Return the number of outputs.
       }],
       /*retTy=*/"unsigned",
       /*methodName=*/"getNumOutputs"
     >,
     //===------------------------------------------------------------------===//
     // Input arguments handling.
     //===------------------------------------------------------------------===//
InterfaceMethod<
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
         auto range = this->getOperation()->getOperands();
         return {range.begin(), range.begin() + getNumInputsAndOutputBuffers()};
       }]
     >,
     InterfaceMethod<
       /*desc=*/[{
+        Return the range over init tensors.
+      }],
+      /*retTy=*/"Operation::operand_range",
+      /*methodName=*/"getInitTensors",
+      /*args=*/(ins),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        auto range = this->getOperation()->getOperands();
+        return {range.begin() + getNumInputsAndOutputBuffers(),
+                range.begin() + getNumInputsAndOutputs()};
+      }]
+    >,
+    InterfaceMethod<
+      /*desc=*/[{
+        Return one single init tensor at position `$i`.
+      }],
+      /*retTy=*/"Value",
+      /*methodName=*/"getInitTensor",
+      /*args=*/(ins "unsigned":$i),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        assert(i < $_op.getNumInitTensors() && "overflowing init tensor index");
+        return getInitTensors()[i];
+      }]
+    >,
+    InterfaceMethod<
+      /*desc=*/[{
+        Return the range over inputs, output buffers and init tensors.
+      }],
+      /*retTy=*/"Operation::operand_range",
+      /*methodName=*/"getShapedOperands",
+      /*args=*/(ins),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        auto range = this->getOperation()->getOperands();
+        return {range.begin(), range.begin() + getNumInputsAndOutputs()};
+      }]
+    >,
+    InterfaceMethod<
+      /*desc=*/[{
         Return the `i`-th shaped type, there are 3 cases:
           1. if `i < $_op.getNumInputs()` then return `getInputShapedType(i)`;
              otherwise
           2. if `i < getNumInputsAndOutputBuffers()` then return the
              `getOutputBufferType(i - $_op.getNumInputs())`; otherwise
           3. return the `i - getNumInputsAndOutputBuffers()` result type.
       }],
       /*retTy=*/"ShapedType",
InterfaceMethod<
       /*desc=*/[{
         Return the indexing maps within the current operation.
       }],
       /*retTy=*/"SmallVector<AffineMap, 4>",
       /*methodName=*/"getIndexingMaps",
       /*args=*/(ins),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
-        return llvm::to_vector<4>($_op.indexing_maps().template getAsValueRange<AffineMapAttr>());
+        return llvm::to_vector<4>(
+          $_op.indexing_maps().template getAsValueRange<AffineMapAttr>());
       }]
     >,
     InterfaceMethod<
       /*desc=*/[{
         Return the input or output indexing map at index `i`.
       }],
       /*retTy=*/"AffineMap",
       /*methodName=*/"getIndexingMap",
let methods = [
     //===------------------------------------------------------------------===//
     StaticInterfaceMethod<
       /*desc=*/[{
         Create an operation of the current type with the given location,
         operands, and attributes.
       }],
       /*retTy=*/"Operation *",
       /*methodName=*/"create",
-      (ins "OpBuilder &":$builder, "Location":$loc,
+      (ins "OpBuilder &":$builder, "Location":$loc, "TypeRange":$resultTypes,
            "ValueRange":$operands,
            "ArrayRef<NamedAttribute>":$attributes), [{
-        return builder.create<ConcreteOp>(loc, TypeRange{}, operands,
-                                          attributes);
+        return builder.create<ConcreteOp>(
+          loc, resultTypes, operands, attributes);
       }]
     >,
     InterfaceMethod<
       /*desc=*/[{
         Clone the current operation with the given location and operands. This
         is used to abstract away the optional underlying region creation.
       }],
       /*retTy=*/"Operation *",
       /*methodName=*/"clone",
-      (ins "OpBuilder &":$b, "Location":$loc, "ValueRange":$operands), [{
+      (ins "OpBuilder &":$b, "Location":$loc, "TypeRange":$resultTypes,
+           "ValueRange":$operands),
+      [{
         BlockAndValueMapping map;
         unsigned numRegions = $_op.getOperation()->getNumRegions();
-        Operation *res = create(b, loc, operands, $_op.getAttrs());
+        Operation *res = create(b, loc, resultTypes, operands, $_op.getAttrs());
         assert(res->getNumRegions() == numRegions && "inconsistent # regions");
         for (unsigned ridx = 0; ridx < numRegions; ++ridx)
           $_op.getOperation()->getRegion(ridx).cloneInto(
             &res->getRegion(ridx), map);
         return res;
       }]
     >
   ];
 }
 
 #endif // LINALG_IR_STRUCTURED_OPS_INTERFACE

mlir/include/mlir/Dialect/Linalg/IR/LinalgTraits.h

 public:
   template <typename ConcreteType>
   class Impl : public OpTrait::TraitBase<ConcreteType, NInputs<N>::Impl> {
   public:
     static unsigned getNumInputs() { return N; }
   };
 };
 
+/// This class provides the API for ops that are known to not have init tensor
+/// operands. Use as a trait as follows:
+///
+///   class CopyOp : public Op<CopyOp, OpTrait::ZeroInitTensors> {
+///
+template <typename ConcreteType>
+class ZeroInitTensors : public TraitBase<ConcreteType, ZeroInitTensors> {
+public:
+  static unsigned getNumInitTensors() { return 0; }
+};
+
 /// This class provides the API for ops that are known to have a specified
 /// number of outputs, all passed as operands. Use as a trait as follows:
 ///
 ///   class DotOp : public Op<DotOp, OpTrait::NOutputs<2>::Impl> {
 ///
 template <unsigned N> class NOutputs {
 public:
   template <typename ConcreteType>

 ///
 template <typename ConcreteType>
 class NamedStructuredOpTrait
     : public OpTrait::TraitBase<ConcreteType, NamedStructuredOpTrait> {
 public:
   unsigned getNumInputs() {
     return cast<ConcreteType>(this->getOperation()).inputs().size();
   }
+  unsigned getNumInitTensors() {
+    return cast<ConcreteType>(this->getOperation()).init_tensors().size();
+  }
   unsigned getNumOutputs() {
     ConcreteType concreteOp = cast<ConcreteType>(this->getOperation());
     return concreteOp.output_buffers().size() +
            concreteOp.result_tensors().size();
   }
   static LogicalResult verifyTrait(Operation *op) {
     ConcreteType concreteOp = cast<ConcreteType>(op);
     unsigned nInputAndBufferOperands =

mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp

for (auto r : viewRanges) {
       strides.push_back(r.stride);
     }
     clonedViews.push_back(
         b.create<SubViewOp>(loc, view, offsets, sizes, strides));
   }
   auto operands = getAssumedNonViewOperands(op);
   clonedViews.append(operands.begin(), operands.end());
 
-  Operation *clonedOp = op.clone(b, loc, clonedViews);
+  Operation *clonedOp = op.clone(b, loc, /*resultTypes*/ {}, clonedViews);
   // When the producer is an IndexedGenercOp, we have to transform its block
   // IV arguments according to the tiling of the consumer, i.e. offset them by
   // the values computed in `loopRanges`.
   if (auto indexedGenericOp = dyn_cast<IndexedGenericOp>(clonedOp)) {
     auto &block = indexedGenericOp.region().front();
 
     OpBuilder::InsertionGuard g(b);
     b.setInsertionPointToStart(&block);

mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp

GenerateLoopNest<LoopTy>::doit(
         // of that one: (d0,d1,d2)->(d2,d0,d1)
         if (!options.interchangeVector.empty())
           ivValues = applyMapToValues(b, loc, invPermutationMap, ivValues);
 
         auto views = makeTiledViews(b, loc, op, viewSizesToLoopsMap, ivValues,
                                     tileSizes, allViewSizes);
         auto operands = getAssumedNonViewOperands(op);
         views.append(operands.begin(), operands.end());
-        res = op.clone(b, loc, views);
+        res = op.clone(b, loc, /*resultTypes*/ {}, views);
         return scf::ValueVector{};
       },
       options.distribution);
 
   // 4. Transforms index arguments of `linalg.generic` w.r.t. to the tiling.
   transformIndexedGenericOpIndices(b, res, ivs, loopIndexToRangeIndex);
 
   // 5. Gather the newly created loops and return them with the new op.