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=https://github.com/tensorflow/mlir/pull/225 from bondhugula:prefetch 4c3b4e93bc64d9a5719504e6d6e1657818a2ead0
PiperOrigin-RevId: 286212997

5 files changed, 146 insertions, 1 deletions

diff --git a/mlir/include/mlir/Dialect/AffineOps/AffineOps.td b/mlir/include/mlir/Dialect/AffineOps/AffineOps.td
index cea44b8dacd..b40990ecb5d 100644
--- a/mlir/include/mlir/Dialect/AffineOps/AffineOps.td
+++ b/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";
diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
index 4deab2e7aef..a8b427e2d94 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
+++ b/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)> {
diff --git a/mlir/include/mlir/Dialect/StandardOps/Ops.td b/mlir/include/mlir/Dialect/StandardOps/Ops.td
index 2cba150560c..76c2ba57ea6 100644
--- a/mlir/include/mlir/Dialect/StandardOps/Ops.td
+++ b/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 = [{
diff --git a/mlir/include/mlir/IR/OpBase.td b/mlir/include/mlir/IR/OpBase.td
index dd7fac27a00..4d5d1fe766f 100644
--- a/mlir/include/mlir/IR/OpBase.td
+++ b/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>,
diff --git a/mlir/include/mlir/IR/OpImplementation.h b/mlir/include/mlir/IR/OpImplementation.h
index 05beaeaee59..97569cc06d9 100644
--- a/mlir/include/mlir/IR/OpImplementation.h
+++ b/mlir/include/mlir/IR/OpImplementation.h
@@ -284,6 +284,12 @@ public:
   /// 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();