[mlir][ptr] Add the `ptradd` and `type_offset` ops, and `generic_space` attr #136434

fabianmcg · 2025-04-19T14:31:28Z

This patch adds the ptr.ptradd and ptr.type_offset operations. Given a ptr value these operations can be used to compute new addresses. For example:

func.func @ops0(%ptr: !ptr.ptr<#ptr.int_space>) -> !ptr.ptr<#ptr.int_space> {
  %off = ptr.type_offset f32 : index
  %res = ptr.ptradd %ptr, %off : !ptr.ptr<#ptr.int_space>, index
  return %res : !ptr.ptr<#ptr.int_space>
}

Additionally, this patch also adds the #ptr.generic_space. This memory space allows loading and storing values to all types.

llvmbot · 2025-04-19T14:32:46Z

@llvm/pr-subscribers-mlir

Author: Fabian Mora (fabianmcg)

Changes

This patch adds the ptr.ptradd and ptr.type_offset operations. Given a ptr value these operations can be used to compute new addresses. For example:

func.func @<!-- -->ops0(%ptr: !ptr.ptr&lt;#ptr.int_space&gt;) -&gt; !ptr.ptr&lt;#ptr.int_space&gt; {
  %off = ptr.type_offset f32 : index
  %res = ptr.ptradd %ptr, %off : !ptr.ptr&lt;#ptr.int_space&gt;, index
  return %res : !ptr.ptr&lt;#ptr.int_space&gt;
}

Additionally, this patch also adds the #ptr.int_space, a memory space that is identified by a integer. This memory space allows loading and storing values to all types.

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

8 Files Affected:

(modified) mlir/include/mlir/Dialect/Ptr/IR/PtrAttrDefs.td (+33)
(modified) mlir/include/mlir/Dialect/Ptr/IR/PtrAttrs.h (+2)
(modified) mlir/include/mlir/Dialect/Ptr/IR/PtrOps.h (+1)
(modified) mlir/include/mlir/Dialect/Ptr/IR/PtrOps.td (+63)
(modified) mlir/lib/Dialect/Ptr/IR/PtrAttrs.cpp (+46)
(modified) mlir/lib/Dialect/Ptr/IR/PtrDialect.cpp (+31)
(added) mlir/test/Dialect/Ptr/canonicalize.mlir (+31)
(added) mlir/test/Dialect/Ptr/ops.mlir (+11)

diff --git a/mlir/include/mlir/Dialect/Ptr/IR/PtrAttrDefs.td b/mlir/include/mlir/Dialect/Ptr/IR/PtrAttrDefs.td
index e75038f300f1a..ac78ee37ab982 100644
--- a/mlir/include/mlir/Dialect/Ptr/IR/PtrAttrDefs.td
+++ b/mlir/include/mlir/Dialect/Ptr/IR/PtrAttrDefs.td
@@ -10,6 +10,7 @@
 #define PTR_ATTRDEFS
 
 include "mlir/Dialect/Ptr/IR/PtrDialect.td"
+include "mlir/Dialect/Ptr/IR/MemorySpaceInterfaces.td"
 include "mlir/IR/AttrTypeBase.td"
 
 // All of the attributes will extend this class.
@@ -20,6 +21,38 @@ class Ptr_Attr<string name, string attrMnemonic,
   let mnemonic = attrMnemonic;
 }
 
+//===----------------------------------------------------------------------===//
+// IAddressSpaceAttr
+//===----------------------------------------------------------------------===//
+
+def Ptr_IAddressSpaceAttr :
+    Ptr_Attr<"IAddressSpace", "int_space", [
+      DeclareAttrInterfaceMethods<MemorySpaceAttrInterface>
+    ]> {
+  let summary = "Int memory space";
+  let description = [{
+    The `int_as` attribute defines a memory space attribute with the following
+    properties:
+    - Load and store operations are always valid, regardless of the type.
+    - Atomic operations are always valid, regardless of the type.
+    - Cast operations are valid between pointers with `int_space` memory space,
+    or between non-scalable `vector` of pointers with `int_space` memory space.
+
+    The default address spaces is 0.
+
+    Example:
+
+    ```mlir
+    // Default address space: 0.
+    #ptr.int_space
+    // Address space 3.
+    #ptr.int_space<3>
+    ```
+  }];
+  let parameters = (ins DefaultValuedParameter<"int64_t", "0">:$value);
+  let assemblyFormat = "(`<` $value^ `>`)?";
+}
+
 //===----------------------------------------------------------------------===//
 // SpecAttr
 //===----------------------------------------------------------------------===//
diff --git a/mlir/include/mlir/Dialect/Ptr/IR/PtrAttrs.h b/mlir/include/mlir/Dialect/Ptr/IR/PtrAttrs.h
index 5ffe23e45fe12..a25885be9915b 100644
--- a/mlir/include/mlir/Dialect/Ptr/IR/PtrAttrs.h
+++ b/mlir/include/mlir/Dialect/Ptr/IR/PtrAttrs.h
@@ -15,6 +15,8 @@
 
 #include "mlir/IR/OpImplementation.h"
 
+#include "mlir/Dialect/Ptr/IR/MemorySpaceInterfaces.h"
+
 #define GET_ATTRDEF_CLASSES
 #include "mlir/Dialect/Ptr/IR/PtrOpsAttrs.h.inc"
 
diff --git a/mlir/include/mlir/Dialect/Ptr/IR/PtrOps.h b/mlir/include/mlir/Dialect/Ptr/IR/PtrOps.h
index 6a0c1429c6be9..43802a07351e5 100644
--- a/mlir/include/mlir/Dialect/Ptr/IR/PtrOps.h
+++ b/mlir/include/mlir/Dialect/Ptr/IR/PtrOps.h
@@ -18,6 +18,7 @@
 #include "mlir/Dialect/Ptr/IR/PtrDialect.h"
 #include "mlir/Dialect/Ptr/IR/PtrTypes.h"
 #include "mlir/IR/OpDefinition.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Ptr/IR/PtrOps.h.inc"
diff --git a/mlir/include/mlir/Dialect/Ptr/IR/PtrOps.td b/mlir/include/mlir/Dialect/Ptr/IR/PtrOps.td
index 02ea71f4322ef..c392765837922 100644
--- a/mlir/include/mlir/Dialect/Ptr/IR/PtrOps.td
+++ b/mlir/include/mlir/Dialect/Ptr/IR/PtrOps.td
@@ -12,6 +12,69 @@
 include "mlir/Dialect/Ptr/IR/PtrDialect.td"
 include "mlir/Dialect/Ptr/IR/PtrAttrDefs.td"
 include "mlir/Dialect/Ptr/IR/MemorySpaceInterfaces.td"
+include "mlir/Interfaces/SideEffectInterfaces.td"
 include "mlir/IR/OpAsmInterface.td"
 
+//===----------------------------------------------------------------------===//
+// PtrAddOp
+//===----------------------------------------------------------------------===//
+
+def Ptr_PtrAddOp : Pointer_Op<"ptradd", [
+    Pure, AllTypesMatch<["base", "result"]>
+  ]> {
+  let summary = "Pointer add operation";
+  let description = [{
+    The `ptradd` operation adds an integer offset to a pointer to produce a new
+    pointer. The input and output pointer types are always the same.
+
+    Example:
+
+    ```mlir
+    %x_off = ptr.ptradd %x, %off : !ptr.ptr<0>, i32
+    ```
+  }];
+
+  let arguments = (ins Ptr_PtrType:$base, AnySignlessIntegerOrIndex:$offset);
+  let results = (outs Ptr_PtrType:$result);
+  let assemblyFormat = [{
+    $base `,` $offset attr-dict `:` type($base) `,` type($offset)
+  }];
+  let hasFolder = 1;
+}
+
+//===----------------------------------------------------------------------===//
+// TypeOffsetOp
+//===----------------------------------------------------------------------===//
+
+def Ptr_TypeOffsetOp : Pointer_Op<"type_offset", [ConstantLike, Pure]> {
+  let summary = "Type offset operation";
+  let description = [{
+    The `type_offset` operation produces an int or index-typed SSA value
+    equal to a target-specific constant representing the offset of a single
+    element of the given type.
+
+    Example:
+
+    ```mlir
+    %0 = ptr.type_offset f32 : index
+    %1 = ptr.type_offset memref<12 x f64> : i32
+    ```
+  }];
+
+  let arguments = (ins TypeAttr:$element_type);
+  let results = (outs AnySignlessIntegerOrIndex:$result);
+  let builders = [
+    OpBuilder<(ins "TypeAttr":$element_type)>
+  ];
+  let assemblyFormat = [{
+    $element_type attr-dict `:` type($result)
+  }];
+  let extraClassDeclaration = [{
+    /// Returns the type offset according to `maybeLayout`. If `maybeLayout` is
+    /// `nullopt` the nearest layout the op will be used for the computation.
+    llvm::TypeSize getTypeSize(std::optional<DataLayout> layout = std::nullopt);
+  }];
+  let hasFolder = 1;
+}
+
 #endif // PTR_OPS
diff --git a/mlir/lib/Dialect/Ptr/IR/PtrAttrs.cpp b/mlir/lib/Dialect/Ptr/IR/PtrAttrs.cpp
index f8ce820d0bcbd..d1f778e72395e 100644
--- a/mlir/lib/Dialect/Ptr/IR/PtrAttrs.cpp
+++ b/mlir/lib/Dialect/Ptr/IR/PtrAttrs.cpp
@@ -11,6 +11,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Ptr/IR/PtrAttrs.h"
+#include "mlir/IR/BuiltinTypes.h"
 #include "llvm/ADT/TypeSwitch.h"
 
 using namespace mlir;
@@ -18,6 +19,51 @@ using namespace mlir::ptr;
 
 constexpr const static unsigned kBitsInByte = 8;
 
+//===----------------------------------------------------------------------===//
+// IAddressSpaceAttr
+//===----------------------------------------------------------------------===//
+
+LogicalResult IAddressSpaceAttr::isValidLoad(
+    Type type, ptr::AtomicOrdering ordering, IntegerAttr alignment,
+    function_ref<InFlightDiagnostic()> emitError) const {
+  return success();
+}
+
+LogicalResult IAddressSpaceAttr::isValidStore(
+    Type type, ptr::AtomicOrdering ordering, IntegerAttr alignment,
+    function_ref<InFlightDiagnostic()> emitError) const {
+  return success();
+}
+
+LogicalResult IAddressSpaceAttr::isValidAtomicOp(
+    ptr::AtomicBinOp op, Type type, ptr::AtomicOrdering ordering,
+    IntegerAttr alignment, function_ref<InFlightDiagnostic()> emitError) const {
+  return success();
+}
+
+LogicalResult IAddressSpaceAttr::isValidAtomicXchg(
+    Type type, ptr::AtomicOrdering successOrdering,
+    ptr::AtomicOrdering failureOrdering, IntegerAttr alignment,
+    function_ref<InFlightDiagnostic()> emitError) const {
+  return success();
+}
+
+LogicalResult IAddressSpaceAttr::isValidAddrSpaceCast(
+    Type tgt, Type src, function_ref<InFlightDiagnostic()> emitError) const {
+  // TODO: update this method once the `addrspace_cast` op is added to the
+  // dialect.
+  assert(false && "unimplemented, see TODO in the source.");
+  return failure();
+}
+
+LogicalResult IAddressSpaceAttr::isValidPtrIntCast(
+    Type intLikeTy, Type ptrLikeTy,
+    function_ref<InFlightDiagnostic()> emitError) const {
+  // TODO: update this method once the int-cast ops are added to the dialect.
+  assert(false && "unimplemented, see TODO in the source.");
+  return failure();
+}
+
 //===----------------------------------------------------------------------===//
 // SpecAttr
 //===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/Ptr/IR/PtrDialect.cpp b/mlir/lib/Dialect/Ptr/IR/PtrDialect.cpp
index ff231dae60c27..e87496cd1696c 100644
--- a/mlir/lib/Dialect/Ptr/IR/PtrDialect.cpp
+++ b/mlir/lib/Dialect/Ptr/IR/PtrDialect.cpp
@@ -12,7 +12,9 @@
 
 #include "mlir/Dialect/Ptr/IR/PtrOps.h"
 #include "mlir/IR/DialectImplementation.h"
+#include "mlir/IR/Matchers.h"
 #include "mlir/IR/PatternMatch.h"
+#include "mlir/Interfaces/DataLayoutInterfaces.h"
 #include "mlir/Transforms/InliningUtils.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/TypeSwitch.h"
@@ -39,6 +41,35 @@ void PtrDialect::initialize() {
       >();
 }
 
+//===----------------------------------------------------------------------===//
+// PtrAddOp
+//===----------------------------------------------------------------------===//
+
+/// Fold the op to the base ptr when the offset is 0.
+OpFoldResult PtrAddOp::fold(FoldAdaptor adaptor) {
+  Attribute attr = adaptor.getOffset();
+  if (!attr)
+    return nullptr;
+  if (llvm::APInt value; m_ConstantInt(&value).match(attr) && value.isZero())
+    return getBase();
+  return nullptr;
+}
+
+//===----------------------------------------------------------------------===//
+// TypeOffsetOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult TypeOffsetOp::fold(FoldAdaptor adaptor) {
+  return TypeAttr::get(getElementType());
+}
+
+llvm::TypeSize TypeOffsetOp::getTypeSize(std::optional<DataLayout> layout) {
+  if (layout)
+    return layout->getTypeSize(getElementType());
+  DataLayout dl = DataLayout::closest(*this);
+  return dl.getTypeSize(getElementType());
+}
+
 //===----------------------------------------------------------------------===//
 // Pointer API.
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/Ptr/canonicalize.mlir b/mlir/test/Dialect/Ptr/canonicalize.mlir
new file mode 100644
index 0000000000000..630040cd42442
--- /dev/null
+++ b/mlir/test/Dialect/Ptr/canonicalize.mlir
@@ -0,0 +1,31 @@
+// RUN: mlir-opt --canonicalize %s | FileCheck %s
+
+/// Check `ptradd` and `type_offset` canonicalizer patterns.
+
+// CHECK-LABEL: @ops0
+func.func @ops0(%ptr: !ptr.ptr<#ptr.int_space<3>>, %c: i1) -> !ptr.ptr<#ptr.int_space<3>> {
+  // CHECK: (%[[PTR_0:.*]]: !ptr.ptr<#ptr.int_space<3>>,
+  // CHECK: %[[F32_OFF:.*]] = ptr.type_offset f32 : index
+  // CHECK: %[[PTR_1:.*]] = ptr.ptradd %[[PTR_0]], %[[F32_OFF]] : <#ptr.int_space<3>>, index
+  // CHECK: %[[PTR_2:.*]] = ptr.ptradd %[[PTR_1]], %[[F32_OFF]] : <#ptr.int_space<3>>, index
+  // CHECK: %[[PTR_3:.*]] = scf.if %{{.*}} -> (!ptr.ptr<#ptr.int_space<3>>) {
+  // CHECK: %[[PTR_4:.*]] = ptr.ptradd %[[PTR_2]], %[[F32_OFF]] : <#ptr.int_space<3>>, index
+  // CHECK: scf.yield %[[PTR_4]] : !ptr.ptr<#ptr.int_space<3>>
+  // CHECK: } else {
+  // CHECK: scf.yield %[[PTR_0]] : !ptr.ptr<#ptr.int_space<3>>
+  // CHECK: }
+  // CHECK: return %[[PTR_3]] : !ptr.ptr<#ptr.int_space<3>>
+  // CHECK: }
+  %off0 = ptr.type_offset f32 : index
+  %res0 = ptr.ptradd %ptr, %off0 : !ptr.ptr<#ptr.int_space<3>>, index
+  %off1 = ptr.type_offset f32 : index
+  %res1 = ptr.ptradd %res0, %off1 : !ptr.ptr<#ptr.int_space<3>>, index
+  %res = scf.if %c -> !ptr.ptr<#ptr.int_space<3>> {
+    %off2 = ptr.type_offset f32 : index
+    %res2 = ptr.ptradd %res1, %off2 : !ptr.ptr<#ptr.int_space<3>>, index
+    scf.yield %res2 : !ptr.ptr<#ptr.int_space<3>>
+  } else {
+    scf.yield %ptr : !ptr.ptr<#ptr.int_space<3>>
+  }
+  return %res : !ptr.ptr<#ptr.int_space<3>>
+}
diff --git a/mlir/test/Dialect/Ptr/ops.mlir b/mlir/test/Dialect/Ptr/ops.mlir
new file mode 100644
index 0000000000000..53008cb84d6b0
--- /dev/null
+++ b/mlir/test/Dialect/Ptr/ops.mlir
@@ -0,0 +1,11 @@
+// RUN: mlir-opt %s | FileCheck %s
+
+/// Check op assembly.
+func.func @ops0(%ptr: !ptr.ptr<#ptr.int_space>) -> !ptr.ptr<#ptr.int_space> {
+  // CHECK-LABEL: @ops0
+  // CHECK: ptr.type_offset f32 : index
+  // CHECK-NEXT: ptr.ptradd %{{.*}}, %{{.*}} : <#ptr.int_space>, index
+  %off = ptr.type_offset f32 : index
+  %res = ptr.ptradd %ptr, %off : !ptr.ptr<#ptr.int_space>, index
+  return %res : !ptr.ptr<#ptr.int_space>
+}

github-actions · 2025-04-19T14:33:10Z

⚠️ We detected that you are using a GitHub private e-mail address to contribute to the repo.
Please turn off Keep my email addresses private setting in your account.
See LLVM Discourse for more information.

gysit

Nice!

mlir/test/Dialect/Ptr/ops.mlir

mlir/test/Dialect/Ptr/canonicalize.mlir

mlir/include/mlir/Dialect/Ptr/IR/PtrOps.td

@OPS0

This patch adds the `ptr.ptradd` and `ptr.type_offset` operations. Given a `ptr` value these operations can be used to compute new addresses. For example: ```mlir func.func @OPS0(%ptr: !ptr.ptr<#ptr.int_space>) -> !ptr.ptr<#ptr.int_space> { %off = ptr.type_offset f32 : index %res = ptr.ptradd %ptr, %off : !ptr.ptr<#ptr.int_space>, index return %res : !ptr.ptr<#ptr.int_space> } ``` Additionally, this patch also adds the `#ptr.int_space`, a memory space that is identified by a integer. This memory space allows loading and storing values to all types.

mlir/include/mlir/Dialect/Ptr/IR/PtrOps.td

mlir/lib/Dialect/Ptr/IR/PtrDialect.cpp

mlir/include/mlir/Dialect/Ptr/IR/PtrAttrDefs.td

mlir/test/Dialect/Ptr/canonicalize.mlir

Co-authored-by: Mehdi Amini <[email protected]>

mlir/include/mlir/Dialect/Ptr/IR/PtrOps.td

mlir/lib/Dialect/Ptr/IR/PtrDialect.cpp

mlir/include/mlir/Dialect/Ptr/IR/PtrEnums.td

gysit

Nice!

LGTM modulo last nits.

Please give others a chance to review before landing.

mlir/include/mlir/Dialect/Ptr/IR/PtrOps.td

mlir/test/Dialect/Ptr/canonicalize.mlir

Co-authored-by: Tobias Gysi <[email protected]>

…c name

@OPS0

…e` attr (llvm#136434) This patch adds the `ptr.ptradd` and `ptr.type_offset` operations. Given a `ptr` value these operations can be used to compute new addresses. For example: ```mlir func.func @OPS0(%ptr: !ptr.ptr<#ptr.int_space>) -> !ptr.ptr<#ptr.int_space> { %off = ptr.type_offset f32 : index %res = ptr.ptradd %ptr, %off : !ptr.ptr<#ptr.int_space>, index return %res : !ptr.ptr<#ptr.int_space> } ``` Additionally, this patch also adds the `#ptr.generic_space`. This memory space allows loading and storing values to all types. --------- Co-authored-by: Mehdi Amini <[email protected]> Co-authored-by: Tobias Gysi <[email protected]>

fabianmcg marked this pull request as ready for review April 19, 2025 14:32

llvmbot added the mlir label Apr 19, 2025

fabianmcg requested review from Dinistro, ftynse, gysit and joker-eph April 19, 2025 14:32

fabianmcg requested a review from rengolin April 19, 2025 14:33

gysit reviewed Apr 19, 2025

View reviewed changes

fabianmcg added 2 commits April 20, 2025 10:53

address reviewer comments

133e3ef

fabianmcg force-pushed the pr-ptradd branch from 2f54230 to 133e3ef Compare April 20, 2025 11:50