[mlir][bufferize] Infer memref types when possible

[springerm]

Summary

Instead of recomputing memref types from tensor types, try to infer them when possible. This results in more precise layout maps.

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

struct BufferizationState {
   /// bufferization. Inplacability decisions can be overridden with the optional
   /// `overrideInPlace` parameter.
   FailureOr<Value>
   getBuffer(RewriterBase &rewriter, OpOperand &opOperand,
             Optional<ForceInPlacability> overrideInPlace = None,
             Optional<Operation *> customCopyInsertionPoint = None);
 
   /// Return the buffer type for a given OpOperand (tensor) after bufferization.
+  ///
+  /// Note: Op implementations should preferrably call `getBuffer()->getType()`.
+  /// This function should only be used if `getBuffer` cannot be used.
   BaseMemRefType getBufferType(OpOperand &opOperand) const;
 
   /// Return a reference to the BufferizationOptions.
   const BufferizationOptions &getOptions() const {
     return analysisState.getOptions();
   }
 
   const AnalysisState &getAnalysisState() const { return analysisState; }

 template <typename OpTy, typename... Args>
 OpTy replaceOpWithNewBufferizedOp(RewriterBase &rewriter, Operation *op,
                                   Args &&...args) {
   auto newOp = rewriter.create<OpTy>(op->getLoc(), std::forward<Args>(args)...);
   replaceOpWithBufferizedValues(rewriter, op, newOp->getResults());
   return newOp;
 }
 
-/// Return a MemRefType to which the `tensorType` can be bufferized in a
-/// composable fashion. The layout must be the most dynamic possible and
-/// canonicalize away once bufferization is finished.
+/// Return a MemRefType to which the `tensorType` can be bufferized.
+///
+/// If possible, op bufferization implementations should not use this function
+/// and instead infer precise memref types for tensor results by themselves.
+///
+/// Unless a layout map was specified, `options` flags determine 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 = {});
 
 /// Try to hoist all new buffer allocations until the next hoisting barrier.
 LogicalResult hoistBufferAllocations(Operation *op,
                                      const BufferizationOptions &options);

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

struct IndexCastOpInterface
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
                                 const AnalysisState &state) const {
     return BufferRelation::Equivalent;
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           BufferizationState &state) const {
     auto castOp = cast<arith::IndexCastOp>(op);
+    auto resultTensorType = castOp.getType().cast<TensorType>();
 
     Value source = *state.getBuffer(rewriter, op->getOpOperand(0) /*in*/);
     auto sourceType = source.getType().cast<BaseMemRefType>();
 
     // Result type should have same layout and address space as the source type.
-    MemRefLayoutAttrInterface layout = {};
-    if (auto rankedMemRefType = sourceType.dyn_cast<MemRefType>())
-      layout = rankedMemRefType.getLayout();
-    Type resultType =
-        getMemRefType(castOp.getType().cast<TensorType>(), state.getOptions(),
-                      layout, sourceType.getMemorySpace());
+    BaseMemRefType resultType;
+    if (auto rankedMemRefType = sourceType.dyn_cast<MemRefType>()) {
+      resultType = MemRefType::get(
+          rankedMemRefType.getShape(), resultTensorType.getElementType(),
+          rankedMemRefType.getLayout(), rankedMemRefType.getMemorySpace());
+    } else {
+      auto unrankedMemrefType = sourceType.cast<UnrankedMemRefType>();
+      resultType = UnrankedMemRefType::get(resultTensorType.getElementType(),
+                                           unrankedMemrefType.getMemorySpace());
+    }
 
     replaceOpWithNewBufferizedOp<arith::IndexCastOp>(rewriter, op, resultType,
                                                      source);
     return success();
   }
 };
 
 /// Bufferization of arith.select. Just replace the operands.
LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
     Value falseBuffer =
         *state.getBuffer(rewriter, selectOp->getOpOperand(2) /*false_value*/);
 
     // The "true" and the "false" operands must have the same type. If the
     // buffers have different types, they differ only in their layout map. Cast
     // both of them to the most dynamic MemRef type.
     if (trueBuffer.getType() != falseBuffer.getType()) {
       auto trueType = trueBuffer.getType().cast<MemRefType>();
-      auto tensorType = selectOp.getTrueValue().getType().cast<TensorType>();
       int64_t dynamicOffset = ShapedType::kDynamicStrideOrOffset;
-      SmallVector<int64_t> dynamicStrides(tensorType.getRank(),
+      SmallVector<int64_t> dynamicStrides(trueType.getRank(),
                                           ShapedType::kDynamicStrideOrOffset);
       AffineMap stridedLayout = makeStridedLinearLayoutMap(
           dynamicStrides, dynamicOffset, op->getContext());
-      BaseMemRefType castedType = bufferization::getMemRefType(
-          tensorType, state.getOptions(), AffineMapAttr::get(stridedLayout),
-          trueType.getMemorySpace());
+      auto castedType =
+          MemRefType::get(trueType.getShape(), trueType.getElementType(),
+                          stridedLayout, trueType.getMemorySpaceAsInt());
       trueBuffer = rewriter.create<memref::CastOp>(loc, castedType, trueBuffer);
       falseBuffer =
           rewriter.create<memref::CastOp>(loc, castedType, falseBuffer);
     }
 
     replaceOpWithNewBufferizedOp<arith::SelectOp>(
         rewriter, op, selectOp.getCondition(), trueBuffer, falseBuffer);
     return success();

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

struct ExecuteRegionOpInterface
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           BufferizationState &state) const {
     auto executeRegionOp = cast<scf::ExecuteRegionOp>(op);
 
     // Compute new result types.
     SmallVector<Type> newResultTypes;
     for (Type type : executeRegionOp->getResultTypes()) {
       if (auto tensorType = type.dyn_cast<TensorType>()) {
+        // TODO: Infer the result type instead of computing it.
         newResultTypes.push_back(getMemRefType(tensorType, state.getOptions()));
       } else {
         newResultTypes.push_back(type);
       }
     }
 
     // Create new op and move over region.
     auto newOp =
struct IfOpInterface
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           BufferizationState &state) const {
     auto ifOp = cast<scf::IfOp>(op);
 
     // Compute new types of the bufferized scf.if op.
     SmallVector<Type> newTypes;
     for (Type returnType : ifOp->getResultTypes()) {
       if (auto tensorType = returnType.dyn_cast<TensorType>()) {
+        // TODO: Infer the result type instead of computing it.
         newTypes.push_back(getMemRefType(tensorType, state.getOptions()));
       } else {
         newTypes.push_back(returnType);
       }
     }
 
     // Create new op.
     auto newIfOp =

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

struct AssumingOpInterface
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           BufferizationState &state) const {
     auto assumingOp = cast<shape::AssumingOp>(op);
 
     // Compute new result types.
     SmallVector<Type> newResultTypes;
     for (Type type : assumingOp->getResultTypes()) {
       if (auto tensorType = type.dyn_cast<TensorType>()) {
+        // TODO: Infer the result type instead of computing it.
         newResultTypes.push_back(getMemRefType(tensorType, state.getOptions()));
       } else {
         newResultTypes.push_back(type);
       }
     }
 
     // Create new op and move over region.
     auto newOp = rewriter.create<shape::AssumingOp>(