[mlir][bufferization][NFC] Change signature of getMemRefType

[springerm]

Summary

These functions now accept unsigned values for memory spaces instead of Attributes.

Depends On D128274

[aartbik]

mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp

606

so, !hasValue() and {} are still distinct situation here I guess?

[springerm]

mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp

606

Correct, None means "not specified" and we do whatever the default constructor of MemRefType/UnrankedMemRefType is doing. AffineMap() or {} means identity layout.

[nicolasvasilache]

We want to go away from affine map to specify layout (see in particular https://discourse.llvm.org/t/rfc-remove-bare-pointer-calling-convention/63419 ) and a bunch of other relates posts.
Can this PR just add the form you want for the address space?

[springerm]

address comments

mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h

 /// copies), the fully dynamic layout map is used by default.
 ///
 /// Note: Canonicalization patterns could clean up layout maps and infer more
 /// precise layout maps after bufferization. However, many possible
 /// canonicalizations are currently not implemented.
 BaseMemRefType getMemRefType(TensorType tensorType,
                              const BufferizationOptions &options,
                              MemRefLayoutAttrInterface layout = {},
-                             Attribute memorySpace = {});
+                             unsigned memorySpace = 0);
 
 /// Return a MemRef type with fully dynamic layout. If the given tensor type
 /// is unranked, return an unranked MemRef type.
 BaseMemRefType getMemRefTypeWithFullyDynamicLayout(TensorType tensorType,
-                                                   Attribute memorySpace = {});
+                                                   unsigned memorySpace = 0);
 
 /// Return a MemRef type with a static identity layout (i.e., no layout map). If
 /// the given tensor type is unranked, return an unranked MemRef type.
-BaseMemRefType
-getMemRefTypeWithStaticIdentityLayout(TensorType tensorType,
-                                      Attribute memorySpace = {});
+BaseMemRefType getMemRefTypeWithStaticIdentityLayout(TensorType tensorType,
+                                                     unsigned memorySpace = 0);
 
 } // namespace bufferization
 } // namespace mlir
 
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h.inc"
 
 #endif // MLIR_DIALECT_BUFFERIZATION_IR_BUFFERIZABLEOPINTERFACE_H_

mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp

bool bufferization::isFunctionArgument(Value value) {
   if (!bbArg)
     return false;
   return isa<func::FuncOp>(bbArg.getOwner()->getParentOp());
 }
 
 BaseMemRefType bufferization::getMemRefType(TensorType tensorType,
                                             const BufferizationOptions &options,
                                             MemRefLayoutAttrInterface layout,
-                                            Attribute memorySpace) {
+                                            unsigned memorySpace) {
+  auto memorySpaceAttr = IntegerAttr::get(
+      IntegerType::get(tensorType.getContext(), 64), memorySpace);
+
   // Case 1: Unranked memref type.
   if (auto unrankedTensorType = tensorType.dyn_cast<UnrankedTensorType>()) {
     assert(!layout && "UnrankedTensorType cannot have a layout map");
     return UnrankedMemRefType::get(unrankedTensorType.getElementType(),
-                                   memorySpace);
+                                   memorySpaceAttr);
   }
 
   // Case 2: Ranked memref type with specified layout.
   auto rankedTensorType = tensorType.cast<RankedTensorType>();
   if (layout) {
     return MemRefType::get(rankedTensorType.getShape(),
                            rankedTensorType.getElementType(), layout,
-                           memorySpace);
+                           memorySpaceAttr);
   }
 
   // Case 3: Configured with "fully dynamic layout maps".
   if (options.unknownTypeConversion ==
       BufferizationOptions::LayoutMapOption::FullyDynamicLayoutMap)
     return getMemRefTypeWithFullyDynamicLayout(tensorType, memorySpace);
 
   // Case 4: Configured with "static identity layout maps".
   if (options.unknownTypeConversion ==
       BufferizationOptions::LayoutMapOption::IdentityLayoutMap)
     return getMemRefTypeWithStaticIdentityLayout(tensorType, memorySpace);
 
   llvm_unreachable("InferLayoutMap is an invalid option");
 }
 
 BaseMemRefType
 bufferization::getMemRefTypeWithFullyDynamicLayout(TensorType tensorType,
-                                                   Attribute memorySpace) {
+                                                   unsigned memorySpace) {
   // Case 1: Unranked memref type.
   if (auto unrankedTensorType = tensorType.dyn_cast<UnrankedTensorType>()) {
     return UnrankedMemRefType::get(unrankedTensorType.getElementType(),
                                    memorySpace);
   }
 
   // Case 2: Ranked memref type.
+  auto memorySpaceAttr = IntegerAttr::get(
+      IntegerType::get(tensorType.getContext(), 64), memorySpace);
   auto rankedTensorType = tensorType.cast<RankedTensorType>();
   int64_t dynamicOffset = ShapedType::kDynamicStrideOrOffset;
   SmallVector<int64_t> dynamicStrides(rankedTensorType.getRank(),
                                       ShapedType::kDynamicStrideOrOffset);
   AffineMap stridedLayout = makeStridedLinearLayoutMap(
       dynamicStrides, dynamicOffset, rankedTensorType.getContext());
   return MemRefType::get(rankedTensorType.getShape(),
                          rankedTensorType.getElementType(), stridedLayout,
-                         memorySpace);
+                         memorySpaceAttr);
 }
 
 /// Return a MemRef type with a static identity layout (i.e., no layout map). If
 /// the given tensor type is unranked, return an unranked MemRef type.
 BaseMemRefType
 bufferization::getMemRefTypeWithStaticIdentityLayout(TensorType tensorType,
-                                                     Attribute memorySpace) {
+                                                     unsigned memorySpace) {
   // Case 1: Unranked memref type.
   if (auto unrankedTensorType = tensorType.dyn_cast<UnrankedTensorType>()) {
     return UnrankedMemRefType::get(unrankedTensorType.getElementType(),
                                    memorySpace);
   }
 
   // Case 2: Ranked memref type.
   auto rankedTensorType = tensorType.cast<RankedTensorType>();
+  auto memorySpaceAttr = IntegerAttr::get(
+      IntegerType::get(tensorType.getContext(), 64), memorySpace);
   MemRefLayoutAttrInterface layout = {};
   return MemRefType::get(rankedTensorType.getShape(),
                          rankedTensorType.getElementType(), layout,
-                         memorySpace);
+                         memorySpaceAttr);
 }

mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp

struct CastOpInterface
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationOptions &options) const {
     auto castOp = cast<tensor::CastOp>(op);
 
     // The result buffer still has the old (pre-cast) type.
     Value resultBuffer = getBuffer(rewriter, castOp.getSource(), options);
     auto sourceMemRefType = resultBuffer.getType().cast<BaseMemRefType>();
-    Attribute memorySpace = sourceMemRefType.getMemorySpace();
     TensorType resultTensorType =
         castOp.getResult().getType().cast<TensorType>();
     MemRefLayoutAttrInterface layout;
 
     if (auto rankedMemRefType = sourceMemRefType.dyn_cast<MemRefType>())
       if (resultTensorType.isa<RankedTensorType>())
         layout = rankedMemRefType.getLayout();
 
     // Compute the new memref type.
     Type resultMemRefType =
-        getMemRefType(resultTensorType, options, layout, memorySpace);
+        getMemRefType(resultTensorType, options, layout,
+                      sourceMemRefType.getMemorySpaceAsInt());
 
     // Replace the op with a memref.cast.
     assert(memref::CastOp::areCastCompatible(resultBuffer.getType(),
                                              resultMemRefType) &&
            "CallOp::bufferize: cast incompatible");
     replaceOpWithNewBufferizedOp<memref::CastOp>(rewriter, op, resultMemRefType,
                                                  resultBuffer);