onnx-mlir

onnx-mlir copied to clipboard

In onnx-mlir, we have --constprop-onnx pass that is used to do constant propagation, e.g. when all inputs of an operation are constants, the pass will replace the operation by a new constant computed during compile time. More info about --constprop-onnx can be found here: https://github.com/onnx/onnx-mlir/blob/main/docs/ConstPropagationPass.md.

All constants in onnx-mlir are represented by DenseElementsAttr. It's important to note that DenseElementsAttr are immortal objects living in MLIR context (See 1 and 2 for more details). Once we create them, we cannot delete it in the same context.

In --constprop-onnx, we maintain a buffer pool to manage all intermediate constants and DenseElementsAttr are only created for the final constant. By this way, we was able to reduce the memory consumption during compile time for resnet152 from 7.1 GB to 0.6GB (See https://github.com/onnx/onnx-mlir/pull/530).

However, there are two issues with the current --constprop-onnx:

1. It is painful to work with the buffer pool. We need to mange memory manually. Furthermore, the buffer pool exists in the pass, so it's not sharable across multiple calls of --constprop-onnx
2. The code in --constprop-onnx is nontrivial to use for the fold mechanism in MLIR that can be invoked directly anywhere with an OpBuilder via OpBuilder::createOrFold to produce a constant.

To overcome the issue of immortal attributes, MLIR recently introduces DenseResourceElementsAttr that works on resource blobs. The important point is that resource blobs supports user-defined deleter that will free the underlying data when the data is not used.

We are using DenseResourceElementsAttr in onnx-mlir to store stickified constants for NNPA. See https://github.com/onnx/onnx-mlir/blob/main/src/Accelerators/NNPA/Transform/ZHigh/ZHighConstPropagation.cpp#L100.

I plan to use DenseResourceElementsAttr for --constprop-onnx. Any comments are welcome and let me know if someone is interested in working on this.

@sorenlassen FYI.

@tungld I don't want to distract from the main conversation too much, but I'm making my way into ConstProp.cpp for a separate issue and based on some background research it appears the following is dead code because nothing else calls it:

/// A helper function to get a value of a given type from an attribute.
template <typename T>
T getAttrValue(Attribute attr) {
  llvm_unreachable("unknown operation");
}

template <>
ATTRIBUTE(unused)
double getAttrValue(Attribute attr) {
  return attr.cast<FloatAttr>().getValueAsDouble();
}

template <>
ATTRIBUTE(unused)
float getAttrValue(Attribute attr) {
  return (float)attr.cast<FloatAttr>().getValueAsDouble();
}

template <>
ATTRIBUTE(unused)
int64_t getAttrValue(Attribute attr) {
  return attr.cast<IntegerAttr>().getInt();
}

template <>
ATTRIBUTE(unused)
int32_t getAttrValue(Attribute attr) {
  return attr.cast<IntegerAttr>().getInt();
}
}

Would you mind I put up a PR to delete it?

thanks for the explanation of DenseResourceElementsAttr, it looks really interesting

resource blobs supports user-defined deleter that will free the underlying data when the data is not used

can you outline how this would work? I looked closely at DenseResourceElementsAttr and AsmResourceBlob and I couldn't figure out where or how AsmResourceBlobs get deleted

Would you mind I put up a PR to delete it?

deleting this unused code sounds good to me, please go ahead and submit a PR

can you outline how this would work? I looked closely at DenseResourceElementsAttr and AsmResourceBlob and I couldn't figure out where or how AsmResourceBlobs get deleted

It looks to me DialectResourceBlobManager is for us to mange AsmResourceBlobs in a dialect. https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/IR/DialectResourceBlobManager.h. Otherwise, we should manage them manually.

From this commit: https://github.com/llvm/llvm-project/commit/5f58e14b36edbdacef86a2f3fc192b5720b2ba62

The DialectResourceBlobManager class provides functionality for managing resource blobs
in a generic, dialect-agnostic fashion. In addition to this class, a dialect interface and custom
resource handle are provided to simplify referencing and interacting with the manager. These
classes intend to simplify the work required for dialects that want to manage resource blobs
during compilation, such as for large elements attrs.  The old manager for the resource example
in the test dialect has been updated to use this, which provides and cleaner and more consistent API.

This commit also adds new HeapAsmResourceBlob and ImmortalAsmResourceBlob to simplify
creating resource blobs in common scenarios.

I cannot find many examples of using DialectResourceBlobManager in MLIR, except ByteCode: https://github.com/llvm/llvm-project/commit/6ab2bcffe45e660a68493e6a7cd04b6f05da51dc. Not sure if it helps.

I found this LLVM discourse thread where you describe onnx-mlir's solution: https://discourse.llvm.org/t/mlircontext-doesnt-free-denseelementsattributes-and-bloats-memory/3691/24

sharing it here because the thread contains a good discussion that's worth reading

@tungld I created a draft PR with a custom attribute DisposableElementsAttr with garbage collection and a few other features to optimize constant propagation: https://github.com/sorenlassen/onnx-mlir-einsum/pull/1

(I created the PR on my onnx-mlir-einsum fork while I was drafting it. I need to figure out how to move the PR to become visible on the main onnx-mlir repo.)

Update: I moved the PR here: https://github.com/onnx/onnx-mlir/pull/1874

@sorenlassen thanks! I will take a look at it hopefully after the Thanksgiving day.

let's continue the quest for constant folding in the new issue #2143