Skip to content

[AutoDiff] Refine pullback generation logic for reference-based types #67867

Open
@asl

Description

@asl

Currently we are generating new adjoint buffers for many operations on class references that do look like projections. For example:

  /// Handle `upcast` instruction.
  ///   Original: y = upcast x
  ///    Adjoint: adj[x] += adj[y]
  ///             (assuming adj[x] and adj[y] have the same type)
  void visitUpcastInst(UpcastInst *ui) {
    auto *bb = ui->getParent();
    assert(ui->getOperand()->getType().isObject());
    assert(getRemappedTangentType(ui->getOperand()->getType()) ==
               getRemappedTangentType(ui->getType()) &&
           "Operand/result must have the same `TangentVector` type");
    switch (getTangentValueCategory(ui)) {
    case SILValueCategory::Object: {
      auto adj = getAdjointValue(bb, ui);
      addAdjointValue(bb, ui->getOperand(), adj, ui->getLoc());
      break;
    }
    case SILValueCategory::Address: {
      auto adjDest = getAdjointBuffer(bb, ui);
      addToAdjointBuffer(bb, ui->getOperand(), adjDest, ui->getLoc());
      builder.emitZeroIntoBuffer(ui->getLoc(), adjDest, IsNotInitialization);
      break;
    }
    }
  }

Note that here x and y represent essentially the same reference-counted objects, so they may share the adjoint buffer instead of allocating separate adjoint buffers for x and y respectfully. Accumulation and zeroing could be omitted as well.

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions