This is an archive of the discontinued LLVM Phabricator instance.

Introduce alloca_scope op
ClosedPublic

Authored by shabalin on Mar 2 2021, 6:45 AM.

Details

Summary

Introduction

This proposal describes the new op to be added to the std (and later moved memref)
dialect called alloca_scope.

Motivation

Alloca operations are easy to misuse, especially if one relies on it while doing
rewriting/conversion passes. For example let's consider a simple example of two
independent dialects, one defines an op that wants to allocate on-stack and
another defines a construct that corresponds to some form of looping:

dialect1.looping_op {
  %x = dialect2.stack_allocating_op
}

Since the dialects might not know about each other they are going to define a
lowering to std/scf/etc independently:

scf.for … {
   %x_temp = std.alloca …
   … // do some domain-specific work using %x_temp buffer 
   … // and store the result into %result
   %x = %result
}

Later on the scf and std.alloca is going to be lowered to llvm using a
combination of llvm.alloca and unstructured control flow.

At this point the use of %x_temp is bound to either be either optimized by
llvm (for example using mem2reg) or in the worst case: perform an independent
stack allocation on each iteration of the loop. While the llvm optimizations are
likely to succeed they are not guaranteed to do so, and they provide
opportunities for surprising issues with unexpected use of stack size.

Proposal

We propose a new operation that defines a finer-grain allocation scope for the
alloca-allocated memory called alloca_scope:

alloca_scope {
   %x_temp = alloca …
   ...
}

Here the lifetime of %x_temp is going to be bound to the narrow annotated
region within alloca_scope. Moreover, one can also return values out of the
alloca_scope with an accompanying alloca_scope.return op (that behaves
similarly to scf.yield):

%result = alloca_scope {
   %x_temp = alloca …
   …
   alloca_scope.return %myvalue
}

Under the hood the alloca_scope is going to lowered to a combination of
llvm.intr.stacksave and llvm.intr.strackrestore that are going to be invoked
automatically as control-flow enters and leaves the body of the alloca_scope.

The key value of the new op is to allow deterministic guaranteed stack use
through an explicit annotation in the code which is finer-grain than the
function-level scope of AutomaticAllocationScope interface. alloca_scope
can be inserted at arbitrary locations and doesn’t require non-trivial
transformations such as outlining.

Which dialect

Before memref dialect is split, alloca_scope can temporarily reside in std
dialect, and later on be moved to memref together with the rest of
memory-related operations.

Implementation

An implementation of the op is available here.

Original commits:

  • Add initial scaffolding for alloca_scope op
  • Add alloca_scope.return op
  • Add no region arguments and variadic results
  • Add op descriptions
  • Add failing test case
  • Add another failing test
  • Initial implementation of lowering for std.alloca_scope
  • Fix backticks
  • Fix getSuccessorRegions implementation

Diff Detail

Event Timeline

shabalin created this revision.Mar 2 2021, 6:45 AM
shabalin requested review of this revision.Mar 2 2021, 6:45 AM

Can you cleanup the summary and provide an actual contextual description instead? In particular the motivation for this op would be interesting to provide here.
(see https://chris.beams.io/posts/git-commit/ for general guidelines)

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td
469 ↗(On Diff #327439)

delimited

@mehdi_amini I was going to post a longer post describing the context and link it here on discourse but it seems to have been lost somewhere in moderation. I'll update update the commit with more info before it shows up.

shabalin retitled this revision from Add alloca_scope op to Introduce alloca_scope op.Mar 2 2021, 12:13 PM
shabalin edited the summary of this revision. (Show Details)

@mehdi_amini Detailed description posted above.

shabalin edited the summary of this revision. (Show Details)Mar 2 2021, 12:16 PM

@mehdi_amini Detailed description posted above.

Thanks! I'll wait for the RFC to show up on Discourse and we can discuss it here :)

I just in general look at the Phabricator description to be what will be committed, and thus follow roughly the guidelines I linked above.

shabalin updated this revision to Diff 327539.Mar 2 2021, 12:20 PM
This comment was removed by shabalin.
shabalin updated this revision to Diff 327540.EditedMar 2 2021, 12:22 PM

Fixed typo in the description

shabalin marked an inline comment as done.Mar 2 2021, 12:23 PM

@mehdi_amini Sounds good to me! Sorry for the initial poor description without any context.

ftynse added inline comments.Mar 3 2021, 1:29 AM
mlir/include/mlir/Dialect/StandardOps/IR/Ops.td
500 ↗(On Diff #327540)

I'm not convinced we want to restrict the body region to only have one block.

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp
2055–2056

Currently, only one block is allowed so beforeBody==afterBody always. I think we should relax that and allow an arbitrary CFG in the body. In which case, we probably want to handle multiple blocks terminated with alloca_scope.return.

mlir/test/Conversion/StandardToLLVM/convert-alloca-scope.mlir
6

Nit: we indent // CHECK lines

mlir/test/Dialect/Standard/ops.mlir
103 ↗(On Diff #327540)

Hmm, don't we actually want to // CHECK the printed op?

107 ↗(On Diff #327540)

Could we also have the test for alloca_scope.return with operands?

shabalin updated this revision to Diff 351134.Jun 10 2021, 4:50 AM

Rebase on top of latest main

I suppose this now belongs to the memref dialect.

shabalin updated this revision to Diff 351394.Jun 11 2021, 3:43 AM

Move the op to memref dialect.

ftynse accepted this revision.Jun 11 2021, 4:12 AM
ftynse added inline comments.
mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
213

Here and below std prefix needs to be updated to memref.

This revision is now accepted and ready to land.Jun 11 2021, 4:12 AM
bondhugula added inline comments.
mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
260

be to -> to be

262

It is customary to include an example here showing the syntax.

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp
2032–2037

You don't need the mlir:: here I think.

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
260

Typo: Parse

shabalin updated this revision to Diff 351427.Jun 11 2021, 6:09 AM
shabalin marked 5 inline comments as done.

Fix doc strings and typos

This revision was automatically updated to reflect the committed changes.