Emit alloc instructions during a conversion pass

MLIR
1

I’m currently writing a conversion pass for a dialect, where I have some operation that I am rewriting using lower level dialects (arith, memref, affine, etc).

Let’s say that a function in the dialect is:

"my_op" (%0, %1, %2) : (memref<5xi64>, memref<5xi64>, memref<5xi64>) -> ()

In one of my passes, I want to emit memref alloc instructions as part of the conversion.

Let’s say part of my pass is to emit some arith instructions, e.g.:

Value lhs_i = rewriter.create<AffineLoadOp>(loc, lhs, ArrayRef<Value>{i});
Value rhs_i = rewriter.create<AffineLoadOp>(loc, rhs, ArrayRef<Value>{i});
Value sum_i = rewriter.create<AddIOp>(loc, intType, lhs_i, rhs_i);

With this alone (and the boilerplate, e.g. erasing the original op), the conversion pass works fine, and my_op is replaced with the above instructions.

However, when I try and add a memref alloc instruction, the conversion fails (silently), and my_op remains with none of the target instructions generated.

MemRefType my_type = MemRefType::get({5}, intType);
Value my_alloc = rewriter.create<AllocOp>(loc, my_type);

As far as I can see I’m not doing anything incorrectly: I define the type/size of the memref, and then generate the alloc instruction.
However, adding those two lines to any of my passes causes them to silently fail.

Is there something obvious I’m missing? As far as I can see looking at some other dialects, e.g. SparseTensor I’m doing the right thing.
Also, looking at this blogpost from early last year.2

I am currently trying to reproduce the failure in a barebones standalone dialect, I will share that once it’s working, assuming I don’t find a fix.

EDIT

I have reproduced this in a simple standalone dialect, defined in this GitHub repo.

If I have some MLIR code such as:

func @my_standalone() {
  %i0 = arith.constant 0 : index
  %0 = memref.alloc() : memref<5xi64>
  %1 = memref.alloc() : memref<5xi64>
  %2 = memref.alloc() : memref<5xi64>
  "standalone.MyOp" (%0, %1, %2) : (memref<5xi64>, memref<5xi64>, memref<5xi64>) -> ()
  return
}

Then if I run $ standalone-opt --standalone-to-loops, then I get the exact same code back.

However, if I comment out the alloc instructions in these lines.

Again, perhaps it’s something obvious I’m missing, but the fact that the failure is silent makes it tricky for me to move forward as a newer user.

2

A good way to debug this kind of things is the -debug option. It should show you in details what the conversion framework is trying to do.

What is possible here is that this AllocOp is not “legal” according to the configuration of your particular conversion and so this pattern application is rolled back.

3

Exactly this, the memref dialect from which AllocOp is taken is not declared as legal here

neither is the op separately.

4

Separately, can you elaborate on your intended use cases?
It is possible you are reinventing wheels.

5

Thanks so much all, the “legal dialect” parameter was what I needed, but wasn’t aware of: I had just blindly copy-pasted that bit of code from another example.

Adding MemRef dialect to the legal dialects allowed me to generate the code I needed.

Regarding my application, I was looking at the Number Theoretic Transform (NTT), for which you can precompute a table of values that can be used in the transformation, assuming you know the shape of your problem ahead of time.