[Mlir-commits] [mlir] [mlir][ptr] Add nontemporal field to ptr.masked_load (PR #204987)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: lonely eagle (linuxlonelyeagle)

<details>
<summary>Changes</summary>

`llvm.intr.masked.load` supports a `nontemporal` flag, but `ptr.masked_load` was missing this field, limiting its expressiveness relative to the underlying LLVM intrinsic. This patch adds a `nontemporal` UnitProp to `ptr.masked_load`, mirroring the existing `nontemporal` support on `ptr.load`, and updates the assembly format accordingly using `oilist` to allow `alignment` and `nontemporal` to appear independently.

---
Full diff: https://github.com/llvm/llvm-project/pull/204987.diff


2 Files Affected:

- (modified) mlir/include/mlir/Dialect/Ptr/IR/PtrOps.td (+7-2) 
- (modified) mlir/test/Dialect/Ptr/ops.mlir (+2) 


``````````diff
diff --git a/mlir/include/mlir/Dialect/Ptr/IR/PtrOps.td b/mlir/include/mlir/Dialect/Ptr/IR/PtrOps.td
index e14f64330c294..1f8780cf81f3d 100644
--- a/mlir/include/mlir/Dialect/Ptr/IR/PtrOps.td
+++ b/mlir/include/mlir/Dialect/Ptr/IR/PtrOps.td
@@ -299,10 +299,15 @@ def Ptr_MaskedLoadOp : Pointer_Op<"masked_load", [
   let arguments = (ins Ptr_PtrType:$ptr,
                        Ptr_Mask1DType:$mask,
                        Ptr_Any1DType:$passthrough,
-                       AlignmentProp:$alignment);
+                       AlignmentProp:$alignment,
+                       UnitProp:$nontemporal);
   let results = (outs Ptr_Any1DType:$result);
   let assemblyFormat = [{
-    $ptr `,` $mask `,` $passthrough (`alignment` `=` $alignment^)?
+    $ptr `,` $mask `,` $passthrough
+    oilist(
+      `alignment` `=` $alignment |
+      `nontemporal` $nontemporal
+    )
     attr-dict `:` qualified(type($ptr)) `->` type($result)
   }];
   let builders = [
diff --git a/mlir/test/Dialect/Ptr/ops.mlir b/mlir/test/Dialect/Ptr/ops.mlir
index 0a906ad559e21..f07295a553a52 100644
--- a/mlir/test/Dialect/Ptr/ops.mlir
+++ b/mlir/test/Dialect/Ptr/ops.mlir
@@ -89,6 +89,7 @@ func.func @scatter_ops_tensor(%value: tensor<8xi64>, %ptrs: tensor<8x!ptr.ptr<#p
 func.func @masked_load_ops(%ptr: !ptr.ptr<#ptr.generic_space>, %mask: vector<4xi1>, %passthrough: vector<4xf32>) -> vector<4xf32> {
   %0 = ptr.masked_load %ptr, %mask, %passthrough : !ptr.ptr<#ptr.generic_space> -> vector<4xf32>
   %1 = ptr.masked_load %ptr, %mask, %passthrough alignment = 16 : !ptr.ptr<#ptr.generic_space> -> vector<4xf32>
+  %2 = ptr.masked_load %ptr, %mask, %passthrough alignment = 16 nontemporal : !ptr.ptr<#ptr.generic_space> -> vector<4xf32> 
   return %0 : vector<4xf32>
 }
 
@@ -96,6 +97,7 @@ func.func @masked_load_ops(%ptr: !ptr.ptr<#ptr.generic_space>, %mask: vector<4xi
 func.func @masked_load_ops_tensor(%ptr: !ptr.ptr<#ptr.generic_space>, %mask: tensor<8xi1>, %passthrough: tensor<8xi32>) -> tensor<8xi32> {
   %0 = ptr.masked_load %ptr, %mask, %passthrough : !ptr.ptr<#ptr.generic_space> -> tensor<8xi32>
   %1 = ptr.masked_load %ptr, %mask, %passthrough alignment = 4 : !ptr.ptr<#ptr.generic_space> -> tensor<8xi32>
+  %2 = ptr.masked_load %ptr, %mask, %passthrough alignment = 4 nontemporal : !ptr.ptr<#ptr.generic_space> -> tensor<8xi32>
   return %0 : tensor<8xi32>
 }
 

``````````

</details>


https://github.com/llvm/llvm-project/pull/204987