address reviewer comments

fabianmcg · fabianmcg · commit 5042d6f67ada · 2024-04-20T11:16:58.000Z
diff --git a/mlir/include/mlir/Dialect/Ptr/IR/MemorySpace.h b/mlir/include/mlir/Dialect/Ptr/IR/MemorySpace.h
@@ -21,14 +21,14 @@
 namespace mlir {
 class Operation;
 namespace ptr {
-/// This method checks if it's valid to perform an `addrspacecast` op in the
+/// Checks if it's valid to perform an `addrspacecast` op in the
 /// memory space.
 /// Compatible types are:
 /// Vectors of rank 1, or scalars of `ptr` type.
 LogicalResult isValidAddrSpaceCastImpl(Type tgt, Type src,
                                        Operation *diagnosticOp);
 
-/// This method checks if it's valid to perform a `ptrtoint` or `inttoptr` op in
+/// Checks if it's valid to perform a `ptrtoint` or `inttoptr` op in
 /// the memory space.
 /// Compatible types are:
 /// IntLikeTy: Vectors of rank 1, or scalars of integer types or `index` type.
@@ -52,28 +52,29 @@ class MemorySpace {
   MemorySpace() = default;
   MemorySpace(std::nullptr_t) {}
   MemorySpace(MemorySpaceAttrInterface memorySpace)
-      : memorySpaceAttr(memorySpace), memorySpace(memorySpace) {}
-  MemorySpace(Attribute memorySpace)
-      : memorySpaceAttr(memorySpace),
+      : underlyingMemorySpace(memorySpace), memorySpace(memorySpace) {}
+  explicit MemorySpace(Attribute memorySpace)
+      : underlyingMemorySpace(memorySpace),
         memorySpace(dyn_cast_or_null<MemorySpaceAttrInterface>(memorySpace)) {}
 
-  operator Attribute() const { return memorySpaceAttr; }
+  operator Attribute() const { return underlyingMemorySpace; }
   operator MemorySpaceAttrInterface() const { return memorySpace; }
   bool operator==(const MemorySpace &memSpace) const {
-    return memSpace.memorySpaceAttr == memorySpaceAttr;
+    return memSpace.underlyingMemorySpace == underlyingMemorySpace;
   }
 
   /// Returns the underlying memory space.
-  Attribute getUnderlyingSpace() const { return memorySpaceAttr; }
+  Attribute getUnderlyingSpace() const { return underlyingMemorySpace; }
 
-  /// Returns true if the underlying memory space is null.
+  /// Returns true if the memory space is null.
   bool isDefaultModel() const { return memorySpace == nullptr; }
 
   /// Returns the memory space as an integer, or 0 if using the default space.
   unsigned getAddressSpace() const {
     if (memorySpace)
       return memorySpace.getAddressSpace();
-    if (auto intAttr = llvm::dyn_cast_or_null<IntegerAttr>(memorySpaceAttr))
+    if (auto intAttr =
+            llvm::dyn_cast_or_null<IntegerAttr>(underlyingMemorySpace))
       return intAttr.getInt();
     return 0;
   }
@@ -84,9 +85,9 @@ class MemorySpace {
     return memorySpace ? memorySpace.getDefaultMemorySpace() : nullptr;
   }
 
-  /// This method checks if it's valid to load a value from the memory space
-  /// with a specific type, alignment, and atomic ordering. The default model
-  /// assumes all values are loadable.
+  /// Checks if it's valid to load a value from the memory space with a specific
+  /// type, alignment, and atomic ordering. The default model assumes all values
+  /// can be loaded.
   LogicalResult isValidLoad(Type type, AtomicOrdering ordering,
                             IntegerAttr alignment,
                             Operation *diagnosticOp = nullptr) const {
@@ -95,9 +96,9 @@ class MemorySpace {
                        : success();
   }
 
-  /// This method checks if it's valid to store a value in the memory space with
-  /// a specific type, alignment, and atomic ordering. The default model assumes
-  /// all values are loadable.
+  /// Checks if it's valid to store a value in the memory space with a specific
+  /// type, alignment, and atomic ordering. The default model assumes all values
+  /// can be stored.
   LogicalResult isValidStore(Type type, AtomicOrdering ordering,
                              IntegerAttr alignment,
                              Operation *diagnosticOp = nullptr) const {
@@ -106,8 +107,8 @@ class MemorySpace {
                        : success();
   }
 
-  /// This method checks if it's valid to perform an atomic operation in the
-  /// memory space with a specific type, alignment, and atomic ordering.
+  /// Checks if it's valid to perform an atomic operation in the memory space
+  /// with a specific type, alignment, and atomic ordering.
   LogicalResult isValidAtomicOp(AtomicBinOp op, Type type,
                                 AtomicOrdering ordering, IntegerAttr alignment,
                                 Operation *diagnosticOp = nullptr) const {
@@ -116,8 +117,8 @@ class MemorySpace {
                        : success();
   }
 
-  /// This method checks if it's valid to perform an atomic operation in the
-  /// memory space with a specific type, alignment, and atomic ordering.
+  /// Checks if it's valid to perform an atomic exchange operation in the memory
+  /// space with a specific type, alignment, and atomic ordering.
   LogicalResult isValidAtomicXchg(Type type, AtomicOrdering successOrdering,
                                   AtomicOrdering failureOrdering,
                                   IntegerAttr alignment,
@@ -128,17 +129,16 @@ class MemorySpace {
                        : success();
   }
 
-  /// This method checks if it's valid to perform an `addrspacecast` op in the
-  /// memory space.
+  /// Checks if it's valid to perform an `addrspacecast` op in the memory space.
   LogicalResult isValidAddrSpaceCast(Type tgt, Type src,
                                      Operation *diagnosticOp = nullptr) const {
     return memorySpace
                ? memorySpace.isValidAddrSpaceCast(tgt, src, diagnosticOp)
                : isValidAddrSpaceCastImpl(tgt, src, diagnosticOp);
   }
 
-  /// This method checks if it's valid to perform a `ptrtoint` or `inttoptr` op
-  /// in the memory space.
+  /// Checks if it's valid to perform a `ptrtoint` or `inttoptr` op in the
+  /// memory space.
   LogicalResult isValidPtrIntCast(Type intLikeTy, Type ptrLikeTy,
                                   Operation *diagnosticOp = nullptr) const {
     return memorySpace
@@ -149,7 +149,7 @@ class MemorySpace {
 
 protected:
   /// Underlying memory space.
-  Attribute memorySpaceAttr{};
+  Attribute underlyingMemorySpace{};
   /// Memory space.
   MemorySpaceAttrInterface memorySpace{};
 };
diff --git a/mlir/lib/Dialect/Ptr/IR/PtrDialect.cpp b/mlir/lib/Dialect/Ptr/IR/PtrDialect.cpp
@@ -39,6 +39,10 @@ void PtrDialect::initialize() {
 // Pointer API.
 //===----------------------------------------------------------------------===//
 
+// Error constants for vector data types.
+constexpr const static unsigned kInvalidRankError = -1;
+constexpr const static unsigned kScalableDimsError = -2;
+
 // Returns a pair containing:
 // The underlying type of a vector or the type itself if it's not a vector.
 // The number of elements in the vector or an error code if the type is not
@@ -49,26 +53,28 @@ static std::pair<Type, int64_t> getVecOrScalarInfo(Type ty) {
     // Vectors of rank greater than one or with scalable dimensions are not
     // supported.
     if (vecTy.getRank() != 1)
-      return {elemTy, -1};
+      return {elemTy, kInvalidRankError};
     else if (vecTy.getScalableDims()[0])
-      return {elemTy, -2};
+      return {elemTy, kScalableDimsError};
     return {elemTy, vecTy.getShape()[0]};
   }
   // `ty` is a scalar type.
   return {ty, 0};
 }
 
-LogicalResult mlir::ptr::isValidAddrSpaceCastImpl(Type tgt, Type src,
-                                                  Operation *op) {
-  std::pair<Type, int64_t> tgtInfo = getVecOrScalarInfo(tgt);
-  std::pair<Type, int64_t> srcInfo = getVecOrScalarInfo(src);
-  if (!isa<PtrType>(tgtInfo.first) || !isa<PtrType>(srcInfo.first))
-    return op ? op->emitError("invalid ptr-like operand") : failure();
+/// Checks whether the shape of the operands is compatible with the operation.
+/// Operands must be scalars or have the same vector shape, additionally only
+/// vectors of rank 1 are supported.
+static LogicalResult verifyShapeInfo(mlir::Operation *op,
+                                     const std::pair<Type, int64_t> &tgtInfo,
+                                     const std::pair<Type, int64_t> &srcInfo) {
   // Check shape validity.
-  if (tgtInfo.second == -1 || srcInfo.second == -1)
+  if (tgtInfo.second == kInvalidRankError ||
+      srcInfo.second == kInvalidRankError)
     return op ? op->emitError("vectors of rank != 1 are not supported")
               : failure();
-  if (tgtInfo.second == -2 || srcInfo.second == -2)
+  if (tgtInfo.second == kScalableDimsError ||
+      srcInfo.second == kScalableDimsError)
     return op ? op->emitError(
                     "vectors with scalable dimensions are not supported")
               : failure();
@@ -77,29 +83,30 @@ LogicalResult mlir::ptr::isValidAddrSpaceCastImpl(Type tgt, Type src,
   return success();
 }
 
+LogicalResult mlir::ptr::isValidAddrSpaceCastImpl(Type tgt, Type src,
+                                                  Operation *op) {
+  std::pair<Type, int64_t> tgtInfo = getVecOrScalarInfo(tgt);
+  std::pair<Type, int64_t> srcInfo = getVecOrScalarInfo(src);
+  if (!isa<PtrType>(tgtInfo.first) || !isa<PtrType>(srcInfo.first))
+    return op ? op->emitError("invalid ptr-like operand") : failure();
+  // Verify shape validity.
+  return verifyShapeInfo(op, tgtInfo, srcInfo);
+}
+
 LogicalResult mlir::ptr::isValidPtrIntCastImpl(Type intLikeTy, Type ptrLikeTy,
                                                Operation *op) {
   // Check int-like type.
   std::pair<Type, int64_t> intInfo = getVecOrScalarInfo(intLikeTy);
+  // The int-like operand is invalid.
   if (!intInfo.first.isSignlessIntOrIndex())
-    /// The int-like operand is invalid.
     return op ? op->emitError("invalid int-like type") : failure();
   // Check ptr-like type.
   std::pair<Type, int64_t> ptrInfo = getVecOrScalarInfo(ptrLikeTy);
+  // The pointer-like operand is invalid.
   if (!isa<PtrType>(ptrInfo.first))
-    /// The pointer-like operand is invalid.
     return op ? op->emitError("invalid ptr-like type") : failure();
-  // Check shape validity.
-  if (intInfo.second == -1 || ptrInfo.second == -1)
-    return op ? op->emitError("vectors of rank != 1 are not supported")
-              : failure();
-  if (intInfo.second == -2 || ptrInfo.second == -2)
-    return op ? op->emitError(
-                    "vectors with scalable dimensions are not supported")
-              : failure();
-  if (intInfo.second != ptrInfo.second)
-    return op ? op->emitError("incompatible operand shapes") : failure();
-  return success();
+  // Verify shape validity.
+  return verifyShapeInfo(op, intInfo, ptrInfo);
 }
 
 #include "mlir/Dialect/Ptr/IR/PtrOpsDialect.cpp.inc"
