Page MenuHomePhabricator

[mlir][Linalg] Allow comprehensive bufferization to use callbacks for alloc/dealloc.
ClosedPublic

Authored by mravishankar on Oct 20 2021, 11:47 AM.

Details

Summary

Using callbacks for allocation/deallocation allows users to override
the default.
Also add an option to comprehensive bufferization pass to use alloca
instead of allocs. Note that this option is just for testing. The
option to use alloca does not work well with the option to allow for
returning memrefs.

Diff Detail

Event Timeline

mravishankar created this revision.Oct 20 2021, 11:47 AM
mravishankar requested review of this revision.Oct 20 2021, 11:47 AM
mehdi_amini requested changes to this revision.Oct 20 2021, 11:53 AM
mehdi_amini added inline comments.
mlir/include/mlir/Dialect/Linalg/Passes.h
67 ↗(On Diff #381046)

I have very strong concerns about C++ injection in passes

I pinged here just yesterday: https://reviews.llvm.org/D110678 on this topic.

This revision now requires changes to proceed.Oct 20 2021, 11:53 AM

@mehdi_amini I see your concerns with reproducers, but what do you suggest here. We want to be able to use an allocator that make sense for the use case. For example, there might be constraints of not being able to do heap allocations. For CUDA you might want to do allocations on scratch pad memory. All this is easier to reason about while generating the allocation, rather than doing it as a later pass.

Maybe a way forward here is to not have a way to override the allocator used at the pass granularity, but rather have the utility function mlir::linalg::bufferize take call backs for allocator and deallocator. That way any pass that calls this utility function will be explicit about its allocation functions, and that should not affect the reproducer?

Maybe a way forward here is to not have a way to override the allocator used at the pass granularity, but rather have the utility function mlir::linalg::bufferize take call backs for allocator and deallocator. That way any pass that calls this utility function will be explicit about its allocation functions, and that should not affect the reproducer?

Yes, thanks.

Remove use of callbacks at pass granularity

springerm accepted this revision.Oct 24 2021, 12:33 AM
springerm added inline comments.
mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp
1718–1722

I'm wondering if all these bufferize functions could become member functions of a new class (or LinalgComprehensiveModuleBufferize) which has builder, aliasInfo etc. as fields, so we don't have to pass this stuff around.

Anyway, not needed for this revision.

2473–2479

I would probably have separate allocationFns: One for alloc and one for alloca. The implementation of these could then call some function that has the common functionality.

3054–3066

If possible, I would rewrite this to something like:

if (!allocationFns) {
  if (useAlloca)
    allocationFns = std::make_unique<AllocationCallbacks>(&AllocaStrategy::allocate, &AllcaStrategy::deallocate);
  } else {
    allocationFns = std::make_unique<AllocationCallbacks>(&AllocStrategy::allocate, &AllocStrategy::deallocate);
}

Where AllocStrategy::allocate etc. are static functions. And drop the field initializers in AllocationCallbacks.

Maybe even a class AllocStrategy with fully virtual functions allocate and deallocate that can be overridden by specific implementations (instead of AllocationCallbacks). (But I people don't like virtual functions here in general, so...)

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
6317

Maybe this also requires updating some CMakeFile?

This revision is now accepted and ready to land.Oct 24 2021, 12:33 AM
mravishankar marked an inline comment as done.

Rebase and address comments.

mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp
1718–1722

I dont have a strong opinion. I am more in favor of modular static functions with well defined interface. The class approach leads to having state across methods and weird coupling.

3054–3066

I can try to clean this up. In terms of a class and inheritance, again the issue for me is the state that you keep. Call backs are restrictive in terms of what they use (well defined by the interface). So it forces you to keep things hermitic. Classes can lead to fairly complex state being passed around. Callbacks are in general how things are being done in Linalg in other places.

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
6317

The CMake file is updated (the builds pass). The CMake file is nested with the source folder and not at the top level like bazel build file is.

mravishankar reopened this revision.Oct 25 2021, 10:55 AM

Reopening for resubmit of reverted change

This revision is now accepted and ready to land.Oct 25 2021, 10:55 AM