Commit 05d5125d authored by Nicolas Vasilache's avatar Nicolas Vasilache
Browse files

[mlir] Generalize OpFoldResult usage in ops with offsets, sizes and operands. 

This revision starts evolving the APIs to manipulate ops with offsets, sizes and operands towards a ValueOrAttr abstraction that is already used in folding under the name OpFoldResult.

The objective, in the future, is to allow such manipulations all the way to the level of ODS to avoid all the genuflexions involved in distinguishing between values and attributes for generic constant foldings.

Once this evolution is accepted, the next step will be a mechanical OpFoldResult -> ValueOrAttr.

Differential Revision: https://reviews.llvm.org/D95310
parent 7163aa99

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

+45 −47
Original line number Diff line number Diff line
@@ -1959,14 +1959,19 @@ def MemRefReinterpretCastOp: 
  let builders = [

    // Build a ReinterpretCastOp with mixed static and dynamic entries.

    OpBuilderDAG<(ins "MemRefType":$resultType, "Value":$source,

      "int64_t":$staticOffset, "ArrayRef<int64_t>":$staticSizes,

      "ArrayRef<int64_t>":$staticStrides, "ValueRange":$offset,

      "ValueRange":$sizes, "ValueRange":$strides,

      "OpFoldResult":$offset, "ArrayRef<OpFoldResult>":$sizes,

      "ArrayRef<OpFoldResult>":$strides,

      CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>,

    // Build a ReinterpretCastOp with all dynamic entries.

    // Build a ReinterpretCastOp with static entries.

    OpBuilderDAG<(ins "MemRefType":$resultType, "Value":$source,

      "Value":$offset, "ValueRange":$sizes, "ValueRange":$strides,

      "int64_t":$offset, "ArrayRef<int64_t>":$sizes,

      "ArrayRef<int64_t>":$strides,

      CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>,

    // Build a ReinterpretCastOp with dynamic entries.

    OpBuilderDAG<(ins "MemRefType":$resultType, "Value":$source,

      "Value":$offset, "ValueRange":$sizes,

      "ValueRange":$strides,

      CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>

  ];



  let extraClassDeclaration = extraBaseClassDeclaration # [{
@@ -2927,23 +2932,33 @@ def SubViewOp : BaseOpWithOffsetSizesAndStrides< 
  let results = (outs AnyMemRef:$result);



  let builders = [

    // Build a SubViewOp with mixed static and dynamic entries.

    OpBuilderDAG<(ins "Value":$source, "ArrayRef<int64_t>":$staticOffsets,

      "ArrayRef<int64_t>":$staticSizes, "ArrayRef<int64_t>":$staticStrides,

      "ValueRange":$offsets, "ValueRange":$sizes, "ValueRange":$strides,

    // Build a SubViewOp with mixed static and dynamic entries and custom

    // result type. If the type passed is nullptr, it is inferred.

    OpBuilderDAG<(ins "Value":$source, "ArrayRef<OpFoldResult>":$offsets,

      "ArrayRef<OpFoldResult>":$sizes, "ArrayRef<OpFoldResult>":$strides,

      CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>,

    // Build a SubViewOp with all dynamic entries.

    OpBuilderDAG<(ins "Value":$source, "ValueRange":$offsets,

      "ValueRange":$sizes, "ValueRange":$strides,

    // Build a SubViewOp with mixed static and dynamic entries and inferred

    // result type.

    OpBuilderDAG<(ins "MemRefType":$resultType, "Value":$source,

      "ArrayRef<OpFoldResult>":$offsets, "ArrayRef<OpFoldResult>":$sizes,

      "ArrayRef<OpFoldResult>":$strides,

      CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>,

    // Build a SubViewOp with mixed static and dynamic entries

    // and custom result type.

    // Build a SubViewOp with static entries and custom result type. If the

    // type passed is nullptr, it is inferred.

    OpBuilderDAG<(ins "Value":$source, "ArrayRef<int64_t>":$offsets,

      "ArrayRef<int64_t>":$sizes, "ArrayRef<int64_t>":$strides,

      CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>,

    // Build a SubViewOp with static entries and inferred result type.

    OpBuilderDAG<(ins "MemRefType":$resultType, "Value":$source,

      "ArrayRef<int64_t>":$staticOffsets, "ArrayRef<int64_t>":$staticSizes,

      "ArrayRef<int64_t>":$staticStrides, "ValueRange":$offsets,

      "ArrayRef<int64_t>":$offsets, "ArrayRef<int64_t>":$sizes,

      "ArrayRef<int64_t>":$strides,

      CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>,

    // Build a SubViewOp with dynamic entries and custom result type. If the

    // type passed is nullptr, it is inferred.

    OpBuilderDAG<(ins "Value":$source, "ValueRange":$offsets,

      "ValueRange":$sizes, "ValueRange":$strides,

      CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>,

    // Build a SubViewOp with all dynamic entries and custom result type.

    // Build a SubViewOp with dynamic entries and inferred result type.

    OpBuilderDAG<(ins "MemRefType":$resultType, "Value":$source,

      "ValueRange":$offsets, "ValueRange":$sizes, "ValueRange":$strides,

      CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>
@@ -3039,26 +3054,6 @@ def SubTensorOp : BaseOpWithOffsetSizesAndStrides< 
  let results = (outs AnyRankedTensor:$result);



  let builders = [

    // Build a SubTensorOp with mixed static and dynamic entries.

    OpBuilderDAG<(ins "Value":$source, "ArrayRef<int64_t>":$staticOffsets,

      "ArrayRef<int64_t>":$staticSizes, "ArrayRef<int64_t>":$staticStrides,

      "ValueRange":$offsets, "ValueRange":$sizes, "ValueRange":$strides,

      CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>,

    // Build a SubTensorOp with all dynamic entries.

    OpBuilderDAG<(ins "Value":$source, "ValueRange":$offsets,

      "ValueRange":$sizes, "ValueRange":$strides,

      CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>,

    // Build a SubTensorOp with mixed static and dynamic entries

    // and custom result type.

    OpBuilderDAG<(ins "RankedTensorType":$resultType, "Value":$source,

      "ArrayRef<int64_t>":$staticOffsets, "ArrayRef<int64_t>":$staticSizes,

      "ArrayRef<int64_t>":$staticStrides, "ValueRange":$offsets,

      "ValueRange":$sizes, "ValueRange":$strides,

      CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>,

    // Build a SubTensorOp with all dynamic entries and custom result type.

    OpBuilderDAG<(ins "RankedTensorType":$resultType, "Value":$source,

      "ValueRange":$offsets, "ValueRange":$sizes, "ValueRange":$strides,

      CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>,

    // Build a SubTensorOp with mixed static and dynamic entries and inferred

    // result type.

    OpBuilderDAG<(ins "Value":$source, "ArrayRef<OpFoldResult>":$offsets,
@@ -3069,6 +3064,15 @@ def SubTensorOp : BaseOpWithOffsetSizesAndStrides< 
    OpBuilderDAG<(ins "RankedTensorType":$resultType, "Value":$source,

      "ArrayRef<OpFoldResult>":$offsets, "ArrayRef<OpFoldResult>":$sizes,

      "ArrayRef<OpFoldResult>":$strides,

      CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>,

    // Build a SubTensorOp with dynamic entries and custom result type. If the

    // type passed is nullptr, it is inferred.

    OpBuilderDAG<(ins "Value":$source, "ValueRange":$offsets,

      "ValueRange":$sizes, "ValueRange":$strides,

      CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>,

    // Build a SubTensorOp with dynamic entries and inferred result type.

    OpBuilderDAG<(ins "RankedTensorType":$resultType, "Value":$source,

      "ValueRange":$offsets, "ValueRange":$sizes, "ValueRange":$strides,

      CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>

  ];
@@ -3157,19 +3161,13 @@ def SubTensorInsertOp : BaseOpWithOffsetSizesAndStrides< 


  let builders = [

    // Build a SubTensorInsertOp with mixed static and dynamic entries.

    OpBuilderDAG<(ins "Value":$source, "Value":$dest,

      "ArrayRef<int64_t>":$staticOffsets, "ArrayRef<int64_t>":$staticSizes,

      "ArrayRef<int64_t>":$staticStrides, "ValueRange":$offsets,

      "ValueRange":$sizes, "ValueRange":$strides,

      CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>,

    // Build a SubTensorInsertOp with all dynamic entries.

    OpBuilderDAG<(ins "Value":$source, "Value":$dest, "ValueRange":$offsets,

      "ValueRange":$sizes, "ValueRange":$strides,

      CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>,

    // Build a SubTensorInsertOp with mixed static and dynamic entries.

    OpBuilderDAG<(ins "Value":$source, "Value":$dest,

      "ArrayRef<OpFoldResult>":$offsets, "ArrayRef<OpFoldResult>":$sizes,

      "ArrayRef<OpFoldResult>":$strides,

      CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>,

    // Build a SubTensorInsertOp with dynamic entries.

    OpBuilderDAG<(ins "Value":$source, "Value":$dest,

      "ValueRange":$offsets, "ValueRange":$sizes, "ValueRange":$strides,

      CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>

  ];

mlir/include/mlir/IR/OpDefinition.h

+14 −0
Original line number Diff line number Diff line
@@ -213,11 +213,25 @@ inline bool operator!=(OpState lhs, OpState rhs) { 
  return lhs.getOperation() != rhs.getOperation();

}



raw_ostream &operator<<(raw_ostream &os, OpFoldResult ofr);



/// This class represents a single result from folding an operation.

class OpFoldResult : public PointerUnion<Attribute, Value> {

  using PointerUnion<Attribute, Value>::PointerUnion;



public:

  void dump() { llvm::errs() << *this << "\n"; }

};



/// Allow printing to a stream.

inline raw_ostream &operator<<(raw_ostream &os, OpFoldResult ofr) {

  if (Value value = ofr.dyn_cast<Value>())

    value.print(os);

  else

    ofr.dyn_cast<Attribute>().print(os);

  return os;

}



/// Allow printing to a stream.

inline raw_ostream &operator<<(raw_ostream &os, OpState &op) {

  op.print(os, OpPrintingFlags().useLocalScope());

mlir/include/mlir/Interfaces/ViewLikeInterface.td

+66 −22
Original line number Diff line number Diff line
@@ -108,28 +108,6 @@ def OffsetSizeAndStrideOpInterface : OpInterface<"OffsetSizeAndStrideOpInterface 
        return $_op.sizes();

      }]

    >,

    InterfaceMethod<

      /*desc=*/[{

        Return a vector of all the static or dynamic sizes of the op.

      }],

      /*retTy=*/"SmallVector<OpFoldResult, 4>",

      /*methodName=*/"getMixedSizes",

      /*args=*/(ins),

      /*methodBody=*/"",

      /*defaultImplementation=*/[{

        SmallVector<OpFoldResult, 4> res;

        std::array<unsigned, 3> ranks = $_op.getArrayAttrRanks();

        unsigned numDynamic = 0;

        unsigned count = ranks[getOffsetOperandGroupPosition()];

        for (unsigned idx = 0; idx < count; ++idx) {

          if (isDynamicSize(idx))

            res.push_back($_op.sizes()[numDynamic++]);

          else

            res.push_back($_op.static_sizes()[idx]);

        }

        return res;

      }]

    >,

    InterfaceMethod<

      /*desc=*/[{

        Return the dynamic stride operands.
@@ -178,6 +156,72 @@ def OffsetSizeAndStrideOpInterface : OpInterface<"OffsetSizeAndStrideOpInterface 
        return $_op.static_strides();

      }]

    >,

    InterfaceMethod<

      /*desc=*/[{

        Return a vector of all the static or dynamic sizes of the op.

      }],

      /*retTy=*/"SmallVector<OpFoldResult, 4>",

      /*methodName=*/"getMixedOffsets",

      /*args=*/(ins),

      /*methodBody=*/"",

      /*defaultImplementation=*/[{

        SmallVector<OpFoldResult, 4> res;

        std::array<unsigned, 3> ranks = $_op.getArrayAttrRanks();

        unsigned numDynamic = 0;

        unsigned count = ranks[getOffsetOperandGroupPosition()];

        for (unsigned idx = 0; idx < count; ++idx) {

          if (isDynamicOffset(idx))

            res.push_back($_op.offsets()[numDynamic++]);

          else

            res.push_back($_op.static_offsets()[idx]);

        }

        return res;

      }]

    >,

    InterfaceMethod<

      /*desc=*/[{

        Return a vector of all the static or dynamic sizes of the op.

      }],

      /*retTy=*/"SmallVector<OpFoldResult, 4>",

      /*methodName=*/"getMixedSizes",

      /*args=*/(ins),

      /*methodBody=*/"",

      /*defaultImplementation=*/[{

        SmallVector<OpFoldResult, 4> res;

        std::array<unsigned, 3> ranks = $_op.getArrayAttrRanks();

        unsigned numDynamic = 0;

        unsigned count = ranks[getSizeOperandGroupPosition()];

        for (unsigned idx = 0; idx < count; ++idx) {

          if (isDynamicSize(idx))

            res.push_back($_op.sizes()[numDynamic++]);

          else

            res.push_back($_op.static_sizes()[idx]);

        }

        return res;

      }]

    >,

    InterfaceMethod<

      /*desc=*/[{

        Return a vector of all the static or dynamic strides of the op.

      }],

      /*retTy=*/"SmallVector<OpFoldResult, 4>",

      /*methodName=*/"getMixedStrides",

      /*args=*/(ins),

      /*methodBody=*/"",

      /*defaultImplementation=*/[{

        SmallVector<OpFoldResult, 4> res;

        std::array<unsigned, 3> ranks = $_op.getArrayAttrRanks();

        unsigned numDynamic = 0;

        unsigned count = ranks[getStrideOperandGroupPosition()];

        for (unsigned idx = 0; idx < count; ++idx) {

          if (isDynamicStride(idx))

            res.push_back($_op.strides()[numDynamic++]);

          else

            res.push_back($_op.static_strides()[idx]);

        }

        return res;

      }]

    >,



    InterfaceMethod<

      /*desc=*/[{

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

+4 −8
Original line number Diff line number Diff line
@@ -241,10 +241,8 @@ public: 
    Value alloc =

        rewriter.create<AllocOp>(op.getLoc(), subviewMemRefType, op.sizes());

    Value subView = rewriter.create<SubViewOp>(

        op.getLoc(), sourceMemref, extractFromI64ArrayAttr(op.static_offsets()),

        extractFromI64ArrayAttr(op.static_sizes()),

        extractFromI64ArrayAttr(op.static_strides()), op.offsets(), op.sizes(),

        op.strides());

        op.getLoc(), sourceMemref, op.getMixedOffsets(), op.getMixedSizes(),

        op.getMixedStrides());

    rewriter.create<linalg::CopyOp>(op.getLoc(), subView, alloc);

    rewriter.replaceOp(op, alloc);

    return success();
@@ -283,10 +281,8 @@ public: 


    // Take a subview to copy the small memref.

    Value subview = rewriter.create<SubViewOp>(

        op.getLoc(), destMemRef, extractFromI64ArrayAttr(op.static_offsets()),

        extractFromI64ArrayAttr(op.static_sizes()),

        extractFromI64ArrayAttr(op.static_strides()), adaptor.offsets(),

        adaptor.sizes(), adaptor.strides());

        op.getLoc(), destMemRef, op.getMixedOffsets(), op.getMixedSizes(),

        op.getMixedStrides());

    // Copy the small memref.

    rewriter.create<linalg::CopyOp>(op.getLoc(), sourceMemRef, subview);

    rewriter.replaceOp(op, destMemRef);

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

+4 −4
Original line number Diff line number Diff line
@@ -61,9 +61,9 @@ using llvm::dbgs; 
// by `permutationMap`.

static void inferShapeComponents(AffineMap permutationMap,

                                 ArrayRef<Range> loopRanges,

                                 SmallVectorImpl<Value> &offsets,

                                 SmallVectorImpl<Value> &sizes,

                                 SmallVectorImpl<Value> &strides) {

                                 SmallVectorImpl<OpFoldResult> &offsets,

                                 SmallVectorImpl<OpFoldResult> &sizes,

                                 SmallVectorImpl<OpFoldResult> &strides) {

  assert(permutationMap.isProjectedPermutation() &&

         "expected some subset of a permutation map");

  SmallVector<Range, 4> shapeRanges(permutationMap.getNumResults());
@@ -101,7 +101,7 @@ static LinalgOp cloneWithLoopRanges(OpBuilder &b, Location loc, LinalgOp op, 
    AffineMap map = op.getIndexingMap(shapedOperandIdx);

    LLVM_DEBUG(llvm::dbgs() << "shapedOperandIdx: " << shapedOperandIdx

                            << " with indexingMap: " << map << "\n");

    SmallVector<Value, 4> offsets, sizes, strides;

    SmallVector<OpFoldResult, 4> offsets, sizes, strides;

    inferShapeComponents(map, loopRanges, offsets, sizes, strides);

    Value shape = en.value();

    Value sub = shape.getType().isa<MemRefType>()