AffineOps.h

Go to the documentation of this file.

//===- AffineOps.h - MLIR Affine Operations -------------------------------===//
//
// Part of the MLIR Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// This file defines convenience types for working with Affine operations
// in the MLIR operation set.
//
//===----------------------------------------------------------------------===//

#ifndef MLIR_DIALECT_AFFINEOPS_AFFINEOPS_H
#define MLIR_DIALECT_AFFINEOPS_AFFINEOPS_H

#include "mlir/IR/AffineMap.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/Dialect.h"
#include "mlir/IR/OpDefinition.h"
#include "mlir/IR/StandardTypes.h"
#include "mlir/Transforms/LoopLikeInterface.h"

namespace mlir {
class AffineBound;
class AffineDimExpr;
class AffineValueMap;
class AffineTerminatorOp;
class FlatAffineConstraints;
class OpBuilder;

bool isTopLevelValue(Value value);

class AffineOpsDialect : public Dialect {
public:
  AffineOpsDialect(MLIRContext *context);
  static StringRef getDialectNamespace() { return "affine"; }

  Operation *materializeConstant(OpBuilder &builder, Attribute value, Type type,
                                 Location loc) override;
};

class AffineApplyOp : public Op<AffineApplyOp, OpTrait::VariadicOperands,
                                OpTrait::OneResult, OpTrait::HasNoSideEffect> {
public:
  using Op::Op;

  static void build(Builder *builder, OperationState &result, AffineMap map,
                    ValueRange operands);

  AffineMap getAffineMap() {
    return getAttrOfType<AffineMapAttr>("map").getValue();
  }

  bool isValidDim();

  bool isValidSymbol();

  static StringRef getOperationName() { return "affine.apply"; }

  operand_range getMapOperands() { return getOperands(); }

  // Hooks to customize behavior of this op.
  static ParseResult parse(OpAsmParser &parser, OperationState &result);
  void print(OpAsmPrinter &p);
  LogicalResult verify();
  OpFoldResult fold(ArrayRef<Attribute> operands);

  static void getCanonicalizationPatterns(OwningRewritePatternList &results,
                                          MLIRContext *context);
};

//
// For example, a DmaStartOp operation that transfers 256 elements of a memref
// '%src' in memory space 0 at indices [%i + 3, %j] to memref '%dst' in memory
// space 1 at indices [%k + 7, %l], would be specified as follows:
//
//   %num_elements = constant 256
//   %idx = constant 0 : index
//   %tag = alloc() : memref<1xi32, 4>
//   affine.dma_start %src[%i + 3, %j], %dst[%k + 7, %l], %tag[%idx],
//     %num_elements :
//       memref<40x128xf32, 0>, memref<2x1024xf32, 1>, memref<1xi32, 2>
//
//   If %stride and %num_elt_per_stride are specified, the DMA is expected to
//   transfer %num_elt_per_stride elements every %stride elements apart from
//   memory space 0 until %num_elements are transferred.
//
//   affine.dma_start %src[%i, %j], %dst[%k, %l], %tag[%idx], %num_elements,
//     %stride, %num_elt_per_stride : ...
//
// TODO(mlir-team): add additional operands to allow source and destination
// striding, and multiple stride levels (possibly using AffineMaps to specify
// multiple levels of striding).
// TODO(andydavis) Consider replacing src/dst memref indices with view memrefs.
class AffineDmaStartOp : public Op<AffineDmaStartOp, OpTrait::VariadicOperands,
                                   OpTrait::ZeroResult> {
public:
  using Op::Op;

  static void build(Builder *builder, OperationState &result, Value srcMemRef,
                    AffineMap srcMap, ValueRange srcIndices, Value destMemRef,
                    AffineMap dstMap, ValueRange destIndices, Value tagMemRef,
                    AffineMap tagMap, ValueRange tagIndices, Value numElements,
                    Value stride = nullptr, Value elementsPerStride = nullptr);

  unsigned getSrcMemRefOperandIndex() { return 0; }

  Value getSrcMemRef() { return getOperand(getSrcMemRefOperandIndex()); }
  MemRefType getSrcMemRefType() {
    return getSrcMemRef()->getType().cast<MemRefType>();
  }

  unsigned getSrcMemRefRank() { return getSrcMemRefType().getRank(); }

  AffineMap getSrcMap() { return getSrcMapAttr().getValue(); }
  AffineMapAttr getSrcMapAttr() {
    return getAttr(getSrcMapAttrName()).cast<AffineMapAttr>();
  }

  operand_range getSrcIndices() {
    return {operand_begin() + getSrcMemRefOperandIndex() + 1,
            operand_begin() + getSrcMemRefOperandIndex() + 1 +
                getSrcMap().getNumInputs()};
  }

  unsigned getSrcMemorySpace() {
    return getSrcMemRef()->getType().cast<MemRefType>().getMemorySpace();
  }

  unsigned getDstMemRefOperandIndex() {
    return getSrcMemRefOperandIndex() + 1 + getSrcMap().getNumInputs();
  }

  Value getDstMemRef() { return getOperand(getDstMemRefOperandIndex()); }
  MemRefType getDstMemRefType() {
    return getDstMemRef()->getType().cast<MemRefType>();
  }

  unsigned getDstMemRefRank() {
    return getDstMemRef()->getType().cast<MemRefType>().getRank();
  }

  unsigned getDstMemorySpace() {
    return getDstMemRef()->getType().cast<MemRefType>().getMemorySpace();
  }

  AffineMap getDstMap() { return getDstMapAttr().getValue(); }
  AffineMapAttr getDstMapAttr() {
    return getAttr(getDstMapAttrName()).cast<AffineMapAttr>();
  }

  operand_range getDstIndices() {
    return {operand_begin() + getDstMemRefOperandIndex() + 1,
            operand_begin() + getDstMemRefOperandIndex() + 1 +
                getDstMap().getNumInputs()};
  }

  unsigned getTagMemRefOperandIndex() {
    return getDstMemRefOperandIndex() + 1 + getDstMap().getNumInputs();
  }

  Value getTagMemRef() { return getOperand(getTagMemRefOperandIndex()); }
  MemRefType getTagMemRefType() {
    return getTagMemRef()->getType().cast<MemRefType>();
  }

  unsigned getTagMemRefRank() {
    return getTagMemRef()->getType().cast<MemRefType>().getRank();
  }

  AffineMap getTagMap() { return getTagMapAttr().getValue(); }
  AffineMapAttr getTagMapAttr() {
    return getAttr(getTagMapAttrName()).cast<AffineMapAttr>();
  }

  operand_range getTagIndices() {
    return {operand_begin() + getTagMemRefOperandIndex() + 1,
            operand_begin() + getTagMemRefOperandIndex() + 1 +
                getTagMap().getNumInputs()};
  }

  Value getNumElements() {
    return getOperand(getTagMemRefOperandIndex() + 1 +
                      getTagMap().getNumInputs());
  }

  NamedAttribute getAffineMapAttrForMemRef(Value memref) {
    if (memref == getSrcMemRef())
      return {Identifier::get(getSrcMapAttrName(), getContext()),
              getSrcMapAttr()};
    else if (memref == getDstMemRef())
      return {Identifier::get(getDstMapAttrName(), getContext()),
              getDstMapAttr()};
    assert(memref == getTagMemRef() &&
           "DmaStartOp expected source, destination or tag memref");
    return {Identifier::get(getTagMapAttrName(), getContext()),
            getTagMapAttr()};
  }

  bool isDestMemorySpaceFaster() {
    return (getSrcMemorySpace() < getDstMemorySpace());
  }

  bool isSrcMemorySpaceFaster() {
    // Assumes that a lower number is for a slower memory space.
    return (getDstMemorySpace() < getSrcMemorySpace());
  }

  unsigned getFasterMemPos() {
    assert(isSrcMemorySpaceFaster() || isDestMemorySpaceFaster());
    return isSrcMemorySpaceFaster() ? 0 : getDstMemRefOperandIndex();
  }

  static StringRef getSrcMapAttrName() { return "src_map"; }
  static StringRef getDstMapAttrName() { return "dst_map"; }
  static StringRef getTagMapAttrName() { return "tag_map"; }

  static StringRef getOperationName() { return "affine.dma_start"; }
  static ParseResult parse(OpAsmParser &parser, OperationState &result);
  void print(OpAsmPrinter &p);
  LogicalResult verify();
  LogicalResult fold(ArrayRef<Attribute> cstOperands,
                     SmallVectorImpl<OpFoldResult> &results);

  bool isStrided() {
    return getNumOperands() !=
           getTagMemRefOperandIndex() + 1 + getTagMap().getNumInputs() + 1;
  }

  Value getStride() {
    if (!isStrided())
      return nullptr;
    return getOperand(getNumOperands() - 1 - 1);
  }

  Value getNumElementsPerStride() {
    if (!isStrided())
      return nullptr;
    return getOperand(getNumOperands() - 1);
  }
};

//
//   affine.dma_start %src[%i, %j], %dst[%k, %l], %tag[%index], %num_elements :
//     memref<2048xf32, 0>, memref<256xf32, 1>, memref<1xi32, 2>
//   ...
//   ...
//   affine.dma_wait %tag[%index], %num_elements : memref<1xi32, 2>
//
class AffineDmaWaitOp : public Op<AffineDmaWaitOp, OpTrait::VariadicOperands,
                                  OpTrait::ZeroResult> {
public:
  using Op::Op;

  static void build(Builder *builder, OperationState &result, Value tagMemRef,
                    AffineMap tagMap, ValueRange tagIndices, Value numElements);

  static StringRef getOperationName() { return "affine.dma_wait"; }

  // Returns the Tag MemRef associated with the DMA operation being waited on.
  Value getTagMemRef() { return getOperand(0); }
  MemRefType getTagMemRefType() {
    return getTagMemRef()->getType().cast<MemRefType>();
  }

  AffineMap getTagMap() { return getTagMapAttr().getValue(); }
  AffineMapAttr getTagMapAttr() {
    return getAttr(getTagMapAttrName()).cast<AffineMapAttr>();
  }

  // Returns the tag memref index for this DMA operation.
  operand_range getTagIndices() {
    return {operand_begin() + 1,
            operand_begin() + 1 + getTagMap().getNumInputs()};
  }

  // Returns the rank (number of indices) of the tag memref.
  unsigned getTagMemRefRank() {
    return getTagMemRef()->getType().cast<MemRefType>().getRank();
  }

  NamedAttribute getAffineMapAttrForMemRef(Value memref) {
    assert(memref == getTagMemRef());
    return {Identifier::get(getTagMapAttrName(), getContext()),
            getTagMapAttr()};
  }

  Value getNumElements() { return getOperand(1 + getTagMap().getNumInputs()); }

  static StringRef getTagMapAttrName() { return "tag_map"; }
  static ParseResult parse(OpAsmParser &parser, OperationState &result);
  void print(OpAsmPrinter &p);
  LogicalResult verify();
  LogicalResult fold(ArrayRef<Attribute> cstOperands,
                     SmallVectorImpl<OpFoldResult> &results);
};

//
//  Example 1:
//
//    %1 = affine.load %0[%i0 + 3, %i1 + 7] : memref<100x100xf32>
//
//  Example 2: Uses 'symbol' keyword for symbols '%n' and '%m'.
//
//    %1 = affine.load %0[%i0 + symbol(%n), %i1 + symbol(%m)]
//      : memref<100x100xf32>
//
class AffineLoadOp : public Op<AffineLoadOp, OpTrait::OneResult,
                               OpTrait::AtLeastNOperands<1>::Impl> {
public:
  using Op::Op;

  static void build(Builder *builder, OperationState &result, AffineMap map,
                    ValueRange operands);
  static void build(Builder *builder, OperationState &result, Value memref,
                    ValueRange indices = {});
  static void build(Builder *builder, OperationState &result, Value memref,
                    AffineMap map, ValueRange mapOperands);

  unsigned getMemRefOperandIndex() { return 0; }

  Value getMemRef() { return getOperand(getMemRefOperandIndex()); }
  void setMemRef(Value value) { setOperand(getMemRefOperandIndex(), value); }
  MemRefType getMemRefType() {
    return getMemRef()->getType().cast<MemRefType>();
  }

  operand_range getMapOperands() { return llvm::drop_begin(getOperands(), 1); }

  AffineMap getAffineMap() { return getAffineMapAttr().getValue(); }
  AffineMapAttr getAffineMapAttr() {
    return getAttr(getMapAttrName()).cast<AffineMapAttr>();
  }

  NamedAttribute getAffineMapAttrForMemRef(Value memref) {
    assert(memref == getMemRef());
    return {Identifier::get(getMapAttrName(), getContext()),
            getAffineMapAttr()};
  }

  static StringRef getMapAttrName() { return "map"; }
  static StringRef getOperationName() { return "affine.load"; }

  // Hooks to customize behavior of this op.
  static ParseResult parse(OpAsmParser &parser, OperationState &result);
  void print(OpAsmPrinter &p);
  LogicalResult verify();
  static void getCanonicalizationPatterns(OwningRewritePatternList &results,
                                          MLIRContext *context);
  OpFoldResult fold(ArrayRef<Attribute> operands);
};

//
//  Example 1:
//
//    affine.store %v0, %0[%i0 + 3, %i1 + 7] : memref<100x100xf32>
//
//  Example 2: Uses 'symbol' keyword for symbols '%n' and '%m'.
//
//    affine.store %v0, %0[%i0 + symbol(%n), %i1 + symbol(%m)]
//      : memref<100x100xf32>
//
class AffineStoreOp : public Op<AffineStoreOp, OpTrait::ZeroResult,
                                OpTrait::AtLeastNOperands<1>::Impl> {
public:
  using Op::Op;

  static void build(Builder *builder, OperationState &result,
                    Value valueToStore, Value memref, ValueRange indices);
  static void build(Builder *builder, OperationState &result,
                    Value valueToStore, Value memref, AffineMap map,
                    ValueRange mapOperands);

  Value getValueToStore() { return getOperand(0); }

  unsigned getMemRefOperandIndex() { return 1; }

  Value getMemRef() { return getOperand(getMemRefOperandIndex()); }
  void setMemRef(Value value) { setOperand(getMemRefOperandIndex(), value); }

  MemRefType getMemRefType() {
    return getMemRef()->getType().cast<MemRefType>();
  }

  operand_range getMapOperands() { return llvm::drop_begin(getOperands(), 2); }

  AffineMap getAffineMap() { return getAffineMapAttr().getValue(); }
  AffineMapAttr getAffineMapAttr() {
    return getAttr(getMapAttrName()).cast<AffineMapAttr>();
  }

  NamedAttribute getAffineMapAttrForMemRef(Value memref) {
    assert(memref == getMemRef());
    return {Identifier::get(getMapAttrName(), getContext()),
            getAffineMapAttr()};
  }

  static StringRef getMapAttrName() { return "map"; }
  static StringRef getOperationName() { return "affine.store"; }

  // Hooks to customize behavior of this op.
  static ParseResult parse(OpAsmParser &parser, OperationState &result);
  void print(OpAsmPrinter &p);
  LogicalResult verify();
  static void getCanonicalizationPatterns(OwningRewritePatternList &results,
                                          MLIRContext *context);
  LogicalResult fold(ArrayRef<Attribute> cstOperands,
                     SmallVectorImpl<OpFoldResult> &results);
};

bool isValidDim(Value value);

bool isValidSymbol(Value value);

void canonicalizeMapAndOperands(AffineMap *map,
                                SmallVectorImpl<Value> *operands);
void canonicalizeSetAndOperands(IntegerSet *set,
                                SmallVectorImpl<Value> *operands);

AffineApplyOp makeComposedAffineApply(OpBuilder &b, Location loc, AffineMap map,
                                      ArrayRef<Value> operands);

void fullyComposeAffineMapAndOperands(AffineMap *map,
                                      SmallVectorImpl<Value> *operands);

#define GET_OP_CLASSES
#include "mlir/Dialect/AffineOps/AffineOps.h.inc"

bool isForInductionVar(Value val);

AffineForOp getForInductionVarOwner(Value val);

void extractForInductionVars(ArrayRef<AffineForOp> forInsts,
                             SmallVectorImpl<Value> *ivs);

class AffineBound {
public:
  AffineForOp getAffineForOp() { return op; }
  AffineMap getMap() { return map; }

  AffineValueMap getAsAffineValueMap();

  unsigned getNumOperands() { return opEnd - opStart; }
  Value getOperand(unsigned idx) { return op.getOperand(opStart + idx); }

  using operand_iterator = AffineForOp::operand_iterator;
  using operand_range = AffineForOp::operand_range;

  operand_iterator operand_begin() { return op.operand_begin() + opStart; }
  operand_iterator operand_end() { return op.operand_begin() + opEnd; }
  operand_range getOperands() { return {operand_begin(), operand_end()}; }

private:
  // 'affine.for' operation that contains this bound.
  AffineForOp op;
  // Start and end positions of this affine bound operands in the list of
  // the containing 'affine.for' operation operands.
  unsigned opStart, opEnd;
  // Affine map for this bound.
  AffineMap map;

  AffineBound(AffineForOp op, unsigned opStart, unsigned opEnd, AffineMap map)
      : op(op), opStart(opStart), opEnd(opEnd), map(map) {}

  friend class AffineForOp;
};

struct AffineApplyNormalizer {
  AffineApplyNormalizer(AffineMap map, ArrayRef<Value> operands);

  AffineMap getAffineMap() { return affineMap; }

  SmallVector<Value, 8> getOperands() {
    SmallVector<Value, 8> res(reorderedDims);
    res.append(concatenatedSymbols.begin(), concatenatedSymbols.end());
    return res;
  }

  unsigned getNumSymbols() { return concatenatedSymbols.size(); }
  unsigned getNumDims() { return reorderedDims.size(); }

  void normalize(AffineMap *otherMap, SmallVectorImpl<Value> *otherOperands);

private:
  AffineDimExpr renumberOneDim(Value v);

  AffineMap renumber(const AffineApplyNormalizer &other);

  DenseMap<Value, unsigned> dimValueToPosition;

  SmallVector<Value, 8> reorderedDims;
  SmallVector<Value, 8> concatenatedSymbols;

  AffineMap affineMap;

  static unsigned &affineApplyDepth() {
    static thread_local unsigned depth = 0;
    return depth;
  }
  static constexpr unsigned kMaxAffineApplyDepth = 1;

  AffineApplyNormalizer() { affineApplyDepth()++; }

public:
  ~AffineApplyNormalizer() { affineApplyDepth()--; }
};

} // end namespace mlir

#endif