[mlir] Add nontemporal field to memref.load/store and convey to llvm.load/store

[guraypp]

Summary

llvm.load op has nonTemporal field which is missing for memref.load and memref.store. This revision first adds nonTemporal field to memref's load/store op, then it lowers the field to llvm.load/store ops.

[guraypp]

cleanup

[guraypp]

rebase

[nicolasvasilache]

Have you checked that foldings and canonicalizations don't just drop the new attribute?

Some of the folding / canonicalizations tests should be updated to cover this.

[guraypp]

Added flag propogation on the folding passes, and added new tests.

[nicolasvasilache]

mlir/test/Dialect/MemRef/fold-memref-alias-ops.mlir

527 ↗

(On Diff #494542)

nit: nl

[guraypp]

rebase, address the comment

[guraypp]

rebase

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

let description = [{
     operations) to precisely analyze references at compile-time using polyhedral
     techniques. This is possible because of the
     [restrictions on dimensions and symbols](Affine.md/#restrictions-on-dimensions-and-symbols)
     in these contexts.
   }];
 
   let arguments = (ins Arg<AnyMemRef, "the reference to load from",
                            [MemRead]>:$memref,
-                       Variadic<Index>:$indices);
+                       Variadic<Index>:$indices,                       
+                       DefaultValuedOptionalAttr<BoolAttr, "false">:$nontemporal);
   let results = (outs AnyType:$result);
 
   let extraClassDeclaration = [{
     Value getMemRef() { return getOperand(0); }
     void setMemRef(Value value) { setOperand(0, value); }
     MemRefType getMemRefType() {
       return getMemRef().getType().cast<MemRefType>();
     }
let description = [{
     techniques. This is possible because of the
     [restrictions on dimensions and symbols](Affine.md/#restrictions-on-dimensions-and-symbols)
     in these contexts.
   }];
 
   let arguments = (ins AnyType:$value,
                        Arg<AnyMemRef, "the reference to store to",
                            [MemWrite]>:$memref,
-                       Variadic<Index>:$indices);
+                       Variadic<Index>:$indices,                       
+                       DefaultValuedOptionalAttr<BoolAttr, "false">:$nontemporal);
 
   let builders = [
     OpBuilder<(ins "Value":$valueToStore, "Value":$memref), [{
       $_state.addOperands(valueToStore);
       $_state.addOperands(memref);
     }]>];
 
   let extraClassDeclaration = [{

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

struct LoadOpLowering : public LoadStoreOpLowering<memref::LoadOp> {
   LogicalResult
   matchAndRewrite(memref::LoadOp loadOp, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     auto type = loadOp.getMemRefType();
 
     Value dataPtr =
         getStridedElementPtr(loadOp.getLoc(), type, adaptor.getMemref(),
                              adaptor.getIndices(), rewriter);
-    rewriter.replaceOpWithNewOp<LLVM::LoadOp>(loadOp, dataPtr);
+    rewriter.replaceOpWithNewOp<LLVM::LoadOp>(loadOp, dataPtr, 0, false,
+                                              loadOp.getNontemporal());
     return success();
   }
 };
 
 // Store operation is lowered to obtaining a pointer to the indexed element,
 // and storing the given value to it.
 struct StoreOpLowering : public LoadStoreOpLowering<memref::StoreOp> {
   using Base::Base;
 
   LogicalResult
   matchAndRewrite(memref::StoreOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     auto type = op.getMemRefType();
 
     Value dataPtr = getStridedElementPtr(op.getLoc(), type, adaptor.getMemref(),
                                          adaptor.getIndices(), rewriter);
-    rewriter.replaceOpWithNewOp<LLVM::StoreOp>(op, adaptor.getValue(), dataPtr);
+    rewriter.replaceOpWithNewOp<LLVM::StoreOp>(op, adaptor.getValue(), dataPtr,
+                                               0, false, op.getNontemporal());
     return success();
   }
 };
 
 // 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<memref::PrefetchOp> {
   using Base::Base;

mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir

func.func @extract_strided_metadata(
   %base, %offset, %sizes:2, %strides:2 =
     memref.extract_strided_metadata %ref : memref<?x?xf32, strided<[?,?], offset: ?>>
     -> memref<f32>, index,
        index, index,
        index, index
 
   return
 }
+
+// -----
+
+// CHECK-LABEL: func @load_non_temporal(
+func.func @load_non_temporal(%arg0 : memref<32xf32, affine_map<(d0) -> (d0)>>) {
+  %0 = memref.alloc() : memref<32xf32, affine_map<(d0) -> (d0)>>
+  %1 = arith.constant 7 : index
+  // CHECK: llvm.load %{{.*}} {nontemporal} : !llvm.ptr<f32>
+  %2 = memref.load %0[%1] {nontemporal = true} : memref<32xf32, affine_map<(d0) -> (d0)>>
+  func.return
+}
+
+// -----
+
+// CHECK-LABEL: func @store_non_temporal(
+func.func @store_non_temporal(%arg0 : memref<32xf32, affine_map<(d0) -> (d0)>>, %data : f32) {
+  %0 = memref.alloc() : memref<32xf32, affine_map<(d0) -> (d0)>>
+  %1 = arith.constant 7 : index
+  // CHECK: llvm.store %{{.*}}, %{{.*}}  {nontemporal} : !llvm.ptr<f32>
+  memref.store %data, %0[%1] {nontemporal = true} : memref<32xf32, affine_map<(d0) -> (d0)>>
+  func.return
+}