diff --git a/mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp b/mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp @@ -364,6 +364,7 @@ isa( [&](auto op) { return OpResult(); }) + // DimOp has no tensor result. + .Case([&](auto op) { return None; }) // ConstantOp is never inplaceable. .Case([&](ConstantOp op) { return op->getResult(0); }) // ExtractSliceOp is different: its result is not inplaceable on op.source diff --git a/mlir/test/Dialect/Linalg/comprehensive-module-bufferize.mlir b/mlir/test/Dialect/Linalg/comprehensive-module-bufferize.mlir --- a/mlir/test/Dialect/Linalg/comprehensive-module-bufferize.mlir +++ b/mlir/test/Dialect/Linalg/comprehensive-module-bufferize.mlir @@ -585,6 +585,49 @@ // ----- +#TILE_MAP = affine_map<(d0)[s0] -> (3, -d0 + s0)> + +// CHECK-DAG: #[[$DYN_MAP:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)> + +// CHECK: func @tiled_fill( +// CHECK-SAME: %[[A:[a-zA-Z0-9]*]]: memref +func @tiled_fill(%A: tensor {linalg.inplaceable = true}) -> tensor { + %c3 = constant 3 : index + %c0 = constant 0 : index + %f0 = constant 0.0 : f32 + + // CHECK: %[[M:.*]] = memref.dim %[[A]], {{.*}} : memref + %0 = tensor.dim %A, %c0 : tensor + + // CHECK: linalg.tiled_loop {{.*}} to (%[[M]]) {{.*}} outs{{.*}}%[[A]] + %1 = linalg.tiled_loop (%arg3) = (%c0) to (%0) step (%c3) + outs (%arg1 = %A: tensor) + iterators["parallel"] + { + // CHECK-NOT: alloc + + %2 = tensor.dim %arg1, %c0 : tensor + %3 = affine.min #TILE_MAP(%arg3)[%2] + + // CHECK: %[[SV_A:.*]] = memref.subview {{.*}} + %4 = tensor.extract_slice %arg1[%arg3] [%3] [1] : tensor to tensor + + // CHECK: linalg.fill(%{{.*}}, %[[SV_A]]) : f32, memref + %5 = linalg.fill(%f0, %4) : f32, tensor -> tensor + %6 = tensor.insert_slice %5 into %arg1[%arg3] [%3] [1] : tensor into tensor + + linalg.yield %6 : tensor + // CHECK: linalg.yield + // CHECK-NOT: tensor + } + + // CHECK: return + // CHECK-NOT: tensor + return %1 : tensor +} + +// ----- + // CHECK: #[[$DYNAMIC:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)> // CHECK: func private @external_func(memref)