OwnershipBasedBufferDeallocation.cpp Source File (original) (raw)

31 namespace mlir {

32 namespace bufferization {

33 #define GEN_PASS_DEF_OWNERSHIPBASEDBUFFERDEALLOCATIONPASS

34 #include "mlir/Dialect/Bufferization/Transforms/Passes.h.inc"

35 }

36 }

38 using namespace mlir;

46 return builder.create arith::ConstantOp(loc, builder.getBoolAttr(value));

47 }

54 if (isa(op))

55 return !hasEffectMemoryEffects::Allocate(op) &&

56 !hasEffectMemoryEffects::Free(op);

62 }

69 return true;

71 if (!region.empty())

72 if (llvm::any_of(region.front().getArguments(), isMemref))

73 return true;

74 return false;

75 }

81 namespace {

85 class Backedges {

86 public:

90 public:

92 Backedges(Operation *op) { recurse(op); }

95 size_t size() const { return edgeSet.size(); }

98 BackedgeSetT::const_iterator begin() const { return edgeSet.begin(); }

100

101 BackedgeSetT::const_iterator end() const { return edgeSet.end(); }

102

103 private:

104

105

106

107 bool enter(Block &current, Block *predecessor) {

108 bool inserted = visited.insert(&current).second;

109 if (!inserted)

110 edgeSet.insert(std::make_pair(predecessor, &current));

111 return inserted;

112 }

113

114

115 void exit(Block &current) { visited.erase(&current); }

116

117

118

121

122

123 if (isa(op)) {

125 recurse(*succ, current);

126 }

127

128

130 if (!region.empty())

131 recurse(region.front(), current);

132 }

133 }

134

135

136

137 void recurse(Block &block, Block *predecessor) {

138

139

140 if (!enter(block, predecessor))

141 return;

142

143

145 recurse(&op);

146

147

148 exit(block);

149 }

150

151

152 BlockSetT visited;

153

154

155 BackedgeSetT edgeSet;

156 };

157

158 }

159

160

161

162

163

164 namespace {

165

166

167 class BufferDeallocation {

168 public:

172

173

174 LogicalResult deallocate(FunctionOpInterface op);

175

176 private:

177

178 template <typename... T>

179 typename std::enable_if<sizeof...(T) == 0, FailureOr<Operation *>>::type

181 return op;

182 }

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

210

211

212 template <typename InterfaceT, typename... InterfacesU>

213 FailureOr<Operation *> handleOp(Operation *op) {

215 if (auto concreteOp = dyn_cast(op)) {

216 FailureOr<Operation *> result = handleInterface(concreteOp);

217 if (failed(result))

218 return failure();

219 next = *result;

220 }

221 if (!next)

222 return FailureOr<Operation *>(nullptr);

223 return handleOp<InterfacesU...>(next);

224 }

225

226

227 FailureOr<Operation *> handleAllInterfaces(Operation *op) {

228 if (auto deallocOpInterface = dyn_cast(op))

229 return deallocOpInterface.process(state, options);

230

231 if (failed(verifyOperationPreconditions(op)))

232 return failure();

233

234 return handleOp<MemoryEffectOpInterface, RegionBranchOpInterface,

235 CallOpInterface, BranchOpInterface,

236 RegionBranchTerminatorOpInterface>(op);

237 }

238

239

240

241

242

243

244

245

246

247

248

249

250

251

252

253

254

255

256

257

258

259 FailureOr<Operation *> handleInterface(BranchOpInterface op);

260

261

262

263

264

265

266

267

268

269

270

271

272

273

274

275

276

277

278

279

280

281

282

283

284

285

286

287

288

289

290

291 FailureOr<Operation *> handleInterface(RegionBranchOpInterface op);

292

293

294

295

296

297

298

299

300

301

302

303

304

305

306

307

308

309

310

311

312

313

314

315

316

317

318

319

320 FailureOr<Operation *> handleInterface(CallOpInterface op);

321

322

323

324

325

326

327

328

329

330

331

332

333

334

335

336

337

338

339

340 FailureOr<Operation *> handleInterface(MemoryEffectOpInterface op);

341

342

343

344

345

346

347

348

349

350

351

352

353

354

355

356

357

358

359

360

361

362

363 FailureOr<Operation *> handleInterface(RegionBranchTerminatorOpInterface op);

364

365

366

367

368

369

370

371

373

374

375

376 template

377 OpTy appendOpResults(OpTy op, ArrayRef types) {

378 return cast(appendOpResults(op.getOperation(), types));

379 }

380

381

382

383

384 LogicalResult deallocate(Block *block);

385

386

387

388

389

390

391

392

393

394

395

396

397 void populateRemainingOwnerships(Operation *op);

398

399

400

401

402 Value materializeMemrefWithGuaranteedOwnership(OpBuilder &builder,

404

405

406

407

408 bool isFunctionWithoutDynamicOwnership(Operation *op);

409

410

411

412

413

414

415

416

417 std::pair<Value, Value>

418 materializeUniqueOwnership(OpBuilder &builder, Value memref, Block *block);

419

420

421

422

423

424 static LogicalResult verifyFunctionPreconditions(FunctionOpInterface op);

425

426

427

428

429

430

431

432

433

434

435

436

437

438

439

440 static LogicalResult verifyOperationPreconditions(Operation *op);

441

442

443

444

445

446

447 static LogicalResult updateFunctionSignature(FunctionOpInterface op);

448

449 private:

450

451

452

454

455

457 };

458

459 }

460

461

462

463

464

465 std::pair<Value, Value>

466 BufferDeallocation::materializeUniqueOwnership(OpBuilder &builder, Value memref,

468

469

471 return state.getMemrefWithUniqueOwnership(builder, memref, block);

472

474 if (!owner)

476

477

478

479 if (auto deallocOpInterface = dyn_cast(owner))

480 return deallocOpInterface.materializeUniqueOwnershipForMemref(

481 state, options, builder, memref);

482

483

484 return state.getMemrefWithUniqueOwnership(builder, memref, block);

485 }

486

487 LogicalResult

488 BufferDeallocation::verifyFunctionPreconditions(FunctionOpInterface op) {

489

490

491 Backedges backedges(op);

492 if (backedges.size()) {

493 op->emitError("Only structured control-flow loops are supported.");

494 return failure();

495 }

496

497 return success();

498 }

499

500 LogicalResult BufferDeallocation::verifyOperationPreconditions(Operation *op) {

501

502

504 return success();

505

506

507

508 if (isa(op))

510 "No deallocation operations must be present when running this pass!");

511

512

513

514

515

516

517

518

519

520 if (!isa(op) &&

522 !isa(op))

524 "ops with unknown memory side effects are not supported");

525

526

528

529

530

531

532

533

534 size_t size = regions.size();

535 if (((size == 1 && !op->getResults().empty()) || size > 1) &&

536 !dyn_cast(op)) {

537 return op->emitError("All operations with attached regions need to "

538 "implement the RegionBranchOpInterface.");

539 }

540

541

542

543

545 return op->emitError("NoTerminator trait is not supported");

546

548

549

550 if (!isa<BranchOpInterface, RegionBranchTerminatorOpInterface>(op) ||

551 (isa(op) &&

552 isa(op)))

553

555 "Terminators must implement either BranchOpInterface or "

556 "RegionBranchTerminatorOpInterface (but not both)!");

557

558

559

561

562 return op->emitError("Terminators with more than one successor "

563 "are not supported!");

564 }

565

566 return success();

567 }

568

569 LogicalResult

570 BufferDeallocation::updateFunctionSignature(FunctionOpInterface op) {

572 op.getFunctionBody().getOps(),

573 [](RegionBranchTerminatorOpInterface op) {

574 return op.getSuccessorOperands(RegionBranchPoint::parent()).getTypes();

575 }));

576 if (!llvm::all_equal(returnOperandTypes))

577 return op->emitError(

578 "there are multiple return operations with different operand types");

579

580 TypeRange resultTypes = op.getResultTypes();

581

582

583

584

585 if (!returnOperandTypes.empty())

586 resultTypes = returnOperandTypes[0];

587

589 op->getContext(), op.getFunctionBody().front().getArgumentTypes(),

590 resultTypes)));

591

592 return success();

593 }

594

595 LogicalResult BufferDeallocation::deallocate(FunctionOpInterface op) {

596

597 if (failed(verifyFunctionPreconditions(op)))

598 return failure();

599

600

602 [&](Block *block) {

603 if (failed(deallocate(block)))

606 });

607 if (result.wasInterrupted())

608 return failure();

609

610

611

612 return updateFunctionSignature(op);

613 }

614

615 LogicalResult BufferDeallocation::deallocate(Block *block) {

617

618

620 state.getLiveMemrefsIn(block, liveMemrefs);

621 for (auto li : liveMemrefs) {

622

623

624

625 if (li.getParentRegion() == block->getParent()) {

626 state.updateOwnership(li, state.getOwnership(li, li.getParentBlock()),

627 block);

628 state.addMemrefToDeallocate(li, block);

629 continue;

630 }

631

632 if (li.getParentRegion()->isProperAncestor(block->getParent())) {

634 state.updateOwnership(li, falseVal, block);

635 }

636 }

637

638 for (unsigned i = 0, e = block->getNumArguments(); i < e; ++i) {

641 continue;

642

643

644

645 if (isa(block->getParentOp()) &&

648 state.updateOwnership(arg, newArg);

649 state.addMemrefToDeallocate(arg, block);

650 continue;

651 }

652

653

655 state.updateOwnership(arg, newArg);

656 state.addMemrefToDeallocate(arg, block);

657 }

658

659

660

661 for (Operation &op : llvm::make_early_inc_range(*block)) {

662 FailureOr<Operation *> result = handleAllInterfaces(&op);

663 if (failed(result))

664 return failure();

665 if (!*result)

666 continue;

667

668 populateRemainingOwnerships(*result);

669 }

670

671

672 return success();

673 }

674

678 newTypes.append(types.begin(), types.end());

683 for (auto [oldRegion, newRegion] :

684 llvm::zip(op->getRegions(), newOp->getRegions()))

685 newRegion.takeBody(oldRegion);

686

690

691 return newOp;

692 }

693

694 FailureOr<Operation *>

695 BufferDeallocation::handleInterface(RegionBranchOpInterface op) {

697

698

699

700

701

702

703

704

705

706

707

708

709

710

713 assert(!regions.empty() && "Must have at least one successor region");

715 op.getEntrySuccessorOperands(regions.front()));

716 unsigned numMemrefOperands = llvm::count_if(entryOperands, isMemref);

717

718

719

721 op->insertOperands(op->getNumOperands(),

723

724 int counter = op->getNumResults();

725 unsigned numMemrefResults = llvm::count_if(op->getResults(), isMemref);

727 RegionBranchOpInterface newOp = appendOpResults(op, ownershipResults);

728

729 for (auto result : llvm::make_filter_range(newOp->getResults(), isMemref)) {

730 state.updateOwnership(result, newOp->getResult(counter++));

731 state.addMemrefToDeallocate(result, newOp->getBlock());

732 }

733

734 return newOp.getOperation();

735 }

736

737 Value BufferDeallocation::materializeMemrefWithGuaranteedOwnership(

739

740 std::pair<Value, Value> newMemrefAndOnwership =

741 materializeUniqueOwnership(builder, memref, block);

742 Value newMemref = newMemrefAndOnwership.first;

743 Value condition = newMemrefAndOnwership.second;

744

745

746

747

749 return newMemref;

750

751

752

753 Value maybeClone =

754 builder

756 memref.getLoc(), condition,

758 builder.createscf::YieldOp(loc, newMemref);

759 },

762 builder.create bufferization::CloneOp(loc, newMemref);

763 builder.create scf::YieldOp(loc, clone);

764 })

765 .getResult(0);

767 state.updateOwnership(maybeClone, trueVal);

768 state.addMemrefToDeallocate(maybeClone, maybeClone.getParentBlock());

769 return maybeClone;

770 }

771

772 FailureOr<Operation *>

773 BufferDeallocation::handleInterface(BranchOpInterface op) {

774 if (op->getNumSuccessors() > 1)

775 return op->emitError("BranchOpInterface operations with multiple "

776 "successors are not supported yet");

777

778 if (op->getNumSuccessors() != 1)

780 "only BranchOpInterface operations with exactly "

781 "one successor are supported yet");

782

783 if (op.getSuccessorOperands(0).getProducedOperandCount() > 0)

784 return op.emitError("produced operands are not supported");

785

786

787

788 Block *block = op->getBlock();

791 if (failed(state.getMemrefsAndConditionsToDeallocate(

792 builder, op.getLoc(), block, memrefs, conditions)))

793 return failure();

794

796 op.getSuccessorOperands(0).getForwardedOperands();

797 state.getMemrefsToRetain(block, op->getSuccessor(0), forwardedOperands,

798 toRetain);

799

800 auto deallocOp = builder.createbufferization::DeallocOp(

801 op.getLoc(), memrefs, conditions, toRetain);

802

803

804

805 state.resetOwnerships(deallocOp.getRetained(), block);

806 for (auto [retained, ownership] :

807 llvm::zip(deallocOp.getRetained(), deallocOp.getUpdatedConditions())) {

808 state.updateOwnership(retained, ownership, block);

809 }

810

811 unsigned numAdditionalReturns = llvm::count_if(forwardedOperands, isMemref);

813 auto additionalConditions =

814 deallocOp.getUpdatedConditions().take_front(numAdditionalReturns);

815 newOperands.append(additionalConditions.begin(), additionalConditions.end());

816 op.getSuccessorOperands(0).getMutableForwardedOperands().assign(newOperands);

817

818 return op.getOperation();

819 }

820

821 FailureOr<Operation *> BufferDeallocation::handleInterface(CallOpInterface op) {

823

824

825

826

827

828 Operation *funcOp = op.resolveCallableInTable(state.getSymbolTable());

829 bool isPrivate = false;

830 if (auto symbol = dyn_cast_or_null(funcOp))

831 isPrivate = symbol.isPrivate() && !symbol.isDeclaration();

832

833

834

835

836

837 if (options.privateFuncDynamicOwnership && isPrivate) {

838 unsigned numMemrefs = llvm::count_if(op->getResults(), isMemref);

839 SmallVector ownershipTypesToAppend(numMemrefs, builder.getI1Type());

840 unsigned ownershipCounter = op->getNumResults();

841 op = appendOpResults(op, ownershipTypesToAppend);

842

843 for (auto result : llvm::make_filter_range(op->getResults(), isMemref)) {

844 state.updateOwnership(result, op->getResult(ownershipCounter++));

845 state.addMemrefToDeallocate(result, result.getParentBlock());

846 }

847

848 return op.getOperation();

849 }

850

851

852

853

855 for (auto result : llvm::make_filter_range(op->getResults(), isMemref)) {

856 state.updateOwnership(result, trueVal);

857 state.addMemrefToDeallocate(result, result.getParentBlock());

858 }

859

860 return op.getOperation();

861 }

862

863 FailureOr<Operation *>

864 BufferDeallocation::handleInterface(MemoryEffectOpInterface op) {

865 auto *block = op->getBlock();

867

868 for (auto operand : llvm::make_filter_range(op->getOperands(), isMemref)) {

870

871

872

873

874

875

876 if (!op->hasAttr(BufferizationDialect::kManualDeallocation))

877 return op->emitError(

878 "memory free side-effect on MemRef value not supported!");

879

880

881

882

883

886 Ownership ownership = state.getOwnership(operand, block);

888 Value ownershipInverted = builder.create arith::XOrIOp(

891 builder.create cf::AssertOp(

892 op.getLoc(), ownershipInverted,

893 "expected that the block does not have ownership");

894 }

895 }

896 }

897

898 for (auto res : llvm::make_filter_range(op->getResults(), isMemref)) {

900 if (allocEffect.has_value()) {

901 if (isaSideEffects::AutomaticAllocationScopeResource(

902 allocEffect->getResource())) {

903

904

905

906

907

908

909

910 state.resetOwnerships(res, block);

911 state.updateOwnership(res, buildBoolValue(builder, op.getLoc(), false));

912 continue;

913 }

914

915 if (op->hasAttr(BufferizationDialect::kManualDeallocation)) {

916

917

918

919 state.resetOwnerships(res, block);

920 state.updateOwnership(res, buildBoolValue(builder, op.getLoc(), false));

921 continue;

922 }

923

924 state.updateOwnership(res, buildBoolValue(builder, op.getLoc(), true));

925 state.addMemrefToDeallocate(res, block);

926 }

927 }

928

929 return op.getOperation();

930 }

931

932 FailureOr<Operation *>

933 BufferDeallocation::handleInterface(RegionBranchTerminatorOpInterface op) {

935

936

937

938

939

940 bool funcWithoutDynamicOwnership =

941 isFunctionWithoutDynamicOwnership(op->getParentOp());

942 if (funcWithoutDynamicOwnership) {

943 for (OpOperand &val : op->getOpOperands()) {

945 continue;

946

947 val.set(materializeMemrefWithGuaranteedOwnership(builder, val.get(),

948 op->getBlock()));

949 }

950 }

951

952

953

954

955

958

962 if (failed(result) || !*result)

963 return result;

964

965

966 if (!funcWithoutDynamicOwnership) {

968 newOperands.append(updatedOwnerships.begin(), updatedOwnerships.end());

969 operands.assign(newOperands);

970 }

971

972 return op.getOperation();

973 }

974

975 bool BufferDeallocation::isFunctionWithoutDynamicOwnership(Operation *op) {

976 auto funcOp = dyn_cast(op);

977 return funcOp && ( options .privateFuncDynamicOwnership ||

978 !funcOp.isPrivate() || funcOp.isExternal());

979 }

980

981 void BufferDeallocation::populateRemainingOwnerships(Operation *op) {

984 continue;

985 if (!state.getOwnership(res, op->getBlock()).isUninitialized())

986 continue;

987

988

989

990

993 continue;

994

995 state.updateOwnership(

996 res, state.getOwnership(operand, operand.getParentBlock()),

998 }

999

1000

1001

1002

1003

1004

1005

1006 if (state.getOwnership(res, op->getBlock()).isUninitialized()) {

1009 }

1010 }

1011 }

1012

1013

1014

1015

1016

1017 namespace {

1018

1019

1020

1021

1022 struct OwnershipBasedBufferDeallocationPass

1023 : public bufferization::impl::OwnershipBasedBufferDeallocationPassBase<

1024 OwnershipBasedBufferDeallocationPass> {

1025 using Base::Base;

1026

1027 void runOnOperation() override {

1029 options.privateFuncDynamicOwnership = privateFuncDynamicOwnership;

1030

1032

1033 auto status = getOperation()->walk([&](func::FuncOp func) {

1034 if (func.isExternal())

1036

1039

1041 });

1042 if (status.wasInterrupted())

1043 signalPassFailure();

1044 }

1045 };

1046

1047 }

1048

1049

1050

1051

1052

1056

1057 BufferDeallocation deallocation(op, options, symbolTables);

1058

1059

1060 return deallocation.deallocate(op);

1061 }

static bool isMemref(Value v)

static Value buildBoolValue(OpBuilder &builder, Location loc, bool value)

static bool hasBufferSemantics(Operation *op)

Return "true" if the given op has buffer semantics.

static bool hasNeitherAllocateNorFreeSideEffect(Operation *op)

Return "true" if the given op is guaranteed to have neither "Allocate" nor "Free" side effects.

static llvm::ManagedStatic< PassManagerOptions > options

This class represents an argument of a Block.

Location getLoc() const

Return the location for this argument.

Block represents an ordered list of Operations.

BlockArgument getArgument(unsigned i)

unsigned getNumArguments()

Region * getParent() const

Provide a 'getParent' method for ilist_node_with_parent methods.

SuccessorRange getSuccessors()

BlockArgument addArgument(Type type, Location loc)

Add one value to the argument list.

OpListType & getOperations()

bool isEntryBlock()

Return if this block is the entry block in the parent region.

Operation * getParentOp()

Returns the closest surrounding operation that contains this block.

BoolAttr getBoolAttr(bool value)

This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...

This class provides a mutable adaptor for a range of operands.

OperandRange getAsOperandRange() const

Explicit conversion to an OperandRange.

void assign(ValueRange values)

Assign this range to the given values.

RAII guard to reset the insertion point of the builder when destroyed.

This class helps build Operations.

static OpBuilder atBlockBegin(Block *block, Listener *listener=nullptr)

Create a builder and set the insertion point to before the first operation in the block but still ins...

void setInsertionPoint(Block *block, Block::iterator insertPoint)

Set the insertion point to the specified location.

Operation * create(const OperationState &state)

Creates an operation given the fields represented as an OperationState.

Operation * insert(Operation *op)

Insert the given operation at the current insertion point and return it.

This class represents an operand of an operation.

This trait indicates that the memory effects of an operation includes the effects of operations neste...

This class provides the API for ops that are known to be terminators.

This class indicates that the regions associated with this op don't have terminators.

This class implements the operand iterators for the Operation class.

Operation is the basic unit of execution within MLIR.

DictionaryAttr getAttrDictionary()

Return all of the attributes on this operation as a DictionaryAttr.

bool hasTrait()

Returns true if the operation was registered with a particular trait, e.g.

unsigned getNumRegions()

Returns the number of regions held by this operation.

Location getLoc()

The source location the operation was defined or derived from.

static Operation * create(Location location, OperationName name, TypeRange resultTypes, ValueRange operands, NamedAttrList &&attributes, OpaqueProperties properties, BlockRange successors, unsigned numRegions)

Create a new Operation with the specific fields.

InFlightDiagnostic emitError(const Twine &message={})

Emit an error about fatal conditions with this operation, reporting up to any diagnostic handlers tha...

Block * getBlock()

Returns the operation block that contains this operation.

MutableArrayRef< Region > getRegions()

Returns the regions held by this operation.

OperationName getName()

The name of an operation is the key identifier for it.

result_type_range getResultTypes()

operand_range getOperands()

Returns an iterator on the underlying Value's.

void replaceAllUsesWith(ValuesT &&values)

Replace all uses of results of this operation with the provided 'values'.

SuccessorRange getSuccessors()

result_range getResults()

OpaqueProperties getPropertiesStorage()

Returns the properties storage.

void erase()

Remove this operation from its parent block and delete it.

unsigned getNumResults()

Return the number of results held by this operation.

static constexpr RegionBranchPoint parent()

Returns an instance of RegionBranchPoint representing the parent operation.

This class contains a list of basic blocks and a link to the parent operation it is attached to.

This class represents a collection of SymbolTables.

This class provides an abstraction over the various different ranges of value types.

This class represents an instance of an SSA value in the MLIR system, representing a computable value...

Type getType() const

Return the type of this value.

Block * getParentBlock()

Return the Block in which this Value is defined.

Location getLoc() const

Return the location of this value.

Operation * getDefiningOp() const

If this value is the result of an operation, return the operation that defines it.

static WalkResult advance()

static WalkResult interrupt()

This class collects all the state that we need to perform the buffer deallocation pass with associate...

This class is used to track the ownership of values.

bool isUnique() const

Check if this ownership value is in the 'Unique' state.

Value getIndicator() const

If this ownership value is in 'Unique' state, this function can be used to get the indicator paramete...

FailureOr< Operation * > insertDeallocOpForReturnLike(DeallocationState &state, Operation *op, ValueRange operands, SmallVectorImpl< Value > &updatedOperandOwnerships)

Insert a bufferization.dealloc operation right before op which has to be a terminator without any suc...

LogicalResult deallocateBuffersOwnershipBased(FunctionOpInterface op, DeallocationOptions options, SymbolTableCollection &symbolTables)

Run the ownership-based buffer deallocation.

Include the generated interface declarations.

bool matchPattern(Value value, const Pattern &pattern)

Entry point for matching a pattern over a Value.

InFlightDiagnostic emitError(Location loc)

Utility method to emit an error message using this location.

detail::constant_int_predicate_matcher m_One()

Matches a constant scalar / vector splat / tensor splat integer one.

Operation * clone(OpBuilder &b, Operation *op, TypeRange newResultTypes, ValueRange newOperands)

auto get(MLIRContext *context, Ts &&...params)

Helper method that injects context only if needed, this helps unify some of the attribute constructio...

This iterator enumerates elements according to their dominance relationship.

The following effect indicates that the operation allocates from some resource.

The following effect indicates that the operation frees some resource that has been allocated.

Options for BufferDeallocationOpInterface-based buffer deallocation.