[mlir][bufferization] Add optional `copy` operand to AllocTensorOp

[springerm]

Summary

If copy is specified, the newly allocated buffer is initialized with the given contents. Also add an optional escape attribute to indicate whether the buffer of the tensor may be returned from the parent block (aka. "escape") after bufferization.

This change is in preparation of connecting One-Shot Bufferize to the sparse compiler.

Depends On D126814

[aartbik]

mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td

41

typo

42

rephrase: should be deallocated elsewhere?

mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp

300

type needs to be verify()-ed as well

[springerm]

address comments

[springerm]

update

mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td

 class Bufferization_Op<string mnemonic, list<Trait> traits = []>
     : Op<Bufferization_Dialect, mnemonic, traits>;
 
 //===----------------------------------------------------------------------===//
 // AllocTensorOp
 //===----------------------------------------------------------------------===//
 
 def Bufferization_AllocTensorOp : Bufferization_Op<"alloc_tensor",
-    [BufferizableOpInterface,
+    [AttrSizedOperandSegments, BufferizableOpInterface,
      DeclareOpInterfaceMethods<ReifyRankedShapedTypeOpInterface>]> {
   let summary = "buffer allocation in tensor land";
 
   let description = [{
     `bufferization.alloc_tensor` materializes an uninitialized tensor with a
     given shape (dynamic or static). It always bufferizes to a new buffer
-    allocation of the given shape. Reading from the result of an `alloc_tensor`
-    op yields an undefined value.
+    allocation of the given shape. The optional `copy` operand specifies the
+    contents of the tensors. If no `copy` operand is specified, reading from the
+    result of an `alloc_tensor` op yields an undefined value.
+
+    If `copy` is specified, no dynamic sizes should be passed, since they are
+    the same as the dynamic sizes of the `copy` operand.
+
+    The optional `escape` attribute indicates whether the buffer escapes the
+    parent block or not. In the latter case, the buffer is deallocated at the
+    of the block (during bufferization). In the former case, the buffer is not
+    deallocated and must be deallocated through some other mechanism.
 
     `alloc_tensor` is a helper op for bufferization. The operation is provided
     as an anchor that marks the beginning of a new tensor SSA use-def chain. It
     can be used to control in-place bufferization decisions during One-Shot
     Bufferize: The bufferized result of a `bufferization.alloc_tensor` does not
     alias with any other buffer, so it can be used to resolve read-after-write
     conflicts that would have been introduced by the in-place bufferization of
     another op.
 
     Both dense and sparse tensor types are supported. The result of a
     `bufferization.alloc_tensor` is a tensor value that can be used like any
     other tensor value. In practice, it is often used as the "out" operand of
     another op. E.g.:
 
     ```mlir
     %c = bufferization.alloc_tensor [%d1, %d2] : tensor<?x?xf32, #SparseMatrix>
     %0 = linalg.matmul
       ins(%a, %b: tensor<?x?xf32, #SparseMatrix>, tensor<?x?xf32, #SparseMatrix>)
       outs(%c: tensor<?x?xf32, #SparseMatrix>) -> tensor<?x?xf32, #SparseMatrix>
     ```
   }];
 
-  let arguments = (ins Variadic<Index>:$dynamicSizes);
+  let arguments = (ins Variadic<Index>:$dynamicSizes,
+                       Optional<AnyTensor>:$copy,
+                       OptionalAttr<BoolAttr>:$escape);
 
   let results = (outs AnyTensor:$result);
 
-  let assemblyFormat = "`(`$dynamicSizes`)` attr-dict `:` type($result)";
-
   let extraClassDeclaration = [{
     LogicalResult bufferize(RewriterBase &rewriter, BufferizationState &state);
 
-    bool isMemoryWrite(OpResult opResult, const AnalysisState &state) const {
-      // AllocTensorOps allocate but do not write.
-      return false;
-    }
+    bool isMemoryWrite(OpResult opResult, const AnalysisState &state);
+
+    bool bufferizesToMemoryRead(OpOperand &opOperand,
+                                const AnalysisState &state);
+
+    bool bufferizesToMemoryWrite(OpOperand &opOperand,
+                                 const AnalysisState &state);
+
+    SmallVector<OpResult> getAliasingOpResult(
+        OpOperand &opOperand, const AnalysisState &state);
 
     RankedTensorType getType() {
       return getResult().getType().cast<RankedTensorType>();
     }
 
     // Return true if the size of the tensor is dynamic at `idx`
     bool isDynamicDim(unsigned idx) {
       return getType().isDynamicDim(idx);
     }
 
     // Return the argument position that contains the dynamic size of
     // the tensor at dimension `idx`. Asserts that the shape is
     // dynamic at that `idx`.
     unsigned getIndexOfDynamicSize(unsigned idx) {
+      assert(!copy() && "no dim sizes specified when copying a tensor");
       assert(isDynamicDim(idx) && "expected dynamic size");
       ArrayRef<int64_t> shape = getType().getShape();
       return std::count_if(
           shape.begin(), shape.begin() + idx,
           [&](int64_t size) { return ShapedType::isDynamic(size); });
     }
 
     // Return the Value of the dynamic size of the tensor at dimension
     // `idx`. Asserts that the shape is dynamic at that `idx.
-    Value getDynamicSize(unsigned idx) {
-      return getOperand(getIndexOfDynamicSize(idx));
-    }
+    Value getDynamicSize(OpBuilder &b, unsigned idx);
 
     // Assert that the size of the result tensor is static at `idx`
     // and return the shape.
     int64_t getStaticSize(unsigned idx) {
       assert(!isDynamicDim(idx) && "expected static size");
       return getType().getShape()[idx];
     }
   }];
 
+  let builders = [
+    // Build an op without `copy` operand and `escape` attribute.
+    OpBuilder<(ins "RankedTensorType":$type, "ValueRange":$dynamicSizes)>,
+
+    // Build an op without `escape` attribute.
+    OpBuilder<(ins "RankedTensorType":$type, "ValueRange":$dynamicSizes,
+                   "Value":$copy)>,
+
+    // Build an op with `copy` and `escape` attribute.
+    OpBuilder<(ins "RankedTensorType":$type, "ValueRange":$dynamicSizes,
+                   "Value":$copy, "bool":$escape)>,
+  ];
+
   let hasCanonicalizer = 1;
+  let hasCustomAssemblyFormat = 1;
   let hasVerifier = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // CloneOp
 //===----------------------------------------------------------------------===//
 
 def Bufferization_CloneOp : Bufferization_Op<"clone", [

mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp

 //===----------------------------------------------------------------------===//
 
 LogicalResult AllocTensorOp::bufferize(RewriterBase &rewriter,
                                        BufferizationState &state) {
   // Nothing to do for dead AllocTensorOps.
   if (getOperation()->getUses().empty())
     return success();
 
-  FailureOr<Value> alloc = state.createAlloc(rewriter, getLoc(), getResult());
+  Optional<bool> dealloc = llvm::None;
+  if (escape().hasValue())
+    dealloc = !*escape();
+  FailureOr<Value> alloc =
+      state.createAlloc(rewriter, getLoc(), getResult(), dealloc);
   if (failed(alloc))
     return failure();
+  if (copy()) {
+    FailureOr<Value> copyValueBuffer = state.getBuffer(
+        rewriter, getOperation()->getOpOperand(getNumOperands() - 1));
+    if (failed(copyValueBuffer))
+      return failure();
+    if (failed(state.getOptions().createMemCpy(rewriter, getLoc(),
+                                               *copyValueBuffer, *alloc)))
+      return failure();
+  }
   replaceOpWithBufferizedValues(rewriter, getOperation(), *alloc);
   return success();
 }
 
+bool AllocTensorOp::isMemoryWrite(OpResult opResult,
+                                  const AnalysisState &state) {
+  // AllocTensorOps do not write unless they have a `copy` value.
+  return static_cast<bool>(copy());
+}
+
+bool AllocTensorOp::bufferizesToMemoryRead(OpOperand &opOperand,
+                                           const AnalysisState &state) {
+  assert(opOperand.getOperandNumber() == getNumOperands() - 1 &&
+         "expected copy operand");
+  return true;
+}
+
+bool AllocTensorOp::bufferizesToMemoryWrite(OpOperand &opOperand,
+                                            const AnalysisState &state) {
+  assert(opOperand.getOperandNumber() == getNumOperands() - 1 &&
+         "expected copy operand");
+  return false;
+}
+
+SmallVector<OpResult>
+AllocTensorOp::getAliasingOpResult(OpOperand &opOperand,
+                                   const AnalysisState &state) {
+  // This is a new allocation. It does not alias with any other buffer.
+  return {};
+}
+
 LogicalResult AllocTensorOp::verify() {
-  if (getType().getNumDynamicDims() !=
+  if (copy() && !dynamicSizes().empty())
+    return emitError("dynamic sizes not needed when copying a tensor");
+  if (!copy() && getType().getNumDynamicDims() !=
                      static_cast<int64_t>(dynamicSizes().size()))
     return emitError("expected ")
            << getType().getNumDynamicDims() << " dynamic sizes";
+  if (copy() && copy().getType() != getType())
+    return emitError("expected that `copy` and return type match");
   return success();
 }
 
+void AllocTensorOp::build(OpBuilder &builder, OperationState &result,
+                          RankedTensorType type, ValueRange dynamicSizes) {
+  build(builder, result, type, dynamicSizes, /*copy=*/Value(),
+        /*escape=*/BoolAttr());
+}
+
+void AllocTensorOp::build(OpBuilder &builder, OperationState &result,
+                          RankedTensorType type, ValueRange dynamicSizes,
+                          Value copy) {
+  build(builder, result, type, dynamicSizes, copy, /*escape=*/BoolAttr());
+}
+
+void AllocTensorOp::build(OpBuilder &builder, OperationState &result,
+                          RankedTensorType type, ValueRange dynamicSizes,
+                          Value copy, bool escape) {
+  build(builder, result, type, dynamicSizes, copy, builder.getBoolAttr(escape));
+}
+
 namespace {
 /// Change the type of the result of a `bufferization.alloc_tensor` by making
 /// the result type statically sized along dimension that in the original
 /// operation where defined as dynamic, but the size was defined using a
 /// `constant` op. For example:
 ///
 ///  %c5 = arith.constant 5: index
 ///  %0 = bufferization.alloc_tensor(%arg0, %c5) : tensor<?x?xf32>
 ///
 ///  to
 ///
 ///  %0 = bufferization.alloc_tensor(%arg0) : tensor<?x5xf32>
 struct ReplaceStaticShapeDims : OpRewritePattern<AllocTensorOp> {
   using OpRewritePattern<AllocTensorOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(AllocTensorOp op,
                                 PatternRewriter &rewriter) const override {
+    if (op.copy())
+      return failure();
     SmallVector<int64_t> newShape = llvm::to_vector(op.getType().getShape());
     SmallVector<Value> newDynamicSizes;
     unsigned int dynValCounter = 0;
     for (int64_t i = 0; i < op.getType().getRank(); ++i) {
       if (!op.isDynamicDim(i))
         continue;
       Value value = op.dynamicSizes()[dynValCounter++];
       APInt intVal;
       if (matchPattern(value, m_ConstantInt(&intVal))) {
         newShape[i] = intVal.getSExtValue();
       } else {
         newDynamicSizes.push_back(value);
       }
     }
     RankedTensorType newType = RankedTensorType::get(
         newShape, op.getType().getElementType(), op.getType().getEncoding());
     if (newType == op.getType())
       return failure();
-    auto newOp =
-        rewriter.create<AllocTensorOp>(op.getLoc(), newType, newDynamicSizes);
+    auto newOp = rewriter.create<AllocTensorOp>(
+        op.getLoc(), newType, newDynamicSizes, /*copy=*/Value(),
+        /*escape=*/op.escapeAttr());
     rewriter.replaceOpWithNewOp<tensor::CastOp>(op, op.getType(), newOp);
     return success();
   }
 };
 
 struct FoldDimOfAllocTensorOp : public OpRewritePattern<tensor::DimOp> {
   using OpRewritePattern<tensor::DimOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(tensor::DimOp dimOp,
                                 PatternRewriter &rewriter) const override {
     Optional<int64_t> maybeConstantIndex = dimOp.getConstantIndex();
     auto allocTensorOp = dimOp.source().getDefiningOp<AllocTensorOp>();
     if (!allocTensorOp || !maybeConstantIndex)
       return failure();
     if (!allocTensorOp.getType().isDynamicDim(*maybeConstantIndex))
       return failure();
-    rewriter.replaceOp(dimOp,
-                       allocTensorOp.getDynamicSize(*maybeConstantIndex));
+    rewriter.replaceOp(
+        dimOp, allocTensorOp.getDynamicSize(rewriter, *maybeConstantIndex));
     return success();
   }
 };
 } // namespace
 
 void AllocTensorOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                 MLIRContext *ctx) {
   results.add<FoldDimOfAllocTensorOp, ReplaceStaticShapeDims>(ctx);
 }
 
 LogicalResult AllocTensorOp::reifyResultShapes(
     OpBuilder &builder, ReifiedRankedShapedTypeDims &reifiedReturnShapes) {
   auto shapes = llvm::to_vector<4>(llvm::map_range(
       llvm::seq<int64_t>(0, getType().getRank()), [&](int64_t dim) -> Value {
         if (isDynamicDim(dim))
-          return getDynamicSize(dim);
+          return getDynamicSize(builder, dim);
         return builder.create<arith::ConstantIndexOp>(getLoc(),
                                                       getStaticSize(dim));
       }));
   reifiedReturnShapes.emplace_back(std::move(shapes));
   return success();
 }
 
+ParseResult AllocTensorOp::parse(OpAsmParser &parser, OperationState &result) {
+  SmallVector<OpAsmParser::UnresolvedOperand> dynamicSizesOperands;
+  if (parser.parseLParen() || parser.parseOperandList(dynamicSizesOperands) ||
+      parser.parseRParen())
+    return failure();
+  ParseResult copyKeyword = parser.parseOptionalKeyword("copy");
+  OpAsmParser::UnresolvedOperand copyOperand;
+  if (copyKeyword.succeeded())
+    if (parser.parseLParen() || parser.parseOperand(copyOperand) ||
+        parser.parseRParen())
+      return failure();
+  if (parser.parseOptionalAttrDict(result.attributes) || parser.parseColon())
+    return failure();
+
+  TensorType type;
+  if (parser.parseCustomTypeWithFallback(type))
+    return failure();
+  result.addTypes(type);
+
+  Type indexType = parser.getBuilder().getIndexType();
+  if (parser.resolveOperands(dynamicSizesOperands, indexType, result.operands))
+    return failure();
+  if (copyKeyword.succeeded())
+    if (parser.resolveOperand(copyOperand, type, result.operands))
+      return failure();
+  result.addAttribute(AllocTensorOp::getOperandSegmentSizeAttr(),
+                      parser.getBuilder().getI32VectorAttr(
+                          {static_cast<int32_t>(dynamicSizesOperands.size()),
+                           static_cast<int32_t>(copyKeyword.succeeded())}));
+  return success();
+}
+
+void AllocTensorOp::print(OpAsmPrinter &p) {
+  p << "(" << dynamicSizes() << ")";
+  if (copy())
+    p << " copy(" << copy() << ")";
+  p.printOptionalAttrDict((*this)->getAttrs(), /*elidedAttrs=*/{
+                              AllocTensorOp::getOperandSegmentSizeAttr()});
+  p << " : ";
+  auto type = result().getType();
+  if (auto validType = type.dyn_cast<::mlir::TensorType>())
+    p.printStrippedAttrOrType(validType);
+  else
+    p << type;
+}
+
+Value AllocTensorOp::getDynamicSize(OpBuilder &b, unsigned idx) {
+  assert(isDynamicDim(idx) && "expected dynamic dim");
+  if (copy())
+    return b.create<tensor::DimOp>(getLoc(), copy(), idx);
+  return getOperand(getIndexOfDynamicSize(idx));
+}
+
 //===----------------------------------------------------------------------===//
 // CloneOp
 //===----------------------------------------------------------------------===//
 
 void CloneOp::getEffects(
     SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
         &effects) {
   effects.emplace_back(MemoryEffects::Read::get(), input(),

mlir/python/mlir/dialects/_bufferization_ops_ext.py

 
 
 class AllocTensorOp:
   """Extends the bufferization.alloc_tensor op."""
 
   def __init__(self,
                tensor_type: Type,
                dynamic_sizes: Sequence[Value],
+               copy: Value,
+               escape: BoolAttr,
                *,
                loc=None,
                ip=None):
     """Constructs an `alloc_tensor` with static and/or dynamic sizes."""
     context = get_default_loc_context(loc)
+    attributes = {}
+    if escape:
+      attributes["escape"] = escape
     op = self.build_generic(
         results=[tensor_type],
-        operands=dynamic_sizes,
-        attributes={},
+        operands=[dynamic_sizes, copy],
+        attributes=attributes,
         loc=loc,
         ip=ip)
     OpView.__init__(self, op)

mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize.mlir

func.func @select_different_tensors(%t: tensor<?xf32>, %sz: index, %c: i1) -> tensor<?xf32> {
   %0 = bufferization.alloc_tensor(%sz) : tensor<?xf32>
 
   // A cast must be inserted because %t and %0 have different memref types.
   // CHECK: %[[casted:.*]] = memref.cast %[[alloc]] : memref<?xf32> to memref<?xf32, #{{.*}}>
   // CHECK: arith.select %{{.*}}, %[[casted]], %[[m]]
   %1 = arith.select %c, %0, %t : tensor<?xf32>
   return %1 : tensor<?xf32>
 }
+
+// -----
+
+// CHECK-LABEL: func @alloc_tensor_with_copy(
+//  CHECK-SAME:     %[[t:.*]]: tensor<5xf32>)
+// TODO: Add a test case with dynamic dim size. This is not possible at the
+// moment because this would create a tensor op during bufferization. That is
+// currently forbidden.
+func.func @alloc_tensor_with_copy(%t: tensor<5xf32>) -> tensor<5xf32> {
+  // CHECK: %[[m:.*]] = bufferization.to_memref %[[t]]
+  // CHECK: %[[alloc:.*]] = memref.alloc() {{.*}} : memref<5xf32>
+  // CHECK: memref.copy %[[m]], %[[alloc]]
+  %0 = bufferization.alloc_tensor() copy(%t) : tensor<5xf32>
+  // CHECK: %[[r:.*]] = bufferization.to_tensor %[[alloc]]
+  // CHECK: memref.dealloc %[[alloc]]
+  // CHECK: return %[[r]]
+  return %0 : tensor<5xf32>
+}
+

mlir/test/Dialect/Bufferization/canonicalize.mlir

 //  CHECK-SAME:   %[[ARG0:.+]]: tensor<4x6x16x32xi8>
 func.func @tensor_cast_to_memref(%arg0 : tensor<4x6x16x32xi8>) ->
   memref<?x?x16x32xi8> {
   %0 = tensor.cast %arg0 : tensor<4x6x16x32xi8> to tensor<?x?x16x32xi8>
   %1 = bufferization.to_memref %0 : memref<?x?x16x32xi8>
   return %1 : memref<?x?x16x32xi8>
 }
 // CHECK:   %[[M:.+]] = bufferization.to_memref %[[ARG0]] : memref<4x6x16x32xi8>
-// CHECK:   %[[M1:.+]] = memref.cast %[[M]] 
+// CHECK:   %[[M1:.+]] = memref.cast %[[M]]
 // CHECK-SAME: memref<4x6x16x32xi8> to memref<?x?x16x32xi8>
 // CHECK:   return %[[M1]] : memref<?x?x16x32xi8>
 
 // -----
 
 // Folding of memref.load(to_memref(%v, %idxs)) -> tensor.extract(%v, %idx)
 // CHECK-LABEL: func @load_from_buffer_cast(
 func.func @load_from_buffer_cast(%arg0: index, %arg1: index,

mlir/test/Dialect/Bufferization/invalid.mlir

 // RUN: mlir-opt %s -split-input-file -verify-diagnostics
 
-func.func @alloc_tensor_err(%arg0 : index)
+func.func @alloc_tensor_missing_dims(%arg0: index)
 {
   // expected-error @+1 {{expected 2 dynamic sizes}}
-  %1 = bufferization.alloc_tensor(%arg0) : tensor<4x?x?x5xf32>
+  %0 = bufferization.alloc_tensor(%arg0) : tensor<4x?x?x5xf32>
+  return
+}
+
+// -----
+
+// expected-note @+1 {{prior use here}}
+func.func @alloc_tensor_type_mismatch(%t: tensor<?xf32>) {
+  // expected-error @+1{{expects different type than prior uses: 'tensor<5xf32>' vs 'tensor<?xf32>'}}
+  %0 = bufferization.alloc_tensor() copy(%t) : tensor<5xf32>
+  return
+}
+
+// -----
+
+func.func @alloc_tensor_copy_and_dims(%t: tensor<?xf32>, %sz: index) {
+  // expected-error @+1{{dynamic sizes not needed when copying a tensor}}
+  %0 = bufferization.alloc_tensor(%sz) copy(%t) : tensor<?xf32>
+  return
+}
+
+// -----
+
+func.func @alloc_tensor_invalid_escape_attr(%sz: index) {
+  // expected-error @+1{{op attribute 'escape' failed to satisfy constraint: bool attribute}}
+  %0 = bufferization.alloc_tensor(%sz) {escape = 5} : tensor<?xf32>
   return
 }

mlir/test/Dialect/Bufferization/ops.mlir

func.func @test_to_memref(%arg0: tensor<?xi64>, %arg1: tensor<*xi64>)
   return %0, %1 : memref<?xi64, affine_map<(d0) -> (d0 + 7)>>, memref<*xi64, 1>
 }
 
 // CHECK-LABEL: func @test_to_tensor
 func.func @test_to_tensor(%buf : memref<2xf32>) -> tensor<2xf32> {
   %tensor = bufferization.to_tensor %buf : memref<2xf32>
   return %tensor : tensor<2xf32>
 }
+
+// CHECK-LABEL: func @test_alloc_tensor_op
+func.func @test_alloc_tensor_op(%t: tensor<?x5xf32>, %sz: index)
+  -> tensor<?x5xf32>
+{
+  // CHECK: bufferization.alloc_tensor(%{{.*}}) : tensor<?x5xf32>
+  %0 = bufferization.alloc_tensor(%sz) : tensor<?x5xf32>
+  // CHECK: bufferization.alloc_tensor() copy(%{{.*}}) : tensor<?x5xf32>
+  %1 = bufferization.alloc_tensor() copy(%t) : tensor<?x5xf32>
+  // CHECK: bufferization.alloc_tensor() : tensor<5x6xf32>
+  %2 = bufferization.alloc_tensor() : tensor<5x6xf32>
+  // CHECK: bufferization.alloc_tensor(%{{.*}}, %{{.*}}) : tensor<?x?xf32>
+  %3 = bufferization.alloc_tensor(%sz, %sz) : tensor<?x?xf32>
+  // CHECK: bufferization.alloc_tensor() copy(%{{.*}}) {escape = true} : tensor<?x5xf32>
+  %4 = bufferization.alloc_tensor() copy(%t) {escape = true} : tensor<?x5xf32>
+  // CHECK: bufferization.alloc_tensor() copy(%{{.*}}) {escape = false} : tensor<?x5xf32>
+  %5 = bufferization.alloc_tensor() copy(%t) {escape = false} : tensor<?x5xf32>
+  return %1 : tensor<?x5xf32>
+}

mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco.py

if self.dst_format is None or self.dst_format.rank() == 0:
       tensor = linalg.InitTensorOp(self.dst_dims, ir_type).result
       zero = arith.ConstantOp(ir_type, 0.0)
       return linalg.fill(zero, outs=[tensor])
 
     # Initialize the sparse tensor.
     mlir_type = _mlir_tensor_type(self.dst_dtype, self.dst_dims,
                                   self.dst_format.mlir_tensor_attr())
     index_type = ir.IndexType.get()
-    return bufferization.AllocTensorOp(mlir_type, [])
+    return bufferization.AllocTensorOp(mlir_type, [], None, None)
 
 
 class _Stats:
   """Information to describe how a tensor expression is implemented.
 
   Currently, we only record the temporary tensors introduced for splitting the
   original expression.
   """