Rewrite transpose/reshape/broadcast between pointwise/reduce operators #3978
+927
−198
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
There was a problem hiding this comment.
Pull Request Overview
This PR rewrites the handling of transpose, reshape, and broadcast operations so that they are not inserted between pointwise and reduction operators. Key changes include:
- New and expanded test cases covering various combinations of pointwise, reduce, squeeze, transpose, broadcast, and reshape operators.
- Introduction of a new matcher (find_op_shape_transform_op) in simplify_reshapes.cpp to better handle shape transformations.
- Updates to the shape_transform_descriptor implementation to support additional cases and the new broadcast flag.
Reviewed Changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated no comments.
Show a summary per file
| File | Description |
|---|---|
| test/simplify_reshapes_test.cpp | Added multiple test cases to exercise new shape transforms |
| src/simplify_reshapes.cpp | Introduced new matcher to catch additional shape transformations |
| src/shape_transform_descriptor.cpp | Updated rebase logic and subdimension handling with a broadcast flag |
| src/include/migraphx/shape_transform_descriptor.hpp | Updated method signatures to support broadcast option |
| src/include/migraphx/rewrite_reshapes.hpp | Minor adjustments with commented debug prints removed |
| src/include/migraphx/matcher.hpp | Updated matcher creation functions to use new names |
Comments suppressed due to low confidence (1)
src/shape_transform_descriptor.cpp:178
- Review the handling of subdimensions when 'broadcast' is true: currently, when the input dimension equals final_dim, subdimensions are only exposed if 'broadcast' is false. Confirm that skipping exposure in the broadcast case is intended and consider adding a clarifying comment to document this behavior.
if(dim == final_dim)
{
This build is not recommended to merge 🔴 |
❌bert-mrpc-tf: ERROR - check error output2025-05-28 19:32:07.612282: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.To enable the following instructions: SSE3 SSE4.1 SSE4.2 AVX AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags. WARNING: All log messages before absl::InitializeLog() is called are written to STDERR I0000 00:00:1748478733.274427 181387 gpu_device.cc:2022] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 62973 MB memory: -> device: 0, name: AMD Instinct MI250X/MI250, pci bus id: 0000:32:00.0 WARNING: All log messages before absl::InitializeLog() is called are written to STDERR I0000 00:00:1748478734.182600 181387 mlir_graph_optimization_pass.cc:401] MLIR V1 optimization pass is not enabled 2025-05-28 19:32:24.130260: E external/local_xla/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc:250] bitcode module is required by this HLO module but was not found at ./opencl.bc 2025-05-28 19:32:24.130304: E external/local_xla/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc:250] bitcode module is required by this HLO module but was not found at ./opencl.bc 2025-05-28 19:32:24.130345: E external/local_xla/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc:250] bitcode module is required by this HLO module but was not found at ./opencl.bc 2025-05-28 19:32:24.130383: E external/local_xla/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc:250] bitcode module is required by this HLO module but was not found at ./opencl.bc 2025-05-28 19:32:24.130412: E external/local_xla/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc:250] bitcode module is required by this HLO module but was not found at ./opencl.bc 2025-05-28 19:32:24.130452: E external/local_xla/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc:250] bitcode module is required by this HLO module but was not found at ./opencl.bc 2025-05-28 19:32:24.130492: E external/local_xla/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc:250] bitcode module is required by this HLO module but was not found at ./opencl.bc 2025-05-28 19:32:24.130717: E external/local_xla/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc:250] bitcode module is required by this HLO module but was not found at ./opencl.bc error: Failure when generating HSACO error: Failure when generating HSACO error: Failure when generating HSACO error: Failure when generating HSACO error: Failure when generating HSACO error: Failure when generating HSACO error: Failure when generating HSACO error: Failure when generating HSACO 2025-05-28 19:32:24.131774: E tensorflow/compiler/mlir/tools/kernel_gen/tf_framework_c_interface.cc:228] INTERNAL: Generating device code failed. 2025-05-28 19:32:24.132911: W tensorflow/core/framework/op_kernel.cc:1829] UNKNOWN: JIT compilation failed. 2025-05-28 19:32:24.132931: I tensorflow/core/framework/local_rendezvous.cc:405] Local rendezvous is aborting with status: UNKNOWN: JIT compilation failed. [[{{node import/bert/embeddings/LayerNorm/moments/SquaredDifference}}]] 2025-05-28 19:32:24.132942: I tensorflow/core/framework/local_rendezvous.cc:405] Local rendezvous is aborting with status: UNKNOWN: JIT compilation failed. [[{{node import/bert/embeddings/LayerNorm/moments/SquaredDifference}}]] [[import/loss/output/_21]] 2025-05-28 19:32:24.132958: I tensorflow/core/framework/local_rendezvous.cc:424] Local rendezvous recv item cancelled. Key hash: 11217777527359497193 Traceback (most recent call last): File "/usr/local/lib/python3.10/dist-packages/tensorflow/python/client/session.py", line 1407, in _do_call return fn(*args) File "/usr/local/lib/python3.10/dist-packages/tensorflow/python/client/session.py", line 1390, in _run_fn return self._call_tf_sessionrun(options, feed_dict, fetch_list, File "/usr/local/lib/python3.10/dist-packages/tensorflow/python/client/session.py", line 1483, in _call_tf_sessionrun return tf_session.TF_SessionRun_wrapper(self._session, options, feed_dict, tensorflow.python.framework.errors_impl.UnknownError: 2 root error(s) found. (0) UNKNOWN: JIT compilation failed. [[{{node import/bert/embeddings/LayerNorm/moments/SquaredDifference}}]] [[import/loss/output/_21]] (1) UNKNOWN: JIT compilation failed. [[{{node import/bert/embeddings/LayerNorm/moments/SquaredDifference}}]] 0 successful operations. 0 derived errors ignored. During handling of the above exception, another exception occurred: Traceback (most recent call last): File "/src/AMDMIGraphX/tools/accuracy/accuracy_checker.py", line 340, in main() File "/src/AMDMIGraphX/tools/accuracy/accuracy_checker.py", line 324, in main y_out = sess.run(y, feed_dict=tf_dict) File "/usr/local/lib/python3.10/dist-packages/tensorflow/python/client/session.py", line 977, in run result = self._run(None, fetches, feed_dict, options_ptr, File "/usr/local/lib/python3.10/dist-packages/tensorflow/python/client/session.py", line 1220, in _run results = self._do_run(handle, final_targets, final_fetches, File "/usr/local/lib/python3.10/dist-packages/tensorflow/python/client/session.py", line 1400, in _do_run return self._do_call(_run_fn, feeds, fetches, targets, options, File "/usr/local/lib/python3.10/dist-packages/tensorflow/python/client/session.py", line 1426, in _do_call raise type(e)(node_def, op, message) # pylint: disable=no-value-for-parameter tensorflow.python.framework.errors_impl.UnknownError: Graph execution error: Detected at node 'import/bert/embeddings/LayerNorm/moments/SquaredDifference' defined at (most recent call last): Node: 'import/bert/embeddings/LayerNorm/moments/SquaredDifference' Detected at node 'import/bert/embeddings/LayerNorm/moments/SquaredDifference' defined at (most recent call last): Node: 'import/bert/embeddings/LayerNorm/moments/SquaredDifference' 2 root error(s) found. (0) UNKNOWN: JIT compilation failed. [[{{node import/bert/embeddings/LayerNorm/moments/SquaredDifference}}]] [[import/loss/output/_21]] (1) UNKNOWN: JIT compilation failed. [[{{node import/bert/embeddings/LayerNorm/moments/SquaredDifference}}]] 0 successful operations. 0 derived errors ignored. Original stack trace for 'import/bert/embeddings/LayerNorm/moments/SquaredDifference': 🔴bert_large_uncased_fp16: FAILED: MIGraphX is not within tolerance - check verbose output🔴distilgpt2_fp16: FAILED: MIGraphX is not within tolerance - check verbose output❌llama2_7b: ERROR - check error outputTraceback (most recent call last):File "/src/AMDMIGraphX/tools/accuracy/accuracy_checker.py", line 340, in main() File "/src/AMDMIGraphX/tools/accuracy/accuracy_checker.py", line 227, in main model.compile( RuntimeError: /src/AMDMIGraphX/src/include/migraphx/check_shapes.hpp:220: same_dims: less: Dimensions do not match 🔴mask-rcnn: FAILED: MIGraphX is not within tolerance - check verbose output❌llama3-8b: ERROR - check error outputTraceback (most recent call last):File "/src/AMDMIGraphX/tools/accuracy/accuracy_checker.py", line 340, in main() File "/src/AMDMIGraphX/tools/accuracy/accuracy_checker.py", line 227, in main model.compile( RuntimeError: /src/AMDMIGraphX/src/include/migraphx/op/multibroadcast.hpp:77: compute_shape: MULTIBROADCAST: input dimensions should <= output size ❌whisper-large-decoder: ERROR - check error outputTraceback (most recent call last):File "/src/AMDMIGraphX/tools/accuracy/accuracy_checker.py", line 340, in main() File "/src/AMDMIGraphX/tools/accuracy/accuracy_checker.py", line 227, in main model.compile( RuntimeError: /src/AMDMIGraphX/src/include/migraphx/check_shapes.hpp:220: same_dims: less: Dimensions do not match ❌mistral-7b: ERROR - check error outputTraceback (most recent call last):File "/src/AMDMIGraphX/tools/accuracy/accuracy_checker.py", line 340, in main() File "/src/AMDMIGraphX/tools/accuracy/accuracy_checker.py", line 227, in main model.compile( RuntimeError: /src/AMDMIGraphX/src/include/migraphx/check_shapes.hpp:220: same_dims: less: Dimensions do not match ❌FLUX.1-schnell: ERROR - check error output |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This will rewrite transpose/reshape/broadcast so they will not be between pointwise and reduction using the shape_transform_descriptor. This is similiar to the
rewrite_reshapesthat is used for pointwise/reduce fusion, however this is done before fusions in simplify_reshapes. It also supports more cases such as reduce->squeeze->pointwise.This could be expanded in the future to support other operators like argmin/argmax/concat etc.