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

[lonely eagle]

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

>From 77cbb815a5a662d72539dd2d65ac872ac08aecc4 Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Sun, 21 Jun 2026 15:40:33 +0000
Subject: [PATCH 1/2] add nontemporal to ptr.masked_load op.

---
 mlir/include/mlir/Dialect/Ptr/IR/PtrOps.td | 9 +++++++--
 mlir/test/Dialect/Ptr/ops.mlir             | 2 ++
 2 files changed, 9 insertions(+), 2 deletions(-)

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>
 }
 

>From 7178c60be3b12c45ccd4b58f458f1957509f43b0 Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Tue, 23 Jun 2026 02:48:44 +0000
Subject: [PATCH 2/2] update code.

---
 .../Target/LLVMIR/Dialect/Ptr/PtrToLLVMIRTranslation.cpp   | 7 +++++++
 mlir/test/Target/LLVMIR/ptr.mlir                           | 3 +++
 2 files changed, 10 insertions(+)

diff --git a/mlir/lib/Target/LLVMIR/Dialect/Ptr/PtrToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/Ptr/PtrToLLVMIRTranslation.cpp
index ca3301aa509a5..2d8a6bf864dd2 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/Ptr/PtrToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/Ptr/PtrToLLVMIRTranslation.cpp
@@ -259,6 +259,13 @@ translateMaskedLoadOp(MaskedLoadOp maskedLoadOp, llvm::IRBuilderBase &builder,
   llvm::Value *result = builder.CreateMaskedLoad(
       resultType, ptr, alignment.valueOrOne(), mask, passthrough);
 
+  if (maskedLoadOp.getNontemporal()) {
+    llvm::MDNode *metadata =
+        llvm::MDNode::get(result->getContext(),
+                          llvm::ConstantAsMetadata::get(builder.getInt32(1)));
+    llvm::cast<llvm::Instruction>(result)->setMetadata(
+        llvm::LLVMContext::MD_nontemporal, metadata);
+  }
   moduleTranslation.mapValue(maskedLoadOp.getResult(), result);
   return success();
 }
diff --git a/mlir/test/Target/LLVMIR/ptr.mlir b/mlir/test/Target/LLVMIR/ptr.mlir
index 94b6628772634..4fb1ca0f7d58a 100644
--- a/mlir/test/Target/LLVMIR/ptr.mlir
+++ b/mlir/test/Target/LLVMIR/ptr.mlir
@@ -118,6 +118,7 @@ llvm.func @gather_ops_i32(%ptrs: vector<8x!ptr.ptr<#llvm.address_space<0>>>, %ma
 // CHECK-SAME: (ptr %[[PTR:.*]], <4 x i1> %[[MASK:.*]], <4 x float> %[[PASSTHROUGH:.*]]) {
 // CHECK-NEXT:   %[[V0:.*]] = call <4 x float> @llvm.masked.load.v4f32.p0(ptr align 1 %[[PTR]], <4 x i1> %[[MASK]], <4 x float> %[[PASSTHROUGH]])
 // CHECK-NEXT:   %[[V1:.*]] = call <4 x float> @llvm.masked.load.v4f32.p0(ptr align 16 %[[PTR]], <4 x i1> %[[MASK]], <4 x float> %[[PASSTHROUGH]])
+// CHECK-NEXT:   %[[V2:.*]] = call <4 x float> @llvm.masked.load.v4f32.p0(ptr align 16 %[[PTR]], <4 x i1> %[[MASK]], <4 x float> %[[PASSTHROUGH]]), !nontemporal
 // CHECK-NEXT:   ret <4 x float> %[[V0]]
 // CHECK-NEXT: }
 llvm.func @masked_load_ops(%ptr: !ptr.ptr<#llvm.address_space<0>>, %mask: vector<4xi1>, %passthrough: vector<4xf32>) -> vector<4xf32> {
@@ -125,6 +126,8 @@ llvm.func @masked_load_ops(%ptr: !ptr.ptr<#llvm.address_space<0>>, %mask: vector
   %0 = ptr.masked_load %ptr, %mask, %passthrough : !ptr.ptr<#llvm.address_space<0>> -> vector<4xf32>
   // Masked load with alignment
   %1 = ptr.masked_load %ptr, %mask, %passthrough alignment = 16 : !ptr.ptr<#llvm.address_space<0>> -> vector<4xf32>
+  // Masked load with nontemporal
+  %2 = ptr.masked_load %ptr, %mask, %passthrough alignment = 16 nontemporal : !ptr.ptr<#llvm.address_space<0>> -> vector<4xf32>
   llvm.return %0 : vector<4xf32>
 }