Skip to content

Conversation

@shelkesagar29
Copy link
Collaborator

@shelkesagar29 shelkesagar29 commented Nov 5, 2024

This PR moves internal mlir-tensorrt project changes
to OSS.
The following commits are included,

NFC: Move "stablehlo-input-preprocessing" and "stablehlo-raise-qdq" passes under the StableHloExt folder
NFC: Don't unnecessarily restrict StableHLO preprocessing passes to 'func' ops

NFC: [StableHloExt] Consolidate scatter/gather utilities into common translation unit and under stablehlo_ext namespace
This change consolidates utility functions related to scatter/gather ops into the
"ScatterGatherUtils.(h|cpp)" translation unit and changes all the functions in those
files to be under the 'stablehlo_ext' namespace.

This is part one of a series of changes that will attempt to consolidate the stablehlo
simplification passes under 'compiler/lib/Transforms' with the transformations under
'compiler/lib/Dialect/StableHloExt/Transforms'. The fact that the simplification
transforms are currently split into these two different sets is a product of the last
major restructuring of the project.

[PlanDialect] Create local instantiation of the upstream ownership-based buffer deallocation pass
This change adds a pass plan-ownership-based-buffer-deallocation which
simply calls the upstream ownership-based buffer deallocation transformation.
The only difference between this and the upstream pass is that we do not
traverse functions within nested modules.

This will allow us to move buffer optimizations and the deallocation
transformation earlier in the pipeline. Currently we require on running
deallocation after all offloading modules (e.g. TensorRT modules) are
processed.

[compiler/StableHloExt] Create general pattern to allow StableHlo ops to absorb generalizing tensor.cast
This change generalizes a rewrite pattern in the stablehlo-ext-constant-folding pass
and exposes the pattern for use in other passes. In many of the early plan segmentation
passes where types are being refined based on shape and value bounds analyses, it is
necessary to create tensor.cast operations since in general it is not legal to update
types of SSA values in-place. It is therefore important that we have a robust set of
patterns to eliminate these tensor.cast operations; otherwise they can cause problems
in the segmentation (e.g. because we don't support converting tensor.cast to
TensorRT like we do with StableHlo ops).

[StablehloToTensorRT] Support additional constant operation types in TensorRT conversion
This change allows for converting arith.constant in addition to stablehlo.constant
during the Stablehlo-to-TensorRT conversion. Since we are allowing this one
additional operation, I didn't generate a whole new pass/pattern set and just
added an instantiation of the templated constant converter that already exists.

[tensorrt] Fix logic in tensorrt.convolution verification routine
The tensorrt.convolution verification routine incorrectly was comparing
shape dimension values without checking if they are dynamic.

[plan] Add integer-range based optimizations to the materialize-shape-calculations pass
This change adds additional simplification patterns based on the IntegerRangeAnalysis.
This is a convenient way to re-use the logic provided in the analysis to refine
operations like tensor.dim into constants where possible.

[tensorrt] Fix tensorrt.reshape canonicalizer when shape operand is present
Fixes a canonicalizer that combines sequential tensorrt.reshape. Previously
it was dropping any shape operand on the consuming reshape operation. Fixes
the issue and adds a test.

NFC: [compiler] Fix trailing whitespace in Plan tablegen files

NFC: remove use of 'is_pointwise' attribute during stablehlo.composite creation
After making the stablehlo.composite shape materialization generic, we
no longer need the is_pointwise attribute that was being attached to Q/DQ
stablehlo.composite operations.

[Dialect/Plan] Make the 'stablehlo.composite' shape propagation completely generic
Makes shape propagation in 'plan-materialize-shape-calculations' completely
generic by simply cloning the composite implementation function body and
let the shape machinery analyze the body operations. The extra implementation
operations end up getting erased as long as the shape machinery can propagate through
them (which means as long as there is no DDS). Otherwise, the extra materialization
of intermediate values will be retained for shape calculation (which
may be removed down the line if the composite op is replaced with its
decomposition and CSE is performed).

[compiler] Fix use-after-free in the plan-create-closed-regions pass
A recent change in the plan-create-closed-regions pass results in calling
op->getLoc() after op has already been replaced/deleted. This change
fixes that issue, resolving ASAN errors.

Copy link
Collaborator

@christopherbate christopherbate left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please add commit list

This PR moves internal `mlir-tensorrt` project changes
to OSS.
The following commits are included,

NFC: Move "stablehlo-input-preprocessing" and "stablehlo-raise-qdq" passes under the StableHloExt folder
NFC: Don't unnecessarily restrict StableHLO preprocessing passes to 'func' ops

NFC: [StableHloExt] Consolidate scatter/gather utilities into common translation unit and under `stablehlo_ext` namespace
This change consolidates utility functions related to scatter/gather ops into the
"ScatterGatherUtils.(h|cpp)" translation unit and changes all the functions in those
files to be under the 'stablehlo_ext' namespace.

This is part one of a series of changes that will attempt to consolidate the stablehlo
simplification passes under 'compiler/lib/Transforms' with the transformations under
'compiler/lib/Dialect/StableHloExt/Transforms'. The fact that the simplification
transforms are currently split into these two different sets is a product of the last
major restructuring of the project.

[PlanDialect] Create local instantiation of the upstream ownership-based buffer deallocation pass
This change adds a pass `plan-ownership-based-buffer-deallocation` which
simply calls the upstream ownership-based buffer deallocation transformation.
The only difference between this and the upstream pass is that we do not
traverse functions within nested modules.

This will allow us to move buffer optimizations and the deallocation
transformation earlier in the pipeline. Currently we require on running
deallocation after all offloading modules (e.g. TensorRT modules) are
processed.

[compiler/StableHloExt] Create general pattern to allow StableHlo ops to absorb generalizing `tensor.cast`
This change generalizes a rewrite pattern in the `stablehlo-ext-constant-folding` pass
and exposes the pattern for use in other passes. In many of the early plan segmentation
passes where types are being refined based on shape and value bounds analyses, it is
necessary to create `tensor.cast` operations since in general it is not legal to update
types of SSA values in-place. It is therefore important that we have a robust set of
patterns to eliminate these `tensor.cast` operations; otherwise they can cause problems
in the segmentation (e.g. because we don't support converting `tensor.cast` to
TensorRT like we do with StableHlo ops).

[StablehloToTensorRT] Support additional constant operation types in TensorRT conversion
This change allows for converting `arith.constant` in addition to `stablehlo.constant`
during the Stablehlo-to-TensorRT conversion. Since we are allowing this one
additional operation, I didn't generate a whole new pass/pattern set and just
added an instantiation of the templated constant converter that already exists.

[tensorrt] Fix logic in `tensorrt.convolution` verification routine
The `tensorrt.convolution` verification routine incorrectly was comparing
shape dimension values without checking if they are dynamic.

[plan] Add integer-range based optimizations to the `materialize-shape-calculations` pass
This change adds additional simplification patterns based on the IntegerRangeAnalysis.
This is a convenient way to re-use the logic provided in the analysis to refine
operations like `tensor.dim` into constants where possible.

[tensorrt] Fix `tensorrt.reshape` canonicalizer when `shape` operand is present
Fixes a canonicalizer that combines sequential `tensorrt.reshape`. Previously
it was dropping any `shape` operand on the consuming reshape operation. Fixes
the issue and adds a test.

NFC: [compiler] Fix trailing whitespace in Plan tablegen files

NFC: remove use of 'is_pointwise' attribute during `stablehlo.composite` creation
After making the `stablehlo.composite` shape materialization generic, we
no longer need the `is_pointwise` attribute that was being attached to Q/DQ
`stablehlo.composite` operations.

[Dialect/Plan] Make the 'stablehlo.composite' shape propagation completely generic
Makes shape propagation in 'plan-materialize-shape-calculations' completely
generic by simply cloning the composite implementation function body and
let the shape machinery analyze the body operations. The extra implementation
operations end up getting erased as long as the shape machinery can propagate through
them (which means as long as there is no DDS). Otherwise, the extra materialization
of intermediate values will be retained for shape calculation (which
may be removed down the line if the composite op is replaced with its
decomposition and CSE is performed).

[compiler] Fix use-after-free in the `plan-create-closed-regions` pass
A recent change in the `plan-create-closed-regions` pass results in calling
`op->getLoc()` after `op` has already been replaced/deleted. This change
fixes that issue, resolving ASAN errors.

GitOrigin-RevId: 3efaf04575e695427f94877e590269764bcdec4a
@shelkesagar29 shelkesagar29 merged commit 274d1c5 into main Nov 5, 2024
1 check passed
@shelkesagar29 shelkesagar29 deleted the move_internal_changes branch November 5, 2024 22:06
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants