[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

 /// what kind of layout map will be used. For best composability (without
 /// 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 = {});
+                             Optional<AffineMap> layout = None,
+                             Optional<unsigned> memorySpace = None);
 
 /// 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 = {});
+BaseMemRefType
+getMemRefTypeWithFullyDynamicLayout(TensorType tensorType,
+                                    Optional<unsigned> memorySpace = None);
 
 /// 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 = {});
+                                      Optional<unsigned> memorySpace = None);
 
 } // 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) {
   auto bbArg = value.dyn_cast<BlockArgument>();
   if (!bbArg)
     return false;
   return isa<func::FuncOp>(bbArg.getOwner()->getParentOp());
 }
 
 BaseMemRefType bufferization::getMemRefType(TensorType tensorType,
                                             const BufferizationOptions &options,
-                                            MemRefLayoutAttrInterface layout,
-                                            Attribute memorySpace) {
+                                            Optional<AffineMap> layout,
+                                            Optional<unsigned> memorySpace) {
   // Case 1: Unranked memref type.
   if (auto unrankedTensorType = tensorType.dyn_cast<UnrankedTensorType>()) {
-    assert(!layout && "UnrankedTensorType cannot have a layout map");
+    assert(!layout.hasValue() && "UnrankedTensorType cannot have a layout map");
+    if (memorySpace.hasValue())
       return UnrankedMemRefType::get(unrankedTensorType.getElementType(),
-                                   memorySpace);
+                                     *memorySpace);
+    return UnrankedMemRefType::get(unrankedTensorType.getElementType(),
+                                   /*memorySpace=*/{});
   }
 
   // Case 2: Ranked memref type with specified layout.
   auto rankedTensorType = tensorType.cast<RankedTensorType>();
-  if (layout) {
+  if (layout.hasValue()) {
+    if (memorySpace.hasValue())
+      return MemRefType::get(rankedTensorType.getShape(),
+                             rankedTensorType.getElementType(), *layout,
+                             *memorySpace);
     return MemRefType::get(rankedTensorType.getShape(),
-                           rankedTensorType.getElementType(), layout,
-                           memorySpace);
+                           rankedTensorType.getElementType(), *layout);
   }
 
   // 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) {
+BaseMemRefType bufferization::getMemRefTypeWithFullyDynamicLayout(
+    TensorType tensorType, Optional<unsigned> memorySpace) {
   // Case 1: Unranked memref type.
   if (auto unrankedTensorType = tensorType.dyn_cast<UnrankedTensorType>()) {
+    if (memorySpace.hasValue())
+      return UnrankedMemRefType::get(unrankedTensorType.getElementType(),
+                                     *memorySpace);
     return UnrankedMemRefType::get(unrankedTensorType.getElementType(),
-                                   memorySpace);
+                                   /*memorySpace=*/{});
   }
 
   // Case 2: Ranked memref type.
   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());
+  if (memorySpace.hasValue())
     return MemRefType::get(rankedTensorType.getShape(),
                            rankedTensorType.getElementType(), stridedLayout,
-                         memorySpace);
+                           *memorySpace);
+  return MemRefType::get(rankedTensorType.getShape(),
+                         rankedTensorType.getElementType(), stridedLayout);
 }
 
 /// 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) {
+BaseMemRefType bufferization::getMemRefTypeWithStaticIdentityLayout(
+    TensorType tensorType, Optional<unsigned> memorySpace) {
   // Case 1: Unranked memref type.
   if (auto unrankedTensorType = tensorType.dyn_cast<UnrankedTensorType>()) {
+    if (memorySpace.hasValue())
       return UnrankedMemRefType::get(unrankedTensorType.getElementType(),
-                                   memorySpace);
+                                     *memorySpace);
+    return UnrankedMemRefType::get(unrankedTensorType.getElementType(),
+                                   /*memorySpace=*/{});
   }
 
   // Case 2: Ranked memref type.
   auto rankedTensorType = tensorType.cast<RankedTensorType>();
-  MemRefLayoutAttrInterface layout = {};
+  if (memorySpace.hasValue())
+    return MemRefType::get(rankedTensorType.getShape(),
+                           rankedTensorType.getElementType(), AffineMap(),
+                           *memorySpace);
   return MemRefType::get(rankedTensorType.getShape(),
-                         rankedTensorType.getElementType(), layout,
-                         memorySpace);
+                         rankedTensorType.getElementType());
 }

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;
+    Optional<AffineMap> layout = None;
 
     if (auto rankedMemRefType = sourceMemRefType.dyn_cast<MemRefType>())
       if (resultTensorType.isa<RankedTensorType>())
-        layout = rankedMemRefType.getLayout();
+        layout = rankedMemRefType.getLayout().getAffineMap();
 
     // 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);