[mlir][tensor][bufferize] Support memory_space for tensor.pad
authorMatthias Springer <springerm@google.com>
Thu, 27 Oct 2022 10:24:02 +0000 (12:24 +0200)
committerMatthias Springer <springerm@google.com>
Thu, 27 Oct 2022 10:29:57 +0000 (12:29 +0200)
commit09dfb44193973db76f33e6518b83e53a762cfb34
treefebe309fed6bc38e721b4c23f4a86ba76f8a4213
parente26f28711c1f81b7f74ce749269e7539492d3934
[mlir][tensor][bufferize] Support memory_space for tensor.pad

This change adds memory space support to tensor.pad. (tensor.generate and tensor.from_elements do not support memory spaces yet.)

The memory space is inferred from the buffer of the source tensor.

Instead of lowering tensor.pad to tensor.generate + tensor.insert_slice, it is now lowered to bufferization.alloc_tensor (with the correct memory space) + linalg.map + tensor.insert_slice.

Memory space support for the remaining two tensor ops is left for a later point, as this requires some more design discussions.

Differential Revision: https://reviews.llvm.org/D136265
mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp
mlir/test/Dialect/Tensor/bufferize.mlir
mlir/test/Dialect/Tensor/one-shot-bufferize.mlir