Introduce prefetch op: affine -> std -> llvm intrinsic

Introduce affine.prefetch: op to prefetch using a multi-dimensional
subscript on a memref; similar to affine.load but has no effect on
semantics, but only on performance.

Provide lowering through std.prefetch, llvm.prefetch and map to llvm's
prefetch instrinsic. All attributes reflected through the lowering -
locality hint, rw, and instr/data cache.

  affine.prefetch %0[%i, %j + 5], false, 3, true : memref<400x400xi32>

Signed-off-by: Uday Bondhugula <uday@polymagelabs.com>

Closes tensorflow/mlir#225

COPYBARA_INTEGRATE_REVIEW=tensorflow/mlir#225 from bondhugula:prefetch 4c3b4e93bc64d9a5719504e6d6e1657818a2ead0
PiperOrigin-RevId: 286212997

Author: bondhugula

mlir/g3doc/OpDefinitions.md

@@ -243,19 +243,21 @@ like `"0.5f"`, and an integer array default value should be specified as like
 `Confined` is provided as a general mechanism to help modelling further
 constraints on attributes beyond the ones brought by value types. You can use
 `Confined` to compose complex constraints out of more primitive ones. For
-example, a 32-bit integer attribute whose minimal value must be 10 can be
+example, a 32-bit integer attribute whose minimum value must be 10 can be
 expressed as `Confined<I32Attr, [IntMinValue<10>]>`.
 
 Right now, the following primitive constraints are supported:
 
-* `IntMinValue<N>`: Specifying an integer attribute to be greater than or equal
-  to `N`
-* `ArrayMinCount<N>`: Specifying an array attribute to have at least `N`
-  elements
-* `IntArrayNthElemEq<I, N>`: Specifying an integer array attribute's `I`-th
-  element to be equal to `N`
-* `IntArrayNthElemMinValue<I, N>`: Specifying an integer array attribute's
-  `I`-th element to be greater than or equal to `N`
+*   `IntMinValue<N>`: Specifying an integer attribute to be greater than or
+    equal to `N`
+*   `IntMaxValue<N>`: Specifying an integer attribute to be less than or equal
+    to `N`
+*   `ArrayMinCount<N>`: Specifying an array attribute to have at least `N`
+    elements
+*   `IntArrayNthElemEq<I, N>`: Specifying an integer array attribute's `I`-th
+    element to be equal to `N`
+*   `IntArrayNthElemMinValue<I, N>`: Specifying an integer array attribute's
+    `I`-th element to be greater than or equal to `N`
 
 TODO: Design and implement more primitive constraints
 

mlir/include/mlir/Dialect/AffineOps/AffineOps.td

@@ -261,6 +261,81 @@ def AffineMinOp : Affine_Op<"min"> {
   let hasFolder = 1;
 }
 
+def AffinePrefetchOp : Affine_Op<"prefetch"> {
+  let summary = "affine prefetch operation";
+  let description = [{
+    The "affine.prefetch" op prefetches data from a memref location described
+    with an affine subscript similar to affine.load, and has three attributes:
+    a read/write specifier, a locality hint, and a cache type specifier as shown
+    below:
+
+      affine.prefetch %0[%i, %j + 5], read, locality<3>, data
+                          : memref<400x400xi32>
+
+    The read/write specifier is either 'read' or 'write', the locality hint
+    specifier ranges from locality<0> (no locality) to locality<3> (extremely
+    local keep in cache). The cache type specifier is either 'data' or 'instr'
+    and specifies whether the prefetch is performed on data cache or on
+    instruction cache.
+  }];
+
+  let arguments = (ins AnyMemRef:$memref, Variadic<Index>:$indices,
+                   BoolAttr:$isWrite,
+                   Confined<I32Attr, [IntMinValue<0>,
+                     IntMaxValue<3>]>:$localityHint,
+                   BoolAttr:$isDataCache);
+
+  let builders = [OpBuilder<
+    "Builder *builder, OperationState &result, Value *memref,"
+    "AffineMap map, ArrayRef<Value *> mapOperands, bool isWrite,"
+    "unsigned localityHint, bool isDataCache",
+    [{
+      assert(map.getNumInputs() == mapOperands.size()
+             && "inconsistent index info");
+      auto localityHintAttr = builder->getI32IntegerAttr(localityHint);
+      auto isWriteAttr = builder->getBoolAttr(isWrite);
+      auto isDataCacheAttr = builder->getBoolAttr(isDataCache);
+      result.addOperands(memref);
+      result.addAttribute(getMapAttrName(), AffineMapAttr::get(map));
+      result.addOperands(mapOperands);
+      result.addAttribute(getLocalityHintAttrName(), localityHintAttr);
+      result.addAttribute(getIsWriteAttrName(), isWriteAttr);
+      result.addAttribute(getIsDataCacheAttrName(), isDataCacheAttr);
+    }]>];
+
+  let extraClassDeclaration = [{
+    MemRefType getMemRefType() {
+      return memref()->getType().cast<MemRefType>();
+    }
+
+    /// Returns the affine map used to index the memref for this operation.
+    AffineMap getAffineMap() { return getAffineMapAttr().getValue(); }
+    AffineMapAttr getAffineMapAttr() {
+      return getAttr(getMapAttrName()).cast<AffineMapAttr>();
+    }
+
+    /// Returns the AffineMapAttr associated with 'memref'.
+    NamedAttribute getAffineMapAttrForMemRef(Value *mref) {
+      assert(mref == memref());
+      return {Identifier::get(getMapAttrName(), getContext()),
+        getAffineMapAttr()};
+    }
+
+    /// Get affine map operands.
+    operand_range getMapOperands() {
+      return {operand_begin() + 1, operand_end()};
+    }
+
+    static StringRef getMapAttrName() { return "map"; }
+    static StringRef getLocalityHintAttrName() { return "localityHint"; }
+    static StringRef getIsWriteAttrName() { return "isWrite"; }
+    static StringRef getIsDataCacheAttrName() { return "isDataCache"; }
+  }];
+
+  let hasCanonicalizer = 1;
+  let hasFolder = 1;
+}
+
 def AffineTerminatorOp :
     Affine_Op<"terminator", [Terminator]> {
   let summary = "affine terminator operation";

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

@@ -674,6 +674,17 @@ def LLVM_FMulAddOp : LLVM_Op<"intr.fmuladd", [NoSideEffect]>,
   }];
 }
 
+def LLVM_Prefetch : LLVM_ZeroResultOp<"intr.prefetch">,
+                    Arguments<(ins LLVM_Type:$addr, LLVM_Type:$rw,
+                    LLVM_Type:$hint, LLVM_Type:$cache)> {
+  let llvmBuilder = [{
+    llvm::Module *module = builder.GetInsertBlock()->getModule();
+    llvm::Function *fn = llvm::Intrinsic::getDeclaration(
+        module, llvm::Intrinsic::prefetch, $addr->getType());
+    builder.CreateCall(fn, {$addr, $rw, $hint, $cache});
+  }];
+}
+
 def LLVM_ExpOp : LLVM_Op<"intr.exp", [NoSideEffect]>,
                    Arguments<(ins LLVM_Type:$in)>,
                    Results<(outs LLVM_Type:$res)> {

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

@@ -928,6 +928,55 @@ def OrOp : IntArithmeticOp<"or", [Commutative]> {
   let hasFolder = 1;
 }
 
+def PrefetchOp : Std_Op<"prefetch"> {
+  let summary = "prefetch operation";
+  let description = [{
+    The "prefetch" op prefetches data from a memref location described with
+    subscript indices similar to std.load, and with three attributes: a
+    read/write specifier, a locality hint, and a cache type specifier as shown
+    below:
+
+      prefetch %0[%i, %j], read, locality<3>, data : memref<400x400xi32>
+
+    The read/write specifier is either 'read' or 'write', the locality hint
+    ranges from locality<0> (no locality) to locality<3> (extremely local keep
+    in cache). The cache type specifier is either 'data' or 'instr'
+    and specifies whether the prefetch is performed on data cache or on
+    instruction cache.
+  }];
+
+  let arguments = (ins AnyMemRef:$memref, Variadic<Index>:$indices,
+                   BoolAttr:$isWrite,
+                   Confined<I32Attr, [IntMinValue<0>,
+                     IntMaxValue<3>]>:$localityHint,
+                   BoolAttr:$isDataCache);
+
+  let builders = [OpBuilder<
+    "Builder *builder, OperationState &result, Value *memref,"
+    "ArrayRef<Value *> indices, bool isWrite, unsigned hint, bool isData",
+    [{
+      auto hintAttr = builder->getI32IntegerAttr(hint);
+      auto isWriteAttr = builder->getBoolAttr(isWrite);
+      auto isDataCacheAttr = builder->getBoolAttr(isData);
+      result.addOperands(memref);
+      result.addOperands(indices);
+      result.addAttribute("localityHint", hintAttr);
+      result.addAttribute("isWrite", isWriteAttr);
+      result.addAttribute("isDataCache", isDataCacheAttr);
+    }]>];
+
+  let extraClassDeclaration = [{
+    MemRefType getMemRefType() {
+      return memref()->getType().cast<MemRefType>();
+    }
+    static StringRef getLocalityHintAttrName() { return "localityHint"; }
+    static StringRef getIsWriteAttrName() { return "isWrite"; }
+    static StringRef getIsDataCacheAttrName() { return "isDataCache"; }
+  }];
+
+  let hasFolder = 1;
+}
+
 def RankOp : Std_Op<"rank", [NoSideEffect]> {
   let summary = "rank operation";
   let description = [{

mlir/include/mlir/IR/OpBase.td

@@ -1242,7 +1242,11 @@ class AllAttrConstraintsOf<list<AttrConstraint> constraints> : AttrConstraint<
 
 class IntMinValue<int n> : AttrConstraint<
     CPred<"$_self.cast<IntegerAttr>().getInt() >= " # n>,
-    "whose minimal value is " # n>;
+    "whose minimum value is " # n>;
+
+class IntMaxValue<int n> : AttrConstraint<
+    CPred<"$_self.cast<IntegerAttr>().getInt() <= " # n>,
+    "whose maximum value is " # n>;
 
 class ArrayMinCount<int n> : AttrConstraint<
     CPred<"$_self.cast<ArrayAttr>().size() >= " # n>,

mlir/include/mlir/IR/OpImplementation.h

@@ -284,6 +284,12 @@ class OpAsmParser {
   /// Parse a `=` token.
   virtual ParseResult parseEqual() = 0;
 
+  /// Parse a '<' token.
+  virtual ParseResult parseLess() = 0;
+
+  /// Parse a '>' token.
+  virtual ParseResult parseGreater() = 0;
+
   /// Parse a given keyword.
   ParseResult parseKeyword(StringRef keyword, const Twine &msg = "") {
     auto loc = getCurrentLocation();

mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp

@@ -405,13 +405,37 @@ class AffineLoadLowering : public OpRewritePattern<AffineLoadOp> {
                                      PatternRewriter &rewriter) const override {
     // Expand affine map from 'affineLoadOp'.
     SmallVector<Value *, 8> indices(op.getMapOperands());
-    auto maybeExpandedMap =
+    auto resultOperands =
         expandAffineMap(rewriter, op.getLoc(), op.getAffineMap(), indices);
-    if (!maybeExpandedMap)
+    if (!resultOperands)
       return matchFailure();
 
     // Build std.load memref[expandedMap.results].
-    rewriter.replaceOpWithNewOp<LoadOp>(op, op.getMemRef(), *maybeExpandedMap);
+    rewriter.replaceOpWithNewOp<LoadOp>(op, op.getMemRef(), *resultOperands);
+    return matchSuccess();
+  }
+};
+
+// Apply the affine map from an 'affine.prefetch' operation to its operands, and
+// feed the results to a newly created 'std.prefetch' operation (which replaces
+// the original 'affine.prefetch').
+class AffinePrefetchLowering : public OpRewritePattern<AffinePrefetchOp> {
+public:
+  using OpRewritePattern<AffinePrefetchOp>::OpRewritePattern;
+
+  PatternMatchResult matchAndRewrite(AffinePrefetchOp op,
+                                     PatternRewriter &rewriter) const override {
+    // Expand affine map from 'affinePrefetchOp'.
+    SmallVector<Value *, 8> indices(op.getMapOperands());
+    auto resultOperands =
+        expandAffineMap(rewriter, op.getLoc(), op.getAffineMap(), indices);
+    if (!resultOperands)
+      return matchFailure();
+
+    // Build std.prefetch memref[expandedMap.results].
+    rewriter.replaceOpWithNewOp<PrefetchOp>(
+        op, op.memref(), *resultOperands, op.isWrite(),
+        op.localityHint().getZExtValue(), op.isDataCache());
     return matchSuccess();
   }
 };
@@ -506,11 +530,10 @@ class AffineDmaWaitLowering : public OpRewritePattern<AffineDmaWaitOp> {
 
 void mlir::populateAffineToStdConversionPatterns(
     OwningRewritePatternList &patterns, MLIRContext *ctx) {
-  patterns
-      .insert<AffineApplyLowering, AffineDmaStartLowering,
-              AffineDmaWaitLowering, AffineLoadLowering, AffineStoreLowering,
-              AffineForLowering, AffineIfLowering, AffineTerminatorLowering>(
-          ctx);
+  patterns.insert<
+      AffineApplyLowering, AffineDmaStartLowering, AffineDmaWaitLowering,
+      AffineLoadLowering, AffinePrefetchLowering, AffineStoreLowering,
+      AffineForLowering, AffineIfLowering, AffineTerminatorLowering>(ctx);
 }
 
 namespace {

mlir/lib/Conversion/StandardToLLVM/ConvertStandardToLLVM.cpp

@@ -1462,6 +1462,39 @@ struct StoreOpLowering : public LoadStoreOpLowering<StoreOp> {
   }
 };
 
+// The prefetch operation is lowered in a way similar to the load operation
+// except that the llvm.prefetch operation is used for replacement.
+struct PrefetchOpLowering : public LoadStoreOpLowering<PrefetchOp> {
+  using Base::Base;
+
+  PatternMatchResult
+  matchAndRewrite(Operation *op, ArrayRef<Value *> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto prefetchOp = cast<PrefetchOp>(op);
+    OperandAdaptor<PrefetchOp> transformed(operands);
+    auto type = prefetchOp.getMemRefType();
+
+    Value *dataPtr = getDataPtr(op->getLoc(), type, transformed.memref(),
+                                transformed.indices(), rewriter, getModule());
+
+    // Replace with llvm.prefetch.
+    auto llvmI32Type = lowering.convertType(rewriter.getIntegerType(32));
+    auto isWrite = rewriter.create<LLVM::ConstantOp>(
+        op->getLoc(), llvmI32Type,
+        rewriter.getI32IntegerAttr(prefetchOp.isWrite()));
+    auto localityHint = rewriter.create<LLVM::ConstantOp>(
+        op->getLoc(), llvmI32Type,
+        rewriter.getI32IntegerAttr(prefetchOp.localityHint().getZExtValue()));
+    auto isData = rewriter.create<LLVM::ConstantOp>(
+        op->getLoc(), llvmI32Type,
+        rewriter.getI32IntegerAttr(prefetchOp.isDataCache()));
+
+    rewriter.replaceOpWithNewOp<LLVM::Prefetch>(op, dataPtr, isWrite,
+                                                localityHint, isData);
+    return matchSuccess();
+  }
+};
+
 // The lowering of index_cast becomes an integer conversion since index becomes
 // an integer.  If the bit width of the source and target integer types is the
 // same, just erase the cast.  If the target type is wider, sign-extend the
@@ -2041,6 +2074,7 @@ void mlir::populateStdToLLVMNonMemoryConversionPatterns(
       MulFOpLowering,
       MulIOpLowering,
       OrOpLowering,
+      PrefetchOpLowering,
       RemFOpLowering,
       RemISOpLowering,
       RemIUOpLowering,