[mlir][tensor][bufferize] Support tensor.rank in BufferizableOpInterfaceImpl

[springerm]

Summary

This is the only op that is not supported via BufferizableOpInterfaceImpl bufferization. Once this op is supported we can switch tensor-bufferize over to the new unified bufferization.

Depends On D117984

[springerm]

rebase

mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.cpp

if (!tensorType) {
       continue;
     }
 
     // Get the buffer type from the bufferized function type.
     Type memrefType = bufferizedFuncType.getInput(idx);
     Value memref = frontBlock.addArgument(memrefType, bbArg.getLoc());
     OpBuilder b(funcOp->getContext());
     b.setInsertionPointToStart(&frontBlock);
-    // Replace all uses of bbArg through a ToMemRefOp by a memref::CastOp.
+    // Replace all uses of bbArg through a ToMemRefOp.
     for (auto &use : llvm::make_early_inc_range(bbArg.getUses())) {
       if (auto toMemrefOp =
               dyn_cast<bufferization::ToMemrefOp>(use.getOwner())) {
+        if (memref.getType() != toMemrefOp.memref().getType()) {
+          // Type has changed, insert a cast.
           assert(memref::CastOp::areCastCompatible(
                      memref.getType(), toMemrefOp.memref().getType()) &&
                  "bufferizeFuncOpBoundary: cast incompatible");
           auto castOp = b.create<memref::CastOp>(
               funcOp.getLoc(), toMemrefOp.memref().getType(), memref);
           toMemrefOp.memref().replaceAllUsesWith(castOp);
+        } else {
+          // Type did not change, replace directly.
+          toMemrefOp.memref().replaceAllUsesWith(memref);
+        }
       }
     }
     // Replace all remaining uses by a to_tensor.
     if (!bbArg.use_empty()) {
       auto toTensorOp =
           b.create<bufferization::ToTensorOp>(funcOp.getLoc(), memref);
       bbArg.replaceAllUsesWith(toTensorOp);
     }

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

if (failed(createMemCpy(rewriter, loc, srcMemref, subView,
                             state.getOptions())))
       return failure();
 
     replaceOpWithBufferizedValues(rewriter, op, *dstMemref);
     return success();
   }
 };
 
+/// Bufferization of tensor.rank. Replace with memref.rank.
+struct RankOpInterface
+    : public BufferizableOpInterface::ExternalModel<RankOpInterface,
+                                                    tensor::RankOp> {
+  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
+                              const BufferizationState &state) const {
+    return true;
+  }
+
+  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
+                               const BufferizationState &state) const {
+    return false;
+  }
+
+  OpResult getAliasingOpResult(Operation *op, OpOperand &opOperand,
+                               const BufferizationState &state) const {
+    return OpResult();
+  }
+
+  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
+                          const BufferizationState &state) const {
+    auto rankOp = cast<tensor::RankOp>(op);
+    Value v = *state.getBuffer(rewriter, rankOp->getOpOperand(0) /*source*/);
+    replaceOpWithNewBufferizedOp<memref::RankOp>(rewriter, op, rankOp.getType(),
+                                                 v);
+    return success();
+  }
+};
+
 } // namespace
 } // namespace tensor
 } // namespace mlir
 
 void mlir::tensor::registerBufferizableOpInterfaceExternalModels(
     DialectRegistry &registry) {
   registry.addOpInterface<CastOp, CastOpInterface>();
   registry.addOpInterface<DimOp, DimOpInterface>();
   registry.addOpInterface<ExtractSliceOp, ExtractSliceOpInterface>();
   registry.addOpInterface<ExtractOp, ExtractOpInterface>();
   registry.addOpInterface<InsertOp, InsertOpInterface>();
   registry.addOpInterface<InsertSliceOp, InsertSliceOpInterface>();
+  registry.addOpInterface<RankOp, RankOpInterface>();
 }

mlir/test/Dialect/Linalg/comprehensive-module-bufferize.mlir

func @write_after_select_read_one(
   %w = tensor.insert %cst into %s[%idx] : tensor<?xf32>
 
   // CHECK: %[[f:.*]] = memref.load %[[t1]]
   %f = tensor.extract %t1[%idx] : tensor<?xf32>
 
   // CHECK: return %[[f]], %[[select]]
   return %f, %w : f32, tensor<?xf32>
 }
+
+// -----
+
+// CHECK-LABEL: func @tensor_rank(
+//  CHECK-SAME:     %[[arg0:.*]]: memref<*xf32>
+func @tensor_rank(%arg0: tensor<*xf32>) -> index {
+  // CHECK: %[[r:.*]] = memref.rank %[[arg0]]
+  %0 = tensor.rank %arg0 : tensor<*xf32>
+  // CHECK: return %[[r]] : index
+  return %0 : index
+}