From e5dc99e642d05b1075b6c5480a5b3f0485a25038 Mon Sep 17 00:00:00 2001
From: Matthias Springer <springerm@google.com>
Date: Wed, 14 Dec 2022 12:26:29 +0100
Subject: [PATCH] [mlir][tensor][bufferize] Improve bufferization of
 DimOp/RankOp

The tensor operands do not bufferize to a memory read.

Differential Revision: https://reviews.llvm.org/D140007
---
 .../Tensor/Transforms/BufferizableOpInterfaceImpl.cpp   |  6 ++++--
 mlir/test/Dialect/Tensor/one-shot-bufferize.mlir        | 17 +++++++++++++++++
 2 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp
index aa5a1d8..3c634e9 100644
--- a/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -205,7 +205,8 @@ struct DimOpInterface
                                                     tensor::DimOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
                               const AnalysisState &state) const {
-    return true;
+    // The op reads the tensor's metadata but not its contents.
+    return false;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
@@ -927,7 +928,8 @@ struct RankOpInterface
                                                     tensor::RankOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
                               const AnalysisState &state) const {
-    return true;
+    // The op reads the tensor's metadata but not its contents.
+    return false;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
diff --git a/mlir/test/Dialect/Tensor/one-shot-bufferize.mlir b/mlir/test/Dialect/Tensor/one-shot-bufferize.mlir
index e1473c8..59fde56 100644
--- a/mlir/test/Dialect/Tensor/one-shot-bufferize.mlir
+++ b/mlir/test/Dialect/Tensor/one-shot-bufferize.mlir
@@ -330,3 +330,20 @@ func.func @insert_slice_full_overwrite(%t: tensor<10xf32>, %b: tensor<10xf32>) -
   %2 = tensor.insert_slice %b into %t[0][10][1] : tensor<10xf32> into tensor<10xf32>
   return %2 : tensor<10xf32>
 }
+
+// -----
+
+// CHECK-LABEL: func @dim_not_reading(
+//  CHECK-SAME:     %[[t:.*]]: memref<?xf32
+func.func @dim_not_reading(%t: tensor<?xf32>, %f: f32, %pos: index) 
+    -> (tensor<?xf32>, index)
+{
+  %c0 = arith.constant 0 : index
+  // CHECK-NOT: memref.alloc
+  // CHECK-NOT: memref.copy
+  //     CHECK: memref.store %{{.*}}, %[[t]]
+  %0 = tensor.insert %f into %t[%pos] : tensor<?xf32>
+  //     CHECK: memref.dim %[[t]]
+  %1 = tensor.dim %t, %c0 : tensor<?xf32>
+  return %0, %1 : tensor<?xf32>, index
+}
-- 
2.7.4