Lowering ::mlir::memref::TensorStoreOp to LLVM IR

MLIR
1

In Toy example of MLIR, ConstantOp lowering is done as follows.

    %c0_47 = constant 0 : index
    %cst_48 = constant 1.000000e+00 : f32
    affine.store %cst_48, %57[%c0_47] : memref<2xf32>
    %c1_49 = constant 1 : index
    %cst_50 = constant 1.000000e+00 : f32
    affine.store %cst_50, %57[%c1_49] : memref<2xf32>

This strategy becomes extremely slow (in terms of compile time) as the size of constant increases.
Therefore, I used memref::TensorStoreOp to lower tensor to memref as follows.

    %cst_38 = constant dense<"0x35FAE73C2387E43CEB1FCF3C764...(skipped)...B86BC0D289A3C"> : tensor<256x128x3x3xf32>
    memref.tensor_store %cst_38, %60 : memref<256x128x3x3xf32>

However, there is no lowering passes provided by MLIR project, which lowers memref::TensorStoreOp to operators of LLVM Dialect.
Do you recommend me to lower memref::TensorStoreOp to LLVM::*Ops,
or other strategy to lower large tensor to memref type?

Sincerely yours,

2

You can see examples here:
https://github.com/llvm/llvm-project/blob/main/mlir/test/Dialect/Arith/bufferize.mlir#L19
https://github.com/llvm/llvm-project/blob/main/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize.mlir#L267

The idea is to hold an arith.constant which is a tensor, and at the bufferization stage large tensors are lowered to memref.global. This is also lowered to llvm. It means you will have a large constant on the data memory and it will be loaded from there when needed, this is the memref.get_global you see.

3

I am trying to convert toy::ConstantOp to Memref::GlobalOp as mentioned here, but not able find a path to do, can anyone give some example passes to do it?

4

The links I attached include examples for such lowering.
Try to run oneShotBufferize pass as seen in the second link:

mlir-opt -one-shot-bufferize="bufferize-function-boundaries=1 allow-return-allocs" -drop-equivalent-buffer-results

and give it the input from one of the tests there, for example:


func.func @basic() -> tensor<3x4xf32> {
  %0 = arith.constant dense<7.0> : tensor<3x4xf32>
  return %0 : tensor<3x4xf32>
}

You should get:

module {
  memref.global "private" constant @__constant_3x4xf32 : memref<3x4xf32> = dense<7.000000e+00> {alignment = 64 : i64}
  func.func @basic() -> memref<3x4xf32> {
    %0 = memref.get_global @__constant_3x4xf32 : memref<3x4xf32>
    return %0 : memref<3x4xf32>
  }
}

Notice this is arith::ConstantOp and not toy::ConstantOp.

5

Thank you !!!