[mlir][linalg] Cleanup left over uses of deprecated LinalgOp methods.

[gysit]

Summary

Replace all remaining uses of deprecated Structured Op Interface methods. This patch is based on https://reviews.llvm.org/D103394.

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

let extraClassDeclaration = [{
     SmallVector<int64_t, 4> computeStaticLoopSizes();
 
     /// Returns all the operands past the inputs, output_buffers and
     /// init_tensors operands. Asserts that these operands are value types to
     /// allow transformations like tiling to just use the values when cloning
     /// `linalgOp`.
     Operation::operand_range getAssumedNonShapedOperands() {
       Operation::operand_range res{
-        getOperation()->getOperands().begin() + getNumShapedOperands(),
+        getOperation()->getOperands().begin() + getNumInputsAndOutputs(),
         getOperation()->getOperands().end()};
       for (Type t : TypeRange{res}) {
         (void)t;
         assert((t.isSignlessIntOrIndexOrFloat() || t.template isa<VectorType>())
                &&"expected scalar or vector type");
       }
       return res;
     }

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

def CopyOp : LinalgStructured_Op<"copy", [CopyOpInterface]> {
 
   let extraClassDeclaration = structuredOpsDecls # [{
     ValueRange inputs() { return getOperands().take_front(); }
     ValueRange outputs() { return getOperands().take_back(); }
 
     // Rank-polymorphic.
     //   filling_value -> O(ivs) with parallel iterators.
     ArrayAttr iterator_types() {
-      unsigned nPar = getInputShapedType(0).getRank();
+      int64_t nPar = getRank(getInputOperand(0));
       return Builder(getContext()).getStrArrayAttr(
         SmallVector<StringRef, 8>(nPar, getParallelIteratorTypeName()));
     }
 
     // I(input_perm(ivs)) -> O(output_perm(ivs))
     ArrayAttr indexing_maps() {
       MLIRContext *context = getContext();
       auto maybeInputMap = inputPermutation();
       auto maybeOutputMap = outputPermutation();
-      unsigned inputRank = getInputShapedType(0).getRank();
-      unsigned outputRank = getOutputShapedType(0).getRank();
+      int64_t inputRank = getRank(getInputOperand(0));
+      int64_t outputRank = getRank(getOutputOperand(0));
       return Builder(getContext()).getAffineMapArrayAttr({
           extractOrIdentityMap(maybeInputMap, inputRank, context),
           extractOrIdentityMap(maybeOutputMap, outputRank, context)});
     }
 
     Value getSource() { return input();}
     Value getTarget() { return output(); }
 
def FillOp : LinalgStructured_Op<"fill", []> {
   let regions = (region AnyRegion:$region);
   let extraClassDeclaration = structuredOpsDecls # [{
     ValueRange inputs() { return {}; }
     ValueRange outputs() { return getOperands().take_front(); }
 
     // Rank-polymorphic.
     //   filling_value -> O(ivs) with parallel iterators.
     ArrayAttr iterator_types() {
-      unsigned nPar = getOutputShapedType(0).getRank();
+      int64_t nPar = getRank(getOutputOperand(0));
       return Builder(getContext()).getStrArrayAttr(
         SmallVector<StringRef, 8>(nPar, getParallelIteratorTypeName()));
     }
 
     ArrayAttr indexing_maps() {
       MLIRContext *context = getContext();
       // filling_value -> O(ivs)
       return Builder(getContext()).getAffineMapArrayAttr({
let extraClassDeclaration = commonUtils # [{
     // too.
     unsigned getNumBatchDimensions() { return 1; }
 
     unsigned getNumInputFeatureDimensions() { return 1; }
 
     unsigned getNumOutputFeatureDimensions() { return 1; }
 
     unsigned getNumSpatialDimensions() {
-      return getOutputShapedType(0).getRank() - getNumBatchDimensions() -
+      return getRank(getOutputOperand(0)) - getNumBatchDimensions() -
              getNumOutputFeatureDimensions();
     }
 
     ArrayAttr iterator_types() {
       // Outer parallel loops are always the number of output dimensions; i.e.
       // [b, xs, q] in the TF notation above.
-      unsigned nPar = getOutputShapedType(0).getRank();
+      int64_t nPar = getRank(getOutputOperand(0));
       unsigned nRed = getNumInputFeatureDimensions();
       // Window loops are a special kind of reduction that is never tiled or
       // parallelized across; i.e. [zs] in the TF notation above whose number
       // match `xs` (i.e. 1 window loop per "image" dimension).
       // This may evolve in the future.
       // Conditionally check nPar is large enough for cases of ill-formed op:
       // this avoids overflows before hitting the verifier.
       assert(nPar > getNumBatchDimensions() + getNumInputFeatureDimensions() &&
let arguments = (ins AnyStridedMemRef:$input,
                    OptionalAttr<I64ElementsAttr>:$padding);
 
   let extraClassDeclaration = commonUtils# [{
     ValueRange inputs() { return getOperands().slice(0, 2); }
     ValueRange outputs() { return getOperands().take_back(); }
 
     ArrayAttr iterator_types() {
       // Outer parallel loops are always the number of output dimensions.
-      unsigned nPar = getOutputShapedType(0).getRank();
+      int64_t nPar = getRank(getOutputOperand(0));
       // The window loops has the same number loops with output dimensions.
       unsigned nWin = nPar;
       SmallVector<StringRef, 8> iters(nPar, getParallelIteratorTypeName());
       iters.reserve(nPar + nWin);
       iters.append(nWin, getWindowIteratorTypeName());
       return Builder(getContext()).getStrArrayAttr(iters);
     }
 

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

for (auto expr : map.getResults()) {
     res.push_back(createFoldedComposedAffineApply(b, loc, map, values));
   }
   return res;
 }
 
 SmallVector<Value, 4> LinalgOp::createFlatListOfOperandDims(OpBuilder &b,
                                                             Location loc) {
   SmallVector<Value, 4> res;
-  for (Value v : getShapedOperands()) {
-    ShapedType t = v.getType().template cast<ShapedType>();
-    for (unsigned i = 0, e = t.getRank(); i < e; ++i)
-      res.push_back(b.createOrFold<memref::DimOp>(loc, v, i));
+  for (OpOperand *opOperand : getInputAndOutputOperands()) {
+    for (int64_t i = 0, e = getRank(opOperand); i < e; ++i)
+      res.push_back(b.createOrFold<memref::DimOp>(loc, opOperand->get(), i));
   }
   return res;
 }
 
 SmallVector<int64_t, 4> LinalgOp::createFlatListOfOperandStaticDims() {
   SmallVector<int64_t, 4> res;
-  for (Value v : getShapedOperands()) {
-    ShapedType t = v.getType().template cast<ShapedType>();
-    assert(t.hasStaticShape() && "expected operands to have static shapes");
-    llvm::append_range(res, t.getShape());
-  }
+  assert(!hasDynamicShape() && "expected operands to have static shapes");
+  for (OpOperand *opOperand : getInputAndOutputOperands())
+    llvm::append_range(res, getShape(opOperand));
   return res;
 }
 
 SmallVector<Range, 4> LinalgOp::createLoopRanges(OpBuilder &b, Location loc) {
   AffineMap map = getLoopsToShapesMap();
   unsigned numDims = map.getNumDims(), numRes = map.getNumResults();
   auto viewSizes = createFlatListOfOperandDims(b, loc);
   SmallVector<Range, 4> res(numDims);
LogicalResult LinalgOp::reifyReturnTypeShapesPerResultDim(
   llvm::SmallSet<unsigned, 4> outputDims;
   llvm::for_each(resultPosRange,
                  [&outputDims](unsigned dim) { outputDims.insert(dim); });
   HasAffineDimExprVisitor checkDimExpr(outputDims);
   Location loc = getOperation()->getLoc();
   auto allResultDimValues =
       applyMapToValues(b, loc, resultShapesFromInputShapesMap,
                        createFlatListOfOperandDims(b, loc));
-  unsigned pos = 0;
+  int64_t pos = 0;
   ArrayRef<AffineExpr> shapeExprs = resultShapesFromInputShapesMap.getResults();
-  for (auto resultIdx : llvm::seq<unsigned>(0, getNumOutputs())) {
-    ShapedType resultType = getOutputShapedType(resultIdx);
+  for (OpOperand *opOperand : getOutputOperands()) {
     SmallVector<Value> shapes;
-    for (unsigned dim : llvm::seq<unsigned>(0, resultType.getRank())) {
+    for (int64_t dim : llvm::seq<int64_t>(0, getRank(opOperand))) {
       if (checkDimExpr.visit(shapeExprs[pos]))
         shapes.push_back(
-            b.createOrFold<memref::DimOp>(loc, getOutput(resultIdx), dim));
+            b.createOrFold<memref::DimOp>(loc, opOperand->get(), dim));
       else
         shapes.push_back(allResultDimValues[pos]);
       pos++;
     }
     reifiedReturnShapes.emplace_back(std::move(shapes));
   }
   return success();
 }

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

for (Value value : outputs) {
     effects.emplace_back(MemoryEffects::Write::get(), value,
                          SideEffects::DefaultResource::get());
   }
 }
 
 void GenericOp::getEffects(
     SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
         &effects) {
-  getGenericEffectsImpl(effects, getOperation()->getResults(),
-                        getInputBuffers(), getOutputBuffers());
+  SmallVector<Value> inputBuffers = getInputBufferOperands();
+  SmallVector<Value> outputBuffers = getOutputBufferOperands();
+  getGenericEffectsImpl(effects, getOperation()->getResults(), inputBuffers,
+                        outputBuffers);
 }
 
 void IndexedGenericOp::getEffects(
     SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
         &effects) {
-  getGenericEffectsImpl(effects, getOperation()->getResults(),
-                        getInputBuffers(), getOutputBuffers());
+  SmallVector<Value> inputBuffers = getInputBufferOperands();
+  SmallVector<Value> outputBuffers = getOutputBufferOperands();
+  getGenericEffectsImpl(effects, getOperation()->getResults(), inputBuffers,
+                        outputBuffers);
 }
 
 template <typename GenericOpType>
 static LogicalResult verifyGenericOp(GenericOpType op) {
   return success();
 }
 
 static LogicalResult verify(GenericOp op) { return verifyGenericOp(op); }

mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-gen.cpp

void {0}::getCanonicalizationPatterns(
       results.add<FoldTensorCastOp>(context);
     }
     LogicalResult {0}::fold(ArrayRef<Attribute>,
                             SmallVectorImpl<OpFoldResult> &) {{
       return foldMemRefCast(*this);
     }
     void {0}::getEffects(SmallVectorImpl<
         SideEffects::EffectInstance<MemoryEffects::Effect> >&effects) {{
+      SmallVector<Value> inputBuffers = getInputBufferOperands();
+      SmallVector<Value> outputBuffers = getOutputBufferOperands();
       getGenericEffectsImpl(effects,
-        getOperation()->getResults(), getInputBuffers(), getOutputBuffers());
+        getOperation()->getResults(), inputBuffers, outputBuffers);
     })FMT";
   os << llvm::formatv(canonicalizersAndFoldersFmt, cppOpName);
 }
 
 // Prints methods for querying whether the current named op has attributes that
 // are used by its indexing maps and for verifying those attributes have the
 // expected type.
 void TCParser::printIndexingMapRequiredAttrMethods(

mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp

void {0}::getCanonicalizationPatterns(
   results.add<FoldTensorCastOp>(context);
 }
 LogicalResult {0}::fold(ArrayRef<Attribute>,
                         SmallVectorImpl<OpFoldResult> &) {{
   return foldMemRefCast(*this);
 }
 void {0}::getEffects(SmallVectorImpl<
     SideEffects::EffectInstance<MemoryEffects::Effect> >&effects) {{
+      SmallVector<Value> inputBuffers = getInputBufferOperands();
+      SmallVector<Value> outputBuffers = getOutputBufferOperands();
       getGenericEffectsImpl(effects,
-    getOperation()->getResults(), getInputBuffers(), getOutputBuffers());
+        getOperation()->getResults(), inputBuffers, outputBuffers);
 }
 )FMT";
 
 static LogicalResult generateNamedGenericOpOds(LinalgOpConfig &opConfig,
                                                GenerationContext &genContext) {
   if (!genContext.shouldGenerateOds())
     return success();