[mlir][ptr] Add ConstantOp with NullAttr and AddressAttr support (#157347)

This patch introduces the `ptr.constant` operation. It also adds the
`NullAttr` and `AddressAttr` for representing null pointers, and integer
raw addresses.

It also implements LLVM IR translation for `ptr.constant` with
`#ptr.null` or `#ptr.address` attributes.

Finally, it extends `FieldParser` to support APInt parsing.

Example:
```mlir
llvm.func @constant_address_op() ->
    !llvm.struct<(!ptr.ptr<#llvm.address_space<0>>,
                  !ptr.ptr<#llvm.address_space<1>>,
                  !ptr.ptr<#llvm.address_space<2>>)> {
  %0 = ptr.constant #ptr.null : !ptr.ptr<#llvm.address_space<0>>
  %1 = ptr.constant #ptr.address<0x1000> : !ptr.ptr<#llvm.address_space<1>>
  %2 = ptr.constant #ptr.address<3735928559> : !ptr.ptr<#llvm.address_space<2>>
  %3 = llvm.mlir.poison : !llvm.struct<(!ptr.ptr<#llvm.address_space<0>>, !ptr.ptr<#llvm.address_space<1>>, !ptr.ptr<#llvm.address_space<2>>)>
  %4 = llvm.insertvalue %0, %3[0] : !llvm.struct<(!ptr.ptr<#llvm.address_space<0>>, !ptr.ptr<#llvm.address_space<1>>, !ptr.ptr<#llvm.address_space<2>>)>
  %5 = llvm.insertvalue %1, %4[1] : !llvm.struct<(!ptr.ptr<#llvm.address_space<0>>, !ptr.ptr<#llvm.address_space<1>>, !ptr.ptr<#llvm.address_space<2>>)>
  %6 = llvm.insertvalue %2, %5[2] : !llvm.struct<(!ptr.ptr<#llvm.address_space<0>>, !ptr.ptr<#llvm.address_space<1>>, !ptr.ptr<#llvm.address_space<2>>)>
  llvm.return %6 : !llvm.struct<(!ptr.ptr<#llvm.address_space<0>>, !ptr.ptr<#llvm.address_space<1>>, !ptr.ptr<#llvm.address_space<2>>)>
}
```
Result of translation to LLVM IR:
```llvm
define { ptr, ptr addrspace(1), ptr addrspace(2) } @constant_address_op() {
  ret { ptr, ptr addrspace(1), ptr addrspace(2) } { ptr null, ptr addrspace(1) inttoptr (i64 4096 to ptr addrspace(1)), ptr addrspace(2) inttoptr (i64 3735928559 to ptr addrspace(2)) }
}
```

This patch also changes all the `convert*` occurrences in function names
or comments to `translate` in the PtrToLLVM file.

---------

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

Author: fabianmcg

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

@@ -22,6 +22,34 @@ class Ptr_Attr<string name, string attrMnemonic,
   let mnemonic = attrMnemonic;
 }
 
+//===----------------------------------------------------------------------===//
+// AddressAttr
+//===----------------------------------------------------------------------===//
+
+def Ptr_AddressAttr : Ptr_Attr<"Address", "address", [
+    DeclareAttrInterfaceMethods<TypedAttrInterface>
+  ]> {
+  let summary = "Address attribute";
+  let description = [{
+    The `address` attribute represents a raw memory address, expressed in bytes.
+
+    Example:
+
+    ```mlir
+      #ptr.address<0x1000> : !ptr.ptr<#ptr.generic_space>
+    ```
+  }];
+  let parameters = (ins AttributeSelfTypeParameter<"", "PtrType">:$type,
+                        APIntParameter<"">:$value);
+  let builders = [
+    AttrBuilderWithInferredContext<(ins "PtrType":$type,
+                                        "const llvm::APInt &":$value), [{
+      return $_get(type.getContext(), type, value);
+    }]>
+  ];
+  let assemblyFormat = "`<` $value `>`";
+}
+
 //===----------------------------------------------------------------------===//
 // GenericSpaceAttr
 //===----------------------------------------------------------------------===//
@@ -37,16 +65,42 @@ def Ptr_GenericSpaceAttr :
     - Load and store operations are always valid, regardless of the type.
     - Atomic operations are always valid, regardless of the type.
     - Cast operations to `generic_space` are always valid.
-  
+
     Example:
 
     ```mlir
-    #ptr.generic_space
+      #ptr.generic_space : !ptr.ptr<#ptr.generic_space>
     ```
   }];
   let assemblyFormat = "";
 }
 
+//===----------------------------------------------------------------------===//
+// NullAttr
+//===----------------------------------------------------------------------===//
+
+def Ptr_NullAttr : Ptr_Attr<"Null", "null", [
+    DeclareAttrInterfaceMethods<TypedAttrInterface>
+  ]> {
+  let summary = "Null pointer attribute";
+  let description = [{
+    The `null` attribute represents a null pointer.
+
+    Example:
+
+    ```mlir
+      #ptr.null
+    ```
+  }];
+  let parameters = (ins AttributeSelfTypeParameter<"", "PtrType">:$type);
+  let builders = [
+    AttrBuilderWithInferredContext<(ins "PtrType":$type), [{
+      return $_get(type.getContext(), type);
+    }]>
+  ];
+  let assemblyFormat = "";
+}
+
 //===----------------------------------------------------------------------===//
 // SpecAttr
 //===----------------------------------------------------------------------===//
@@ -62,7 +116,7 @@ def Ptr_SpecAttr : Ptr_Attr<"Spec", "spec"> {
      - [Optional] index: bitwidth that should be used when performing index
      computations for the type. Setting the field to `kOptionalSpecValue`, means
      the field is optional.
-    
+
     Furthermore, the attribute will verify that all present values are divisible
     by 8 (number of bits in a byte), and that `preferred` > `abi`.
 

mlir/include/mlir/Dialect/Ptr/IR/PtrAttrs.h

@@ -21,6 +21,12 @@
 #include "mlir/Dialect/Ptr/IR/MemorySpaceInterfaces.h"
 #include "mlir/Dialect/Ptr/IR/PtrEnums.h"
 
+namespace mlir {
+namespace ptr {
+class PtrType;
+} // namespace ptr
+} // namespace mlir
+
 #define GET_ATTRDEF_CLASSES
 #include "mlir/Dialect/Ptr/IR/PtrOpsAttrs.h.inc"
 

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

@@ -36,7 +36,7 @@ class Ptr_ShapedValueType<list<Type> allowedTypes, list<Pred> preds = []> :
     /*cppType=*/"::mlir::ShapedType">;
 
 // A ptr-like type, either scalar or shaped type with value semantics.
-def Ptr_PtrLikeType : 
+def Ptr_PtrLikeType :
   AnyTypeOf<[Ptr_ShapedValueType<[Ptr_PtrType], [HasRankPred]>, Ptr_PtrType]>;
 
 // An int-like type, either scalar or shaped type with value semantics.
@@ -57,6 +57,31 @@ def Ptr_Mask1DType :
 def Ptr_Ptr1DType :
   Ptr_ShapedValueType<[Ptr_PtrType], [HasAnyRankOfPred<[1]>]>;
 
+//===----------------------------------------------------------------------===//
+// ConstantOp
+//===----------------------------------------------------------------------===//
+
+def Ptr_ConstantOp : Pointer_Op<"constant", [
+    ConstantLike, Pure, AllTypesMatch<["value", "result"]>
+  ]> {
+  let summary = "Pointer constant operation";
+  let description = [{
+    The `constant` operation produces a pointer constant. The attribute must be
+    a typed attribute of pointer type.
+
+    Example:
+
+    ```mlir
+    // Create a null pointer
+    %null = ptr.constant #ptr.null : !ptr.ptr<#ptr.generic_space>
+    ```
+  }];
+  let arguments = (ins TypedAttrInterface:$value);
+  let results = (outs Ptr_PtrType:$result);
+  let assemblyFormat = "attr-dict $value";
+  let hasFolder = 1;
+}
+
 //===----------------------------------------------------------------------===//
 // FromPtrOp
 //===----------------------------------------------------------------------===//
@@ -81,7 +106,7 @@ def Ptr_FromPtrOp : Pointer_Op<"from_ptr", [
     ```mlir
     %typed_ptr = ptr.from_ptr %ptr : !ptr.ptr<#ptr.generic_space> -> !my.ptr<f32, #ptr.generic_space>
     %memref = ptr.from_ptr %ptr metadata %md : !ptr.ptr<#ptr.generic_space> -> memref<f32, #ptr.generic_space>
-  
+
     // Cast the `%ptr` to a memref without utilizing metadata.
     %memref = ptr.from_ptr %ptr : !ptr.ptr<#ptr.generic_space> -> memref<f32, #ptr.generic_space>
     ```
@@ -361,13 +386,13 @@ def Ptr_PtrAddOp : Pointer_Op<"ptr_add", [
     // Scalar base and offset
     %x_off  = ptr.ptr_add %x, %off : !ptr.ptr<#ptr.generic_space>, i32
     %x_off0 = ptr.ptr_add nusw %x, %off : !ptr.ptr<#ptr.generic_space>, i32
-    
+
     // Shaped base with scalar offset
     %ptrs_off = ptr.ptr_add %ptrs, %off : vector<4x!ptr.ptr<#ptr.generic_space>>, i32
-    
+
     // Scalar base with shaped offset
     %x_offs = ptr.ptr_add %x, %offs : !ptr.ptr<#ptr.generic_space>, vector<4xi32>
-    
+
     // Both base and offset are shaped
     %ptrs_offs = ptr.ptr_add %ptrs, %offs : vector<4x!ptr.ptr<#ptr.generic_space>>, vector<4xi32>
     ```
@@ -382,7 +407,7 @@ def Ptr_PtrAddOp : Pointer_Op<"ptr_add", [
   }];
   let hasFolder = 1;
   let extraClassDeclaration = [{
-    /// `ViewLikeOp::getViewSource` method. 
+    /// `ViewLikeOp::getViewSource` method.
     Value getViewSource() { return getBase(); }
 
     /// Returns the ptr type of the operation.
@@ -418,7 +443,7 @@ def Ptr_ScatterOp : Pointer_Op<"scatter", [
     // Scatter values to multiple memory locations
     ptr.scatter %value, %ptrs, %mask :
       vector<4xf32>, vector<4x!ptr.ptr<#ptr.generic_space>>
-    
+
     // Scatter with alignment
     ptr.scatter %value, %ptrs, %mask alignment = 8 :
       vector<4xf32>, vector<4x!ptr.ptr<#ptr.generic_space>>

mlir/include/mlir/IR/DialectImplementation.h

@@ -103,10 +103,11 @@ struct FieldParser<
 
 /// Parse any integer.
 template <typename IntT>
-struct FieldParser<IntT,
-                   std::enable_if_t<std::is_integral<IntT>::value, IntT>> {
+struct FieldParser<IntT, std::enable_if_t<(std::is_integral<IntT>::value ||
+                                           std::is_same_v<IntT, llvm::APInt>),
+                                          IntT>> {
   static FailureOr<IntT> parse(AsmParser &parser) {
-    IntT value = 0;
+    IntT value{};
     if (parser.parseInteger(value))
       return failure();
     return value;

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

@@ -56,6 +56,12 @@ verifyAlignment(std::optional<int64_t> alignment,
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// ConstantOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult ConstantOp::fold(FoldAdaptor adaptor) { return getValue(); }
+
 //===----------------------------------------------------------------------===//
 // FromPtrOp
 //===----------------------------------------------------------------------===//