@@ -13,6 +13,7 @@ include "mlir/Dialect/Ptr/IR/PtrDialect.td"
1313include "mlir/Dialect/Ptr/IR/PtrAttrDefs.td"
1414include "mlir/Dialect/Ptr/IR/PtrEnums.td"
1515include "mlir/Dialect/Ptr/IR/MemorySpaceInterfaces.td"
16+ include "mlir/Interfaces/InferTypeOpInterface.td"
1617include "mlir/Interfaces/SideEffectInterfaces.td"
1718include "mlir/Interfaces/ViewLikeInterface.td"
1819include "mlir/IR/OpAsmInterface.td"
@@ -34,8 +35,15 @@ class Ptr_ShapedValueType<list<Type> allowedTypes, list<Pred> preds = []> :
3435 /*descr=*/[{A shaped type with value semantics and rank.}],
3536 /*cppType=*/"::mlir::ShapedType">;
3637
37- // A shaped pointer type with value semantics and rank.
38- class Ptr_ShapedPtrType : Ptr_ShapedValueType<[Ptr_PtrType], [HasRankPred]>;
38+ // A ptr-like type, either scalar or shaped type with value semantics.
39+ def Ptr_PtrLikeType :
40+ AnyTypeOf<[Ptr_ShapedValueType<[Ptr_PtrType], [HasRankPred]>, Ptr_PtrType]>;
41+
42+ // An int-like type, either scalar or shaped type with value semantics.
43+ def Ptr_IntLikeType :AnyTypeOf<[
44+ Ptr_ShapedValueType<[AnySignlessIntegerOrIndex], [HasRankPred]>,
45+ AnySignlessIntegerOrIndex
46+ ]>;
3947
4048// A shaped value type of rank 1 of any element type.
4149def Ptr_Any1DType :
@@ -167,41 +175,6 @@ def Ptr_GetMetadataOp : Pointer_Op<"get_metadata", [
167175 }];
168176}
169177
170- //===----------------------------------------------------------------------===//
171- // PtrAddOp
172- //===----------------------------------------------------------------------===//
173-
174- def Ptr_PtrAddOp : Pointer_Op<"ptr_add", [
175- Pure, AllTypesMatch<["base", "result"]>, ViewLikeOpInterface
176- ]> {
177- let summary = "Pointer add operation";
178- let description = [{
179- The `ptr_add` operation adds an integer offset to a pointer to produce a new
180- pointer. The input and output pointer types are always the same.
181-
182- Example:
183-
184- ```mlir
185- %x_off = ptr.ptr_add %x, %off : !ptr.ptr<#ptr.generic_space>, i32
186- %x_off0 = ptr.ptr_add nusw %x, %off : !ptr.ptr<#ptr.generic_space>, i32
187- ```
188- }];
189-
190- let arguments = (ins
191- Ptr_PtrType:$base,
192- AnySignlessIntegerOrIndex:$offset,
193- DefaultValuedProp<EnumProp<Ptr_PtrAddFlags>, "PtrAddFlags::none">:$flags);
194- let results = (outs Ptr_PtrType:$result);
195- let assemblyFormat = [{
196- ($flags^)? $base `,` $offset attr-dict `:` type($base) `,` type($offset)
197- }];
198- let hasFolder = 1;
199- let extraClassDeclaration = [{
200- /// `ViewLikeOp::getViewSource` method.
201- Value getViewSource() { return getBase(); }
202- }];
203- }
204-
205178//===----------------------------------------------------------------------===//
206179// LoadOp
207180//===----------------------------------------------------------------------===//
@@ -361,6 +334,62 @@ def Ptr_MaskedStoreOp : Pointer_Op<"masked_store", [
361334 let hasVerifier = 1;
362335}
363336
337+ //===----------------------------------------------------------------------===//
338+ // PtrAddOp
339+ //===----------------------------------------------------------------------===//
340+
341+ def Ptr_PtrAddOp : Pointer_Op<"ptr_add", [
342+ Pure, ViewLikeOpInterface,
343+ DeclareOpInterfaceMethods<InferTypeOpInterface>
344+ ]> {
345+ let summary = "Pointer add operation";
346+ let description = [{
347+ The `ptr_add` operation adds an int-like offset to one or more pointers to produce one or more new pointers.
348+
349+ The operation supports both scalar and shaped types with value semantics:
350+ - When both base and offset are scalar: produces a single new pointer
351+ - When base is shaped and offset is scalar: adds the same offset to each
352+ pointer in the base
353+ - When base is scalar and offset is shaped: adds the single pointer to each
354+ offset in the shaped value
355+ - When both are shaped: performs element-wise addition (shapes must be
356+ compatible)
357+
358+ Example:
359+
360+ ```mlir
361+ // Scalar base and offset
362+ %x_off = ptr.ptr_add %x, %off : !ptr.ptr<#ptr.generic_space>, i32
363+ %x_off0 = ptr.ptr_add nusw %x, %off : !ptr.ptr<#ptr.generic_space>, i32
364+
365+ // Shaped base with scalar offset
366+ %ptrs_off = ptr.ptr_add %ptrs, %off : vector<4x!ptr.ptr<#ptr.generic_space>>, i32
367+
368+ // Scalar base with shaped offset
369+ %x_offs = ptr.ptr_add %x, %offs : !ptr.ptr<#ptr.generic_space>, vector<4xi32>
370+
371+ // Both base and offset are shaped
372+ %ptrs_offs = ptr.ptr_add %ptrs, %offs : vector<4x!ptr.ptr<#ptr.generic_space>>, vector<4xi32>
373+ ```
374+ }];
375+ let arguments = (ins
376+ Ptr_PtrLikeType:$base,
377+ Ptr_IntLikeType:$offset,
378+ DefaultValuedProp<EnumProp<Ptr_PtrAddFlags>, "PtrAddFlags::none">:$flags);
379+ let results = (outs Ptr_PtrLikeType:$result);
380+ let assemblyFormat = [{
381+ ($flags^)? $base `,` $offset attr-dict `:` type($base) `,` type($offset)
382+ }];
383+ let hasFolder = 1;
384+ let extraClassDeclaration = [{
385+ /// `ViewLikeOp::getViewSource` method.
386+ Value getViewSource() { return getBase(); }
387+
388+ /// Returns the ptr type of the operation.
389+ ptr::PtrType getPtrType();
390+ }];
391+ }
392+
364393//===----------------------------------------------------------------------===//
365394// ScatterOp
366395//===----------------------------------------------------------------------===//
0 commit comments