diff --git a/.bazelrc b/.bazelrc index 2bd923c26db74c..92ddbf13548c4e 100644 --- a/.bazelrc +++ b/.bazelrc @@ -56,6 +56,7 @@ # # rbe_linux_cpu: RBE options to build with only CPU support. # rbe_linux_cuda: RBE options to build with GPU support using clang. +# rbe_linux_cuda_nvcc: RBE options to build with GPU support using nvcc. # # rbe_win_py39: Windows Python 3.9 RBE config # @@ -238,9 +239,12 @@ build:cuda_clang --@local_config_cuda//:cuda_compiler=clang # Select supported compute capabilities (supported graphics cards). # This is the same as the official TensorFlow builds. # See https://developer.nvidia.com/cuda-gpus#compute -# TODO(angerson, perfinion): What does sm_ vs compute_ mean? How can users -# select a good value for this? See go/tf-pip-cuda -build:cuda_clang --repo_env=TF_CUDA_COMPUTE_CAPABILITIES="sm_50,sm_60,sm_70,sm_75,compute_80" +# `compute_XY` enables PTX embedding in addition to SASS. PTX +# is forward compatible beyond the current compute capability major +# release while SASS is only forward compatible inside the current +# major release. Example: sm_80 kernels can run on sm_89 GPUs but +# not on sm_90 GPUs. compute_80 kernels though can also run on sm_90 GPUs. +build:cuda_clang --repo_env=TF_CUDA_COMPUTE_CAPABILITIES="sm_50,sm_60,sm_70,sm_80,compute_90" # Set up compilation CUDA version and paths and use the CUDA Clang toolchain. build:cuda_clang_official --config=cuda_clang @@ -250,7 +254,7 @@ build:cuda_clang_official --action_env=CUDA_TOOLKIT_PATH="/usr/local/cuda-12.2" build:cuda_clang_official --action_env=GCC_HOST_COMPILER_PATH="/dt9/usr/bin/gcc" build:cuda_clang_official --action_env=CLANG_CUDA_COMPILER_PATH="/usr/lib/llvm-17/bin/clang" build:cuda_clang_official --action_env=LD_LIBRARY_PATH="/usr/local/cuda/lib64:/usr/local/cuda/extras/CUPTI/lib64" -build:cuda_clang_official --crosstool_top="@sigbuild-r2.14-clang_config_cuda//crosstool:toolchain" +build:cuda_clang_official --crosstool_top="@sigbuild-r2.16-clang_config_cuda//crosstool:toolchain" # Debug config build:dbg -c dbg @@ -488,12 +492,12 @@ build:rbe_linux --host_linkopt=-lm build:rbe_linux_cpu --config=rbe_linux # Linux cpu and cuda builds share the same toolchain now. -build:rbe_linux_cpu --host_crosstool_top="@sigbuild-r2.14-clang_config_cuda//crosstool:toolchain" -build:rbe_linux_cpu --crosstool_top="@sigbuild-r2.14-clang_config_cuda//crosstool:toolchain" -build:rbe_linux_cpu --extra_toolchains="@sigbuild-r2.14-clang_config_cuda//crosstool:toolchain-linux-x86_64" -build:rbe_linux_cpu --extra_execution_platforms="@sigbuild-r2.14-clang_config_platform//:platform" -build:rbe_linux_cpu --host_platform="@sigbuild-r2.14-clang_config_platform//:platform" -build:rbe_linux_cpu --platforms="@sigbuild-r2.14-clang_config_platform//:platform" +build:rbe_linux_cpu --host_crosstool_top="@sigbuild-r2.16-clang_config_cuda//crosstool:toolchain" +build:rbe_linux_cpu --crosstool_top="@sigbuild-r2.16-clang_config_cuda//crosstool:toolchain" +build:rbe_linux_cpu --extra_toolchains="@sigbuild-r2.16-clang_config_cuda//crosstool:toolchain-linux-x86_64" +build:rbe_linux_cpu --extra_execution_platforms="@sigbuild-r2.16-clang_config_platform//:platform" +build:rbe_linux_cpu --host_platform="@sigbuild-r2.16-clang_config_platform//:platform" +build:rbe_linux_cpu --platforms="@sigbuild-r2.16-clang_config_platform//:platform" # This is needed for all Clang17 builds but must not be present in GCC builds. build:rbe_linux_cpu --copt=-Wno-error=unused-command-line-argument # This was added in clang-16 by https://reviews.llvm.org/D133574. @@ -502,7 +506,7 @@ build:rbe_linux_cpu --copt=-Wno-error=unused-command-line-argument # See https://github.com/protocolbuffers/upb/blob/9effcbcb27f0a665f9f345030188c0b291e32482/upb/upb.c#L183. build:rbe_linux_cpu --copt=-Wno-gnu-offsetof-extensions # Python config is the same across all containers because the binary is the same -build:rbe_linux_cpu --repo_env=TF_PYTHON_CONFIG_REPO="@sigbuild-r2.14-clang_config_python" +build:rbe_linux_cpu --repo_env=TF_PYTHON_CONFIG_REPO="@sigbuild-r2.16-clang_config_python" build:rbe_linux_cpu --python_path="/usr/bin/python3" # These you may need to change for your own GCP project. common:rbe_linux_cpu --remote_instance_name=projects/tensorflow-testing/instances/default_instance @@ -523,9 +527,9 @@ build:rbe_linux_cuda --config=cuda_clang_official build:rbe_linux_cuda --config=rbe_linux_cpu # For Remote build execution -- GPU configuration build:rbe_linux_cuda --repo_env=REMOTE_GPU_TESTING=1 -build:rbe_linux_cuda --repo_env=TF_CUDA_CONFIG_REPO="@sigbuild-r2.14-clang_config_cuda" -build:rbe_linux_cuda --repo_env=TF_TENSORRT_CONFIG_REPO="@sigbuild-r2.14-clang_config_tensorrt" -build:rbe_linux_cuda --repo_env=TF_NCCL_CONFIG_REPO="@sigbuild-r2.14-clang_config_nccl" +build:rbe_linux_cuda --repo_env=TF_CUDA_CONFIG_REPO="@sigbuild-r2.16-clang_config_cuda" +build:rbe_linux_cuda --repo_env=TF_TENSORRT_CONFIG_REPO="@sigbuild-r2.16-clang_config_tensorrt" +build:rbe_linux_cuda --repo_env=TF_NCCL_CONFIG_REPO="@sigbuild-r2.16-clang_config_nccl" test:rbe_linux_cuda --test_env=LD_LIBRARY_PATH="/usr/local/cuda/lib64:/usr/local/cuda/extras/CUPTI/lib64" # ROCm @@ -539,6 +543,35 @@ build:rbe_linux_rocm_base --platforms="@ubuntu20.04-gcc9_manylinux2014-rocm_conf build:rbe_linux_rocm_base --action_env=TF_ROCM_CONFIG_REPO="@ubuntu20.04-gcc9_manylinux2014-rocm_config_rocm" build:rbe_linux_rocm_py3.9 --config=rbe_linux_rocm_base --repo_env=TF_PYTHON_CONFIG_REPO="@ubuntu20.04-gcc9_manylinux2014-rocm_config_python3.9" +build:rbe_linux_cuda_nvcc --config=cuda +build:rbe_linux_cuda_nvcc --repo_env TF_NCCL_USE_STUB=1 +build:rbe_linux_cuda_nvcc --@local_xla//xla/python:enable_gpu=true +build:rbe_linux_cuda_nvcc --@local_xla//xla/python:jax_cuda_pip_rpaths=true +build:rbe_linux_cuda_nvcc --define=xla_python_enable_gpu=true +build:rbe_linux_cuda_nvcc --config=tensorrt +build:rbe_linux_cuda_nvcc --repo_env=TF_CUDA_COMPUTE_CAPABILITIES="sm_50,sm_60,sm_70,sm_75,compute_80" +build:rbe_linux_cuda_nvcc --action_env=TF_CUDA_VERSION="12" +build:rbe_linux_cuda_nvcc --action_env=TF_CUDNN_VERSION="8" +build:rbe_linux_cuda_nvcc --action_env=CUDA_TOOLKIT_PATH="/usr/local/cuda-12.2" +build:rbe_linux_cuda_nvcc --action_env=GCC_HOST_COMPILER_PATH="/dt9/usr/bin/gcc" +build:rbe_linux_cuda_nvcc --action_env=LD_LIBRARY_PATH="/usr/local/cuda/lib64:/usr/local/cuda/extras/CUPTI/lib64" +build:rbe_linux_cuda_nvcc --crosstool_top="@ubuntu20.04-clang_manylinux2014-cuda12.2-cudnn8.9_config_cuda//crosstool:toolchain" +build:rbe_linux_cuda_nvcc --config=rbe_linux +build:rbe_linux_cuda_nvcc --host_crosstool_top="@ubuntu20.04-clang_manylinux2014-cuda12.2-cudnn8.9_config_cuda//crosstool:toolchain" +build:rbe_linux_cuda_nvcc --extra_toolchains="@ubuntu20.04-clang_manylinux2014-cuda12.2-cudnn8.9_config_cuda//crosstool:toolchain-linux-x86_64" +build:rbe_linux_cuda_nvcc --extra_execution_platforms="@ubuntu20.04-clang_manylinux2014-cuda12.2-cudnn8.9_config_platform//:platform" +build:rbe_linux_cuda_nvcc --host_platform="@ubuntu20.04-clang_manylinux2014-cuda12.2-cudnn8.9_config_platform//:platform" +build:rbe_linux_cuda_nvcc --platforms="@ubuntu20.04-clang_manylinux2014-cuda12.2-cudnn8.9_config_platform//:platform" +build:rbe_linux_cuda_nvcc --repo_env=TF_PYTHON_CONFIG_REPO="@ubuntu20.04-clang_manylinux2014-cuda12.2-cudnn8.9_config_python3.9" +build:rbe_linux_cuda_nvcc --python_path="/usr/bin/python3" +# These you may need to change for your own GCP project. +common:rbe_linux_cuda_nvcc --remote_instance_name=projects/tensorflow-testing/instances/default_instance +build:rbe_linux_cuda_nvcc --repo_env=REMOTE_GPU_TESTING=1 +build:rbe_linux_cuda_nvcc --repo_env=TF_CUDA_CONFIG_REPO="@ubuntu20.04-gcc9_manylinux2014-cuda12.2-cudnn8.9_config_cuda" +build:rbe_linux_cuda_nvcc --repo_env=TF_TENSORRT_CONFIG_REPO="@ubuntu20.04-gcc9_manylinux2014-cuda12.2-cudnn8.9_config_tensorrt" +build:rbe_linux_cuda_nvcc --repo_env=TF_NCCL_CONFIG_REPO="@ubuntu20.04-gcc9_manylinux2014-cuda12.2-cudnn8.9_config_nccl" +test:rbe_linux_cuda_nvcc --test_env=LD_LIBRARY_PATH="/usr/local/cuda/lib64:/usr/local/cuda/extras/CUPTI/lib64" + # TODO(kanglan): Remove rbe_win and rbe_win_py3* after b/289091160 is fixed build:rbe_win --config=rbe_base build:rbe_win --crosstool_top="//tensorflow/tools/toolchains/win/tf_win_05022023:toolchain" @@ -593,8 +626,6 @@ try-import %workspace%/.bazelrc.user # Here are bazelrc configs for release builds # Build TensorFlow v2. test:release_base --test_size_filters=small,medium -# TODO(b/294367488) disable after 2.15 brancut -test:release_base --flaky_test_attempts=3 # Target the AVX instruction set build:release_linux_base --config=avx_linux @@ -632,7 +663,7 @@ test:release_linux_base --test_summary=short # Use the Clang toolchain to compile build:release_cpu_linux --config=release_linux_base -build:release_cpu_linux --crosstool_top="@sigbuild-r2.14-clang_config_cuda//crosstool:toolchain" +build:release_cpu_linux --crosstool_top="@sigbuild-r2.16-clang_config_cuda//crosstool:toolchain" build:release_gpu_linux --config=release_cpu_linux # Set up compilation CUDA version and paths and use the CUDA Clang toolchain. diff --git a/.github/workflows/osv-scanner-scheduled.yml b/.github/workflows/osv-scanner-scheduled.yml new file mode 100644 index 00000000000000..bb39d60168e08d --- /dev/null +++ b/.github/workflows/osv-scanner-scheduled.yml @@ -0,0 +1,39 @@ +# Copyright 2023 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +name: OSV-Scanner Scheduled Scan + +on: + schedule: + - cron: 0 4 * * 1 + +permissions: + # Require writing security events to upload SARIF file to security tab + security-events: write + # Only need to read contents + contents: read + +jobs: + scan-scheduled: + uses: "google/osv-scanner/.github/workflows/osv-scanner-reusable.yml@main" + with: + scan-args: |- + --lockfile=requirements.txt:./requirements_lock_3_9.txt + --lockfile=requirements.txt:./requirements_lock_3_10.txt + --lockfile=requirements.txt:./requirements_lock_3_11.txt + --lockfile=requirements.txt:./requirements_lock_3_12.txt + --lockfile=requirements.txt:./ci/official/containers/linux_arm64/devel.requirements.txt + --lockfile=requirements.txt:./ci/official/containers/linux_arm64/jax.requirements.txt + --lockfile=requirements.txt:./ci/official/containers/linux_arm64/devel.usertools/test.requirements.txt \ No newline at end of file diff --git a/.github/workflows/sigbuild-docker-branch.yml b/.github/workflows/sigbuild-docker-branch.yml index 108fe471efa2db..9f842f9fb27c11 100644 --- a/.github/workflows/sigbuild-docker-branch.yml +++ b/.github/workflows/sigbuild-docker-branch.yml @@ -34,7 +34,7 @@ jobs: runs-on: ubuntu-latest strategy: matrix: - python-version: [python3.9, python3.10, python3.11] + python-version: [python3.9, python3.10, python3.11, python3.12] steps: - name: Delete unnecessary tools folder run: rm -rf /opt/hostedtoolcache diff --git a/.github/workflows/sigbuild-docker-presubmit.yml b/.github/workflows/sigbuild-docker-presubmit.yml index c61e65e7d834c0..03ae6f1dadf63f 100644 --- a/.github/workflows/sigbuild-docker-presubmit.yml +++ b/.github/workflows/sigbuild-docker-presubmit.yml @@ -32,7 +32,7 @@ jobs: runs-on: ubuntu-latest strategy: matrix: - python-version: [python3.9, python3.10, python3.11] + python-version: [python3.9, python3.10, python3.11, python3.12] permissions: contents: read pull-requests: write @@ -87,6 +87,7 @@ jobs: message: | I pushed these containers: + - `gcr.io/tensorflow-sigs/build:${{ github.event.number }}-python3.12` - `gcr.io/tensorflow-sigs/build:${{ github.event.number }}-python3.11` - `gcr.io/tensorflow-sigs/build:${{ github.event.number }}-python3.10` - `gcr.io/tensorflow-sigs/build:${{ github.event.number }}-python3.9` diff --git a/.github/workflows/sigbuild-docker.yml b/.github/workflows/sigbuild-docker.yml index ce9b99c494fc5e..5549f2995ac80f 100644 --- a/.github/workflows/sigbuild-docker.yml +++ b/.github/workflows/sigbuild-docker.yml @@ -37,7 +37,7 @@ jobs: runs-on: ubuntu-latest strategy: matrix: - python-version: [python3.9, python3.10, python3.11] + python-version: [python3.9, python3.10, python3.11, python3.12] steps: - name: Delete unnecessary tools folder run: rm -rf /opt/hostedtoolcache diff --git a/.github/workflows/update-rbe.yml b/.github/workflows/update-rbe.yml index ca22041782fa4f..1b421effec8198 100644 --- a/.github/workflows/update-rbe.yml +++ b/.github/workflows/update-rbe.yml @@ -105,6 +105,18 @@ jobs: map sigbuild-r2.14-clang-python3.9 2.14-python3.9 map sigbuild-r2.14-clang-python3.10 2.14-python3.10 map sigbuild-r2.14-clang-python3.11 2.14-python3.11 + # TF 2.16 + map sigbuild-r2.16 2.16-python3.9 + map sigbuild-r2.16-python3.9 2.16-python3.9 + map sigbuild-r2.16-python3.10 2.16-python3.10 + map sigbuild-r2.16-python3.11 2.16-python3.11 + map sigbuild-r2.16-python3.12 2.16-python3.12 + # TF 2.16 + Clang (containers are the same, but env vars in configs.bzl are different) + map sigbuild-r2.16-clang 2.16-python3.9 + map sigbuild-r2.16-clang-python3.9 2.16-python3.9 + map sigbuild-r2.16-clang-python3.10 2.16-python3.10 + map sigbuild-r2.16-clang-python3.11 2.16-python3.11 + map sigbuild-r2.16-clang-python3.12 2.16-python3.12 - name: Create Pull Request with changes uses: peter-evans/create-pull-request@2b011faafdcbc9ceb11414d64d0573f37c774b04 # v4.2.3 with: diff --git a/RELEASE.md b/RELEASE.md index 632c3652af1eff..94c795442d8992 100644 --- a/RELEASE.md +++ b/RELEASE.md @@ -28,6 +28,8 @@ * `tf.lite` * Added support for `stablehlo.gather`. + * Added support for `stablehlo.add`. + * Added support for `stablehlo.multiply`. ## Keras diff --git a/ci/official/containers/linux_arm64/devel.packages.txt b/ci/official/containers/linux_arm64/devel.packages.txt index a8a9cb442c8b0b..efbae80eefacee 100644 --- a/ci/official/containers/linux_arm64/devel.packages.txt +++ b/ci/official/containers/linux_arm64/devel.packages.txt @@ -3,6 +3,8 @@ autoconf automake build-essential ca-certificates +# TODO(b/308399490) Remove CMake once dm-tree (Keras dependency) has 3.12 wheels +cmake llvm-17 clang-17 clang-format-12 diff --git a/ci/official/envs/ci_default b/ci/official/envs/ci_default index 183b2048ce1a5f..eb7938c8b3449d 100644 --- a/ci/official/envs/ci_default +++ b/ci/official/envs/ci_default @@ -16,6 +16,7 @@ TFCI_NIGHTLY_UPDATE_VERSION_ENABLE= TFCI_NVIDIA_SMI_ENABLE= TFCI_OUTPUT_DIR=build_output TFCI_LIBTPU_DOWNLOAD_ENABLE=0 +TFCI_LIBTPU_DOWNLOAD_NIGHTLY_ENABLE=0 TFCI_LIBTPU_DOWNLOAD_URL= TFCI_UPLOAD_LIB_ENABLE= TFCI_UPLOAD_LIB_LATEST_ENABLE= diff --git a/ci/official/envs/ci_nightly_uploads b/ci/official/envs/ci_nightly_uploads index e35a712d034966..ca6671f5ea3c59 100644 --- a/ci/official/envs/ci_nightly_uploads +++ b/ci/official/envs/ci_nightly_uploads @@ -1,6 +1,8 @@ -TFCI_UPLOAD_LIB_ENABLE=0 -TFCI_UPLOAD_LIB_LATEST_ENABLE= -TFCI_UPLOAD_WHL_GCS_ENABLE= +TFCI_UPLOAD_LIB_ENABLE=1 +TFCI_UPLOAD_LIB_GCS_URI="gs://libtensorflow-nightly/$(date -I)" +TFCI_UPLOAD_LIB_LATEST_ENABLE=1 +TFCI_UPLOAD_LIB_LATEST_GCS_URI="gs://libtensorflow-nightly/latest" +TFCI_UPLOAD_WHL_GCS_ENABLE=0 TFCI_UPLOAD_WHL_GCS_URI= -TFCI_UPLOAD_WHL_PYPI_ARGS= -TFCI_UPLOAD_WHL_PYPI_ENABLE= +TFCI_UPLOAD_WHL_PYPI_ARGS=(--config-file="$KOKORO_KEYSTORE_DIR/73361_tensorflow_pypirc_using_global_api_token" --repository pypi-warehouse) +TFCI_UPLOAD_WHL_PYPI_ENABLE=1 diff --git a/ci/official/envs/disable_all_uploads b/ci/official/envs/disable_all_uploads index b09169d677fa8f..6559f80e7f7c12 100644 --- a/ci/official/envs/disable_all_uploads +++ b/ci/official/envs/disable_all_uploads @@ -1,7 +1,9 @@ -TFCI_UPLOAD_LIB_ENABLE=0 +TFCI_DOCKER_REBUILD_UPLOAD_ENABLE=0 +TFCI_UPLOAD_LIB_ENABLE= +TFCI_UPLOAD_LIB_GCS_URI= TFCI_UPLOAD_LIB_LATEST_ENABLE= -TFCI_DOCKER_REBUILD_UPLOAD_ENABLE= +TFCI_UPLOAD_LIB_LATEST_GCS_URI= TFCI_UPLOAD_WHL_GCS_ENABLE= TFCI_UPLOAD_WHL_GCS_URI= -TFCI_UPLOAD_WHL_PYPI_ENABLE= TFCI_UPLOAD_WHL_PYPI_ARGS= +TFCI_UPLOAD_WHL_PYPI_ENABLE= diff --git a/ci/official/envs/nightly_linux_arm64_cpu_py310 b/ci/official/envs/nightly_linux_arm64_cpu_py310 index 7f1b040dff89b5..5b7900c43423b2 100644 --- a/ci/official/envs/nightly_linux_arm64_cpu_py310 +++ b/ci/official/envs/nightly_linux_arm64_cpu_py310 @@ -1,5 +1,6 @@ source ci/official/envs/ci_default -source ci/official/envs/ci_nightly_uploads +# Disable arm64 uploads while being worked on +source ci/official/envs/disable_all_uploads TFCI_PYTHON_VERSION=3.10 TFCI_BAZEL_COMMON_ARGS=(--config release_arm64_linux --config tf_public_cache_push --config resultstore --repo_env=TF_PYTHON_VERSION=$TFCI_PYTHON_VERSION) TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=linux_arm64 diff --git a/ci/official/envs/nightly_linux_arm64_cpu_py311 b/ci/official/envs/nightly_linux_arm64_cpu_py311 index 7779d84f0493ca..6edb93ba0bdf73 100644 --- a/ci/official/envs/nightly_linux_arm64_cpu_py311 +++ b/ci/official/envs/nightly_linux_arm64_cpu_py311 @@ -1,5 +1,6 @@ source ci/official/envs/ci_default -source ci/official/envs/ci_nightly_uploads +# Disable arm64 uploads while being worked on +source ci/official/envs/disable_all_uploads TFCI_PYTHON_VERSION=3.11 TFCI_BAZEL_COMMON_ARGS=(--config release_arm64_linux --config tf_public_cache_push --config resultstore --repo_env=TF_PYTHON_VERSION=$TFCI_PYTHON_VERSION) TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=linux_arm64 diff --git a/ci/official/envs/nightly_linux_arm64_cpu_py312 b/ci/official/envs/nightly_linux_arm64_cpu_py312 new file mode 100644 index 00000000000000..dfe96fafb5568e --- /dev/null +++ b/ci/official/envs/nightly_linux_arm64_cpu_py312 @@ -0,0 +1,10 @@ +source ci/official/envs/ci_default +# Disable arm64 uploads while being worked on +source ci/official/envs/disable_all_uploads +TFCI_PYTHON_VERSION=3.12 +TFCI_BAZEL_COMMON_ARGS=(--config release_arm64_linux --config tf_public_cache_push --config resultstore --repo_env=TF_PYTHON_VERSION=$TFCI_PYTHON_VERSION) +TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=linux_arm64 +TFCI_BUILD_PIP_PACKAGE_ARGS=(--cpu --nightly_flag) +TFCI_DOCKER_IMAGE=gcr.io/tensorflow-sigs/build-arm64:tf-latest-multi-python +TFCI_DOCKER_REBUILD_ARGS=(--target=tf ci/official/containers/linux_arm64) +TFCI_NIGHTLY_UPDATE_VERSION_ENABLE=1 diff --git a/ci/official/envs/nightly_linux_arm64_cpu_py39 b/ci/official/envs/nightly_linux_arm64_cpu_py39 index 274f777423e132..e3b516111fdc85 100644 --- a/ci/official/envs/nightly_linux_arm64_cpu_py39 +++ b/ci/official/envs/nightly_linux_arm64_cpu_py39 @@ -1,5 +1,6 @@ source ci/official/envs/ci_default -source ci/official/envs/ci_nightly_uploads +# Disable arm64 uploads while being worked on +source ci/official/envs/disable_all_uploads TFCI_PYTHON_VERSION=3.9 TFCI_BAZEL_COMMON_ARGS=(--config release_arm64_linux --config tf_public_cache_push --config resultstore --repo_env=TF_PYTHON_VERSION=$TFCI_PYTHON_VERSION) TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=linux_arm64 diff --git a/ci/official/envs/nightly_linux_x86_cpu_py312 b/ci/official/envs/nightly_linux_x86_cpu_py312 new file mode 100644 index 00000000000000..586fd92e5d703c --- /dev/null +++ b/ci/official/envs/nightly_linux_x86_cpu_py312 @@ -0,0 +1,10 @@ +source ci/official/envs/ci_default +# Disable 3.12 uploads while being worked on +source ci/official/envs/disable_all_uploads +TFCI_PYTHON_VERSION=3.12 +TFCI_BAZEL_COMMON_ARGS=(--config release_cpu_linux --config tf_public_cache_push --config resultstore --repo_env=TF_PYTHON_VERSION=$TFCI_PYTHON_VERSION) +TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=linux_cpu +TFCI_BUILD_PIP_PACKAGE_ARGS=(--cpu --nightly_flag) +TFCI_DOCKER_IMAGE=tensorflow/build:latest-python${TFCI_PYTHON_VERSION} +TFCI_DOCKER_REBUILD_ARGS=(--build-arg PYTHON_VERSION=$TFCI_PYTHON_VERSION --target=devel tools/tf_sig_build_dockerfiles) +TFCI_NIGHTLY_UPDATE_VERSION_ENABLE=1 diff --git a/ci/official/envs/nightly_linux_x86_cuda_py312 b/ci/official/envs/nightly_linux_x86_cuda_py312 new file mode 100644 index 00000000000000..4767f6dbdd6483 --- /dev/null +++ b/ci/official/envs/nightly_linux_x86_cuda_py312 @@ -0,0 +1,11 @@ +source ci/official/envs/ci_default +# Disable 3.12 uploads while being worked on +source ci/official/envs/disable_all_uploads +TFCI_PYTHON_VERSION=3.12 +TFCI_BAZEL_COMMON_ARGS=(--config release_gpu_linux --config tf_public_cache_push --config resultstore --repo_env=TF_PYTHON_VERSION=$TFCI_PYTHON_VERSION) +TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=linux_cuda +TFCI_BUILD_PIP_PACKAGE_ARGS=(--nightly_flag) +TFCI_DOCKER_ARGS=(--gpus all) +TFCI_DOCKER_IMAGE=tensorflow/build:latest-python${TFCI_PYTHON_VERSION} +TFCI_DOCKER_REBUILD_ARGS=(--build-arg PYTHON_VERSION=$TFCI_PYTHON_VERSION --target=devel tools/tf_sig_build_dockerfiles) +TFCI_NIGHTLY_UPDATE_VERSION_ENABLE=1 \ No newline at end of file diff --git a/ci/official/envs/nightly_linux_x86_tpu_py310 b/ci/official/envs/nightly_linux_x86_tpu_py310 index da77f5ab3668ae..4e8014120f3762 100644 --- a/ci/official/envs/nightly_linux_x86_tpu_py310 +++ b/ci/official/envs/nightly_linux_x86_tpu_py310 @@ -1,11 +1,11 @@ source ci/official/envs/ci_default -source ci/official/envs/ci_nightly_uploads +# Disable tpu uploads while being worked on +source ci/official/envs/disable_all_uploads TFCI_PYTHON_VERSION=3.10 TFCI_BAZEL_COMMON_ARGS=(--config tf_public_cache_push --config resultstore --repo_env=TF_PYTHON_VERSION=$TFCI_PYTHON_VERSION --config=tpu) TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=linux_tpu TFCI_BUILD_PIP_PACKAGE_ARGS=(--tpu --nightly_flag) TFCI_DOCKER_IMAGE=tensorflow/build:latest-python${TFCI_PYTHON_VERSION} TFCI_DOCKER_REBUILD_ARGS=(--build-arg PYTHON_VERSION=$TFCI_PYTHON_VERSION --target=devel tools/tf_sig_build_dockerfiles) -TFCI_LIBTPU_DOWNLOAD_ENABLE=1 -TFCI_LIBTPU_DOWNLOAD_URL=https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/libtpu/1.8.0/libtpu.so +TFCI_LIBTPU_DOWNLOAD_NIGHTLY_ENABLE=1 TFCI_NIGHTLY_UPDATE_VERSION_ENABLE=1 diff --git a/ci/official/envs/nightly_linux_x86_tpu_py311 b/ci/official/envs/nightly_linux_x86_tpu_py311 index 8f95c5df576b18..e4ae8cccf4fd46 100644 --- a/ci/official/envs/nightly_linux_x86_tpu_py311 +++ b/ci/official/envs/nightly_linux_x86_tpu_py311 @@ -1,11 +1,11 @@ source ci/official/envs/ci_default -source ci/official/envs/ci_nightly_uploads +# Disable tpu uploads while being worked on +source ci/official/envs/disable_all_uploads TFCI_PYTHON_VERSION=3.11 TFCI_BAZEL_COMMON_ARGS=(--config tf_public_cache_push --config resultstore --repo_env=TF_PYTHON_VERSION=$TFCI_PYTHON_VERSION --config=tpu) TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=linux_tpu TFCI_BUILD_PIP_PACKAGE_ARGS=(--tpu --nightly_flag) TFCI_DOCKER_IMAGE=tensorflow/build:latest-python${TFCI_PYTHON_VERSION} TFCI_DOCKER_REBUILD_ARGS=(--build-arg PYTHON_VERSION=$TFCI_PYTHON_VERSION --target=devel tools/tf_sig_build_dockerfiles) -TFCI_LIBTPU_DOWNLOAD_ENABLE=1 -TFCI_LIBTPU_DOWNLOAD_URL=https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/libtpu/1.8.0/libtpu.so +TFCI_LIBTPU_DOWNLOAD_NIGHTLY_ENABLE=1 TFCI_NIGHTLY_UPDATE_VERSION_ENABLE=1 diff --git a/ci/official/envs/nightly_linux_x86_tpu_py312 b/ci/official/envs/nightly_linux_x86_tpu_py312 new file mode 100644 index 00000000000000..54d96b16548a4a --- /dev/null +++ b/ci/official/envs/nightly_linux_x86_tpu_py312 @@ -0,0 +1,11 @@ +source ci/official/envs/ci_default +# Disable tpu uploads while being worked on +source ci/official/envs/disable_all_uploads +TFCI_PYTHON_VERSION=3.12 +TFCI_BAZEL_COMMON_ARGS=(--config tf_public_cache_push --config resultstore --repo_env=TF_PYTHON_VERSION=$TFCI_PYTHON_VERSION --config=tpu) +TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=linux_tpu +TFCI_BUILD_PIP_PACKAGE_ARGS=(--tpu --nightly_flag) +TFCI_DOCKER_IMAGE=tensorflow/build:latest-python${TFCI_PYTHON_VERSION} +TFCI_DOCKER_REBUILD_ARGS=(--build-arg PYTHON_VERSION=$TFCI_PYTHON_VERSION --target=devel tools/tf_sig_build_dockerfiles) +TFCI_LIBTPU_DOWNLOAD_NIGHTLY_ENABLE=1 +TFCI_NIGHTLY_UPDATE_VERSION_ENABLE=1 diff --git a/ci/official/envs/nightly_linux_x86_tpu_py39 b/ci/official/envs/nightly_linux_x86_tpu_py39 index b75e450cda1a0a..4adaa8b216fbba 100644 --- a/ci/official/envs/nightly_linux_x86_tpu_py39 +++ b/ci/official/envs/nightly_linux_x86_tpu_py39 @@ -1,11 +1,11 @@ source ci/official/envs/ci_default -source ci/official/envs/ci_nightly_uploads +# Disable tpu uploads while being worked on +source ci/official/envs/disable_all_uploads TFCI_PYTHON_VERSION=3.9 TFCI_BAZEL_COMMON_ARGS=(--config tf_public_cache_push --config resultstore --repo_env=TF_PYTHON_VERSION=$TFCI_PYTHON_VERSION --config=tpu) TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=linux_tpu TFCI_BUILD_PIP_PACKAGE_ARGS=(--tpu --nightly_flag) TFCI_DOCKER_IMAGE=tensorflow/build:latest-python${TFCI_PYTHON_VERSION} TFCI_DOCKER_REBUILD_ARGS=(--build-arg PYTHON_VERSION=$TFCI_PYTHON_VERSION --target=devel tools/tf_sig_build_dockerfiles) -TFCI_LIBTPU_DOWNLOAD_ENABLE=1 -TFCI_LIBTPU_DOWNLOAD_URL=https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/libtpu/1.8.0/libtpu.so +TFCI_LIBTPU_DOWNLOAD_NIGHTLY_ENABLE=1 TFCI_NIGHTLY_UPDATE_VERSION_ENABLE=1 diff --git a/ci/official/envs/nightly_macos_arm64_py312 b/ci/official/envs/nightly_macos_arm64_py312 new file mode 100644 index 00000000000000..21432f076f6283 --- /dev/null +++ b/ci/official/envs/nightly_macos_arm64_py312 @@ -0,0 +1,9 @@ +source ci/official/envs/ci_default +source ci/official/envs/disable_all_uploads +TFCI_PYTHON_VERSION=3.12 +TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=macos_arm64 +TFCI_BAZEL_COMMON_ARGS=(--config release_macos_arm64 --config tf_public_cache_push --config resultstore --repo_env=TF_PYTHON_VERSION=$TFCI_PYTHON_VERSION) +TFCI_BUILD_PIP_PACKAGE_ARGS=(--cpu --nightly_flag) +TFCI_DOCKER_ENABLE=0 +TFCI_NIGHTLY_UPDATE_VERSION_ENABLE=1 +TFCI_UPLOAD_WHL_GCS_ENABLE=1 diff --git a/ci/official/requirements_updater/README.md b/ci/official/requirements_updater/README.md index 292cb072a367aa..ad2e350dec8dc6 100644 --- a/ci/official/requirements_updater/README.md +++ b/ci/official/requirements_updater/README.md @@ -1,75 +1,128 @@ -### Hermetic Python +# Hermetic Python -Hermetic Python allows us not to rely on system-installed python and -system-installed python packages, instead we register our own python toolchain. +Hermetic Python allows not to rely on system-installed Python, and +system-installed Python packages. \ +Instead, an independent Python toolchain is registered, ensuring the right +dependencies are always used. \ See https://github.com/bazelbuild/rules_python/ for more details. -#### Hermetic Python toolchain details +### Specifying the Python version -By default, Python 3.9 is used. +Note: Only a number of minor Python versions are supported at any given time. -To set your own version for hermetic Python toolchain, use `TF_PYTHON_VERSION` -environment variable, e.g. +By default, the lowest supported version is used. + +To set a different version, use the `TF_PYTHON_VERSION` environment variable, +e.g. ``` -export TF_PYTHON_VERSION=3.10 +export TF_PYTHON_VERSION=3.11 ``` -To set a version from argument line, add to your command +To specify the version via a Bazel command argument, use the following: ``` ---repo_env=TF_PYTHON_VERSION=3.10 +--repo_env=TF_PYTHON_VERSION=3.11 ``` -### Requirements updater - -Requirements updater is a standalone tool intended to simplify process of -updating requirements for multiple versions of Python. +## Requirements updater -#### How to update/add requirements +Requirements updater is a standalone tool, intended to simplify process of +updating requirements for multiple minor versions of Python. -By default, the name of the input requirements file is `requirements.in`, -but it can be set using the `REQUIREMENTS_FILE_NAME` variable, for example: -``` -export REQUIREMENTS_FILE_NAME=`my_requirements.in` -``` +It takes in a file with a set of dependencies, and produces a more detailed +requirements file for each version, with hashes specified for each +dependency required, as well as their sub-dependencies. -To set a version from the argument line, add to your command -``` ---repo_env=REQUIREMENTS_FILE_NAME=`my_requirements.in` -``` +### How to update/add requirements -#### How to run the updater +By default, the name of the base requirements file is `requirements.in`, but it +can be set using the `REQUIREMENTS_FILE_NAME` variable. \ +For example: ``` -bash updater.sh +export REQUIREMENTS_FILE_NAME=my_requirements.in ``` -### How to add a new Python version - -1) In the `WORKSPACE` file add a new version to `python_versions` argument of -the `python_register_multi_toolchains` function. - -2) In `BUILD.bazel` file add a load statement for the new version, e.g. +To specify the file via a Bazel command argument, use the following: ``` -load("@python//3.11:defs.bzl", - compile_pip_requirements_3_11 = "compile_pip_requirements") +--repo_env=REQUIREMENTS_FILE_NAME=my_requirements.in ``` -Add a new entry for the loaded `compile_pip_requirements`, e.g. +### How to run the updater ``` -compile_pip_requirements_3_11( - name = "requirements_3_11", - extra_args = ["--allow-unsafe"], - requirements_in = "requirements.in", - requirements_txt = "requirements_lock_3_11.txt", -) +bash updater.sh ``` -3) Add the version to `SUPPORTED_VERSIONS` in `updater.sh`, after that run the - requirements updater tool. - -4) As a result, a new `requirements_lock_3_11.txt` file should appear under the -root of tensorflow directory. +## How to add a new Python version + +Note: Updating the +[rules-python](https://github.com/bazelbuild/rules_python/releases) version may +be required before going through the steps below. This is due to the new Python +versions becoming available through `rules-python`. \ +See +[here](https://github.com/tensorflow/tensorflow/commit/f91457f258fdd78f693044a57efa63a38335d1de), +and +[here](https://github.com/tensorflow/tensorflow/commit/052445e04ce20fd747657e0198a1bcec2b6dff5b), +for an example. + +See +[this commit](https://github.com/tensorflow/tensorflow/commit/5f7f05a80aac9b01325a78ec3fcff0dbedb1cc23) +as a rough example of the steps below. + +All the files referenced below are located in the same directory as this README, +unless indicated otherwise. + +1) Add the new version to the `VERSIONS` variable inside + `tensorflow/tools/toolchains/python/python_repo.bzl`. \ + While this isn't necessary for running the updater, it is required for + actually using the new version with Tensorflow. + +2) In the `WORKSPACE` file, add the new version to the `python_versions` + parameter of the `python_register_multi_toolchains` function. + +3) In the `BUILD.bazel` file, add a load statement for the new version, e.g. + + ``` + load("@python//3.11:defs.bzl", + compile_pip_requirements_3_11 = "compile_pip_requirements") + ``` + + Add a new entry for the loaded `compile_pip_requirements`, e.g. + + ``` + compile_pip_requirements_3_11( + name = "requirements_3_11", + extra_args = ["--allow-unsafe"], + requirements_in = "requirements.in", + requirements_txt = "requirements_lock_3_11.txt", + ) + ``` + + ``` + compile_pip_requirements_3_11( + name = "requirements_3_11_release", + extra_args = [ + "--allow-unsafe", + "-P keras-nightly", + "-P tb-nightly", + "-P tf-estimator-nightly", + ], + requirements_in = "requirements.in", + requirements_txt = "requirements_lock_3_11.txt", + ) + ``` + +4) Add the version to `SUPPORTED_VERSIONS` in `updater.sh`, and + `release_updater.sh` + +5) Run the `updater.sh` shell script. \ + If the base requirements file hasn't yet been updated to account for the new + Python version, which will require different versions for at least some + dependencies, it will need to be updated now, for the script to run + successfully. + +6) A new `requirements_lock_3_11.txt` file should appear under the root of the + `tensorflow` directory. diff --git a/ci/official/utilities/rename_and_verify_wheels.sh b/ci/official/utilities/rename_and_verify_wheels.sh index cd02b829194a98..4388329ae6edd7 100755 --- a/ci/official/utilities/rename_and_verify_wheels.sh +++ b/ci/official/utilities/rename_and_verify_wheels.sh @@ -20,14 +20,14 @@ set -euxo pipefail DIR=$1 -find $DIR -iname "*.whl" | while read wheel; do +find "$DIR" -iname "*.whl" | while read wheel; do echo "Checking and renaming $wheel..." wheel=$(realpath "$wheel") # Repair wheel based upon name/architecture, fallback to x86 if [[ $wheel == *"aarch64.whl" ]]; then - time python3 -m auditwheel repair --plat manylinux2014_aarch64 "$wheel" --wheel-dir build 2>&1 | tee check.txt + time python3 -m auditwheel repair --plat manylinux2014_aarch64 "$wheel" --wheel-dir "$DIR" 2>&1 | tee check.txt else - time python3 -m auditwheel repair --plat manylinux2014_x86_64 "$wheel" --wheel-dir build 2>&1 | tee check.txt + time python3 -m auditwheel repair --plat manylinux2014_x86_64 "$wheel" --wheel-dir "$DIR" 2>&1 | tee check.txt fi # We don't need the original wheel if it was renamed @@ -38,5 +38,5 @@ find $DIR -iname "*.whl" | while read wheel; do fi rm check.txt - TF_WHEEL="$wheel" bats ./ci/official/utilities/wheel_verification.bats --timing + TF_WHEEL="$wheel" BUILD_DIR="$DIR" bats ./ci/official/utilities/wheel_verification.bats --timing done diff --git a/ci/official/utilities/setup.sh b/ci/official/utilities/setup.sh index 2d1e3300da728c..fc352671274d99 100755 --- a/ci/official/utilities/setup.sh +++ b/ci/official/utilities/setup.sh @@ -57,6 +57,16 @@ else echo 'If you have not, you will see a lot of undefined variable errors.' fi +# Force-disable uploads if the job initiator is not Kokoro +# This is temporary: it's currently standard practice for employees to +# run nightly jobs for testing purposes. We're aiming to move away from +# this with more convenient methods, but as long as it's possible to do, +# we want to make sure those extra jobs don't upload anything. +# TODO(angerson) Remove this once it's no longer relevant +if [[ "${KOKORO_BUILD_INITIATOR:-}" != "kokoro" ]]; then + source ./ci/official/envs/disable_all_uploads +fi + # Create and expand to the full path of TFCI_OUTPUT_DIR export TFCI_OUTPUT_DIR=$(realpath "$TFCI_OUTPUT_DIR") mkdir -p "$TFCI_OUTPUT_DIR" diff --git a/ci/official/utilities/wheel_verification.bats b/ci/official/utilities/wheel_verification.bats index e63afc43ac2b1f..99d0f32e35162e 100644 --- a/ci/official/utilities/wheel_verification.bats +++ b/ci/official/utilities/wheel_verification.bats @@ -12,13 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -# Suite of verification tests for the SINGLE TensorFlow wheel in the "build" -# directory, or whatever path is set as $TF_WHEEL. +# Suite of verification tests for the SINGLE TensorFlow wheel in the +# $BUILD_DIR directory, or whatever path is set as $TF_WHEEL. setup_file() { - cd build + cd "$BUILD_DIR" if [[ -z "$TF_WHEEL" ]]; then - export TF_WHEEL=$(find build -iname "*.whl") + export TF_WHEEL=$(find "$BUILD_DIR" -iname "*.whl") fi # Setup the env for the python import testing diff --git a/ci/official/wheel.sh b/ci/official/wheel.sh index e3e569f4c112c2..20c6f2637d7e12 100755 --- a/ci/official/wheel.sh +++ b/ci/official/wheel.sh @@ -29,6 +29,14 @@ fi if [[ "$TFCI_LIBTPU_DOWNLOAD_ENABLE" == 1 ]]; then wget -P ./tensorflow/lib/ "$TFCI_LIBTPU_DOWNLOAD_URL" fi +if [[ "$TFCI_LIBTPU_DOWNLOAD_NIGHTLY_ENABLE" == 1 ]]; then + # For nightly jobs, libtpu.so comes from the latest nightly libtpu build. + # Note: expects a working wheel for today + DATE=$(TZ='America/Los_Angeles' date '+%Y%m%d') + tfrun wget "https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/wheels/libtpu-nightly/libtpu_nightly-0.1.dev${DATE}-py3-none-any.whl" -O libtpu.whl + # -j to discard intermediate directories; -o to overwrite if exists; -d to set output dir + tfrun unzip libtpu.whl libtpu/libtpu.so -j -o -d ./tensorflow/lib +fi tfrun bazel "${TFCI_BAZEL_BAZELRC_ARGS[@]}" build "${TFCI_BAZEL_COMMON_ARGS[@]}" //tensorflow/tools/pip_package:build_pip_package tfrun ./bazel-bin/tensorflow/tools/pip_package/build_pip_package "$TFCI_OUTPUT_DIR" "${TFCI_BUILD_PIP_PACKAGE_ARGS[@]}" diff --git a/configure.py b/configure.py index cbb74beb00c22f..c1cb20162012f6 100644 --- a/configure.py +++ b/configure.py @@ -878,11 +878,12 @@ def retrieve_clang_version(clang_executable): # Disable clang extension that rejects type definitions within offsetof. # This was added in clang-16 by https://reviews.llvm.org/D133574. +# Still required for clang-17. # Can be removed once upb is updated, since a type definition is used within # offset of in the current version of ubp. See # https://github.com/protocolbuffers/upb/blob/9effcbcb27f0a665f9f345030188c0b291e32482/upb/upb.c#L183. -def disable_clang16_offsetof_extension(clang_version): - if int(clang_version.split('.')[0]) == 16: +def disable_clang_offsetof_extension(clang_version): + if int(clang_version.split('.')[0]) in (16, 17): write_to_bazelrc('build --copt=-Wno-gnu-offsetof-extensions') @@ -1399,7 +1400,7 @@ def main(): # Set up which clang we should use as the cuda / host compiler. clang_cuda_compiler_path = set_clang_cuda_compiler_path(environ_cp) clang_version = retrieve_clang_version(clang_cuda_compiler_path) - disable_clang16_offsetof_extension(clang_version) + disable_clang_offsetof_extension(clang_version) else: # Set up which gcc nvcc should use as the host compiler # No need to set this on Windows @@ -1413,7 +1414,7 @@ def main(): if environ_cp.get('TF_NEED_CLANG') == '1': clang_compiler_path = set_clang_compiler_path(environ_cp) clang_version = retrieve_clang_version(clang_compiler_path) - disable_clang16_offsetof_extension(clang_version) + disable_clang_offsetof_extension(clang_version) # ROCm / CUDA are mutually exclusive. # At most 1 GPU platform can be configured. diff --git a/tensorflow/BUILD b/tensorflow/BUILD index 18ed141388c020..ef01b603800a71 100644 --- a/tensorflow/BUILD +++ b/tensorflow/BUILD @@ -1056,6 +1056,7 @@ package_group( "//third_party/py/envlogger/...", "//third_party/py/gldm/...", "//third_party/py/guesslang/...", + "//third_party/py/keras/...", "//third_party/py/tf_keras/...", "//third_party/yggdrasil_decision_forests/...", "//waymo/ml/cn/...", diff --git a/tensorflow/c/c_api.cc b/tensorflow/c/c_api.cc index bf598c4d57c148..5490149bc905b1 100644 --- a/tensorflow/c/c_api.cc +++ b/tensorflow/c/c_api.cc @@ -1817,6 +1817,21 @@ TF_ImportGraphDefResults* TF_GraphImportGraphDefWithResults( return results; } +TF_ImportGraphDefResults* TF_GraphImportGraphDefWithResultsNoSerialization( + TF_Graph* graph, const TF_Buffer* graph_def, + const TF_ImportGraphDefOptions* options, TF_Status* status) { + const GraphDef* graph_def_ptr = + reinterpret_cast(graph_def->data); + auto results = new TF_ImportGraphDefResults(); + mutex_lock l(graph->mu); + GraphImportGraphDefLocked(graph, *graph_def_ptr, options, results, status); + if (!status->status.ok()) { + delete results; + return nullptr; + } + return results; +} + void TF_GraphImportGraphDefWithReturnOutputs( TF_Graph* graph, const TF_Buffer* graph_def, const TF_ImportGraphDefOptions* options, TF_Output* return_outputs, diff --git a/tensorflow/c/c_api.h b/tensorflow/c/c_api.h index 2f4cf6062de04c..9812b0a7dfcef3 100644 --- a/tensorflow/c/c_api.h +++ b/tensorflow/c/c_api.h @@ -834,6 +834,14 @@ TF_GraphImportGraphDefWithResults(TF_Graph* graph, const TF_Buffer* graph_def, const TF_ImportGraphDefOptions* options, TF_Status* status); +// Has the same behavior as TF_GraphImportGraphDefWithResults, but instead of +// taking in a serialized tensorflow::GraphDef, it takes in a *pointer* to the +// C++ *in memory representation* of the GraphDef, stored in `graph_def->data` +TF_CAPI_EXPORT extern TF_ImportGraphDefResults* +TF_GraphImportGraphDefWithResultsNoSerialization( + TF_Graph* graph, const TF_Buffer* graph_def, + const TF_ImportGraphDefOptions* options, TF_Status* status); + // Import the graph serialized in `graph_def` into `graph`. // Convenience function for when only return outputs are needed. // diff --git a/tensorflow/c/c_api_test.cc b/tensorflow/c/c_api_test.cc index 008e2d772a31e7..e50221aadeebfb 100644 --- a/tensorflow/c/c_api_test.cc +++ b/tensorflow/c/c_api_test.cc @@ -23,6 +23,7 @@ limitations under the License. #include "tensorflow/c/c_api_internal.h" #include "tensorflow/c/c_test_util.h" +#include "tensorflow/c/tf_buffer.h" #include "tensorflow/c/tf_buffer_internal.h" #include "tensorflow/c/tf_status.h" #include "tensorflow/cc/saved_model/signature_constants.h" @@ -764,8 +765,15 @@ TEST(CAPI, ImportGraphDef) { EXPECT_EQ(2, TF_ImportGraphDefOptionsNumReturnOutputs(opts)); TF_ImportGraphDefOptionsAddReturnOperation(opts, "scalar"); EXPECT_EQ(1, TF_ImportGraphDefOptionsNumReturnOperations(opts)); + tensorflow::GraphDef graph_def_proto; + ASSERT_TRUE(tensorflow::ParseProtoUnlimited(&graph_def_proto, graph_def->data, + graph_def->length)); + TF_Buffer graph_def_buffer; + graph_def_buffer.data = reinterpret_cast(&graph_def_proto); + graph_def_buffer.length = sizeof(tensorflow::GraphDef*); TF_ImportGraphDefResults* results = - TF_GraphImportGraphDefWithResults(graph, graph_def, opts, s); + TF_GraphImportGraphDefWithResultsNoSerialization(graph, &graph_def_buffer, + opts, s); ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); TF_Operation* scalar2 = TF_GraphOperationByName(graph, "imported2/scalar"); @@ -956,8 +964,16 @@ TEST(CAPI, ImportGraphDef_MissingUnusedInputMappings) { TF_ImportGraphDefOptionsSetPrefix(opts, "imported"); TF_ImportGraphDefOptionsAddInputMapping(opts, "scalar", 0, {scalar, 0}); TF_ImportGraphDefOptionsAddInputMapping(opts, "fake", 0, {scalar, 0}); + + tensorflow::GraphDef graph_def_proto; + ASSERT_TRUE(tensorflow::ParseProtoUnlimited(&graph_def_proto, graph_def->data, + graph_def->length)); + TF_Buffer graph_def_buffer; + graph_def_buffer.data = reinterpret_cast(&graph_def_proto); + graph_def_buffer.length = sizeof(tensorflow::GraphDef*); TF_ImportGraphDefResults* results = - TF_GraphImportGraphDefWithResults(graph, graph_def, opts, s); + TF_GraphImportGraphDefWithResultsNoSerialization(graph, &graph_def_buffer, + opts, s); ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); // Check unused input mappings diff --git a/tensorflow/c/experimental/next_pluggable_device/BUILD b/tensorflow/c/experimental/next_pluggable_device/BUILD index c5cd8d15c3642a..5c7bbddc3af6f2 100644 --- a/tensorflow/c/experimental/next_pluggable_device/BUILD +++ b/tensorflow/c/experimental/next_pluggable_device/BUILD @@ -57,6 +57,7 @@ cc_library( hdrs = ["tensor_pjrt_buffer_util.h"], visibility = ["//visibility:public"], deps = [ + "//tensorflow/compiler/jit:pjrt_tensor_buffer_util", "//tensorflow/core:framework", "//tensorflow/core/tfrt/common:async_value_tensor", "//tensorflow/core/tfrt/common:global_state", @@ -64,9 +65,11 @@ cc_library( "//tensorflow/core/tfrt/common:pjrt_util", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/strings", "@local_tsl//tsl/platform:errors", "@local_tsl//tsl/platform:statusor", "@local_xla//xla/pjrt:pjrt_c_api_client", + "@local_xla//xla/pjrt:pjrt_client", "@local_xla//xla/pjrt/c:pjrt_c_api_hdrs", ], ) @@ -80,14 +83,18 @@ tf_cc_test( "//tensorflow/core:framework_types_hdr", "//tensorflow/core/tfrt/common:async_value_tensor", "//tensorflow/core/tfrt/common:pjrt_util", + "@com_google_absl//absl/log:check", "@com_google_googletest//:gtest_main", "@local_tsl//tsl/lib/core:status_test_util", + "@local_tsl//tsl/platform:casts", "@local_tsl//tsl/platform:status_matchers", "@local_tsl//tsl/protobuf:error_codes_proto_impl_cc", + "@local_xla//xla:shape_util", "@local_xla//xla/pjrt:pjrt_api", "@local_xla//xla/pjrt:pjrt_c_api_client", "@local_xla//xla/pjrt:tfrt_cpu_pjrt_client", "@local_xla//xla/pjrt/c:pjrt_c_api_cpu", + "@local_xla//xla/pjrt/c:pjrt_c_api_hdrs", "@local_xla//xla/pjrt/c:pjrt_c_api_wrapper_impl", ], ) diff --git a/tensorflow/c/experimental/next_pluggable_device/tensor_pjrt_buffer_util.cc b/tensorflow/c/experimental/next_pluggable_device/tensor_pjrt_buffer_util.cc index 02f9388aa247b3..18a851e394aea7 100644 --- a/tensorflow/c/experimental/next_pluggable_device/tensor_pjrt_buffer_util.cc +++ b/tensorflow/c/experimental/next_pluggable_device/tensor_pjrt_buffer_util.cc @@ -15,11 +15,15 @@ limitations under the License. #include "tensorflow/c/experimental/next_pluggable_device/tensor_pjrt_buffer_util.h" #include +#include #include "absl/status/status.h" #include "absl/status/statusor.h" +#include "absl/strings/str_cat.h" +#include "tensorflow/compiler/jit/pjrt_tensor_buffer_util.h" #include "xla/pjrt/c/pjrt_c_api.h" #include "xla/pjrt/pjrt_c_api_client.h" +#include "xla/pjrt/pjrt_client.h" #include "tensorflow/core/framework/resource_mgr.h" #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/framework/types.h" @@ -50,14 +54,16 @@ absl::StatusOr GetPjRtCBufferFromTensor(const Tensor* tensor) { absl::Status SetPjRtCBufferToTensor(PJRT_Buffer* c_buffer, xla::PjRtCApiClient* c_api_client, Tensor* tensor) { + auto buffer = std::make_unique(c_api_client, c_buffer); tensorflow::AsyncValueTensor* av_tensor = tensorflow::AsyncValueTensor::FromTensor(tensor); if (av_tensor == nullptr) { - return absl::InternalError( - "The tensor to set PjRtBuffer is not an AsyncValueTensor."); + TF_ASSIGN_OR_RETURN( + *tensor, MakeTensorFromPjRtBuffer(tensor->dtype(), tensor->shape(), + std::move(buffer))); + } else { + av_tensor->SetBuffer(std::move(buffer)); } - av_tensor->SetBuffer( - std::make_unique(c_api_client, c_buffer)); return absl::OkStatus(); } diff --git a/tensorflow/c/experimental/next_pluggable_device/tensor_pjrt_buffer_util_test.cc b/tensorflow/c/experimental/next_pluggable_device/tensor_pjrt_buffer_util_test.cc index 06fbd7b69d293d..c72f0cfafa6ead 100644 --- a/tensorflow/c/experimental/next_pluggable_device/tensor_pjrt_buffer_util_test.cc +++ b/tensorflow/c/experimental/next_pluggable_device/tensor_pjrt_buffer_util_test.cc @@ -14,20 +14,27 @@ limitations under the License. ==============================================================================*/ #include "tensorflow/c/experimental/next_pluggable_device/tensor_pjrt_buffer_util.h" +#include #include +#include #include #include #include +#include "absl/log/check.h" +#include "xla/pjrt/c/pjrt_c_api.h" #include "xla/pjrt/c/pjrt_c_api_cpu.h" #include "xla/pjrt/c/pjrt_c_api_wrapper_impl.h" #include "xla/pjrt/pjrt_api.h" #include "xla/pjrt/pjrt_c_api_client.h" #include "xla/pjrt/tfrt_cpu_pjrt_client.h" +#include "xla/shape.h" +#include "xla/shape_util.h" #include "tensorflow/core/framework/types.h" #include "tensorflow/core/tfrt/common/async_value_tensor.h" #include "tensorflow/core/tfrt/common/pjrt_util.h" #include "tsl/lib/core/status_test_util.h" +#include "tsl/platform/casts.h" #include "tsl/platform/status_matchers.h" #include "tsl/protobuf/error_codes.pb.h" @@ -38,6 +45,27 @@ using ::testing::HasSubstr; using ::testing::NotNull; using ::tsl::testing::StatusIs; +PJRT_Buffer* CreateCBuffer() { + auto status = pjrt::PjrtApi(DEVICE_CPU); + if (!status.ok()) { + CHECK_OK(pjrt::SetPjrtApi(DEVICE_CPU, GetPjrtApi())); + } + auto pjrt_client = xla::GetCApiClient(DEVICE_CPU); + CHECK_OK(pjrt_client.status()); + auto c_api_client = down_cast(pjrt_client->get()); + std::vector data(1, 0); + xla::Shape shape = xla::ShapeUtil::MakeShape(xla::S32, {1}); + + auto buffer = c_api_client->pjrt_c_client()->client->BufferFromHostBuffer( + data.data(), shape.element_type(), shape.dimensions(), + /*byte_strides=*/std::nullopt, + xla::PjRtClient::HostBufferSemantics::kImmutableOnlyDuringCall, nullptr, + c_api_client->pjrt_c_client()->client->addressable_devices()[0]); + CHECK_OK(buffer.status()); + + return new PJRT_Buffer{std::move(*buffer), c_api_client->pjrt_c_client()}; +} + TEST(TensorPjRtBufferUtilTest, GetPjRtCBufferFromTensorNoBuffer) { auto allocator = std::make_unique(); tensorflow::Tensor tensor(allocator.get(), DT_FLOAT, {1}); @@ -103,36 +131,18 @@ TEST(TensorPjRtBufferUtilTest, GetPjRtCBufferFromTensorSuccess) { TEST(TensorPjRtBufferUtilTest, SetPjRtCBufferToTensorNotAsyncValueTensor) { tensorflow::Tensor tensor(DT_FLOAT, {1}); + TF_ASSERT_OK_AND_ASSIGN(auto pjrt_client, xla::GetCApiClient(DEVICE_CPU)); + PJRT_Buffer* c_buffer = CreateCBuffer(); - EXPECT_THAT( - SetPjRtCBufferToTensor(nullptr, nullptr, &tensor), - StatusIs( - error::INTERNAL, - HasSubstr(absl::StrCat( - "The tensor to set PjRtBuffer is not an AsyncValueTensor")))); + TF_EXPECT_OK(SetPjRtCBufferToTensor( + c_buffer, down_cast(pjrt_client.get()), &tensor)); } TEST(TensorPjRtBufferUtilTest, SetPjRtCBufferToTensorSuccess) { auto allocator = std::make_unique(); - auto status = pjrt::PjrtApi(DEVICE_CPU); - if (!status.ok()) { - TF_ASSERT_OK(pjrt::SetPjrtApi(DEVICE_CPU, GetPjrtApi())); - } - TF_ASSERT_OK_AND_ASSIGN(auto pjrt_client, xla::GetCApiClient(DEVICE_CPU)); - auto c_api_client = down_cast(pjrt_client.get()); - std::vector data(1, 0); - xla::Shape shape = xla::ShapeUtil::MakeShape(xla::S32, {1}); - TF_ASSERT_OK_AND_ASSIGN( - auto buffer, - c_api_client->pjrt_c_client()->client->BufferFromHostBuffer( - data.data(), shape.element_type(), shape.dimensions(), - /*byte_strides=*/std::nullopt, - xla::PjRtClient::HostBufferSemantics::kImmutableOnlyDuringCall, - nullptr, - c_api_client->pjrt_c_client()->client->addressable_devices()[0])); tensorflow::Tensor tensor(allocator.get(), DT_FLOAT, {1}); - auto c_buffer = - new PJRT_Buffer{std::move(buffer), c_api_client->pjrt_c_client()}; + TF_ASSERT_OK_AND_ASSIGN(auto pjrt_client, xla::GetCApiClient(DEVICE_CPU)); + PJRT_Buffer* c_buffer = CreateCBuffer(); TF_EXPECT_OK(SetPjRtCBufferToTensor( c_buffer, down_cast(pjrt_client.get()), &tensor)); diff --git a/tensorflow/cc/framework/cc_op_gen_util.cc b/tensorflow/cc/framework/cc_op_gen_util.cc index d0c65d10b2184d..0a64525f477c8a 100644 --- a/tensorflow/cc/framework/cc_op_gen_util.cc +++ b/tensorflow/cc/framework/cc_op_gen_util.cc @@ -15,6 +15,7 @@ limitations under the License. #include "tensorflow/cc/framework/cc_op_gen_util.h" +#include #include #include #include @@ -29,6 +30,7 @@ limitations under the License. #include "tensorflow/core/framework/tensor_shape.pb.h" #include "tensorflow/core/lib/gtl/map_util.h" #include "tensorflow/core/lib/io/path.h" +#include "tensorflow/core/platform/strcat.h" #include "tensorflow/core/platform/types.h" #include "tsl/platform/statusor.h" @@ -206,7 +208,12 @@ string PrintAttrValue(const string& op, const AttrValue& attr_value) { return strings::StrCat(attr_value.i()); case AttrValue::kF: { const float f = attr_value.f(); - return strings::StrCat(attr_value.f(), floorf(f) == f ? ".0" : "", "f"); + if (std::isinf(f)) { + return strings::StrCat(f < 0.0f ? "-" : "+", + "std::numeric_limits::infinity()"); + } else { + return strings::StrCat(attr_value.f(), floorf(f) == f ? ".0" : "", "f"); + } } case AttrValue::kB: return attr_value.b() ? "true" : "false"; diff --git a/tensorflow/cc/saved_model/BUILD b/tensorflow/cc/saved_model/BUILD index b09ebe0bd0944c..935b37b37aa5c0 100644 --- a/tensorflow/cc/saved_model/BUILD +++ b/tensorflow/cc/saved_model/BUILD @@ -522,6 +522,7 @@ cc_library( visibility = [ "//learning/brain/contrib/hub/server/distro:__subpackages__", "//learning/brain/contrib/tpu_modeling:__subpackages__", + "//learning/metadata/artifactoid/cc:__subpackages__", "//learning/tfx/pipeline/util:__subpackages__", "//tensorflow/python/saved_model:__subpackages__", ], diff --git a/tensorflow/compiler/jit/device_compiler_client.cc b/tensorflow/compiler/jit/device_compiler_client.cc index 747aa7de090125..5d0042c3a4ed76 100644 --- a/tensorflow/compiler/jit/device_compiler_client.cc +++ b/tensorflow/compiler/jit/device_compiler_client.cc @@ -37,8 +37,6 @@ xla::ExecutableBuildOptions GetExecutableBuildOptions( build_options.set_alias_passthrough_params(options.alias_passthrough_params); build_options.mutable_debug_options()->set_xla_detailed_logging( options.detailed_logging); - build_options.mutable_debug_options()->set_xla_enable_dumping( - options.detailed_logging); if (tensorflow::OpDeterminismRequired()) { build_options.mutable_debug_options()->set_xla_gpu_deterministic_ops(true); } diff --git a/tensorflow/compiler/jit/device_compiler_client_test.cc b/tensorflow/compiler/jit/device_compiler_client_test.cc index 4ac2e7f1a205d1..f42ae36f50735b 100644 --- a/tensorflow/compiler/jit/device_compiler_client_test.cc +++ b/tensorflow/compiler/jit/device_compiler_client_test.cc @@ -60,5 +60,17 @@ TEST(GetExecutableOptionTest, DeviceOrdinalNotSet) { EXPECT_EQ(build_option.device_ordinal(), -1); } +TEST(GetExecutableOptionTest, DumpingWithoutDetailedLogging) { + XlaCompiler::Options options; + options.detailed_logging = false; + XlaCompiler::CompilationResult result; + + auto build_option = + GetExecutableBuildOptions(options, result, /*default_device_ordinal=*/-1); + + EXPECT_FALSE(build_option.debug_options().xla_detailed_logging()); + EXPECT_TRUE(build_option.debug_options().xla_enable_dumping()); +} + } // namespace } // namespace tensorflow diff --git a/tensorflow/compiler/jit/get_compiler_ir.cc b/tensorflow/compiler/jit/get_compiler_ir.cc index 37987cb55ecf38..2614d9909bbd80 100644 --- a/tensorflow/compiler/jit/get_compiler_ir.cc +++ b/tensorflow/compiler/jit/get_compiler_ir.cc @@ -74,8 +74,6 @@ static StatusOr> BuildExecutable( build_options.set_alias_passthrough_params(options.alias_passthrough_params); build_options.mutable_debug_options()->set_xla_detailed_logging( options.detailed_logging); - build_options.mutable_debug_options()->set_xla_enable_dumping( - options.detailed_logging); // If the embed_ir_in_executable is set, hlo_proto will be dumped in // executable. The hlo_proto contains HLO modules and buffer assignment. build_options.mutable_debug_options()->set_xla_embed_ir_in_executable( @@ -148,7 +146,7 @@ static StatusOr> BuildXlaCompilerArgumentFromTensorSpec( const FunctionBody* fbody, absl::Span must_be_constant_idxs, absl::Span inputs, - absl::Span variable_args, Device* device, + absl::Span variable_args, absl::Span flat_arg_shape_and_dtype) { TF_RET_CHECK(fbody != nullptr); auto& input_args = fbody->fdef.signature().input_arg(); @@ -326,7 +324,7 @@ StatusOr GetCompilerIr( if (compiler_arg_source == CompilerArgSource::TENSOR_SPEC) { args = BuildXlaCompilerArgumentFromTensorSpec(fbody, constant_arg_indices, - inputs, variable_infos, dev, + inputs, variable_infos, input_arg_shape_and_dtype); } else if (compiler_arg_source == CompilerArgSource::CONCRETE_INPUT) { args = XlaComputationLaunchContext::BuildXlaCompilerArguments( diff --git a/tensorflow/compiler/mlir/BUILD b/tensorflow/compiler/mlir/BUILD index 28e2753c68649c..0bde3279c1d7f6 100644 --- a/tensorflow/compiler/mlir/BUILD +++ b/tensorflow/compiler/mlir/BUILD @@ -61,6 +61,7 @@ cc_library( "//tensorflow/compiler/mlir/tensorflow/transforms:tf_saved_model_passes", # buildcleaner:keep "//tensorflow/compiler/mlir/tensorflow/transforms/host_runtime:lower_cluster_to_runtime_ops", "//tensorflow/compiler/mlir/tf2xla:compile_mlir_util", + "//tensorflow/compiler/mlir/tf2xla/internal/passes:clustering_passes", "//tensorflow/compiler/mlir/tf2xla/transforms:tf_xla_passes", "//tensorflow/compiler/mlir/tf2xla/transforms:xla_legalize_tf", "//tensorflow/compiler/mlir/tosa:tf_passes", diff --git a/tensorflow/compiler/mlir/lite/BUILD b/tensorflow/compiler/mlir/lite/BUILD index 87de9412f32756..8117705b0fac2b 100644 --- a/tensorflow/compiler/mlir/lite/BUILD +++ b/tensorflow/compiler/mlir/lite/BUILD @@ -12,6 +12,7 @@ package( "//learning/brain/mobile/programmability:__subpackages__", "//tensorflow/lite/experimental/tf_runtime:__subpackages__", "//tensorflow/lite/testing:__subpackages__", + "//third_party/odml/infra/genai/conversion/per_layer:__subpackages__", ], licenses = ["notice"], ) @@ -873,8 +874,8 @@ cc_library( "transforms/lift_tflite_flex_ops.h", ], deps = [ + ":tensorflow_lite", ":tensorflow_lite_passes_inc_gen", - "//tensorflow/compiler/mlir/lite:tensorflow_lite", "//tensorflow/compiler/mlir/tensorflow", "//tensorflow/compiler/mlir/tensorflow:convert_attr", "//tensorflow/compiler/mlir/tensorflow:tensorflow_attributes", @@ -883,6 +884,7 @@ cc_library( "//tensorflow/core:framework", "//tensorflow/core:lib", "//tensorflow/core:protos_all_cc", + "@com_google_absl//absl/strings", "@flatbuffers", "@llvm-project//mlir:FuncDialect", "@llvm-project//mlir:IR", diff --git a/tensorflow/compiler/mlir/lite/experimental/tac/tests/device-transform-gpu.mlir b/tensorflow/compiler/mlir/lite/experimental/tac/tests/device-transform-gpu.mlir index b53508f01efa2d..cbd2dd6618f93d 100644 --- a/tensorflow/compiler/mlir/lite/experimental/tac/tests/device-transform-gpu.mlir +++ b/tensorflow/compiler/mlir/lite/experimental/tac/tests/device-transform-gpu.mlir @@ -153,8 +153,8 @@ func.func @padSliceTo4D(%arg0: tensor<4x384x32xf32>) -> tensor<1x384x32xf32> { } // CHECK: func @padSliceTo4D(%[[VAL_0:.*]]: tensor<4x384x32xf32>) -> tensor<1x384x32xf32> { -// CHECK-DAG: %[[VAL_1:.*]] = "tf.Const"() {value = dense<0> : tensor<4xi32>} : () -> tensor<4xi32> -// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() {value = dense<[1, 1, 384, 32]> : tensor<4xi32>} : () -> tensor<4xi32> +// CHECK-DAG: %[[VAL_1:.*]] = "tf.Const"() <{value = dense<0> : tensor<4xi32>}> : () -> tensor<4xi32> +// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() <{value = dense<[1, 1, 384, 32]> : tensor<4xi32>}> : () -> tensor<4xi32> // CHECK-DAG: %[[VAL_3:.*]] = arith.constant dense<[1, 4, 384, 32]> : tensor<4xi32> // CHECK-DAG: %[[VAL_4:.*]] = "tfl.pseudo_const"() {value = dense<[1, 384, 32]> : tensor<3xi32> // CHECK: %[[VAL_5:.*]] = "tfl.reshape"(%[[VAL_0]], %[[VAL_3]]) : (tensor<4x384x32xf32>, tensor<4xi32>) -> tensor<1x4x384x32xf32> diff --git a/tensorflow/compiler/mlir/lite/flatbuffer_export.cc b/tensorflow/compiler/mlir/lite/flatbuffer_export.cc index 7070f5795ca3a0..be4aae3d95db08 100644 --- a/tensorflow/compiler/mlir/lite/flatbuffer_export.cc +++ b/tensorflow/compiler/mlir/lite/flatbuffer_export.cc @@ -675,6 +675,11 @@ class Translator { std::optional>> CreateMetadataVector(); + // Encodes the `tfl.metadata_buffer` array attribute of the module to the + // metadata_buffer section in the final model. Returns empty if there isn't + // such attribute in the mlir module. + VectorBufferOffset CreateMetadataBufferVector(); + // Builds and returns list of tfl.SignatureDef sections in the model. std::optional>> CreateSignatureDefs(const std::vector& signature_defs); @@ -1635,6 +1640,14 @@ std::optional> Translator::BuildOperator( if (auto shlo_op = llvm::dyn_cast(inst)) { return BuildStablehloGatherOp(shlo_op, operands, results); } + if (auto shlo_op = llvm::dyn_cast(inst)) { + return BuildStablehloOperatorwithoutOptions( + inst, operands, results, tflite::BuiltinOperator_STABLEHLO_ADD); + } + if (auto shlo_op = llvm::dyn_cast(inst)) { + return BuildStablehloOperatorwithoutOptions( + inst, operands, results, tflite::BuiltinOperator_STABLEHLO_MULTIPLY); + } // for ops don't have kernels, only serialize when conversion is set to true if (convert_stablehlo_) { if (auto shlo_op = llvm::dyn_cast(inst)) { @@ -1643,17 +1656,6 @@ std::optional> Translator::BuildOperator( tflite::BuiltinOperator_STABLEHLO_LOGISTIC); } - if (auto shlo_op = llvm::dyn_cast(inst)) { - return BuildStablehloOperatorwithoutOptions( - inst, operands, results, tflite::BuiltinOperator_STABLEHLO_ADD); - } - - if (auto shlo_op = llvm::dyn_cast(inst)) { - return BuildStablehloOperatorwithoutOptions( - inst, operands, results, - tflite::BuiltinOperator_STABLEHLO_MULTIPLY); - } - if (auto shlo_op = llvm::dyn_cast(inst)) { return BuildStablehloOperatorwithoutOptions( inst, operands, results, tflite::BuiltinOperator_STABLEHLO_DIVIDE); @@ -2608,6 +2610,18 @@ Translator::CreateMetadataVector() { return builder_.CreateVector(metadata); } +VectorBufferOffset Translator::CreateMetadataBufferVector() { + auto array_attr = + module_->getAttrOfType("tfl.metadata_buffer"); + std::vector metadata_buffer; + if (!array_attr) return 0; + for (auto value : array_attr.getAsValueRange()) { + metadata_buffer.push_back(value.getSExtValue()); + } + + return builder_.CreateVector(metadata_buffer); +} + // Helper method that returns list of all strings in a StringAttr identified // by 'attr_key' and values are separated by a comma. llvm::SmallVector GetStringsFromAttrWithSeparator( @@ -3010,7 +3024,8 @@ std::optional Translator::TranslateInternal() { // Build the model and finish the model building process. auto description = builder_.CreateString(model_description.data()); - VectorBufferOffset metadata_buffer = 0; // Deprecated + VectorBufferOffset metadata_buffer = + CreateMetadataBufferVector(); // Deprecated auto metadata = CreateMetadataVector(); if (!metadata) return std::nullopt; diff --git a/tensorflow/compiler/mlir/lite/flatbuffer_import.cc b/tensorflow/compiler/mlir/lite/flatbuffer_import.cc index 40943740ba1f00..6eb2aee99aacb9 100644 --- a/tensorflow/compiler/mlir/lite/flatbuffer_import.cc +++ b/tensorflow/compiler/mlir/lite/flatbuffer_import.cc @@ -1904,6 +1904,11 @@ OwningOpRef tflite::FlatBufferToMlir( mlir::UnitAttr::get(builder.getContext())); } + if (!model->metadata_buffer.empty()) { + module->setAttr("tfl.metadata_buffer", + builder.getI32ArrayAttr(model->metadata_buffer)); + } + if (use_stablehlo_constant) { module->setAttr("tfl.metadata", builder.getDictionaryAttr(builder.getNamedAttr( diff --git a/tensorflow/compiler/mlir/lite/metrics/error_collector_inst_test.cc b/tensorflow/compiler/mlir/lite/metrics/error_collector_inst_test.cc index 83d0a0e6cf1ccf..798473986f64fd 100644 --- a/tensorflow/compiler/mlir/lite/metrics/error_collector_inst_test.cc +++ b/tensorflow/compiler/mlir/lite/metrics/error_collector_inst_test.cc @@ -160,7 +160,7 @@ TEST(ErrorCollectorTest, TessFailurePass) { EXPECT_EQ(collected_errors.count(NewConverterErrorData( "MockFailurePass", "Failed at tf.Const op\nsee current operation: %0 = " - "\"tf.Const\"() {value = dense<1> : tensor<4xi32>} : () -> " + "\"tf.Const\"() <{value = dense<1> : tensor<4xi32>}> : () -> " "tensor<4xi32>\nError code: ERROR_NEEDS_FLEX_OPS", ConverterErrorData::ERROR_NEEDS_FLEX_OPS, "tf.Const", mlir::FileLineColLoc::get(input_file_id, 2, 9))), @@ -168,22 +168,23 @@ TEST(ErrorCollectorTest, TessFailurePass) { EXPECT_EQ(collected_errors.count(NewConverterErrorData( "MockFailurePass", "Failed at tf.Const op\nsee current operation: %1 = " - "\"tf.Const\"() {value = dense<0> : tensor<4xi32>} : () -> " + "\"tf.Const\"() <{value = dense<0> : tensor<4xi32>}> : () -> " "tensor<4xi32>\nError code: ERROR_NEEDS_FLEX_OPS", ConverterErrorData::ERROR_NEEDS_FLEX_OPS, "tf.Const", mlir::FileLineColLoc::get(input_file_id, 2, 9))), 1); - EXPECT_EQ(collected_errors.count(NewConverterErrorData( - "MockFailurePass", - "Failed at tf.StridedSlice op\nsee current operation: %2 = " - "\"tf.StridedSlice\"(%arg0, %1, %1, %0) {begin_mask = 11 : " - "i64, device = \"\", ellipsis_mask = 0 : i64, end_mask = 11 : " - "i64, new_axis_mask = 4 : i64, shrink_axis_mask = 0 : i64} : " - "(tensor<*xf32>, tensor<4xi32>, tensor<4xi32>, tensor<4xi32>) " - "-> tensor<*xf32>\nError code: ERROR_NEEDS_FLEX_OPS", - ConverterErrorData::ERROR_NEEDS_FLEX_OPS, "tf.StridedSlice", - mlir::FileLineColLoc::get(input_file_id, 4, 10))), - 1); + EXPECT_EQ( + collected_errors.count(NewConverterErrorData( + "MockFailurePass", + "Failed at tf.StridedSlice op\nsee current operation: %2 = " + "\"tf.StridedSlice\"(%arg0, %1, %1, %0) <{begin_mask = 11 : " + "i64, ellipsis_mask = 0 : i64, end_mask = 11 : i64, new_axis_mask = " + "4 : i64, shrink_axis_mask = 0 : i64}> {device = \"\"} : " + "(tensor<*xf32>, tensor<4xi32>, tensor<4xi32>, tensor<4xi32>) " + "-> tensor<*xf32>\nError code: ERROR_NEEDS_FLEX_OPS", + ConverterErrorData::ERROR_NEEDS_FLEX_OPS, "tf.StridedSlice", + mlir::FileLineColLoc::get(input_file_id, 4, 10))), + 1); // Check the location information. std::vector locations; diff --git a/tensorflow/compiler/mlir/lite/quantization/tensorflow/tests/fallback_to_flex_ops_default.mlir b/tensorflow/compiler/mlir/lite/quantization/tensorflow/tests/fallback_to_flex_ops_default.mlir index 589d438311ac50..9c6d9b8aa8059b 100644 --- a/tensorflow/compiler/mlir/lite/quantization/tensorflow/tests/fallback_to_flex_ops_default.mlir +++ b/tensorflow/compiler/mlir/lite/quantization/tensorflow/tests/fallback_to_flex_ops_default.mlir @@ -4,7 +4,7 @@ func.func @bias_add(%arg0: tensor<1x10x10x32xf32>, %arg1: tensor<32xf32>) -> tensor<1x10x10x32xf32> { %0 = "tf.BiasAdd"(%arg0, %arg1) {T = "tfdtype$DT_FLOAT", data_format = "NHWC"} : (tensor<1x10x10x32xf32>, tensor<32xf32>) -> tensor<1x10x10x32xf32> func.return %0 : tensor<1x10x10x32xf32> -// CHECK: %[[BIASADD_0:.*]] = "tf.BiasAdd"(%arg0, %arg1) {T = "tfdtype$DT_FLOAT", data_format = "NHWC"} : (tensor<1x10x10x32xf32>, tensor<32xf32>) -> tensor<1x10x10x32xf32> +// CHECK: %[[BIASADD_0:.*]] = "tf.BiasAdd"(%arg0, %arg1) <{data_format = "NHWC"}> {T = "tfdtype$DT_FLOAT"} : (tensor<1x10x10x32xf32>, tensor<32xf32>) -> tensor<1x10x10x32xf32> // CHECK: return %[[BIASADD_0]] : tensor<1x10x10x32xf32> } @@ -30,8 +30,8 @@ func.func @conv2d_backprop_input_with_add(%arg0: tensor<4xi32>, %arg1: tensor<3x %1 = "tf.Const"() {value = dense<1.000000e+00> : tensor<1xf32>} : () -> tensor<1xf32> %2 = "tf.AddV2"(%0, %1): (tensor<15x28x28x1xf32>, tensor<1xf32>) -> tensor<15x28x28x1xf32> func.return %2 : tensor<15x28x28x1xf32> -// CHECK: %[[CONST_0:.*]] = "tf.Const"() {value = dense<1.000000e+00> : tensor<1xf32>} : () -> tensor<1xf32> -// CHECK: %[[CONV2DBACKPROPINPUT_0:.*]] = "tf.Conv2DBackpropInput"(%arg0, %arg1, %arg2) {dilations = [1, 1, 1, 1], padding = "SAME", strides = [1, 2, 2, 1]} : (tensor<4xi32>, tensor<3x3x1x32xf32>, tensor<15x14x14x32xf32>) -> tensor<15x28x28x1xf32> +// CHECK: %[[CONST_0:.*]] = "tf.Const"() <{value = dense<1.000000e+00> : tensor<1xf32>}> : () -> tensor<1xf32> +// CHECK: %[[CONV2DBACKPROPINPUT_0:.*]] = "tf.Conv2DBackpropInput"(%arg0, %arg1, %arg2) <{dilations = [1, 1, 1, 1], padding = "SAME", strides = [1, 2, 2, 1]}> : (tensor<4xi32>, tensor<3x3x1x32xf32>, tensor<15x14x14x32xf32>) -> tensor<15x28x28x1xf32> // CHECK: %[[ADDV2_0:.*]] = "tf.AddV2"(%[[CONV2DBACKPROPINPUT_0]], %[[CONST_0]]) : (tensor<15x28x28x1xf32>, tensor<1xf32>) -> tensor<15x28x28x1xf32> // CHECK: return %[[ADDV2_0]] : tensor<15x28x28x1xf32> } @@ -42,8 +42,8 @@ func.func @conv2d_backprop_input_with_sub(%arg0: tensor<4xi32>, %arg1: tensor<3x %1 = "tf.Const"() {value = dense<1.000000e+00> : tensor<1xf32>} : () -> tensor<1xf32> %2 = "tf.Sub"(%0, %1): (tensor<15x28x28x1xf32>, tensor<1xf32>) -> tensor<15x28x28x1xf32> func.return %2 : tensor<15x28x28x1xf32> -// CHECK: %[[CONST_0:.*]] = "tf.Const"() {value = dense<1.000000e+00> : tensor<1xf32>} : () -> tensor<1xf32> -// CHECK: %[[CONV2DBACKPROPINPUT_0:.*]] = "tf.Conv2DBackpropInput"(%arg0, %arg1, %arg2) {dilations = [1, 1, 1, 1], padding = "SAME", strides = [1, 2, 2, 1]} : (tensor<4xi32>, tensor<3x3x1x32xf32>, tensor<15x14x14x32xf32>) -> tensor<15x28x28x1xf32> +// CHECK: %[[CONST_0:.*]] = "tf.Const"() <{value = dense<1.000000e+00> : tensor<1xf32>}> : () -> tensor<1xf32> +// CHECK: %[[CONV2DBACKPROPINPUT_0:.*]] = "tf.Conv2DBackpropInput"(%arg0, %arg1, %arg2) <{dilations = [1, 1, 1, 1], padding = "SAME", strides = [1, 2, 2, 1]}> : (tensor<4xi32>, tensor<3x3x1x32xf32>, tensor<15x14x14x32xf32>) -> tensor<15x28x28x1xf32> // CHECK: %[[SUB_0:.*]] = "tf.Sub"(%[[CONV2DBACKPROPINPUT_0]], %[[CONST_0]]) : (tensor<15x28x28x1xf32>, tensor<1xf32>) -> tensor<15x28x28x1xf32> // CHECK: return %[[SUB_0]] : tensor<15x28x28x1xf32> } @@ -71,7 +71,7 @@ func.func @identity_with_const() -> tensor<*xf32> { %0 = "tf.Identity"(%cst) {device = ""} : (tensor<2xf32>) -> tensor<*xf32> %1 = "tf.AddV2"(%0, %cst_1) {device = ""} : (tensor<*xf32>, tensor) -> tensor<*xf32> func.return %1 : tensor<*xf32> -// CHECK: %[[CONST_0:.*]] = "tf.Const"() {value = dense<[2.177590e-01, 2.89503098]> : tensor<2xf32>} : () -> tensor<*xf32> +// CHECK: %[[CONST_0:.*]] = "tf.Const"() <{value = dense<[2.177590e-01, 2.89503098]> : tensor<2xf32>}> : () -> tensor<*xf32> // CHECK: return %[[CONST_0]] : tensor<*xf32> } @@ -80,7 +80,7 @@ func.func @identity(%arg0: tensor<2xf32>) -> tensor<*xf32> { %0 = "tf.Identity"(%arg0) {device = ""} : (tensor<2xf32>) -> tensor<*xf32> %1 = "tf.AddV2"(%0, %cst_1) {device = ""} : (tensor<*xf32>, tensor) -> tensor<*xf32> func.return %1 : tensor<*xf32> -// CHECK: %[[CONST_0:.*]] = "tf.Const"() {device = "", value = dense<1.000000e-03> : tensor} : () -> tensor +// CHECK: %[[CONST_0:.*]] = "tf.Const"() <{value = dense<1.000000e-03> : tensor}> {device = ""} : () -> tensor // CHECK: %[[IDENTITY_0:.*]] = "tf.Identity"(%arg0) {device = ""} : (tensor<2xf32>) -> tensor<*xf32> // CHECK: %[[ADDV2_0:.*]] = "tfl.custom"(%0, %cst) {custom_code = "FlexAddV2", custom_option = #tfl} : (tensor<*xf32>, tensor) -> tensor<*xf32> // CHECK: return %[[ADDV2_0]] : tensor<*xf32> @@ -95,7 +95,7 @@ func.func @bias_add_with_identity(%arg0: tensor<4xi32>, %arg1: tensor<3x3x1x32xf %2 = "tf.Conv2DBackpropInput"(%arg0, %arg1, %arg2) {strides = [1, 2, 2, 1], padding="SAME", dilations=[1, 1, 1, 1]}: (tensor<4xi32>, tensor<3x3x1x32xf32>, tensor<15x14x14x32xf32>) -> tensor<15x28x28x1xf32> %3 = "tf.AddV2"(%2, %1): (tensor<15x28x28x1xf32>, tensor<1xf32>) -> tensor<15x28x28x1xf32> func.return %2 : tensor<15x28x28x1xf32> -// CHECK: %[[CONV2DBACKPROPINPUT_0:.*]] = "tf.Conv2DBackpropInput"(%arg0, %arg1, %arg2) {dilations = [1, 1, 1, 1], padding = "SAME", strides = [1, 2, 2, 1]} : (tensor<4xi32>, tensor<3x3x1x32xf32>, tensor<15x14x14x32xf32>) -> tensor<15x28x28x1xf32> +// CHECK: %[[CONV2DBACKPROPINPUT_0:.*]] = "tf.Conv2DBackpropInput"(%arg0, %arg1, %arg2) <{dilations = [1, 1, 1, 1], padding = "SAME", strides = [1, 2, 2, 1]}> : (tensor<4xi32>, tensor<3x3x1x32xf32>, tensor<15x14x14x32xf32>) -> tensor<15x28x28x1xf32> // CHECK: return %[[CONV2DBACKPROPINPUT_0]] : tensor<15x28x28x1xf32> } @@ -108,10 +108,10 @@ func.func @conv_with_relu1_pattern1(%arg0: tensor<1x3x4x3xf32>) -> (tensor<1x3x4 %1 = "tf.Maximum"(%0, %cst_0) : (tensor<1x3x4x2xf32>, tensor) -> tensor<1x3x4x2xf32> %2 = "tf.Minimum"(%1, %cst_1) : (tensor<1x3x4x2xf32>, tensor) -> tensor<1x3x4x2xf32> func.return %2 : tensor<1x3x4x2xf32> -// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() {value = dense<{{.*}}> : tensor<1x1x3x2xf32>} : () -> tensor<1x1x3x2xf32> -// CHECK-DAG: %[[CONST_1:.*]] = "tf.Const"() {value = dense<-1.000000e+00> : tensor} : () -> tensor -// CHECK-DAG: %[[CONST_2:.*]] = "tf.Const"() {value = dense<1.000000e+00> : tensor} : () -> tensor -// CHECK: %[[CONV2D_0:.*]] = "tf.Conv2D"(%arg0, %[[CONST_0]]) {data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 1, 1]} : (tensor<1x3x4x3xf32>, tensor<1x1x3x2xf32>) -> tensor<1x3x4x2xf32> +// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() <{value = dense<{{.*}}> : tensor<1x1x3x2xf32>}> : () -> tensor<1x1x3x2xf32> +// CHECK-DAG: %[[CONST_1:.*]] = "tf.Const"() <{value = dense<-1.000000e+00> : tensor}> : () -> tensor +// CHECK-DAG: %[[CONST_2:.*]] = "tf.Const"() <{value = dense<1.000000e+00> : tensor}> : () -> tensor +// CHECK: %[[CONV2D_0:.*]] = "tf.Conv2D"(%arg0, %[[CONST_0]]) <{data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 1, 1]}> : (tensor<1x3x4x3xf32>, tensor<1x1x3x2xf32>) -> tensor<1x3x4x2xf32> // CHECK: %[[MAXIMUM_0:.*]] = "tf.Maximum"(%[[CONV2D_0]], %[[CONST_1]]) : (tensor<1x3x4x2xf32>, tensor) -> tensor<1x3x4x2xf32> // CHECK: %[[MINIMUM_0:.*]] = "tf.Minimum"(%[[MAXIMUM_0]], %[[CONST_2]]) : (tensor<1x3x4x2xf32>, tensor) -> tensor<1x3x4x2xf32> // CHECK: return %[[MINIMUM_0]] : tensor<1x3x4x2xf32> @@ -126,10 +126,10 @@ func.func @conv_with_relu1_pattern2(%arg0: tensor<1x3x4x3xf32>) -> (tensor<1x3x4 %1 = "tf.Minimum"(%0, %cst_1) : (tensor<1x3x4x2xf32>, tensor) -> tensor<1x3x4x2xf32> %2 = "tf.Maximum"(%1, %cst_0) : (tensor<1x3x4x2xf32>, tensor) -> tensor<1x3x4x2xf32> func.return %2 : tensor<1x3x4x2xf32> -// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() {value = dense<{{.*}}> : tensor<1x1x3x2xf32>} : () -> tensor<1x1x3x2xf32> -// CHECK-DAG: %[[CONST_1:.*]] = "tf.Const"() {value = dense<-1.000000e+00> : tensor} : () -> tensor -// CHECK-DAG: %[[CONST_2:.*]] = "tf.Const"() {value = dense<1.000000e+00> : tensor} : () -> tensor -// CHECK: %[[CONV2D_0:.*]] = "tf.Conv2D"(%arg0, %[[CONST_0]]) {data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 1, 1]} : (tensor<1x3x4x3xf32>, tensor<1x1x3x2xf32>) -> tensor<1x3x4x2xf32> +// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() <{value = dense<{{.*}}> : tensor<1x1x3x2xf32>}> : () -> tensor<1x1x3x2xf32> +// CHECK-DAG: %[[CONST_1:.*]] = "tf.Const"() <{value = dense<-1.000000e+00> : tensor}> : () -> tensor +// CHECK-DAG: %[[CONST_2:.*]] = "tf.Const"() <{value = dense<1.000000e+00> : tensor}> : () -> tensor +// CHECK: %[[CONV2D_0:.*]] = "tf.Conv2D"(%arg0, %[[CONST_0]]) <{data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 1, 1]}> : (tensor<1x3x4x3xf32>, tensor<1x1x3x2xf32>) -> tensor<1x3x4x2xf32> // CHECK: %[[MINIMUM_0:.*]] = "tf.Minimum"(%[[CONV2D_0]], %[[CONST_2]]) : (tensor<1x3x4x2xf32>, tensor) -> tensor<1x3x4x2xf32> // CHECK: %[[MAXIMUM_0:.*]] = "tf.Maximum"(%[[MINIMUM_0]], %[[CONST_1]]) : (tensor<1x3x4x2xf32>, tensor) -> tensor<1x3x4x2xf32> // CHECK: return %[[MAXIMUM_0]] : tensor<1x3x4x2xf32> @@ -144,10 +144,10 @@ func.func @conv_with_relu1_invalid_pattern(%arg0: tensor<1x3x4x3xf32>) -> (tenso %1 = "tf.Minimum"(%0, %cst_1) : (tensor<1x3x4x2xf32>, tensor<2xf32>) -> tensor<1x3x4x2xf32> %2 = "tf.Maximum"(%1, %cst_0) : (tensor<1x3x4x2xf32>, tensor<2xf32>) -> tensor<1x3x4x2xf32> func.return %2 : tensor<1x3x4x2xf32> -// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() {value = dense<{{.*}}> : tensor<1x1x3x2xf32>} : () -> tensor<1x1x3x2xf32> -// CHECK-DAG: %[[CONST_1:.*]] = "tf.Const"() {value = dense<[-1.000000e+00, -3.000000e+00]> : tensor<2xf32>} : () -> tensor<2xf32> -// CHECK-DAG: %[[CONST_2:.*]] = "tf.Const"() {value = dense<[1.000000e+00, 3.000000e+00]> : tensor<2xf32>} : () -> tensor<2xf32> -// CHECK: %[[CONV2D_0:.*]] = "tf.Conv2D"(%arg0, %[[CONST_0]]) {data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 1, 1]} : (tensor<1x3x4x3xf32>, tensor<1x1x3x2xf32>) -> tensor<1x3x4x2xf32> +// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() <{value = dense<{{.*}}> : tensor<1x1x3x2xf32>}> : () -> tensor<1x1x3x2xf32> +// CHECK-DAG: %[[CONST_1:.*]] = "tf.Const"() <{value = dense<[-1.000000e+00, -3.000000e+00]> : tensor<2xf32>}> : () -> tensor<2xf32> +// CHECK-DAG: %[[CONST_2:.*]] = "tf.Const"() <{value = dense<[1.000000e+00, 3.000000e+00]> : tensor<2xf32>}> : () -> tensor<2xf32> +// CHECK: %[[CONV2D_0:.*]] = "tf.Conv2D"(%arg0, %[[CONST_0]]) <{data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 1, 1]}> : (tensor<1x3x4x3xf32>, tensor<1x1x3x2xf32>) -> tensor<1x3x4x2xf32> // CHECK: %[[CUSTOM_0:.*]] = "tfl.custom"(%[[CONV2D_0]], %[[CONST_2]]) {custom_code = "FlexMinimum", custom_option = #tfl} : (tensor<1x3x4x2xf32>, tensor<2xf32>) -> tensor<1x3x4x2xf32> // CHECK: %[[CUSTOM_1:.*]] = "tfl.custom"(%[[CUSTOM_0]], %[[CONST_1]]) {custom_code = "FlexMaximum", custom_option = #tfl} : (tensor<1x3x4x2xf32>, tensor<2xf32>) -> tensor<1x3x4x2xf32> // CHECK: return %[[CUSTOM_1]] : tensor<1x3x4x2xf32> diff --git a/tensorflow/compiler/mlir/lite/quantization/tensorflow/tests/fallback_to_flex_ops_legacy.mlir b/tensorflow/compiler/mlir/lite/quantization/tensorflow/tests/fallback_to_flex_ops_legacy.mlir index dea1a9cc823f0b..5835d7d107cef5 100644 --- a/tensorflow/compiler/mlir/lite/quantization/tensorflow/tests/fallback_to_flex_ops_legacy.mlir +++ b/tensorflow/compiler/mlir/lite/quantization/tensorflow/tests/fallback_to_flex_ops_legacy.mlir @@ -4,7 +4,7 @@ func.func @bias_add(%arg0: tensor<1x10x10x32xf32>, %arg1: tensor<32xf32>) -> tensor<1x10x10x32xf32> { %0 = "tf.BiasAdd"(%arg0, %arg1) {T = "tfdtype$DT_FLOAT", data_format = "NHWC"} : (tensor<1x10x10x32xf32>, tensor<32xf32>) -> tensor<1x10x10x32xf32> func.return %0 : tensor<1x10x10x32xf32> -// CHECK: %[[BIASADD_0:.*]] = "tf.BiasAdd"(%arg0, %arg1) {T = "tfdtype$DT_FLOAT", data_format = "NHWC"} : (tensor<1x10x10x32xf32>, tensor<32xf32>) -> tensor<1x10x10x32xf32> +// CHECK: %[[BIASADD_0:.*]] = "tf.BiasAdd"(%arg0, %arg1) <{data_format = "NHWC"}> {T = "tfdtype$DT_FLOAT"} : (tensor<1x10x10x32xf32>, tensor<32xf32>) -> tensor<1x10x10x32xf32> // CHECK: return %[[BIASADD_0]] : tensor<1x10x10x32xf32> } @@ -30,8 +30,8 @@ func.func @conv2d_backprop_input_with_add(%arg0: tensor<4xi32>, %arg1: tensor<3x %1 = "tf.Const"() {value = dense<1.000000e+00> : tensor<1xf32>} : () -> tensor<1xf32> %2 = "tf.AddV2"(%0, %1): (tensor<15x28x28x1xf32>, tensor<1xf32>) -> tensor<15x28x28x1xf32> func.return %2 : tensor<15x28x28x1xf32> -// CHECK: %[[CONST_0:.*]] = "tf.Const"() {value = dense<1.000000e+00> : tensor<1xf32>} : () -> tensor<1xf32> -// CHECK: %[[CONV2DBACKPROPINPUT_0:.*]] = "tf.Conv2DBackpropInput"(%arg0, %arg1, %arg2) {dilations = [1, 1, 1, 1], padding = "SAME", strides = [1, 2, 2, 1]} : (tensor<4xi32>, tensor<3x3x1x32xf32>, tensor<15x14x14x32xf32>) -> tensor<15x28x28x1xf32> +// CHECK: %[[CONST_0:.*]] = "tf.Const"() <{value = dense<1.000000e+00> : tensor<1xf32>}> : () -> tensor<1xf32> +// CHECK: %[[CONV2DBACKPROPINPUT_0:.*]] = "tf.Conv2DBackpropInput"(%arg0, %arg1, %arg2) <{dilations = [1, 1, 1, 1], padding = "SAME", strides = [1, 2, 2, 1]}> : (tensor<4xi32>, tensor<3x3x1x32xf32>, tensor<15x14x14x32xf32>) -> tensor<15x28x28x1xf32> // CHECK: %[[ADDV2_0:.*]] = "tf.AddV2"(%[[CONV2DBACKPROPINPUT_0]], %[[CONST_0]]) {no_fallback} : (tensor<15x28x28x1xf32>, tensor<1xf32>) -> tensor<15x28x28x1xf32> // CHECK: return %[[ADDV2_0]] : tensor<15x28x28x1xf32> } @@ -42,8 +42,8 @@ func.func @conv2d_backprop_input_with_sub(%arg0: tensor<4xi32>, %arg1: tensor<3x %1 = "tf.Const"() {value = dense<1.000000e+00> : tensor<1xf32>} : () -> tensor<1xf32> %2 = "tf.Sub"(%0, %1): (tensor<15x28x28x1xf32>, tensor<1xf32>) -> tensor<15x28x28x1xf32> func.return %2 : tensor<15x28x28x1xf32> -// CHECK: %[[CONST_0:.*]] = "tf.Const"() {value = dense<1.000000e+00> : tensor<1xf32>} : () -> tensor<1xf32> -// CHECK: %[[CONV2DBACKPROPINPUT_0:.*]] = "tf.Conv2DBackpropInput"(%arg0, %arg1, %arg2) {dilations = [1, 1, 1, 1], padding = "SAME", strides = [1, 2, 2, 1]} : (tensor<4xi32>, tensor<3x3x1x32xf32>, tensor<15x14x14x32xf32>) -> tensor<15x28x28x1xf32> +// CHECK: %[[CONST_0:.*]] = "tf.Const"() <{value = dense<1.000000e+00> : tensor<1xf32>}> : () -> tensor<1xf32> +// CHECK: %[[CONV2DBACKPROPINPUT_0:.*]] = "tf.Conv2DBackpropInput"(%arg0, %arg1, %arg2) <{dilations = [1, 1, 1, 1], padding = "SAME", strides = [1, 2, 2, 1]}> : (tensor<4xi32>, tensor<3x3x1x32xf32>, tensor<15x14x14x32xf32>) -> tensor<15x28x28x1xf32> // CHECK: %[[SUB_0:.*]] = "tf.Sub"(%[[CONV2DBACKPROPINPUT_0]], %[[CONST_0]]) {no_fallback} : (tensor<15x28x28x1xf32>, tensor<1xf32>) -> tensor<15x28x28x1xf32> // CHECK: return %[[SUB_0]] : tensor<15x28x28x1xf32> } @@ -73,6 +73,6 @@ func.func @bias_add_with_identity(%arg0: tensor<4xi32>, %arg1: tensor<3x3x1x32xf %2 = "tf.Conv2DBackpropInput"(%arg0, %arg1, %arg2) {strides = [1, 2, 2, 1], padding="SAME", dilations=[1, 1, 1, 1]}: (tensor<4xi32>, tensor<3x3x1x32xf32>, tensor<15x14x14x32xf32>) -> tensor<15x28x28x1xf32> %3 = "tf.AddV2"(%2, %1): (tensor<15x28x28x1xf32>, tensor<1xf32>) -> tensor<15x28x28x1xf32> func.return %2 : tensor<15x28x28x1xf32> -// CHECK: %[[CONV2DBACKPROPINPUT_0:.*]] = "tf.Conv2DBackpropInput"(%arg0, %arg1, %arg2) {dilations = [1, 1, 1, 1], padding = "SAME", strides = [1, 2, 2, 1]} : (tensor<4xi32>, tensor<3x3x1x32xf32>, tensor<15x14x14x32xf32>) -> tensor<15x28x28x1xf32> +// CHECK: %[[CONV2DBACKPROPINPUT_0:.*]] = "tf.Conv2DBackpropInput"(%arg0, %arg1, %arg2) <{dilations = [1, 1, 1, 1], padding = "SAME", strides = [1, 2, 2, 1]}> : (tensor<4xi32>, tensor<3x3x1x32xf32>, tensor<15x14x14x32xf32>) -> tensor<15x28x28x1xf32> // CHECK: return %[[CONV2DBACKPROPINPUT_0]] : tensor<15x28x28x1xf32> } diff --git a/tensorflow/compiler/mlir/lite/quantization/tensorflow/tests/tf_to_quant.mlir b/tensorflow/compiler/mlir/lite/quantization/tensorflow/tests/tf_to_quant.mlir index 83f07de73cc5c7..dd93ae25948fef 100644 --- a/tensorflow/compiler/mlir/lite/quantization/tensorflow/tests/tf_to_quant.mlir +++ b/tensorflow/compiler/mlir/lite/quantization/tensorflow/tests/tf_to_quant.mlir @@ -36,7 +36,7 @@ func.func @fakeQuantForActivationNoDuplication(tensor<8xf32>) -> (tensor<8x!quan %1 = "quantfork.qcast"(%0) : (tensor<8xf32>) -> tensor<8x!quant.uniform> func.return %1 : tensor<8x!quant.uniform> -// CHECK: %0 = "tf.FakeQuantWithMinMaxVars"(%arg0, %cst, %cst_0) {narrow_range = false, num_bits = 5 : i64} +// CHECK: %0 = "tf.FakeQuantWithMinMaxVars"(%arg0, %cst, %cst_0) <{narrow_range = false, num_bits = 5 : i64}> // CHECK: %1 = "quantfork.qcast"(%0) : (tensor<8xf32>) -> tensor<8x!quant.uniform> // CHECK: return %1 } @@ -51,7 +51,7 @@ func.func @fakeQuantFolded() -> (tensor<8xf32>) { %rst = "tf.FakeQuantWithMinMaxVars"(%in, %mini, %maxi) {num_bits = 5, narrow_range = false} : (tensor<8xf32>, tensor, tensor) -> tensor<8xf32> func.return %rst : tensor<8xf32> -// CHECK: %[[CONSTANT:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<8xf32>} +// CHECK: %[[CONSTANT:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<8xf32>}> // CHECK: %[[QUANTIZE:.*]] = "quantfork.qcast"(%[[CONSTANT]]) : (tensor<8xf32>) -> tensor<8x!quant.uniform> // CHECK: %[[DEQUANTIZE:.*]] = "quantfork.dcast"(%[[QUANTIZE]]) // CHECK: return %[[DEQUANTIZE]] : tensor<8xf32> @@ -79,7 +79,7 @@ func.func @fakeQuantWithConv2D(tensor<256x32x32x3xf32>) -> (tensor<256x8x7x16xf3 %rst = "tf.Conv2D"(%arg, %fq) {T = "tfdtype$DT_FLOAT", data_format = "NHWC", dilations = [1, 2, 3, 1], padding = "SAME", strides = [1, 4, 5, 1]} : (tensor<256x32x32x3xf32>, tensor<3x3x3x16xf32>) -> tensor<256x8x7x16xf32> func.return %rst : tensor<256x8x7x16xf32> -// CHECK: %[[CONSTANT0:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<3x3x3x16xf32>} +// CHECK: %[[CONSTANT0:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<3x3x3x16xf32>}> // CHECK: %[[QUANTIZE:.*]] = "quantfork.qcast"(%[[CONSTANT0]]) : (tensor<3x3x3x16xf32>) -> tensor<3x3x3x16x!quant.uniform> // CHECK: %[[DEQUANTIZE:.*]] = "quantfork.dcast"(%[[QUANTIZE]]) // CHECK: %[[CONV:.*]] = "tf.Conv2D"(%arg0, %[[DEQUANTIZE]]) @@ -98,7 +98,7 @@ func.func @perChannelFakeQuantWithConv2D(tensor<256x32x32x3xf32>) -> (tensor<256 %rst = "tf.Conv2D"(%arg, %fq) {T = "tfdtype$DT_FLOAT", data_format = "NHWC", dilations = [1, 2, 3, 1], padding = "SAME", strides = [1, 4, 5, 1]} : (tensor<256x32x32x3xf32>, tensor<3x3x3x16xf32>) -> tensor<256x8x7x16xf32> func.return %rst : tensor<256x8x7x16xf32> -// CHECK: %[[CONSTANT0:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<3x3x3x16xf32>} +// CHECK: %[[CONSTANT0:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<3x3x3x16xf32>}> // CHECK: %[[QUANTIZE:.*]] = "quantfork.qcast"(%[[CONSTANT0]]) : (tensor<3x3x3x16xf32>) -> tensor<3x3x3x16x!quant.uniform> @@ -119,7 +119,7 @@ func.func @fakeQuantWithDepthwiseConv2D(tensor<256x32x32x3xf32>) -> (tensor<256x %rst = "tf.DepthwiseConv2dNative"(%arg, %fq) {T = "tfdtype$DT_FLOAT", data_format = "NHWC", dilations = [1, 2, 3, 1], padding = "SAME", strides = [1, 4, 5, 1]} : (tensor<256x32x32x3xf32>, tensor<3x3x3x16xf32>) -> tensor<256x30x30x16xf32> func.return %rst : tensor<256x30x30x16xf32> -// CHECK: %[[CONSTANT0:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<3x3x3x16xf32>} +// CHECK: %[[CONSTANT0:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<3x3x3x16xf32>}> // CHECK: %[[QUANTIZE:.*]] = "quantfork.qcast"(%[[CONSTANT0]]) : (tensor<3x3x3x16xf32>) -> tensor<3x3x3x16x!quant.uniform> // CHECK: %[[DEQUANTIZE:.*]] = "quantfork.dcast"(%[[QUANTIZE]]) // CHECK: %[[CONV:.*]] = "tf.DepthwiseConv2dNative"(%arg0, %[[DEQUANTIZE]]) @@ -138,7 +138,7 @@ func.func @perChannelFakeQuantWithDepthwiseConv2D(tensor<256x32x32x3xf32>) -> (t %rst = "tf.DepthwiseConv2dNative"(%arg, %fq) {T = "tfdtype$DT_FLOAT", data_format = "NHWC", dilations = [1, 2, 3, 1], padding = "SAME", strides = [1, 4, 5, 1]} : (tensor<256x32x32x3xf32>, tensor<3x3x3x16xf32>) -> tensor<256x30x30x16xf32> func.return %rst : tensor<256x30x30x16xf32> -// CHECK: %[[CONSTANT0:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<3x3x3x16xf32>} +// CHECK: %[[CONSTANT0:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<3x3x3x16xf32>}> // CHECK: %[[QUANTIZE:.*]] = "quantfork.qcast"(%[[CONSTANT0]]) : (tensor<3x3x3x16xf32>) -> tensor<3x3x3x16x!quant.uniform> diff --git a/tensorflow/compiler/mlir/lite/quantization/tensorflow/tests/tf_to_quant_4bit.mlir b/tensorflow/compiler/mlir/lite/quantization/tensorflow/tests/tf_to_quant_4bit.mlir index 54c3de4bacd7a3..519226a7755b5b 100644 --- a/tensorflow/compiler/mlir/lite/quantization/tensorflow/tests/tf_to_quant_4bit.mlir +++ b/tensorflow/compiler/mlir/lite/quantization/tensorflow/tests/tf_to_quant_4bit.mlir @@ -36,7 +36,7 @@ func.func @fakeQuantForActivationNoDuplication(tensor<8xf32>) -> (tensor<8x!quan %1 = "quantfork.qcast"(%0) : (tensor<8xf32>) -> tensor<8x!quant.uniform> func.return %1 : tensor<8x!quant.uniform> -// CHECK: %0 = "tf.FakeQuantWithMinMaxVars"(%arg0, %cst, %cst_0) {narrow_range = false, num_bits = 3 : i64} +// CHECK: %0 = "tf.FakeQuantWithMinMaxVars"(%arg0, %cst, %cst_0) <{narrow_range = false, num_bits = 3 : i64}> // CHECK: %1 = "quantfork.qcast"(%0) : (tensor<8xf32>) -> tensor<8x!quant.uniform> // CHECK: return %1 } @@ -51,7 +51,7 @@ func.func @fakeQuantFolded() -> (tensor<8xf32>) { %rst = "tf.FakeQuantWithMinMaxVars"(%in, %mini, %maxi) {num_bits = 3, narrow_range = false} : (tensor<8xf32>, tensor, tensor) -> tensor<8xf32> func.return %rst : tensor<8xf32> -// CHECK: %[[CONSTANT:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<8xf32>} +// CHECK: %[[CONSTANT:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<8xf32>}> // CHECK: %[[QUANTIZE:.*]] = "quantfork.qcast"(%[[CONSTANT]]) : (tensor<8xf32>) -> tensor<8x!quant.uniform> // CHECK: %[[DEQUANTIZE:.*]] = "quantfork.dcast"(%[[QUANTIZE]]) // CHECK: return %[[DEQUANTIZE]] : tensor<8xf32> @@ -79,7 +79,7 @@ func.func @fakeQuantWithConv2D(tensor<256x32x32x3xf32>) -> (tensor<256x8x7x16xf3 %rst = "tf.Conv2D"(%arg, %fq) {T = "tfdtype$DT_FLOAT", data_format = "NHWC", dilations = [1, 2, 3, 1], padding = "SAME", strides = [1, 4, 5, 1]} : (tensor<256x32x32x3xf32>, tensor<3x3x3x16xf32>) -> tensor<256x8x7x16xf32> func.return %rst : tensor<256x8x7x16xf32> -// CHECK: %[[CONSTANT0:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<3x3x3x16xf32>} +// CHECK: %[[CONSTANT0:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<3x3x3x16xf32>}> // CHECK: %[[QUANTIZE:.*]] = "quantfork.qcast"(%[[CONSTANT0]]) : (tensor<3x3x3x16xf32>) -> tensor<3x3x3x16x!quant.uniform> // CHECK: %[[DEQUANTIZE:.*]] = "quantfork.dcast"(%[[QUANTIZE]]) // CHECK: %[[CONV:.*]] = "tf.Conv2D"(%arg0, %[[DEQUANTIZE]]) @@ -98,7 +98,7 @@ func.func @perChannelFakeQuantWithConv2D(tensor<256x32x32x3xf32>) -> (tensor<256 %rst = "tf.Conv2D"(%arg, %fq) {T = "tfdtype$DT_FLOAT", data_format = "NHWC", dilations = [1, 2, 3, 1], padding = "SAME", strides = [1, 4, 5, 1]} : (tensor<256x32x32x3xf32>, tensor<3x3x3x16xf32>) -> tensor<256x8x7x16xf32> func.return %rst : tensor<256x8x7x16xf32> -// CHECK: %[[CONSTANT0:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<3x3x3x16xf32>} +// CHECK: %[[CONSTANT0:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<3x3x3x16xf32>}> // CHECK: %[[QUANTIZE:.*]] = "quantfork.qcast"(%[[CONSTANT0]]) : (tensor<3x3x3x16xf32>) -> tensor<3x3x3x16x!quant.uniform> @@ -119,7 +119,7 @@ func.func @fakeQuantWithDepthwiseConv2D(tensor<256x32x32x3xf32>) -> (tensor<256x %rst = "tf.DepthwiseConv2dNative"(%arg, %fq) {T = "tfdtype$DT_FLOAT", data_format = "NHWC", dilations = [1, 2, 3, 1], padding = "SAME", strides = [1, 4, 5, 1]} : (tensor<256x32x32x3xf32>, tensor<3x3x3x16xf32>) -> tensor<256x30x30x16xf32> func.return %rst : tensor<256x30x30x16xf32> -// CHECK: %[[CONSTANT0:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<3x3x3x16xf32>} +// CHECK: %[[CONSTANT0:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<3x3x3x16xf32>}> // CHECK: %[[QUANTIZE:.*]] = "quantfork.qcast"(%[[CONSTANT0]]) : (tensor<3x3x3x16xf32>) -> tensor<3x3x3x16x!quant.uniform> // CHECK: %[[DEQUANTIZE:.*]] = "quantfork.dcast"(%[[QUANTIZE]]) // CHECK: %[[CONV:.*]] = "tf.DepthwiseConv2dNative"(%arg0, %[[DEQUANTIZE]]) @@ -138,7 +138,7 @@ func.func @perChannelFakeQuantWithDepthwiseConv2D(tensor<256x32x32x3xf32>) -> (t %rst = "tf.DepthwiseConv2dNative"(%arg, %fq) {T = "tfdtype$DT_FLOAT", data_format = "NHWC", dilations = [1, 2, 3, 1], padding = "SAME", strides = [1, 4, 5, 1]} : (tensor<256x32x32x3xf32>, tensor<3x3x3x16xf32>) -> tensor<256x30x30x16xf32> func.return %rst : tensor<256x30x30x16xf32> -// CHECK: %[[CONSTANT0:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<3x3x3x16xf32>} +// CHECK: %[[CONSTANT0:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<3x3x3x16xf32>}> // CHECK: %[[QUANTIZE:.*]] = "quantfork.qcast"(%[[CONSTANT0]]) : (tensor<3x3x3x16xf32>) -> tensor<3x3x3x16x!quant.uniform> diff --git a/tensorflow/compiler/mlir/lite/stablehlo/BUILD b/tensorflow/compiler/mlir/lite/stablehlo/BUILD index ce3de9b523e54d..2459f3d214d13a 100644 --- a/tensorflow/compiler/mlir/lite/stablehlo/BUILD +++ b/tensorflow/compiler/mlir/lite/stablehlo/BUILD @@ -193,6 +193,7 @@ cc_library( ":drop_savedmodel_semantics", ":fold_broadcast_pass", ":fuse_convolution_pass", + ":legalize_tf_xla_call_module_to_stablehlo_pass", ":optimize", ":rename_entrypoint_to_main", ":smuggle_disallowed_ops", diff --git a/tensorflow/compiler/mlir/lite/stablehlo/tests/call_xla_module_to_stablehlo.mlir b/tensorflow/compiler/mlir/lite/stablehlo/tests/call_xla_module_to_stablehlo.mlir new file mode 100644 index 00000000000000..795840247cab93 --- /dev/null +++ b/tensorflow/compiler/mlir/lite/stablehlo/tests/call_xla_module_to_stablehlo.mlir @@ -0,0 +1,26 @@ +//RUN: tf_tfl_translate --enable-stablehlo-conversion --input-mlir %s -o - | flatbuffer_translate --tflite-flatbuffer-to-mlir - -o - | FileCheck %s + + +module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 1660 : i32}} { + func.func @main(%arg0: tensor<2x3xi32>) -> tensor<2x3xi32> attributes {tf.entry_function = {control_outputs = "", inputs = "args_tf_0", outputs = "Identity"}} { + %0 = tf_executor.graph { + %outputs, %control = tf_executor.island wraps "tf.Identity"(%arg0) {device = ""} : (tensor<2x3xi32>) -> tensor<2x3xi32> + %outputs_0, %control_1 = tf_executor.island wraps "tf.XlaSharding"(%outputs) {_XlaSharding = "", device = "", sharding = "", unspecified_dims = []} : (tensor<2x3xi32>) -> tensor<2x3xi32> + %outputs_2, %control_3 = tf_executor.island wraps "tf.XlaCallModule"(%outputs_0) {Sout = [#tf_type.shape<2x3>], device = "", dim_args_spec = [], disabled_checks = [], function_list = [], has_token_input_output = false, module = "ML\EFR\01StableHLO_v0.9.0\00\01\17\05\01\03\01\03\05\03\07\07\09\0B\03]?\0B\01)\07\0F\0B+\0B\0F\0B\0B\0B3\0B\0B\0B\0B\0F\0B\0F\0B\13\0B\03\17\0F\13\0B\0B\0B\0F\13\0B\0B\0B\0B\01\05\0B\0F\03\07\17\17\07\02\D7\1F\11\03\05\05\0D\03\09\09\0B\0D\03\0F\03\05\11\05\0F\11\01\00\05\11\05\13\05\15\03\0B\15)\171\193\05;\1B=\05\17\05\19\05\1B\05\1D\1D\1F\01\05\1F\1D#%\05!\17'\A9\01\05#\03\03+\0D\03-/\1D%\1D'#\07\03\035\0D\0379\1D)\1D+\1D-\1D/\01\09\01\02\02)\05\09\0D\09\11\03\05\03\05\1B\04C\05\01\11\01\07\07\03\01\05\03\11\01\13\07\03\05\0B\03\05\1D\05\06!\03\05\05\01\01\07\04\01\03\03\06\03\01\05\01\00f\051\0F\0B\03!\1B\1D[;\05\1F\15\1D\15\1D%)9\13\15\19\11\0F\0B\11builtin\00vhlo\00module\00func_v1\00multiply_v1\00return_v1\00sym_name\00jax.uses_shape_polymorphism\00mhlo.num_partitions\00mhlo.num_replicas\00jit_jax_model\00arg_attrs\00function_type\00res_attrs\00sym_visibility\00x\00jit(jax_model)/jit(main)/mul\00experimental/users/ypang/lite/convert_ulm.py\00mhlo.sharding\00{replicated}\00jax.result_info\00\00main\00public\00", platforms = ["CPU"], version = 8 : i64} : (tensor<2x3xi32>) -> tensor<2x3xi32> + %control_4 = tf_executor.island(%control_3) wraps "tf.NoOp"() {device = ""} : () -> () + %outputs_5, %control_6 = tf_executor.island wraps "tf.PreventGradient"(%outputs_2) {device = "", message = "The jax2tf-converted function does not support gradients. Use `with_gradient` parameter to enable gradients"} : (tensor<2x3xi32>) -> tensor<2x3xi32> + %outputs_7, %control_8 = tf_executor.island wraps "tf.Identity"(%outputs_5) {device = ""} : (tensor<2x3xi32>) -> tensor<2x3xi32> + %outputs_9, %control_10 = tf_executor.island(%control_4) wraps "tf.Identity"(%outputs_7) {device = ""} : (tensor<2x3xi32>) -> tensor<2x3xi32> + tf_executor.fetch %outputs_9 : tensor<2x3xi32> + } + return %0 : tensor<2x3xi32> + } +} + +// CHECK: module attributes {tfl.description = "MLIR Converted.", tfl.metadata = {keep_stablehlo_constant = "true"}, tfl.schema_version = 3 : i32} { +// CHECK-NEXT: func.func @main(%arg0: tensor<2x3xi32>) -> tensor<2x3xi32> attributes {tf.entry_function = {inputs = "args_tf_0", outputs = "Identity"}} { +// CHECK-NEXT: %0 = stablehlo.custom_call @Sharding(%arg0) {mhlo.sharding = ""} : (tensor<2x3xi32>) -> tensor<2x3xi32> +// CHECK-NEXT: %1 = stablehlo.multiply %0, %0 : tensor<2x3xi32> +// CHECK-NEXT: return %1 : tensor<2x3xi32> +// CHECK-NEXT: } +// CHECK-NEXT: } \ No newline at end of file diff --git a/tensorflow/compiler/mlir/lite/stablehlo/tests/legalize_hlo.mlir b/tensorflow/compiler/mlir/lite/stablehlo/tests/legalize_hlo.mlir index 18cb7cebfd60b7..96a7152b550dcf 100644 --- a/tensorflow/compiler/mlir/lite/stablehlo/tests/legalize_hlo.mlir +++ b/tensorflow/compiler/mlir/lite/stablehlo/tests/legalize_hlo.mlir @@ -562,15 +562,15 @@ func.func @pow_dynamic(%arg0: tensor) -> tensor { // CHECK-LABEL: func @floordiv_broadcast_i32( // CHECK-SAME: %[[VAL_0:.*]]: tensor<2x3xi32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<3xi32>) -> tensor<2x3xi32> { -// CHECK: %[[VAL_2:.*]] = "tf.Const"() {value = dense<0> : tensor<2x3xi32>} : () -> tensor<2x3xi32> +// CHECK: %[[VAL_2:.*]] = "tf.Const"() <{value = dense<0> : tensor<2x3xi32>}> : () -> tensor<2x3xi32> // CHECK: %[[VAL_3:.*]] = "tf.Less"(%[[VAL_0]], %[[VAL_2]]) : (tensor<2x3xi32>, tensor<2x3xi32>) -> tensor<2x3xi1> -// CHECK: %[[VAL_4:.*]] = "tf.Const"() {value = dense<0> : tensor<3xi32>} : () -> tensor<3xi32> +// CHECK: %[[VAL_4:.*]] = "tf.Const"() <{value = dense<0> : tensor<3xi32>}> : () -> tensor<3xi32> // CHECK: %[[VAL_5:.*]] = "tf.Less"(%[[VAL_1]], %[[VAL_4]]) : (tensor<3xi32>, tensor<3xi32>) -> tensor<3xi1> -// CHECK: %[[VAL_6:.*]] = "tf.Equal"(%[[VAL_3]], %[[VAL_5]]) {incompatible_shape_error = true} : (tensor<2x3xi1>, tensor<3xi1>) -> tensor<2x3xi1> +// CHECK: %[[VAL_6:.*]] = "tf.Equal"(%[[VAL_3]], %[[VAL_5]]) <{incompatible_shape_error = true}> : (tensor<2x3xi1>, tensor<3xi1>) -> tensor<2x3xi1> // CHECK: %[[VAL_7:.*]] = "tf.Div"(%[[VAL_0]], %[[VAL_1]]) : (tensor<2x3xi32>, tensor<3xi32>) -> tensor<2x3xi32> // CHECK: %[[VAL_8:.*]] = "tf.Abs"(%[[VAL_0]]) : (tensor<2x3xi32>) -> tensor<2x3xi32> // CHECK: %[[VAL_9:.*]] = "tf.Abs"(%[[VAL_1]]) : (tensor<3xi32>) -> tensor<3xi32> -// CHECK: %[[VAL_10:.*]] = "tf.Const"() {value = dense<1> : tensor<3xi32>} : () -> tensor<3xi32> +// CHECK: %[[VAL_10:.*]] = "tf.Const"() <{value = dense<1> : tensor<3xi32>}> : () -> tensor<3xi32> // CHECK: %[[VAL_11:.*]] = "tf.Sub"(%[[VAL_9]], %[[VAL_10]]) : (tensor<3xi32>, tensor<3xi32>) -> tensor<3xi32> // CHECK: %[[VAL_12:.*]] = "tf.AddV2"(%[[VAL_8]], %[[VAL_11]]) : (tensor<2x3xi32>, tensor<3xi32>) -> tensor<2x3xi32> // CHECK: %[[VAL_13:.*]] = "tf.Neg"(%[[VAL_12]]) : (tensor<2x3xi32>) -> tensor<2x3xi32> @@ -601,15 +601,15 @@ func.func @floordiv_broadcast_i32(%arg0: tensor<2x3xi32>, %arg1: tensor<3xi32>) // CHECK-LABEL: func @floordiv_reverse_broadcast_i32( // CHECK-SAME: %[[VAL_0:.*]]: tensor<3xi32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<2x3xi32>) -> tensor<2x3xi32> { -// CHECK: %[[VAL_2:.*]] = "tf.Const"() {value = dense<0> : tensor<3xi32>} : () -> tensor<3xi32> +// CHECK: %[[VAL_2:.*]] = "tf.Const"() <{value = dense<0> : tensor<3xi32>}> : () -> tensor<3xi32> // CHECK: %[[VAL_3:.*]] = "tf.Less"(%[[VAL_0]], %[[VAL_2]]) : (tensor<3xi32>, tensor<3xi32>) -> tensor<3xi1> -// CHECK: %[[VAL_4:.*]] = "tf.Const"() {value = dense<0> : tensor<2x3xi32>} : () -> tensor<2x3xi32> +// CHECK: %[[VAL_4:.*]] = "tf.Const"() <{value = dense<0> : tensor<2x3xi32>}> : () -> tensor<2x3xi32> // CHECK: %[[VAL_5:.*]] = "tf.Less"(%[[VAL_1]], %[[VAL_4]]) : (tensor<2x3xi32>, tensor<2x3xi32>) -> tensor<2x3xi1> -// CHECK: %[[VAL_6:.*]] = "tf.Equal"(%[[VAL_3]], %[[VAL_5]]) {incompatible_shape_error = true} : (tensor<3xi1>, tensor<2x3xi1>) -> tensor<2x3xi1> +// CHECK: %[[VAL_6:.*]] = "tf.Equal"(%[[VAL_3]], %[[VAL_5]]) <{incompatible_shape_error = true}> : (tensor<3xi1>, tensor<2x3xi1>) -> tensor<2x3xi1> // CHECK: %[[VAL_7:.*]] = "tf.Div"(%[[VAL_0]], %[[VAL_1]]) : (tensor<3xi32>, tensor<2x3xi32>) -> tensor<2x3xi32> // CHECK: %[[VAL_8:.*]] = "tf.Abs"(%[[VAL_0]]) : (tensor<3xi32>) -> tensor<3xi32> // CHECK: %[[VAL_9:.*]] = "tf.Abs"(%[[VAL_1]]) : (tensor<2x3xi32>) -> tensor<2x3xi32> -// CHECK: %[[VAL_10:.*]] = "tf.Const"() {value = dense<1> : tensor<2x3xi32>} : () -> tensor<2x3xi32> +// CHECK: %[[VAL_10:.*]] = "tf.Const"() <{value = dense<1> : tensor<2x3xi32>}> : () -> tensor<2x3xi32> // CHECK: %[[VAL_11:.*]] = "tf.Sub"(%[[VAL_9]], %[[VAL_10]]) : (tensor<2x3xi32>, tensor<2x3xi32>) -> tensor<2x3xi32> // CHECK: %[[VAL_12:.*]] = "tf.AddV2"(%[[VAL_8]], %[[VAL_11]]) : (tensor<3xi32>, tensor<2x3xi32>) -> tensor<2x3xi32> // CHECK: %[[VAL_13:.*]] = "tf.Neg"(%[[VAL_12]]) : (tensor<2x3xi32>) -> tensor<2x3xi32> @@ -669,7 +669,7 @@ func.func @floordiv_f16_broadcast(%arg0: tensor<2x3xf16>, %arg1: tensor<3xf16>) // CHECK-LABEL: func @equal( // CHECK-SAME: %[[VAL_0:.*]]: tensor<2xi32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<2xi32>) -> tensor<2xi1> { -// CHECK: %[[VAL_2:.*]] = "tf.Equal"(%[[VAL_0]], %[[VAL_1]]) {incompatible_shape_error = true} : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi1> +// CHECK: %[[VAL_2:.*]] = "tf.Equal"(%[[VAL_0]], %[[VAL_1]]) <{incompatible_shape_error = true}> : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi1> // CHECK: return %[[VAL_2]] : tensor<2xi1> // CHECK: } func.func @equal(%arg0: tensor<2xi32>, %arg1: tensor<2xi32>) -> tensor<2xi1> { @@ -680,7 +680,7 @@ func.func @equal(%arg0: tensor<2xi32>, %arg1: tensor<2xi32>) -> tensor<2xi1> { // CHECK-LABEL: func @equal_dynamic( // CHECK-SAME: %[[VAL_0:.*]]: tensor, // CHECK-SAME: %[[VAL_1:.*]]: tensor<1xi32>) -> tensor { -// CHECK: %[[VAL_2:.*]] = "tf.Equal"(%[[VAL_0]], %[[VAL_1]]) {incompatible_shape_error = true} : (tensor, tensor<1xi32>) -> tensor +// CHECK: %[[VAL_2:.*]] = "tf.Equal"(%[[VAL_0]], %[[VAL_1]]) <{incompatible_shape_error = true}> : (tensor, tensor<1xi32>) -> tensor // CHECK: return %[[VAL_2]] : tensor // CHECK: } func.func @equal_dynamic(%arg0: tensor, %arg1: tensor<1xi32>) -> tensor { @@ -691,7 +691,7 @@ func.func @equal_dynamic(%arg0: tensor, %arg1: tensor<1xi32>) -> tensor, // CHECK-SAME: %[[VAL_1:.*]]: tensor<1x2xi32>) -> tensor<1x2xi1> { -// CHECK: %[[VAL_2:.*]] = "tf.Equal"(%[[VAL_0]], %[[VAL_1]]) {incompatible_shape_error = true} : (tensor<1x1xi32>, tensor<1x2xi32>) -> tensor<1x2xi1> +// CHECK: %[[VAL_2:.*]] = "tf.Equal"(%[[VAL_0]], %[[VAL_1]]) <{incompatible_shape_error = true}> : (tensor<1x1xi32>, tensor<1x2xi32>) -> tensor<1x2xi1> // CHECK: return %[[VAL_2]] : tensor<1x2xi1> // CHECK: } func.func @equal_broadcast(%arg0: tensor<1x1xi32>, %arg1: tensor<1x2xi32>) -> tensor<1x2xi1> { @@ -703,7 +703,7 @@ func.func @equal_broadcast(%arg0: tensor<1x1xi32>, %arg1: tensor<1x2xi32>) -> te // CHECK-LABEL: func @equal_broadcast_chlo( // CHECK-SAME: %[[VAL_0:.*]]: tensor<1xi32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<1x2xi32>) -> tensor<1x2xi1> { -// CHECK: %[[VAL_2:.*]] = "tf.Equal"(%[[VAL_0]], %[[VAL_1]]) {incompatible_shape_error = true} : (tensor<1xi32>, tensor<1x2xi32>) -> tensor<1x2xi1> +// CHECK: %[[VAL_2:.*]] = "tf.Equal"(%[[VAL_0]], %[[VAL_1]]) <{incompatible_shape_error = true}> : (tensor<1xi32>, tensor<1x2xi32>) -> tensor<1x2xi1> // CHECK: return %[[VAL_2]] : tensor<1x2xi1> // CHECK: } func.func @equal_broadcast_chlo(%arg0: tensor<1xi32>, %arg1: tensor<1x2xi32>) -> tensor<1x2xi1> { @@ -714,7 +714,7 @@ func.func @equal_broadcast_chlo(%arg0: tensor<1xi32>, %arg1: tensor<1x2xi32>) -> // CHECK-LABEL: func @equal_broadcast_no_incompatible_shapes_error( // CHECK-SAME: %[[VAL_0:.*]]: tensor<2xi32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<1x2xi32>) -> tensor<1x2xi1> { -// CHECK: %[[VAL_2:.*]] = "tf.Equal"(%[[VAL_0]], %[[VAL_1]]) {incompatible_shape_error = true} : (tensor<2xi32>, tensor<1x2xi32>) -> tensor<1x2xi1> +// CHECK: %[[VAL_2:.*]] = "tf.Equal"(%[[VAL_0]], %[[VAL_1]]) <{incompatible_shape_error = true}> : (tensor<2xi32>, tensor<1x2xi32>) -> tensor<1x2xi1> // CHECK: return %[[VAL_2]] : tensor<1x2xi1> // CHECK: } func.func @equal_broadcast_no_incompatible_shapes_error(%arg0: tensor<2xi32>, %arg1: tensor<1x2xi32>) -> tensor<1x2xi1> { @@ -725,7 +725,7 @@ func.func @equal_broadcast_no_incompatible_shapes_error(%arg0: tensor<2xi32>, %a // CHECK-LABEL: func @equal_incompatible_shape_broadcastable( // CHECK-SAME: %[[VAL_0:.*]]: tensor, // CHECK-SAME: %[[VAL_1:.*]]: tensor<1xi32>) -> tensor { -// CHECK: %[[VAL_2:.*]] = "tf.Equal"(%[[VAL_0]], %[[VAL_1]]) {incompatible_shape_error = true} : (tensor, tensor<1xi32>) -> tensor +// CHECK: %[[VAL_2:.*]] = "tf.Equal"(%[[VAL_0]], %[[VAL_1]]) <{incompatible_shape_error = true}> : (tensor, tensor<1xi32>) -> tensor // CHECK: return %[[VAL_2]] : tensor // CHECK: } func.func @equal_incompatible_shape_broadcastable(%arg0: tensor, %arg1: tensor<1xi32>) -> tensor { @@ -743,7 +743,7 @@ func.func @equal_unsupported_compare_type(%arg0: tensor<1xf32>, %arg1: tensor<1x // CHECK-LABEL: func @notequal( // CHECK-SAME: %[[VAL_0:.*]]: tensor<2xi32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<2xi32>) -> tensor<2xi1> { -// CHECK: %[[VAL_2:.*]] = "tf.NotEqual"(%[[VAL_0]], %[[VAL_1]]) {incompatible_shape_error = true} : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi1> +// CHECK: %[[VAL_2:.*]] = "tf.NotEqual"(%[[VAL_0]], %[[VAL_1]]) <{incompatible_shape_error = true}> : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi1> // CHECK: return %[[VAL_2]] : tensor<2xi1> // CHECK: } func.func @notequal(%arg0: tensor<2xi32>, %arg1: tensor<2xi32>) -> tensor<2xi1> { @@ -754,7 +754,7 @@ func.func @notequal(%arg0: tensor<2xi32>, %arg1: tensor<2xi32>) -> tensor<2xi1> // CHECK-LABEL: func @notequal_broadcast( // CHECK-SAME: %[[VAL_0:.*]]: tensor<1x1xi32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<1x2xi32>) -> tensor<1x2xi1> { -// CHECK: %[[VAL_2:.*]] = "tf.NotEqual"(%[[VAL_0]], %[[VAL_1]]) {incompatible_shape_error = true} : (tensor<1x1xi32>, tensor<1x2xi32>) -> tensor<1x2xi1> +// CHECK: %[[VAL_2:.*]] = "tf.NotEqual"(%[[VAL_0]], %[[VAL_1]]) <{incompatible_shape_error = true}> : (tensor<1x1xi32>, tensor<1x2xi32>) -> tensor<1x2xi1> // CHECK: return %[[VAL_2]] : tensor<1x2xi1> // CHECK: } func.func @notequal_broadcast(%arg0: tensor<1x1xi32>, %arg1: tensor<1x2xi32>) -> tensor<1x2xi1> { @@ -766,7 +766,7 @@ func.func @notequal_broadcast(%arg0: tensor<1x1xi32>, %arg1: tensor<1x2xi32>) -> // CHECK-LABEL: func @notequal_broadcast_chlo( // CHECK-SAME: %[[VAL_0:.*]]: tensor<1xi32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<1x2xi32>) -> tensor<1x2xi1> { -// CHECK: %[[VAL_2:.*]] = "tf.NotEqual"(%[[VAL_0]], %[[VAL_1]]) {incompatible_shape_error = true} : (tensor<1xi32>, tensor<1x2xi32>) -> tensor<1x2xi1> +// CHECK: %[[VAL_2:.*]] = "tf.NotEqual"(%[[VAL_0]], %[[VAL_1]]) <{incompatible_shape_error = true}> : (tensor<1xi32>, tensor<1x2xi32>) -> tensor<1x2xi1> // CHECK: return %[[VAL_2]] : tensor<1x2xi1> // CHECK: } func.func @notequal_broadcast_chlo(%arg0: tensor<1xi32>, %arg1: tensor<1x2xi32>) -> tensor<1x2xi1> { @@ -777,7 +777,7 @@ func.func @notequal_broadcast_chlo(%arg0: tensor<1xi32>, %arg1: tensor<1x2xi32>) // CHECK-LABEL: func @notequal_broadcast_no_incompatible_shapes_error( // CHECK-SAME: %[[VAL_0:.*]]: tensor<2xi32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<1x2xi32>) -> tensor<1x2xi1> { -// CHECK: %[[VAL_2:.*]] = "tf.NotEqual"(%[[VAL_0]], %[[VAL_1]]) {incompatible_shape_error = true} : (tensor<2xi32>, tensor<1x2xi32>) -> tensor<1x2xi1> +// CHECK: %[[VAL_2:.*]] = "tf.NotEqual"(%[[VAL_0]], %[[VAL_1]]) <{incompatible_shape_error = true}> : (tensor<2xi32>, tensor<1x2xi32>) -> tensor<1x2xi1> // CHECK: return %[[VAL_2]] : tensor<1x2xi1> // CHECK: } func.func @notequal_broadcast_no_incompatible_shapes_error(%arg0: tensor<2xi32>, %arg1: tensor<1x2xi32>) -> tensor<1x2xi1> { @@ -788,7 +788,7 @@ func.func @notequal_broadcast_no_incompatible_shapes_error(%arg0: tensor<2xi32>, // CHECK-LABEL: func @notequal_incompatible_shape_broadcastable( // CHECK-SAME: %[[VAL_0:.*]]: tensor, // CHECK-SAME: %[[VAL_1:.*]]: tensor<1xi32>) -> tensor { -// CHECK: %[[VAL_2:.*]] = "tf.NotEqual"(%[[VAL_0]], %[[VAL_1]]) {incompatible_shape_error = true} : (tensor, tensor<1xi32>) -> tensor +// CHECK: %[[VAL_2:.*]] = "tf.NotEqual"(%[[VAL_0]], %[[VAL_1]]) <{incompatible_shape_error = true}> : (tensor, tensor<1xi32>) -> tensor // CHECK: return %[[VAL_2]] : tensor // CHECK: } func.func @notequal_incompatible_shape_broadcastable(%arg0: tensor, %arg1: tensor<1xi32>) -> tensor { @@ -942,7 +942,7 @@ func.func @broadcast_less_equal_chlo(%arg0: tensor<1xi32>, %arg1: tensor<1x2xi32 // CHECK-LABEL: func @concat_v2( // CHECK-SAME: %[[VAL_0:.*]]: tensor<3x3xf32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<3x3xf32>) -> tensor<6x3xf32> { -// CHECK: %[[VAL_2:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor +// CHECK: %[[VAL_2:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK: %[[VAL_3:.*]] = "tf.ConcatV2"(%[[VAL_0]], %[[VAL_1]], %[[VAL_2]]) : (tensor<3x3xf32>, tensor<3x3xf32>, tensor) -> tensor<6x3xf32> // CHECK: return %[[VAL_3]] : tensor<6x3xf32> // CHECK: } @@ -954,7 +954,7 @@ func.func @concat_v2(%arg0: tensor<3x3xf32>, %arg1: tensor<3x3xf32>) -> tensor<6 // CHECK-LABEL: func @concat_v2_1d_axis( // CHECK-SAME: %[[VAL_0:.*]]: tensor<3x3xf32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<3x3xf32>) -> tensor<3x6xf32> { -// CHECK: %[[VAL_2:.*]] = "tf.Const"() {value = dense<1> : tensor} : () -> tensor +// CHECK: %[[VAL_2:.*]] = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor // CHECK: %[[VAL_3:.*]] = "tf.ConcatV2"(%[[VAL_0]], %[[VAL_1]], %[[VAL_2]]) : (tensor<3x3xf32>, tensor<3x3xf32>, tensor) -> tensor<3x6xf32> // CHECK: return %[[VAL_3]] : tensor<3x6xf32> // CHECK: } @@ -964,7 +964,7 @@ func.func @concat_v2_1d_axis(%arg0: tensor<3x3xf32>, %arg1: tensor<3x3xf32>) -> } // CHECK-LABEL: func @const() -> tensor<2xi32> { -// CHECK: %[[VAL_0:.*]] = "tf.Const"() {value = dense<0> : tensor<2xi32>} : () -> tensor<2xi32> +// CHECK: %[[VAL_0:.*]] = "tf.Const"() <{value = dense<0> : tensor<2xi32>}> : () -> tensor<2xi32> // CHECK: return %[[VAL_0]] : tensor<2xi32> // CHECK: } func.func @const() -> tensor<2xi32> { @@ -974,7 +974,7 @@ func.func @const() -> tensor<2xi32> { // CHECK-LABEL: func @relu( // CHECK-SAME: %[[VAL_0:.*]]: tensor<1xi32>) -> tensor<1xi32> { -// CHECK: %[[VAL_1:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor +// CHECK: %[[VAL_1:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK: %[[VAL_2:.*]] = "tf.Maximum"(%[[VAL_0]], %[[VAL_1]]) : (tensor<1xi32>, tensor) -> tensor<1xi32> // CHECK: return %[[VAL_2]] : tensor<1xi32> // CHECK: } @@ -986,7 +986,7 @@ func.func @relu(%arg0: tensor<1xi32>) -> tensor<1xi32> { // CHECK-LABEL: func @relu_unranked( // CHECK-SAME: %[[VAL_0:.*]]: tensor) -> tensor { -// CHECK: %[[VAL_1:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor +// CHECK: %[[VAL_1:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK: %[[VAL_2:.*]] = "tf.Maximum"(%[[VAL_0]], %[[VAL_1]]) : (tensor, tensor) -> tensor // CHECK: return %[[VAL_2]] : tensor // CHECK: } @@ -998,8 +998,8 @@ func.func @relu_unranked(%arg0: tensor) -> tensor { // CHECK-LABEL: func @relu6( // CHECK-SAME: %[[VAL_0:.*]]: tensor<1xi32>) -> tensor<1xi32> { -// CHECK-DAG: %[[VAL_1:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor -// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() {value = dense<6> : tensor} : () -> tensor +// CHECK-DAG: %[[VAL_1:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor +// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() <{value = dense<6> : tensor}> : () -> tensor // CHECK: %[[VAL_3:.*]] = "tf.Minimum"(%[[VAL_0]], %[[VAL_2]]) : (tensor<1xi32>, tensor) -> tensor<1xi32> // CHECK: %[[VAL_4:.*]] = "tf.Maximum"(%[[VAL_3]], %[[VAL_1]]) : (tensor<1xi32>, tensor) -> tensor<1xi32> // CHECK: return %[[VAL_4]] : tensor<1xi32> @@ -1014,8 +1014,8 @@ func.func @relu6(%arg0: tensor<1xi32>) -> tensor<1xi32> { // CHECK-LABEL: func @relu6_unranked( // CHECK-SAME: %[[VAL_0:.*]]: tensor) -> tensor { -// CHECK-DAG: %[[VAL_1:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor -// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() {value = dense<6> : tensor} : () -> tensor +// CHECK-DAG: %[[VAL_1:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor +// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() <{value = dense<6> : tensor}> : () -> tensor // CHECK: %[[VAL_3:.*]] = "tf.Minimum"(%[[VAL_0]], %[[VAL_2]]) : (tensor, tensor) -> tensor // CHECK: %[[VAL_4:.*]] = "tf.Maximum"(%[[VAL_3]], %[[VAL_1]]) : (tensor, tensor) -> tensor // CHECK: return %[[VAL_4]] : tensor @@ -1031,9 +1031,9 @@ func.func @relu6_unranked(%arg0: tensor) -> tensor { // CHECK-LABEL: func @relu_grad( // CHECK-SAME: %[[VAL_0:.*]]: tensor<4x8xf32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor) -> tensor<4x8xf32> { -// CHECK: %[[VAL_2:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor} : () -> tensor +// CHECK: %[[VAL_2:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor}> : () -> tensor // CHECK: %[[VAL_3:.*]] = "tf.Greater"(%[[VAL_1]], %[[VAL_2]]) : (tensor, tensor) -> tensor -// CHECK: %[[VAL_4:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<4x8xf32>} : () -> tensor<4x8xf32> +// CHECK: %[[VAL_4:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<4x8xf32>}> : () -> tensor<4x8xf32> // CHECK: %[[VAL_5:.*]] = "tf.Select"(%[[VAL_3]], %[[VAL_0]], %[[VAL_4]]) : (tensor, tensor<4x8xf32>, tensor<4x8xf32>) -> tensor<4x8xf32> // CHECK: return %[[VAL_5]] : tensor<4x8xf32> // CHECK: } @@ -1133,9 +1133,9 @@ func.func @selectv2_broadcasted_condition(%arg0: tensor<1x1xi1>, %arg1: tensor<1 // CHECK-LABEL: func @transpose_2d( // CHECK-SAME: %[[VAL_0:.*]]: tensor<2x3xf32>) -> tensor<3x2xf32> { -// CHECK-DAG: %[[VAL_1:.*]] = "tf.Const"() {value = dense<[1, 0]> : tensor<2xi64>} : () -> tensor<2xi64> -// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() {value = dense<[1, 0]> : tensor<2xi64>} : () -> tensor<2xi64> -// CHECK-DAG: %[[VAL_3:.*]] = "tf.Const"() {value = dense<[1, 0]> : tensor<2xi64>} : () -> tensor<2xi64> +// CHECK-DAG: %[[VAL_1:.*]] = "tf.Const"() <{value = dense<[1, 0]> : tensor<2xi64>}> : () -> tensor<2xi64> +// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() <{value = dense<[1, 0]> : tensor<2xi64>}> : () -> tensor<2xi64> +// CHECK-DAG: %[[VAL_3:.*]] = "tf.Const"() <{value = dense<[1, 0]> : tensor<2xi64>}> : () -> tensor<2xi64> // CHECK: %[[VAL_4:.*]] = "tf.Transpose"(%[[VAL_0]], %[[VAL_3]]) : (tensor<2x3xf32>, tensor<2xi64>) -> tensor<3x2xf32> // CHECK: return %[[VAL_4]] : tensor<3x2xf32> // CHECK: } @@ -1148,9 +1148,9 @@ func.func @transpose_2d(%arg0: tensor<2x3xf32>) -> tensor<3x2xf32> { // CHECK-LABEL: func @transpose_3d_int32( // CHECK-SAME: %[[VAL_0:.*]]: tensor<1x2x3xf32>) -> tensor<3x2x1xf32> { -// CHECK-DAG: %[[VAL_1:.*]] = "tf.Const"() {value = dense<[2, 1, 0]> : tensor<3xi32>} : () -> tensor<3xi32> -// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() {value = dense<[2, 1, 0]> : tensor<3xi64>} : () -> tensor<3xi64> -// CHECK-DAG: %[[VAL_3:.*]] = "tf.Const"() {value = dense<[2, 1, 0]> : tensor<3xi64>} : () -> tensor<3xi64> +// CHECK-DAG: %[[VAL_1:.*]] = "tf.Const"() <{value = dense<[2, 1, 0]> : tensor<3xi32>}> : () -> tensor<3xi32> +// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() <{value = dense<[2, 1, 0]> : tensor<3xi64>}> : () -> tensor<3xi64> +// CHECK-DAG: %[[VAL_3:.*]] = "tf.Const"() <{value = dense<[2, 1, 0]> : tensor<3xi64>}> : () -> tensor<3xi64> // CHECK: %[[VAL_4:.*]] = "tf.Transpose"(%[[VAL_0]], %[[VAL_3]]) : (tensor<1x2x3xf32>, tensor<3xi64>) -> tensor<3x2x1xf32> // CHECK: return %[[VAL_4]] : tensor<3x2x1xf32> // CHECK: } @@ -1163,9 +1163,9 @@ func.func @transpose_3d_int32(%arg0: tensor<1x2x3xf32>) -> tensor<3x2x1xf32> { // CHECK-LABEL: func @transpose_3d( // CHECK-SAME: %[[VAL_0:.*]]: tensor<1x2x3xf32>) -> tensor<3x2x1xf32> { -// CHECK-DAG: %[[VAL_1:.*]] = "tf.Const"() {value = dense<[2, 1, 0]> : tensor<3xi64>} : () -> tensor<3xi64> -// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() {value = dense<[2, 1, 0]> : tensor<3xi64>} : () -> tensor<3xi64> -// CHECK-DAG: %[[VAL_3:.*]] = "tf.Const"() {value = dense<[2, 1, 0]> : tensor<3xi64>} : () -> tensor<3xi64> +// CHECK-DAG: %[[VAL_1:.*]] = "tf.Const"() <{value = dense<[2, 1, 0]> : tensor<3xi64>}> : () -> tensor<3xi64> +// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() <{value = dense<[2, 1, 0]> : tensor<3xi64>}> : () -> tensor<3xi64> +// CHECK-DAG: %[[VAL_3:.*]] = "tf.Const"() <{value = dense<[2, 1, 0]> : tensor<3xi64>}> : () -> tensor<3xi64> // CHECK: %[[VAL_4:.*]] = "tf.Transpose"(%[[VAL_0]], %[[VAL_3]]) : (tensor<1x2x3xf32>, tensor<3xi64>) -> tensor<3x2x1xf32> // CHECK: return %[[VAL_4]] : tensor<3x2x1xf32> // CHECK: } @@ -1178,9 +1178,9 @@ func.func @transpose_3d(%arg0: tensor<1x2x3xf32>) -> tensor<3x2x1xf32> { // CHECK-LABEL: func @transpose_dynamic_2d( // CHECK-SAME: %[[VAL_0:.*]]: tensor) -> tensor<4x?xf32> { -// CHECK-DAG: %[[VAL_1:.*]] = "tf.Const"() {value = dense<[1, 0]> : tensor<2xi64>} : () -> tensor<2xi64> -// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() {value = dense<[1, 0]> : tensor<2xi64>} : () -> tensor<2xi64> -// CHECK-DAG: %[[VAL_3:.*]] = "tf.Const"() {value = dense<[1, 0]> : tensor<2xi64>} : () -> tensor<2xi64> +// CHECK-DAG: %[[VAL_1:.*]] = "tf.Const"() <{value = dense<[1, 0]> : tensor<2xi64>}> : () -> tensor<2xi64> +// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() <{value = dense<[1, 0]> : tensor<2xi64>}> : () -> tensor<2xi64> +// CHECK-DAG: %[[VAL_3:.*]] = "tf.Const"() <{value = dense<[1, 0]> : tensor<2xi64>}> : () -> tensor<2xi64> // CHECK: %[[VAL_4:.*]] = "tf.Transpose"(%[[VAL_0]], %[[VAL_3]]) : (tensor, tensor<2xi64>) -> tensor<4x?xf32> // CHECK: return %[[VAL_4]] : tensor<4x?xf32> // CHECK: } @@ -1193,9 +1193,9 @@ func.func @transpose_dynamic_2d(%arg0: tensor) -> tensor<4x?xf32> { // CHECK-LABEL: func @transpose_unranked_2d( // CHECK-SAME: %[[VAL_0:.*]]: tensor<*xf32>) -> tensor<*xf32> { -// CHECK-DAG: %[[VAL_1:.*]] = "tf.Const"() {value = dense<[1, 0]> : tensor<2xi64>} : () -> tensor<2xi64> -// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() {value = dense<[1, 0]> : tensor<2xi64>} : () -> tensor<2xi64> -// CHECK-DAG: %[[VAL_3:.*]] = "tf.Const"() {value = dense<[1, 0]> : tensor<2xi64>} : () -> tensor<2xi64> +// CHECK-DAG: %[[VAL_1:.*]] = "tf.Const"() <{value = dense<[1, 0]> : tensor<2xi64>}> : () -> tensor<2xi64> +// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() <{value = dense<[1, 0]> : tensor<2xi64>}> : () -> tensor<2xi64> +// CHECK-DAG: %[[VAL_3:.*]] = "tf.Const"() <{value = dense<[1, 0]> : tensor<2xi64>}> : () -> tensor<2xi64> // CHECK: %[[VAL_4:.*]] = "tf.Transpose"(%[[VAL_0]], %[[VAL_3]]) : (tensor<*xf32>, tensor<2xi64>) -> tensor<*xf32> // CHECK: return %[[VAL_4]] : tensor<*xf32> // CHECK: } @@ -1488,9 +1488,9 @@ func.func @neg_unranked(%arg0: tensor<*xf32>) -> tensor<*xf32> { // CHECK-LABEL: func @sigmoid( // CHECK-SAME: %[[VAL_0:.*]]: tensor<2xf32>) -> tensor<2xf32> { -// CHECK-DAG: %[[VAL_1:.*]] = "tf.Const"() {value = dense<5.000000e-01> : tensor} : () -> tensor -// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() {value = dense<2> : tensor<1xi64>} : () -> tensor<1xi64> -// CHECK-DAG: %[[VAL_3:.*]] = "tf.Const"() {value = dense<5.000000e-01> : tensor<2xf32>} : () -> tensor<2xf32> +// CHECK-DAG: %[[VAL_1:.*]] = "tf.Const"() <{value = dense<5.000000e-01> : tensor}> : () -> tensor +// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() <{value = dense<2> : tensor<1xi64>}> : () -> tensor<1xi64> +// CHECK-DAG: %[[VAL_3:.*]] = "tf.Const"() <{value = dense<5.000000e-01> : tensor<2xf32>}> : () -> tensor<2xf32> // CHECK: %[[VAL_4:.*]] = "tf.Mul"(%[[VAL_0]], %[[VAL_3]]) : (tensor<2xf32>, tensor<2xf32>) -> tensor<2xf32> // CHECK: %[[VAL_5:.*]] = "tf.Tanh"(%[[VAL_4]]) : (tensor<2xf32>) -> tensor<2xf32> // CHECK: %[[VAL_6:.*]] = "tf.Mul"(%[[VAL_5]], %[[VAL_3]]) : (tensor<2xf32>, tensor<2xf32>) -> tensor<2xf32> @@ -1671,10 +1671,10 @@ func.func @bitcast_same_widths(%arg0: tensor<2xf32>) -> tensor<2xi32> { // CHECK-LABEL: func @sign( // CHECK-SAME: %[[VAL_0:.*]]: tensor<1x2x3x4xf32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<1x2x3x4xf32>) -> tensor<1x2x3x4xf32> { -// CHECK: %[[VAL_2:.*]] = "tf.NotEqual"(%[[VAL_0]], %[[VAL_1]]) {incompatible_shape_error = true} : (tensor<1x2x3x4xf32>, tensor<1x2x3x4xf32>) -> tensor<1x2x3x4xi1> -// CHECK: %[[VAL_3:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<1x2x3x4xf32>} : () -> tensor<1x2x3x4xf32> -// CHECK: %[[VAL_4:.*]] = "tf.NotEqual"(%[[VAL_0]], %[[VAL_1]]) {incompatible_shape_error = true} : (tensor<1x2x3x4xf32>, tensor<1x2x3x4xf32>) -> tensor<1x2x3x4xi1> -// CHECK: %[[VAL_5:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<1x2x3x4xf32>} : () -> tensor<1x2x3x4xf32> +// CHECK: %[[VAL_2:.*]] = "tf.NotEqual"(%[[VAL_0]], %[[VAL_1]]) <{incompatible_shape_error = true}> : (tensor<1x2x3x4xf32>, tensor<1x2x3x4xf32>) -> tensor<1x2x3x4xi1> +// CHECK: %[[VAL_3:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<1x2x3x4xf32>}> : () -> tensor<1x2x3x4xf32> +// CHECK: %[[VAL_4:.*]] = "tf.NotEqual"(%[[VAL_0]], %[[VAL_1]]) <{incompatible_shape_error = true}> : (tensor<1x2x3x4xf32>, tensor<1x2x3x4xf32>) -> tensor<1x2x3x4xi1> +// CHECK: %[[VAL_5:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<1x2x3x4xf32>}> : () -> tensor<1x2x3x4xf32> // CHECK: %[[VAL_6:.*]] = "tf.Sign"(%[[VAL_0]]) : (tensor<1x2x3x4xf32>) -> tensor<1x2x3x4xf32> // CHECK: %[[VAL_7:.*]] = "tf.Select"(%[[VAL_4]], %[[VAL_5]], %[[VAL_6]]) : (tensor<1x2x3x4xi1>, tensor<1x2x3x4xf32>, tensor<1x2x3x4xf32>) -> tensor<1x2x3x4xf32> // CHECK: %[[VAL_8:.*]] = "tf.Select"(%[[VAL_2]], %[[VAL_3]], %[[VAL_7]]) : (tensor<1x2x3x4xi1>, tensor<1x2x3x4xf32>, tensor<1x2x3x4xf32>) -> tensor<1x2x3x4xf32> @@ -1693,7 +1693,7 @@ func.func @sign(%arg0: tensor<1x2x3x4xf32>, %arg1: tensor<1x2x3x4xf32>) -> tenso // CHECK-LABEL: func @size_rank_one_i32( // CHECK-SAME: %[[VAL_0:.*]]: tensor) -> tensor { -// CHECK: %[[VAL_1:.*]] = "tf.Const"() {value = dense<1> : tensor} : () -> tensor +// CHECK: %[[VAL_1:.*]] = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor // CHECK: return %[[VAL_1]] : tensor // CHECK: } func.func @size_rank_one_i32(%arg0: tensor) -> tensor { @@ -1703,7 +1703,7 @@ func.func @size_rank_one_i32(%arg0: tensor) -> tensor { // CHECK-LABEL: func @size_rank_one_i64( // CHECK-SAME: %[[VAL_0:.*]]: tensor) -> tensor { -// CHECK: %[[VAL_1:.*]] = "tf.Const"() {value = dense<1> : tensor} : () -> tensor +// CHECK: %[[VAL_1:.*]] = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor // CHECK: return %[[VAL_1]] : tensor // CHECK: } func.func @size_rank_one_i64(%arg0: tensor) -> tensor { @@ -1724,7 +1724,7 @@ func.func @complex(%arg0: tensor<3xf32>, %arg1: tensor<3xf32>) -> tensor<3xcompl // CHECK-LABEL: func @convert_i32_f32( // CHECK-SAME: %[[VAL_0:.*]]: tensor<2xi32>) -> tensor<2xf32> { -// CHECK: %[[VAL_1:.*]] = "tf.Cast"(%[[VAL_0]]) {Truncate = false} : (tensor<2xi32>) -> tensor<2xf32> +// CHECK: %[[VAL_1:.*]] = "tf.Cast"(%[[VAL_0]]) <{Truncate = false}> : (tensor<2xi32>) -> tensor<2xf32> // CHECK: return %[[VAL_1]] : tensor<2xf32> // CHECK: } func.func @convert_i32_f32(%arg0: tensor<2xi32>) -> tensor<2xf32> { @@ -1734,11 +1734,11 @@ func.func @convert_i32_f32(%arg0: tensor<2xi32>) -> tensor<2xf32> { // CHECK-LABEL: func @convert_slice( // CHECK-SAME: %[[VAL_0:.*]]: tensor<1x4672xf32>) -> tensor<1x519xf32> { -// CHECK-DAG: %[[VAL_1:.*]] = "tf.Const"() {value = dense<[0, 4153]> : tensor<2xi64>} : () -> tensor<2xi64> -// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() {value = dense<[1, 4672]> : tensor<2xi64>} : () -> tensor<2xi64> -// CHECK-DAG: %[[VAL_3:.*]] = "tf.Const"() {value = dense<1> : tensor<2xi64>} : () -> tensor<2xi64> +// CHECK-DAG: %[[VAL_1:.*]] = "tf.Const"() <{value = dense<[0, 4153]> : tensor<2xi64>}> : () -> tensor<2xi64> +// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() <{value = dense<[1, 4672]> : tensor<2xi64>}> : () -> tensor<2xi64> +// CHECK-DAG: %[[VAL_3:.*]] = "tf.Const"() <{value = dense<1> : tensor<2xi64>}> : () -> tensor<2xi64> // CHECK: %[[VAL_4:.*]] = "tf.StridedSlice"(%[[VAL_0]], %[[VAL_1]], %[[VAL_2]], %[[VAL_3]]) -// CHECK-SAME: {begin_mask = 0 : i64, ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 0 : i64} +// CHECK-SAME: <{begin_mask = 0 : i64, ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 0 : i64}> // CHECK-SAME: (tensor<1x4672xf32>, tensor<2xi64>, tensor<2xi64>, tensor<2xi64>) -> tensor<1x519xf32> // CHECK: return %[[VAL_4]] : tensor<1x519xf32> // CHECK: } @@ -1786,7 +1786,7 @@ func.func @round_nearest_even(%arg0: tensor<2xf32>) -> tensor<2xf32> { // CHECK-SAME: %[[VAL_1:.*]]: tensor<256xf32>) -> tensor<1xf32> { // CHECK: %[[VAL_2:.*]] = arith.constant dense<[256, 1]> : tensor<2xi64> // CHECK: %[[VAL_3:.*]] = "tf.Reshape"(%[[VAL_1]], %[[VAL_2]]) : (tensor<256xf32>, tensor<2xi64>) -> tensor<256x1xf32> -// CHECK: %[[VAL_4:.*]] = "tf.BatchMatMulV3"(%[[VAL_0]], %[[VAL_3]]) {adj_x = false, adj_y = false} : (tensor<1x256xf32>, tensor<256x1xf32>) -> tensor<1x1xf32> +// CHECK: %[[VAL_4:.*]] = "tf.BatchMatMulV3"(%[[VAL_0]], %[[VAL_3]]) <{adj_x = false, adj_y = false}> : (tensor<1x256xf32>, tensor<256x1xf32>) -> tensor<1x1xf32> // CHECK: %[[VAL_5:.*]] = arith.constant dense<1> : tensor<1xi64> // CHECK: %[[VAL_6:.*]] = "tf.Reshape"(%[[VAL_4]], %[[VAL_5]]) : (tensor<1x1xf32>, tensor<1xi64>) -> tensor<1xf32> // CHECK: return %[[VAL_6]] : tensor<1xf32> @@ -1803,7 +1803,7 @@ func.func @convert_dot_2d_1d(%arg0: tensor<1x256xf32>, %arg1: tensor<256xf32>) - // CHECK: %[[VAL_3:.*]] = "tf.Reshape"(%[[VAL_0]], %[[VAL_2]]) : (tensor<256xf32>, tensor<2xi64>) -> tensor<1x256xf32> // CHECK: %[[VAL_4:.*]] = arith.constant dense<[256, 1]> : tensor<2xi64> // CHECK: %[[VAL_5:.*]] = "tf.Reshape"(%[[VAL_1]], %[[VAL_4]]) : (tensor<256xf32>, tensor<2xi64>) -> tensor<256x1xf32> -// CHECK: %[[VAL_6:.*]] = "tf.BatchMatMulV3"(%[[VAL_3]], %[[VAL_5]]) {adj_x = false, adj_y = false} : (tensor<1x256xf32>, tensor<256x1xf32>) -> tensor<1x1xf32> +// CHECK: %[[VAL_6:.*]] = "tf.BatchMatMulV3"(%[[VAL_3]], %[[VAL_5]]) <{adj_x = false, adj_y = false}> : (tensor<1x256xf32>, tensor<256x1xf32>) -> tensor<1x1xf32> // CHECK: %[[VAL_7:.*]] = arith.constant dense<> : tensor<0xi64> // CHECK: %[[VAL_8:.*]] = "tf.Reshape"(%[[VAL_6]], %[[VAL_7]]) : (tensor<1x1xf32>, tensor<0xi64>) -> tensor // CHECK: return %[[VAL_8]] : tensor @@ -1816,7 +1816,7 @@ func.func @convert_dot_1d_1d(%arg0: tensor<256xf32>, %arg1: tensor<256xf32>) -> // CHECK-LABEL: func @convert_dot_2d_2d( // CHECK-SAME: %[[VAL_0:.*]]: tensor<1x256xf32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<256x1xf32>) -> tensor<1x1xf32> { -// CHECK: %[[VAL_2:.*]] = "tf.BatchMatMulV3"(%[[VAL_0]], %[[VAL_1]]) {adj_x = false, adj_y = false} : (tensor<1x256xf32>, tensor<256x1xf32>) -> tensor<1x1xf32> +// CHECK: %[[VAL_2:.*]] = "tf.BatchMatMulV3"(%[[VAL_0]], %[[VAL_1]]) <{adj_x = false, adj_y = false}> : (tensor<1x256xf32>, tensor<256x1xf32>) -> tensor<1x1xf32> // CHECK: return %[[VAL_2]] : tensor<1x1xf32> // CHECK: } func.func @convert_dot_2d_2d(%arg0: tensor<1x256xf32>, %arg1: tensor<256x1xf32>) -> tensor<1x1xf32> { @@ -1861,9 +1861,9 @@ func.func @dynamic_broadcast_in_dim_tf_style(%arg0: tensor, %arg1 // CHECK-LABEL: func @dynamic_broadcast_in_dim_general_case_expand_back_dims( // CHECK-SAME: %[[ARG_0:.*]]: tensor, // CHECK-SAME: %[[ARG_1:.*]]: tensor<4xi32>) -> tensor { -// CHECK %[[CST_0:.*]] = "tf.Const"() {value = dense<2> : tensor} : () -> tensor +// CHECK %[[CST_0:.*]] = "tf.Const"() <{value = dense<2> : tensor}> : () -> tensor // CHECK %[[VAL_0:.*]] = "tf.ExpandDims"(%[[ARG_0]], %[[CST_0]]) : (tensor, tensor) -> tensor -// CHECK %[[CST_1:.*]] = "tf.Const"() {value = dense<3> : tensor} : () -> tensor +// CHECK %[[CST_1:.*]] = "tf.Const"() <{value = dense<3> : tensor}> : () -> tensor // CHECK %[[VAL_1:.*]] = "tf.ExpandDims"(%[[VAL_0]], %[[CST_1]]) : (tensor, tensor) -> tensor // CHECK %[[VAL_2:.*]] = "tf.BroadcastTo"(%[[VAL_1]], %[[ARG_1]]) : (tensor, tensor<4xi32>) -> tensor // CHECK return %[[VAL_2]] : tensor @@ -1875,7 +1875,7 @@ func.func @dynamic_broadcast_in_dim_general_case_expand_back_dims(%arg0: tensor< // CHECK-LABEL: func @dynamic_broadcast_in_dim_general_case_expand_middle_dim( // CHECK-SAME: %[[ARG_0:.*]]: tensor, // CHECK-SAME: %[[ARG_1:.*]]: tensor<4xi32>) -> tensor { -// CHECK %[[CST_0:.*]] = "tf.Const"() {value = dense<2> : tensor} : () -> tensor +// CHECK %[[CST_0:.*]] = "tf.Const"() <{value = dense<2> : tensor}> : () -> tensor // CHECK %[[VAL_0:.*]] = "tf.ExpandDims"(%[[ARG_0]], %[[CST_0]]) : (tensor, tensor) -> tensor // CHECK %[[VAL_1:.*]] = "tf.BroadcastTo"(%[[VAL_0]], %[[ARG_1]]) : (tensor, tensor<4xi32>) -> tensor // CHECK return %[[VAL_1]] : tensor @@ -1887,15 +1887,15 @@ func.func @dynamic_broadcast_in_dim_general_case_expand_middle_dim(%arg0: tensor // CHECK-LABEL: func @convert_dot_general( // CHECK-SAME: %[[VAL_0:.*]]: tensor<3x2x6x5x1xf32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<3x2x4x6xf32>) -> tensor<3x5x1x4xf32> { -// CHECK: %[[VAL_2:.*]] = "tf.Const"() {value = dense<[0, 3, 4, 1, 2]> : tensor<5xi64>} : () -> tensor<5xi64> +// CHECK: %[[VAL_2:.*]] = "tf.Const"() <{value = dense<[0, 3, 4, 1, 2]> : tensor<5xi64>}> : () -> tensor<5xi64> // CHECK: %[[VAL_3:.*]] = "tf.Transpose"(%[[VAL_0]], %[[VAL_2]]) : (tensor<3x2x6x5x1xf32>, tensor<5xi64>) -> tensor<3x5x1x2x6xf32> -// CHECK: %[[VAL_4:.*]] = "tf.Const"() {value = dense<[0, 1, 3, 2]> : tensor<4xi64>} : () -> tensor<4xi64> +// CHECK: %[[VAL_4:.*]] = "tf.Const"() <{value = dense<[0, 1, 3, 2]> : tensor<4xi64>}> : () -> tensor<4xi64> // CHECK: %[[VAL_5:.*]] = "tf.Transpose"(%[[VAL_1]], %[[VAL_4]]) : (tensor<3x2x4x6xf32>, tensor<4xi64>) -> tensor<3x2x6x4xf32> // CHECK: %[[VAL_6:.*]] = arith.constant dense<[3, 5, 12]> : tensor<3xi64> // CHECK: %[[VAL_7:.*]] = "tf.Reshape"(%[[VAL_3]], %[[VAL_6]]) : (tensor<3x5x1x2x6xf32>, tensor<3xi64>) -> tensor<3x5x12xf32> // CHECK: %[[VAL_8:.*]] = arith.constant dense<[3, 12, 4]> : tensor<3xi64> // CHECK: %[[VAL_9:.*]] = "tf.Reshape"(%[[VAL_5]], %[[VAL_8]]) : (tensor<3x2x6x4xf32>, tensor<3xi64>) -> tensor<3x12x4xf32> -// CHECK: %[[VAL_10:.*]] = "tf.BatchMatMulV3"(%[[VAL_7]], %[[VAL_9]]) {adj_x = false, adj_y = false} : (tensor<3x5x12xf32>, tensor<3x12x4xf32>) -> tensor<3x5x4xf32> +// CHECK: %[[VAL_10:.*]] = "tf.BatchMatMulV3"(%[[VAL_7]], %[[VAL_9]]) <{adj_x = false, adj_y = false}> : (tensor<3x5x12xf32>, tensor<3x12x4xf32>) -> tensor<3x5x4xf32> // CHECK: %[[VAL_11:.*]] = arith.constant dense<[3, 5, 1, 4]> : tensor<4xi64> // CHECK: %[[VAL_12:.*]] = "tf.Reshape"(%[[VAL_10]], %[[VAL_11]]) : (tensor<3x5x4xf32>, tensor<4xi64>) -> tensor<3x5x1x4xf32> // CHECK: return %[[VAL_12]] : tensor<3x5x1x4xf32> @@ -1929,7 +1929,7 @@ func.func @quantized_dot_general_not_converted(%arg0: tensor<1x1x512xf32>, %arg1 // CHECK-SAME: %[[VAL_1:.*]]: tensor<1024x1024xf32>) -> tensor<1x1x1024xf32> { // CHECK: %[[VAL_2:.*]] = arith.constant dense<[1, 1024]> : tensor<2xi64> // CHECK: %[[VAL_3:.*]] = "tf.Reshape"(%[[VAL_0]], %[[VAL_2]]) : {{.*}} -> tensor<1x1024xf32> -// CHECK: %[[VAL_4:.*]] = "tf.BatchMatMulV3"(%[[VAL_3]], %[[VAL_1]]) {adj_x = false, adj_y = false} : {{.*}} -> tensor<1x1024xf32> +// CHECK: %[[VAL_4:.*]] = "tf.BatchMatMulV3"(%[[VAL_3]], %[[VAL_1]]) <{adj_x = false, adj_y = false}> : {{.*}} -> tensor<1x1024xf32> // CHECK: %[[VAL_5:.*]] = arith.constant dense<[1, 1, 1024]> : tensor<3xi64> // CHECK: %[[VAL_6:.*]] = "tf.Reshape"(%[[VAL_4]], %[[VAL_5]]) : {{.*}} -> tensor<1x1x1024xf32> // CHECK: return %[[VAL_6]] : tensor<1x1x1024xf32> @@ -1952,7 +1952,7 @@ func.func @convert_dot_general_repeated(%arg0: tensor<1x1x1024xf32>, %arg1: tens // CHECK-SAME: %[[VAL_1:.*]]: tensor<256x8xi8>) -> tensor<8xi32> { // CHECK: %[[VAL_2:.*]] = arith.constant dense<[1, 256]> : tensor<2xi64> // CHECK: %[[VAL_3:.*]] = "tf.Reshape"(%[[VAL_0]], %[[VAL_2]]) : (tensor<256xi8>, tensor<2xi64>) -> tensor<1x256xi8> -// CHECK: %[[VAL_4:.*]] = "tf.BatchMatMulV3"(%[[VAL_3]], %[[VAL_1]]) {adj_x = false, adj_y = false} : (tensor<1x256xi8>, tensor<256x8xi8>) -> tensor<1x8xi32> +// CHECK: %[[VAL_4:.*]] = "tf.BatchMatMulV3"(%[[VAL_3]], %[[VAL_1]]) <{adj_x = false, adj_y = false}> : (tensor<1x256xi8>, tensor<256x8xi8>) -> tensor<1x8xi32> // CHECK: %[[VAL_5:.*]] = arith.constant dense<8> : tensor<1xi64> // CHECK: %[[VAL_6:.*]] = "tf.Reshape"(%[[VAL_4]], %[[VAL_5]]) : (tensor<1x8xi32>, tensor<1xi64>) -> tensor<8xi32> // CHECK: return %[[VAL_6]] : tensor<8xi32> @@ -1970,26 +1970,26 @@ func.func @convert_dot_general_int8(%arg0: tensor<256xi8>, %arg1: tensor<256x8xi // CHECK-LABEL: func @convert_dot_general_dynamic_rhs_out_dim( // CHECK-SAME: %arg0: tensor<4x4x256xf32>, // CHECK-SAME: %arg1: tensor<4x?x256xf32>) -> tensor<4x4x?xf32> { -// CHECK-DAG: %cst = "tf.Const"() {value = dense<[0, 2, 1]> : tensor<3xi64>} : () -> tensor<3xi64> +// CHECK-DAG: %cst = "tf.Const"() <{value = dense<[0, 2, 1]> : tensor<3xi64>}> : () -> tensor<3xi64> // CHECK: %0 = "tf.Transpose"(%arg1, %cst) : (tensor<4x?x256xf32>, tensor<3xi64>) -> tensor<4x256x?xf32> // CHECK: %1 = "tf.Shape"(%arg1) : (tensor<4x?x256xf32>) -> tensor<3xi32> -// CHECK-DAG: %cst_0 = "tf.Const"() {value = dense<[-1, 0, -1]> : tensor<3xi32>} : () -> tensor<3xi32> -// CHECK-DAG: %cst_1 = "tf.Const"() {value = dense<[-1, -1, 0]> : tensor<3xi32>} : () -> tensor<3xi32> -// CHECK-DAG: %cst_2 = "tf.Const"() {value = dense<1> : tensor} : () -> tensor +// CHECK-DAG: %cst_0 = "tf.Const"() <{value = dense<[-1, 0, -1]> : tensor<3xi32>}> : () -> tensor<3xi32> +// CHECK-DAG: %cst_1 = "tf.Const"() <{value = dense<[-1, -1, 0]> : tensor<3xi32>}> : () -> tensor<3xi32> +// CHECK-DAG: %cst_2 = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor // CHECK: %2 = "tf.UnsortedSegmentProd"(%1, %cst_0, %cst_2) : (tensor<3xi32>, tensor<3xi32>, tensor) -> tensor<1xi32> // CHECK: %3 = "tf.UnsortedSegmentProd"(%1, %cst_1, %cst_2) : (tensor<3xi32>, tensor<3xi32>, tensor) -> tensor<1xi32> -// CHECK-DAG: %cst_3 = "tf.Const"() {value = dense<4> : tensor<1xi32>} : () -> tensor<1xi32> -// CHECK-DAG: %cst_4 = "tf.Const"() {value = dense<0> : tensor} : () -> tensor +// CHECK-DAG: %cst_3 = "tf.Const"() <{value = dense<4> : tensor<1xi32>}> : () -> tensor<1xi32> +// CHECK-DAG: %cst_4 = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK: %4 = "tf.Concat"(%cst_4, %cst_3, %3, %2) : (tensor, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<3xi32> // CHECK: %5 = "tf.Reshape"(%0, %4) : (tensor<4x256x?xf32>, tensor<3xi32>) -> tensor<4x256x?xf32> -// CHECK: %6 = "tf.BatchMatMulV3"(%arg0, %5) {adj_x = false, adj_y = false} : (tensor<4x4x256xf32>, tensor<4x256x?xf32>) -> tensor<4x4x?xf32> +// CHECK: %6 = "tf.BatchMatMulV3"(%arg0, %5) <{adj_x = false, adj_y = false}> : (tensor<4x4x256xf32>, tensor<4x256x?xf32>) -> tensor<4x4x?xf32> // CHECK: %7 = "tf.Shape"(%arg0) : (tensor<4x4x256xf32>) -> tensor<3xi32> // CHECK: %8 = "tf.Shape"(%arg1) : (tensor<4x?x256xf32>) -> tensor<3xi32> -// CHECK-DAG: %cst_5 = "tf.Const"() {value = dense<[0, 1]> : tensor<2xi64>} : () -> tensor<2xi64> -// CHECK: %9 = "tf.Gather"(%7, %cst_5) {validate_indices = true} : (tensor<3xi32>, tensor<2xi64>) -> tensor<2xi32> -// CHECK-DAG: %cst_6 = "tf.Const"() {value = dense<1> : tensor<1xi64>} : () -> tensor<1xi64> -// CHECK: %10 = "tf.Gather"(%8, %cst_6) {validate_indices = true} : (tensor<3xi32>, tensor<1xi64>) -> tensor<1xi32> -// CHECK-DAG: %cst_7 = "tf.Const"() {value = dense<0> : tensor} : () -> tensor +// CHECK-DAG: %cst_5 = "tf.Const"() <{value = dense<[0, 1]> : tensor<2xi64>}> : () -> tensor<2xi64> +// CHECK: %9 = "tf.Gather"(%7, %cst_5) <{validate_indices = true}> : (tensor<3xi32>, tensor<2xi64>) -> tensor<2xi32> +// CHECK-DAG: %cst_6 = "tf.Const"() <{value = dense<1> : tensor<1xi64>}> : () -> tensor<1xi64> +// CHECK: %10 = "tf.Gather"(%8, %cst_6) <{validate_indices = true}> : (tensor<3xi32>, tensor<1xi64>) -> tensor<1xi32> +// CHECK-DAG: %cst_7 = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK: %11 = "tf.Concat"(%cst_7, %9, %10) : (tensor, tensor<2xi32>, tensor<1xi32>) -> tensor<3xi32> // CHECK: %12 = "tf.Reshape"(%6, %11) : (tensor<4x4x?xf32>, tensor<3xi32>) -> tensor<4x4x?xf32> // CHECK: return %12 : tensor<4x4x?xf32> @@ -2008,38 +2008,38 @@ func.return %0 : tensor<4x4x?xf32> // CHECK-LABEL: func @convert_dot_general_dynamic_batch_dim( // CHECK-SAME: %arg0: tensor<2x?x2x3xf32>, // CHECK-SAME: %arg1: tensor<2x?x4x3xf32>) -> tensor<2x?x2x4xf32> { -// CHECK-DAG: %cst = "tf.Const"() {value = dense<[0, 1, 3, 2]> : tensor<4xi64>} : () -> tensor<4xi64> +// CHECK-DAG: %cst = "tf.Const"() <{value = dense<[0, 1, 3, 2]> : tensor<4xi64>}> : () -> tensor<4xi64> // CHECK: %0 = "tf.Transpose"(%arg1, %cst) : (tensor<2x?x4x3xf32>, tensor<4xi64>) -> tensor<2x?x3x4xf32> // CHECK: %1 = "tf.Shape"(%arg0) : (tensor<2x?x2x3xf32>) -> tensor<4xi32> -// CHECK-DAG: %cst_0 = "tf.Const"() {value = dense<[-1, -1, 0, -1]> : tensor<4xi32>} : () -> tensor<4xi32> -// CHECK-DAG: %cst_1 = "tf.Const"() {value = dense<[-1, -1, -1, 0]> : tensor<4xi32>} : () -> tensor<4xi32> -// CHECK-DAG: %cst_2 = "tf.Const"() {value = dense<1> : tensor} : () -> tensor +// CHECK-DAG: %cst_0 = "tf.Const"() <{value = dense<[-1, -1, 0, -1]> : tensor<4xi32>}> : () -> tensor<4xi32> +// CHECK-DAG: %cst_1 = "tf.Const"() <{value = dense<[-1, -1, -1, 0]> : tensor<4xi32>}> : () -> tensor<4xi32> +// CHECK-DAG: %cst_2 = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor // CHECK: %2 = "tf.UnsortedSegmentProd"(%1, %cst_0, %cst_2) : (tensor<4xi32>, tensor<4xi32>, tensor) -> tensor<1xi32> // CHECK: %3 = "tf.UnsortedSegmentProd"(%1, %cst_1, %cst_2) : (tensor<4xi32>, tensor<4xi32>, tensor) -> tensor<1xi32> -// CHECK-DAG: %cst_3 = "tf.Const"() {value = dense<[0, 1]> : tensor<2xi64>} : () -> tensor<2xi64> -// CHECK: %4 = "tf.Gather"(%1, %cst_3) {validate_indices = true} : (tensor<4xi32>, tensor<2xi64>) -> tensor<2xi32> -// CHECK-DAG: %cst_4 = "tf.Const"() {value = dense<0> : tensor} : () -> tensor +// CHECK-DAG: %cst_3 = "tf.Const"() <{value = dense<[0, 1]> : tensor<2xi64>}> : () -> tensor<2xi64> +// CHECK: %4 = "tf.Gather"(%1, %cst_3) <{validate_indices = true}> : (tensor<4xi32>, tensor<2xi64>) -> tensor<2xi32> +// CHECK-DAG: %cst_4 = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK: %5 = "tf.Concat"(%cst_4, %4, %2, %3) : (tensor, tensor<2xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<4xi32> // CHECK: %6 = "tf.Reshape"(%arg0, %5) : (tensor<2x?x2x3xf32>, tensor<4xi32>) -> tensor<2x?x2x3xf32> // CHECK: %7 = "tf.Shape"(%arg1) : (tensor<2x?x4x3xf32>) -> tensor<4xi32> -// CHECK-DAG: %cst_5 = "tf.Const"() {value = dense<[-1, -1, 0, -1]> : tensor<4xi32>} : () -> tensor<4xi32> -// CHECK-DAG: %cst_6 = "tf.Const"() {value = dense<[-1, -1, -1, 0]> : tensor<4xi32>} : () -> tensor<4xi32> -// CHECK-DAG: %cst_7 = "tf.Const"() {value = dense<1> : tensor} : () -> tensor +// CHECK-DAG: %cst_5 = "tf.Const"() <{value = dense<[-1, -1, 0, -1]> : tensor<4xi32>}> : () -> tensor<4xi32> +// CHECK-DAG: %cst_6 = "tf.Const"() <{value = dense<[-1, -1, -1, 0]> : tensor<4xi32>}> : () -> tensor<4xi32> +// CHECK-DAG: %cst_7 = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor // CHECK: %8 = "tf.UnsortedSegmentProd"(%7, %cst_5, %cst_7) : (tensor<4xi32>, tensor<4xi32>, tensor) -> tensor<1xi32> // CHECK: %9 = "tf.UnsortedSegmentProd"(%7, %cst_6, %cst_7) : (tensor<4xi32>, tensor<4xi32>, tensor) -> tensor<1xi32> -// CHECK-DAG: %cst_8 = "tf.Const"() {value = dense<[0, 1]> : tensor<2xi64>} : () -> tensor<2xi64> -// CHECK: %10 = "tf.Gather"(%7, %cst_8) {validate_indices = true} : (tensor<4xi32>, tensor<2xi64>) -> tensor<2xi32> -// CHECK-DAG: %cst_9 = "tf.Const"() {value = dense<0> : tensor} : () -> tensor +// CHECK-DAG: %cst_8 = "tf.Const"() <{value = dense<[0, 1]> : tensor<2xi64>}> : () -> tensor<2xi64> +// CHECK: %10 = "tf.Gather"(%7, %cst_8) <{validate_indices = true}> : (tensor<4xi32>, tensor<2xi64>) -> tensor<2xi32> +// CHECK-DAG: %cst_9 = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK: %11 = "tf.Concat"(%cst_9, %10, %9, %8) : (tensor, tensor<2xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<4xi32> // CHECK: %12 = "tf.Reshape"(%0, %11) : (tensor<2x?x3x4xf32>, tensor<4xi32>) -> tensor<2x?x3x4xf32> -// CHECK: %13 = "tf.BatchMatMulV3"(%6, %12) {adj_x = false, adj_y = false} : (tensor<2x?x2x3xf32>, tensor<2x?x3x4xf32>) -> tensor<2x?x2x4xf32> +// CHECK: %13 = "tf.BatchMatMulV3"(%6, %12) <{adj_x = false, adj_y = false}> : (tensor<2x?x2x3xf32>, tensor<2x?x3x4xf32>) -> tensor<2x?x2x4xf32> // CHECK: %14 = "tf.Shape"(%arg0) : (tensor<2x?x2x3xf32>) -> tensor<4xi32> // CHECK: %15 = "tf.Shape"(%arg1) : (tensor<2x?x4x3xf32>) -> tensor<4xi32> -// CHECK-DAG: %cst_10 = "tf.Const"() {value = dense<[0, 1, 2]> : tensor<3xi64>} : () -> tensor<3xi64> -// CHECK: %16 = "tf.Gather"(%14, %cst_10) {validate_indices = true} : (tensor<4xi32>, tensor<3xi64>) -> tensor<3xi32> -// CHECK: %cst_11 = "tf.Const"() {value = dense<2> : tensor<1xi64>} : () -> tensor<1xi64> -// CHECK: %17 = "tf.Gather"(%15, %cst_11) {validate_indices = true} : (tensor<4xi32>, tensor<1xi64>) -> tensor<1xi32> -// CHECK-DAG: %cst_12 = "tf.Const"() {value = dense<0> : tensor} : () -> tensor +// CHECK-DAG: %cst_10 = "tf.Const"() <{value = dense<[0, 1, 2]> : tensor<3xi64>}> : () -> tensor<3xi64> +// CHECK: %16 = "tf.Gather"(%14, %cst_10) <{validate_indices = true}> : (tensor<4xi32>, tensor<3xi64>) -> tensor<3xi32> +// CHECK: %cst_11 = "tf.Const"() <{value = dense<2> : tensor<1xi64>}> : () -> tensor<1xi64> +// CHECK: %17 = "tf.Gather"(%15, %cst_11) <{validate_indices = true}> : (tensor<4xi32>, tensor<1xi64>) -> tensor<1xi32> +// CHECK-DAG: %cst_12 = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK: %18 = "tf.Concat"(%cst_12, %16, %17) : (tensor, tensor<3xi32>, tensor<1xi32>) -> tensor<4xi32> // CHECK: %19 = "tf.Reshape"(%13, %18) : (tensor<2x?x2x4xf32>, tensor<4xi32>) -> tensor<2x?x2x4xf32> // CHECK: return %19 : tensor<2x?x2x4xf32> @@ -2058,36 +2058,36 @@ func.return %0 : tensor<2x?x2x4xf32> // CHECK-LABEL: func @convert_dot_general_dynamic_lhs_rhs_out_dims( // CHECK-SAME: %arg0: tensor<2x2x?x3xf32>, // CHECK-SAME: %arg1: tensor<2x4x?x3xf32>) -> tensor<2x2x?x4x?xf32> { -// CHECK-DAG: %cst = "tf.Const"() {value = dense<[0, 3, 1, 2]> : tensor<4xi64>} : () -> tensor<4xi64> +// CHECK-DAG: %cst = "tf.Const"() <{value = dense<[0, 3, 1, 2]> : tensor<4xi64>}> : () -> tensor<4xi64> // CHECK: %0 = "tf.Transpose"(%arg1, %cst) : (tensor<2x4x?x3xf32>, tensor<4xi64>) -> tensor<2x3x4x?xf32> // CHECK: %1 = "tf.Shape"(%arg0) : (tensor<2x2x?x3xf32>) -> tensor<4xi32> -// CHECK-DAG: %cst_0 = "tf.Const"() {value = dense<[-1, 0, 0, -1]> : tensor<4xi32>} : () -> tensor<4xi32> -// CHECK-DAG: %cst_1 = "tf.Const"() {value = dense<[-1, -1, -1, 0]> : tensor<4xi32>} : () -> tensor<4xi32> -// CHECK-DAG: %cst_2 = "tf.Const"() {value = dense<1> : tensor} : () -> tensor +// CHECK-DAG: %cst_0 = "tf.Const"() <{value = dense<[-1, 0, 0, -1]> : tensor<4xi32>}> : () -> tensor<4xi32> +// CHECK-DAG: %cst_1 = "tf.Const"() <{value = dense<[-1, -1, -1, 0]> : tensor<4xi32>}> : () -> tensor<4xi32> +// CHECK-DAG: %cst_2 = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor // CHECK: %2 = "tf.UnsortedSegmentProd"(%1, %cst_0, %cst_2) : (tensor<4xi32>, tensor<4xi32>, tensor) -> tensor<1xi32> // CHECK: %3 = "tf.UnsortedSegmentProd"(%1, %cst_1, %cst_2) : (tensor<4xi32>, tensor<4xi32>, tensor) -> tensor<1xi32> -// CHECK-DAG: %cst_3 = "tf.Const"() {value = dense<2> : tensor<1xi32>} : () -> tensor<1xi32> -// CHECK-DAG: %cst_4 = "tf.Const"() {value = dense<0> : tensor} : () -> tensor +// CHECK-DAG: %cst_3 = "tf.Const"() <{value = dense<2> : tensor<1xi32>}> : () -> tensor<1xi32> +// CHECK-DAG: %cst_4 = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK: %4 = "tf.Concat"(%cst_4, %cst_3, %2, %3) : (tensor, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<3xi32> // CHECK: %5 = "tf.Reshape"(%arg0, %4) : (tensor<2x2x?x3xf32>, tensor<3xi32>) -> tensor<2x?x3xf32> // CHECK: %6 = "tf.Shape"(%arg1) : (tensor<2x4x?x3xf32>) -> tensor<4xi32> -// CHECK-DAG: %cst_5 = "tf.Const"() {value = dense<[-1, 0, 0, -1]> : tensor<4xi32>} : () -> tensor<4xi32> -// CHECK-DAG: %cst_6 = "tf.Const"() {value = dense<[-1, -1, -1, 0]> : tensor<4xi32>} : () -> tensor<4xi32> -// CHECK-DAG: %cst_7 = "tf.Const"() {value = dense<1> : tensor} : () -> tensor +// CHECK-DAG: %cst_5 = "tf.Const"() <{value = dense<[-1, 0, 0, -1]> : tensor<4xi32>}> : () -> tensor<4xi32> +// CHECK-DAG: %cst_6 = "tf.Const"() <{value = dense<[-1, -1, -1, 0]> : tensor<4xi32>}> : () -> tensor<4xi32> +// CHECK-DAG: %cst_7 = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor // CHECK: %7 = "tf.UnsortedSegmentProd"(%6, %cst_5, %cst_7) : (tensor<4xi32>, tensor<4xi32>, tensor) -> tensor<1xi32> // CHECK: %8 = "tf.UnsortedSegmentProd"(%6, %cst_6, %cst_7) : (tensor<4xi32>, tensor<4xi32>, tensor) -> tensor<1xi32> -// CHECK-DAG: %cst_8 = "tf.Const"() {value = dense<2> : tensor<1xi32>} : () -> tensor<1xi32> -// CHECK-DAG: %cst_9 = "tf.Const"() {value = dense<0> : tensor} : () -> tensor +// CHECK-DAG: %cst_8 = "tf.Const"() <{value = dense<2> : tensor<1xi32>}> : () -> tensor<1xi32> +// CHECK-DAG: %cst_9 = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK: %9 = "tf.Concat"(%cst_9, %cst_8, %8, %7) : (tensor, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<3xi32> // CHECK: %10 = "tf.Reshape"(%0, %9) : (tensor<2x3x4x?xf32>, tensor<3xi32>) -> tensor<2x3x?xf32> -// CHECK: %11 = "tf.BatchMatMulV3"(%5, %10) {adj_x = false, adj_y = false} : (tensor<2x?x3xf32>, tensor<2x3x?xf32>) -> tensor<2x?x?xf32> +// CHECK: %11 = "tf.BatchMatMulV3"(%5, %10) <{adj_x = false, adj_y = false}> : (tensor<2x?x3xf32>, tensor<2x3x?xf32>) -> tensor<2x?x?xf32> // CHECK: %12 = "tf.Shape"(%arg0) : (tensor<2x2x?x3xf32>) -> tensor<4xi32> // CHECK: %13 = "tf.Shape"(%arg1) : (tensor<2x4x?x3xf32>) -> tensor<4xi32> -// CHECK-DAG: %cst_10 = "tf.Const"() {value = dense<[0, 1, 2]> : tensor<3xi64>} : () -> tensor<3xi64> -// CHECK: %14 = "tf.Gather"(%12, %cst_10) {validate_indices = true} : (tensor<4xi32>, tensor<3xi64>) -> tensor<3xi32> -// CHECK-DAG: %cst_11 = "tf.Const"() {value = dense<[1, 2]> : tensor<2xi64>} : () -> tensor<2xi64> -// CHECK: %15 = "tf.Gather"(%13, %cst_11) {validate_indices = true} : (tensor<4xi32>, tensor<2xi64>) -> tensor<2xi32> -// CHECK-DAG: %cst_12 = "tf.Const"() {value = dense<0> : tensor} : () -> tensor +// CHECK-DAG: %cst_10 = "tf.Const"() <{value = dense<[0, 1, 2]> : tensor<3xi64>}> : () -> tensor<3xi64> +// CHECK: %14 = "tf.Gather"(%12, %cst_10) <{validate_indices = true}> : (tensor<4xi32>, tensor<3xi64>) -> tensor<3xi32> +// CHECK-DAG: %cst_11 = "tf.Const"() <{value = dense<[1, 2]> : tensor<2xi64>}> : () -> tensor<2xi64> +// CHECK: %15 = "tf.Gather"(%13, %cst_11) <{validate_indices = true}> : (tensor<4xi32>, tensor<2xi64>) -> tensor<2xi32> +// CHECK-DAG: %cst_12 = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK: %16 = "tf.Concat"(%cst_12, %14, %15) : (tensor, tensor<3xi32>, tensor<2xi32>) -> tensor<5xi32> // CHECK: %17 = "tf.Reshape"(%11, %16) : (tensor<2x?x?xf32>, tensor<5xi32>) -> tensor<2x2x?x4x?xf32> // CHECK: return %17 : tensor<2x2x?x4x?xf32> @@ -2107,26 +2107,26 @@ func.return %0 : tensor<2x2x?x4x?xf32> // CHECK-SAME: %arg0: tensor<4x4x?xf32>, // CHECK-SAME: %arg1: tensor<4x?x256xf32>) -> tensor<4x4x256xf32> { // CHECK: %0 = "tf.Shape"(%arg0) : (tensor<4x4x?xf32>) -> tensor<3xi32> -// CHECK-DAG: %cst = "tf.Const"() {value = dense<[-1, 0, -1]> : tensor<3xi32>} : () -> tensor<3xi32> -// CHECK-DAG: %cst_0 = "tf.Const"() {value = dense<[-1, -1, 0]> : tensor<3xi32>} : () -> tensor<3xi32> -// CHECK-DAG: %cst_1 = "tf.Const"() {value = dense<1> : tensor} : () -> tensor +// CHECK-DAG: %cst = "tf.Const"() <{value = dense<[-1, 0, -1]> : tensor<3xi32>}> : () -> tensor<3xi32> +// CHECK-DAG: %cst_0 = "tf.Const"() <{value = dense<[-1, -1, 0]> : tensor<3xi32>}> : () -> tensor<3xi32> +// CHECK-DAG: %cst_1 = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor // CHECK: %1 = "tf.UnsortedSegmentProd"(%0, %cst, %cst_1) : (tensor<3xi32>, tensor<3xi32>, tensor) -> tensor<1xi32> // CHECK: %2 = "tf.UnsortedSegmentProd"(%0, %cst_0, %cst_1) : (tensor<3xi32>, tensor<3xi32>, tensor) -> tensor<1xi32> -// CHECK-DAG: %cst_2 = "tf.Const"() {value = dense<4> : tensor<1xi32>} : () -> tensor<1xi32> -// CHECK-DAG: %cst_3 = "tf.Const"() {value = dense<0> : tensor} : () -> tensor +// CHECK-DAG: %cst_2 = "tf.Const"() <{value = dense<4> : tensor<1xi32>}> : () -> tensor<1xi32> +// CHECK-DAG: %cst_3 = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK: %3 = "tf.Concat"(%cst_3, %cst_2, %1, %2) : (tensor, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<3xi32> // CHECK: %4 = "tf.Reshape"(%arg0, %3) : (tensor<4x4x?xf32>, tensor<3xi32>) -> tensor<4x4x?xf32> // CHECK: %5 = "tf.Shape"(%arg1) : (tensor<4x?x256xf32>) -> tensor<3xi32> -// CHECK-DAG: %cst_4 = "tf.Const"() {value = dense<[-1, -1, 0]> : tensor<3xi32>} : () -> tensor<3xi32> -// CHECK-DAG: %cst_5 = "tf.Const"() {value = dense<[-1, 0, -1]> : tensor<3xi32>} : () -> tensor<3xi32> -// CHECK-DAG: %cst_6 = "tf.Const"() {value = dense<1> : tensor} : () -> tensor +// CHECK-DAG: %cst_4 = "tf.Const"() <{value = dense<[-1, -1, 0]> : tensor<3xi32>}> : () -> tensor<3xi32> +// CHECK-DAG: %cst_5 = "tf.Const"() <{value = dense<[-1, 0, -1]> : tensor<3xi32>}> : () -> tensor<3xi32> +// CHECK-DAG: %cst_6 = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor // CHECK: %6 = "tf.UnsortedSegmentProd"(%5, %cst_4, %cst_6) : (tensor<3xi32>, tensor<3xi32>, tensor) -> tensor<1xi32> // CHECK: %7 = "tf.UnsortedSegmentProd"(%5, %cst_5, %cst_6) : (tensor<3xi32>, tensor<3xi32>, tensor) -> tensor<1xi32> -// CHECK-DAG: %cst_7 = "tf.Const"() {value = dense<4> : tensor<1xi32>} : () -> tensor<1xi32> -// CHECK-DAG: %cst_8 = "tf.Const"() {value = dense<0> : tensor} : () -> tensor +// CHECK-DAG: %cst_7 = "tf.Const"() <{value = dense<4> : tensor<1xi32>}> : () -> tensor<1xi32> +// CHECK-DAG: %cst_8 = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK: %8 = "tf.Concat"(%cst_8, %cst_7, %7, %6) : (tensor, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<3xi32> // CHECK: %9 = "tf.Reshape"(%arg1, %8) : (tensor<4x?x256xf32>, tensor<3xi32>) -> tensor<4x?x256xf32> -// CHECK: %10 = "tf.BatchMatMulV3"(%4, %9) {adj_x = false, adj_y = false} : (tensor<4x4x?xf32>, tensor<4x?x256xf32>) -> tensor<4x4x256xf32> +// CHECK: %10 = "tf.BatchMatMulV3"(%4, %9) <{adj_x = false, adj_y = false}> : (tensor<4x4x?xf32>, tensor<4x?x256xf32>) -> tensor<4x4x256xf32> // CHECK: return %10 : tensor<4x4x256xf32> // CHECK: } func.func @convert_dot_general_dynamic_contracting_dim(%arg0: tensor<4x4x?xf32>, %arg1: tensor<4x?x256xf32>) -> tensor<4x4x256xf32> { @@ -2145,14 +2145,14 @@ func.return %0 : tensor<4x4x256xf32> // CHECK-SAME: %[[VAL_1:.*]]: tensor<1x256x256xbf16>) -> tensor<16x32x256xbf16> { // CHECK-DAG: %[[VAL_2:.*]] = arith.constant dense<[16, 32, 256, 1]> : tensor<4xi64> // CHECK: %[[VAL_3:.*]] = "tf.Reshape"(%[[VAL_0]], %[[VAL_2]]) : (tensor<16x32x256xbf16>, tensor<4xi64>) -> tensor<16x32x256x1xbf16> -// CHECK-DAG: %[[VAL_4:.*]] = "tf.Const"() {value = dense<[0, 1, 3, 2]> : tensor<4xi64>} : () -> tensor<4xi64> +// CHECK-DAG: %[[VAL_4:.*]] = "tf.Const"() <{value = dense<[0, 1, 3, 2]> : tensor<4xi64>}> : () -> tensor<4xi64> // CHECK: %[[VAL_5:.*]] = "tf.Transpose"(%[[VAL_3]], %[[VAL_4]]) : (tensor<16x32x256x1xbf16>, tensor<4xi64>) -> tensor<16x32x1x256xbf16> // CHECK-DAG: %[[VAL_6:.*]] = arith.constant dense<[1, 256, 256, 1]> : tensor<4xi64> // CHECK: %[[VAL_7:.*]] = "tf.Reshape"(%[[VAL_1]], %[[VAL_6]]) : (tensor<1x256x256xbf16>, tensor<4xi64>) -> tensor<1x256x256x1xbf16> -// CHECK-DAG: %[[VAL_8:.*]] = "tf.Const"() {value = dense<[0, 3, 1, 2]> : tensor<4xi64>} : () -> tensor<4xi64> +// CHECK-DAG: %[[VAL_8:.*]] = "tf.Const"() <{value = dense<[0, 3, 1, 2]> : tensor<4xi64>}> : () -> tensor<4xi64> // CHECK: %[[VAL_9:.*]] = "tf.Transpose"(%[[VAL_7]], %[[VAL_8]]) : (tensor<1x256x256x1xbf16>, tensor<4xi64>) -> tensor<1x1x256x256xbf16> -// CHECK: %[[VAL_10:.*]] = "tf.Conv2D"(%[[VAL_5]], %[[VAL_9]]) {data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "VALID", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true} : (tensor<16x32x1x256xbf16>, tensor<1x1x256x256xbf16>) -> tensor<16x32x1x256xbf16> -// CHECK: %[[VAL_11:.*]] = "tf.Const"() {value = dense<[0, 1, 3, 2]> : tensor<4xi64>} : () -> tensor<4xi64> +// CHECK: %[[VAL_10:.*]] = "tf.Conv2D"(%[[VAL_5]], %[[VAL_9]]) <{data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "VALID", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true}> : (tensor<16x32x1x256xbf16>, tensor<1x1x256x256xbf16>) -> tensor<16x32x1x256xbf16> +// CHECK: %[[VAL_11:.*]] = "tf.Const"() <{value = dense<[0, 1, 3, 2]> : tensor<4xi64>}> : () -> tensor<4xi64> // CHECK: %[[VAL_12:.*]] = "tf.Transpose"(%[[VAL_10]], %[[VAL_11]]) : (tensor<16x32x1x256xbf16>, tensor<4xi64>) -> tensor<16x32x256x1xbf16> // CHECK: %[[VAL_13:.*]] = arith.constant dense<[16, 32, 256]> : tensor<3xi64> // CHECK: %[[VAL_14:.*]] = "tf.Reshape"(%[[VAL_12]], %[[VAL_13]]) : (tensor<16x32x256x1xbf16>, tensor<3xi64>) -> tensor<16x32x256xbf16> @@ -2177,14 +2177,14 @@ func.func @convert_conv1d(%arg0: tensor<16x32x256xbf16>, %arg1: tensor<1x256x256 // CHECK-SAME: %[[VAL_1:.*]]: tensor<1x256x256xbf16>) -> tensor { // CHECK-DAG: %[[VAL_2:.*]] = arith.constant dense<[-9223372036854775808, 32, 256, 1]> : tensor<4xi64> // CHECK: %[[VAL_3:.*]] = "tf.Reshape"(%[[VAL_0]], %[[VAL_2]]) : (tensor, tensor<4xi64>) -> tensor -// CHECK-DAG: %[[VAL_4:.*]] = "tf.Const"() {value = dense<[0, 1, 3, 2]> : tensor<4xi64>} : () -> tensor<4xi64> +// CHECK-DAG: %[[VAL_4:.*]] = "tf.Const"() <{value = dense<[0, 1, 3, 2]> : tensor<4xi64>}> : () -> tensor<4xi64> // CHECK: %[[VAL_5:.*]] = "tf.Transpose"(%[[VAL_3]], %[[VAL_4]]) : (tensor, tensor<4xi64>) -> tensor // CHECK-DAG: %[[VAL_6:.*]] = arith.constant dense<[1, 256, 256, 1]> : tensor<4xi64> // CHECK: %[[VAL_7:.*]] = "tf.Reshape"(%[[VAL_1]], %[[VAL_6]]) : (tensor<1x256x256xbf16>, tensor<4xi64>) -> tensor<1x256x256x1xbf16> -// CHECK-DAG: %[[VAL_8:.*]] = "tf.Const"() {value = dense<[0, 3, 1, 2]> : tensor<4xi64>} : () -> tensor<4xi64> +// CHECK-DAG: %[[VAL_8:.*]] = "tf.Const"() <{value = dense<[0, 3, 1, 2]> : tensor<4xi64>}> : () -> tensor<4xi64> // CHECK: %[[VAL_9:.*]] = "tf.Transpose"(%[[VAL_7]], %[[VAL_8]]) : (tensor<1x256x256x1xbf16>, tensor<4xi64>) -> tensor<1x1x256x256xbf16> -// CHECK: %[[VAL_10:.*]] = "tf.Conv2D"(%[[VAL_5]], %[[VAL_9]]) {data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "VALID", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true} : (tensor, tensor<1x1x256x256xbf16>) -> tensor -// CHECK: %[[VAL_11:.*]] = "tf.Const"() {value = dense<[0, 1, 3, 2]> : tensor<4xi64>} : () -> tensor<4xi64> +// CHECK: %[[VAL_10:.*]] = "tf.Conv2D"(%[[VAL_5]], %[[VAL_9]]) <{data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "VALID", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true}> : (tensor, tensor<1x1x256x256xbf16>) -> tensor +// CHECK: %[[VAL_11:.*]] = "tf.Const"() <{value = dense<[0, 1, 3, 2]> : tensor<4xi64>}> : () -> tensor<4xi64> // CHECK: %[[VAL_12:.*]] = "tf.Transpose"(%[[VAL_10]], %[[VAL_11]]) : (tensor, tensor<4xi64>) -> tensor // CHECK: %[[VAL_13:.*]] = arith.constant dense<[-9223372036854775808, 32, 256]> : tensor<3xi64> // CHECK: %[[VAL_14:.*]] = "tf.Reshape"(%[[VAL_12]], %[[VAL_13]]) : (tensor, tensor<3xi64>) -> tensor @@ -2211,14 +2211,14 @@ func.func @convert_conv1d_dynamic_batch(%arg0: tensor, %arg1: ten // CHECK-SAME: %[[VAL_1:.*]]: tensor<1x256x256xbf16>) -> tensor<16x32x256xbf16> { // CHECK-DAG: %[[VAL_2:.*]] = arith.constant dense<[16, 32, 256, 1]> : tensor<4xi64> // CHECK: %[[VAL_3:.*]] = "tf.Reshape"(%[[VAL_0]], %[[VAL_2]]) : (tensor<16x32x256xbf16>, tensor<4xi64>) -> tensor<16x32x256x1xbf16> -// CHECK-DAG: %[[VAL_4:.*]] = "tf.Const"() {value = dense<[0, 1, 3, 2]> : tensor<4xi64>} : () -> tensor<4xi64> +// CHECK-DAG: %[[VAL_4:.*]] = "tf.Const"() <{value = dense<[0, 1, 3, 2]> : tensor<4xi64>}> : () -> tensor<4xi64> // CHECK: %[[VAL_5:.*]] = "tf.Transpose"(%[[VAL_3]], %[[VAL_4]]) : (tensor<16x32x256x1xbf16>, tensor<4xi64>) -> tensor<16x32x1x256xbf16> // CHECK-DAG: %[[VAL_6:.*]] = arith.constant dense<[1, 256, 256, 1]> : tensor<4xi64> // CHECK: %[[VAL_7:.*]] = "tf.Reshape"(%[[VAL_1]], %[[VAL_6]]) : (tensor<1x256x256xbf16>, tensor<4xi64>) -> tensor<1x256x256x1xbf16> -// CHECK-DAG: %[[VAL_8:.*]] = "tf.Const"() {value = dense<[0, 3, 1, 2]> : tensor<4xi64>} : () -> tensor<4xi64> +// CHECK-DAG: %[[VAL_8:.*]] = "tf.Const"() <{value = dense<[0, 3, 1, 2]> : tensor<4xi64>}> : () -> tensor<4xi64> // CHECK: %[[VAL_9:.*]] = "tf.Transpose"(%[[VAL_7]], %[[VAL_8]]) : (tensor<1x256x256x1xbf16>, tensor<4xi64>) -> tensor<1x1x256x256xbf16> -// CHECK: %[[VAL_10:.*]] = "tf.Conv2D"(%[[VAL_5]], %[[VAL_9]]) {data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "VALID", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true} : (tensor<16x32x1x256xbf16>, tensor<1x1x256x256xbf16>) -> tensor<16x32x1x256xbf16> -// CHECK: %[[VAL_11:.*]] = "tf.Const"() {value = dense<[0, 1, 3, 2]> : tensor<4xi64>} : () -> tensor<4xi64> +// CHECK: %[[VAL_10:.*]] = "tf.Conv2D"(%[[VAL_5]], %[[VAL_9]]) <{data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "VALID", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true}> : (tensor<16x32x1x256xbf16>, tensor<1x1x256x256xbf16>) -> tensor<16x32x1x256xbf16> +// CHECK: %[[VAL_11:.*]] = "tf.Const"() <{value = dense<[0, 1, 3, 2]> : tensor<4xi64>}> : () -> tensor<4xi64> // CHECK: %[[VAL_12:.*]] = "tf.Transpose"(%[[VAL_10]], %[[VAL_11]]) : (tensor<16x32x1x256xbf16>, tensor<4xi64>) -> tensor<16x32x256x1xbf16> // CHECK: %[[VAL_13:.*]] = arith.constant dense<[16, 32, 256]> : tensor<3xi64> // CHECK: %[[VAL_14:.*]] = "tf.Reshape"(%[[VAL_12]], %[[VAL_13]]) : (tensor<16x32x256x1xbf16>, tensor<3xi64>) -> tensor<16x32x256xbf16> @@ -2240,14 +2240,14 @@ func.func @convert_conv1d_no_lhs_dil_rhs_dil_precision_conf(%arg0: tensor<16x32x // CHECK-SAME: %[[VAL_1:.*]]: tensor<1x256x256xbf16>) -> tensor { // CHECK-DAG: %[[VAL_2:.*]] = arith.constant dense<[-9223372036854775808, 32, 256, 1]> : tensor<4xi64> // CHECK: %[[VAL_3:.*]] = "tf.Reshape"(%[[VAL_0]], %[[VAL_2]]) : (tensor, tensor<4xi64>) -> tensor -// CHECK-DAG: %[[VAL_4:.*]] = "tf.Const"() {value = dense<[0, 1, 3, 2]> : tensor<4xi64>} : () -> tensor<4xi64> +// CHECK-DAG: %[[VAL_4:.*]] = "tf.Const"() <{value = dense<[0, 1, 3, 2]> : tensor<4xi64>}> : () -> tensor<4xi64> // CHECK: %[[VAL_5:.*]] = "tf.Transpose"(%[[VAL_3]], %[[VAL_4]]) : (tensor, tensor<4xi64>) -> tensor // CHECK-DAG: %[[VAL_6:.*]] = arith.constant dense<[1, 256, 256, 1]> : tensor<4xi64> // CHECK: %[[VAL_7:.*]] = "tf.Reshape"(%[[VAL_1]], %[[VAL_6]]) : (tensor<1x256x256xbf16>, tensor<4xi64>) -> tensor<1x256x256x1xbf16> -// CHECK-DAG: %[[VAL_8:.*]] = "tf.Const"() {value = dense<[0, 3, 1, 2]> : tensor<4xi64>} : () -> tensor<4xi64> +// CHECK-DAG: %[[VAL_8:.*]] = "tf.Const"() <{value = dense<[0, 3, 1, 2]> : tensor<4xi64>}> : () -> tensor<4xi64> // CHECK: %[[VAL_9:.*]] = "tf.Transpose"(%[[VAL_7]], %[[VAL_8]]) : (tensor<1x256x256x1xbf16>, tensor<4xi64>) -> tensor<1x1x256x256xbf16> -// CHECK: %[[VAL_10:.*]] = "tf.Conv2D"(%[[VAL_5]], %[[VAL_9]]) {data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "VALID", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true} : (tensor, tensor<1x1x256x256xbf16>) -> tensor -// CHECK: %[[VAL_11:.*]] = "tf.Const"() {value = dense<[0, 1, 3, 2]> : tensor<4xi64>} : () -> tensor<4xi64> +// CHECK: %[[VAL_10:.*]] = "tf.Conv2D"(%[[VAL_5]], %[[VAL_9]]) <{data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "VALID", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true}> : (tensor, tensor<1x1x256x256xbf16>) -> tensor +// CHECK: %[[VAL_11:.*]] = "tf.Const"() <{value = dense<[0, 1, 3, 2]> : tensor<4xi64>}> : () -> tensor<4xi64> // CHECK: %[[VAL_12:.*]] = "tf.Transpose"(%[[VAL_10]], %[[VAL_11]]) : (tensor, tensor<4xi64>) -> tensor // CHECK: %[[VAL_13:.*]] = arith.constant dense<[-9223372036854775808, 32, 256]> : tensor<3xi64> // CHECK: %[[VAL_14:.*]] = "tf.Reshape"(%[[VAL_12]], %[[VAL_13]]) : (tensor, tensor<3xi64>) -> tensor @@ -2271,14 +2271,14 @@ func.func @convert_conv1d_no_lhs_dil_rhs_dil_precision_conf_dynamic_batch(%arg0: // CHECK-SAME: %[[VAL_1:.*]]: tensor<256x1x256xbf16>) -> tensor<256x16x32xbf16> { // CHECK-DAG: %[[VAL_2:.*]] = arith.constant dense<[32, 16, 256, 1]> : tensor<4xi64> // CHECK: %[[VAL_3:.*]] = "tf.Reshape"(%[[VAL_0]], %[[VAL_2]]) : (tensor<32x16x256xbf16>, tensor<4xi64>) -> tensor<32x16x256x1xbf16> -// CHECK-DAG: %[[VAL_4:.*]] = "tf.Const"() {value = dense<[1, 0, 3, 2]> : tensor<4xi64>} : () -> tensor<4xi64> +// CHECK-DAG: %[[VAL_4:.*]] = "tf.Const"() <{value = dense<[1, 0, 3, 2]> : tensor<4xi64>}> : () -> tensor<4xi64> // CHECK: %[[VAL_5:.*]] = "tf.Transpose"(%[[VAL_3]], %[[VAL_4]]) : (tensor<32x16x256x1xbf16>, tensor<4xi64>) -> tensor<16x32x1x256xbf16> // CHECK: %[[VAL_6:.*]] = arith.constant dense<[256, 1, 256, 1]> : tensor<4xi64> // CHECK: %[[VAL_7:.*]] = "tf.Reshape"(%[[VAL_1]], %[[VAL_6]]) : (tensor<256x1x256xbf16>, tensor<4xi64>) -> tensor<256x1x256x1xbf16> -// CHECK-DAG: %[[VAL_8:.*]] = "tf.Const"() {value = dense<[1, 3, 2, 0]> : tensor<4xi64>} : () -> tensor<4xi64> +// CHECK-DAG: %[[VAL_8:.*]] = "tf.Const"() <{value = dense<[1, 3, 2, 0]> : tensor<4xi64>}> : () -> tensor<4xi64> // CHECK: %[[VAL_9:.*]] = "tf.Transpose"(%[[VAL_7]], %[[VAL_8]]) : (tensor<256x1x256x1xbf16>, tensor<4xi64>) -> tensor<1x1x256x256xbf16> -// CHECK: %[[VAL_10:.*]] = "tf.Conv2D"(%[[VAL_5]], %[[VAL_9]]) {data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "VALID", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true} : (tensor<16x32x1x256xbf16>, tensor<1x1x256x256xbf16>) -> tensor<16x32x1x256xbf16> -// CHECK-DAG: %[[VAL_11:.*]] = "tf.Const"() {value = dense<[3, 0, 1, 2]> : tensor<4xi64>} : () -> tensor<4xi64> +// CHECK: %[[VAL_10:.*]] = "tf.Conv2D"(%[[VAL_5]], %[[VAL_9]]) <{data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "VALID", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true}> : (tensor<16x32x1x256xbf16>, tensor<1x1x256x256xbf16>) -> tensor<16x32x1x256xbf16> +// CHECK-DAG: %[[VAL_11:.*]] = "tf.Const"() <{value = dense<[3, 0, 1, 2]> : tensor<4xi64>}> : () -> tensor<4xi64> // CHECK: %[[VAL_12:.*]] = "tf.Transpose"(%[[VAL_10]], %[[VAL_11]]) : (tensor<16x32x1x256xbf16>, tensor<4xi64>) -> tensor<256x16x32x1xbf16> // CHECK: %[[VAL_13:.*]] = arith.constant dense<[256, 16, 32]> : tensor<3xi64> // CHECK: %[[VAL_14:.*]] = "tf.Reshape"(%[[VAL_12]], %[[VAL_13]]) : (tensor<256x16x32x1xbf16>, tensor<3xi64>) -> tensor<256x16x32xbf16> @@ -2343,14 +2343,14 @@ func.func @no_convert_conv1d_feature_group_gt_1(%arg0: tensor<16x32x256xbf16>, % // CHECK-SAME: %[[VAL_1:.*]]: tensor<1x256x256xbf16>) -> tensor<16x32x256xbf16> { // CHECK-DAG: %[[VAL_2:.*]] = arith.constant dense<[16, 32, 256, 1]> : tensor<4xi64> // CHECK: %[[VAL_3:.*]] = "tf.Reshape"(%[[VAL_0]], %[[VAL_2]]) : (tensor<16x32x256xbf16>, tensor<4xi64>) -> tensor<16x32x256x1xbf16> -// CHECK-DAG: %[[VAL_4:.*]] = "tf.Const"() {value = dense<[0, 1, 3, 2]> : tensor<4xi64>} : () -> tensor<4xi64> +// CHECK-DAG: %[[VAL_4:.*]] = "tf.Const"() <{value = dense<[0, 1, 3, 2]> : tensor<4xi64>}> : () -> tensor<4xi64> // CHECK: %[[VAL_5:.*]] = "tf.Transpose"(%[[VAL_3]], %[[VAL_4]]) : (tensor<16x32x256x1xbf16>, tensor<4xi64>) -> tensor<16x32x1x256xbf16> // CHECK-DAG: %[[VAL_6:.*]] = arith.constant dense<[1, 256, 256, 1]> : tensor<4xi64> // CHECK: %[[VAL_7:.*]] = "tf.Reshape"(%[[VAL_1]], %[[VAL_6]]) : (tensor<1x256x256xbf16>, tensor<4xi64>) -> tensor<1x256x256x1xbf16> -// CHECK-DAG: %[[VAL_8:.*]] = "tf.Const"() {value = dense<[0, 3, 1, 2]> : tensor<4xi64>} : () -> tensor<4xi64> +// CHECK-DAG: %[[VAL_8:.*]] = "tf.Const"() <{value = dense<[0, 3, 1, 2]> : tensor<4xi64>}> : () -> tensor<4xi64> // CHECK: %[[VAL_9:.*]] = "tf.Transpose"(%[[VAL_7]], %[[VAL_8]]) : (tensor<1x256x256x1xbf16>, tensor<4xi64>) -> tensor<1x1x256x256xbf16> -// CHECK: %[[VAL_10:.*]] = "tf.Conv2D"(%[[VAL_5]], %[[VAL_9]]) {data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "VALID", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true} : (tensor<16x32x1x256xbf16>, tensor<1x1x256x256xbf16>) -> tensor<16x32x1x256xbf16> -// CHECK: %[[VAL_11:.*]] = "tf.Const"() {value = dense<[0, 1, 3, 2]> : tensor<4xi64>} : () -> tensor<4xi64> +// CHECK: %[[VAL_10:.*]] = "tf.Conv2D"(%[[VAL_5]], %[[VAL_9]]) <{data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "VALID", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true}> : (tensor<16x32x1x256xbf16>, tensor<1x1x256x256xbf16>) -> tensor<16x32x1x256xbf16> +// CHECK: %[[VAL_11:.*]] = "tf.Const"() <{value = dense<[0, 1, 3, 2]> : tensor<4xi64>}> : () -> tensor<4xi64> // CHECK: %[[VAL_12:.*]] = "tf.Transpose"(%[[VAL_10]], %[[VAL_11]]) : (tensor<16x32x1x256xbf16>, tensor<4xi64>) -> tensor<16x32x256x1xbf16> // CHECK: %[[VAL_13:.*]] = arith.constant dense<[16, 32, 256]> : tensor<3xi64> // CHECK: %[[VAL_14:.*]] = "tf.Reshape"(%[[VAL_12]], %[[VAL_13]]) : (tensor<16x32x256x1xbf16>, tensor<3xi64>) -> tensor<16x32x256xbf16> @@ -2372,7 +2372,7 @@ func.func @convert_conv1d_missing_windows_strides_fallback(%arg0: tensor<16x32x2 // CHECK-LABEL: func.func @convert_conv1d_missing_windows_strides_fallback_2( // CHECK-SAME: %[[VAL_0:.*]]: tensor<1x64x64x4xbf16>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<3x3x4x320xbf16>) -> tensor<1x62x62x320xbf16> { -// CHECK: %[[VAL_2:.*]] = "tf.Conv2D"(%[[VAL_0]], %[[VAL_1]]) {data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "VALID", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true} : (tensor<1x64x64x4xbf16>, tensor<3x3x4x320xbf16>) -> tensor<1x62x62x320xbf16> +// CHECK: %[[VAL_2:.*]] = "tf.Conv2D"(%[[VAL_0]], %[[VAL_1]]) <{data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "VALID", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true}> : (tensor<1x64x64x4xbf16>, tensor<3x3x4x320xbf16>) -> tensor<1x62x62x320xbf16> // CHECK: return %[[VAL_2]] : tensor<1x62x62x320xbf16> // CHECK: } func.func @convert_conv1d_missing_windows_strides_fallback_2(%arg0: tensor<1x64x64x4xbf16>, %arg1: tensor<3x3x4x320xbf16>) -> tensor<1x62x62x320xbf16> { @@ -2391,7 +2391,7 @@ func.func @convert_conv1d_missing_windows_strides_fallback_2(%arg0: tensor<1x64x // CHECK-LABEL: func @convert_conv2d( // CHECK-SAME: %[[VAL_0:.*]]: tensor<1x8x8x207xf32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> { -// CHECK: %[[VAL_2:.*]] = "tf.Conv2D"(%[[VAL_0]], %[[VAL_1]]) {data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true} : (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> +// CHECK: %[[VAL_2:.*]] = "tf.Conv2D"(%[[VAL_0]], %[[VAL_1]]) <{data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true}> : (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> // CHECK: return %[[VAL_2]] : tensor<1x8x8x16xf32> // CHECK: } func.func @convert_conv2d(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> { @@ -2414,7 +2414,7 @@ func.func @convert_conv2d(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16 // CHECK-LABEL: func @convert_group_conv2d( // CHECK-SAME: %[[VAL_0:.*]]: tensor<1x14x14x2240xf32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<3x3x112x2240xf32>) -> tensor<1x7x7x2240xf32> { -// CHECK: %[[VAL_2:.*]] = "tf.Conv2D"(%[[VAL_0]], %[[VAL_1]]) {data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [0, 0, 1, 1, 1, 1, 0, 0], padding = "EXPLICIT", strides = [1, 2, 2, 1], use_cudnn_on_gpu = true} : (tensor<1x14x14x2240xf32>, tensor<3x3x112x2240xf32>) -> tensor<1x7x7x2240xf32> +// CHECK: %[[VAL_2:.*]] = "tf.Conv2D"(%[[VAL_0]], %[[VAL_1]]) <{data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [0, 0, 1, 1, 1, 1, 0, 0], padding = "EXPLICIT", strides = [1, 2, 2, 1], use_cudnn_on_gpu = true}> : (tensor<1x14x14x2240xf32>, tensor<3x3x112x2240xf32>) -> tensor<1x7x7x2240xf32> // CHECk: return %[[VAL_2]] : tensor<1x7x7x2240xf32> // CHECK: } func.func @convert_group_conv2d(%arg0: tensor<1x14x14x2240xf32>, %arg1: tensor<3x3x112x2240xf32>) -> tensor<1x7x7x2240xf32> { @@ -2428,13 +2428,13 @@ func.func @convert_group_conv2d(%arg0: tensor<1x14x14x2240xf32>, %arg1: tensor<3 // CHECK-LABEL: func.func @convert_transpose_conv_with_transpose( // CHECK-SAME: %[[VAL_0:.*]]: tensor<1x256x64x64xf32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<2x2x64x256xf32>) -> tensor<1x64x128x128xf32> { -// CHECK: %[[VAL_2:.*]] = "tf.Const"() {value = dense<[0, 2, 3, 1]> : tensor<4xi64>} : () -> tensor<4xi64> +// CHECK: %[[VAL_2:.*]] = "tf.Const"() <{value = dense<[0, 2, 3, 1]> : tensor<4xi64>}> : () -> tensor<4xi64> // CHECK: %[[VAL_3:.*]] = "tf.Transpose"(%[[VAL_0:.*]], %[[VAL_2:.*]]) : (tensor<1x256x64x64xf32>, tensor<4xi64>) -> tensor<1x64x64x256xf32> -// CHECK: %[[VAL_4:.*]] = "tf.Const"() {value = dense<[0, 1]> : tensor<2xi64>} : () -> tensor<2xi64> +// CHECK: %[[VAL_4:.*]] = "tf.Const"() <{value = dense<[0, 1]> : tensor<2xi64>}> : () -> tensor<2xi64> // CHECK: %[[VAL_5:.*]] = "tf.ReverseV2"(%[[VAL_1:.*]], %[[VAL_4:.*]]) : (tensor<2x2x64x256xf32>, tensor<2xi64>) -> tensor<2x2x64x256xf32> -// CHECK: %[[VAL_6:.*]] = "tf.Const"() {value = dense<[1, 128, 128, 64]> : tensor<4xi32>} : () -> tensor<4xi32> -// CHECK: %[[VAL_7:.*]] = "tf.Conv2DBackpropInput"(%[[VAL_6:.*]], %[[VAL_5:.*]], %[[VAL_3:.*]]) {data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 2, 2, 1], use_cudnn_on_gpu = true} : (tensor<4xi32>, tensor<2x2x64x256xf32>, tensor<1x64x64x256xf32>) -> tensor<1x128x128x64xf32> -// CHECK: %[[VAL_8:.*]] = "tf.Const"() {value = dense<[0, 3, 1, 2]> : tensor<4xi64>} : () -> tensor<4xi64> +// CHECK: %[[VAL_6:.*]] = "tf.Const"() <{value = dense<[1, 128, 128, 64]> : tensor<4xi32>}> : () -> tensor<4xi32> +// CHECK: %[[VAL_7:.*]] = "tf.Conv2DBackpropInput"(%[[VAL_6:.*]], %[[VAL_5:.*]], %[[VAL_3:.*]]) <{data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 2, 2, 1], use_cudnn_on_gpu = true}> : (tensor<4xi32>, tensor<2x2x64x256xf32>, tensor<1x64x64x256xf32>) -> tensor<1x128x128x64xf32> +// CHECK: %[[VAL_8:.*]] = "tf.Const"() <{value = dense<[0, 3, 1, 2]> : tensor<4xi64>}> : () -> tensor<4xi64> // CHECK: %[[VAL_9:.*]] = "tf.Transpose"(%[[VAL_7:.*]], %[[VAL_8:.*]]) : (tensor<1x128x128x64xf32>, tensor<4xi64>) -> tensor<1x64x128x128xf32> // CHECK: return %[[VAL_9:.*]] : tensor<1x64x128x128xf32> // CHECK: } @@ -2450,13 +2450,13 @@ func.func @convert_transpose_conv_with_transpose(%arg0: tensor<1x256x64x64xf32>, // CHECK-LABEL: func.func @convert_transpose_conv_with_transpose2( // CHECK-SAME: %[[VAL_0:.*]]: tensor<64x64x1x256xf32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<2x2x64x256xf32>) -> tensor<128x128x1x64xf32> { -// CHECK: %[[VAL_2:.*]] = "tf.Const"() {value = dense<[2, 0, 1, 3]> : tensor<4xi64>} : () -> tensor<4xi64> +// CHECK: %[[VAL_2:.*]] = "tf.Const"() <{value = dense<[2, 0, 1, 3]> : tensor<4xi64>}> : () -> tensor<4xi64> // CHECK: %[[VAL_3:.*]] = "tf.Transpose"(%[[VAL_0:.*]], %[[VAL_2:.*]]) : (tensor<64x64x1x256xf32>, tensor<4xi64>) -> tensor<1x64x64x256xf32> -// CHECK: %[[VAL_4:.*]] = "tf.Const"() {value = dense<[0, 1]> : tensor<2xi64>} : () -> tensor<2xi64> +// CHECK: %[[VAL_4:.*]] = "tf.Const"() <{value = dense<[0, 1]> : tensor<2xi64>}> : () -> tensor<2xi64> // CHECK: %[[VAL_5:.*]] = "tf.ReverseV2"(%[[VAL_1:.*]], %[[VAL_4:.*]]) : (tensor<2x2x64x256xf32>, tensor<2xi64>) -> tensor<2x2x64x256xf32> -// CHECK: %[[VAL_6:.*]] = "tf.Const"() {value = dense<[1, 128, 128, 64]> : tensor<4xi32>} : () -> tensor<4xi32> -// CHECK: %[[VAL_7:.*]] = "tf.Conv2DBackpropInput"(%[[VAL_6:.*]], %[[VAL_5:.*]], %[[VAL_3:.*]]) {data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 2, 2, 1], use_cudnn_on_gpu = true} : (tensor<4xi32>, tensor<2x2x64x256xf32>, tensor<1x64x64x256xf32>) -> tensor<1x128x128x64xf32> -// CHECK: %[[VAL_8:.*]] = "tf.Const"() {value = dense<[1, 2, 0, 3]> : tensor<4xi64>} : () -> tensor<4xi64> +// CHECK: %[[VAL_6:.*]] = "tf.Const"() <{value = dense<[1, 128, 128, 64]> : tensor<4xi32>}> : () -> tensor<4xi32> +// CHECK: %[[VAL_7:.*]] = "tf.Conv2DBackpropInput"(%[[VAL_6:.*]], %[[VAL_5:.*]], %[[VAL_3:.*]]) <{data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 2, 2, 1], use_cudnn_on_gpu = true}> : (tensor<4xi32>, tensor<2x2x64x256xf32>, tensor<1x64x64x256xf32>) -> tensor<1x128x128x64xf32> +// CHECK: %[[VAL_8:.*]] = "tf.Const"() <{value = dense<[1, 2, 0, 3]> : tensor<4xi64>}> : () -> tensor<4xi64> // CHECK: %[[VAL_9:.*]] = "tf.Transpose"(%[[VAL_7:.*]], %[[VAL_8:.*]]) : (tensor<1x128x128x64xf32>, tensor<4xi64>) -> tensor<128x128x1x64xf32> // CHECK: return %[[VAL_9:.*]] : tensor<128x128x1x64xf32> // CHECK: } @@ -2473,7 +2473,7 @@ func.func @convert_transpose_conv_with_transpose2(%arg0: tensor<64x64x1x256xf32> // CHECK-LABEL: func @convert_conv2d_dynamic_batch( // CHECK-SAME: %[[VAL_0:.*]]: tensor, // CHECK-SAME: %[[VAL_1:.*]]: tensor<3x3x207x16xf32>) -> tensor { -// CHECK: %[[VAL_2:.*]] = "tf.Conv2D"(%[[VAL_0]], %[[VAL_1]]) {data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true} : (tensor, tensor<3x3x207x16xf32>) -> tensor +// CHECK: %[[VAL_2:.*]] = "tf.Conv2D"(%[[VAL_0]], %[[VAL_1]]) <{data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true}> : (tensor, tensor<3x3x207x16xf32>) -> tensor // CHECK: return %[[VAL_2]] : tensor // CHECK: } func.func @convert_conv2d_dynamic_batch(%arg0: tensor, %arg1: tensor<3x3x207x16xf32>) -> tensor { @@ -2496,7 +2496,7 @@ func.func @convert_conv2d_dynamic_batch(%arg0: tensor, %arg1: ten // CHECK-LABEL: func @convert_conv2d_no_padding( // CHECK-SAME: %[[VAL_0:.*]]: tensor<1x6x6x207xf32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<3x3x207x16xf32>) -> tensor<1x4x4x16xf32> { -// CHECK: %[[VAL_2:.*]] = "tf.Conv2D"(%[[VAL_0]], %[[VAL_1]]) {data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "VALID", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true} : (tensor<1x6x6x207xf32>, tensor<3x3x207x16xf32>) -> tensor<1x4x4x16xf32> +// CHECK: %[[VAL_2:.*]] = "tf.Conv2D"(%[[VAL_0]], %[[VAL_1]]) <{data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "VALID", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true}> : (tensor<1x6x6x207xf32>, tensor<3x3x207x16xf32>) -> tensor<1x4x4x16xf32> // CHECK: return %[[VAL_2]] : tensor<1x4x4x16xf32> // CHECK: } func.func @convert_conv2d_no_padding(%arg0: tensor<1x6x6x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x4x4x16xf32> { @@ -2519,7 +2519,7 @@ func.func @convert_conv2d_no_padding(%arg0: tensor<1x6x6x207xf32>, %arg1: tensor // CHECK-LABEL: func @convert_conv2d_no_rhs_dilation( // CHECK-SAME: %[[VAL_0:.*]]: tensor<1x8x8x207xf32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> { -// CHECK: %[[VAL_2:.*]] = "tf.Conv2D"(%[[VAL_0]], %[[VAL_1]]) {data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true} : (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> +// CHECK: %[[VAL_2:.*]] = "tf.Conv2D"(%[[VAL_0]], %[[VAL_1]]) <{data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true}> : (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> // CHECK: return %[[VAL_2]] : tensor<1x8x8x16xf32> // CHECK: } func.func @convert_conv2d_no_rhs_dilation(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> { @@ -2542,7 +2542,7 @@ func.func @convert_conv2d_no_rhs_dilation(%arg0: tensor<1x8x8x207xf32>, %arg1: t // CHECK-LABEL: func @convert_conv2d_no_window_strides( // CHECK-SAME: %[[VAL_0:.*]]: tensor<1x8x8x207xf32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> { -// CHECK: %[[VAL_2:.*]] = "tf.Conv2D"(%[[VAL_0]], %[[VAL_1]]) {data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true} : (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> +// CHECK: %[[VAL_2:.*]] = "tf.Conv2D"(%[[VAL_0]], %[[VAL_1]]) <{data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true}> : (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> // CHECK: return %[[VAL_2]] : tensor<1x8x8x16xf32> // CHECK: } func.func @convert_conv2d_no_window_strides(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> { @@ -2565,7 +2565,7 @@ func.func @convert_conv2d_no_window_strides(%arg0: tensor<1x8x8x207xf32>, %arg1: // CHECK-LABEL: func @convert_conv2d_no_lhs_dilation( // CHECK-SAME: %[[VAL_0:.*]]: tensor<1x8x8x207xf32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> { -// CHECK: %[[VAL_2:.*]] = "tf.Conv2D"(%[[VAL_0]], %[[VAL_1]]) {data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true} : (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> +// CHECK: %[[VAL_2:.*]] = "tf.Conv2D"(%[[VAL_0]], %[[VAL_1]]) <{data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true}> : (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> // CHECK: return %[[VAL_2]] : tensor<1x8x8x16xf32> // CHECK: } func.func @convert_conv2d_no_lhs_dilation(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> { @@ -2588,12 +2588,12 @@ func.func @convert_conv2d_no_lhs_dilation(%arg0: tensor<1x8x8x207xf32>, %arg1: t // CHECK-LABEL: func @convert_conv2d_with_transpose( // CHECK-SAME: %[[VAL_0:.*]]: tensor<8x8x1x207xf32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<3x3x16x207xf32>) -> tensor<16x8x8x1xf32> { -// CHECK: %[[VAL_2:.*]] = "tf.Const"() {value = dense<[2, 0, 1, 3]> : tensor<4xi64>} : () -> tensor<4xi64> +// CHECK: %[[VAL_2:.*]] = "tf.Const"() <{value = dense<[2, 0, 1, 3]> : tensor<4xi64>}> : () -> tensor<4xi64> // CHECK: %[[VAL_3:.*]] = "tf.Transpose"(%[[VAL_0]], %[[VAL_2]]) : (tensor<8x8x1x207xf32>, tensor<4xi64>) -> tensor<1x8x8x207xf32> -// CHECK: %[[VAL_4:.*]] = "tf.Const"() {value = dense<[0, 1, 3, 2]> : tensor<4xi64>} : () -> tensor<4xi64> +// CHECK: %[[VAL_4:.*]] = "tf.Const"() <{value = dense<[0, 1, 3, 2]> : tensor<4xi64>}> : () -> tensor<4xi64> // CHECK: %[[VAL_5:.*]] = "tf.Transpose"(%[[VAL_1]], %[[VAL_4]]) : (tensor<3x3x16x207xf32>, tensor<4xi64>) -> tensor<3x3x207x16xf32> -// CHECK: %[[VAL_6:.*]] = "tf.Conv2D"(%[[VAL_3]], %[[VAL_5]]) {data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true} : (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> -// CHECK: %[[VAL_7:.*]] = "tf.Const"() {value = dense<[3, 1, 2, 0]> : tensor<4xi64>} : () -> tensor<4xi64> +// CHECK: %[[VAL_6:.*]] = "tf.Conv2D"(%[[VAL_3]], %[[VAL_5]]) <{data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true}> : (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<1x8x8x16xf32> +// CHECK: %[[VAL_7:.*]] = "tf.Const"() <{value = dense<[3, 1, 2, 0]> : tensor<4xi64>}> : () -> tensor<4xi64> // CHECK: %[[VAL_8:.*]] = "tf.Transpose"(%[[VAL_6]], %[[VAL_7]]) : (tensor<1x8x8x16xf32>, tensor<4xi64>) -> tensor<16x8x8x1xf32> // CHECK: return %[[VAL_8]] : tensor<16x8x8x1xf32> // CHECK: } @@ -2617,12 +2617,12 @@ func.func @convert_conv2d_with_transpose(%arg0: tensor<8x8x1x207xf32>, %arg1: te // CHECK-LABEL: func @convert_conv2d_with_transpose_dynamic_batch( // CHECK-SAME: %[[VAL_0:.*]]: tensor<8x8x?x207xf32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<3x3x16x207xf32>) -> tensor<16x8x8x?xf32> { -// CHECK: %[[VAL_2:.*]] = "tf.Const"() {value = dense<[2, 0, 1, 3]> : tensor<4xi64>} : () -> tensor<4xi64> +// CHECK: %[[VAL_2:.*]] = "tf.Const"() <{value = dense<[2, 0, 1, 3]> : tensor<4xi64>}> : () -> tensor<4xi64> // CHECK: %[[VAL_3:.*]] = "tf.Transpose"(%[[VAL_0]], %[[VAL_2]]) : (tensor<8x8x?x207xf32>, tensor<4xi64>) -> tensor -// CHECK: %[[VAL_4:.*]] = "tf.Const"() {value = dense<[0, 1, 3, 2]> : tensor<4xi64>} : () -> tensor<4xi64> +// CHECK: %[[VAL_4:.*]] = "tf.Const"() <{value = dense<[0, 1, 3, 2]> : tensor<4xi64>}> : () -> tensor<4xi64> // CHECK: %[[VAL_5:.*]] = "tf.Transpose"(%[[VAL_1]], %[[VAL_4]]) : (tensor<3x3x16x207xf32>, tensor<4xi64>) -> tensor<3x3x207x16xf32> -// CHECK: %[[VAL_6:.*]] = "tf.Conv2D"(%[[VAL_3]], %[[VAL_5]]) {data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true} : (tensor, tensor<3x3x207x16xf32>) -> tensor -// CHECK: %[[VAL_7:.*]] = "tf.Const"() {value = dense<[3, 1, 2, 0]> : tensor<4xi64>} : () -> tensor<4xi64> +// CHECK: %[[VAL_6:.*]] = "tf.Conv2D"(%[[VAL_3]], %[[VAL_5]]) <{data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true}> : (tensor, tensor<3x3x207x16xf32>) -> tensor +// CHECK: %[[VAL_7:.*]] = "tf.Const"() <{value = dense<[3, 1, 2, 0]> : tensor<4xi64>}> : () -> tensor<4xi64> // CHECK: %[[VAL_8:.*]] = "tf.Transpose"(%[[VAL_6]], %[[VAL_7]]) : (tensor, tensor<4xi64>) -> tensor<16x8x8x?xf32> // CHECK: return %[[VAL_8]] : tensor<16x8x8x?xf32> // CHECK: } @@ -2646,7 +2646,7 @@ func.func @convert_conv2d_with_transpose_dynamic_batch(%arg0: tensor<8x8x?x207xf // CHECK-LABEL: func @convert_conv2d_explicit_padding( // CHECK-SAME: %[[VAL_0:.*]]: tensor<64x8x8x8xf32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<8x8x8x64xf32>) -> tensor<64x3x3x64xf32> { -// CHECK: %[[VAL_2:.*]] = "tf.Conv2D"(%[[VAL_0]], %[[VAL_1]]) {data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [0, 0, 1, 1, 1, 1, 0, 0], padding = "EXPLICIT", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true} : (tensor<64x8x8x8xf32>, tensor<8x8x8x64xf32>) -> tensor<64x3x3x64xf32> +// CHECK: %[[VAL_2:.*]] = "tf.Conv2D"(%[[VAL_0]], %[[VAL_1]]) <{data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [0, 0, 1, 1, 1, 1, 0, 0], padding = "EXPLICIT", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true}> : (tensor<64x8x8x8xf32>, tensor<8x8x8x64xf32>) -> tensor<64x3x3x64xf32> // CHECK: return %[[VAL_2]] : tensor<64x3x3x64xf32> // CHECK: } func.func @convert_conv2d_explicit_padding(%arg0: tensor<64x8x8x8xf32>, %arg1: tensor<8x8x8x64xf32>) -> tensor<64x3x3x64xf32> { @@ -2670,7 +2670,7 @@ func.func @convert_conv2d_explicit_padding(%arg0: tensor<64x8x8x8xf32>, %arg1: t // CHECK-LABEL: func @convert_conv2d_explicit_padding_dynamic_batch( // CHECK-SAME: %[[VAL_0:.*]]: tensor, // CHECK-SAME: %[[VAL_1:.*]]: tensor<8x8x8x64xf32>) -> tensor { -// CHECK: %[[VAL_2:.*]] = "tf.Conv2D"(%[[VAL_0]], %[[VAL_1]]) {data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [0, 0, 1, 1, 1, 1, 0, 0], padding = "EXPLICIT", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true} : (tensor, tensor<8x8x8x64xf32>) -> tensor +// CHECK: %[[VAL_2:.*]] = "tf.Conv2D"(%[[VAL_0]], %[[VAL_1]]) <{data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [0, 0, 1, 1, 1, 1, 0, 0], padding = "EXPLICIT", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true}> : (tensor, tensor<8x8x8x64xf32>) -> tensor // CHECK: return %[[VAL_2]] : tensor // CHECK: } func.func @convert_conv2d_explicit_padding_dynamic_batch(%arg0: tensor, %arg1: tensor<8x8x8x64xf32>) -> tensor { @@ -2694,8 +2694,8 @@ func.func @convert_conv2d_explicit_padding_dynamic_batch(%arg0: tensor, // CHECK-SAME: %[[ARG1:.*]]: tensor<3x2x64x4xf32>) -> tensor<128x4x3x4xf32> { -// CHECK-DAG: %[[START:.*]] = "tf.Const"() {value = dense<[0, 0, 5, 0]> : tensor<4xi64>} : () -> tensor<4xi64> -// CHECK-DAG: %[[SIZE:.*]] = "tf.Const"() {value = dense<[128, 5, 4, 64]> : tensor<4xi64>} : () -> tensor<4xi64> +// CHECK-DAG: %[[START:.*]] = "tf.Const"() <{value = dense<[0, 0, 5, 0]> : tensor<4xi64>}> : () -> tensor<4xi64> +// CHECK-DAG: %[[SIZE:.*]] = "tf.Const"() <{value = dense<[128, 5, 4, 64]> : tensor<4xi64>}> : () -> tensor<4xi64> // CHECK: %[[SLICED_ARG0:.*]] = "tf.Slice"(%[[ARG0]], %[[START]], %[[SIZE]]) // CHECK-SAME: (tensor<128x7x9x64xf32>, tensor<4xi64>, tensor<4xi64>) -> tensor<128x5x4x64xf32> // CHECK: %[[CONV:.*]] = "tf.Conv2D"(%[[SLICED_ARG0]], %[[ARG1]]) @@ -2724,8 +2724,8 @@ func.func @convert_conv2d_negative_explicit_padding(%arg0: tensor<128x7x9x64xf32 // CHECK-LABEL: func @convert_conv2d_negative_explicit_padding_dynamic_batch( // CHECK-SAME: %[[ARG0:.*]]: tensor, // CHECK-SAME: %[[ARG1:.*]]: tensor<3x2x64x4xf32>) -> tensor { -// CHECK-DAG: %[[START:.*]] = "tf.Const"() {value = dense<[0, 0, 5, 0]> : tensor<4xi64>} : () -> tensor<4xi64> -// CHECK-DAG: %[[SIZE:.*]] = "tf.Const"() {value = dense<[-9223372036854775808, 5, 4, 64]> : tensor<4xi64>} : () -> tensor<4xi64> +// CHECK-DAG: %[[START:.*]] = "tf.Const"() <{value = dense<[0, 0, 5, 0]> : tensor<4xi64>}> : () -> tensor<4xi64> +// CHECK-DAG: %[[SIZE:.*]] = "tf.Const"() <{value = dense<[-9223372036854775808, 5, 4, 64]> : tensor<4xi64>}> : () -> tensor<4xi64> // CHECK: %[[SLICED_ARG0:.*]] = "tf.Slice"(%[[ARG0]], %[[START]], %[[SIZE]]) // CHECK-SAME: (tensor, tensor<4xi64>, tensor<4xi64>) -> tensor // CHECK: %[[CONV:.*]] = "tf.Conv2D"(%[[SLICED_ARG0]], %[[ARG1]]) @@ -2756,7 +2756,7 @@ func.func @convert_conv2d_negative_explicit_padding_dynamic_batch(%arg0: tensor< // CHECK-SAME: %[[VAL_1:.*]]: tensor<3x3x1x3312xf32>) -> tensor<1x8x8x3312xf32> { // CHECK: %[[CST:.*]] = arith.constant dense<[3, 3, 207, 16]> : tensor<4xi64> // CHECK: %[[VAL_2:.*]] = "tf.Reshape"(%[[VAL_1]], %[[CST]]) : (tensor<3x3x1x3312xf32>, tensor<4xi64>) -> tensor<3x3x207x16xf32> -// CHECK: %[[VAL_3:.*]] = "tf.DepthwiseConv2dNative"(%[[VAL_0]], %[[VAL_2]]) {data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 1, 1]} : (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<1x8x8x3312xf32> +// CHECK: %[[VAL_3:.*]] = "tf.DepthwiseConv2dNative"(%[[VAL_0]], %[[VAL_2]]) <{data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 1, 1]}> : (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<1x8x8x3312xf32> // CHECK: return %[[VAL_3]] : tensor<1x8x8x3312xf32> // CHECK: } func.func @convert_depthwise_conv2d(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x1x3312xf32>) -> tensor<1x8x8x3312xf32> { @@ -2779,8 +2779,8 @@ func.func @convert_depthwise_conv2d(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor< // CHECK-LABEL: func @convert_conv2d_to_resize( // CHECK-SAME: %[[VAL_0:.*]]: tensor<1x56x624x16xf32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<1x257x16x1xf32>) -> tensor<1x56x904x16xf32> { -// CHECK-DAG: %[[SIZE:.*]] = "tf.Const"() {value = dense<[56, 904]> : tensor<2xi32>} : () -> tensor<2xi32> -// CHECK: %[[VAL_2:.*]] = "tf.ResizeBilinear"(%[[VAL_0]], %[[SIZE]]) {align_corners = true, half_pixel_centers = false} : (tensor<1x56x624x16xf32>, tensor<2xi32>) -> tensor<1x56x904x16xf32> +// CHECK-DAG: %[[SIZE:.*]] = "tf.Const"() <{value = dense<[56, 904]> : tensor<2xi32>}> : () -> tensor<2xi32> +// CHECK: %[[VAL_2:.*]] = "tf.ResizeBilinear"(%[[VAL_0]], %[[SIZE]]) <{align_corners = true, half_pixel_centers = false}> : (tensor<1x56x624x16xf32>, tensor<2xi32>) -> tensor<1x56x904x16xf32> // CHECK: return %[[VAL_2]] : tensor<1x56x904x16xf32> // CHECK: } func.func @convert_conv2d_to_resize(%arg0: tensor<1x56x624x16xf32>, %arg1: tensor<1x257x16x1xf32>) -> tensor<1x56x904x16xf32> { @@ -2798,8 +2798,8 @@ func.func @convert_conv2d_to_resize(%arg0: tensor<1x56x624x16xf32>, %arg1: tenso // CHECK-LABEL: func @convert_conv2d_resize_perferred( // CHECK-SAME: %[[VAL_0:.*]]: tensor<1x56x1248x16xf32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<3x1x16x1xf32>) -> tensor<1x111x1248x16xf32> { -// CHECK-DAG: %[[SIZE:.*]] = "tf.Const"() {value = dense<[111, 1248]> : tensor<2xi32>} : () -> tensor<2xi32> -// CHECK: %[[VAL_2:.*]] = "tf.ResizeBilinear"(%[[VAL_0]], %[[SIZE]]) {align_corners = true, half_pixel_centers = false} : (tensor<1x56x1248x16xf32>, tensor<2xi32>) -> tensor<1x111x1248x16xf32> +// CHECK-DAG: %[[SIZE:.*]] = "tf.Const"() <{value = dense<[111, 1248]> : tensor<2xi32>}> : () -> tensor<2xi32> +// CHECK: %[[VAL_2:.*]] = "tf.ResizeBilinear"(%[[VAL_0]], %[[SIZE]]) <{align_corners = true, half_pixel_centers = false}> : (tensor<1x56x1248x16xf32>, tensor<2xi32>) -> tensor<1x111x1248x16xf32> // CHECK: return %[[VAL_2]] : tensor<1x111x1248x16xf32> // CHECK: } func.func @convert_conv2d_resize_perferred(%arg0: tensor<1x56x1248x16xf32>, %arg1: tensor<3x1x16x1xf32>) -> tensor<1x111x1248x16xf32> { @@ -2817,10 +2817,10 @@ func.func @convert_conv2d_resize_perferred(%arg0: tensor<1x56x1248x16xf32>, %arg // CHECK-LABEL: func @convert_conv2d_back_prop_input_same_pad( // CHECK-SAME: %[[VAL_0:.*]]: tensor<1x256x256x2xf32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<4x4x2x2xf32>) -> tensor<1x512x512x2xf32> { -// CHECK: %[[VAL_3:.*]] = "tf.Const"() {value = dense<[0, 1]> : tensor<2xi64>} : () -> tensor<2xi64> +// CHECK: %[[VAL_3:.*]] = "tf.Const"() <{value = dense<[0, 1]> : tensor<2xi64>}> : () -> tensor<2xi64> // CHECK: %[[VAL_4:.*]] = "tf.ReverseV2"(%[[VAL_1]], %[[VAL_3]]) : (tensor<4x4x2x2xf32>, tensor<2xi64>) -> tensor<4x4x2x2xf32> -// CHECK: %[[VAL_2:.*]] = "tf.Const"() {value = dense<[1, 512, 512, 2]> : tensor<4xi32>} : () -> tensor<4xi32> -// CHECK: %[[VAL_5:.*]] = "tf.Conv2DBackpropInput"(%[[VAL_2]], %[[VAL_4]], %[[VAL_0]]) {data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 2, 2, 1], use_cudnn_on_gpu = true} : (tensor<4xi32>, tensor<4x4x2x2xf32>, tensor<1x256x256x2xf32>) -> tensor<1x512x512x2xf32> +// CHECK: %[[VAL_2:.*]] = "tf.Const"() <{value = dense<[1, 512, 512, 2]> : tensor<4xi32>}> : () -> tensor<4xi32> +// CHECK: %[[VAL_5:.*]] = "tf.Conv2DBackpropInput"(%[[VAL_2]], %[[VAL_4]], %[[VAL_0]]) <{data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 2, 2, 1], use_cudnn_on_gpu = true}> : (tensor<4xi32>, tensor<4x4x2x2xf32>, tensor<1x256x256x2xf32>) -> tensor<1x512x512x2xf32> // CHECK: return %[[VAL_5]] : tensor<1x512x512x2xf32> // CHECK: } func.func @convert_conv2d_back_prop_input_same_pad(%arg0: tensor<1x256x256x2xf32>, %arg1: tensor<4x4x2x2xf32>) -> tensor<1x512x512x2xf32> { @@ -2845,10 +2845,10 @@ func.func @convert_conv2d_back_prop_input_negative_pad(%arg0: tensor<1x256x256x2 // CHECK-LABEL: func @convert_conv2d_back_prop_input( // CHECK-SAME: %[[VAL_0:.*]]: tensor<8x4x4x32xf32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<3x3x64x32xf32>) -> tensor<8x8x8x64xf32> { -// CHECK: %[[VAL_2:.*]] = "tf.Const"() {value = dense<[0, 1]> : tensor<2xi64>} : () -> tensor<2xi64> +// CHECK: %[[VAL_2:.*]] = "tf.Const"() <{value = dense<[0, 1]> : tensor<2xi64>}> : () -> tensor<2xi64> // CHECK: %[[VAL_3:.*]] = "tf.ReverseV2"(%[[VAL_1]], %[[VAL_2]]) : (tensor<3x3x64x32xf32>, tensor<2xi64>) -> tensor<3x3x64x32xf32> -// CHECK: %[[VAL_4:.*]] = "tf.Const"() {value = dense<[8, 8, 8, 64]> : tensor<4xi32>} : () -> tensor<4xi32> -// CHECK: %[[VAL_5:.*]] = "tf.Conv2DBackpropInput"(%[[VAL_4]], %[[VAL_3]], %[[VAL_0]]) {data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 2, 2, 1], use_cudnn_on_gpu = true} : (tensor<4xi32>, tensor<3x3x64x32xf32>, tensor<8x4x4x32xf32>) -> tensor<8x8x8x64xf32> +// CHECK: %[[VAL_4:.*]] = "tf.Const"() <{value = dense<[8, 8, 8, 64]> : tensor<4xi32>}> : () -> tensor<4xi32> +// CHECK: %[[VAL_5:.*]] = "tf.Conv2DBackpropInput"(%[[VAL_4]], %[[VAL_3]], %[[VAL_0]]) <{data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 2, 2, 1], use_cudnn_on_gpu = true}> : (tensor<4xi32>, tensor<3x3x64x32xf32>, tensor<8x4x4x32xf32>) -> tensor<8x8x8x64xf32> // CHECK: return %[[VAL_5]] : tensor<8x8x8x64xf32> // CHECK: } func.func @convert_conv2d_back_prop_input(%arg0: tensor<8x4x4x32xf32>, %arg1: tensor<3x3x64x32xf32>) -> tensor<8x8x8x64xf32> { @@ -2872,12 +2872,12 @@ func.func @convert_conv2d_back_prop_input(%arg0: tensor<8x4x4x32xf32>, %arg1: te // CHECK-LABEL: func @convert_conv2d_back_prop_input_transpose_filter( // CHECK-SAME: %[[VAL_0:.*]]: tensor<8x4x4x32xf32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<3x3x32x64xf32>) -> tensor<8x8x8x64xf32> { -// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() {value = dense<[0, 1]> : tensor<2xi64>} : () -> tensor<2xi64> -// CHECK-DAG: %[[VAL_3:.*]] = "tf.Const"() {value = dense<[0, 1, 3, 2]> : tensor<4xi64>} : () -> tensor<4xi64> +// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() <{value = dense<[0, 1]> : tensor<2xi64>}> : () -> tensor<2xi64> +// CHECK-DAG: %[[VAL_3:.*]] = "tf.Const"() <{value = dense<[0, 1, 3, 2]> : tensor<4xi64>}> : () -> tensor<4xi64> // CHECK-DAG: %[[VAL_4:.*]] = "tf.Transpose"(%[[VAL_1]], %[[VAL_3]]) : (tensor<3x3x32x64xf32>, tensor<4xi64>) -> tensor<3x3x64x32xf32> // CHECK: %[[VAL_5:.*]] = "tf.ReverseV2"(%[[VAL_4]], %[[VAL_2]]) : (tensor<3x3x64x32xf32>, tensor<2xi64>) -> tensor<3x3x64x32xf32> -// CHECK: %[[VAL_6:.*]] = "tf.Const"() {value = dense<[8, 8, 8, 64]> : tensor<4xi32>} : () -> tensor<4xi32> -// CHECK: %[[VAL_7:.*]] = "tf.Conv2DBackpropInput"(%[[VAL_6]], %[[VAL_5]], %[[VAL_0]]) {data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 2, 2, 1], use_cudnn_on_gpu = true} : (tensor<4xi32>, tensor<3x3x64x32xf32>, tensor<8x4x4x32xf32>) -> tensor<8x8x8x64xf32> +// CHECK: %[[VAL_6:.*]] = "tf.Const"() <{value = dense<[8, 8, 8, 64]> : tensor<4xi32>}> : () -> tensor<4xi32> +// CHECK: %[[VAL_7:.*]] = "tf.Conv2DBackpropInput"(%[[VAL_6]], %[[VAL_5]], %[[VAL_0]]) <{data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 2, 2, 1], use_cudnn_on_gpu = true}> : (tensor<4xi32>, tensor<3x3x64x32xf32>, tensor<8x4x4x32xf32>) -> tensor<8x8x8x64xf32> // CHECK: return %[[VAL_7]] : tensor<8x8x8x64xf32> // CHECK: } func.func @convert_conv2d_back_prop_input_transpose_filter(%arg0: tensor<8x4x4x32xf32>, %arg1: tensor<3x3x32x64xf32>) -> tensor<8x8x8x64xf32> { @@ -2901,7 +2901,7 @@ func.func @convert_conv2d_back_prop_input_transpose_filter(%arg0: tensor<8x4x4x3 // CHECK-LABEL: func @convert_conv2d_valid_padding( // CHECK-SAME: %[[VAL_0:.*]]: tensor<1x8x8x207xf32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<3x3x207x16xf32>) -> tensor<1x6x6x16xf32> { -// CHECK: %[[VAL_2:.*]] = "tf.Conv2D"(%[[VAL_0]], %[[VAL_1]]) {data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "VALID", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true} : (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<1x6x6x16xf32> +// CHECK: %[[VAL_2:.*]] = "tf.Conv2D"(%[[VAL_0]], %[[VAL_1]]) <{data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "VALID", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true}> : (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<1x6x6x16xf32> // CHECK: return %[[VAL_2]] : tensor<1x6x6x16xf32> // CHECK: } func.func @convert_conv2d_valid_padding(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x6x6x16xf32> { @@ -2924,7 +2924,7 @@ func.func @convert_conv2d_valid_padding(%arg0: tensor<1x8x8x207xf32>, %arg1: ten // CHECK-LABEL: func @convert_conv2d_valid_padding_dynamic_batch( // CHECK-SAME: %[[VAL_0:.*]]: tensor, // CHECK-SAME: %[[VAL_1:.*]]: tensor<3x3x207x16xf32>) -> tensor { -// CHECK: %[[VAL_2:.*]] = "tf.Conv2D"(%[[VAL_0]], %[[VAL_1]]) {data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "VALID", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true} : (tensor, tensor<3x3x207x16xf32>) -> tensor +// CHECK: %[[VAL_2:.*]] = "tf.Conv2D"(%[[VAL_0]], %[[VAL_1]]) <{data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "VALID", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true}> : (tensor, tensor<3x3x207x16xf32>) -> tensor // CHECK: return %[[VAL_2]] : tensor // CHECK: } func.func @convert_conv2d_valid_padding_dynamic_batch(%arg0: tensor, %arg1: tensor<3x3x207x16xf32>) -> tensor { @@ -2946,9 +2946,9 @@ func.func @convert_conv2d_valid_padding_dynamic_batch(%arg0: tensor) -> tensor<1xf32> { -// CHECK-DAG: %[[VAL_1:.*]] = "tf.Const"() {value = dense<1.000000e+00> : tensor} : () -> tensor -// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi64>} : () -> tensor<1xi64> -// CHECK: %[[VAL_3:.*]] = "tf.Prod"(%[[VAL_0]], %[[VAL_2]]) {keep_dims = false} : (tensor<1x256xf32>, tensor<1xi64>) -> tensor<1xf32> +// CHECK-DAG: %[[VAL_1:.*]] = "tf.Const"() <{value = dense<1.000000e+00> : tensor}> : () -> tensor +// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi64>}> : () -> tensor<1xi64> +// CHECK: %[[VAL_3:.*]] = "tf.Prod"(%[[VAL_0]], %[[VAL_2]]) <{keep_dims = false}> : (tensor<1x256xf32>, tensor<1xi64>) -> tensor<1xf32> // CHECK: return %[[VAL_3]] : tensor<1xf32> // CHECK: } func.func @convert_reduce_to_prod(%arg0: tensor<1x256xf32>) -> tensor<1xf32> { @@ -2963,9 +2963,9 @@ func.func @convert_reduce_to_prod(%arg0: tensor<1x256xf32>) -> tensor<1xf32> { // CHECK-LABEL: func @convert_reduce_to_sum( // CHECK-SAME: %[[VAL_0:.*]]: tensor<1x256xf32>) -> tensor<1xf32> { -// CHECK-DAG: %[[VAL_1:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor} : () -> tensor -// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi64>} : () -> tensor<1xi64> -// CHECK: %[[VAL_3:.*]] = "tf.Sum"(%[[VAL_0]], %[[VAL_2]]) {keep_dims = false} : (tensor<1x256xf32>, tensor<1xi64>) -> tensor<1xf32> +// CHECK-DAG: %[[VAL_1:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor}> : () -> tensor +// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi64>}> : () -> tensor<1xi64> +// CHECK: %[[VAL_3:.*]] = "tf.Sum"(%[[VAL_0]], %[[VAL_2]]) <{keep_dims = false}> : (tensor<1x256xf32>, tensor<1xi64>) -> tensor<1xf32> // CHECK: return %[[VAL_3]] : tensor<1xf32> // CHECK: } func.func @convert_reduce_to_sum(%arg0: tensor<1x256xf32>) -> tensor<1xf32> { @@ -2981,8 +2981,8 @@ func.func @convert_reduce_to_sum(%arg0: tensor<1x256xf32>) -> tensor<1xf32> { // CHECK-LABEL: func @convert_reduce_to_prod_non_constant_init( // CHECK-SAME: %[[ARG_0:.*]]: tensor<1x256xf32>, // CHECK-SAME: %[[ARG_1:.*]]: tensor) -> tensor<1xf32> { -// CHECK-DAG: %[[VAL_0:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi64>} : () -> tensor<1xi64> -// CHECK: %[[VAL_1:.*]] = "tf.Prod"(%[[ARG_0]], %[[VAL_0]]) {keep_dims = false} : (tensor<1x256xf32>, tensor<1xi64>) -> tensor<1xf32> +// CHECK-DAG: %[[VAL_0:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi64>}> : () -> tensor<1xi64> +// CHECK: %[[VAL_1:.*]] = "tf.Prod"(%[[ARG_0]], %[[VAL_0]]) <{keep_dims = false}> : (tensor<1x256xf32>, tensor<1xi64>) -> tensor<1xf32> // CHECK: %[[VAL_2:.*]] = "tf.Mul"(%[[VAL_1]], %[[ARG_1]]) : (tensor<1xf32>, tensor) -> tensor<1xf32> // CHECK: return %[[VAL_2]] : tensor<1xf32> // CHECK: } @@ -2999,8 +2999,8 @@ func.func @convert_reduce_to_prod_non_constant_init(%arg0: tensor<1x256xf32>, %a // CHECK-LABEL: func @convert_reduce_to_sum_non_constant_init( // CHECK-SAME: %[[ARG_0:.*]]: tensor<1x256xf32>, // CHECK-SAME: %[[ARG_1:.*]]: tensor) -> tensor<1xf32> { -// CHECK-DAG: %[[VAL_0:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi64>} : () -> tensor<1xi64> -// CHECK: %[[VAL_1:.*]] = "tf.Sum"(%[[ARG_0]], %[[VAL_0]]) {keep_dims = false} : (tensor<1x256xf32>, tensor<1xi64>) -> tensor<1xf32> +// CHECK-DAG: %[[VAL_0:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi64>}> : () -> tensor<1xi64> +// CHECK: %[[VAL_1:.*]] = "tf.Sum"(%[[ARG_0]], %[[VAL_0]]) <{keep_dims = false}> : (tensor<1x256xf32>, tensor<1xi64>) -> tensor<1xf32> // CHECK: %[[VAL_2:.*]] = "tf.Add"(%[[VAL_1]], %[[ARG_1]]) : (tensor<1xf32>, tensor) -> tensor<1xf32> // CHECK: return %[[VAL_2]] : tensor<1xf32> // CHECK: } @@ -3015,9 +3015,9 @@ func.func @convert_reduce_to_sum_non_constant_init(%arg0: tensor<1x256xf32>, %ar // CHECK-LABEL: func @convert_int_reduce_to_prod( // CHECK-SAME: %[[VAL_0:.*]]: tensor<1x256xi32>) -> tensor<1xi32> { -// CHECK-DAG: %[[VAL_1:.*]] = "tf.Const"() {value = dense<1> : tensor} : () -> tensor -// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi64>} : () -> tensor<1xi64> -// CHECK: %[[VAL_3:.*]] = "tf.Prod"(%[[VAL_0]], %[[VAL_2]]) {keep_dims = false} : (tensor<1x256xi32>, tensor<1xi64>) -> tensor<1xi32> +// CHECK-DAG: %[[VAL_1:.*]] = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor +// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi64>}> : () -> tensor<1xi64> +// CHECK: %[[VAL_3:.*]] = "tf.Prod"(%[[VAL_0]], %[[VAL_2]]) <{keep_dims = false}> : (tensor<1x256xi32>, tensor<1xi64>) -> tensor<1xi32> // CHECK: return %[[VAL_3]] : tensor<1xi32> // CHECK: } func.func @convert_int_reduce_to_prod(%arg0: tensor<1x256xi32>) -> tensor<1xi32> { @@ -3033,9 +3033,9 @@ func.func @convert_int_reduce_to_prod(%arg0: tensor<1x256xi32>) -> tensor<1xi32> // CHECK-LABEL: func @convert_int_reduce_to_sum( // CHECK-SAME: %[[VAL_0:.*]]: tensor<1x256xi32>) -> tensor<1xi32> { -// CHECK-DAG: %[[VAL_1:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor -// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi64>} : () -> tensor<1xi64> -// CHECK: %[[VAL_3:.*]] = "tf.Sum"(%[[VAL_0]], %[[VAL_2]]) {keep_dims = false} : (tensor<1x256xi32>, tensor<1xi64>) -> tensor<1xi32> +// CHECK-DAG: %[[VAL_1:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor +// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi64>}> : () -> tensor<1xi64> +// CHECK: %[[VAL_3:.*]] = "tf.Sum"(%[[VAL_0]], %[[VAL_2]]) <{keep_dims = false}> : (tensor<1x256xi32>, tensor<1xi64>) -> tensor<1xi32> // CHECK: return %[[VAL_3]] : tensor<1xi32> // CHECK: } func.func @convert_int_reduce_to_sum(%arg0: tensor<1x256xi32>) -> tensor<1xi32> { @@ -3050,8 +3050,8 @@ func.func @convert_int_reduce_to_sum(%arg0: tensor<1x256xi32>) -> tensor<1xi32> // CHECK-LABEL: func @convert_reduce_to_max( // CHECK-SAME: %[[VAL_0:.*]]: tensor<1x256xf32>) -> tensor<1xf32> { -// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi64>} : () -> tensor<1xi64> -// CHECK: %[[VAL_3:.*]] = "tf.Max"(%[[VAL_0]], %[[VAL_2]]) {keep_dims = false} : (tensor<1x256xf32>, tensor<1xi64>) -> tensor<1xf32> +// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi64>}> : () -> tensor<1xi64> +// CHECK: %[[VAL_3:.*]] = "tf.Max"(%[[VAL_0]], %[[VAL_2]]) <{keep_dims = false}> : (tensor<1x256xf32>, tensor<1xi64>) -> tensor<1xf32> // CHECK: return %[[VAL_3]] : tensor<1xf32> // CHECK: } func.func @convert_reduce_to_max(%arg0: tensor<1x256xf32>) -> tensor<1xf32> { @@ -3067,8 +3067,8 @@ func.func @convert_reduce_to_max(%arg0: tensor<1x256xf32>) -> tensor<1xf32> { // CHECK-LABEL: func @convert_reduce_to_max_int( // CHECK-SAME: %[[VAL_0:.*]]: tensor<1x4xi32>) -> tensor<1xi32> { -// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi64>} : () -> tensor<1xi64> -// CHECK: %[[VAL_3:.*]] = "tf.Max"(%[[VAL_0]], %[[VAL_2]]) {keep_dims = false} : (tensor<1x4xi32>, tensor<1xi64>) -> tensor<1xi32> +// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi64>}> : () -> tensor<1xi64> +// CHECK: %[[VAL_3:.*]] = "tf.Max"(%[[VAL_0]], %[[VAL_2]]) <{keep_dims = false}> : (tensor<1x4xi32>, tensor<1xi64>) -> tensor<1xi32> // CHECK: return %[[VAL_3]] : tensor<1xi32> func.func @convert_reduce_to_max_int(%arg0: tensor<1x4xi32>) -> tensor<1xi32> { // -2147483648 is MIN for INT32 @@ -3083,8 +3083,8 @@ func.func @convert_reduce_to_max_int(%arg0: tensor<1x4xi32>) -> tensor<1xi32> { // CHECK-LABEL: func @convert_reduce_to_min( // CHECK-SAME: %[[VAL_0:.*]]: tensor<1x256xf32>) -> tensor<1xf32> { -// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi64>} : () -> tensor<1xi64> -// CHECK: %[[VAL_3:.*]] = "tf.Min"(%[[VAL_0]], %[[VAL_2]]) {keep_dims = false} : (tensor<1x256xf32>, tensor<1xi64>) -> tensor<1xf32> +// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi64>}> : () -> tensor<1xi64> +// CHECK: %[[VAL_3:.*]] = "tf.Min"(%[[VAL_0]], %[[VAL_2]]) <{keep_dims = false}> : (tensor<1x256xf32>, tensor<1xi64>) -> tensor<1xf32> // CHECK: return %[[VAL_3]] : tensor<1xf32> // CHECK: } func.func @convert_reduce_to_min(%arg0: tensor<1x256xf32>) -> tensor<1xf32> { @@ -3100,8 +3100,8 @@ func.func @convert_reduce_to_min(%arg0: tensor<1x256xf32>) -> tensor<1xf32> { // CHECK-LABEL: func @convert_reduce_to_min_int( // CHECK-SAME: %[[VAL_0:.*]]: tensor<1x4xi32>) -> tensor<1xi32> { -// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi64>} : () -> tensor<1xi64> -// CHECK: %[[VAL_3:.*]] = "tf.Min"(%[[VAL_0]], %[[VAL_2]]) {keep_dims = false} : (tensor<1x4xi32>, tensor<1xi64>) -> tensor<1xi32> +// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi64>}> : () -> tensor<1xi64> +// CHECK: %[[VAL_3:.*]] = "tf.Min"(%[[VAL_0]], %[[VAL_2]]) <{keep_dims = false}> : (tensor<1x4xi32>, tensor<1xi64>) -> tensor<1xi32> // CHECK: return %[[VAL_3]] : tensor<1xi32> func.func @convert_reduce_to_min_int(%arg0: tensor<1x4xi32>) -> tensor<1xi32> { // 2147483647 is MAX for INT32 @@ -3115,9 +3115,9 @@ func.func @convert_reduce_to_min_int(%arg0: tensor<1x4xi32>) -> tensor<1xi32> { } // CHECK-LABEL: func @convert_iota_1d() -> tensor<123xf32> { -// CHECK-DAG: %[[VAL_0:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor} : () -> tensor -// CHECK-DAG: %[[VAL_1:.*]] = "tf.Const"() {value = dense<1.230000e+02> : tensor} : () -> tensor -// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() {value = dense<1.000000e+00> : tensor} : () -> tensor +// CHECK-DAG: %[[VAL_0:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor}> : () -> tensor +// CHECK-DAG: %[[VAL_1:.*]] = "tf.Const"() <{value = dense<1.230000e+02> : tensor}> : () -> tensor +// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() <{value = dense<1.000000e+00> : tensor}> : () -> tensor // CHECK: %[[VAL_3:.*]] = "tf.Range"(%[[VAL_0]], %[[VAL_1]], %[[VAL_2]]) : (tensor, tensor, tensor) -> tensor<123xf32> // CHECK: return %[[VAL_3]] : tensor<123xf32> // CHECK: } @@ -3127,13 +3127,13 @@ func.func @convert_iota_1d() -> tensor<123xf32> { } // CHECK-LABEL: func @convert_iota_3d() -> tensor<5x7x9xi32> { -// CHECK-DAG: %[[VAL_0:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor -// CHECK-DAG: %[[VAL_1:.*]] = "tf.Const"() {value = dense<7> : tensor} : () -> tensor -// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() {value = dense<1> : tensor} : () -> tensor +// CHECK-DAG: %[[VAL_0:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor +// CHECK-DAG: %[[VAL_1:.*]] = "tf.Const"() <{value = dense<7> : tensor}> : () -> tensor +// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor // CHECK: %[[VAL_3:.*]] = "tf.Range"(%[[VAL_0]], %[[VAL_1]], %[[VAL_2]]) : (tensor, tensor, tensor) -> tensor<7xi32> -// CHECK: %[[VAL_4:.*]] = "tf.Const"() {value = dense<[1, 7, 1]> : tensor<3xi64>} : () -> tensor<3xi64> +// CHECK: %[[VAL_4:.*]] = "tf.Const"() <{value = dense<[1, 7, 1]> : tensor<3xi64>}> : () -> tensor<3xi64> // CHECK: %[[VAL_5:.*]] = "tf.Reshape"(%[[VAL_3]], %[[VAL_4]]) : (tensor<7xi32>, tensor<3xi64>) -> tensor<1x7x1xi32> -// CHECK: %[[VAL_6:.*]] = "tf.Const"() {value = dense<[5, 7, 9]> : tensor<3xi64>} : () -> tensor<3xi64> +// CHECK: %[[VAL_6:.*]] = "tf.Const"() <{value = dense<[5, 7, 9]> : tensor<3xi64>}> : () -> tensor<3xi64> // CHECK: %[[VAL_7:.*]] = "tf.BroadcastTo"(%[[VAL_5]], %[[VAL_6]]) : (tensor<1x7x1xi32>, tensor<3xi64>) -> tensor<5x7x9xi32> // CHECK: return %[[VAL_7]] : tensor<5x7x9xi32> // CHECK: } @@ -3144,7 +3144,7 @@ func.func @convert_iota_3d() -> tensor<5x7x9xi32> { // CHECK-LABEL: func @convert_avgpool_valid( // CHECK-SAME: %[[VAL_0:.*]]: tensor<4x16x16x8xf32>) -> tensor<4x7x7x8xf32> { -// CHECK: %[[VAL_1:.*]] = "tf.AvgPool"(%[[VAL_0]]) {data_format = "NHWC", ksize = [1, 3, 3, 1], padding = "VALID", strides = [1, 2, 2, 1]} : (tensor<4x16x16x8xf32>) -> tensor<4x7x7x8xf32> +// CHECK: %[[VAL_1:.*]] = "tf.AvgPool"(%[[VAL_0]]) <{data_format = "NHWC", ksize = [1, 3, 3, 1], padding = "VALID", strides = [1, 2, 2, 1]}> : (tensor<4x16x16x8xf32>) -> tensor<4x7x7x8xf32> // CHECK: return %[[VAL_1]] : tensor<4x7x7x8xf32> // CHECK: } func.func @convert_avgpool_valid(%arg0: tensor<4x16x16x8xf32>) -> tensor<4x7x7x8xf32> { @@ -3166,7 +3166,7 @@ func.func @convert_avgpool_valid(%arg0: tensor<4x16x16x8xf32>) -> tensor<4x7x7x8 // CHECK-LABEL: func @convert_avgpool_valid_broadcasted_divisor( // CHECK-SAME: %[[VAL_0:.*]]: tensor<4x16x16x8xf32>) -> tensor<4x7x7x8xf32> { -// CHECK: %[[VAL_1:.*]] = "tf.AvgPool"(%[[VAL_0]]) {data_format = "NHWC", ksize = [1, 3, 3, 1], padding = "VALID", strides = [1, 2, 2, 1]} : (tensor<4x16x16x8xf32>) -> tensor<4x7x7x8xf32> +// CHECK: %[[VAL_1:.*]] = "tf.AvgPool"(%[[VAL_0]]) <{data_format = "NHWC", ksize = [1, 3, 3, 1], padding = "VALID", strides = [1, 2, 2, 1]}> : (tensor<4x16x16x8xf32>) -> tensor<4x7x7x8xf32> // CHECK: return %[[VAL_1]] : tensor<4x7x7x8xf32> // CHECK: } func.func @convert_avgpool_valid_broadcasted_divisor(%arg0: tensor<4x16x16x8xf32>) -> tensor<4x7x7x8xf32> { @@ -3189,7 +3189,7 @@ func.func @convert_avgpool_valid_broadcasted_divisor(%arg0: tensor<4x16x16x8xf32 // CHECK-LABEL: func @convert_avgpool_valid_channel_first( // CHECK-SAME: %[[VAL_0:.*]]: tensor<4x3x16x16xf32>) -> tensor<4x3x7x7xf32> { -// CHECK: %[[VAL_1:.*]] = "tf.AvgPool"(%[[VAL_0]]) {data_format = "NCHW", ksize = [1, 1, 3, 3], padding = "VALID", strides = [1, 1, 2, 2]} : (tensor<4x3x16x16xf32>) -> tensor<4x3x7x7xf32> +// CHECK: %[[VAL_1:.*]] = "tf.AvgPool"(%[[VAL_0]]) <{data_format = "NCHW", ksize = [1, 1, 3, 3], padding = "VALID", strides = [1, 1, 2, 2]}> : (tensor<4x3x16x16xf32>) -> tensor<4x3x7x7xf32> // CHECK: return %[[VAL_1]] : tensor<4x3x7x7xf32> // CHECK: } func.func @convert_avgpool_valid_channel_first(%arg0: tensor<4x3x16x16xf32>) -> tensor<4x3x7x7xf32> { @@ -3211,7 +3211,7 @@ func.func @convert_avgpool_valid_channel_first(%arg0: tensor<4x3x16x16xf32>) -> // CHECK-LABEL: func @convert_avgpool_valid_rw( // CHECK-SAME: %[[VAL_0:.*]]: tensor<4x16x16x8xf32>) -> tensor<4x7x7x8xf32> { -// CHECK: %[[VAL_1:.*]] = "tf.AvgPool"(%[[VAL_0]]) {data_format = "NHWC", ksize = [1, 3, 3, 1], padding = "VALID", strides = [1, 2, 2, 1]} : (tensor<4x16x16x8xf32>) -> tensor<4x7x7x8xf32> +// CHECK: %[[VAL_1:.*]] = "tf.AvgPool"(%[[VAL_0]]) <{data_format = "NHWC", ksize = [1, 3, 3, 1], padding = "VALID", strides = [1, 2, 2, 1]}> : (tensor<4x16x16x8xf32>) -> tensor<4x7x7x8xf32> // CHECK: return %[[VAL_1]] : tensor<4x7x7x8xf32> // CHECK: } func.func @convert_avgpool_valid_rw(%arg0: tensor<4x16x16x8xf32>) -> tensor<4x7x7x8xf32> { @@ -3243,7 +3243,7 @@ func.func @convert_avgpool_valid_rw(%arg0: tensor<4x16x16x8xf32>) -> tensor<4x7x // CHECK-LABEL: func @convert_avgpool_valid_rw_broadcasted_const_lhs( // CHECK-SAME: %[[VAL_0:.*]]: tensor<4x16x16x8xf32>) -> tensor<4x7x7x8xf32> { -// CHECK: %[[VAL_1:.*]] = "tf.AvgPool"(%[[VAL_0]]) {data_format = "NHWC", ksize = [1, 3, 3, 1], padding = "VALID", strides = [1, 2, 2, 1]} : (tensor<4x16x16x8xf32>) -> tensor<4x7x7x8xf32> +// CHECK: %[[VAL_1:.*]] = "tf.AvgPool"(%[[VAL_0]]) <{data_format = "NHWC", ksize = [1, 3, 3, 1], padding = "VALID", strides = [1, 2, 2, 1]}> : (tensor<4x16x16x8xf32>) -> tensor<4x7x7x8xf32> // CHECK: return %[[VAL_1]] : tensor<4x7x7x8xf32> // CHECK: } func.func @convert_avgpool_valid_rw_broadcasted_const_lhs(%arg0: tensor<4x16x16x8xf32>) -> tensor<4x7x7x8xf32> { @@ -3276,7 +3276,7 @@ func.func @convert_avgpool_valid_rw_broadcasted_const_lhs(%arg0: tensor<4x16x16x // CHECK-LABEL: func @convert_avgpool_valid_3d( // CHECK-SAME: %[[VAL_0:.*]]: tensor<4x16x16x16x8xf32>) -> tensor<4x7x7x7x8xf32> { -// CHECK: %[[VAL_1:.*]] = "tf.AvgPool3D"(%[[VAL_0]]) {data_format = "NDHWC", ksize = [1, 3, 3, 3, 1], padding = "VALID", strides = [1, 2, 2, 2, 1]} : (tensor<4x16x16x16x8xf32>) -> tensor<4x7x7x7x8xf32> +// CHECK: %[[VAL_1:.*]] = "tf.AvgPool3D"(%[[VAL_0]]) <{data_format = "NDHWC", ksize = [1, 3, 3, 3, 1], padding = "VALID", strides = [1, 2, 2, 2, 1]}> : (tensor<4x16x16x16x8xf32>) -> tensor<4x7x7x7x8xf32> // CHECK: return %[[VAL_1]] : tensor<4x7x7x7x8xf32> // CHECK: } func.func @convert_avgpool_valid_3d(%arg0: tensor<4x16x16x16x8xf32>) -> tensor<4x7x7x7x8xf32> { @@ -3298,7 +3298,7 @@ func.func @convert_avgpool_valid_3d(%arg0: tensor<4x16x16x16x8xf32>) -> tensor<4 // CHECK-LABEL: func @convert_avgpool_valid_3d_channel_first( // CHECK-SAME: %[[VAL_0:.*]]: tensor<4x3x16x16x16xf32>) -> tensor<4x3x7x7x7xf32> { -// CHECK: %[[VAL_1:.*]] = "tf.AvgPool3D"(%[[VAL_0]]) {data_format = "NCDHW", ksize = [1, 1, 3, 3, 3], padding = "VALID", strides = [1, 1, 2, 2, 2]} : (tensor<4x3x16x16x16xf32>) -> tensor<4x3x7x7x7xf32> +// CHECK: %[[VAL_1:.*]] = "tf.AvgPool3D"(%[[VAL_0]]) <{data_format = "NCDHW", ksize = [1, 1, 3, 3, 3], padding = "VALID", strides = [1, 1, 2, 2, 2]}> : (tensor<4x3x16x16x16xf32>) -> tensor<4x3x7x7x7xf32> // CHECK: return %[[VAL_1]] : tensor<4x3x7x7x7xf32> // CHECK: } func.func @convert_avgpool_valid_3d_channel_first(%arg0: tensor<4x3x16x16x16xf32>) -> tensor<4x3x7x7x7xf32> { @@ -3320,7 +3320,7 @@ func.func @convert_avgpool_valid_3d_channel_first(%arg0: tensor<4x3x16x16x16xf32 // CHECK-LABEL: func @convert_avgpool_same( // CHECK-SAME: %[[VAL_0:.*]]: tensor<4x16x16x8xf32>) -> tensor<4x8x8x8xf32> { -// CHECK: %[[VAL_1:.*]] = "tf.AvgPool"(%[[VAL_0]]) {data_format = "NHWC", ksize = [1, 3, 3, 1], padding = "SAME", strides = [1, 2, 2, 1]} : (tensor<4x16x16x8xf32>) -> tensor<4x8x8x8xf32> +// CHECK: %[[VAL_1:.*]] = "tf.AvgPool"(%[[VAL_0]]) <{data_format = "NHWC", ksize = [1, 3, 3, 1], padding = "SAME", strides = [1, 2, 2, 1]}> : (tensor<4x16x16x8xf32>) -> tensor<4x8x8x8xf32> // CHECK: return %[[VAL_1]] : tensor<4x8x8x8xf32> // CHECK: } func.func @convert_avgpool_same(%arg0: tensor<4x16x16x8xf32>) -> tensor<4x8x8x8xf32> { @@ -3352,7 +3352,7 @@ func.func @convert_avgpool_same(%arg0: tensor<4x16x16x8xf32>) -> tensor<4x8x8x8x // CHECK-LABEL: func @convert_avgpool_reshape_broadcast( // CHECK-SAME: %[[VAL_0:.*]]: tensor<4x16x16x8xf32>) -> tensor<4x8x8x8xf32> { -// CHECK: %[[VAL_1:.*]] = "tf.AvgPool"(%[[VAL_0]]) {data_format = "NHWC", ksize = [1, 3, 3, 1], padding = "SAME", strides = [1, 2, 2, 1]} : (tensor<4x16x16x8xf32>) -> tensor<4x8x8x8xf32> +// CHECK: %[[VAL_1:.*]] = "tf.AvgPool"(%[[VAL_0]]) <{data_format = "NHWC", ksize = [1, 3, 3, 1], padding = "SAME", strides = [1, 2, 2, 1]}> : (tensor<4x16x16x8xf32>) -> tensor<4x8x8x8xf32> // CHECK: return %[[VAL_1]] : tensor<4x8x8x8xf32> // CHECK: } func.func @convert_avgpool_reshape_broadcast(%arg0: tensor<4x16x16x8xf32>) -> tensor<4x8x8x8xf32> { @@ -3376,7 +3376,7 @@ func.func @convert_avgpool_reshape_broadcast(%arg0: tensor<4x16x16x8xf32>) -> te // CHECK-LABEL: func @convert_maxpool_valid( // CHECK-SAME: %[[VAL_0:.*]]: tensor<4x16x16x8xf32>) -> tensor<4x7x7x8xf32> { -// CHECK: %[[VAL_1:.*]] = "tf.MaxPool"(%[[VAL_0]]) {data_format = "NHWC", explicit_paddings = [], ksize = [1, 3, 3, 1], padding = "VALID", strides = [1, 2, 2, 1]} : (tensor<4x16x16x8xf32>) -> tensor<4x7x7x8xf32> +// CHECK: %[[VAL_1:.*]] = "tf.MaxPool"(%[[VAL_0]]) <{data_format = "NHWC", explicit_paddings = [], ksize = [1, 3, 3, 1], padding = "VALID", strides = [1, 2, 2, 1]}> : (tensor<4x16x16x8xf32>) -> tensor<4x7x7x8xf32> // CHECK: return %[[VAL_1]] : tensor<4x7x7x8xf32> // CHECK: } func.func @convert_maxpool_valid(%arg0: tensor<4x16x16x8xf32>) -> tensor<4x7x7x8xf32> { @@ -3397,7 +3397,7 @@ func.func @convert_maxpool_valid(%arg0: tensor<4x16x16x8xf32>) -> tensor<4x7x7x8 // CHECK-LABEL: func @convert_maxpool_valid_channel_first( // CHECK-SAME: %[[VAL_0:.*]]: tensor<4x3x16x16xf32>) -> tensor<4x3x7x7xf32> { -// CHECK: %[[VAL_1:.*]] = "tf.MaxPool"(%[[VAL_0]]) {data_format = "NCHW", explicit_paddings = [], ksize = [1, 1, 3, 3], padding = "VALID", strides = [1, 1, 2, 2]} : (tensor<4x3x16x16xf32>) -> tensor<4x3x7x7xf32> +// CHECK: %[[VAL_1:.*]] = "tf.MaxPool"(%[[VAL_0]]) <{data_format = "NCHW", explicit_paddings = [], ksize = [1, 1, 3, 3], padding = "VALID", strides = [1, 1, 2, 2]}> : (tensor<4x3x16x16xf32>) -> tensor<4x3x7x7xf32> // CHECK: return %[[VAL_1]] : tensor<4x3x7x7xf32> // CHECK: } func.func @convert_maxpool_valid_channel_first(%arg0: tensor<4x3x16x16xf32>) -> tensor<4x3x7x7xf32> { @@ -3418,7 +3418,7 @@ func.func @convert_maxpool_valid_channel_first(%arg0: tensor<4x3x16x16xf32>) -> // CHECK-LABEL: func @convert_maxpool_valid_3d( // CHECK-SAME: %[[VAL_0:.*]]: tensor<4x16x16x16x8xf32>) -> tensor<4x7x7x7x8xf32> { -// CHECK: %[[VAL_1:.*]] = "tf.MaxPool3D"(%[[VAL_0]]) {data_format = "NDHWC", ksize = [1, 3, 3, 3, 1], padding = "VALID", strides = [1, 2, 2, 2, 1]} : (tensor<4x16x16x16x8xf32>) -> tensor<4x7x7x7x8xf32> +// CHECK: %[[VAL_1:.*]] = "tf.MaxPool3D"(%[[VAL_0]]) <{data_format = "NDHWC", ksize = [1, 3, 3, 3, 1], padding = "VALID", strides = [1, 2, 2, 2, 1]}> : (tensor<4x16x16x16x8xf32>) -> tensor<4x7x7x7x8xf32> // CHECK: return %[[VAL_1]] : tensor<4x7x7x7x8xf32> // CHECK: } func.func @convert_maxpool_valid_3d(%arg0: tensor<4x16x16x16x8xf32>) -> tensor<4x7x7x7x8xf32> { @@ -3439,7 +3439,7 @@ func.func @convert_maxpool_valid_3d(%arg0: tensor<4x16x16x16x8xf32>) -> tensor<4 // CHECK-LABEL: func @convert_maxpool_valid_3d_channel_first( // CHECK-SAME: %[[VAL_0:.*]]: tensor<4x3x16x16x16xf32>) -> tensor<4x3x7x7x7xf32> { -// CHECK: %[[VAL_1:.*]] = "tf.MaxPool3D"(%[[VAL_0]]) {data_format = "NCDHW", ksize = [1, 1, 3, 3, 3], padding = "VALID", strides = [1, 1, 2, 2, 2]} : (tensor<4x3x16x16x16xf32>) -> tensor<4x3x7x7x7xf32> +// CHECK: %[[VAL_1:.*]] = "tf.MaxPool3D"(%[[VAL_0]]) <{data_format = "NCDHW", ksize = [1, 1, 3, 3, 3], padding = "VALID", strides = [1, 1, 2, 2, 2]}> : (tensor<4x3x16x16x16xf32>) -> tensor<4x3x7x7x7xf32> // CHECK: return %[[VAL_1]] : tensor<4x3x7x7x7xf32> // CHECK: } func.func @convert_maxpool_valid_3d_channel_first(%arg0: tensor<4x3x16x16x16xf32>) -> tensor<4x3x7x7x7xf32> { @@ -3460,7 +3460,7 @@ func.func @convert_maxpool_valid_3d_channel_first(%arg0: tensor<4x3x16x16x16xf32 // CHECK-LABEL: func @convert_maxpool_same( // CHECK-SAME: %[[VAL_0:.*]]: tensor<4x16x16x8xf32>) -> tensor<4x8x8x8xf32> { -// CHECK: %[[VAL_1:.*]] = "tf.MaxPool"(%[[VAL_0]]) {data_format = "NHWC", explicit_paddings = [], ksize = [1, 3, 3, 1], padding = "SAME", strides = [1, 2, 2, 1]} : (tensor<4x16x16x8xf32>) -> tensor<4x8x8x8xf32> +// CHECK: %[[VAL_1:.*]] = "tf.MaxPool"(%[[VAL_0]]) <{data_format = "NHWC", explicit_paddings = [], ksize = [1, 3, 3, 1], padding = "SAME", strides = [1, 2, 2, 1]}> : (tensor<4x16x16x8xf32>) -> tensor<4x8x8x8xf32> // CHECK: return %[[VAL_1]] : tensor<4x8x8x8xf32> // CHECK: } func.func @convert_maxpool_same(%arg0: tensor<4x16x16x8xf32>) -> tensor<4x8x8x8xf32> { @@ -3574,8 +3574,8 @@ func.func @convert_floor_mod_int(%arg0: tensor<192x8xi32>, %arg1: tensor<192x8xi } // CHECK-LABEL: func @convert_floor_mod_float_cst -// CHECK-DAG: %[[CST1:.*]] = "tf.Const"() {value = dense<2.000000e+00> : tensor<192x8xbf16>} : () -> tensor<192x8xbf16> -// CHECK-DAG: %[[CST2:.*]] = "tf.Const"() {value = dense<2.000000e+00> : tensor<192x8xbf16>} : () -> tensor<192x8xbf16> +// CHECK-DAG: %[[CST1:.*]] = "tf.Const"() <{value = dense<2.000000e+00> : tensor<192x8xbf16>}> : () -> tensor<192x8xbf16> +// CHECK-DAG: %[[CST2:.*]] = "tf.Const"() <{value = dense<2.000000e+00> : tensor<192x8xbf16>}> : () -> tensor<192x8xbf16> // CHECK: %[[RESULT:.*]] = "tf.FloorMod"(%arg0, %[[CST2]]) : (tensor<192x8xbf16>, tensor<192x8xbf16>) -> tensor<192x8xbf16> // CHECK: return %[[RESULT]] : tensor<192x8xbf16> // CHECK: } @@ -3592,8 +3592,8 @@ func.func @convert_floor_mod_float_cst(%arg0: tensor<192x8xbf16>) -> tensor<192x } // CHECK-LABEL: func @convert_floor_mod_int_cst -// CHECK-DAG: %[[CST1:.*]] = "tf.Const"() {value = dense<2> : tensor<192x8xi32>} : () -> tensor<192x8xi32> -// CHECK-DAG: %[[CST2:.*]] = "tf.Const"() {value = dense<2> : tensor<192x8xi32>} : () -> tensor<192x8xi32> +// CHECK-DAG: %[[CST1:.*]] = "tf.Const"() <{value = dense<2> : tensor<192x8xi32>}> : () -> tensor<192x8xi32> +// CHECK-DAG: %[[CST2:.*]] = "tf.Const"() <{value = dense<2> : tensor<192x8xi32>}> : () -> tensor<192x8xi32> // CHECK: %[[RESULT:.*]] = "tf.FloorMod"(%arg0, %[[CST2]]) : (tensor<192x8xi32>, tensor<192x8xi32>) -> tensor<192x8xi32> // CHECK: return %[[RESULT]] : tensor<192x8xi32> // CHECK: } @@ -3649,7 +3649,7 @@ func.func @convert_floor_div(%arg0: tensor<10x10xbf16>, %arg1: tensor<10x10xbf16 } // CHECK-LABEL: func @convert_floor_div_cst -// CHECK: %[[CST2:.*]] = "tf.Const"() {value = dense<2.000000e+00> : tensor<10x10xbf16>} : () -> tensor<10x10xbf16> +// CHECK: %[[CST2:.*]] = "tf.Const"() <{value = dense<2.000000e+00> : tensor<10x10xbf16>}> : () -> tensor<10x10xbf16> // CHECK: %[[RESULT:.*]] = "tf.FloorDiv"(%arg0, %[[CST2]]) : (tensor<10x10xbf16>, tensor<10x10xbf16>) -> tensor<10x10xbf16> // CHECK: return %[[RESULT]] // CHECK: } @@ -3674,7 +3674,7 @@ func.func @convert_floor_div_cst(%arg0: tensor<10x10xbf16>) -> tensor<10x10xbf16 } // CHECK-LABEL: func @convert_floor_div_cst2 -// CHECK: %[[CST2:.*]] = "tf.Const"() {value = dense<2.000000e+00> : tensor<10x10xbf16>} : () -> tensor<10x10xbf16> +// CHECK: %[[CST2:.*]] = "tf.Const"() <{value = dense<2.000000e+00> : tensor<10x10xbf16>}> : () -> tensor<10x10xbf16> // CHECK: %[[RESULT:.*]] = "tf.FloorDiv"(%arg0, %[[CST2]]) : (tensor<10x10xbf16>, tensor<10x10xbf16>) -> tensor<10x10xbf16> // CHECK: return %[[RESULT]] // CHECK: } @@ -3791,10 +3791,10 @@ func.func @convert_gather_transpose(%arg0: tensor<128x256xf32>, %arg1: tensor<4x // CHECK-LABEL: func @convert_gather_offset( // CHECK-SAME: %[[VAL_0:.*]]: tensor<1x20xi32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<1x1xi32>) -> tensor<1x1xi32> { -// CHECK: %[[VAL_2:.*]] = "tf.Const"() {value = dense<[1, 0]> : tensor<2xi64>} : () -> tensor<2xi64> +// CHECK: %[[VAL_2:.*]] = "tf.Const"() <{value = dense<[1, 0]> : tensor<2xi64>}> : () -> tensor<2xi64> // CHECK: %[[VAL_3:.*]] = "tf.Transpose"(%[[VAL_0]], %[[VAL_2]]) : (tensor<1x20xi32>, tensor<2xi64>) -> tensor<20x1xi32> // CHECK: %[[VAL_4:.*]] = "tf.GatherNd"(%[[VAL_3]], %[[VAL_1]]) : (tensor<20x1xi32>, tensor<1x1xi32>) -> tensor<1x1xi32> -// CHECK: %[[VAL_5:.*]] = "tf.Const"() {value = dense<[1, 0]> : tensor<2xi64>} : () -> tensor<2xi64> +// CHECK: %[[VAL_5:.*]] = "tf.Const"() <{value = dense<[1, 0]> : tensor<2xi64>}> : () -> tensor<2xi64> // CHECK: %[[VAL_6:.*]] = "tf.Transpose"(%[[VAL_4]], %[[VAL_5]]) : (tensor<1x1xi32>, tensor<2xi64>) -> tensor<1x1xi32> // CHECK: return %[[VAL_6]] : tensor<1x1xi32> // CHECK: } @@ -3815,12 +3815,12 @@ func.func @convert_gather_offset(%arg0: tensor<1x20xi32>, %arg1: tensor<1x1xi32> // CHECK-LABEL: func @convert_gather_to_slice_batch_size_1( // CHECK-SAME: %[[ARG_0:.*]]: tensor<1x2944xi32>, // CHECK-SAME: %[[ARG_1:.*]]: tensor<1x2xi32>) -// CHECK-DAG: %[[CST:.*]] = "tf.Const"() {value = dense<[0, 1440]> : tensor<2xi32>} : () -> tensor<2xi32> -// CHECK-DAG: %[[CST_0:.*]] = "tf.Const"() {value = dense<0> : tensor<2xi32>} : () -> tensor<2xi32> +// CHECK-DAG: %[[CST:.*]] = "tf.Const"() <{value = dense<[0, 1440]> : tensor<2xi32>}> : () -> tensor<2xi32> +// CHECK-DAG: %[[CST_0:.*]] = "tf.Const"() <{value = dense<0> : tensor<2xi32>}> : () -> tensor<2xi32> // CHECK: %[[VAL_0:.*]] = "tf.Maximum"(%[[ARG_1]], %[[CST_0:.*]]) : (tensor<1x2xi32>, tensor<2xi32>) -> tensor<1x2xi32> // CHECK: %[[VAL_1:.*]] = "tf.Minimum"(%[[VAL_0]], %[[CST]]) : (tensor<1x2xi32>, tensor<2xi32>) -> tensor<1x2xi32> -// CHECK-DAG: %[[CST_1:.*]] = "tf.Const"() {value = dense<[1, 1504]> : tensor<2xi32>} : () -> tensor<2xi32> -// CHECK: %[[VAL_2:.*]] = "tf.Squeeze"(%[[VAL_1]]) {squeeze_dims = [0]} : (tensor<1x2xi32>) -> tensor<2xi32> +// CHECK-DAG: %[[CST_1:.*]] = "tf.Const"() <{value = dense<[1, 1504]> : tensor<2xi32>}> : () -> tensor<2xi32> +// CHECK: %[[VAL_2:.*]] = "tf.Squeeze"(%[[VAL_1]]) <{squeeze_dims = [0]}> : (tensor<1x2xi32>) -> tensor<2xi32> // CHECK: %[[VAL_3:.*]] = "tf.Slice"(%[[ARG_0]], %[[VAL_2]], %[[CST_1]]) : (tensor<1x2944xi32>, tensor<2xi32>, tensor<2xi32>) -> tensor<1x1504xi32> // CHECK: return %[[VAL_3]] // CHECK: } @@ -3878,27 +3878,27 @@ func.func @convert_gather_scalar_dynamic_indices(%arg0: tensor<256000xf32>, %arg // CHECK-LABEL: func @convert_gather_to_slice( // CHECK-SAME: %[[ARG_0:.*]]: tensor<3x2944xi32>, // CHECK-SAME: %[[ARG_1:.*]]: tensor<3x2xi32>) -// CHECK-DAG: %[[CST:.*]] = "tf.Const"() {value = dense<[2, 1440]> : tensor<2xi32>} : () -> tensor<2xi32> -// CHECK-DAG: %[[CST_0:.*]] = "tf.Const"() {value = dense<0> : tensor<2xi32>} : () -> tensor<2xi32> +// CHECK-DAG: %[[CST:.*]] = "tf.Const"() <{value = dense<[2, 1440]> : tensor<2xi32>}> : () -> tensor<2xi32> +// CHECK-DAG: %[[CST_0:.*]] = "tf.Const"() <{value = dense<0> : tensor<2xi32>}> : () -> tensor<2xi32> // CHECK: %[[VAL_0:.*]] = "tf.Maximum"(%[[ARG_1]], %[[CST_0]]) : (tensor<3x2xi32>, tensor<2xi32>) -> tensor<3x2xi32> // CHECK: %[[VAL_1:.*]] = "tf.Minimum"(%[[VAL_0]], %[[CST]]) : (tensor<3x2xi32>, tensor<2xi32>) -> tensor<3x2xi32> -// CHECK-DAG: %[[CST_1:.*]] = "tf.Const"() {value = dense<[1, 1504]> : tensor<2xi32>} : () -> tensor<2xi32> -// CHECK-DAG: %[[CST_2:.*]] = "tf.Const"() {value = dense<0> : tensor<2xi32>} : () -> tensor<2xi32> -// CHECK-DAG: %[[CST_3:.*]] = "tf.Const"() {value = dense<[1, 2]> : tensor<2xi32>} : () -> tensor<2xi32> +// CHECK-DAG: %[[CST_1:.*]] = "tf.Const"() <{value = dense<[1, 1504]> : tensor<2xi32>}> : () -> tensor<2xi32> +// CHECK-DAG: %[[CST_2:.*]] = "tf.Const"() <{value = dense<0> : tensor<2xi32>}> : () -> tensor<2xi32> +// CHECK-DAG: %[[CST_3:.*]] = "tf.Const"() <{value = dense<[1, 2]> : tensor<2xi32>}> : () -> tensor<2xi32> // CHECK: %[[VAL_2:.*]] = "tf.Slice"(%[[VAL_1]], %[[CST_2]], %[[CST_3]]) : (tensor<3x2xi32>, tensor<2xi32>, tensor<2xi32>) -> tensor<1x2xi32> -// CHECK: %[[VAL_3:.*]] = "tf.Squeeze"(%[[VAL_2]]) {squeeze_dims = [0]} : (tensor<1x2xi32>) -> tensor<2xi32> +// CHECK: %[[VAL_3:.*]] = "tf.Squeeze"(%[[VAL_2]]) <{squeeze_dims = [0]}> : (tensor<1x2xi32>) -> tensor<2xi32> // CHECK: %[[VAL_4:.*]] = "tf.Slice"(%[[ARG_0]], %[[VAL_3]], %[[CST_1]]) : (tensor<3x2944xi32>, tensor<2xi32>, tensor<2xi32>) -> tensor<1x1504xi32> -// CHECK-DAG: %[[CST_4:.*]] = "tf.Const"() {value = dense<[1, 0]> : tensor<2xi32>} : () -> tensor<2xi32> -// CHECK-DAG: %[[CST_5:.*]] = "tf.Const"() {value = dense<[1, 2]> : tensor<2xi32>} : () -> tensor<2xi32> +// CHECK-DAG: %[[CST_4:.*]] = "tf.Const"() <{value = dense<[1, 0]> : tensor<2xi32>}> : () -> tensor<2xi32> +// CHECK-DAG: %[[CST_5:.*]] = "tf.Const"() <{value = dense<[1, 2]> : tensor<2xi32>}> : () -> tensor<2xi32> // CHECK: %[[VAL_5:.*]] = "tf.Slice"(%[[VAL_1]], %[[CST_4]], %[[CST_5]]) : (tensor<3x2xi32>, tensor<2xi32>, tensor<2xi32>) -> tensor<1x2xi32> -// CHECK: %[[VAL_6:.*]] = "tf.Squeeze"(%[[VAL_5]]) {squeeze_dims = [0]} : (tensor<1x2xi32>) -> tensor<2xi32> +// CHECK: %[[VAL_6:.*]] = "tf.Squeeze"(%[[VAL_5]]) <{squeeze_dims = [0]}> : (tensor<1x2xi32>) -> tensor<2xi32> // CHECK: %[[VAL_7:.*]] = "tf.Slice"(%[[ARG_0]], %[[VAL_6]], %[[CST_1]]) : (tensor<3x2944xi32>, tensor<2xi32>, tensor<2xi32>) -> tensor<1x1504xi32> -// CHECK-DAG: %[[CST_6:.*]] = "tf.Const"() {value = dense<[2, 0]> : tensor<2xi32>} : () -> tensor<2xi32> -// CHECK-DAG: %[[CST_7:.*]] = "tf.Const"() {value = dense<[1, 2]> : tensor<2xi32>} : () -> tensor<2xi32> +// CHECK-DAG: %[[CST_6:.*]] = "tf.Const"() <{value = dense<[2, 0]> : tensor<2xi32>}> : () -> tensor<2xi32> +// CHECK-DAG: %[[CST_7:.*]] = "tf.Const"() <{value = dense<[1, 2]> : tensor<2xi32>}> : () -> tensor<2xi32> // CHECK: %[[VAL_8:.*]] = "tf.Slice"(%[[VAL_1]], %[[CST_6]], %[[CST_7]]) : (tensor<3x2xi32>, tensor<2xi32>, tensor<2xi32>) -> tensor<1x2xi32> -// CHECK: %[[VAL_9:.*]] = "tf.Squeeze"(%[[VAL_8]]) {squeeze_dims = [0]} : (tensor<1x2xi32>) -> tensor<2xi32> +// CHECK: %[[VAL_9:.*]] = "tf.Squeeze"(%[[VAL_8]]) <{squeeze_dims = [0]}> : (tensor<1x2xi32>) -> tensor<2xi32> // CHECK: %[[VAL_10:.*]] = "tf.Slice"(%[[ARG_0]], %[[VAL_9]], %[[CST_1]]) : (tensor<3x2944xi32>, tensor<2xi32>, tensor<2xi32>) -> tensor<1x1504xi32> -// CHECK-DAG: %[[CST_8:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor +// CHECK-DAG: %[[CST_8:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK: %[[VAL_11:.*]] = "tf.ConcatV2"(%[[VAL_4]], %[[VAL_7]], %[[VAL_10]], %[[CST_8]]) : (tensor<1x1504xi32>, tensor<1x1504xi32>, tensor<1x1504xi32>, tensor) -> tensor<3x1504xi32> // CHECK: return %[[VAL_11]] // CHECK: } @@ -3936,17 +3936,17 @@ func.func @convert_gather_to_slice_dynamic_error(%arg0: tensor<3x?xi32>, %arg1: // CHECK-SAME: %[[VAL_0:.*]]: tensor<7x3xf32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor, // CHECK-SAME: %[[VAL_2:.*]]: tensor) -> tensor<4x2xf32> { -// CHECK: %[[VAL_3:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor -// CHECK: %[[VAL_4:.*]] = "tf.Cast"(%[[VAL_1]]) {Truncate = false} : (tensor) -> tensor -// CHECK: %[[VAL_5:.*]] = "tf.Const"() {value = dense<3> : tensor} : () -> tensor +// CHECK: %[[VAL_3:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor +// CHECK: %[[VAL_4:.*]] = "tf.Cast"(%[[VAL_1]]) <{Truncate = false}> : (tensor) -> tensor +// CHECK: %[[VAL_5:.*]] = "tf.Const"() <{value = dense<3> : tensor}> : () -> tensor // CHECK: %[[VAL_6:.*]] = "tf.Minimum"(%[[VAL_4]], %[[VAL_5]]) : (tensor, tensor) -> tensor // CHECK: %[[VAL_7:.*]] = "tf.Maximum"(%[[VAL_6]], %[[VAL_3]]) : (tensor, tensor) -> tensor -// CHECK: %[[VAL_8:.*]] = "tf.Cast"(%[[VAL_2]]) {Truncate = false} : (tensor) -> tensor -// CHECK: %[[VAL_9:.*]] = "tf.Const"() {value = dense<1> : tensor} : () -> tensor +// CHECK: %[[VAL_8:.*]] = "tf.Cast"(%[[VAL_2]]) <{Truncate = false}> : (tensor) -> tensor +// CHECK: %[[VAL_9:.*]] = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor // CHECK: %[[VAL_10:.*]] = "tf.Minimum"(%[[VAL_8]], %[[VAL_9]]) : (tensor, tensor) -> tensor // CHECK: %[[VAL_11:.*]] = "tf.Maximum"(%[[VAL_10]], %[[VAL_3]]) : (tensor, tensor) -> tensor -// CHECK: %[[VAL_12:.*]] = "tf.Pack"(%[[VAL_7]], %[[VAL_11]]) {axis = 0 : i64} : (tensor, tensor) -> tensor<2xi32> -// CHECK: %[[VAL_13:.*]] = "tf.Const"() {value = dense<[4, 2]> : tensor<2xi64>} : () -> tensor<2xi64> +// CHECK: %[[VAL_12:.*]] = "tf.Pack"(%[[VAL_7]], %[[VAL_11]]) <{axis = 0 : i64}> : (tensor, tensor) -> tensor<2xi32> +// CHECK: %[[VAL_13:.*]] = "tf.Const"() <{value = dense<[4, 2]> : tensor<2xi64>}> : () -> tensor<2xi64> // CHECK: %[[VAL_14:.*]] = "tf.Slice"(%[[VAL_0]], %[[VAL_12]], %[[VAL_13]]) : (tensor<7x3xf32>, tensor<2xi32>, tensor<2xi64>) -> tensor<4x2xf32> // CHECK: return %[[VAL_14]] : tensor<4x2xf32> // CHECK: } @@ -3959,17 +3959,17 @@ func.func @convert_dynamic_slice(%arg0: tensor<7x3xf32>, %arg1: tensor, %ar // CHECK-SAME: %[[VAL_0:.*]]: tensor<7x3xf32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor, // CHECK-SAME: %[[VAL_2:.*]]: tensor) -> tensor<4x2xf32> { -// CHECK: %[[VAL_3:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor -// CHECK: %[[VAL_4:.*]] = "tf.Cast"(%[[VAL_1]]) {Truncate = false} : (tensor) -> tensor -// CHECK: %[[VAL_5:.*]] = "tf.Const"() {value = dense<3> : tensor} : () -> tensor +// CHECK: %[[VAL_3:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor +// CHECK: %[[VAL_4:.*]] = "tf.Cast"(%[[VAL_1]]) <{Truncate = false}> : (tensor) -> tensor +// CHECK: %[[VAL_5:.*]] = "tf.Const"() <{value = dense<3> : tensor}> : () -> tensor // CHECK: %[[VAL_6:.*]] = "tf.Minimum"(%[[VAL_4]], %[[VAL_5]]) : (tensor, tensor) -> tensor // CHECK: %[[VAL_7:.*]] = "tf.Maximum"(%[[VAL_6]], %[[VAL_3]]) : (tensor, tensor) -> tensor -// CHECK: %[[VAL_8:.*]] = "tf.Cast"(%[[VAL_2]]) {Truncate = false} : (tensor) -> tensor -// CHECK: %[[VAL_9:.*]] = "tf.Const"() {value = dense<1> : tensor} : () -> tensor +// CHECK: %[[VAL_8:.*]] = "tf.Cast"(%[[VAL_2]]) <{Truncate = false}> : (tensor) -> tensor +// CHECK: %[[VAL_9:.*]] = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor // CHECK: %[[VAL_10:.*]] = "tf.Minimum"(%[[VAL_8]], %[[VAL_9]]) : (tensor, tensor) -> tensor // CHECK: %[[VAL_11:.*]] = "tf.Maximum"(%[[VAL_10]], %[[VAL_3]]) : (tensor, tensor) -> tensor -// CHECK: %[[VAL_12:.*]] = "tf.Pack"(%[[VAL_7]], %[[VAL_11]]) {axis = 0 : i64} : (tensor, tensor) -> tensor<2xi32> -// CHECK: %[[VAL_13:.*]] = "tf.Const"() {value = dense<[4, 2]> : tensor<2xi64>} : () -> tensor<2xi64> +// CHECK: %[[VAL_12:.*]] = "tf.Pack"(%[[VAL_7]], %[[VAL_11]]) <{axis = 0 : i64}> : (tensor, tensor) -> tensor<2xi32> +// CHECK: %[[VAL_13:.*]] = "tf.Const"() <{value = dense<[4, 2]> : tensor<2xi64>}> : () -> tensor<2xi64> // CHECK: %[[VAL_14:.*]] = "tf.Slice"(%[[VAL_0]], %[[VAL_12]], %[[VAL_13]]) : (tensor<7x3xf32>, tensor<2xi32>, tensor<2xi64>) -> tensor<4x2xf32> // CHECK: return %[[VAL_14]] : tensor<4x2xf32> // CHECK: } @@ -4206,7 +4206,7 @@ func.func @convert_scatter_sub(%arg0: tensor<20x6xf32>, %arg1: tensor<4x1xi32>, // CHECK-LABEL: func @convert_argmax( // CHECK-SAME: %[[VAL_0:.*]]: tensor<4x32x256xf32>) -> (tensor<4x32xf32>, tensor<4x32xi32>) { // CHECK: %[[VAL_9:.*]] = "tf.Const"{{.*}}value = dense<2> : tensor<1xi64> -// CHECK: %[[VAL_10:.*]] = "tf.Max"(%[[VAL_0]], %[[VAL_9]]) {keep_dims = false} : {{.*}} -> tensor<4x32xf32> +// CHECK: %[[VAL_10:.*]] = "tf.Max"(%[[VAL_0]], %[[VAL_9]]) <{keep_dims = false}> : {{.*}} -> tensor<4x32xf32> // CHECK: %[[VAL_11:.*]] = "tf.ArgMax"(%[[VAL_0]], %[[VAL_9]]) : {{.*}} -> tensor<4x32xi32> // CHECK: return %[[VAL_10]], %[[VAL_11]] // CHECK: } @@ -4233,11 +4233,11 @@ func.func @convert_argmax(%arg0: tensor<4x32x256xf32>) -> (tensor<4x32xf32>, ten // CHECK-LABEL: func @convert_argmax_constant( // CHECK-SAME: %[[VAL_0:.*]]: tensor<2x2x4xf32>) -> (tensor<2x2xf32>, tensor<2x2xi32>) { -// CHECK-DAG: %[[VAL_1:.*]] = "tf.Const"() {value = dense<0xFF800000> : tensor} : () -> tensor -// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor -// CHECK-DAG: %[[VAL_3:.*]] = "tf.Const"() {value = dense<{{\[\[}}[0, 1, 2, 3], [0, 1, 2, 3]], {{\[\[}}0, 1, 2, 3], [0, 1, 2, 3]]]> : tensor<2x2x4xi32>} : () -> tensor<2x2x4xi32> -// CHECK-DAG: %[[VAL_4:.*]] = "tf.Const"() {value = dense<2> : tensor<1xi64>} : () -> tensor<1xi64> -// CHECK: %[[VAL_5:.*]] = "tf.Max"(%[[VAL_0]], %[[VAL_4]]) {keep_dims = false} : (tensor<2x2x4xf32>, tensor<1xi64>) -> tensor<2x2xf32> +// CHECK-DAG: %[[VAL_1:.*]] = "tf.Const"() <{value = dense<0xFF800000> : tensor}> : () -> tensor +// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor +// CHECK-DAG: %[[VAL_3:.*]] = "tf.Const"() <{value = dense<{{\[\[}}[0, 1, 2, 3], [0, 1, 2, 3]], {{\[\[}}0, 1, 2, 3], [0, 1, 2, 3]]]> : tensor<2x2x4xi32>}> : () -> tensor<2x2x4xi32> +// CHECK-DAG: %[[VAL_4:.*]] = "tf.Const"() <{value = dense<2> : tensor<1xi64>}> : () -> tensor<1xi64> +// CHECK: %[[VAL_5:.*]] = "tf.Max"(%[[VAL_0]], %[[VAL_4]]) <{keep_dims = false}> : (tensor<2x2x4xf32>, tensor<1xi64>) -> tensor<2x2xf32> // CHECK: %[[VAL_6:.*]] = "tf.ArgMax"(%[[VAL_0]], %[[VAL_4]]) : (tensor<2x2x4xf32>, tensor<1xi64>) -> tensor<2x2xi32> // CHECK: return %[[VAL_5]], %[[VAL_6]] : tensor<2x2xf32>, tensor<2x2xi32> // CHECK: } @@ -4263,11 +4263,11 @@ func.func @convert_argmax_constant(%arg0: tensor<2x2x4xf32>) -> (tensor<2x2xf32> // CHECK-LABEL: func @convert_argmax_constant_non_z_axis( // CHECK-SAME: %[[VAL_0:.*]]: tensor<4x4xf32>) -> (tensor<4xf32>, tensor<4xi32>) { -// CHECK-DAG: %[[VAL_1:.*]] = "tf.Const"() {value = dense<0xFF800000> : tensor} : () -> tensor -// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor -// CHECK-DAG: %[[VAL_3:.*]] = "tf.Const"() {value = dense<{{\[\[}}0, 0, 0, 0], [1, 1, 1, 1], [2, 2, 2, 2], [3, 3, 3, 3]]> : tensor<4x4xi32>} : () -> tensor<4x4xi32> -// CHECK-DAG: %[[VAL_4:.*]] = "tf.Const"() {value = dense<0> : tensor<1xi64>} : () -> tensor<1xi64> -// CHECK: %[[VAL_5:.*]] = "tf.Max"(%[[VAL_0]], %[[VAL_4]]) {keep_dims = false} : (tensor<4x4xf32>, tensor<1xi64>) -> tensor<4xf32> +// CHECK-DAG: %[[VAL_1:.*]] = "tf.Const"() <{value = dense<0xFF800000> : tensor}> : () -> tensor +// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor +// CHECK-DAG: %[[VAL_3:.*]] = "tf.Const"() <{value = dense<{{\[\[}}0, 0, 0, 0], [1, 1, 1, 1], [2, 2, 2, 2], [3, 3, 3, 3]]> : tensor<4x4xi32>}> : () -> tensor<4x4xi32> +// CHECK-DAG: %[[VAL_4:.*]] = "tf.Const"() <{value = dense<0> : tensor<1xi64>}> : () -> tensor<1xi64> +// CHECK: %[[VAL_5:.*]] = "tf.Max"(%[[VAL_0]], %[[VAL_4]]) <{keep_dims = false}> : (tensor<4x4xf32>, tensor<1xi64>) -> tensor<4xf32> // CHECK: %[[VAL_6:.*]] = "tf.ArgMax"(%[[VAL_0]], %[[VAL_4]]) : (tensor<4x4xf32>, tensor<1xi64>) -> tensor<4xi32> // CHECK: return %[[VAL_5]], %[[VAL_6]] : tensor<4xf32>, tensor<4xi32> // CHECK: } @@ -4293,14 +4293,14 @@ func.func @convert_argmax_constant_non_z_axis(%arg0: tensor<4x4xf32>) -> (tensor // CHECK-LABEL: func.func @convert_argmax_bool( // CHECK-SAME: %[[VAL_0:.*]]: tensor<2xi1>) -> tensor { -// CHECK-DAG: %[[VAL_1:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor -// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() {value = dense<2> : tensor} : () -> tensor -// CHECK-DAG: %[[VAL_3:.*]] = "tf.Const"() {value = dense<1> : tensor} : () -> tensor +// CHECK-DAG: %[[VAL_1:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor +// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() <{value = dense<2> : tensor}> : () -> tensor +// CHECK-DAG: %[[VAL_3:.*]] = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor // CHECK: %[[VAL_4:.*]] = "tf.Range"(%[[VAL_1]], %[[VAL_2]], %[[VAL_3]]) : (tensor, tensor, tensor) -> tensor<2xi32> -// CHECK-DAG: %[[VAL_5:.*]] = "tf.Const"() {value = dense : tensor} : () -> tensor -// CHECK-DAG: %[[VAL_6:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor -// CHECK-DAG: %[[VAL_7:.*]] = "tf.Const"() {value = dense<0> : tensor<1xi64>} : () -> tensor<1xi64> -// CHECK: %[[VAL_8:.*]] = "tf.Any"(%[[VAL_0]], %[[VAL_7]]) {keep_dims = false} : (tensor<2xi1>, tensor<1xi64>) -> tensor +// CHECK-DAG: %[[VAL_5:.*]] = "tf.Const"() <{value = dense : tensor}> : () -> tensor +// CHECK-DAG: %[[VAL_6:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor +// CHECK-DAG: %[[VAL_7:.*]] = "tf.Const"() <{value = dense<0> : tensor<1xi64>}> : () -> tensor<1xi64> +// CHECK: %[[VAL_8:.*]] = "tf.Any"(%[[VAL_0]], %[[VAL_7]]) <{keep_dims = false}> : (tensor<2xi1>, tensor<1xi64>) -> tensor // CHECK: %[[VAL_9:.*]] = "tf.ArgMax"(%[[VAL_0]], %[[VAL_7]]) : (tensor<2xi1>, tensor<1xi64>) -> tensor // CHECK: return %[[VAL_9]] : tensor // CHECK: } @@ -4326,7 +4326,7 @@ func.func @convert_argmax_bool(%arg0: tensor<2xi1>) -> tensor { // CHECK-LABEL: func @convert_argmin( // CHECK-SAME: %[[VAL_0:.*]]: tensor<4x32x256xf32>) -> (tensor<4x32xf32>, tensor<4x32xi32>) { // CHECK: %[[VAL_9:.*]] = "tf.Const"{{.*}}value = dense<2> : tensor<1xi64> -// CHECK: %[[VAL_10:.*]] = "tf.Min"(%[[VAL_0]], %[[VAL_9]]) {keep_dims = false} : {{.*}} -> tensor<4x32xf32> +// CHECK: %[[VAL_10:.*]] = "tf.Min"(%[[VAL_0]], %[[VAL_9]]) <{keep_dims = false}> : {{.*}} -> tensor<4x32xf32> // CHECK: %[[VAL_11:.*]] = "tf.ArgMin"(%[[VAL_0]], %[[VAL_9]]) : {{.*}} -> tensor<4x32xi32> // CHECK: return %[[VAL_10]], %[[VAL_11]] // CHECK: } @@ -4354,7 +4354,7 @@ func.func @convert_argmin(%arg0: tensor<4x32x256xf32>) -> (tensor<4x32xf32>, ten // CHECK-LABEL: func @convert_argmin_i16( // CHECK-SAME: %[[VAL_0:.*]]: tensor<2xi16>) -> (tensor, tensor) { // CHECK: %[[VAL_9:.*]] = "tf.Const"{{.*}}value = dense<0> : tensor<1xi64> -// CHECK: %[[VAL_10:.*]] = "tf.Min"(%[[VAL_0]], %[[VAL_9]]) {keep_dims = false} : {{.*}} -> tensor +// CHECK: %[[VAL_10:.*]] = "tf.Min"(%[[VAL_0]], %[[VAL_9]]) <{keep_dims = false}> : {{.*}} -> tensor // CHECK: %[[VAL_11:.*]] = "tf.ArgMin"(%[[VAL_0]], %[[VAL_9]]) : {{.*}} -> tensor // CHECK: return %[[VAL_10]], %[[VAL_11]] // CHECK: } @@ -4381,11 +4381,11 @@ func.func @convert_argmin_i16(%arg0: tensor<2xi16>) -> (tensor, tensor // CHECK-LABEL: func @convert_argmin_constant( // CHECK-SAME: %[[VAL_0:.*]]: tensor<2x2x4xf32>) -> (tensor<2x2xf32>, tensor<2x2xi32>) { -// CHECK-DAG: %[[VAL_1:.*]] = "tf.Const"() {value = dense<0x7F800000> : tensor} : () -> tensor -// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor -// CHECK-DAG: %[[VAL_3:.*]] = "tf.Const"() {value = dense<{{\[\[}}[0, 1, 2, 3], [0, 1, 2, 3]], {{\[\[}}0, 1, 2, 3], [0, 1, 2, 3]]]> : tensor<2x2x4xi32>} : () -> tensor<2x2x4xi32> -// CHECK-DAG: %[[VAL_4:.*]] = "tf.Const"() {value = dense<2> : tensor<1xi64>} : () -> tensor<1xi64> -// CHECK: %[[VAL_5:.*]] = "tf.Min"(%[[VAL_0]], %[[VAL_4]]) {keep_dims = false} : (tensor<2x2x4xf32>, tensor<1xi64>) -> tensor<2x2xf32> +// CHECK-DAG: %[[VAL_1:.*]] = "tf.Const"() <{value = dense<0x7F800000> : tensor}> : () -> tensor +// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor +// CHECK-DAG: %[[VAL_3:.*]] = "tf.Const"() <{value = dense<{{\[\[}}[0, 1, 2, 3], [0, 1, 2, 3]], {{\[\[}}0, 1, 2, 3], [0, 1, 2, 3]]]> : tensor<2x2x4xi32>}> : () -> tensor<2x2x4xi32> +// CHECK-DAG: %[[VAL_4:.*]] = "tf.Const"() <{value = dense<2> : tensor<1xi64>}> : () -> tensor<1xi64> +// CHECK: %[[VAL_5:.*]] = "tf.Min"(%[[VAL_0]], %[[VAL_4]]) <{keep_dims = false}> : (tensor<2x2x4xf32>, tensor<1xi64>) -> tensor<2x2xf32> // CHECK: %[[VAL_6:.*]] = "tf.ArgMin"(%[[VAL_0]], %[[VAL_4]]) : (tensor<2x2x4xf32>, tensor<1xi64>) -> tensor<2x2xi32> // CHECK: return %[[VAL_5]], %[[VAL_6]] : tensor<2x2xf32>, tensor<2x2xi32> // CHECK: } @@ -4411,14 +4411,14 @@ func.func @convert_argmin_constant(%arg0: tensor<2x2x4xf32>) -> (tensor<2x2xf32> // CHECK-LABEL: func.func @convert_argmin_bool( // CHECK-SAME: %[[VAL_0:.*]]: tensor<2xi1>) -> tensor { -// CHECK-DAG: %[[VAL_1:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor -// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() {value = dense<2> : tensor} : () -> tensor -// CHECK-DAG: %[[VAL_3:.*]] = "tf.Const"() {value = dense<1> : tensor} : () -> tensor +// CHECK-DAG: %[[VAL_1:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor +// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() <{value = dense<2> : tensor}> : () -> tensor +// CHECK-DAG: %[[VAL_3:.*]] = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor // CHECK: %[[VAL_4:.*]] = "tf.Range"(%[[VAL_1]], %[[VAL_2]], %[[VAL_3]]) : (tensor, tensor, tensor) -> tensor<2xi32> -// CHECK-DAG: %[[VAL_5:.*]] = "tf.Const"() {value = dense : tensor} : () -> tensor -// CHECK-DAG: %[[VAL_6:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor -// CHECK-DAG: %[[VAL_7:.*]] = "tf.Const"() {value = dense<0> : tensor<1xi64>} : () -> tensor<1xi64> -// CHECK: %[[VAL_8:.*]] = "tf.All"(%[[VAL_0]], %[[VAL_7]]) {keep_dims = false} : (tensor<2xi1>, tensor<1xi64>) -> tensor +// CHECK-DAG: %[[VAL_5:.*]] = "tf.Const"() <{value = dense : tensor}> : () -> tensor +// CHECK-DAG: %[[VAL_6:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor +// CHECK-DAG: %[[VAL_7:.*]] = "tf.Const"() <{value = dense<0> : tensor<1xi64>}> : () -> tensor<1xi64> +// CHECK: %[[VAL_8:.*]] = "tf.All"(%[[VAL_0]], %[[VAL_7]]) <{keep_dims = false}> : (tensor<2xi1>, tensor<1xi64>) -> tensor // CHECK: %[[VAL_9:.*]] = "tf.ArgMin"(%[[VAL_0]], %[[VAL_7]]) : (tensor<2xi1>, tensor<1xi64>) -> tensor // CHECK: return %[[VAL_9]] : tensor // CHECK: } @@ -4442,16 +4442,16 @@ func.func @convert_argmin_bool(%arg0: tensor<2xi1>) -> tensor { // CHECK-LABEL: func @convert_argmax_with_reshaped_iota( // CHECK-SAME: %[[VAL_0:.*]]: tensor<1x32x1xf32>) -> (tensor<1x1xf32>, tensor<1x1xi32>) { -// CHECK-DAG: %[[VAL_1:.*]] = "tf.Const"() {value = dense<0xFF800000> : tensor} : () -> tensor -// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor -// CHECK-DAG: %[[VAL_3:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor -// CHECK-DAG: %[[VAL_4:.*]] = "tf.Const"() {value = dense<32> : tensor} : () -> tensor -// CHECK-DAG: %[[VAL_5:.*]] = "tf.Const"() {value = dense<1> : tensor} : () -> tensor +// CHECK-DAG: %[[VAL_1:.*]] = "tf.Const"() <{value = dense<0xFF800000> : tensor}> : () -> tensor +// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor +// CHECK-DAG: %[[VAL_3:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor +// CHECK-DAG: %[[VAL_4:.*]] = "tf.Const"() <{value = dense<32> : tensor}> : () -> tensor +// CHECK-DAG: %[[VAL_5:.*]] = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor // CHECK: %[[VAL_6:.*]] = "tf.Range"(%[[VAL_3]], %[[VAL_4]], %[[VAL_5]]) : (tensor, tensor, tensor) -> tensor<32xi32> // CHECK-DAG: %[[VAL_7:.*]] = arith.constant dense<[1, 32, 1]> : tensor<3xi64> // CHECK: %[[VAL_8:.*]] = "tf.Reshape"(%[[VAL_6]], %[[VAL_7]]) : (tensor<32xi32>, tensor<3xi64>) -> tensor<1x32x1xi32> -// CHECK-DAG: %[[VAL_9:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi64>} : () -> tensor<1xi64> -// CHECK: %[[VAL_10:.*]] = "tf.Max"(%[[VAL_0]], %[[VAL_9]]) {keep_dims = false} : (tensor<1x32x1xf32>, tensor<1xi64>) -> tensor<1x1xf32> +// CHECK-DAG: %[[VAL_9:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi64>}> : () -> tensor<1xi64> +// CHECK: %[[VAL_10:.*]] = "tf.Max"(%[[VAL_0]], %[[VAL_9]]) <{keep_dims = false}> : (tensor<1x32x1xf32>, tensor<1xi64>) -> tensor<1x1xf32> // CHECK: %[[VAL_11:.*]] = "tf.ArgMax"(%[[VAL_0]], %[[VAL_9]]) : (tensor<1x32x1xf32>, tensor<1xi64>) -> tensor<1x1xi32> // CHECK: return %[[VAL_10]], %[[VAL_11]] : tensor<1x1xf32>, tensor<1x1xi32> // CHECK: } @@ -4488,7 +4488,7 @@ func.func @convert_not(%arg0: tensor<5x3x1xi1>) -> tensor<5x3x1xi1> { // CHECK-LABEL: func @convert_not_i8( // CHECK-SAME: %[[ARG:.*]]: tensor<7x9x11xi8>) -> tensor<7x9x11xi8> { -// CHECK: %[[CST:.*]] = "tf.Const"() {value = dense<-1> : tensor} : () -> tensor +// CHECK: %[[CST:.*]] = "tf.Const"() <{value = dense<-1> : tensor}> : () -> tensor // CHECK: %[[RES:.*]] = "tf.BitwiseXor"(%[[ARG]], %[[CST]]) : (tensor<7x9x11xi8>, tensor) -> tensor<7x9x11xi8> // CHECK: return %[[RES]] : tensor<7x9x11xi8> // CHECK: } @@ -4499,7 +4499,7 @@ func.func @convert_not_i8(%arg0: tensor<7x9x11xi8>) -> tensor<7x9x11xi8> { // CHECK-LABEL: func @convert_not_i16( // CHECK-SAME: %[[ARG:.*]]: tensor<7x9x11xi16>) -> tensor<7x9x11xi16> { -// CHECK: %[[CST:.*]] = "tf.Const"() {value = dense<-1> : tensor} : () -> tensor +// CHECK: %[[CST:.*]] = "tf.Const"() <{value = dense<-1> : tensor}> : () -> tensor // CHECK: %[[RES:.*]] = "tf.BitwiseXor"(%[[ARG]], %[[CST]]) : (tensor<7x9x11xi16>, tensor) -> tensor<7x9x11xi16> // CHECK: return %[[RES]] : tensor<7x9x11xi16> // CHECK: } @@ -4510,7 +4510,7 @@ func.func @convert_not_i16(%arg0: tensor<7x9x11xi16>) -> tensor<7x9x11xi16> { // CHECK-LABEL: func @convert_not_i32( // CHECK-SAME: %[[ARG:.*]]: tensor<7x9x11xi32>) -> tensor<7x9x11xi32> { -// CHECK: %[[CST:.*]] = "tf.Const"() {value = dense<-1> : tensor} : () -> tensor +// CHECK: %[[CST:.*]] = "tf.Const"() <{value = dense<-1> : tensor}> : () -> tensor // CHECK: %[[RES:.*]] = "tf.BitwiseXor"(%[[ARG]], %[[CST]]) : (tensor<7x9x11xi32>, tensor) -> tensor<7x9x11xi32> // CHECK: return %[[RES]] : tensor<7x9x11xi32> // CHECK: } @@ -4521,7 +4521,7 @@ func.func @convert_not_i32(%arg0: tensor<7x9x11xi32>) -> tensor<7x9x11xi32> { // CHECK-LABEL: func @convert_not_i64( // CHECK-SAME: %[[ARG:.*]]: tensor<7x9x11xi64>) -> tensor<7x9x11xi64> { -// CHECK: %[[CST:.*]] = "tf.Const"() {value = dense<-1> : tensor} : () -> tensor +// CHECK: %[[CST:.*]] = "tf.Const"() <{value = dense<-1> : tensor}> : () -> tensor // CHECK: %[[RES:.*]] = "tf.BitwiseXor"(%[[ARG]], %[[CST]]) : (tensor<7x9x11xi64>, tensor) -> tensor<7x9x11xi64> // CHECK: return %[[RES]] : tensor<7x9x11xi64> // CHECK: } @@ -4532,7 +4532,7 @@ func.func @convert_not_i64(%arg0: tensor<7x9x11xi64>) -> tensor<7x9x11xi64> { // CHECK-LABEL: func @convert_not_ui8( // CHECK-SAME: %[[ARG:.*]]: tensor<7x9x11xui8>) -> tensor<7x9x11xui8> { -// CHECK: %[[CST:.*]] = "tf.Const"() {value = dense<255> : tensor} : () -> tensor +// CHECK: %[[CST:.*]] = "tf.Const"() <{value = dense<255> : tensor}> : () -> tensor // CHECK: %[[RES:.*]] = "tf.BitwiseXor"(%[[ARG]], %[[CST]]) : (tensor<7x9x11xui8>, tensor) -> tensor<7x9x11xui8> // CHECK: return %[[RES]] : tensor<7x9x11xui8> // CHECK: } @@ -4543,7 +4543,7 @@ func.func @convert_not_ui8(%arg0: tensor<7x9x11xui8>) -> tensor<7x9x11xui8> { // CHECK-LABEL: func @convert_not_ui16( // CHECK-SAME: %[[ARG:.*]]: tensor<7x9x11xui16>) -> tensor<7x9x11xui16> { -// CHECK: %[[CST:.*]] = "tf.Const"() {value = dense<65535> : tensor} : () -> tensor +// CHECK: %[[CST:.*]] = "tf.Const"() <{value = dense<65535> : tensor}> : () -> tensor // CHECK: %[[RES:.*]] = "tf.BitwiseXor"(%[[ARG]], %[[CST]]) : (tensor<7x9x11xui16>, tensor) -> tensor<7x9x11xui16> // CHECK: return %[[RES]] : tensor<7x9x11xui16> // CHECK: } @@ -4554,7 +4554,7 @@ func.func @convert_not_ui16(%arg0: tensor<7x9x11xui16>) -> tensor<7x9x11xui16> { // CHECK-LABEL: func @convert_not_ui32( // CHECK-SAME: %[[ARG:.*]]: tensor<7x9x11xui32>) -> tensor<7x9x11xui32> { -// CHECK: %[[CST:.*]] = "tf.Const"() {value = dense<4294967295> : tensor} : () -> tensor +// CHECK: %[[CST:.*]] = "tf.Const"() <{value = dense<4294967295> : tensor}> : () -> tensor // CHECK: %[[RES:.*]] = "tf.BitwiseXor"(%[[ARG]], %[[CST]]) : (tensor<7x9x11xui32>, tensor) -> tensor<7x9x11xui32> // CHECK: return %[[RES]] : tensor<7x9x11xui32> // CHECK: } @@ -4565,7 +4565,7 @@ func.func @convert_not_ui32(%arg0: tensor<7x9x11xui32>) -> tensor<7x9x11xui32> { // CHECK-LABEL: func @convert_not_ui64( // CHECK-SAME: %[[ARG:.*]]: tensor<7x9x11xui64>) -> tensor<7x9x11xui64> { -// CHECK: %[[CST:.*]] = "tf.Const"() {value = dense<18446744073709551615> : tensor} : () -> tensor +// CHECK: %[[CST:.*]] = "tf.Const"() <{value = dense<18446744073709551615> : tensor}> : () -> tensor // CHECK: %[[RES:.*]] = "tf.BitwiseXor"(%[[ARG]], %[[CST]]) : (tensor<7x9x11xui64>, tensor) -> tensor<7x9x11xui64> // CHECK: return %[[RES]] : tensor<7x9x11xui64> // CHECK: } @@ -4580,7 +4580,7 @@ func.func @convert_not_ui64(%arg0: tensor<7x9x11xui64>) -> tensor<7x9x11xui64> { // CHECK-DAG: %[[CST_0:.*]] = arith.constant dense<1> : tensor // CHECK-DAG: %[[CST_1:.*]] = arith.constant dense<0> : tensor // CHECK-DAG: %[[CST_2:.*]] = arith.constant dense<1000> : tensor -// CHECK: %[[WHILEREGION_0:.*]]:3 = "tf.WhileRegion"(%[[CST_1]], %[[CST_0]], %[[CST_2]]) ({ +// CHECK: %[[WHILEREGION_0:.*]]:3 = "tf.WhileRegion"(%[[CST_1]], %[[CST_0]], %[[CST_2]]) <{is_stateless = false, parallel_iterations = 10 : i64}> ({ // CHECK: ^bb0(%arg0: tensor, %arg1: tensor, %arg2: tensor): // CHECK: %[[LESS_0:.*]] = "tf.Less"(%arg0, %arg2) : (tensor, tensor) -> tensor // CHECK: "tf.Yield"(%[[LESS_0]]) : (tensor) -> () @@ -4588,7 +4588,7 @@ func.func @convert_not_ui64(%arg0: tensor<7x9x11xui64>) -> tensor<7x9x11xui64> { // CHECK: ^bb0(%arg0: tensor, %arg1: tensor, %arg2: tensor): // CHECK: %[[ADDV2_0:.*]] = "tf.AddV2"(%arg0, %arg1) : (tensor, tensor) -> tensor // CHECK: "tf.Yield"(%[[ADDV2_0]], %arg1, %arg2) : (tensor, tensor, tensor) -> () -// CHECK: }) {is_stateless = false, parallel_iterations = 10 : i64} : (tensor, tensor, tensor) -> (tensor, tensor, tensor) +// CHECK: }) : (tensor, tensor, tensor) -> (tensor, tensor, tensor) // CHECK: return %[[WHILEREGION_0]]#0, %[[WHILEREGION_0]]#1, %[[WHILEREGION_0]]#2 : tensor, tensor, tensor // CHECK: } func.func @while_with_variadic() -> (tensor, tensor, tensor) { @@ -4613,19 +4613,19 @@ func.func @while_with_variadic() -> (tensor, tensor, tensor) { // CHECK-DAG: %[[CST_0:.*]] = arith.constant dense<1> : tensor // CHECK-DAG: %[[CST_1:.*]] = arith.constant dense<0> : tensor // CHECK-DAG: %[[CST_2:.*]] = arith.constant dense<1000> : tensor -// CHECK: %[[WHILEREGION_0:.*]]:5 = "tf.WhileRegion"(%[[CST_1]], %[[CST_0]], %[[CST_2]], %arg0, %arg1) ({ +// CHECK: %[[WHILEREGION_0:.*]]:5 = "tf.WhileRegion"(%[[CST_1]], %[[CST_0]], %[[CST_2]], %arg0, %arg1) <{is_stateless = false, parallel_iterations = 10 : i64}> ({ // CHECK: ^bb0(%arg2: tensor, %arg3: tensor, %arg4: tensor, %arg5: tensor<1x256xf32>, %arg6: tensor<1xf32>): // CHECK: %[[LESS_0:.*]] = "tf.Less"(%arg2, %arg4) : (tensor, tensor) -> tensor // CHECK: "tf.Yield"(%[[LESS_0]]) : (tensor) -> () // CHECK: }, { // CHECK: ^bb0(%arg2: tensor, %arg3: tensor, %arg4: tensor, %arg5: tensor<1x256xf32>, %arg6: tensor<1xf32>): // CHECK: %[[ADDV2_0:.*]] = "tf.AddV2"(%arg2, %arg3) : (tensor, tensor) -> tensor -// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor} : () -> tensor -// CHECK-DAG: %[[CONST_1:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi64>} : () -> tensor<1xi64> -// CHECK: %[[SUM_0:.*]] = "tf.Sum"(%arg5, %[[CONST_1]]) {keep_dims = false} : (tensor<1x256xf32>, tensor<1xi64>) -> tensor<1xf32> +// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor}> : () -> tensor +// CHECK-DAG: %[[CONST_1:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi64>}> : () -> tensor<1xi64> +// CHECK: %[[SUM_0:.*]] = "tf.Sum"(%arg5, %[[CONST_1]]) <{keep_dims = false}> : (tensor<1x256xf32>, tensor<1xi64>) -> tensor<1xf32> // CHECK: %[[ADDV2_1:.*]] = "tf.AddV2"(%[[SUM_0]], %arg6) : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> // CHECK: "tf.Yield"(%[[ADDV2_0]], %arg3, %arg4, %arg5, %[[ADDV2_1]]) : (tensor, tensor, tensor, tensor<1x256xf32>, tensor<1xf32>) -> () -// CHECK: }) {is_stateless = false, parallel_iterations = 10 : i64} : (tensor, tensor, tensor, tensor<1x256xf32>, tensor<1xf32>) -> (tensor, tensor, tensor, tensor<1x256xf32>, tensor<1xf32>) +// CHECK: }) : (tensor, tensor, tensor, tensor<1x256xf32>, tensor<1xf32>) -> (tensor, tensor, tensor, tensor<1x256xf32>, tensor<1xf32>) // CHECK: return %[[WHILEREGION_0]]#0, %[[WHILEREGION_0]]#1, %[[WHILEREGION_0]]#2, %[[WHILEREGION_0]]#4 : tensor, tensor, tensor, tensor<1xf32> // CHECK: } func.func @while_with_reduce(%arg0: tensor<1x256xf32>, %arg1: tensor<1xf32>) -> (tensor, tensor, tensor, tensor<1xf32>) { @@ -4656,11 +4656,11 @@ func.func @while_with_reduce(%arg0: tensor<1x256xf32>, %arg1: tensor<1xf32>) -> // CHECK-LABEL: func @if // CHECK-DAG: %[[CST_0:.*]] = arith.constant dense<0> : tensor // CHECK-DAG: %[[CST_1:.*]] = arith.constant dense<1000> : tensor -// CHECK: %[[RES:.*]] = "tf.IfRegion"(%arg0) ({ +// CHECK: %[[RES:.*]] = "tf.IfRegion"(%arg0) <{is_stateless = false}> ({ // CHECK: "tf.Yield"(%[[CST_0]]) : (tensor) -> () // CHECK: }, { // CHECK: "tf.Yield"(%[[CST_1]]) : (tensor) -> () -// CHECK: }) {is_stateless = false} : (tensor) -> tensor +// CHECK: }) : (tensor) -> tensor // CHECK: return %[[RES]] func.func @if(%arg0: tensor) -> (tensor) { %cst_0 = arith.constant dense<0> : tensor @@ -4679,7 +4679,7 @@ func.func @if(%arg0: tensor) -> (tensor) { // CHECK-SAME: %[[VAL_2:[a-z0-9]*]]: tensor, // CHECK-SAME: %[[VAL_3:[a-z0-9]*]]: tensor, // CHECK-SAME: %[[VAL_4:[a-z0-9]*]]: tensor) -> tensor<28x1x100xf32> { -// CHECK: %0 = "tf.Pack"(%arg2, %arg3, %arg4) {axis = 0 : i64} : (tensor, tensor, tensor) -> tensor<3xi32> +// CHECK: %0 = "tf.Pack"(%arg2, %arg3, %arg4) <{axis = 0 : i64}> : (tensor, tensor, tensor) -> tensor<3xi32> // CHECK: %1 = "tf.XlaDynamicUpdateSlice"(%arg0, %arg1, %0) : (tensor<28x1x100xf32>, tensor<1x1x100xf32>, tensor<3xi32>) -> tensor<28x1x100xf32> // CHECK: return %1 : tensor<28x1x100xf32> func.func @convert_dynamic_update_slice(%arg0: tensor<28x1x100xf32>, %arg1: tensor<1x1x100xf32>, %arg2: tensor, %arg3: tensor, %arg4: tensor) -> tensor<28x1x100xf32> { @@ -4692,7 +4692,7 @@ func.func @convert_dynamic_update_slice(%arg0: tensor<28x1x100xf32>, %arg1: tens // CHECK-SAME: %arg1: tensor, // CHECK-SAME: %arg2: tensor, // CHECK-SAME: %arg3: tensor) -> tensor { -// CHECK: %0 = "tf.Pack"(%arg2, %arg3) {axis = 0 : i64} : (tensor, tensor) -> tensor<2xi32> +// CHECK: %0 = "tf.Pack"(%arg2, %arg3) <{axis = 0 : i64}> : (tensor, tensor) -> tensor<2xi32> // CHECK: %1 = "tf.XlaDynamicUpdateSlice"(%arg0, %arg1, %0) : (tensor, tensor, tensor<2xi32>) -> tensor // CHECK: return %1 : tensor // CHECK: } @@ -4707,7 +4707,7 @@ func.func @dynamic_update_slice_inputs_have_dynamic_dim(%arg0: tensor, // CHECK-SAME: %arg2: tensor, // CHECK-SAME: %arg3: tensor, // CHECK-SAME: %arg4: tensor) -> tensor<1x?x256xf32> { -// CHECK: %0 = "tf.Pack"(%arg2, %arg3, %arg4) {axis = 0 : i64} : (tensor, tensor, tensor) -> tensor<3xi32> +// CHECK: %0 = "tf.Pack"(%arg2, %arg3, %arg4) <{axis = 0 : i64}> : (tensor, tensor, tensor) -> tensor<3xi32> // CHECK: %1 = "tf.XlaDynamicUpdateSlice"(%arg0, %arg1, %0) : (tensor<1x?x256xf32>, tensor<1x1x256xf32>, tensor<3xi32>) -> tensor<1x?x256xf32> // CHECK: return %1 : tensor<1x?x256xf32> // CHECK: } @@ -4719,9 +4719,9 @@ func.func @dynamic_update_slice_operand_has_dynamic_dim(%arg0: tensor<1x?x256xf3 // CHECK-LABEL: func @convert_reduce_to_all( // CHECK-SAME: %[[ARG_0:.*]]: tensor<1x2x3x4x5xi1>, // CHECK-SAME: %[[ARG_1:.*]]: tensor<2xi64>) -> tensor<2x4x5xi1> { -// CHECK-DAG: %[[TRUE_CST:.*]] = "tf.Const"() {value = dense : tensor} : () -> tensor -// CHECK-DAG: %[[DIMENSIONS:.*]] = "tf.Const"() {value = dense<[0, 2]> : tensor<2xi64>} : () -> tensor<2xi64> -// CHECK: %[[VAL_0:.*]] = "tf.All"(%[[ARG_0]], %[[DIMENSIONS]]) {keep_dims = false} : (tensor<1x2x3x4x5xi1>, tensor<2xi64>) -> tensor<2x4x5xi1> +// CHECK-DAG: %[[TRUE_CST:.*]] = "tf.Const"() <{value = dense : tensor}> : () -> tensor +// CHECK-DAG: %[[DIMENSIONS:.*]] = "tf.Const"() <{value = dense<[0, 2]> : tensor<2xi64>}> : () -> tensor<2xi64> +// CHECK: %[[VAL_0:.*]] = "tf.All"(%[[ARG_0]], %[[DIMENSIONS]]) <{keep_dims = false}> : (tensor<1x2x3x4x5xi1>, tensor<2xi64>) -> tensor<2x4x5xi1> // CHECK: return %[[VAL_0:.*]] : tensor<2x4x5xi1> // CHECK: } func.func @convert_reduce_to_all(%arg0: tensor<1x2x3x4x5xi1>, %arg1: tensor<2xi64>) -> tensor<2x4x5xi1> { @@ -4738,8 +4738,8 @@ func.func @convert_reduce_to_all(%arg0: tensor<1x2x3x4x5xi1>, %arg1: tensor<2xi6 // CHECK-SAME: %[[ARG_0:.*]]: tensor, // CHECK-SAME: %[[ARG_1:.*]]: tensor<1x2x3x4x5xi1>, // CHECK-SAME: %[[ARG_2:.*]]: tensor<2xi64>) -> tensor<2x4x5xi1> { -// CHECK-DAG: %[[DIMENSIONS:.*]] = "tf.Const"() {value = dense<[0, 2]> : tensor<2xi64>} : () -> tensor<2xi64> -// CHECK: %[[VAL_0:.*]] = "tf.All"(%[[ARG_1]], %[[DIMENSIONS]]) {keep_dims = false} : (tensor<1x2x3x4x5xi1>, tensor<2xi64>) -> tensor<2x4x5xi1> +// CHECK-DAG: %[[DIMENSIONS:.*]] = "tf.Const"() <{value = dense<[0, 2]> : tensor<2xi64>}> : () -> tensor<2xi64> +// CHECK: %[[VAL_0:.*]] = "tf.All"(%[[ARG_1]], %[[DIMENSIONS]]) <{keep_dims = false}> : (tensor<1x2x3x4x5xi1>, tensor<2xi64>) -> tensor<2x4x5xi1> // CHECK: %[[VAL_1:.*]] = "tf.LogicalAnd"(%[[VAL_0]], %[[ARG_0]]) : (tensor<2x4x5xi1>, tensor) -> tensor<2x4x5xi1> // CHECK: return %[[VAL_1:.*]] : tensor<2x4x5xi1> // CHECK: } @@ -4755,9 +4755,9 @@ func.func @convert_reduce_to_all_non_constant_init(%arg0: tensor, %arg1: ten // CHECK-LABEL: func @convert_reduce_to_any( // CHECK-SAME: %[[ARG_0:.*]]: tensor<1x2x3x4x5xi1>, // CHECK-SAME: %[[ARG_1:.*]]: tensor<2xi64>) -> tensor<2x4x5xi1> { -// CHECK-DAG: %[[FALSE_CST:.*]] = "tf.Const"() {value = dense : tensor} : () -> tensor -// CHECK-DAG: %[[DIMENSIONS:.*]] = "tf.Const"() {value = dense<[0, 2]> : tensor<2xi64>} : () -> tensor<2xi64> -// CHECK: %[[VAL_0:.*]] = "tf.Any"(%[[ARG_0]], %[[DIMENSIONS]]) {keep_dims = false} : (tensor<1x2x3x4x5xi1>, tensor<2xi64>) -> tensor<2x4x5xi1> +// CHECK-DAG: %[[FALSE_CST:.*]] = "tf.Const"() <{value = dense : tensor}> : () -> tensor +// CHECK-DAG: %[[DIMENSIONS:.*]] = "tf.Const"() <{value = dense<[0, 2]> : tensor<2xi64>}> : () -> tensor<2xi64> +// CHECK: %[[VAL_0:.*]] = "tf.Any"(%[[ARG_0]], %[[DIMENSIONS]]) <{keep_dims = false}> : (tensor<1x2x3x4x5xi1>, tensor<2xi64>) -> tensor<2x4x5xi1> // CHECK: return %[[VAL_0:.*]] : tensor<2x4x5xi1> // CHECK: } func.func @convert_reduce_to_any(%arg0: tensor<1x2x3x4x5xi1>, %arg1: tensor<2xi64>) -> tensor<2x4x5xi1> { @@ -4774,8 +4774,8 @@ func.func @convert_reduce_to_any(%arg0: tensor<1x2x3x4x5xi1>, %arg1: tensor<2xi6 // CHECK-SAME: %[[ARG_0:.*]]: tensor, // CHECK-SAME: %[[ARG_1:.*]]: tensor<1x2x3x4x5xi1>, // CHECK-SAME: %[[ARG_2:.*]]: tensor<2xi64>) -> tensor<2x4x5xi1> { -// CHECK-DAG: %[[DIMENSIONS:.*]] = "tf.Const"() {value = dense<[0, 2]> : tensor<2xi64>} : () -> tensor<2xi64> -// CHECK: %[[VAL_0:.*]] = "tf.Any"(%[[ARG_1]], %[[DIMENSIONS]]) {keep_dims = false} : (tensor<1x2x3x4x5xi1>, tensor<2xi64>) -> tensor<2x4x5xi1> +// CHECK-DAG: %[[DIMENSIONS:.*]] = "tf.Const"() <{value = dense<[0, 2]> : tensor<2xi64>}> : () -> tensor<2xi64> +// CHECK: %[[VAL_0:.*]] = "tf.Any"(%[[ARG_1]], %[[DIMENSIONS]]) <{keep_dims = false}> : (tensor<1x2x3x4x5xi1>, tensor<2xi64>) -> tensor<2x4x5xi1> // CHECK: %[[VAL_1:.*]] = "tf.LogicalOr"(%[[VAL_0]], %[[ARG_0]]) : (tensor<2x4x5xi1>, tensor) -> tensor<2x4x5xi1> // CHECK: return %[[VAL_1:.*]] : tensor<2x4x5xi1> // CHECK: } @@ -4790,14 +4790,14 @@ func.func @convert_reduce_to_any_non_constant_init(%arg0: tensor, %arg1: ten // CHECK-LABEL: func @convert_sort_to_topk_iota_broadcast( // CHECK-SAME: %[[ARG_0:.*]]: tensor<3x6xf32>) -> (tensor<3x6xf32>, tensor<3x6xi32>) { -// CHECK-DAG: %[[VAL_0:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor -// CHECK-DAG: %[[VAL_1:.*]] = "tf.Const"() {value = dense<6> : tensor} : () -> tensor -// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() {value = dense<1> : tensor} : () -> tensor +// CHECK-DAG: %[[VAL_0:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor +// CHECK-DAG: %[[VAL_1:.*]] = "tf.Const"() <{value = dense<6> : tensor}> : () -> tensor +// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor // CHECK: %[[VAL_3:.*]] = "tf.Range"(%cst, %cst_0, %cst_1) : (tensor, tensor, tensor) -> tensor<6xi32> // CHECK: %[[VAL_4:.*]] = arith.constant dense<[3, 6]> : tensor<2xi64> // CHECK: %[[VAL_5:.*]] = "tf.BroadcastTo"(%0, %cst_2) : (tensor<6xi32>, tensor<2xi64>) -> tensor<3x6xi32> -// CHECK: %[[K:.*]] = "tf.Const"() {value = dense<6> : tensor} : () -> tensor -// CHECK: %[[VALUES:.*]], %[[INDICES:.*]] = "tf.TopKV2"(%[[ARG_0]], %[[K]]) {sorted = true} : (tensor<3x6xf32>, tensor) -> (tensor<3x6xf32>, tensor<3x6xi32>) +// CHECK: %[[K:.*]] = "tf.Const"() <{value = dense<6> : tensor}> : () -> tensor +// CHECK: %[[VALUES:.*]], %[[INDICES:.*]] = "tf.TopKV2"(%[[ARG_0]], %[[K]]) <{sorted = true}> : (tensor<3x6xf32>, tensor) -> (tensor<3x6xf32>, tensor<3x6xi32>) // CHECK: return %[[VALUES]], %[[INDICES]] : tensor<3x6xf32>, tensor<3x6xi32> // CHECK: } func.func @convert_sort_to_topk_iota_broadcast(%arg0: tensor<3x6xf32>) -> (tensor<3x6xf32>, tensor<3x6xi32>) { @@ -4813,11 +4813,11 @@ func.func @convert_sort_to_topk_iota_broadcast(%arg0: tensor<3x6xf32>) -> (tenso // CHECK-LABEL: func @convert_sort_to_topk_iotacst_broadcast( // CHECK-SAME: %[[ARG_0:.*]]: tensor<3x6xf32>) -> (tensor<3x6xf32>, tensor<3x6xi32>) { -// CHECK-DAG: %[[VAL_0:.*]] = "tf.Const"() {value = dense<[0, 1, 2, 3, 4, 5]> : tensor<6xi32>} : () -> tensor<6xi32> +// CHECK-DAG: %[[VAL_0:.*]] = "tf.Const"() <{value = dense<[0, 1, 2, 3, 4, 5]> : tensor<6xi32>}> : () -> tensor<6xi32> // CHECK-DAG: %[[VAL_1:.*]] = arith.constant dense<[3, 6]> : tensor<2xi64> // CHECK: %[[VAL_2:.*]] = "tf.BroadcastTo"(%cst, %cst_0) : (tensor<6xi32>, tensor<2xi64>) -> tensor<3x6xi32> -// CHECK: %[[K:.*]] = "tf.Const"() {value = dense<6> : tensor} : () -> tensor -// CHECK: %[[VALUES:.*]], %[[INDICES:.*]] = "tf.TopKV2"(%[[ARG_0]], %[[K]]) {sorted = true} : (tensor<3x6xf32>, tensor) -> (tensor<3x6xf32>, tensor<3x6xi32>) +// CHECK: %[[K:.*]] = "tf.Const"() <{value = dense<6> : tensor}> : () -> tensor +// CHECK: %[[VALUES:.*]], %[[INDICES:.*]] = "tf.TopKV2"(%[[ARG_0]], %[[K]]) <{sorted = true}> : (tensor<3x6xf32>, tensor) -> (tensor<3x6xf32>, tensor<3x6xi32>) // CHECK: return %[[VALUES]], %[[INDICES]] : tensor<3x6xf32>, tensor<3x6xi32> // CHECK: } func.func @convert_sort_to_topk_iotacst_broadcast(%arg0: tensor<3x6xf32>) -> (tensor<3x6xf32>, tensor<3x6xi32>) { @@ -4833,9 +4833,9 @@ func.func @convert_sort_to_topk_iotacst_broadcast(%arg0: tensor<3x6xf32>) -> (te // CHECK-LABEL: func @convert_sort_to_topk_const( // CHECK-SAME: %[[ARG_0:.*]]: tensor<3x6xf32>) -> (tensor<3x6xf32>, tensor<3x6xi32>) { -// CHECK-DAG: %[[VAL_0:.*]] = "tf.Const"() {value = dense<{{.*}}> : tensor<3x6xi32>} : () -> tensor<3x6xi32> -// CHECK-DAG: %[[K:.*]] = "tf.Const"() {value = dense<6> : tensor} : () -> tensor -// CHECK: %[[VALUES:.*]], %[[INDICES:.*]] = "tf.TopKV2"(%[[ARG_0]], %[[K]]) {sorted = true} : (tensor<3x6xf32>, tensor) -> (tensor<3x6xf32>, tensor<3x6xi32>) +// CHECK-DAG: %[[VAL_0:.*]] = "tf.Const"() <{value = dense<{{.*}}> : tensor<3x6xi32>}> : () -> tensor<3x6xi32> +// CHECK-DAG: %[[K:.*]] = "tf.Const"() <{value = dense<6> : tensor}> : () -> tensor +// CHECK: %[[VALUES:.*]], %[[INDICES:.*]] = "tf.TopKV2"(%[[ARG_0]], %[[K]]) <{sorted = true}> : (tensor<3x6xf32>, tensor) -> (tensor<3x6xf32>, tensor<3x6xi32>) // CHECK: return %[[VALUES]], %[[INDICES]] : tensor<3x6xf32>, tensor<3x6xi32> // CHECK: } func.func @convert_sort_to_topk_const(%arg0: tensor<3x6xf32>) -> (tensor<3x6xf32>, tensor<3x6xi32>) { @@ -4913,7 +4913,7 @@ func.func @not_convert_remainder_for_uint64(%arg0: tensor<10x8xui64>, %arg1: ten // CHECK-LABEL: func @convert_population_count_i32( // CHECK-SAME: %[[ARG_0:.*]]: tensor<8xi32> // CHECK: %[[POP_CNT:.*]] = "tf.PopulationCount"(%[[ARG_0]]) : (tensor<8xi32>) -> tensor<8xui8> -// CHECK: %[[RES:.*]] = "tf.Cast"(%[[POP_CNT]]) {Truncate = false} : (tensor<8xui8>) -> tensor<8xi32> +// CHECK: %[[RES:.*]] = "tf.Cast"(%[[POP_CNT]]) <{Truncate = false}> : (tensor<8xui8>) -> tensor<8xi32> // CHECK: return %[[RES]] // CHECK: } func.func @convert_population_count_i32(%arg0: tensor<8xi32>) -> tensor<8xi32> { @@ -4932,8 +4932,8 @@ func.func @convert_population_count_ui8(%arg0: tensor<8xui8>) -> tensor<8xui8> { } // CHECK-LABEL: func @torch_index_select( -// CHECK: %[[AXIS:.+]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor -// CHECK: %[[RES:.+]] = "tf.GatherV2"(%arg0, %arg1, %[[AXIS]]) {batch_dims = 0 : i64} +// CHECK: %[[AXIS:.+]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor +// CHECK: %[[RES:.+]] = "tf.GatherV2"(%arg0, %arg1, %[[AXIS]]) <{batch_dims = 0 : i64}> // CHECK: return %[[RES]] func.func @torch_index_select(%arg0: tensor<2x1xf32>, %arg1: tensor<2xi32>) -> tensor<2x1xf32> { @@ -4945,9 +4945,9 @@ func.func @torch_index_select(%arg0: tensor<2x1xf32>, %arg1: tensor<2xi32>) -> t // CHECK-LABEL: func @lowered_cumsum( // CHECK-SAME: %[[VAL_0:.*]]: tensor<4x12xf32>) -> tensor<4x12xf32> { -// CHECK-DAG: %[[VAL_1:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor} : () -> tensor -// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor -// CHECK: %[[VAL_3:.*]] = "tf.Cumsum"(%[[VAL_0]], %[[VAL_2]]) {exclusive = false, reverse = false} : (tensor<4x12xf32>, tensor) -> tensor<4x12xf32> +// CHECK-DAG: %[[VAL_1:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor}> : () -> tensor +// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor +// CHECK: %[[VAL_3:.*]] = "tf.Cumsum"(%[[VAL_0]], %[[VAL_2]]) <{exclusive = false, reverse = false}> : (tensor<4x12xf32>, tensor) -> tensor<4x12xf32> // CHECK: return %[[VAL_3]] : tensor<4x12xf32> // CHECK: } func.func @lowered_cumsum(%arg0: tensor<4x12xf32>) -> tensor<4x12xf32> { @@ -4962,9 +4962,9 @@ func.func @lowered_cumsum(%arg0: tensor<4x12xf32>) -> tensor<4x12xf32> { // CHECK-LABEL: func @lowered_cumsum_trivial_attrs( // CHECK-SAME: %[[VAL_0:.*]]: tensor<4x12xf32>) -> tensor<4x12xf32> { -// CHECK-DAG: %[[VAL_1:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor} : () -> tensor -// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor -// CHECK: %[[VAL_3:.*]] = "tf.Cumsum"(%[[VAL_0]], %[[VAL_2]]) {exclusive = false, reverse = false} : (tensor<4x12xf32>, tensor) -> tensor<4x12xf32> +// CHECK-DAG: %[[VAL_1:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor}> : () -> tensor +// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor +// CHECK: %[[VAL_3:.*]] = "tf.Cumsum"(%[[VAL_0]], %[[VAL_2]]) <{exclusive = false, reverse = false}> : (tensor<4x12xf32>, tensor) -> tensor<4x12xf32> // CHECK: return %[[VAL_3]] : tensor<4x12xf32> // CHECK: } func.func @lowered_cumsum_trivial_attrs(%arg0: tensor<4x12xf32>) -> tensor<4x12xf32> { @@ -4979,9 +4979,9 @@ func.func @lowered_cumsum_trivial_attrs(%arg0: tensor<4x12xf32>) -> tensor<4x12x // CHECK-LABEL: func @lowered_cumprod( // CHECK-SAME: %[[VAL_0:.*]]: tensor<4x12xf32>) -> tensor<4x12xf32> { -// CHECK-DAG: %[[VAL_1:.*]] = "tf.Const"() {value = dense<1.000000e+00> : tensor} : () -> tensor -// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() {value = dense<1> : tensor} : () -> tensor -// CHECK: %[[VAL_3:.*]] = "tf.Cumprod"(%[[VAL_0]], %[[VAL_2]]) {exclusive = false, reverse = false} : (tensor<4x12xf32>, tensor) -> tensor<4x12xf32> +// CHECK-DAG: %[[VAL_1:.*]] = "tf.Const"() <{value = dense<1.000000e+00> : tensor}> : () -> tensor +// CHECK-DAG: %[[VAL_2:.*]] = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor +// CHECK: %[[VAL_3:.*]] = "tf.Cumprod"(%[[VAL_0]], %[[VAL_2]]) <{exclusive = false, reverse = false}> : (tensor<4x12xf32>, tensor) -> tensor<4x12xf32> // CHECK: return %[[VAL_3]] : tensor<4x12xf32> // CHECK: } func.func @lowered_cumprod(%arg0: tensor<4x12xf32>) -> tensor<4x12xf32> { @@ -5022,7 +5022,7 @@ func.func @convert_dot_quant_type(%arg0: tensor<1x256xf32>, %arg1: tensor<256x!q // CHECK-LABEL: func @get_dimension_size( // CHECK-SAME: %[[ARG_0:.*]]: tensor<4x256x?xf32>) -> tensor { -// CHECK %[[CST_0:.*]] = "tf.Const"() {value = dense<256> : tensor} : () -> tensor +// CHECK %[[CST_0:.*]] = "tf.Const"() <{value = dense<256> : tensor}> : () -> tensor // CHECK return %[[CST_0]] : tensor func.func @get_dimension_size(%arg0: tensor<4x256x?xf32>) -> tensor { %0 = "mhlo.get_dimension_size"(%arg0) {dimension = 1 : i64} : (tensor<4x256x?xf32>) -> tensor @@ -5032,10 +5032,10 @@ func.func @get_dimension_size(%arg0: tensor<4x256x?xf32>) -> tensor { // CHECK-LABEL: func @get_dimension_size_dynamic( // CHECK-SAME: %[[ARG_0:.*]]: tensor<4x256x?xf32>) -> tensor { // CHECK %[[VAL_0:.*]] = "tf.Shape"(%[[ARG_0]]) : (tensor<4x256x?xf32>) -> tensor<3xi32> -// CHECK %[[CST_0:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32> -// CHECK %[[CST_1:.*]] = "tf.Const"() {value = dense<2> : tensor<1xi64>} : () -> tensor<1xi64> +// CHECK %[[CST_0:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}> : () -> tensor<1xi32> +// CHECK %[[CST_1:.*]] = "tf.Const"() <{value = dense<2> : tensor<1xi64>}> : () -> tensor<1xi64> // CHECK %[[VAL_1:.*]] = "tf.Slice"(%[[VAL_0]], %[[CST_1]], %[[CST_0]]) : (tensor<3xi32>, tensor<1xi64>, tensor<1xi32>) -> tensor<1xi32> -// CHECK %[[VAL_2:.*]] = "tf.Squeeze"(%[[VAL_1]]) {squeeze_dims = [0]} : (tensor<1xi32>) -> tensor +// CHECK %[[VAL_2:.*]] = "tf.Squeeze"(%[[VAL_1]]) <{squeeze_dims = [0]}> : (tensor<1xi32>) -> tensor // CHECK return %[[VAL_2]] : tensor func.func @get_dimension_size_dynamic(%arg0: tensor<4x256x?xf32>) -> tensor { %0 = "mhlo.get_dimension_size"(%arg0) {dimension = 2 : i64} : (tensor<4x256x?xf32>) -> tensor @@ -5044,10 +5044,10 @@ func.func @get_dimension_size_dynamic(%arg0: tensor<4x256x?xf32>) -> tensor // CHECK-LABEL: func @dynamic_iota_i32_1d( // CHECK-SAME: %[[ARG_0:.*]]: tensor<1xi32>) -> tensor { -// CHECK-DAG: %[[CST_0:.*]] = "tf.Const"() {value = dense<> : tensor<0xi32>} : () -> tensor<0xi32> +// CHECK-DAG: %[[CST_0:.*]] = "tf.Const"() <{value = dense<> : tensor<0xi32>}> : () -> tensor<0xi32> // CHECK: %[[VAL_0:.*]] = "tf.Reshape"(%arg0, %[[CST_0]]) : (tensor<1xi32>, tensor<0xi32>) -> tensor -// CHECK-DAG: %[[CST_1:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor -// CHECK-DAG: %[[CST_2:.*]] = "tf.Const"() {value = dense<1> : tensor} : () -> tensor +// CHECK-DAG: %[[CST_1:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor +// CHECK-DAG: %[[CST_2:.*]] = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor // CHECK: %[[VAL_1:.*]] = "tf.Range"(%[[CST_1]], %[[VAL_0]], %[[CST_2]]) : (tensor, tensor, tensor) -> tensor // CHECK: return %[[VAL_1]] : tensor func.func @dynamic_iota_i32_1d(%arg0: tensor<1xi32>) -> tensor { @@ -5057,11 +5057,11 @@ func.func @dynamic_iota_i32_1d(%arg0: tensor<1xi32>) -> tensor { // CHECK-LABEL: func @dynamic_iota_f32_1d( // CHECK-SAME: %[[ARG_0:.*]]: tensor<1xi32>) -> tensor { -// CHECK: %[[VAL_0:.*]] = "tf.Cast"(%arg0) {Truncate = false} : (tensor<1xi32>) -> tensor<1xf32> -// CHECK-DAG: %[[CST_0:.*]] = "tf.Const"() {value = dense<> : tensor<0xi32>} : () -> tensor<0xi32> +// CHECK: %[[VAL_0:.*]] = "tf.Cast"(%arg0) <{Truncate = false}> : (tensor<1xi32>) -> tensor<1xf32> +// CHECK-DAG: %[[CST_0:.*]] = "tf.Const"() <{value = dense<> : tensor<0xi32>}> : () -> tensor<0xi32> // CHECK: %[[VAL_1:.*]] = "tf.Reshape"(%[[VAL_0]], %[[CST_0]]) : (tensor<1xf32>, tensor<0xi32>) -> tensor -// CHECK-DAG: %[[CST_1:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor} : () -> tensor -// CHECK-DAG: %[[CST_2:.*]] = "tf.Const"() {value = dense<1.000000e+00> : tensor} : () -> tensor +// CHECK-DAG: %[[CST_1:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor}> : () -> tensor +// CHECK-DAG: %[[CST_2:.*]] = "tf.Const"() <{value = dense<1.000000e+00> : tensor}> : () -> tensor // CHECK: %[[VAL_2:.*]] = "tf.Range"(%[[CST_1]], %[[VAL_1]], %[[CST_2]]) : (tensor, tensor, tensor) -> tensor // CHECK: return %[[VAL_2]] : tensor func.func @dynamic_iota_f32_1d(%arg0: tensor<1xi32>) -> tensor { @@ -5073,8 +5073,8 @@ func.func @dynamic_iota_f32_1d(%arg0: tensor<1xi32>) -> tensor { // CHECK-SAME: %arg0: tensor<1x?x4x256xf32>, // CHECK-SAME: %arg1: tensor<4xi32>, // CHECK-SAME: %arg2: tensor<4xi32>) -> tensor<1x?x4x128xf32> { -// CHECK: %cst = "tf.Const"() {value = dense<1> : tensor<4xi32>} : () -> tensor<4xi32> -// CHECK: %0 = "tf.StridedSlice"(%arg0, %arg1, %arg2, %cst) {begin_mask = 0 : i64, ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 0 : i64} : (tensor<1x?x4x256xf32>, tensor<4xi32>, tensor<4xi32>, tensor<4xi32>) -> tensor<1x?x4x128xf32> +// CHECK: %cst = "tf.Const"() <{value = dense<1> : tensor<4xi32>}> : () -> tensor<4xi32> +// CHECK: %0 = "tf.StridedSlice"(%arg0, %arg1, %arg2, %cst) <{begin_mask = 0 : i64, ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 0 : i64}> : (tensor<1x?x4x256xf32>, tensor<4xi32>, tensor<4xi32>, tensor<4xi32>) -> tensor<1x?x4x128xf32> // CHECK: return %0 : tensor<1x?x4x128xf32> func.func @real_dynamic_slice_strides_equal_to_1_signed(%arg0: tensor<1x?x4x256xf32>, %arg1: tensor<4xi32>, %arg2: tensor<4xi32>) -> tensor<1x?x4x128xf32> { %cst = mhlo.constant dense<1> : tensor<4xi32> @@ -5086,8 +5086,8 @@ func.return %0 : tensor<1x?x4x128xf32> // CHECK-SAME: %arg0: tensor<1x?x2x4xf32>, // CHECK-SAME: %arg1: tensor<4xi32>, // CHECK-SAME: %arg2: tensor<4xi32>) -> tensor<1x?x1x2xf32> { -// CHECK %cst = "tf.Const"() {value = dense<2> : tensor<4xi32>} : () -> tensor<4xi32> -// CHECK %0 = "tf.StridedSlice"(%arg0, %arg1, %arg2, %cst) {begin_mask = 0 : i64, ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 0 : i64} : (tensor<1x?x2x4xf32>, tensor<4xi32>, tensor<4xi32>, tensor<4xi32>) -> tensor<1x?x1x2xf32> +// CHECK %cst = "tf.Const"() <{value = dense<2> : tensor<4xi32>}> : () -> tensor<4xi32> +// CHECK %0 = "tf.StridedSlice"(%arg0, %arg1, %arg2, %cst) <{begin_mask = 0 : i64, ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 0 : i64}> : (tensor<1x?x2x4xf32>, tensor<4xi32>, tensor<4xi32>, tensor<4xi32>) -> tensor<1x?x1x2xf32> // CHECK return %0 : tensor<1x?x1x2xf32> func.func @real_dynamic_slice_strides_not_equal_to_1(%arg0: tensor<1x?x2x4xf32>, %arg1: tensor<4xi32>, %arg2: tensor<4xi32>) -> tensor<1x?x1x2xf32> { %cst = mhlo.constant dense<2> : tensor<4xi32> @@ -5113,7 +5113,7 @@ func.func @remove_shape_assertion_custom_call(%arg1: tensor) -> tensor< // CHECK-SAME: %[[ARG_1:.*]]: tensor<1x4xi32>, // CHECK-SAME: %[[ARG_2:.*]]: tensor, // CHECK-SAME: %[[ARG_3:.*]]: tensor) -> (tensor<1x4xf32>, tensor<1x4xi32>) { -// CHECK: %[[VALUES:.*]], %[[INDICES:.*]] = "tf.ApproxTopK"(%[[ARG_0]]) {aggregate_to_topk = true, is_max_k = true, k = 4 : i64, recall_target = 8.500000e-01 : f32, reduction_dimension = 1 : i64, reduction_input_size_override = -1 : i64} : (tensor<1x4xf32>) -> (tensor<1x4xf32>, tensor<1x4xi32>) +// CHECK: %[[VALUES:.*]], %[[INDICES:.*]] = "tf.ApproxTopK"(%[[ARG_0]]) <{aggregate_to_topk = true, is_max_k = true, k = 4 : i64, recall_target = 8.500000e-01 : f32, reduction_dimension = 1 : i64, reduction_input_size_override = -1 : i64}> : (tensor<1x4xf32>) -> (tensor<1x4xf32>, tensor<1x4xi32>) // CHECK: return %[[VALUES]], %[[INDICES]] : tensor<1x4xf32>, tensor<1x4xi32> // CHECK: } func.func @convert_approx_top_k_custom_call(%arg0: tensor<1x4xf32>, %arg1: tensor<1x4xi32>, %arg2: tensor, %arg3: tensor) -> (tensor<1x4xf32>, tensor<1x4xi32>) { diff --git a/tensorflow/compiler/mlir/lite/stablehlo/tests/tf-tfl-translate-serialize-stablehlo.mlir b/tensorflow/compiler/mlir/lite/stablehlo/tests/tf-tfl-translate-serialize-stablehlo.mlir index e822963984cf62..4e12ffd931c5f2 100644 --- a/tensorflow/compiler/mlir/lite/stablehlo/tests/tf-tfl-translate-serialize-stablehlo.mlir +++ b/tensorflow/compiler/mlir/lite/stablehlo/tests/tf-tfl-translate-serialize-stablehlo.mlir @@ -1,9 +1,22 @@ -//RUN: tf_tfl_translate --enable-stablehlo-conversion --input-mlir %s -o /tmp/temp.stablehlo; [ -f /tmp/temp.stablehlo ] +//RUN: tf_tfl_translate --enable-stablehlo-conversion --input-mlir %s -o - | flatbuffer_translate --tflite-flatbuffer-to-mlir - -o - | FileCheck %s module { -func.func @main(%arg0: tensor<2xi32>) -> tensor<2xi32> { - %0 = "tf.Add"(%arg0, %arg0) : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi32> - func.return %0 : tensor<2xi32> +func.func @tfInplaceUpdate(%arg0: tensor<2x1x2xf32>) -> tensor<2x1x2xf32> { + %1 = arith.constant dense<1> : tensor<1xi32> + %2 = arith.constant dense<2.0> : tensor<1x1x2xf32> + %3 = "tf.InplaceUpdate"(%arg0, %1, %2) {device = ""} + : (tensor<2x1x2xf32>, tensor<1xi32>, tensor<1x1x2xf32>) -> tensor<2x1x2xf32> + func.return %3 : tensor<2x1x2xf32> } -} \ No newline at end of file +} + +//CHECK: module attributes {tfl.description = "MLIR Converted.", tfl.metadata = {keep_stablehlo_constant = "true"}, tfl.schema_version = 3 : i32} { +//CHECK-NEXT: func.func @main(%arg0: tensor<2x1x2xf32>) -> tensor<2x1x2xf32> attributes {tf.entry_function = {inputs = "arg0", outputs = "stablehlo.dynamic_update_slice"}} { +//CHECK-DAG: %0 = stablehlo.constant dense<2.000000e+00> : tensor<1x1x2xf32> +//CHECK-DAG: %1 = stablehlo.constant dense<1> : tensor +//CHECK-DAG: %2 = stablehlo.constant dense<0> : tensor +//CHECK-NEXT: %3 = stablehlo.dynamic_update_slice %arg0, %0, %1, %2, %2 : (tensor<2x1x2xf32>, tensor<1x1x2xf32>, tensor, tensor, tensor) -> tensor<2x1x2xf32> +//CHECK-NEXT: return %3 : tensor<2x1x2xf32> +//CHECK-NEXT: } +//CHECK-NEXT:} \ No newline at end of file diff --git a/tensorflow/compiler/mlir/lite/stablehlo/tests/unfuse_mhlo_batch_norm.mlir b/tensorflow/compiler/mlir/lite/stablehlo/tests/unfuse_mhlo_batch_norm.mlir index bd2f56dd930bb5..073f31e39786d9 100644 --- a/tensorflow/compiler/mlir/lite/stablehlo/tests/unfuse_mhlo_batch_norm.mlir +++ b/tensorflow/compiler/mlir/lite/stablehlo/tests/unfuse_mhlo_batch_norm.mlir @@ -132,13 +132,13 @@ func.func @batchNormInference_f16_overflow( func.func @batchNormTraining_4D_middle_features( %x: tensor<3x4x256x6xf32>, %scale: tensor<256xf32>, %offset: tensor<256xf32>) -> (tensor<3x4x256x6xf32>) { - // CHECK-DAG: %[[CST_AXIS:.+]] = "tf.Const"() {value = dense<[0, 1, 3]> : tensor<3xi32>} : () -> tensor<3xi32> + // CHECK-DAG: %[[CST_AXIS:.+]] = "tf.Const"() <{value = dense<[0, 1, 3]> : tensor<3xi32>}> : () -> tensor<3xi32> // CHECK-DAG: %[[X_SHAPE:.+]] = shape.const_shape [3, 4, 256, 6] : tensor<4xindex> // CHECK-DAG: %[[EPS:.+]] = mhlo.constant dense<1.000000e+00> : tensor<256xf32> - // CHECK-DAG: %[[MEAN:.+]] = "tf.Mean"(%arg0, %[[CST_AXIS]]) {keep_dims = false} : (tensor<3x4x256x6xf32>, tensor<3xi32>) -> tensor<256xf32> + // CHECK-DAG: %[[MEAN:.+]] = "tf.Mean"(%arg0, %[[CST_AXIS]]) <{keep_dims = false}> : (tensor<3x4x256x6xf32>, tensor<3xi32>) -> tensor<256xf32> // CHECK-DAG: %[[MEAN_BCAST:.+]] = "mhlo.dynamic_broadcast_in_dim"(%[[MEAN]], %[[X_SHAPE]]) {broadcast_dimensions = dense<2> : tensor<1xi64>} : (tensor<256xf32>, tensor<4xindex>) -> tensor<3x4x256x6xf32> // CHECK-DAG: %[[SQ_DIFF:.+]] = "tf.SquaredDifference"(%arg0, %[[MEAN_BCAST]]) : (tensor<3x4x256x6xf32>, tensor<3x4x256x6xf32>) -> tensor<3x4x256x6xf32> - // CHECK-DAG: %[[VARIANCE:.+]] = "tf.Mean"(%[[SQ_DIFF]], %[[CST_AXIS]]) {keep_dims = false} : (tensor<3x4x256x6xf32>, tensor<3xi32>) -> tensor<256xf32> + // CHECK-DAG: %[[VARIANCE:.+]] = "tf.Mean"(%[[SQ_DIFF]], %[[CST_AXIS]]) <{keep_dims = false}> : (tensor<3x4x256x6xf32>, tensor<3xi32>) -> tensor<256xf32> // CHECK-DAG: %[[VARIANCE_EPS:.+]] = mhlo.add %[[VARIANCE]], %[[EPS]] : tensor<256xf32> // CHECK-DAG: %[[VARIANCE_EPS_RSQRT:.+]] = mhlo.rsqrt %[[VARIANCE_EPS]] : tensor<256xf32> // CHECK-DAG: %[[MULTIPLIER:.+]] = mhlo.multiply %[[VARIANCE_EPS_RSQRT]], %[[SCALE]] : tensor<256xf32> @@ -152,4 +152,4 @@ func.func @batchNormTraining_4D_middle_features( {epsilon = 1.0 : f32, feature_index = 2 : i64} : (tensor<3x4x256x6xf32>, tensor<256xf32>, tensor<256xf32>) -> (tensor<3x4x256x6xf32>, tensor<256xf32>, tensor<256xf32>) func.return %0 : tensor<3x4x256x6xf32> -} \ No newline at end of file +} diff --git a/tensorflow/compiler/mlir/lite/stablehlo/transforms/transforms.cc b/tensorflow/compiler/mlir/lite/stablehlo/transforms/transforms.cc index 562afd53f6f76c..df3a5f62e8ff59 100644 --- a/tensorflow/compiler/mlir/lite/stablehlo/transforms/transforms.cc +++ b/tensorflow/compiler/mlir/lite/stablehlo/transforms/transforms.cc @@ -18,6 +18,7 @@ limitations under the License. #include "mlir/Pass/PassManager.h" // from @llvm-project #include "mlir/Transforms/Passes.h" // from @llvm-project #include "tensorflow/compiler/mlir/lite/stablehlo/transforms/drop_savedmodel_semantics.h" +#include "tensorflow/compiler/mlir/lite/stablehlo/transforms/legalize_tf_xla_call_module_to_stablehlo_pass.h" #include "tensorflow/compiler/mlir/lite/stablehlo/transforms/passes.h" #include "tensorflow/compiler/mlir/lite/stablehlo/transforms/rename_entrypoint_to_main.h" #include "tensorflow/compiler/mlir/lite/stablehlo/transforms/smuggle_disallowed_ops.h" @@ -34,6 +35,9 @@ namespace odml { void AddTFToStablehloPasses(OpPassManager& pm, bool skip_resize, bool smuggle_disallowed_ops) { pm.addPass(CreateRenameEntrypointToMainPass()); + + // if the input is a call_xla_module, then unwrap the content + pm.addPass(mlir::odml::CreateLegalizeTFXlaCallModuleToStablehloPass()); // TODO(b/230572023): Consider improving shape inference for While op instead // of dropping the attribute. This need not be correct for models not trained // on TPU. diff --git a/tensorflow/compiler/mlir/lite/tests/dilated-conv.mlir b/tensorflow/compiler/mlir/lite/tests/dilated-conv.mlir index ccb5507a8ee5c5..378ed7fb2a46e1 100644 --- a/tensorflow/compiler/mlir/lite/tests/dilated-conv.mlir +++ b/tensorflow/compiler/mlir/lite/tests/dilated-conv.mlir @@ -10,7 +10,7 @@ func.func @testDilatedConv(%arg0: tensor<1x128x128x3xf32>, %arg1: tensor<5x5x3x8 // CHECK-LABEL: testDilatedConv // CHECK-SAME: ([[INPUT:%.*]]: tensor<1x128x128x3xf32>, [[FILTER:%.*]]: tensor<5x5x3x8xf32>) - // CHECK-NEXT: [[RESULT:%.*]] = "tf.Conv2D"([[INPUT]], [[FILTER]]) {dilations = [1, 2, 2, 1], padding = "VALID", strides = [1, 1, 1, 1]} : (tensor<1x128x128x3xf32>, tensor<5x5x3x8xf32>) -> tensor<1x120x120x8xf32> + // CHECK-NEXT: [[RESULT:%.*]] = "tf.Conv2D"([[INPUT]], [[FILTER]]) <{dilations = [1, 2, 2, 1], padding = "VALID", strides = [1, 1, 1, 1]}> : (tensor<1x128x128x3xf32>, tensor<5x5x3x8xf32>) -> tensor<1x120x120x8xf32> // CHECK-NEXT: return [[RESULT]] : tensor<1x120x120x8xf32> } @@ -24,7 +24,7 @@ func.func @testDilatedConvWithNonConstantPadAndCrops(%arg0: tensor<1x128x128x3xf // CHECK-LABEL: testDilatedConvWithNonConstantPadAndCrops // CHECK-SAME: ([[INPUT:%.*]]: tensor<1x128x128x3xf32>, [[FILTER:%.*]]: tensor<5x5x3x8xf32>) - // CHECK-NEXT: [[RESULT:%.*]] = "tf.Conv2D"([[INPUT]], [[FILTER]]) {dilations = [1, 2, 2, 1], padding = "VALID", strides = [1, 1, 1, 1]} : (tensor<1x128x128x3xf32>, tensor<5x5x3x8xf32>) -> tensor<1x120x120x8xf32> + // CHECK-NEXT: [[RESULT:%.*]] = "tf.Conv2D"([[INPUT]], [[FILTER]]) <{dilations = [1, 2, 2, 1], padding = "VALID", strides = [1, 1, 1, 1]}> : (tensor<1x128x128x3xf32>, tensor<5x5x3x8xf32>) -> tensor<1x120x120x8xf32> // CHECK-NEXT: return [[RESULT]] : tensor<1x120x120x8xf32> } @@ -39,7 +39,7 @@ func.func @testDilatedConvWithNonZeroBasePadding(%arg0: tensor<1x128x128x3xf32>, // CHECK-LABEL: testDilatedConvWithNonZeroBasePadding // CHECK-SAME: ([[INPUT:%.*]]: tensor<1x128x128x3xf32>, [[FILTER:%.*]]: tensor<5x5x3x8xf32>) - // CHECK-NEXT: [[RESULT:%.*]] = "tf.Conv2D"([[INPUT]], [[FILTER]]) {dilations = [1, 2, 2, 1], padding = "SAME", strides = [1, 1, 1, 1]} : (tensor<1x128x128x3xf32>, tensor<5x5x3x8xf32>) -> tensor<1x128x128x8xf32> + // CHECK-NEXT: [[RESULT:%.*]] = "tf.Conv2D"([[INPUT]], [[FILTER]]) <{dilations = [1, 2, 2, 1], padding = "SAME", strides = [1, 1, 1, 1]}> : (tensor<1x128x128x3xf32>, tensor<5x5x3x8xf32>) -> tensor<1x128x128x8xf32> // CHECK-NEXT: return [[RESULT]] : tensor<1x128x128x8xf32> } @@ -54,7 +54,7 @@ func.func @testDilatedConvWithFp16(%arg0 : tensor<1x20x30x40xf16>, %arg1: tensor // CHECK-LABEL: testDilatedConvWithFp16 // CHECK-SAME: ([[INPUT:%.*]]: tensor<1x20x30x40xf16>, [[FILTER:%.*]]: tensor<5x5x40x32xf16>) - // CHECK-NEXT: [[RESULT:%.*]] = "tf.Conv2D"([[INPUT]], [[FILTER]]) {data_format = "NHWC", dilations = [1, 2, 2, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 1, 1]} : (tensor<1x20x30x40xf16>, tensor<5x5x40x32xf16>) -> tensor<1x20x30x32xf16> + // CHECK-NEXT: [[RESULT:%.*]] = "tf.Conv2D"([[INPUT]], [[FILTER]]) <{data_format = "NHWC", dilations = [1, 2, 2, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 1, 1]}> : (tensor<1x20x30x40xf16>, tensor<5x5x40x32xf16>) -> tensor<1x20x30x32xf16> // CHECK-NEXT: return [[RESULT]] : tensor<1x20x30x32xf16> } @@ -85,7 +85,7 @@ func.func @testDilatedDepthWiseConv(%arg0: tensor<1x128x128x3xf32>, %arg1: tenso // CHECK-LABEL: testDilatedDepthWiseConv // CHECK-SAME: ([[INPUT:%.*]]: tensor<1x128x128x3xf32>, [[FILTER:%.*]]: tensor<5x5x3x8xf32>) - // CHECK-NEXT: [[RESULT:%.*]] = "tf.DepthwiseConv2dNative"([[INPUT]], [[FILTER]]) {dilations = [1, 2, 2, 1], padding = "SAME", strides = [1, 1, 1, 1]} : (tensor<1x128x128x3xf32>, tensor<5x5x3x8xf32>) -> tensor<1x128x128x8xf32> + // CHECK-NEXT: [[RESULT:%.*]] = "tf.DepthwiseConv2dNative"([[INPUT]], [[FILTER]]) <{dilations = [1, 2, 2, 1], padding = "SAME", strides = [1, 1, 1, 1]}> : (tensor<1x128x128x3xf32>, tensor<5x5x3x8xf32>) -> tensor<1x128x128x8xf32> // CHECK-NEXT: return [[RESULT]] : tensor<1x128x128x8xf32> } @@ -103,7 +103,7 @@ func.func @testDilatedConvWithPad(%arg0: tensor<1x128x128x3xf32>, %arg1: tensor< // CHECK-LABEL: testDilatedConvWithPad // CHECK-SAME: ([[INPUT:%.*]]: tensor<1x128x128x3xf32>, [[FILTER:%.*]]: tensor<5x5x3x8xf32>, [[BIAS:%.*]]: tensor<8xf32>) - // CHECK-NEXT: [[CONV:%.*]] = "tf.Conv2D"([[INPUT]], [[FILTER]]) {dilations = [1, 2, 2, 1], padding = "SAME", strides = [1, 1, 1, 1]} : (tensor<1x128x128x3xf32>, tensor<5x5x3x8xf32>) -> tensor<1x128x128x8xf32> + // CHECK-NEXT: [[CONV:%.*]] = "tf.Conv2D"([[INPUT]], [[FILTER]]) <{dilations = [1, 2, 2, 1], padding = "SAME", strides = [1, 1, 1, 1]}> : (tensor<1x128x128x3xf32>, tensor<5x5x3x8xf32>) -> tensor<1x128x128x8xf32> // CHECK-NEXT: [[RESULT:%.*]] = "tf.BiasAdd"([[CONV]], [[BIAS]]) : (tensor<1x128x128x8xf32>, tensor<8xf32>) -> tensor<1x128x128x8xf32> // CHECK-NEXT: return [[RESULT]] : tensor<1x128x128x8xf32> } @@ -122,7 +122,7 @@ func.func @testDilatedDepthWiseConvWithPad(%arg0: tensor<1x128x128x3xf32>, %arg1 // CHECK-LABEL: testDilatedDepthWiseConvWithPad // CHECK-SAME: ([[INPUT:%.*]]: tensor<1x128x128x3xf32>, [[FILTER:%.*]]: tensor<5x5x3x8xf32>, [[BIAS:%.*]]: tensor<8xf32>) - // CHECK-NEXT: [[CONV:%.*]] = "tf.DepthwiseConv2dNative"([[INPUT]], [[FILTER]]) {dilations = [1, 2, 2, 1], padding = "SAME", strides = [1, 1, 1, 1]} : (tensor<1x128x128x3xf32>, tensor<5x5x3x8xf32>) -> tensor<1x128x128x8xf32> + // CHECK-NEXT: [[CONV:%.*]] = "tf.DepthwiseConv2dNative"([[INPUT]], [[FILTER]]) <{dilations = [1, 2, 2, 1], padding = "SAME", strides = [1, 1, 1, 1]}> : (tensor<1x128x128x3xf32>, tensor<5x5x3x8xf32>) -> tensor<1x128x128x8xf32> // CHECK-NEXT: [[RESULT:%.*]] = "tf.BiasAdd"([[CONV]], [[BIAS]]) : (tensor<1x128x128x8xf32>, tensor<8xf32>) -> tensor<1x128x128x8xf32> // CHECK-NEXT: return [[RESULT]] : tensor<1x128x128x8xf32> } @@ -139,7 +139,7 @@ func.func @testDilatedConvWithBiasAdd(%arg0: tensor<1x128x128x3xf32>, %arg1: ten // CHECK-LABEL: testDilatedConvWithBiasAdd // CHECK-SAME: ([[INPUT:%.*]]: tensor<1x128x128x3xf32>, [[FILTER:%.*]]: tensor<5x5x3x8xf32>, [[BIAS:%.*]]: tensor<8xf32>) - // CHECK-NEXT: [[CONV:%.*]] = "tf.Conv2D"([[INPUT]], [[FILTER]]) {dilations = [1, 2, 2, 1], padding = "SAME", strides = [1, 1, 1, 1]} : (tensor<1x128x128x3xf32>, tensor<5x5x3x8xf32>) -> tensor<1x128x128x8xf32> + // CHECK-NEXT: [[CONV:%.*]] = "tf.Conv2D"([[INPUT]], [[FILTER]]) <{dilations = [1, 2, 2, 1], padding = "SAME", strides = [1, 1, 1, 1]}> : (tensor<1x128x128x3xf32>, tensor<5x5x3x8xf32>) -> tensor<1x128x128x8xf32> // CHECK-NEXT: [[RESULT:%.*]] = "tf.BiasAdd"([[CONV]], [[BIAS]]) : (tensor<1x128x128x8xf32>, tensor<8xf32>) -> tensor<1x128x128x8xf32> // CHECK-NEXT: return [[RESULT]] : tensor<1x128x128x8xf32> } @@ -156,7 +156,7 @@ func.func @testDilatedDepthWiseConvWithBiasAdd(%arg0: tensor<1x128x128x3xf32>, % // CHECK-LABEL: testDilatedDepthWiseConvWithBiasAdd // CHECK-SAME: ([[INPUT:%.*]]: tensor<1x128x128x3xf32>, [[FILTER:%.*]]: tensor<5x5x3x8xf32>, [[BIAS:%.*]]: tensor<8xf32>) - // CHECK-NEXT: [[CONV:%.*]] = "tf.DepthwiseConv2dNative"([[INPUT]], [[FILTER]]) {dilations = [1, 2, 2, 1], padding = "SAME", strides = [1, 1, 1, 1]} : (tensor<1x128x128x3xf32>, tensor<5x5x3x8xf32>) -> tensor<1x128x128x8xf32> + // CHECK-NEXT: [[CONV:%.*]] = "tf.DepthwiseConv2dNative"([[INPUT]], [[FILTER]]) <{dilations = [1, 2, 2, 1], padding = "SAME", strides = [1, 1, 1, 1]}> : (tensor<1x128x128x3xf32>, tensor<5x5x3x8xf32>) -> tensor<1x128x128x8xf32> // CHECK-NEXT: [[RESULT:%.*]] = "tf.BiasAdd"([[CONV]], [[BIAS]]) : (tensor<1x128x128x8xf32>, tensor<8xf32>) -> tensor<1x128x128x8xf32> // CHECK-NEXT: return [[RESULT]] : tensor<1x128x128x8xf32> } @@ -176,10 +176,10 @@ func.func @testDilatedConvWithExpandSqueeze1(%arg0: tensor<1x128x128xf32>, %arg1 // CHECK-LABEL: testDilatedConvWithExpandSqueeze1 // CHECK-SAME: ([[INPUT:%.*]]: tensor<1x128x128xf32>, [[FILTER:%.*]]: tensor<5x5x1x1xf32>, [[BIAS:%.*]]: tensor<128xf32>) - // CHECK-NEXT: [[AXIS:%.*]] = "tf.Const"() {value = dense<3> : tensor} : () -> tensor + // CHECK-NEXT: [[AXIS:%.*]] = "tf.Const"() <{value = dense<3> : tensor}> : () -> tensor // CHECK-NEXT: [[EXPAND:%.*]] = "tf.ExpandDims"([[INPUT]], [[AXIS]]) : (tensor<1x128x128xf32>, tensor) -> tensor<1x128x128x1xf32> - // CHECK-NEXT: [[CONV:%.*]] = "tf.Conv2D"([[EXPAND]], [[FILTER]]) {dilations = [1, 2, 2, 1], padding = "SAME", strides = [1, 1, 1, 1]} : (tensor<1x128x128x1xf32>, tensor<5x5x1x1xf32>) -> tensor<1x128x128x1xf32> - // CHECK-NEXT: [[SQUEEZE:%.*]] = "tf.Squeeze"([[CONV]]) {squeeze_dims = [3]} : (tensor<1x128x128x1xf32>) -> tensor<1x128x128xf32> + // CHECK-NEXT: [[CONV:%.*]] = "tf.Conv2D"([[EXPAND]], [[FILTER]]) <{dilations = [1, 2, 2, 1], padding = "SAME", strides = [1, 1, 1, 1]}> : (tensor<1x128x128x1xf32>, tensor<5x5x1x1xf32>) -> tensor<1x128x128x1xf32> + // CHECK-NEXT: [[SQUEEZE:%.*]] = "tf.Squeeze"([[CONV]]) <{squeeze_dims = [3]}> : (tensor<1x128x128x1xf32>) -> tensor<1x128x128xf32> // CHECK-NEXT: [[RESULT:%.*]] = "tf.BiasAdd"([[SQUEEZE]], [[BIAS]]) : (tensor<1x128x128xf32>, tensor<128xf32>) -> tensor<1x128x128xf32> // CHECK-NEXT: return [[RESULT]] : tensor<1x128x128xf32> } @@ -199,10 +199,10 @@ func.func @testDilatedDepthWiseConvWithExpandSqueeze1(%arg0: tensor<1x128x128xf3 // CHECK-LABEL: testDilatedDepthWiseConvWithExpandSqueeze1 // CHECK-SAME: ([[INPUT:%.*]]: tensor<1x128x128xf32>, [[FILTER:%.*]]: tensor<5x5x1x1xf32>, [[BIAS:%.*]]: tensor<128xf32>) - // CHECK-NEXT: [[AXIS:%.*]] = "tf.Const"() {value = dense<3> : tensor} : () -> tensor + // CHECK-NEXT: [[AXIS:%.*]] = "tf.Const"() <{value = dense<3> : tensor}> : () -> tensor // CHECK-NEXT: [[EXPAND:%.*]] = "tf.ExpandDims"([[INPUT]], [[AXIS]]) : (tensor<1x128x128xf32>, tensor) -> tensor<1x128x128x1xf32> - // CHECK-NEXT: [[CONV:%.*]] = "tf.DepthwiseConv2dNative"([[EXPAND]], [[FILTER]]) {dilations = [1, 2, 2, 1], padding = "SAME", strides = [1, 1, 1, 1]} : (tensor<1x128x128x1xf32>, tensor<5x5x1x1xf32>) -> tensor<1x128x128x1xf32> - // CHECK-NEXT: [[SQUEEZE:%.*]] = "tf.Squeeze"([[CONV]]) {squeeze_dims = [3]} : (tensor<1x128x128x1xf32>) -> tensor<1x128x128xf32> + // CHECK-NEXT: [[CONV:%.*]] = "tf.DepthwiseConv2dNative"([[EXPAND]], [[FILTER]]) <{dilations = [1, 2, 2, 1], padding = "SAME", strides = [1, 1, 1, 1]}> : (tensor<1x128x128x1xf32>, tensor<5x5x1x1xf32>) -> tensor<1x128x128x1xf32> + // CHECK-NEXT: [[SQUEEZE:%.*]] = "tf.Squeeze"([[CONV]]) <{squeeze_dims = [3]}> : (tensor<1x128x128x1xf32>) -> tensor<1x128x128xf32> // CHECK-NEXT: [[RESULT:%.*]] = "tf.BiasAdd"([[SQUEEZE]], [[BIAS]]) : (tensor<1x128x128xf32>, tensor<128xf32>) -> tensor<1x128x128xf32> // CHECK-NEXT: return [[RESULT]] : tensor<1x128x128xf32> } @@ -222,10 +222,10 @@ func.func @testDilatedConvWithExpandSqueeze2(%arg0: tensor<1x128x128xf32>, %arg1 // CHECK-LABEL: testDilatedConvWithExpandSqueeze2 // CHECK-SAME: ([[INPUT:%.*]]: tensor<1x128x128xf32>, [[FILTER:%.*]]: tensor<5x5x1x1xf32>, [[BIAS:%.*]]: tensor) - // CHECK-NEXT: [[AXIS:%.*]] = "tf.Const"() {value = dense<3> : tensor} : () -> tensor + // CHECK-NEXT: [[AXIS:%.*]] = "tf.Const"() <{value = dense<3> : tensor}> : () -> tensor // CHECK-NEXT: [[EXPAND:%.*]] = "tf.ExpandDims"([[INPUT]], [[AXIS]]) : (tensor<1x128x128xf32>, tensor) -> tensor<1x128x128x1xf32> - // CHECK-NEXT: [[CONV:%.*]] = "tf.Conv2D"([[EXPAND]], [[FILTER]]) {dilations = [1, 2, 2, 1], padding = "SAME", strides = [1, 1, 1, 1]} : (tensor<1x128x128x1xf32>, tensor<5x5x1x1xf32>) -> tensor<1x128x128x1xf32> - // CHECK-NEXT: [[SQUEEZE:%.*]] = "tf.Squeeze"([[CONV]]) {squeeze_dims = [3]} : (tensor<1x128x128x1xf32>) -> tensor<1x128x128xf32> + // CHECK-NEXT: [[CONV:%.*]] = "tf.Conv2D"([[EXPAND]], [[FILTER]]) <{dilations = [1, 2, 2, 1], padding = "SAME", strides = [1, 1, 1, 1]}> : (tensor<1x128x128x1xf32>, tensor<5x5x1x1xf32>) -> tensor<1x128x128x1xf32> + // CHECK-NEXT: [[SQUEEZE:%.*]] = "tf.Squeeze"([[CONV]]) <{squeeze_dims = [3]}> : (tensor<1x128x128x1xf32>) -> tensor<1x128x128xf32> // CHECK-NEXT: [[RESULT:%.*]] = "tf.BiasAdd"([[SQUEEZE]], [[BIAS]]) : (tensor<1x128x128xf32>, tensor) -> tensor<1x128x128xf32> // CHECK-NEXT: return [[RESULT]] : tensor<1x128x128xf32> } @@ -245,10 +245,10 @@ func.func @testDilatedDepthWiseConvWithExpandSqueeze2(%arg0: tensor<1x128x128xf3 // CHECK-LABEL: testDilatedDepthWiseConvWithExpandSqueeze2 // CHECK-SAME: ([[INPUT:%.*]]: tensor<1x128x128xf32>, [[FILTER:%.*]]: tensor<5x5x1x1xf32>, [[BIAS:%.*]]: tensor) - // CHECK-NEXT: [[AXIS:%.*]] = "tf.Const"() {value = dense<3> : tensor} : () -> tensor + // CHECK-NEXT: [[AXIS:%.*]] = "tf.Const"() <{value = dense<3> : tensor}> : () -> tensor // CHECK-NEXT: [[EXPAND:%.*]] = "tf.ExpandDims"([[INPUT]], [[AXIS]]) : (tensor<1x128x128xf32>, tensor) -> tensor<1x128x128x1xf32> - // CHECK-NEXT: [[CONV:%.*]] = "tf.DepthwiseConv2dNative"([[EXPAND]], [[FILTER]]) {dilations = [1, 2, 2, 1], padding = "SAME", strides = [1, 1, 1, 1]} : (tensor<1x128x128x1xf32>, tensor<5x5x1x1xf32>) -> tensor<1x128x128x1xf32> - // CHECK-NEXT: [[SQUEEZE:%.*]] = "tf.Squeeze"([[CONV]]) {squeeze_dims = [3]} : (tensor<1x128x128x1xf32>) -> tensor<1x128x128xf32> + // CHECK-NEXT: [[CONV:%.*]] = "tf.DepthwiseConv2dNative"([[EXPAND]], [[FILTER]]) <{dilations = [1, 2, 2, 1], padding = "SAME", strides = [1, 1, 1, 1]}> : (tensor<1x128x128x1xf32>, tensor<5x5x1x1xf32>) -> tensor<1x128x128x1xf32> + // CHECK-NEXT: [[SQUEEZE:%.*]] = "tf.Squeeze"([[CONV]]) <{squeeze_dims = [3]}> : (tensor<1x128x128x1xf32>) -> tensor<1x128x128xf32> // CHECK-NEXT: [[RESULT:%.*]] = "tf.BiasAdd"([[SQUEEZE]], [[BIAS]]) : (tensor<1x128x128xf32>, tensor) -> tensor<1x128x128xf32> // CHECK-NEXT: return [[RESULT]] : tensor<1x128x128xf32> } @@ -270,10 +270,10 @@ func.func @testDilatedConvWithExpandSqueeze3(%arg0: tensor<1x128x128xf32>, %arg1 // CHECK-LABEL: testDilatedConvWithExpandSqueeze3 // CHECK-SAME: ([[INPUT:%.*]]: tensor<1x128x128xf32>, [[FILTER:%.*]]: tensor<5x5x1x1xf32>, [[BIAS:%.*]]: tensor<128xf32>) - // CHECK-NEXT: [[AXIS:%.*]] = "tf.Const"() {value = dense<3> : tensor} : () -> tensor + // CHECK-NEXT: [[AXIS:%.*]] = "tf.Const"() <{value = dense<3> : tensor}> : () -> tensor // CHECK-NEXT: [[EXPAND:%.*]] = "tf.ExpandDims"([[INPUT]], [[AXIS]]) : (tensor<1x128x128xf32>, tensor) -> tensor<1x128x128x1xf32> - // CHECK-NEXT: [[CONV:%.*]] = "tf.Conv2D"([[EXPAND]], [[FILTER]]) {dilations = [1, 2, 2, 1], padding = "SAME", strides = [1, 1, 1, 1]} : (tensor<1x128x128x1xf32>, tensor<5x5x1x1xf32>) -> tensor<1x128x128x1xf32> - // CHECK-NEXT: [[SQUEEZE:%.*]] = "tf.Squeeze"([[CONV]]) {squeeze_dims = [3]} : (tensor<1x128x128x1xf32>) -> tensor<1x128x128xf32> + // CHECK-NEXT: [[CONV:%.*]] = "tf.Conv2D"([[EXPAND]], [[FILTER]]) <{dilations = [1, 2, 2, 1], padding = "SAME", strides = [1, 1, 1, 1]}> : (tensor<1x128x128x1xf32>, tensor<5x5x1x1xf32>) -> tensor<1x128x128x1xf32> + // CHECK-NEXT: [[SQUEEZE:%.*]] = "tf.Squeeze"([[CONV]]) <{squeeze_dims = [3]}> : (tensor<1x128x128x1xf32>) -> tensor<1x128x128xf32> // CHECK-NEXT: [[RESULT:%.*]] = "tf.BiasAdd"([[SQUEEZE]], [[BIAS]]) : (tensor<1x128x128xf32>, tensor<128xf32>) -> tensor<1x128x128xf32> // CHECK-NEXT: return [[RESULT]] : tensor<1x128x128xf32> } @@ -295,10 +295,10 @@ func.func @testDilatedDepthWiseConvWithExpandSqueeze3(%arg0: tensor<1x128x128xf3 // CHECK-LABEL: testDilatedDepthWiseConvWithExpandSqueeze3 // CHECK-SAME: ([[INPUT:%.*]]: tensor<1x128x128xf32>, [[FILTER:%.*]]: tensor<5x5x1x1xf32>, [[BIAS:%.*]]: tensor<128xf32>) - // CHECK-NEXT: [[AXIS:%.*]] = "tf.Const"() {value = dense<3> : tensor} : () -> tensor + // CHECK-NEXT: [[AXIS:%.*]] = "tf.Const"() <{value = dense<3> : tensor}> : () -> tensor // CHECK-NEXT: [[EXPAND:%.*]] = "tf.ExpandDims"([[INPUT]], [[AXIS]]) : (tensor<1x128x128xf32>, tensor) -> tensor<1x128x128x1xf32> - // CHECK-NEXT: [[CONV:%.*]] = "tf.DepthwiseConv2dNative"([[EXPAND]], [[FILTER]]) {dilations = [1, 2, 2, 1], padding = "SAME", strides = [1, 1, 1, 1]} : (tensor<1x128x128x1xf32>, tensor<5x5x1x1xf32>) -> tensor<1x128x128x1xf32> - // CHECK-NEXT: [[SQUEEZE:%.*]] = "tf.Squeeze"([[CONV]]) {squeeze_dims = [3]} : (tensor<1x128x128x1xf32>) -> tensor<1x128x128xf32> + // CHECK-NEXT: [[CONV:%.*]] = "tf.DepthwiseConv2dNative"([[EXPAND]], [[FILTER]]) <{dilations = [1, 2, 2, 1], padding = "SAME", strides = [1, 1, 1, 1]}> : (tensor<1x128x128x1xf32>, tensor<5x5x1x1xf32>) -> tensor<1x128x128x1xf32> + // CHECK-NEXT: [[SQUEEZE:%.*]] = "tf.Squeeze"([[CONV]]) <{squeeze_dims = [3]}> : (tensor<1x128x128x1xf32>) -> tensor<1x128x128xf32> // CHECK-NEXT: [[RESULT:%.*]] = "tf.BiasAdd"([[SQUEEZE]], [[BIAS]]) : (tensor<1x128x128xf32>, tensor<128xf32>) -> tensor<1x128x128xf32> // CHECK-NEXT: return [[RESULT]] : tensor<1x128x128xf32> } @@ -407,10 +407,10 @@ func.func @testDilatedConv1DExpandH(%arg0: tensor<1x128x3xf32>, %arg1: tensor<1x // CHECK-LABEL: testDilatedConv1DExpandH // CHECK-SAME: ([[INPUT:%.*]]: tensor<1x128x3xf32>, [[FILTER:%.*]]: tensor<1x5x3x8xf32>) - // CHECK-NEXT: [[AXIS:%.*]] = "tf.Const"() {value = dense<-3> : tensor} : () -> tensor + // CHECK-NEXT: [[AXIS:%.*]] = "tf.Const"() <{value = dense<-3> : tensor}> : () -> tensor // CHECK-NEXT: [[EXPAND:%.*]] = "tf.ExpandDims"([[INPUT]], [[AXIS]]) : (tensor<1x128x3xf32>, tensor) -> tensor<1x1x128x3xf32> - // CHECK-NEXT: [[CONV:%.*]] = "tf.Conv2D"([[EXPAND]], [[FILTER]]) {dilations = [1, 1, 2, 1], padding = "SAME", strides = [1, 1, 1, 1]} : (tensor<1x1x128x3xf32>, tensor<1x5x3x8xf32>) -> tensor<1x1x128x8xf32> - // CHECK-NEXT: [[RESULT:%.*]] = "tf.Squeeze"([[CONV]]) {squeeze_dims = [-3]} : (tensor<1x1x128x8xf32>) -> tensor<1x128x8xf32> + // CHECK-NEXT: [[CONV:%.*]] = "tf.Conv2D"([[EXPAND]], [[FILTER]]) <{dilations = [1, 1, 2, 1], padding = "SAME", strides = [1, 1, 1, 1]}> : (tensor<1x1x128x3xf32>, tensor<1x5x3x8xf32>) -> tensor<1x1x128x8xf32> + // CHECK-NEXT: [[RESULT:%.*]] = "tf.Squeeze"([[CONV]]) <{squeeze_dims = [-3]}> : (tensor<1x1x128x8xf32>) -> tensor<1x128x8xf32> // CHECK-NEXT: return [[RESULT]] : tensor<1x128x8xf32> } @@ -429,10 +429,10 @@ func.func @testDilatedConv1DExpandHWithBiasAdd(%arg0: tensor<1x128x3xf32>, %arg1 // CHECK-LABEL: testDilatedConv1DExpandHWithBiasAdd // CHECK-SAME: ([[INPUT:%.*]]: tensor<1x128x3xf32>, [[FILTER:%.*]]: tensor<1x5x3x8xf32>, [[BIAS:%.*]]: tensor<8xf32>) - // CHECK-NEXT: [[AXIS:%.*]] = "tf.Const"() {value = dense<-3> : tensor} : () -> tensor + // CHECK-NEXT: [[AXIS:%.*]] = "tf.Const"() <{value = dense<-3> : tensor}> : () -> tensor // CHECK-NEXT: [[EXPAND:%.*]] = "tf.ExpandDims"([[INPUT]], [[AXIS]]) : (tensor<1x128x3xf32>, tensor) -> tensor<1x1x128x3xf32> - // CHECK-NEXT: [[CONV:%.*]] = "tf.Conv2D"([[EXPAND]], [[FILTER]]) {dilations = [1, 1, 2, 1], padding = "SAME", strides = [1, 1, 1, 1]} : (tensor<1x1x128x3xf32>, tensor<1x5x3x8xf32>) -> tensor<1x1x128x8xf32> - // CHECK-NEXT: [[SQUEEZE:%.*]] = "tf.Squeeze"([[CONV]]) {squeeze_dims = [-3]} : (tensor<1x1x128x8xf32>) -> tensor<1x128x8xf32> + // CHECK-NEXT: [[CONV:%.*]] = "tf.Conv2D"([[EXPAND]], [[FILTER]]) <{dilations = [1, 1, 2, 1], padding = "SAME", strides = [1, 1, 1, 1]}> : (tensor<1x1x128x3xf32>, tensor<1x5x3x8xf32>) -> tensor<1x1x128x8xf32> + // CHECK-NEXT: [[SQUEEZE:%.*]] = "tf.Squeeze"([[CONV]]) <{squeeze_dims = [-3]}> : (tensor<1x1x128x8xf32>) -> tensor<1x128x8xf32> // CHECK-NEXT: [[RESULT:%.*]] = "tf.BiasAdd"([[SQUEEZE]], [[BIAS]]) : (tensor<1x128x8xf32>, tensor<8xf32>) -> tensor<1x128x8xf32> // CHECK-NEXT: return [[RESULT]] : tensor<1x128x8xf32> } @@ -451,10 +451,10 @@ func.func @testDilatedConv1DExpandW(%arg0: tensor<1x128x3xf32>, %arg1: tensor<5x // CHECK-LABEL: testDilatedConv1DExpandW // CHECK-SAME: ([[INPUT:%.*]]: tensor<1x128x3xf32>, [[FILTER:%.*]]: tensor<5x1x3x8xf32>) - // CHECK-NEXT: [[AXIS:%.*]] = "tf.Const"() {value = dense<-2> : tensor} : () -> tensor + // CHECK-NEXT: [[AXIS:%.*]] = "tf.Const"() <{value = dense<-2> : tensor}> : () -> tensor // CHECK-NEXT: [[EXPAND:%.*]] = "tf.ExpandDims"([[INPUT]], [[AXIS]]) : (tensor<1x128x3xf32>, tensor) -> tensor<1x128x1x3xf32> - // CHECK-NEXT: [[CONV:%.*]] = "tf.Conv2D"([[EXPAND]], [[FILTER]]) {dilations = [1, 2, 1, 1], padding = "SAME", strides = [1, 1, 1, 1]} : (tensor<1x128x1x3xf32>, tensor<5x1x3x8xf32>) -> tensor<1x128x1x8xf32> - // CHECK-NEXT: [[RESULT:%.*]] = "tf.Squeeze"([[CONV]]) {squeeze_dims = [-2]} : (tensor<1x128x1x8xf32>) -> tensor<1x128x8xf32> + // CHECK-NEXT: [[CONV:%.*]] = "tf.Conv2D"([[EXPAND]], [[FILTER]]) <{dilations = [1, 2, 1, 1], padding = "SAME", strides = [1, 1, 1, 1]}> : (tensor<1x128x1x3xf32>, tensor<5x1x3x8xf32>) -> tensor<1x128x1x8xf32> + // CHECK-NEXT: [[RESULT:%.*]] = "tf.Squeeze"([[CONV]]) <{squeeze_dims = [-2]}> : (tensor<1x128x1x8xf32>) -> tensor<1x128x8xf32> // CHECK-NEXT: return [[RESULT]] : tensor<1x128x8xf32> } @@ -473,10 +473,10 @@ func.func @testDilatedConv1DExpandWWithBiasAdd(%arg0: tensor<1x128x3xf32>, %arg1 // CHECK-LABEL: testDilatedConv1DExpandWWithBiasAdd // CHECK-SAME: ([[INPUT:%.*]]: tensor<1x128x3xf32>, [[FILTER:%.*]]: tensor<5x1x3x8xf32>, [[BIAS:%.*]]: tensor<8xf32>) - // CHECK-NEXT: [[AXIS:%.*]] = "tf.Const"() {value = dense<-2> : tensor} : () -> tensor + // CHECK-NEXT: [[AXIS:%.*]] = "tf.Const"() <{value = dense<-2> : tensor}> : () -> tensor // CHECK-NEXT: [[EXPAND:%.*]] = "tf.ExpandDims"([[INPUT]], [[AXIS]]) : (tensor<1x128x3xf32>, tensor) -> tensor<1x128x1x3xf32> - // CHECK-NEXT: [[CONV:%.*]] = "tf.Conv2D"([[EXPAND]], [[FILTER]]) {dilations = [1, 2, 1, 1], padding = "SAME", strides = [1, 1, 1, 1]} : (tensor<1x128x1x3xf32>, tensor<5x1x3x8xf32>) -> tensor<1x128x1x8xf32> - // CHECK-NEXT: [[SQUEEZE:%.*]] = "tf.Squeeze"([[CONV]]) {squeeze_dims = [-2]} : (tensor<1x128x1x8xf32>) -> tensor<1x128x8xf32> + // CHECK-NEXT: [[CONV:%.*]] = "tf.Conv2D"([[EXPAND]], [[FILTER]]) <{dilations = [1, 2, 1, 1], padding = "SAME", strides = [1, 1, 1, 1]}> : (tensor<1x128x1x3xf32>, tensor<5x1x3x8xf32>) -> tensor<1x128x1x8xf32> + // CHECK-NEXT: [[SQUEEZE:%.*]] = "tf.Squeeze"([[CONV]]) <{squeeze_dims = [-2]}> : (tensor<1x128x1x8xf32>) -> tensor<1x128x8xf32> // CHECK-NEXT: [[RESULT:%.*]] = "tf.BiasAdd"([[SQUEEZE]], [[BIAS]]) : (tensor<1x128x8xf32>, tensor<8xf32>) -> tensor<1x128x8xf32> // CHECK-NEXT: return [[RESULT]] : tensor<1x128x8xf32> } @@ -495,10 +495,10 @@ func.func @testDilatedConv1DWithMixedPostiveAndNegativeAxis(%arg0: tensor<1x128x // CHECK-LABEL: testDilatedConv1DWithMixedPostiveAndNegativeAxis // CHECK-SAME: ([[INPUT:%.*]]: tensor<1x128x3xf32>, [[FILTER:%.*]]: tensor<1x5x3x8xf32>) - // CHECK-NEXT: [[AXIS:%.*]] = "tf.Const"() {value = dense<1> : tensor} : () -> tensor + // CHECK-NEXT: [[AXIS:%.*]] = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor // CHECK-NEXT: [[EXPAND:%.*]] = "tf.ExpandDims"([[INPUT]], [[AXIS]]) : (tensor<1x128x3xf32>, tensor) -> tensor<1x1x128x3xf32> - // CHECK-NEXT: [[CONV:%.*]] = "tf.Conv2D"([[EXPAND]], [[FILTER]]) {dilations = [1, 1, 2, 1], padding = "SAME", strides = [1, 1, 1, 1]} : (tensor<1x1x128x3xf32>, tensor<1x5x3x8xf32>) -> tensor<1x1x128x8xf32> - // CHECK-NEXT: [[RESULT:%.*]] = "tf.Squeeze"([[CONV]]) {squeeze_dims = [-3]} : (tensor<1x1x128x8xf32>) -> tensor<1x128x8xf32> + // CHECK-NEXT: [[CONV:%.*]] = "tf.Conv2D"([[EXPAND]], [[FILTER]]) <{dilations = [1, 1, 2, 1], padding = "SAME", strides = [1, 1, 1, 1]}> : (tensor<1x1x128x3xf32>, tensor<1x5x3x8xf32>) -> tensor<1x1x128x8xf32> + // CHECK-NEXT: [[RESULT:%.*]] = "tf.Squeeze"([[CONV]]) <{squeeze_dims = [-3]}> : (tensor<1x1x128x8xf32>) -> tensor<1x128x8xf32> // CHECK-NEXT: return [[RESULT]] : tensor<1x128x8xf32> } @@ -518,11 +518,11 @@ func.func @testPaddedDilatedConv(%arg0 : tensor<2x1920x64xf32>) -> tensor<2x192 // CHECK-LABEL: testPaddedDilatedConv // CHECK-SAME: ([[INPUT:%.*]]: tensor<2x1920x64xf32>) - // CHECK-NEXT: [[AXIS:%.*]] = "tf.Const"() {value = dense<2> : tensor} : () -> tensor - // CHECK-NEXT: [[FILTER:%.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<3x1x64x128xf32>} : () -> tensor<3x1x64x128xf32> + // CHECK-NEXT: [[AXIS:%.*]] = "tf.Const"() <{value = dense<2> : tensor}> : () -> tensor + // CHECK-NEXT: [[FILTER:%.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<3x1x64x128xf32>}> : () -> tensor<3x1x64x128xf32> // CHECK-NEXT: [[EXPAND:%.*]] = "tf.ExpandDims"([[INPUT]], [[AXIS]]) {device = ""} : (tensor<2x1920x64xf32>, tensor) -> tensor<2x1920x1x64xf32> - // CHECK-NEXT: [[CONV:%.*]] = "tf.Conv2D"([[EXPAND]], [[FILTER]]) {data_format = "NHWC", device = "", dilations = [1, 2, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true} : (tensor<2x1920x1x64xf32>, tensor<3x1x64x128xf32>) -> tensor<2x1920x1x128xf32> - // CHECK-NEXT: [[RESULT:%.*]] = "tf.Squeeze"([[CONV]]) {device = "", squeeze_dims = [2]} : (tensor<2x1920x1x128xf32>) -> tensor<2x1920x128xf32> + // CHECK-NEXT: [[CONV:%.*]] = "tf.Conv2D"([[EXPAND]], [[FILTER]]) <{data_format = "NHWC", dilations = [1, 2, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true}> {device = ""} : (tensor<2x1920x1x64xf32>, tensor<3x1x64x128xf32>) -> tensor<2x1920x1x128xf32> + // CHECK-NEXT: [[RESULT:%.*]] = "tf.Squeeze"([[CONV]]) <{squeeze_dims = [2]}> {device = ""} : (tensor<2x1920x1x128xf32>) -> tensor<2x1920x128xf32> // CHECK-NEXT: return [[RESULT]] : tensor<2x1920x128xf32> } @@ -539,7 +539,7 @@ func.func @testDilatedConvInterleaved(%arg0: tensor<1x128x128x3xf32>, %arg1: ten // CHECK-LABEL: testDilatedConvInterleaved // CHECK-SAME: ([[INPUT:%.*]]: tensor<1x128x128x3xf32>, [[FILTER:%.*]]: tensor<5x5x3x8xf32>) - // CHECK-NEXT: [[RESULT0:%.*]] = "tf.Conv2D"([[INPUT]], [[FILTER]]) {dilations = [1, 2, 2, 1], padding = "VALID", strides = [1, 1, 1, 1]} : (tensor<1x128x128x3xf32>, tensor<5x5x3x8xf32>) -> tensor<1x120x120x8xf32> - // CHECK-NEXT: [[RESULT1:%.*]] = "tf.Conv2D"([[INPUT]], [[FILTER]]) {dilations = [1, 2, 2, 1], padding = "VALID", strides = [1, 1, 1, 1]} : (tensor<1x128x128x3xf32>, tensor<5x5x3x8xf32>) -> tensor<1x120x120x8xf32> + // CHECK-NEXT: [[RESULT0:%.*]] = "tf.Conv2D"([[INPUT]], [[FILTER]]) <{dilations = [1, 2, 2, 1], padding = "VALID", strides = [1, 1, 1, 1]}> : (tensor<1x128x128x3xf32>, tensor<5x5x3x8xf32>) -> tensor<1x120x120x8xf32> + // CHECK-NEXT: [[RESULT1:%.*]] = "tf.Conv2D"([[INPUT]], [[FILTER]]) <{dilations = [1, 2, 2, 1], padding = "VALID", strides = [1, 1, 1, 1]}> : (tensor<1x128x128x3xf32>, tensor<5x5x3x8xf32>) -> tensor<1x120x120x8xf32> // CHECK-NEXT: return [[RESULT0]], [[RESULT1]] : tensor<1x120x120x8xf32>, tensor<1x120x120x8xf32> } diff --git a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/if_op.mlir b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/if_op.mlir index f29afb30846142..7ea7e48777522e 100644 --- a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/if_op.mlir +++ b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/if_op.mlir @@ -1,7 +1,7 @@ // RUN: flatbuffer_translate -mlir-to-tflite-flatbuffer %s -o - | flatbuffer_translate --tflite-flatbuffer-to-mlir - -o - | FileCheck %s // Confirm function references in if ops are preserved func.func @main(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>) -> tensor<1xf32> { -// CHECK: %{{.*}} = "tf.If"(%{{.*}}, %{{.*}}, %{{.*}}) {else_branch = @cond_false, is_stateless = false, then_branch = @cond_true} : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> +// CHECK: %{{.*}} = "tf.If"(%{{.*}}, %{{.*}}, %{{.*}}) <{else_branch = @cond_false, is_stateless = false, then_branch = @cond_true}> : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> %0 = "tfl.less"(%arg0, %arg1) : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1> %1 = "tf.If"(%0, %arg0, %arg1) {else_branch = @cond_false, then_branch = @cond_true, is_stateless = false} : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> func.return %1 : tensor<1xf32> diff --git a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/metadata_buffer.mlir b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/metadata_buffer.mlir new file mode 100644 index 00000000000000..6b76b31c9a52bf --- /dev/null +++ b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/metadata_buffer.mlir @@ -0,0 +1,9 @@ +// RUN: flatbuffer_translate -mlir-to-tflite-flatbuffer %s -o - | flatbuffer_translate --tflite-flatbuffer-to-mlir - -o - | FileCheck %s + +// CHECK: tfl.metadata_buffer = [3 : i32, 7 : i32] +module attributes {tfl.metadata_buffer = [3 : i32, 7 : i32]} { + func.func @main(%arg0: tensor, %arg1: tensor<3x2xi32>) -> tensor<3x2xi32> { + %0 = "tfl.add" (%arg0, %arg1) {fused_activation_function = "NONE"} : (tensor, tensor<3x2xi32>) -> tensor<3x2xi32> + func.return %0 : tensor<3x2xi32> + } +} \ No newline at end of file diff --git a/tensorflow/compiler/mlir/lite/tests/legalize-tf-while.mlir b/tensorflow/compiler/mlir/lite/tests/legalize-tf-while.mlir index a807cc84ee0a78..ab9b39bd94cb97 100644 --- a/tensorflow/compiler/mlir/lite/tests/legalize-tf-while.mlir +++ b/tensorflow/compiler/mlir/lite/tests/legalize-tf-while.mlir @@ -66,7 +66,7 @@ func.func @while_cond_10_frozen0(%arg0: tensor<*xi32>, %arg1: tensor<*xi32>, %ar // CANON: }, { // CANON: ^bb0([[VAL_11:%.*]]: tensor<*xi32>, [[VAL_12:%.*]]: tensor<*xi32>, [[VAL_13:%.*]]: tensor<*xf32>): // CANON-DAG: [[VAL_4:%.*]] = arith.constant dense<1> : tensor -// CANON-DAG: [[VAL_5:%.*]] = "tf.Const"() {value = dense<2.560000e+02> : tensor<256x256xf32>} : () -> tensor +// CANON-DAG: [[VAL_5:%.*]] = "tf.Const"() <{value = dense<2.560000e+02> : tensor<256x256xf32>}> : () -> tensor // CANON: [[VAL_14:%.*]] = "tf.AddV2"([[VAL_12]], [[VAL_4]]) // CANON: [[VAL_15:%.*]] = "tf.AddV2"([[VAL_13]], [[VAL_5]]) // CANON: [[VAL_16:%.*]] = "tf.AddV2"([[VAL_11]], [[VAL_4]]) diff --git a/tensorflow/compiler/mlir/lite/tests/legalize-tf.mlir b/tensorflow/compiler/mlir/lite/tests/legalize-tf.mlir index 444a494f73769a..685efd5be0ca2d 100644 --- a/tensorflow/compiler/mlir/lite/tests/legalize-tf.mlir +++ b/tensorflow/compiler/mlir/lite/tests/legalize-tf.mlir @@ -198,7 +198,7 @@ func.func @fakeQuantVarsTrue(%arg0: tensor<8x8x8x8xf32>, %arg1: tensor, %ar func.return %0 : tensor<8x8x8x8xf32> // CHECK-LABEL: fakeQuantVarsTrue - // CHECK: "tf.FakeQuantWithMinMaxVars"(%arg0, %arg1, %arg2) {max = 1.000000e+00 : f32, min = 0.000000e+00 : f32, narrow_range = true, num_bits = 5 : i64} + // CHECK: "tf.FakeQuantWithMinMaxVars"(%arg0, %arg1, %arg2) <{narrow_range = true, num_bits = 5 : i64}> {max = 1.000000e+00 : f32, min = 0.000000e+00 : f32} } func.func @fakeQuantArgsFalse4Bits(%arg0: tensor<8x8x8x8xf32>) -> tensor<8x8x8x8xf32> { @@ -235,7 +235,7 @@ func.func @fakeQuantVarsTrue4Bits(%arg0: tensor<8x8x8x8xf32>, %arg1: tensor func.return %0 : tensor<8x8x8x8xf32> // CHECK-LABEL: fakeQuantVarsTrue - // CHECK: "tf.FakeQuantWithMinMaxVars"(%arg0, %arg1, %arg2) {max = 1.000000e+00 : f32, min = 0.000000e+00 : f32, narrow_range = true, num_bits = 3 : i64} + // CHECK: "tf.FakeQuantWithMinMaxVars"(%arg0, %arg1, %arg2) <{narrow_range = true, num_bits = 3 : i64}> {max = 1.000000e+00 : f32, min = 0.000000e+00 : f32} } func.func @const() -> tensor<2xi32> { @@ -1421,7 +1421,7 @@ func.func @strided_slice_big_dims(%arg0: tensor<5x6x7xf32>, %arg1: tensor<3xi32> %0 = "tf.StridedSlice"(%arg0, %arg1, %arg2, %arg3) {begin_mask = 0 : i64, ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 7 : i64, shrink_axis_mask = 0 : i64, offset = false} : (tensor<5x6x7xf32>, tensor<3xi32>, tensor<3xi32>, tensor<3xi32>) -> tensor<1x1x5x6x7xf32> func.return %0 : tensor<1x1x5x6x7xf32> // CHECK-LABEL: strided_slice_big_dims - // CHECK: %0 = "tf.StridedSlice"(%arg0, %arg1, %arg2, %arg3) {begin_mask = 0 : i64, ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 7 : i64, offset = false, shrink_axis_mask = 0 : i64} : (tensor<5x6x7xf32>, tensor<3xi32>, tensor<3xi32>, tensor<3xi32>) -> tensor<1x1x5x6x7xf32> + // CHECK: %0 = "tf.StridedSlice"(%arg0, %arg1, %arg2, %arg3) <{begin_mask = 0 : i64, ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 7 : i64, shrink_axis_mask = 0 : i64}> {offset = false} : (tensor<5x6x7xf32>, tensor<3xi32>, tensor<3xi32>, tensor<3xi32>) -> tensor<1x1x5x6x7xf32> } func.func @slice1Tensor(%arg0: tensor<2x3x5xf32>, %arg1: tensor<3xi32>, %arg2: tensor<3xi32>) -> tensor { @@ -1606,7 +1606,7 @@ func.func @sparse_to_dense_with_2d_sparse_indices_and_second_dim_greater_than_4( %0 = "tf.SparseToDense"(%arg0, %arg1, %arg2, %arg3) {validate_indices = true}: (tensor<3x5xi32>, tensor<3xi32>, tensor<2xf32>, tensor) -> tensor func.return %0 : tensor // CHECK-LABEL: sparse_to_dense_with_2d_sparse_indices_and_second_dim_greater_than_4 - // CHECK: "tf.SparseToDense"(%arg0, %arg1, %arg2, %arg3) {validate_indices = true} : (tensor<3x5xi32>, tensor<3xi32>, tensor<2xf32>, tensor) -> tensor + // CHECK: "tf.SparseToDense"(%arg0, %arg1, %arg2, %arg3) <{validate_indices = true}> : (tensor<3x5xi32>, tensor<3xi32>, tensor<2xf32>, tensor) -> tensor } func.func @where(%arg0: tensor<3x5xi1>) -> tensor { @@ -2311,7 +2311,7 @@ func.func @conv3d_invalid_strides(%arg0: tensor,%arg1: tensor, tensor) -> tensor func.return %0: tensor // CHECK-LABEL: conv3d_invalid_strides - // CHECK: [[BCT:%.*]] = "tf.Conv3D"(%arg0, %arg1) {padding = "SAME", strides = [2, 1, 1, 1, 1]} : (tensor, tensor) -> tensor + // CHECK: [[BCT:%.*]] = "tf.Conv3D"(%arg0, %arg1) <{padding = "SAME", strides = [2, 1, 1, 1, 1]}> : (tensor, tensor) -> tensor // CHECK: return [[BCT]] : tensor } @@ -2705,7 +2705,7 @@ func.func @approx_top_k_with_min_k(%arg0: tensor<1x4xf32>) -> (tensor<1x4xf32>, func.return %values, %indices: tensor<1x4xf32>, tensor<1x4xi32> // CHECK-LABEL: approx_top_k_with_min_k - // CHECK: %values, %indices = "tf.ApproxTopK"(%arg0) {aggregate_to_topk = true, is_max_k = false, k = 4 : i64, recall_target = 8.500000e-01 : f32, reduction_dimension = 1 : i64, reduction_input_size_override = -1 : i64} : (tensor<1x4xf32>) -> (tensor<1x4xf32>, tensor<1x4xi32>) + // CHECK: %values, %indices = "tf.ApproxTopK"(%arg0) <{aggregate_to_topk = true, is_max_k = false, k = 4 : i64, recall_target = 8.500000e-01 : f32, reduction_dimension = 1 : i64, reduction_input_size_override = -1 : i64}> : (tensor<1x4xf32>) -> (tensor<1x4xf32>, tensor<1x4xi32>) // CHECK: return %values, %indices : tensor<1x4xf32>, tensor<1x4xi32> } @@ -2714,7 +2714,7 @@ func.func @approx_top_k_reduction_dimension_not_last_dim(%arg0: tensor<1x4xf32>) func.return %values, %indices: tensor<1x4xf32>, tensor<1x4xi32> // CHECK-LABEL: approx_top_k_reduction_dimension_not_last_dim - // CHECK: %values, %indices = "tf.ApproxTopK"(%arg0) {aggregate_to_topk = true, is_max_k = true, k = 4 : i64, recall_target = 8.500000e-01 : f32, reduction_dimension = 0 : i64, reduction_input_size_override = -1 : i64} : (tensor<1x4xf32>) -> (tensor<1x4xf32>, tensor<1x4xi32>) + // CHECK: %values, %indices = "tf.ApproxTopK"(%arg0) <{aggregate_to_topk = true, is_max_k = true, k = 4 : i64, recall_target = 8.500000e-01 : f32, reduction_dimension = 0 : i64, reduction_input_size_override = -1 : i64}> : (tensor<1x4xf32>) -> (tensor<1x4xf32>, tensor<1x4xi32>) // CHECK: return %values, %indices : tensor<1x4xf32>, tensor<1x4xi32> } diff --git a/tensorflow/compiler/mlir/lite/tests/lift_tflite_flex_ops.mlir b/tensorflow/compiler/mlir/lite/tests/lift_tflite_flex_ops.mlir index a03519fbdd1d6f..8ed0fe8cac86cf 100644 --- a/tensorflow/compiler/mlir/lite/tests/lift_tflite_flex_ops.mlir +++ b/tensorflow/compiler/mlir/lite/tests/lift_tflite_flex_ops.mlir @@ -20,7 +20,7 @@ func.func @TfBatchMatMulV2(%arg0: tensor<4x128x2xf32>, %arg1: tensor<2x1xf32>) custom_option = #tfl } : (tensor<4x128x2xf32>, tensor<2x1xf32>) -> tensor<4x128x1xf32> -// CHECK: "tf.BatchMatMulV2"(%arg0, %arg1) {T = f32, adj_x = false, adj_y = false} : (tensor<4x128x2xf32>, tensor<2x1xf32>) -> tensor<4x128x1xf32> +// CHECK: "tf.BatchMatMulV2"(%arg0, %arg1) <{adj_x = false, adj_y = false}> {T = f32} : (tensor<4x128x2xf32>, tensor<2x1xf32>) -> tensor<4x128x1xf32> func.return %0 : tensor<4x128x1xf32> } @@ -66,7 +66,7 @@ func.func @TfMapDataset(%arg0: tensor) -> (tensor // CHECK: "tf.MapDataset"( -// CHECK-SAME: {Targuments = [], f = @{{.*}}, metadata = "", output_shapes = [#tf_type.shape<>], output_types = [!tf_type.string], preserve_cardinality = true, use_inter_op_parallelism = true} +// CHECK-SAME: <{f = @{{.*}}, metadata = "", output_shapes = [#tf_type.shape<>], output_types = [!tf_type.string], preserve_cardinality = true, use_inter_op_parallelism = true}> {Targuments = []} } // CHECK-LABEL: TfTakeWhileDataset @@ -78,7 +78,7 @@ func.func @TfTakeWhileDataset(%arg0: tensor, %arg1: tensor // CHECK: "tf.TakeWhileDataset"( -// CHECK-SAME: {Targuments = [!tf_type.resource, !tf_type.resource, i64, !tf_type.resource, !tf_type.resource, !tf_type.resource, !tf_type.resource, i64], metadata = "", output_shapes = [#tf_type.shape<>], output_types = [!tf_type.string], predicate = @{{.*}}} +// CHECK-SAME: <{metadata = "", output_shapes = [#tf_type.shape<>], output_types = [!tf_type.string], predicate = @{{.*}}}> {Targuments = [!tf_type.resource, !tf_type.resource, i64, !tf_type.resource, !tf_type.resource, !tf_type.resource, !tf_type.resource, i64]} } // CHECK-LABEL: FailureOnInvalidOp diff --git a/tensorflow/compiler/mlir/lite/tests/lower-static-tensor-list.mlir b/tensorflow/compiler/mlir/lite/tests/lower-static-tensor-list.mlir index 3bfeceee049626..79d969106d868a 100644 --- a/tensorflow/compiler/mlir/lite/tests/lower-static-tensor-list.mlir +++ b/tensorflow/compiler/mlir/lite/tests/lower-static-tensor-list.mlir @@ -4,9 +4,9 @@ // CHECK-LABEL: tensorlistConst func.func @tensorlistConst(%arg0 : tensor<1xi32>) -> tensor<2x3xi32> { - // CHECK-DAG: %[[ELEMENT0:.*]] = "tf.Const"() {value = dense<[0, 1, 2]> : tensor<3xi32>} : () -> tensor<3xi32> - // CHECK-DAG: %[[ELEMENT1:.*]] = "tf.Const"() {value = dense<[3, 4, 5]> : tensor<3xi32>} : () -> tensor<3xi32> - // CHECK: %[[LIST:.*]] = "tf.Pack"(%[[ELEMENT0]], %[[ELEMENT1]]) {axis = 0 : i64} : (tensor<3xi32>, tensor<3xi32>) -> tensor<2x3xi32> + // CHECK-DAG: %[[ELEMENT0:.*]] = "tf.Const"() <{value = dense<[0, 1, 2]> : tensor<3xi32>}> : () -> tensor<3xi32> + // CHECK-DAG: %[[ELEMENT1:.*]] = "tf.Const"() <{value = dense<[3, 4, 5]> : tensor<3xi32>}> : () -> tensor<3xi32> + // CHECK: %[[LIST:.*]] = "tf.Pack"(%[[ELEMENT0]], %[[ELEMENT1]]) <{axis = 0 : i64}> : (tensor<3xi32>, tensor<3xi32>) -> tensor<2x3xi32> %0 = "tf.Const"() {value = #tf_type : tensor} : () -> tensor>> // CHECK: return %[[LIST]] @@ -20,7 +20,7 @@ func.func @tensorlistConst(%arg0 : tensor<1xi32>) -> tensor<2x3xi32> { func.func @emptyTensorlistConst(%arg0 : tensor<1xi32>) -> tensor<0x3xi32> { %0 = "tf.Const"() {value = #tf_type : tensor} : () -> tensor>> - // CHECK: "tf.Const"() {value = dense<> : tensor<0x3xi32>} : () -> tensor<0x3xi32> + // CHECK: "tf.Const"() <{value = dense<> : tensor<0x3xi32>}> : () -> tensor<0x3xi32> // CHECK-NOT: tf.TensorListStack %1 = "tf.TensorListStack"(%0, %arg0) : (tensor>>, tensor<1xi32>) -> tensor<0x3xi32> func.return %1 : tensor<0x3xi32> @@ -35,7 +35,7 @@ func.func @tensorlistGetItem(%arg0: tensor<3x10xf32>, %arg1: tensor<1xi32>, %arg %2 = "tf.TensorListStack"(%0, %arg1) : (tensor>>, tensor<1xi32>) -> tensor<3x10xf32> func.return %1, %2 : tensor<10xf32>, tensor<3x10xf32> -// CHECK: %0 = "tf.Gather"(%arg0, %arg2) {validate_indices = true} : (tensor<3x10xf32>, tensor) -> tensor<10xf32> +// CHECK: %0 = "tf.Gather"(%arg0, %arg2) <{validate_indices = true}> : (tensor<3x10xf32>, tensor) -> tensor<10xf32> // CHECK: return %0, %arg0 : tensor<10xf32>, tensor<3x10xf32> } @@ -48,7 +48,7 @@ func.func @tensorlistGetItemWithUnknownRank(%arg0: tensor<*xf32>, %arg1: tensor< %2 = "tf.TensorListStack"(%0, %arg1) : (tensor>>, tensor<1xi32>) -> tensor<*xf32> func.return %1, %2 : tensor<*xf32>, tensor<*xf32> -// CHECK: %0 = "tf.Gather"(%arg0, %arg2) {validate_indices = true} : (tensor<*xf32>, tensor) -> tensor<*xf32> +// CHECK: %0 = "tf.Gather"(%arg0, %arg2) <{validate_indices = true}> : (tensor<*xf32>, tensor) -> tensor<*xf32> // CHECK: return %0, %arg0 : tensor<*xf32>, tensor<*xf32> } @@ -175,7 +175,7 @@ func.func @tensorlistReserve(%arg0: tensor<3xi32>, %arg1: tensor, %arg2: te // CHECK-DAG: [[SHAPE:%.*]] = "tf.Concat"([[ZERO2]], [[DIM0]], %arg0) : (tensor, tensor<1xi32>, tensor<3xi32>) -> tensor<4xi32> // CHECK-DAG: [[VALUES:%.*]] = arith.constant dense<0.000000e+00> : tensor // CHECK: [[LIST:%.*]] = "tf.Fill"([[SHAPE]], [[VALUES]]) : (tensor<4xi32>, tensor) -> tensor -// CHECK: [[RESULT:%.*]] = "tf.Gather"([[LIST]], %arg2) {validate_indices = true} : (tensor, tensor) -> tensor +// CHECK: [[RESULT:%.*]] = "tf.Gather"([[LIST]], %arg2) <{validate_indices = true}> : (tensor, tensor) -> tensor // CHECK: return [[RESULT]] : tensor } @@ -188,7 +188,7 @@ func.func @tensorlistReserveUnrankedElements(%arg0: tensor, %arg1: tensor func.return %1 : tensor<*xf32> // CHECK: [[RESULT:%[0-9]+]] = "tf.Fill"{{.*}}(tensor, tensor) -> tensor<*xf32> -// CHECK: [[RESULT2:%[0-9]+]] = "tf.Gather"{{.*}}{validate_indices = true} : (tensor<*xf32>, tensor) -> tensor<*xf32> +// CHECK: [[RESULT2:%[0-9]+]] = "tf.Gather"{{.*}}<{validate_indices = true}> : (tensor<*xf32>, tensor) -> tensor<*xf32> // CHECK: return [[RESULT2]] : tensor<*xf32> } @@ -208,7 +208,7 @@ func.func @tensorlistReserveConstantUnknownElementShapeDim(%arg0: tensor, % // CHECK-DAG: [[SHAPE:%.*]] = "tf.Concat"([[ZERO2]], [[DIM0]], [[ELEMENT_SHAPE]]) : (tensor, tensor<1xi32>, tensor<2xi32>) -> tensor<3xi32> // CHECK-DAG: [[VALUES:%.*]] = arith.constant dense<0.000000e+00> : tensor // CHECK: [[LIST:%.*]] = "tf.Fill"([[SHAPE]], [[VALUES]]) : (tensor<3xi32>, tensor) -> tensor -// CHECK: [[RESULT:%.*]] = "tf.Gather"([[LIST]], %arg1) {validate_indices = true} : (tensor, tensor) -> tensor +// CHECK: [[RESULT:%.*]] = "tf.Gather"([[LIST]], %arg1) <{validate_indices = true}> : (tensor, tensor) -> tensor // CHECK: return [[RESULT]] : tensor } @@ -245,7 +245,7 @@ func.func @tensorlistReserveUnrankedElementShape(%arg0: tensor<*xi32>, %arg1: te // CHECK-DAG: [[CONCAT:%.*]] = "tf.Concat"([[AXIS_1]], [[EXPAND_DIM]], %arg0) : (tensor, tensor<1xi32>, tensor<*xi32>) -> tensor // CHECK: [[CST:%.*]] = arith.constant dense<0.000000e+00> : tensor // CHECK: [[FILL:%.*]] = "tf.Fill"([[CONCAT]], [[CST]]) : (tensor, tensor) -> tensor<*xf32> -// CHECK: [[GATHER:%.*]] = "tf.Gather"([[FILL]], %arg2) {validate_indices = true} : (tensor<*xf32>, tensor) -> tensor<*xf32> +// CHECK: [[GATHER:%.*]] = "tf.Gather"([[FILL]], %arg2) <{validate_indices = true}> : (tensor<*xf32>, tensor) -> tensor<*xf32> // CHECK: return [[GATHER]] : tensor<*xf32> } @@ -263,7 +263,7 @@ func.func @EmptyTensorList(%arg0: tensor<3xi32>, %arg1: tensor, %arg2: tens // CHECK-DAG: [[SHAPE:%.*]] = "tf.Concat"([[ZERO]], [[DIM0]], [[ELEM_SHAPE]]) : (tensor, tensor<1xi32>, tensor<3xi32>) -> tensor<4xi32> // CHECK-DAG: [[VALUES:%.*]] = arith.constant dense<0.000000e+00> : tensor // CHECK: [[LIST:%.*]] = "tf.Fill"([[SHAPE]], [[VALUES]]) : (tensor<4xi32>, tensor) -> tensor<0x?x?x?xf32> -// CHECK: [[RESULT:%.*]] = "tf.Gather"([[LIST]], [[IDX]]) {validate_indices = true} : (tensor<0x?x?x?xf32>, tensor) -> tensor +// CHECK: [[RESULT:%.*]] = "tf.Gather"([[LIST]], [[IDX]]) <{validate_indices = true}> : (tensor<0x?x?x?xf32>, tensor) -> tensor // CHECK: return [[RESULT]] : tensor } @@ -294,7 +294,7 @@ func.func @tensorlistLength(%arg0: tensor<3x10xf32>, %arg1: tensor<1xi32>) -> (t // CHECK-SAME: ([[INPUT:%.*]]: tensor<3x10xf32>, [[ELEM_SHAPE:%.*]]: tensor<1xi32>) // CHECK-DAG: [[SHAPE:%.*]] = "tf.Shape"([[INPUT]]) {{.*}} -> tensor<2xi32> // CHECK-DAG: [[ZERO:%cst.*]] = arith.constant dense<0> : tensor -// CHECK: [[RESULT:%.*]] = "tf.Gather"([[SHAPE]], [[ZERO]]) {validate_indices = true} : (tensor<2xi32>, tensor) -> tensor +// CHECK: [[RESULT:%.*]] = "tf.Gather"([[SHAPE]], [[ZERO]]) <{validate_indices = true}> : (tensor<2xi32>, tensor) -> tensor // CHECK: return [[RESULT]] : tensor } @@ -352,7 +352,8 @@ func.func @tensorlistWhileRegion(%arg0: tensor<2x3xf32>) -> tensor<*xf32> { %cst_1 = arith.constant dense<-1> : tensor %0 = "tf.TensorListFromTensor"(%arg0, %cst) : (tensor<2x3xf32>, tensor<1xi32>) -> tensor>> // CHECK: "tf.WhileRegion" - %1:2 = "tf.WhileRegion"(%cst_0, %0) ({ + // CHECK: <{is_stateless = false}> + %1:2 = "tf.WhileRegion"(%cst_0, %0) <{is_stateless = false}> ({ ^bb0(%carg0: tensor, %carg1: tensor): %cst_2 = arith.constant dense<2> : tensor %1 = "tf.Less"(%carg0, %cst_2) : (tensor, tensor) -> tensor @@ -376,9 +377,9 @@ func.func @tensorlistWhileRegion(%arg0: tensor<2x3xf32>) -> tensor<*xf32> { // CHECK-NOT: tensor // CHECK: "tf.Yield"(%[[LEN]], %[[BARG1]]) : (tensor, tensor<*xf32>) -> () - }) {is_stateless = false} : (tensor, tensor>>) -> (tensor, tensor>>) + }) : (tensor, tensor>>) -> (tensor, tensor>>) // make sure the variant types in input/output have been updated - // CHECK: {is_stateless = false} : (tensor, tensor<2x3xf32>) -> (tensor, tensor<*xf32>) + // : (tensor, tensor<2x3xf32>) -> (tensor, tensor<*xf32>) %2 = "tf.TensorListStack"(%1#1, %cst_1) : (tensor>>, tensor) -> tensor<*xf32> // CHECK: return %0#1 : tensor<*xf32> func.return %2 : tensor<*xf32> @@ -443,11 +444,11 @@ func.func @tensorlistResize(%arg0: tensor<3x10xf32>, %arg1: tensor<1xi32>, %arg2 // CHECK: [[ZERO:%.*]] = arith.constant dense<0> : tensor // CHECK: [[SHAPE:%.*]] = "tf.Shape"([[INPUT]]) : (tensor<3x10xf32>) -> tensor<2xi32> // CHECK: [[ZERO_1:%.*]] = arith.constant dense<0> : tensor -// CHECK: [[INPUT_SIZE:%.*]] = "tf.Gather"([[SHAPE]], [[ZERO_1]]) {validate_indices = true} : (tensor<2xi32>, tensor) -> tensor +// CHECK: [[INPUT_SIZE:%.*]] = "tf.Gather"([[SHAPE]], [[ZERO_1]]) <{validate_indices = true}> : (tensor<2xi32>, tensor) -> tensor // CHECK: [[SIZE_DIFF:%.*]] = "tf.Sub"([[SIZE]], [[INPUT_SIZE]]) : (tensor, tensor) -> tensor // CHECK: [[DIFF_RES:%.*]] = "tf.Greater"([[SIZE_DIFF]], [[ZERO]]) : (tensor, tensor) -> tensor // CHECK: [[SHAPE_1:%.*]] = "tf.Shape"([[INPUT]]) : (tensor<3x10xf32>) -> tensor -// CHECK: [[RESULT:%.*]] = "tf.If"([[DIFF_RES]], [[INPUT]], [[SHAPE_1]], [[SIZE_DIFF]], [[SIZE]]) {else_branch = @cond_false, is_stateless = true, then_branch = @cond_true} : (tensor, tensor<3x10xf32>, tensor, tensor, tensor) -> tensor +// CHECK: [[RESULT:%.*]] = "tf.If"([[DIFF_RES]], [[INPUT]], [[SHAPE_1]], [[SIZE_DIFF]], [[SIZE]]) <{else_branch = @cond_false, is_stateless = true, then_branch = @cond_true}> : (tensor, tensor<3x10xf32>, tensor, tensor, tensor) -> tensor // CHECK: return [[RESULT]] : tensor } @@ -510,7 +511,7 @@ func.func @tensorlistConcat(%arg0: tensor<3x2x2xf32>, %lead: tensor) -> (te func.return %t#0, %t#1 : tensor, tensor<0xi64> // CHECK: [[ELEMENT_SHAPE:%.*]] = arith.constant dense<2> : tensor<2xi32> -// CHECK: [[UNPACK:%.*]]:3 = "tf.Unpack"(%arg0) {axis = 0 : i64} : (tensor<3x2x2xf32>) -> (tensor<2x2xf32>, tensor<2x2xf32>, tensor<2x2xf32>) +// CHECK: [[UNPACK:%.*]]:3 = "tf.Unpack"(%arg0) <{axis = 0 : i64}> : (tensor<3x2x2xf32>) -> (tensor<2x2xf32>, tensor<2x2xf32>, tensor<2x2xf32>) // CHECK: [[SCALAR_ZERO:%.*]] = arith.constant dense<0> : tensor // CHECK: [[CONCAT:%.*]] = "tf.Concat"([[SCALAR_ZERO]], [[UNPACK]]#0, [[UNPACK]]#1, [[UNPACK]]#2) : (tensor, tensor<2x2xf32>, tensor<2x2xf32>, tensor<2x2xf32>) -> tensor // CHECK: [[LENGTHS:%.*]] = arith.constant dense<0> : tensor<0xi64> @@ -567,7 +568,7 @@ func.func @tensorListIf(%arg0: tensor<3x10xf32>, %arg1: tensor<1xi32>, %arg2: te // CHECK: func @tensorListIf // CHECK-NEXT: %cst = arith.constant dense<2> : tensor // CHECK-NEXT: %0 = "tf.Less"(%arg2, %cst) : (tensor, tensor) -> tensor -// CHECK-NEXT: %1 = "tf.If"(%0, %arg0) {else_branch = @tensorListIfCondFalse, is_stateless = true, then_branch = @tensorListIfCondTrue} : (tensor, tensor<3x10xf32>) -> tensor<3x10xf32> +// CHECK-NEXT: %1 = "tf.If"(%0, %arg0) <{else_branch = @tensorListIfCondFalse, is_stateless = true, then_branch = @tensorListIfCondTrue}> : (tensor, tensor<3x10xf32>) -> tensor<3x10xf32> // CHECK-NEXT: return %1 : tensor<3x10xf32> } diff --git a/tensorflow/compiler/mlir/lite/tests/mlir2flatbuffer/metadata_buffer.mlir b/tensorflow/compiler/mlir/lite/tests/mlir2flatbuffer/metadata_buffer.mlir new file mode 100644 index 00000000000000..f53f3954f14211 --- /dev/null +++ b/tensorflow/compiler/mlir/lite/tests/mlir2flatbuffer/metadata_buffer.mlir @@ -0,0 +1,11 @@ +// RUN: flatbuffer_translate -mlir-to-tflite-flatbuffer %s -o - | flatbuffer_to_string - | FileCheck %s + +module attributes {tfl.metadata_buffer = [3 : i32, 7 : i32]} { + func.func @main(%arg0: tensor, %arg1: tensor<3x2xi32>) -> tensor<3x2xi32> { + %0 = "tfl.add" (%arg0, %arg1) {fused_activation_function = "NONE"} : (tensor, tensor<3x2xi32>) -> tensor<3x2xi32> + func.return %0 : tensor<3x2xi32> + } +} + +// CHECK: metadata_buffer: [ 3, 7 ], +// CHECK-NEXT: metadata: \ No newline at end of file diff --git a/tensorflow/compiler/mlir/lite/tests/prepare-composite-functions-tf.mlir b/tensorflow/compiler/mlir/lite/tests/prepare-composite-functions-tf.mlir index 7aac8662a8373a..27d98c7599c93d 100644 --- a/tensorflow/compiler/mlir/lite/tests/prepare-composite-functions-tf.mlir +++ b/tensorflow/compiler/mlir/lite/tests/prepare-composite-functions-tf.mlir @@ -194,24 +194,24 @@ func.func @inference_standard_lstm_time_major(%arg0: tensor, %arg1: t // CHECK: [[VAL_7:%.*]] = "tf.Transpose"([[VAL_3]], [[VAL_6]]) : (tensor<8x40xf32>, tensor<2xi32>) -> tensor<40x8xf32> // CHECK: [[VAL_8:%.*]] = arith.constant dense<[1, 0]> : tensor<2xi32> // CHECK: [[VAL_9:%.*]] = "tf.Transpose"([[VAL_4]], [[VAL_8]]) : (tensor<10x40xf32>, tensor<2xi32>) -> tensor<40x10xf32> -// CHECK-DAG: [[VAL_10:%.*]] = "tf.Const"() {value = dense<10> : tensor<4xi32>} : () -> tensor<4xi32> -// CHECK-DAG: [[VAL_11:%.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor +// CHECK-DAG: [[VAL_10:%.*]] = "tf.Const"() <{value = dense<10> : tensor<4xi32>}> : () -> tensor<4xi32> +// CHECK-DAG: [[VAL_11:%.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK: [[VAL_12:%.*]]:4 = "tf.SplitV"([[VAL_7]], [[VAL_10]], [[VAL_11]]) : (tensor<40x8xf32>, tensor<4xi32>, tensor) -> (tensor<10x8xf32>, tensor<10x8xf32>, tensor<10x8xf32>, tensor<10x8xf32>) -// CHECK-DAG: [[VAL_13:%.*]] = "tf.Const"() {value = dense<10> : tensor<4xi32>} : () -> tensor<4xi32> -// CHECK-DAG: [[VAL_14:%.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor +// CHECK-DAG: [[VAL_13:%.*]] = "tf.Const"() <{value = dense<10> : tensor<4xi32>}> : () -> tensor<4xi32> +// CHECK-DAG: [[VAL_14:%.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK: [[VAL_15:%.*]]:4 = "tf.SplitV"([[VAL_9]], [[VAL_13]], [[VAL_14]]) : (tensor<40x10xf32>, tensor<4xi32>, tensor) -> (tensor<10x10xf32>, tensor<10x10xf32>, tensor<10x10xf32>, tensor<10x10xf32>) -// CHECK-DAG: [[VAL_16:%.*]] = "tf.Const"() {value = dense<10> : tensor<4xi32>} : () -> tensor<4xi32> -// CHECK-DAG: [[VAL_17:%.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor +// CHECK-DAG: [[VAL_16:%.*]] = "tf.Const"() <{value = dense<10> : tensor<4xi32>}> : () -> tensor<4xi32> +// CHECK-DAG: [[VAL_17:%.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK: [[VAL_18:%.*]]:4 = "tf.SplitV"([[VAL_5]], [[VAL_16]], [[VAL_17]]) : (tensor<40xf32>, tensor<4xi32>, tensor) -> (tensor<10xf32>, tensor<10xf32>, tensor<10xf32>, tensor<10xf32>) // CHECK: [[VAL_19:%.*]] = "tfl.no_value"() {value} : () -> none // CHECK: [[VAL_20:%.*]] = "tfl.unidirectional_sequence_lstm"([[VAL_0]], [[VAL_12]]#0, [[VAL_12]]#1, [[VAL_12]]#2, [[VAL_12]]#3, [[VAL_15]]#0, [[VAL_15]]#1, [[VAL_15]]#2, [[VAL_15]]#3, [[VAL_19]], [[VAL_19]], [[VAL_19]], [[VAL_18]]#0, [[VAL_18]]#1, [[VAL_18]]#2, [[VAL_18]]#3, [[VAL_19]], [[VAL_19]], [[VAL_1]], [[VAL_2]], [[VAL_19]], [[VAL_19]], [[VAL_19]], [[VAL_19]]) {cell_clip = 1.000000e+01 : f32, diagonal_recurrent_tensors = false, fused_activation_function = "TANH", proj_clip = 0.000000e+00 : f32, time_major = true} : (tensor, tensor<10x8xf32>, tensor<10x8xf32>, tensor<10x8xf32>, tensor<10x8xf32>, tensor<10x10xf32>, tensor<10x10xf32>, tensor<10x10xf32>, tensor<10x10xf32>, none, none, none, tensor<10xf32>, tensor<10xf32>, tensor<10xf32>, tensor<10xf32>, none, none, tensor<8x10xf32>, tensor<8x10xf32>, none, none, none, none) -> tensor // CHECK-DAG: [[VAL_21:%.*]] = arith.constant dense<[-1, 0, 0]> : tensor<3xi32> // CHECK-DAG: [[VAL_22:%.*]] = arith.constant dense<0> : tensor<3xi32> // CHECK-DAG: [[VAL_23:%.*]] = arith.constant dense<1> : tensor<3xi32> -// CHECK: [[VAL_24:%.*]] = "tf.StridedSlice"([[VAL_20]], [[VAL_21]], [[VAL_22]], [[VAL_23]]) {begin_mask = 6 : i64, ellipsis_mask = 0 : i64, end_mask = 6 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 1 : i64} : (tensor, tensor<3xi32>, tensor<3xi32>, tensor<3xi32>) -> tensor<8x10xf32> -// CHECK-DAG: [[VAL_25:%.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<1xf32>} : () -> tensor<8x10xf32> -// CHECK-DAG: [[VAL_26:%.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<1xf32>} : () -> tensor<8x10xf32> -// CHECK-DAG: [[VAL_27:%.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<1xf32>} : () -> tensor +// CHECK: [[VAL_24:%.*]] = "tf.StridedSlice"([[VAL_20]], [[VAL_21]], [[VAL_22]], [[VAL_23]]) <{begin_mask = 6 : i64, ellipsis_mask = 0 : i64, end_mask = 6 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 1 : i64}> : (tensor, tensor<3xi32>, tensor<3xi32>, tensor<3xi32>) -> tensor<8x10xf32> +// CHECK-DAG: [[VAL_25:%.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<1xf32>}> : () -> tensor<8x10xf32> +// CHECK-DAG: [[VAL_26:%.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<1xf32>}> : () -> tensor<8x10xf32> +// CHECK-DAG: [[VAL_27:%.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<1xf32>}> : () -> tensor // CHECK: return [[VAL_24]], [[VAL_20]], [[VAL_25]], [[VAL_26]], [[VAL_27]] : tensor<8x10xf32>, tensor, tensor<8x10xf32>, tensor<8x10xf32>, tensor // CHECK: } } @@ -240,32 +240,32 @@ func.func @inference_standard_indy_lstm_time_major(%arg0: tensor<8x8x8xf32>, %ar // CHECK: [[VAL_7:%.*]] = "tf.Transpose"([[VAL_3]], [[VAL_6]]) : (tensor<8x40xf32>, tensor<2xi32>) -> tensor<40x8xf32> // CHECK: [[VAL_8:%.*]] = arith.constant dense<[1, 0]> : tensor<2xi32> // CHECK: [[VAL_9:%.*]] = "tf.Transpose"([[VAL_4]], [[VAL_8]]) : (tensor<10x4xf32>, tensor<2xi32>) -> tensor<4x10xf32> -// CHECK-DAG: [[VAL_10:%.*]] = "tf.Const"() {value = dense<10> : tensor<4xi32>} : () -> tensor<4xi32> -// CHECK-DAG: [[VAL_11:%.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor +// CHECK-DAG: [[VAL_10:%.*]] = "tf.Const"() <{value = dense<10> : tensor<4xi32>}> : () -> tensor<4xi32> +// CHECK-DAG: [[VAL_11:%.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK: [[VAL_12:%.*]]:4 = "tf.SplitV"([[VAL_7]], [[VAL_10]], [[VAL_11]]) : (tensor<40x8xf32>, tensor<4xi32>, tensor) -> (tensor<10x8xf32>, tensor<10x8xf32>, tensor<10x8xf32>, tensor<10x8xf32>) -// CHECK-DAG: [[VAL_13:%.*]] = "tf.Const"() {value = dense<1> : tensor<4xi32>} : () -> tensor<4xi32> -// CHECK-DAG: [[VAL_14:%.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor +// CHECK-DAG: [[VAL_13:%.*]] = "tf.Const"() <{value = dense<1> : tensor<4xi32>}> : () -> tensor<4xi32> +// CHECK-DAG: [[VAL_14:%.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK: [[VAL_15:%.*]]:4 = "tf.SplitV"([[VAL_9]], [[VAL_13]], [[VAL_14]]) : (tensor<4x10xf32>, tensor<4xi32>, tensor) -> (tensor<1x10xf32>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<1x10xf32>) -// CHECK-DAG: [[VAL_20:%.*]] = "tf.Const"() {value = dense<-1> : tensor<1xi32>} : () -> tensor<1xi32> +// CHECK-DAG: [[VAL_20:%.*]] = "tf.Const"() <{value = dense<-1> : tensor<1xi32>}> : () -> tensor<1xi32> // CHECK: [[VAL_21:%.*]] = "tf.Reshape"([[VAL_15]]#0, [[VAL_20]]) : (tensor<1x10xf32>, tensor<1xi32>) -> tensor<10xf32> -// CHECK-DAG: [[VAL_22:%.*]] = "tf.Const"() {value = dense<-1> : tensor<1xi32>} : () -> tensor<1xi32> +// CHECK-DAG: [[VAL_22:%.*]] = "tf.Const"() <{value = dense<-1> : tensor<1xi32>}> : () -> tensor<1xi32> // CHECK: [[VAL_23:%.*]] = "tf.Reshape"([[VAL_15]]#1, [[VAL_22]]) : (tensor<1x10xf32>, tensor<1xi32>) -> tensor<10xf32> -// CHECK-DAG: [[VAL_24:%.*]] = "tf.Const"() {value = dense<-1> : tensor<1xi32>} : () -> tensor<1xi32> +// CHECK-DAG: [[VAL_24:%.*]] = "tf.Const"() <{value = dense<-1> : tensor<1xi32>}> : () -> tensor<1xi32> // CHECK: [[VAL_25:%.*]] = "tf.Reshape"([[VAL_15]]#2, [[VAL_24]]) : (tensor<1x10xf32>, tensor<1xi32>) -> tensor<10xf32> -// CHECK-DAG: [[VAL_26:%.*]] = "tf.Const"() {value = dense<-1> : tensor<1xi32>} : () -> tensor<1xi32> +// CHECK-DAG: [[VAL_26:%.*]] = "tf.Const"() <{value = dense<-1> : tensor<1xi32>}> : () -> tensor<1xi32> // CHECK: [[VAL_27:%.*]] = "tf.Reshape"([[VAL_15]]#3, [[VAL_26]]) : (tensor<1x10xf32>, tensor<1xi32>) -> tensor<10xf32> -// CHECK-DAG: [[VAL_28:%.*]] = "tf.Const"() {value = dense<10> : tensor<4xi32>} : () -> tensor<4xi32> -// CHECK-DAG: [[VAL_29:%.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor +// CHECK-DAG: [[VAL_28:%.*]] = "tf.Const"() <{value = dense<10> : tensor<4xi32>}> : () -> tensor<4xi32> +// CHECK-DAG: [[VAL_29:%.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK: [[VAL_30:%.*]]:4 = "tf.SplitV"([[VAL_5]], [[VAL_28]], [[VAL_29]]) : (tensor<40xf32>, tensor<4xi32>, tensor) -> (tensor<10xf32>, tensor<10xf32>, tensor<10xf32>, tensor<10xf32>) // CHECK: [[VAL_31:%.*]] = "tfl.no_value"() {value} : () -> none // CHECK: [[VAL_32:%.*]] = "tfl.unidirectional_sequence_lstm"([[VAL_0]], [[VAL_12]]#0, [[VAL_12]]#1, [[VAL_12]]#2, [[VAL_12]]#3, [[VAL_21]], [[VAL_23]], [[VAL_25]], [[VAL_27]], [[VAL_31]], [[VAL_31]], [[VAL_31]], [[VAL_30]]#0, [[VAL_30]]#1, [[VAL_30]]#2, [[VAL_30]]#3, [[VAL_31]], [[VAL_31]], [[VAL_1]], [[VAL_2]], [[VAL_31]], [[VAL_31]], [[VAL_31]], [[VAL_31]]) {cell_clip = 1.000000e+01 : f32, diagonal_recurrent_tensors = true, fused_activation_function = "TANH", proj_clip = 0.000000e+00 : f32, time_major = true} : (tensor<8x8x8xf32>, tensor<10x8xf32>, tensor<10x8xf32>, tensor<10x8xf32>, tensor<10x8xf32>, tensor<10xf32>, tensor<10xf32>, tensor<10xf32>, tensor<10xf32>, none, none, none, tensor<10xf32>, tensor<10xf32>, tensor<10xf32>, tensor<10xf32>, none, none, tensor<8x10xf32>, tensor<8x10xf32>, none, none, none, none) -> tensor<8x8x10xf32> // CHECK-DAG: [[VAL_33:%.*]] = arith.constant dense<[-1, 0, 0]> : tensor<3xi32> // CHECK-DAG: [[VAL_34:%.*]] = arith.constant dense<0> : tensor<3xi32> // CHECK-DAG: [[VAL_35:%.*]] = arith.constant dense<1> : tensor<3xi32> -// CHECK: [[VAL_36:%.*]] = "tf.StridedSlice"([[VAL_32]], [[VAL_33]], [[VAL_34]], [[VAL_35]]) {begin_mask = 6 : i64, ellipsis_mask = 0 : i64, end_mask = 6 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 1 : i64} : (tensor<8x8x10xf32>, tensor<3xi32>, tensor<3xi32>, tensor<3xi32>) -> tensor<8x10xf32> -// CHECK-DAG: [[VAL_37:%.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<1xf32>} : () -> tensor<8x10xf32> -// CHECK-DAG: [[VAL_38:%.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<1xf32>} : () -> tensor<8x10xf32> -// CHECK-DAG: [[VAL_39:%.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<1xf32>} : () -> tensor +// CHECK: [[VAL_36:%.*]] = "tf.StridedSlice"([[VAL_32]], [[VAL_33]], [[VAL_34]], [[VAL_35]]) <{begin_mask = 6 : i64, ellipsis_mask = 0 : i64, end_mask = 6 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 1 : i64}> : (tensor<8x8x10xf32>, tensor<3xi32>, tensor<3xi32>, tensor<3xi32>) -> tensor<8x10xf32> +// CHECK-DAG: [[VAL_37:%.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<1xf32>}> : () -> tensor<8x10xf32> +// CHECK-DAG: [[VAL_38:%.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<1xf32>}> : () -> tensor<8x10xf32> +// CHECK-DAG: [[VAL_39:%.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<1xf32>}> : () -> tensor // CHECK: return [[VAL_36]], [[VAL_32]], [[VAL_37]], [[VAL_38]], [[VAL_39]] : tensor<8x10xf32>, tensor<8x8x10xf32>, tensor<8x10xf32>, tensor<8x10xf32>, tensor // CHECK: } @@ -290,24 +290,24 @@ func.func @inference_standard_lstm_non_time_major(%arg0: tensor<8x8x8xf32>, %arg // CHECK: [[VAL_7:%.*]] = "tf.Transpose"([[VAL_3]], [[VAL_6]]) : (tensor<8x40xf32>, tensor<2xi32>) -> tensor<40x8xf32> // CHECK: [[VAL_8:%.*]] = arith.constant dense<[1, 0]> : tensor<2xi32> // CHECK: [[VAL_9:%.*]] = "tf.Transpose"([[VAL_4]], [[VAL_8]]) : (tensor<10x40xf32>, tensor<2xi32>) -> tensor<40x10xf32> -// CHECK-DAG: [[VAL_10:%.*]] = "tf.Const"() {value = dense<10> : tensor<4xi32>} : () -> tensor<4xi32> -// CHECK-DAG: [[VAL_11:%.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor +// CHECK-DAG: [[VAL_10:%.*]] = "tf.Const"() <{value = dense<10> : tensor<4xi32>}> : () -> tensor<4xi32> +// CHECK-DAG: [[VAL_11:%.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK: [[VAL_12:%.*]]:4 = "tf.SplitV"([[VAL_7]], [[VAL_10]], [[VAL_11]]) : (tensor<40x8xf32>, tensor<4xi32>, tensor) -> (tensor<10x8xf32>, tensor<10x8xf32>, tensor<10x8xf32>, tensor<10x8xf32>) -// CHECK-DAG: [[VAL_13:%.*]] = "tf.Const"() {value = dense<10> : tensor<4xi32>} : () -> tensor<4xi32> -// CHECK-DAG: [[VAL_14:%.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor +// CHECK-DAG: [[VAL_13:%.*]] = "tf.Const"() <{value = dense<10> : tensor<4xi32>}> : () -> tensor<4xi32> +// CHECK-DAG: [[VAL_14:%.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK: [[VAL_15:%.*]]:4 = "tf.SplitV"([[VAL_9]], [[VAL_13]], [[VAL_14]]) : (tensor<40x10xf32>, tensor<4xi32>, tensor) -> (tensor<10x10xf32>, tensor<10x10xf32>, tensor<10x10xf32>, tensor<10x10xf32>) -// CHECK-DAG: [[VAL_16:%.*]] = "tf.Const"() {value = dense<10> : tensor<4xi32>} : () -> tensor<4xi32> -// CHECK-DAG: [[VAL_17:%.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor +// CHECK-DAG: [[VAL_16:%.*]] = "tf.Const"() <{value = dense<10> : tensor<4xi32>}> : () -> tensor<4xi32> +// CHECK-DAG: [[VAL_17:%.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK: [[VAL_18:%.*]]:4 = "tf.SplitV"([[VAL_5]], [[VAL_16]], [[VAL_17]]) : (tensor<40xf32>, tensor<4xi32>, tensor) -> (tensor<10xf32>, tensor<10xf32>, tensor<10xf32>, tensor<10xf32>) // CHECK: [[VAL_19:%.*]] = "tfl.no_value"() {value} : () -> none // CHECK: [[VAL_20:%.*]] = "tfl.unidirectional_sequence_lstm"([[VAL_0]], [[VAL_12]]#0, [[VAL_12]]#1, [[VAL_12]]#2, [[VAL_12]]#3, [[VAL_15]]#0, [[VAL_15]]#1, [[VAL_15]]#2, [[VAL_15]]#3, [[VAL_19]], [[VAL_19]], [[VAL_19]], [[VAL_18]]#0, [[VAL_18]]#1, [[VAL_18]]#2, [[VAL_18]]#3, [[VAL_19]], [[VAL_19]], [[VAL_1]], [[VAL_2]], [[VAL_19]], [[VAL_19]], [[VAL_19]], [[VAL_19]]) {cell_clip = 1.000000e+01 : f32, diagonal_recurrent_tensors = false, fused_activation_function = "TANH", proj_clip = 0.000000e+00 : f32, time_major = false} : (tensor<8x8x8xf32>, tensor<10x8xf32>, tensor<10x8xf32>, tensor<10x8xf32>, tensor<10x8xf32>, tensor<10x10xf32>, tensor<10x10xf32>, tensor<10x10xf32>, tensor<10x10xf32>, none, none, none, tensor<10xf32>, tensor<10xf32>, tensor<10xf32>, tensor<10xf32>, none, none, tensor<8x10xf32>, tensor<8x10xf32>, none, none, none, none) -> tensor<8x8x10xf32> // CHECK-DAG: [[VAL_21:%.*]] = arith.constant dense<[0, -1, 0]> : tensor<3xi32> // CHECK-DAG: [[VAL_22:%.*]] = arith.constant dense<0> : tensor<3xi32> // CHECK-DAG: [[VAL_23:%.*]] = arith.constant dense<1> : tensor<3xi32> -// CHECK: [[VAL_24:%.*]] = "tf.StridedSlice"([[VAL_20]], [[VAL_21]], [[VAL_22]], [[VAL_23]]) {begin_mask = 5 : i64, ellipsis_mask = 0 : i64, end_mask = 5 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 2 : i64} : (tensor<8x8x10xf32>, tensor<3xi32>, tensor<3xi32>, tensor<3xi32>) -> tensor<8x10xf32> -// CHECK-DAG: [[VAL_25:%.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<1xf32>} : () -> tensor<8x10xf32> -// CHECK-DAG: [[VAL_26:%.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<1xf32>} : () -> tensor<8x10xf32> -// CHECK-DAG: [[VAL_27:%.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<1xf32>} : () -> tensor +// CHECK: [[VAL_24:%.*]] = "tf.StridedSlice"([[VAL_20]], [[VAL_21]], [[VAL_22]], [[VAL_23]]) <{begin_mask = 5 : i64, ellipsis_mask = 0 : i64, end_mask = 5 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 2 : i64}> : (tensor<8x8x10xf32>, tensor<3xi32>, tensor<3xi32>, tensor<3xi32>) -> tensor<8x10xf32> +// CHECK-DAG: [[VAL_25:%.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<1xf32>}> : () -> tensor<8x10xf32> +// CHECK-DAG: [[VAL_26:%.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<1xf32>}> : () -> tensor<8x10xf32> +// CHECK-DAG: [[VAL_27:%.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<1xf32>}> : () -> tensor // CHECK: return [[VAL_24]], [[VAL_20]], [[VAL_25]], [[VAL_26]], [[VAL_27]] : tensor<8x10xf32>, tensor<8x8x10xf32>, tensor<8x10xf32>, tensor<8x10xf32>, tensor // CHECK: } @@ -337,32 +337,32 @@ func.func @inference_standard_indy_lstm_non_time_major(%arg0: tensor<8x8x8xf32>, // CHECK: [[VAL_7:%.*]] = "tf.Transpose"([[VAL_3]], [[VAL_6]]) : (tensor<8x40xf32>, tensor<2xi32>) -> tensor<40x8xf32> // CHECK: [[VAL_8:%.*]] = arith.constant dense<[1, 0]> : tensor<2xi32> // CHECK: [[VAL_9:%.*]] = "tf.Transpose"([[VAL_4]], [[VAL_8]]) : (tensor<10x4xf32>, tensor<2xi32>) -> tensor<4x10xf32> -// CHECK-DAG: [[VAL_10:%.*]] = "tf.Const"() {value = dense<10> : tensor<4xi32>} : () -> tensor<4xi32> -// CHECK-DAG: [[VAL_11:%.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor +// CHECK-DAG: [[VAL_10:%.*]] = "tf.Const"() <{value = dense<10> : tensor<4xi32>}> : () -> tensor<4xi32> +// CHECK-DAG: [[VAL_11:%.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK: [[VAL_12:%.*]]:4 = "tf.SplitV"([[VAL_7]], [[VAL_10]], [[VAL_11]]) : (tensor<40x8xf32>, tensor<4xi32>, tensor) -> (tensor<10x8xf32>, tensor<10x8xf32>, tensor<10x8xf32>, tensor<10x8xf32>) -// CHECK-DAG: [[VAL_13:%.*]] = "tf.Const"() {value = dense<1> : tensor<4xi32>} : () -> tensor<4xi32> -// CHECK-DAG: [[VAL_14:%.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor +// CHECK-DAG: [[VAL_13:%.*]] = "tf.Const"() <{value = dense<1> : tensor<4xi32>}> : () -> tensor<4xi32> +// CHECK-DAG: [[VAL_14:%.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK: [[VAL_15:%.*]]:4 = "tf.SplitV"([[VAL_9]], [[VAL_13]], [[VAL_14]]) : (tensor<4x10xf32>, tensor<4xi32>, tensor) -> (tensor<1x10xf32>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<1x10xf32>) -// CHECK-DAG: [[VAL_20:%.*]] = "tf.Const"() {value = dense<-1> : tensor<1xi32>} : () -> tensor<1xi32> +// CHECK-DAG: [[VAL_20:%.*]] = "tf.Const"() <{value = dense<-1> : tensor<1xi32>}> : () -> tensor<1xi32> // CHECK: [[VAL_21:%.*]] = "tf.Reshape"([[VAL_15]]#0, [[VAL_20]]) : (tensor<1x10xf32>, tensor<1xi32>) -> tensor<10xf32> -// CHECK-DAG: [[VAL_22:%.*]] = "tf.Const"() {value = dense<-1> : tensor<1xi32>} : () -> tensor<1xi32> +// CHECK-DAG: [[VAL_22:%.*]] = "tf.Const"() <{value = dense<-1> : tensor<1xi32>}> : () -> tensor<1xi32> // CHECK: [[VAL_23:%.*]] = "tf.Reshape"([[VAL_15]]#1, [[VAL_22]]) : (tensor<1x10xf32>, tensor<1xi32>) -> tensor<10xf32> -// CHECK-DAG: [[VAL_24:%.*]] = "tf.Const"() {value = dense<-1> : tensor<1xi32>} : () -> tensor<1xi32> +// CHECK-DAG: [[VAL_24:%.*]] = "tf.Const"() <{value = dense<-1> : tensor<1xi32>}> : () -> tensor<1xi32> // CHECK: [[VAL_25:%.*]] = "tf.Reshape"([[VAL_15]]#2, [[VAL_24]]) : (tensor<1x10xf32>, tensor<1xi32>) -> tensor<10xf32> -// CHECK-DAG: [[VAL_26:%.*]] = "tf.Const"() {value = dense<-1> : tensor<1xi32>} : () -> tensor<1xi32> +// CHECK-DAG: [[VAL_26:%.*]] = "tf.Const"() <{value = dense<-1> : tensor<1xi32>}> : () -> tensor<1xi32> // CHECK: [[VAL_27:%.*]] = "tf.Reshape"([[VAL_15]]#3, [[VAL_26]]) : (tensor<1x10xf32>, tensor<1xi32>) -> tensor<10xf32> -// CHECK-DAG: [[VAL_28:%.*]] = "tf.Const"() {value = dense<10> : tensor<4xi32>} : () -> tensor<4xi32> -// CHECK-DAG: [[VAL_29:%.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor +// CHECK-DAG: [[VAL_28:%.*]] = "tf.Const"() <{value = dense<10> : tensor<4xi32>}> : () -> tensor<4xi32> +// CHECK-DAG: [[VAL_29:%.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK: [[VAL_30:%.*]]:4 = "tf.SplitV"([[VAL_5]], [[VAL_28]], [[VAL_29]]) : (tensor<40xf32>, tensor<4xi32>, tensor) -> (tensor<10xf32>, tensor<10xf32>, tensor<10xf32>, tensor<10xf32>) // CHECK: [[VAL_31:%.*]] = "tfl.no_value"() {value} : () -> none // CHECK: [[VAL_32:%.*]] = "tfl.unidirectional_sequence_lstm"([[VAL_0]], [[VAL_12]]#0, [[VAL_12]]#1, [[VAL_12]]#2, [[VAL_12]]#3, [[VAL_21]], [[VAL_23]], [[VAL_25]], [[VAL_27]], [[VAL_31]], [[VAL_31]], [[VAL_31]], [[VAL_30]]#0, [[VAL_30]]#1, [[VAL_30]]#2, [[VAL_30]]#3, [[VAL_31]], [[VAL_31]], [[VAL_1]], [[VAL_2]], [[VAL_31]], [[VAL_31]], [[VAL_31]], [[VAL_31]]) {cell_clip = 1.000000e+01 : f32, diagonal_recurrent_tensors = true, fused_activation_function = "TANH", proj_clip = 0.000000e+00 : f32, time_major = false} : (tensor<8x8x8xf32>, tensor<10x8xf32>, tensor<10x8xf32>, tensor<10x8xf32>, tensor<10x8xf32>, tensor<10xf32>, tensor<10xf32>, tensor<10xf32>, tensor<10xf32>, none, none, none, tensor<10xf32>, tensor<10xf32>, tensor<10xf32>, tensor<10xf32>, none, none, tensor<8x10xf32>, tensor<8x10xf32>, none, none, none, none) -> tensor<8x8x10xf32> // CHECK-DAG: [[VAL_33:%.*]] = arith.constant dense<[0, -1, 0]> : tensor<3xi32> // CHECK-DAG: [[VAL_34:%.*]] = arith.constant dense<0> : tensor<3xi32> // CHECK-DAG: [[VAL_35:%.*]] = arith.constant dense<1> : tensor<3xi32> -// CHECK: [[VAL_36:%.*]] = "tf.StridedSlice"([[VAL_32]], [[VAL_33]], [[VAL_34]], [[VAL_35]]) {begin_mask = 5 : i64, ellipsis_mask = 0 : i64, end_mask = 5 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 2 : i64} : (tensor<8x8x10xf32>, tensor<3xi32>, tensor<3xi32>, tensor<3xi32>) -> tensor<8x10xf32> -// CHECK-DAG: [[VAL_37:%.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<1xf32>} : () -> tensor<8x10xf32> -// CHECK-DAG: [[VAL_38:%.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<1xf32>} : () -> tensor<8x10xf32> -// CHECK-DAG: [[VAL_39:%.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<1xf32>} : () -> tensor +// CHECK: [[VAL_36:%.*]] = "tf.StridedSlice"([[VAL_32]], [[VAL_33]], [[VAL_34]], [[VAL_35]]) <{begin_mask = 5 : i64, ellipsis_mask = 0 : i64, end_mask = 5 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 2 : i64}> : (tensor<8x8x10xf32>, tensor<3xi32>, tensor<3xi32>, tensor<3xi32>) -> tensor<8x10xf32> +// CHECK-DAG: [[VAL_37:%.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<1xf32>}> : () -> tensor<8x10xf32> +// CHECK-DAG: [[VAL_38:%.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<1xf32>}> : () -> tensor<8x10xf32> +// CHECK-DAG: [[VAL_39:%.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<1xf32>}> : () -> tensor // CHECK: return [[VAL_36]], [[VAL_32]], [[VAL_37]], [[VAL_38]], [[VAL_39]] : tensor<8x10xf32>, tensor<8x8x10xf32>, tensor<8x10xf32>, tensor<8x10xf32>, tensor // CHECK: } @@ -389,24 +389,24 @@ func.func @inference_standard_lstm_time_major_go_backwards(%arg0: tensor, tensor<2xi32>) -> tensor<40x8xf32> // CHECK: [[VAL_10:%.*]] = arith.constant dense<[1, 0]> : tensor<2xi32> // CHECK: [[VAL_11:%.*]] = "tf.Transpose"([[VAL_4]], [[VAL_10]]) : (tensor<10x40xf32>, tensor<2xi32>) -> tensor<40x10xf32> -// CHECK-DAG: [[VAL_12:%.*]] = "tf.Const"() {value = dense<10> : tensor<4xi32>} : () -> tensor<4xi32> -// CHECK-DAG: [[VAL_13:%.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor +// CHECK-DAG: [[VAL_12:%.*]] = "tf.Const"() <{value = dense<10> : tensor<4xi32>}> : () -> tensor<4xi32> +// CHECK-DAG: [[VAL_13:%.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK: [[VAL_14:%.*]]:4 = "tf.SplitV"([[VAL_9]], [[VAL_12]], [[VAL_13]]) : (tensor<40x8xf32>, tensor<4xi32>, tensor) -> (tensor<10x8xf32>, tensor<10x8xf32>, tensor<10x8xf32>, tensor<10x8xf32>) -// CHECK-DAG: [[VAL_15:%.*]] = "tf.Const"() {value = dense<10> : tensor<4xi32>} : () -> tensor<4xi32> -// CHECK-DAG: [[VAL_16:%.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor +// CHECK-DAG: [[VAL_15:%.*]] = "tf.Const"() <{value = dense<10> : tensor<4xi32>}> : () -> tensor<4xi32> +// CHECK-DAG: [[VAL_16:%.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK: [[VAL_17:%.*]]:4 = "tf.SplitV"([[VAL_11]], [[VAL_15]], [[VAL_16]]) : (tensor<40x10xf32>, tensor<4xi32>, tensor) -> (tensor<10x10xf32>, tensor<10x10xf32>, tensor<10x10xf32>, tensor<10x10xf32>) -// CHECK-DAG: [[VAL_18:%.*]] = "tf.Const"() {value = dense<10> : tensor<4xi32>} : () -> tensor<4xi32> -// CHECK-DAG: [[VAL_19:%.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor +// CHECK-DAG: [[VAL_18:%.*]] = "tf.Const"() <{value = dense<10> : tensor<4xi32>}> : () -> tensor<4xi32> +// CHECK-DAG: [[VAL_19:%.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK: [[VAL_20:%.*]]:4 = "tf.SplitV"([[VAL_5]], [[VAL_18]], [[VAL_19]]) : (tensor<40xf32>, tensor<4xi32>, tensor) -> (tensor<10xf32>, tensor<10xf32>, tensor<10xf32>, tensor<10xf32>) // CHECK: [[VAL_21:%.*]] = "tfl.no_value"() {value} : () -> none // CHECK: [[VAL_22:%.*]] = "tfl.unidirectional_sequence_lstm"([[VAL_7]], [[VAL_14]]#0, [[VAL_14]]#1, [[VAL_14]]#2, [[VAL_14]]#3, [[VAL_17]]#0, [[VAL_17]]#1, [[VAL_17]]#2, [[VAL_17]]#3, [[VAL_21]], [[VAL_21]], [[VAL_21]], [[VAL_20]]#0, [[VAL_20]]#1, [[VAL_20]]#2, [[VAL_20]]#3, [[VAL_21]], [[VAL_21]], [[VAL_1]], [[VAL_2]], [[VAL_21]], [[VAL_21]], [[VAL_21]], [[VAL_21]]) {cell_clip = 1.000000e+01 : f32, diagonal_recurrent_tensors = false, fused_activation_function = "TANH", proj_clip = 0.000000e+00 : f32, time_major = true} : (tensor, tensor<10x8xf32>, tensor<10x8xf32>, tensor<10x8xf32>, tensor<10x8xf32>, tensor<10x10xf32>, tensor<10x10xf32>, tensor<10x10xf32>, tensor<10x10xf32>, none, none, none, tensor<10xf32>, tensor<10xf32>, tensor<10xf32>, tensor<10xf32>, none, none, tensor<8x10xf32>, tensor<8x10xf32>, none, none, none, none) -> tensor // CHECK-DAG: [[VAL_23:%.*]] = arith.constant dense<[-1, 0, 0]> : tensor<3xi32> // CHECK-DAG: [[VAL_24:%.*]] = arith.constant dense<0> : tensor<3xi32> // CHECK-DAG: [[VAL_25:%.*]] = arith.constant dense<1> : tensor<3xi32> -// CHECK: [[VAL_26:%.*]] = "tf.StridedSlice"([[VAL_22]], [[VAL_23]], [[VAL_24]], [[VAL_25]]) {begin_mask = 6 : i64, ellipsis_mask = 0 : i64, end_mask = 6 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 1 : i64} : (tensor, tensor<3xi32>, tensor<3xi32>, tensor<3xi32>) -> tensor<8x10xf32> -// CHECK-DAG: [[VAL_27:%.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<1xf32>} : () -> tensor<8x10xf32> -// CHECK-DAG: [[VAL_28:%.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<1xf32>} : () -> tensor<8x10xf32> -// CHECK-DAG: [[VAL_29:%.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<1xf32>} : () -> tensor +// CHECK: [[VAL_26:%.*]] = "tf.StridedSlice"([[VAL_22]], [[VAL_23]], [[VAL_24]], [[VAL_25]]) <{begin_mask = 6 : i64, ellipsis_mask = 0 : i64, end_mask = 6 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 1 : i64}> : (tensor, tensor<3xi32>, tensor<3xi32>, tensor<3xi32>) -> tensor<8x10xf32> +// CHECK-DAG: [[VAL_27:%.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<1xf32>}> : () -> tensor<8x10xf32> +// CHECK-DAG: [[VAL_28:%.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<1xf32>}> : () -> tensor<8x10xf32> +// CHECK-DAG: [[VAL_29:%.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<1xf32>}> : () -> tensor // CHECK: return [[VAL_26]], [[VAL_22]], [[VAL_27]], [[VAL_28]], [[VAL_29]] : tensor<8x10xf32>, tensor, tensor<8x10xf32>, tensor<8x10xf32>, tensor // CHECK: } @@ -438,32 +438,32 @@ func.func @inference_standard_indy_lstm_time_major_go_backwards(%arg0: tensor<8x // CHECK: [[VAL_7:%.*]] = "tf.Transpose"([[VAL_3]], [[VAL_6]]) : (tensor<8x40xf32>, tensor<2xi32>) -> tensor<40x8xf32> // CHECK: [[VAL_8:%.*]] = arith.constant dense<[1, 0]> : tensor<2xi32> // CHECK: [[VAL_9:%.*]] = "tf.Transpose"([[VAL_4]], [[VAL_8]]) : (tensor<10x4xf32>, tensor<2xi32>) -> tensor<4x10xf32> -// CHECK-DAG: [[VAL_10:%.*]] = "tf.Const"() {value = dense<10> : tensor<4xi32>} : () -> tensor<4xi32> -// CHECK-DAG: [[VAL_11:%.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor +// CHECK-DAG: [[VAL_10:%.*]] = "tf.Const"() <{value = dense<10> : tensor<4xi32>}> : () -> tensor<4xi32> +// CHECK-DAG: [[VAL_11:%.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK: [[VAL_12:%.*]]:4 = "tf.SplitV"([[VAL_7]], [[VAL_10]], [[VAL_11]]) : (tensor<40x8xf32>, tensor<4xi32>, tensor) -> (tensor<10x8xf32>, tensor<10x8xf32>, tensor<10x8xf32>, tensor<10x8xf32>) -// CHECK-DAG: [[VAL_13:%.*]] = "tf.Const"() {value = dense<1> : tensor<4xi32>} : () -> tensor<4xi32> -// CHECK-DAG: [[VAL_14:%.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor +// CHECK-DAG: [[VAL_13:%.*]] = "tf.Const"() <{value = dense<1> : tensor<4xi32>}> : () -> tensor<4xi32> +// CHECK-DAG: [[VAL_14:%.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK: [[VAL_15:%.*]]:4 = "tf.SplitV"([[VAL_9]], [[VAL_13]], [[VAL_14]]) : (tensor<4x10xf32>, tensor<4xi32>, tensor) -> (tensor<1x10xf32>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<1x10xf32>) -// CHECK-DAG: [[VAL_20:%.*]] = "tf.Const"() {value = dense<-1> : tensor<1xi32>} : () -> tensor<1xi32> +// CHECK-DAG: [[VAL_20:%.*]] = "tf.Const"() <{value = dense<-1> : tensor<1xi32>}> : () -> tensor<1xi32> // CHECK: [[VAL_21:%.*]] = "tf.Reshape"([[VAL_15]]#0, [[VAL_20]]) : (tensor<1x10xf32>, tensor<1xi32>) -> tensor<10xf32> -// CHECK-DAG: [[VAL_22:%.*]] = "tf.Const"() {value = dense<-1> : tensor<1xi32>} : () -> tensor<1xi32> +// CHECK-DAG: [[VAL_22:%.*]] = "tf.Const"() <{value = dense<-1> : tensor<1xi32>}> : () -> tensor<1xi32> // CHECK: [[VAL_23:%.*]] = "tf.Reshape"([[VAL_15]]#1, [[VAL_22]]) : (tensor<1x10xf32>, tensor<1xi32>) -> tensor<10xf32> -// CHECK-DAG: [[VAL_24:%.*]] = "tf.Const"() {value = dense<-1> : tensor<1xi32>} : () -> tensor<1xi32> +// CHECK-DAG: [[VAL_24:%.*]] = "tf.Const"() <{value = dense<-1> : tensor<1xi32>}> : () -> tensor<1xi32> // CHECK: [[VAL_25:%.*]] = "tf.Reshape"([[VAL_15]]#2, [[VAL_24]]) : (tensor<1x10xf32>, tensor<1xi32>) -> tensor<10xf32> -// CHECK-DAG: [[VAL_26:%.*]] = "tf.Const"() {value = dense<-1> : tensor<1xi32>} : () -> tensor<1xi32> +// CHECK-DAG: [[VAL_26:%.*]] = "tf.Const"() <{value = dense<-1> : tensor<1xi32>}> : () -> tensor<1xi32> // CHECK: [[VAL_27:%.*]] = "tf.Reshape"([[VAL_15]]#3, [[VAL_26]]) : (tensor<1x10xf32>, tensor<1xi32>) -> tensor<10xf32> -// CHECK-DAG: [[VAL_28:%.*]] = "tf.Const"() {value = dense<10> : tensor<4xi32>} : () -> tensor<4xi32> -// CHECK-DAG: [[VAL_29:%.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor +// CHECK-DAG: [[VAL_28:%.*]] = "tf.Const"() <{value = dense<10> : tensor<4xi32>}> : () -> tensor<4xi32> +// CHECK-DAG: [[VAL_29:%.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK: [[VAL_30:%.*]]:4 = "tf.SplitV"([[VAL_5]], [[VAL_28]], [[VAL_29]]) : (tensor<40xf32>, tensor<4xi32>, tensor) -> (tensor<10xf32>, tensor<10xf32>, tensor<10xf32>, tensor<10xf32>) // CHECK: [[VAL_31:%.*]] = "tfl.no_value"() {value} : () -> none // CHECK: [[VAL_32:%.*]] = "tfl.unidirectional_sequence_lstm"([[VAL_41]], [[VAL_12]]#0, [[VAL_12]]#1, [[VAL_12]]#2, [[VAL_12]]#3, [[VAL_21]], [[VAL_23]], [[VAL_25]], [[VAL_27]], [[VAL_31]], [[VAL_31]], [[VAL_31]], [[VAL_30]]#0, [[VAL_30]]#1, [[VAL_30]]#2, [[VAL_30]]#3, [[VAL_31]], [[VAL_31]], [[VAL_1]], [[VAL_2]], [[VAL_31]], [[VAL_31]], [[VAL_31]], [[VAL_31]]) {cell_clip = 1.000000e+01 : f32, diagonal_recurrent_tensors = true, fused_activation_function = "TANH", proj_clip = 0.000000e+00 : f32, time_major = true} : (tensor<8x8x8xf32>, tensor<10x8xf32>, tensor<10x8xf32>, tensor<10x8xf32>, tensor<10x8xf32>, tensor<10xf32>, tensor<10xf32>, tensor<10xf32>, tensor<10xf32>, none, none, none, tensor<10xf32>, tensor<10xf32>, tensor<10xf32>, tensor<10xf32>, none, none, tensor<8x10xf32>, tensor<8x10xf32>, none, none, none, none) -> tensor<8x8x10xf32> // CHECK-DAG: [[VAL_33:%.*]] = arith.constant dense<[-1, 0, 0]> : tensor<3xi32> // CHECK-DAG: [[VAL_34:%.*]] = arith.constant dense<0> : tensor<3xi32> // CHECK-DAG: [[VAL_35:%.*]] = arith.constant dense<1> : tensor<3xi32> -// CHECK: [[VAL_36:%.*]] = "tf.StridedSlice"([[VAL_32]], [[VAL_33]], [[VAL_34]], [[VAL_35]]) {begin_mask = 6 : i64, ellipsis_mask = 0 : i64, end_mask = 6 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 1 : i64} : (tensor<8x8x10xf32>, tensor<3xi32>, tensor<3xi32>, tensor<3xi32>) -> tensor<8x10xf32> -// CHECK-DAG: [[VAL_37:%.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<1xf32>} : () -> tensor<8x10xf32> -// CHECK-DAG: [[VAL_38:%.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<1xf32>} : () -> tensor<8x10xf32> -// CHECK-DAG: [[VAL_39:%.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<1xf32>} : () -> tensor +// CHECK: [[VAL_36:%.*]] = "tf.StridedSlice"([[VAL_32]], [[VAL_33]], [[VAL_34]], [[VAL_35]]) <{begin_mask = 6 : i64, ellipsis_mask = 0 : i64, end_mask = 6 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 1 : i64}> : (tensor<8x8x10xf32>, tensor<3xi32>, tensor<3xi32>, tensor<3xi32>) -> tensor<8x10xf32> +// CHECK-DAG: [[VAL_37:%.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<1xf32>}> : () -> tensor<8x10xf32> +// CHECK-DAG: [[VAL_38:%.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<1xf32>}> : () -> tensor<8x10xf32> +// CHECK-DAG: [[VAL_39:%.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<1xf32>}> : () -> tensor // CHECK: return [[VAL_36]], [[VAL_32]], [[VAL_37]], [[VAL_38]], [[VAL_39]] : tensor<8x10xf32>, tensor<8x8x10xf32>, tensor<8x10xf32>, tensor<8x10xf32>, tensor // CHECK: } @@ -490,24 +490,24 @@ func.func @inference_standard_lstm_non_time_major_go_backwards(%arg0: tensor<8x8 // CHECK: [[VAL_9:%.*]] = "tf.Transpose"([[VAL_3]], [[VAL_8]]) : (tensor<8x40xf32>, tensor<2xi32>) -> tensor<40x8xf32> // CHECK: [[VAL_10:%.*]] = arith.constant dense<[1, 0]> : tensor<2xi32> // CHECK: [[VAL_11:%.*]] = "tf.Transpose"([[VAL_4]], [[VAL_10]]) : (tensor<10x40xf32>, tensor<2xi32>) -> tensor<40x10xf32> -// CHECK-DAG: [[VAL_12:%.*]] = "tf.Const"() {value = dense<10> : tensor<4xi32>} : () -> tensor<4xi32> -// CHECK-DAG: [[VAL_13:%.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor +// CHECK-DAG: [[VAL_12:%.*]] = "tf.Const"() <{value = dense<10> : tensor<4xi32>}> : () -> tensor<4xi32> +// CHECK-DAG: [[VAL_13:%.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK: [[VAL_14:%.*]]:4 = "tf.SplitV"([[VAL_9]], [[VAL_12]], [[VAL_13]]) : (tensor<40x8xf32>, tensor<4xi32>, tensor) -> (tensor<10x8xf32>, tensor<10x8xf32>, tensor<10x8xf32>, tensor<10x8xf32>) -// CHECK-DAG: [[VAL_15:%.*]] = "tf.Const"() {value = dense<10> : tensor<4xi32>} : () -> tensor<4xi32> -// CHECK-DAG: [[VAL_16:%.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor +// CHECK-DAG: [[VAL_15:%.*]] = "tf.Const"() <{value = dense<10> : tensor<4xi32>}> : () -> tensor<4xi32> +// CHECK-DAG: [[VAL_16:%.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK: [[VAL_17:%.*]]:4 = "tf.SplitV"([[VAL_11]], [[VAL_15]], [[VAL_16]]) : (tensor<40x10xf32>, tensor<4xi32>, tensor) -> (tensor<10x10xf32>, tensor<10x10xf32>, tensor<10x10xf32>, tensor<10x10xf32>) -// CHECK-DAG: [[VAL_18:%.*]] = "tf.Const"() {value = dense<10> : tensor<4xi32>} : () -> tensor<4xi32> -// CHECK-DAG: [[VAL_19:%.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor +// CHECK-DAG: [[VAL_18:%.*]] = "tf.Const"() <{value = dense<10> : tensor<4xi32>}> : () -> tensor<4xi32> +// CHECK-DAG: [[VAL_19:%.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK: [[VAL_20:%.*]]:4 = "tf.SplitV"([[VAL_5]], [[VAL_18]], [[VAL_19]]) : (tensor<40xf32>, tensor<4xi32>, tensor) -> (tensor<10xf32>, tensor<10xf32>, tensor<10xf32>, tensor<10xf32>) // CHECK: [[VAL_21:%.*]] = "tfl.no_value"() {value} : () -> none // CHECK: [[VAL_22:%.*]] = "tfl.unidirectional_sequence_lstm"([[VAL_7]], [[VAL_14]]#0, [[VAL_14]]#1, [[VAL_14]]#2, [[VAL_14]]#3, [[VAL_17]]#0, [[VAL_17]]#1, [[VAL_17]]#2, [[VAL_17]]#3, [[VAL_21]], [[VAL_21]], [[VAL_21]], [[VAL_20]]#0, [[VAL_20]]#1, [[VAL_20]]#2, [[VAL_20]]#3, [[VAL_21]], [[VAL_21]], [[VAL_1]], [[VAL_2]], [[VAL_21]], [[VAL_21]], [[VAL_21]], [[VAL_21]]) {cell_clip = 1.000000e+01 : f32, diagonal_recurrent_tensors = false, fused_activation_function = "TANH", proj_clip = 0.000000e+00 : f32, time_major = false} : (tensor<8x8x8xf32>, tensor<10x8xf32>, tensor<10x8xf32>, tensor<10x8xf32>, tensor<10x8xf32>, tensor<10x10xf32>, tensor<10x10xf32>, tensor<10x10xf32>, tensor<10x10xf32>, none, none, none, tensor<10xf32>, tensor<10xf32>, tensor<10xf32>, tensor<10xf32>, none, none, tensor<8x10xf32>, tensor<8x10xf32>, none, none, none, none) -> tensor<8x8x10xf32> // CHECK-DAG: [[VAL_23:%.*]] = arith.constant dense<[0, -1, 0]> : tensor<3xi32> // CHECK-DAG: [[VAL_24:%.*]] = arith.constant dense<0> : tensor<3xi32> // CHECK-DAG: [[VAL_25:%.*]] = arith.constant dense<1> : tensor<3xi32> -// CHECK: [[VAL_26:%.*]] = "tf.StridedSlice"([[VAL_22]], [[VAL_23]], [[VAL_24]], [[VAL_25]]) {begin_mask = 5 : i64, ellipsis_mask = 0 : i64, end_mask = 5 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 2 : i64} : (tensor<8x8x10xf32>, tensor<3xi32>, tensor<3xi32>, tensor<3xi32>) -> tensor<8x10xf32> -// CHECK-DAG: [[VAL_27:%.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<1xf32>} : () -> tensor<8x10xf32> -// CHECK-DAG: [[VAL_28:%.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<1xf32>} : () -> tensor<8x10xf32> -// CHECK-DAG: [[VAL_29:%.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<1xf32>} : () -> tensor +// CHECK: [[VAL_26:%.*]] = "tf.StridedSlice"([[VAL_22]], [[VAL_23]], [[VAL_24]], [[VAL_25]]) <{begin_mask = 5 : i64, ellipsis_mask = 0 : i64, end_mask = 5 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 2 : i64}> : (tensor<8x8x10xf32>, tensor<3xi32>, tensor<3xi32>, tensor<3xi32>) -> tensor<8x10xf32> +// CHECK-DAG: [[VAL_27:%.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<1xf32>}> : () -> tensor<8x10xf32> +// CHECK-DAG: [[VAL_28:%.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<1xf32>}> : () -> tensor<8x10xf32> +// CHECK-DAG: [[VAL_29:%.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<1xf32>}> : () -> tensor // CHECK: return [[VAL_26]], [[VAL_22]], [[VAL_27]], [[VAL_28]], [[VAL_29]] : tensor<8x10xf32>, tensor<8x8x10xf32>, tensor<8x10xf32>, tensor<8x10xf32>, tensor // CHECK: } @@ -539,32 +539,32 @@ func.func @inference_standard_indy_lstm_non_time_major_go_backwards(%arg0: tenso // CHECK: [[VAL_7:%.*]] = "tf.Transpose"([[VAL_3]], [[VAL_6]]) : (tensor<8x40xf32>, tensor<2xi32>) -> tensor<40x8xf32> // CHECK: [[VAL_8:%.*]] = arith.constant dense<[1, 0]> : tensor<2xi32> // CHECK: [[VAL_9:%.*]] = "tf.Transpose"([[VAL_4]], [[VAL_8]]) : (tensor<10x4xf32>, tensor<2xi32>) -> tensor<4x10xf32> -// CHECK-DAG: [[VAL_10:%.*]] = "tf.Const"() {value = dense<10> : tensor<4xi32>} : () -> tensor<4xi32> -// CHECK-DAG: [[VAL_11:%.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor +// CHECK-DAG: [[VAL_10:%.*]] = "tf.Const"() <{value = dense<10> : tensor<4xi32>}> : () -> tensor<4xi32> +// CHECK-DAG: [[VAL_11:%.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK: [[VAL_12:%.*]]:4 = "tf.SplitV"([[VAL_7]], [[VAL_10]], [[VAL_11]]) : (tensor<40x8xf32>, tensor<4xi32>, tensor) -> (tensor<10x8xf32>, tensor<10x8xf32>, tensor<10x8xf32>, tensor<10x8xf32>) -// CHECK-DAG: [[VAL_13:%.*]] = "tf.Const"() {value = dense<1> : tensor<4xi32>} : () -> tensor<4xi32> -// CHECK-DAG: [[VAL_14:%.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor +// CHECK-DAG: [[VAL_13:%.*]] = "tf.Const"() <{value = dense<1> : tensor<4xi32>}> : () -> tensor<4xi32> +// CHECK-DAG: [[VAL_14:%.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK: [[VAL_15:%.*]]:4 = "tf.SplitV"([[VAL_9]], [[VAL_13]], [[VAL_14]]) : (tensor<4x10xf32>, tensor<4xi32>, tensor) -> (tensor<1x10xf32>, tensor<1x10xf32>, tensor<1x10xf32>, tensor<1x10xf32>) -// CHECK-DAG: [[VAL_20:%.*]] = "tf.Const"() {value = dense<-1> : tensor<1xi32>} : () -> tensor<1xi32> +// CHECK-DAG: [[VAL_20:%.*]] = "tf.Const"() <{value = dense<-1> : tensor<1xi32>}> : () -> tensor<1xi32> // CHECK: [[VAL_21:%.*]] = "tf.Reshape"([[VAL_15]]#0, [[VAL_20]]) : (tensor<1x10xf32>, tensor<1xi32>) -> tensor<10xf32> -// CHECK-DAG: [[VAL_22:%.*]] = "tf.Const"() {value = dense<-1> : tensor<1xi32>} : () -> tensor<1xi32> +// CHECK-DAG: [[VAL_22:%.*]] = "tf.Const"() <{value = dense<-1> : tensor<1xi32>}> : () -> tensor<1xi32> // CHECK: [[VAL_23:%.*]] = "tf.Reshape"([[VAL_15]]#1, [[VAL_22]]) : (tensor<1x10xf32>, tensor<1xi32>) -> tensor<10xf32> -// CHECK-DAG: [[VAL_24:%.*]] = "tf.Const"() {value = dense<-1> : tensor<1xi32>} : () -> tensor<1xi32> +// CHECK-DAG: [[VAL_24:%.*]] = "tf.Const"() <{value = dense<-1> : tensor<1xi32>}> : () -> tensor<1xi32> // CHECK: [[VAL_25:%.*]] = "tf.Reshape"([[VAL_15]]#2, [[VAL_24]]) : (tensor<1x10xf32>, tensor<1xi32>) -> tensor<10xf32> -// CHECK-DAG: [[VAL_26:%.*]] = "tf.Const"() {value = dense<-1> : tensor<1xi32>} : () -> tensor<1xi32> +// CHECK-DAG: [[VAL_26:%.*]] = "tf.Const"() <{value = dense<-1> : tensor<1xi32>}> : () -> tensor<1xi32> // CHECK: [[VAL_27:%.*]] = "tf.Reshape"([[VAL_15]]#3, [[VAL_26]]) : (tensor<1x10xf32>, tensor<1xi32>) -> tensor<10xf32> -// CHECK-DAG: [[VAL_28:%.*]] = "tf.Const"() {value = dense<10> : tensor<4xi32>} : () -> tensor<4xi32> -// CHECK-DAG: [[VAL_29:%.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor +// CHECK-DAG: [[VAL_28:%.*]] = "tf.Const"() <{value = dense<10> : tensor<4xi32>}> : () -> tensor<4xi32> +// CHECK-DAG: [[VAL_29:%.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK: [[VAL_30:%.*]]:4 = "tf.SplitV"([[VAL_5]], [[VAL_28]], [[VAL_29]]) : (tensor<40xf32>, tensor<4xi32>, tensor) -> (tensor<10xf32>, tensor<10xf32>, tensor<10xf32>, tensor<10xf32>) // CHECK: [[VAL_31:%.*]] = "tfl.no_value"() {value} : () -> none // CHECK: [[VAL_32:%.*]] = "tfl.unidirectional_sequence_lstm"([[VAL_41]], [[VAL_12]]#0, [[VAL_12]]#1, [[VAL_12]]#2, [[VAL_12]]#3, [[VAL_21]], [[VAL_23]], [[VAL_25]], [[VAL_27]], [[VAL_31]], [[VAL_31]], [[VAL_31]], [[VAL_30]]#0, [[VAL_30]]#1, [[VAL_30]]#2, [[VAL_30]]#3, [[VAL_31]], [[VAL_31]], [[VAL_1]], [[VAL_2]], [[VAL_31]], [[VAL_31]], [[VAL_31]], [[VAL_31]]) {cell_clip = 1.000000e+01 : f32, diagonal_recurrent_tensors = true, fused_activation_function = "TANH", proj_clip = 0.000000e+00 : f32, time_major = false} : (tensor<8x8x8xf32>, tensor<10x8xf32>, tensor<10x8xf32>, tensor<10x8xf32>, tensor<10x8xf32>, tensor<10xf32>, tensor<10xf32>, tensor<10xf32>, tensor<10xf32>, none, none, none, tensor<10xf32>, tensor<10xf32>, tensor<10xf32>, tensor<10xf32>, none, none, tensor<8x10xf32>, tensor<8x10xf32>, none, none, none, none) -> tensor<8x8x10xf32> // CHECK-DAG: [[VAL_33:%.*]] = arith.constant dense<[0, -1, 0]> : tensor<3xi32> // CHECK-DAG: [[VAL_34:%.*]] = arith.constant dense<0> : tensor<3xi32> // CHECK-DAG: [[VAL_35:%.*]] = arith.constant dense<1> : tensor<3xi32> -// CHECK: [[VAL_36:%.*]] = "tf.StridedSlice"([[VAL_32]], [[VAL_33]], [[VAL_34]], [[VAL_35]]) {begin_mask = 5 : i64, ellipsis_mask = 0 : i64, end_mask = 5 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 2 : i64} : (tensor<8x8x10xf32>, tensor<3xi32>, tensor<3xi32>, tensor<3xi32>) -> tensor<8x10xf32> -// CHECK-DAG: [[VAL_37:%.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<1xf32>} : () -> tensor<8x10xf32> -// CHECK-DAG: [[VAL_38:%.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<1xf32>} : () -> tensor<8x10xf32> -// CHECK-DAG: [[VAL_39:%.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<1xf32>} : () -> tensor +// CHECK: [[VAL_36:%.*]] = "tf.StridedSlice"([[VAL_32]], [[VAL_33]], [[VAL_34]], [[VAL_35]]) <{begin_mask = 5 : i64, ellipsis_mask = 0 : i64, end_mask = 5 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 2 : i64}> : (tensor<8x8x10xf32>, tensor<3xi32>, tensor<3xi32>, tensor<3xi32>) -> tensor<8x10xf32> +// CHECK-DAG: [[VAL_37:%.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<1xf32>}> : () -> tensor<8x10xf32> +// CHECK-DAG: [[VAL_38:%.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<1xf32>}> : () -> tensor<8x10xf32> +// CHECK-DAG: [[VAL_39:%.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<1xf32>}> : () -> tensor // CHECK: return [[VAL_36]], [[VAL_32]], [[VAL_37]], [[VAL_38]], [[VAL_39]] : tensor<8x10xf32>, tensor<8x8x10xf32>, tensor<8x10xf32>, tensor<8x10xf32>, tensor // CHECK: } @@ -596,24 +596,24 @@ func.func @inference_standard_lstm_time_major_can_fuse(%arg0: tensor, // CHECK: [[VAL_7:%.*]] = "tf.Transpose"([[VAL_3]], [[VAL_6]]) : (tensor<8x40xf32>, tensor<2xi32>) -> tensor<40x8xf32> // CHECK: [[VAL_8:%.*]] = arith.constant dense<[1, 0]> : tensor<2xi32> // CHECK: [[VAL_9:%.*]] = "tf.Transpose"([[VAL_4]], [[VAL_8]]) : (tensor<10x40xf32>, tensor<2xi32>) -> tensor<40x10xf32> -// CHECK-DAG: [[VAL_10:%.*]] = "tf.Const"() {value = dense<10> : tensor<4xi32>} : () -> tensor<4xi32> -// CHECK-DAG: [[VAL_11:%.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor +// CHECK-DAG: [[VAL_10:%.*]] = "tf.Const"() <{value = dense<10> : tensor<4xi32>}> : () -> tensor<4xi32> +// CHECK-DAG: [[VAL_11:%.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK: [[VAL_12:%.*]]:4 = "tf.SplitV"([[VAL_7]], [[VAL_10]], [[VAL_11]]) : (tensor<40x8xf32>, tensor<4xi32>, tensor) -> (tensor<10x8xf32>, tensor<10x8xf32>, tensor<10x8xf32>, tensor<10x8xf32>) -// CHECK-DAG: [[VAL_13:%.*]] = "tf.Const"() {value = dense<10> : tensor<4xi32>} : () -> tensor<4xi32> -// CHECK-DAG: [[VAL_14:%.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor +// CHECK-DAG: [[VAL_13:%.*]] = "tf.Const"() <{value = dense<10> : tensor<4xi32>}> : () -> tensor<4xi32> +// CHECK-DAG: [[VAL_14:%.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK: [[VAL_15:%.*]]:4 = "tf.SplitV"([[VAL_9]], [[VAL_13]], [[VAL_14]]) : (tensor<40x10xf32>, tensor<4xi32>, tensor) -> (tensor<10x10xf32>, tensor<10x10xf32>, tensor<10x10xf32>, tensor<10x10xf32>) -// CHECK-DAG: [[VAL_16:%.*]] = "tf.Const"() {value = dense<10> : tensor<4xi32>} : () -> tensor<4xi32> -// CHECK-DAG: [[VAL_17:%.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor +// CHECK-DAG: [[VAL_16:%.*]] = "tf.Const"() <{value = dense<10> : tensor<4xi32>}> : () -> tensor<4xi32> +// CHECK-DAG: [[VAL_17:%.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK: [[VAL_18:%.*]]:4 = "tf.SplitV"([[VAL_5]], [[VAL_16]], [[VAL_17]]) : (tensor<40xf32>, tensor<4xi32>, tensor) -> (tensor<10xf32>, tensor<10xf32>, tensor<10xf32>, tensor<10xf32>) // CHECK: [[VAL_19:%.*]] = "tfl.no_value"() {value} : () -> none // CHECK: [[VAL_20:%.*]] = "tfl.unidirectional_sequence_lstm"([[VAL_0]], [[VAL_12]]#0, [[VAL_12]]#1, [[VAL_12]]#2, [[VAL_12]]#3, [[VAL_15]]#0, [[VAL_15]]#1, [[VAL_15]]#2, [[VAL_15]]#3, [[VAL_19]], [[VAL_19]], [[VAL_19]], [[VAL_18]]#0, [[VAL_18]]#1, [[VAL_18]]#2, [[VAL_18]]#3, [[VAL_19]], [[VAL_19]], [[VAL_1]], [[VAL_2]], [[VAL_19]], [[VAL_19]], [[VAL_19]], [[VAL_19]]) {cell_clip = 1.000000e+01 : f32, diagonal_recurrent_tensors = false, fused_activation_function = "TANH", proj_clip = 0.000000e+00 : f32, time_major = true} : (tensor, tensor<10x8xf32>, tensor<10x8xf32>, tensor<10x8xf32>, tensor<10x8xf32>, tensor<10x10xf32>, tensor<10x10xf32>, tensor<10x10xf32>, tensor<10x10xf32>, none, none, none, tensor<10xf32>, tensor<10xf32>, tensor<10xf32>, tensor<10xf32>, none, none, tensor<8x10xf32>, tensor<8x10xf32>, none, none, none, none) -> tensor // CHECK-DAG: [[VAL_21:%.*]] = arith.constant dense<[-1, 0, 0]> : tensor<3xi32> // CHECK-DAG: [[VAL_22:%.*]] = arith.constant dense<0> : tensor<3xi32> // CHECK-DAG: [[VAL_23:%.*]] = arith.constant dense<1> : tensor<3xi32> -// CHECK: [[VAL_24:%.*]] = "tf.StridedSlice"([[VAL_20]], [[VAL_21]], [[VAL_22]], [[VAL_23]]) {begin_mask = 6 : i64, ellipsis_mask = 0 : i64, end_mask = 6 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 1 : i64} : (tensor, tensor<3xi32>, tensor<3xi32>, tensor<3xi32>) -> tensor<8x10xf32> -// CHECK-DAG: [[VAL_25:%.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<1xf32>} : () -> tensor<8x10xf32> -// CHECK-DAG: [[VAL_26:%.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<1xf32>} : () -> tensor<8x10xf32> -// CHECK-DAG: [[VAL_27:%.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<1xf32>} : () -> tensor +// CHECK: [[VAL_24:%.*]] = "tf.StridedSlice"([[VAL_20]], [[VAL_21]], [[VAL_22]], [[VAL_23]]) <{begin_mask = 6 : i64, ellipsis_mask = 0 : i64, end_mask = 6 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 1 : i64}> : (tensor, tensor<3xi32>, tensor<3xi32>, tensor<3xi32>) -> tensor<8x10xf32> +// CHECK-DAG: [[VAL_25:%.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<1xf32>}> : () -> tensor<8x10xf32> +// CHECK-DAG: [[VAL_26:%.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<1xf32>}> : () -> tensor<8x10xf32> +// CHECK-DAG: [[VAL_27:%.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<1xf32>}> : () -> tensor // CHECK: return [[VAL_24]], [[VAL_20]], [[VAL_25]], [[VAL_26]], [[VAL_27]] : tensor<8x10xf32>, tensor, tensor<8x10xf32>, tensor<8x10xf32>, tensor // CHECK: } @@ -646,24 +646,24 @@ func.func @inference_standard_lstm_time_major_can_fuse_last_output(%arg0: tensor // CHECK: [[VAL_7:%.*]] = "tf.Transpose"([[VAL_3]], [[VAL_6]]) : (tensor<8x40xf32>, tensor<2xi32>) -> tensor<40x8xf32> // CHECK: [[VAL_8:%.*]] = arith.constant dense<[1, 0]> : tensor<2xi32> // CHECK: [[VAL_9:%.*]] = "tf.Transpose"([[VAL_4]], [[VAL_8]]) : (tensor<10x40xf32>, tensor<2xi32>) -> tensor<40x10xf32> -// CHECK-DAG: [[VAL_10:%.*]] = "tf.Const"() {value = dense<10> : tensor<4xi32>} : () -> tensor<4xi32> -// CHECK-DAG: [[VAL_11:%.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor +// CHECK-DAG: [[VAL_10:%.*]] = "tf.Const"() <{value = dense<10> : tensor<4xi32>}> : () -> tensor<4xi32> +// CHECK-DAG: [[VAL_11:%.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK: [[VAL_12:%.*]]:4 = "tf.SplitV"([[VAL_7]], [[VAL_10]], [[VAL_11]]) : (tensor<40x8xf32>, tensor<4xi32>, tensor) -> (tensor<10x8xf32>, tensor<10x8xf32>, tensor<10x8xf32>, tensor<10x8xf32>) -// CHECK-DAG: [[VAL_13:%.*]] = "tf.Const"() {value = dense<10> : tensor<4xi32>} : () -> tensor<4xi32> -// CHECK-DAG: [[VAL_14:%.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor +// CHECK-DAG: [[VAL_13:%.*]] = "tf.Const"() <{value = dense<10> : tensor<4xi32>}> : () -> tensor<4xi32> +// CHECK-DAG: [[VAL_14:%.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK: [[VAL_15:%.*]]:4 = "tf.SplitV"([[VAL_9]], [[VAL_13]], [[VAL_14]]) : (tensor<40x10xf32>, tensor<4xi32>, tensor) -> (tensor<10x10xf32>, tensor<10x10xf32>, tensor<10x10xf32>, tensor<10x10xf32>) -// CHECK-DAG: [[VAL_16:%.*]] = "tf.Const"() {value = dense<10> : tensor<4xi32>} : () -> tensor<4xi32> -// CHECK-DAG: [[VAL_17:%.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor +// CHECK-DAG: [[VAL_16:%.*]] = "tf.Const"() <{value = dense<10> : tensor<4xi32>}> : () -> tensor<4xi32> +// CHECK-DAG: [[VAL_17:%.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK: [[VAL_18:%.*]]:4 = "tf.SplitV"([[VAL_5]], [[VAL_16]], [[VAL_17]]) : (tensor<40xf32>, tensor<4xi32>, tensor) -> (tensor<10xf32>, tensor<10xf32>, tensor<10xf32>, tensor<10xf32>) // CHECK: [[VAL_19:%.*]] = "tfl.no_value"() {value} : () -> none // CHECK: [[VAL_20:%.*]] = "tfl.unidirectional_sequence_lstm"([[VAL_0]], [[VAL_12]]#0, [[VAL_12]]#1, [[VAL_12]]#2, [[VAL_12]]#3, [[VAL_15]]#0, [[VAL_15]]#1, [[VAL_15]]#2, [[VAL_15]]#3, [[VAL_19]], [[VAL_19]], [[VAL_19]], [[VAL_18]]#0, [[VAL_18]]#1, [[VAL_18]]#2, [[VAL_18]]#3, [[VAL_19]], [[VAL_19]], [[VAL_1]], [[VAL_2]], [[VAL_19]], [[VAL_19]], [[VAL_19]], [[VAL_19]]) {cell_clip = 1.000000e+01 : f32, diagonal_recurrent_tensors = false, fused_activation_function = "TANH", proj_clip = 0.000000e+00 : f32, time_major = true} : (tensor, tensor<10x8xf32>, tensor<10x8xf32>, tensor<10x8xf32>, tensor<10x8xf32>, tensor<10x10xf32>, tensor<10x10xf32>, tensor<10x10xf32>, tensor<10x10xf32>, none, none, none, tensor<10xf32>, tensor<10xf32>, tensor<10xf32>, tensor<10xf32>, none, none, tensor<8x10xf32>, tensor<8x10xf32>, none, none, none, none) -> tensor // CHECK-DAG: [[VAL_21:%.*]] = arith.constant dense<[-1, 0, 0]> : tensor<3xi32> // CHECK-DAG: [[VAL_22:%.*]] = arith.constant dense<0> : tensor<3xi32> // CHECK-DAG: [[VAL_23:%.*]] = arith.constant dense<1> : tensor<3xi32> -// CHECK: [[VAL_24:%.*]] = "tf.StridedSlice"([[VAL_20]], [[VAL_21]], [[VAL_22]], [[VAL_23]]) {begin_mask = 6 : i64, ellipsis_mask = 0 : i64, end_mask = 6 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 1 : i64} : (tensor, tensor<3xi32>, tensor<3xi32>, tensor<3xi32>) -> tensor<8x10xf32> -// CHECK-DAG: [[VAL_25:%.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<1xf32>} : () -> tensor<8x10xf32> -// CHECK-DAG: [[VAL_26:%.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<1xf32>} : () -> tensor<8x10xf32> -// CHECK-DAG: [[VAL_27:%.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<1xf32>} : () -> tensor +// CHECK: [[VAL_24:%.*]] = "tf.StridedSlice"([[VAL_20]], [[VAL_21]], [[VAL_22]], [[VAL_23]]) <{begin_mask = 6 : i64, ellipsis_mask = 0 : i64, end_mask = 6 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 1 : i64}> : (tensor, tensor<3xi32>, tensor<3xi32>, tensor<3xi32>) -> tensor<8x10xf32> +// CHECK-DAG: [[VAL_25:%.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<1xf32>}> : () -> tensor<8x10xf32> +// CHECK-DAG: [[VAL_26:%.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<1xf32>}> : () -> tensor<8x10xf32> +// CHECK-DAG: [[VAL_27:%.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<1xf32>}> : () -> tensor // CHECK: return [[VAL_24]], [[VAL_20]], [[VAL_25]], [[VAL_26]], [[VAL_27]] : tensor<8x10xf32>, tensor, tensor<8x10xf32>, tensor<8x10xf32>, tensor // CHECK: } @@ -684,13 +684,13 @@ func.func @inference_standard_lstm_with_mask(%arg0: tensor, %arg1: te } // CHECK: func @inference_standard_lstm_with_mask([[ARG_0:%.*]]: tensor, [[ARG_1:%.*]]: tensor<8x10xf32>, [[ARG_2:%.*]]: tensor<8x10xf32>, [[ARG_3:%.*]]: tensor<8x40xf32>, [[ARG_4:%.*]]: tensor<10x40xf32>, [[ARG_5:%.*]]: tensor<40xf32>, [[ARG_6:%.*]]: tensor) -> (tensor<8x10xf32>, tensor, tensor<8x10xf32>, tensor<8x10xf32>, tensor) attributes {tf._input_shapes = ["tfshape$dim { size: -1 } dim { size: 8 } dim { size: 8 }", "tfshape$dim { size: 8 } dim { size: 10 }", "tfshape$dim { size: 8 } dim { size: 10 }", "tfshape$unknown_rank: true", "tfshape$unknown_rank: false", "tfshape$unknown_rank: false", "tfshape$dim { size: -1 } dim { size: 8 }"], tf.api_implements = "lstm_b4e9f0e7-ac55-42bc-8ef2-8496419a608c", tf.api_preferred_device = "CPU", tf.go_backwards = false, tf.time_major = true} { -// CHECK: [[VAL_0:%.*]] = "tf.BatchMatMulV2"([[ARG_0]], [[ARG_3]]) {adj_x = false, adj_y = false} : (tensor, tensor<8x40xf32>) -> tensor +// CHECK: [[VAL_0:%.*]] = "tf.BatchMatMulV2"([[ARG_0]], [[ARG_3]]) <{adj_x = false, adj_y = false}> : (tensor, tensor<8x40xf32>) -> tensor // CHECK: [[VAL_1:%.*]] = "tf.Add"([[VAL_0]], [[ARG_5]]) : (tensor, tensor<40xf32>) -> tensor -// CHECK: [[VAL_2:%.*]] = "tf.BatchMatMulV2"([[VAL_1]], [[ARG_4]]) {adj_x = false, adj_y = true} : (tensor, tensor<10x40xf32>) -> tensor +// CHECK: [[VAL_2:%.*]] = "tf.BatchMatMulV2"([[VAL_1]], [[ARG_4]]) <{adj_x = false, adj_y = true}> : (tensor, tensor<10x40xf32>) -> tensor // CHECK: [[VAL_3:%.*]] = "tf.Add"([[VAL_2]], [[ARG_1]]) : (tensor, tensor<8x10xf32>) -> tensor // CHECK: [[VAL_4:%.*]] = "tf.Add"([[VAL_2]], [[ARG_2]]) : (tensor, tensor<8x10xf32>) -> tensor // CHECK: [[VAL_5:%.*]] = "tf.Add"([[ARG_1]], [[ARG_2]]) : (tensor<8x10xf32>, tensor<8x10xf32>) -> tensor<8x10xf32> -// CHECK: [[VAL_6:%.*]] = "tf.Const"() {_output_shapes = ["tfshape$"], device = "/device:CPU:0", dtype = f32, value = dense<1.000000e+00> : tensor} : () -> tensor +// CHECK: [[VAL_6:%.*]] = "tf.Const"() <{value = dense<1.000000e+00> : tensor}> {_output_shapes = ["tfshape$"], device = "/device:CPU:0", dtype = f32} : () -> tensor // CHECK: return [[VAL_5]], [[VAL_4]], [[VAL_5]], [[VAL_5]], [[VAL_6]] : tensor<8x10xf32>, tensor, tensor<8x10xf32>, tensor<8x10xf32>, tensor // CHECK: } @@ -718,13 +718,13 @@ func.func @inference_standard_lstm_time_major_cannot_fuse(%arg0: tensor, [[VAL_1:%.*]]: tensor, [[VAL_2:%.*]]: tensor, [[VAL_3:%.*]]: tensor<8x40xf32>, [[VAL_4:%.*]]: tensor<10x40xf32>, [[VAL_5:%.*]]: tensor<40xf32>) -> (tensor, tensor, tensor, tensor, tensor) attributes {tf._input_shapes = ["tfshape$dim { size: -1 } dim { size: 8 } dim { size: 8 }", "tfshape$dim { size: -1 } dim { size: 10 }", "tfshape$dim { size: -1 } dim { size: 10 }", "tfshape$unknown_rank: true", "tfshape$unknown_rank: true", "tfshape$unknown_rank: true"], tf.api_implements = "lstm_b4e9f0e7-ac55-42bc-8ef2-8496419a608c", tf.api_preferred_device = "CPU", tf.go_backwards = false, tf.time_major = true} { -// CHECK: [[VAL_6:%.*]] = "tf.BatchMatMulV2"([[VAL_0]], [[VAL_3]]) {adj_x = false, adj_y = false} : (tensor, tensor<8x40xf32>) -> tensor +// CHECK: [[VAL_6:%.*]] = "tf.BatchMatMulV2"([[VAL_0]], [[VAL_3]]) <{adj_x = false, adj_y = false}> : (tensor, tensor<8x40xf32>) -> tensor // CHECK: [[VAL_7:%.*]] = "tf.Add"([[VAL_6]], [[VAL_5]]) : (tensor, tensor<40xf32>) -> tensor -// CHECK: [[VAL_8:%.*]] = "tf.BatchMatMulV2"([[VAL_7]], [[VAL_4]]) {adj_x = false, adj_y = true} : (tensor, tensor<10x40xf32>) -> tensor +// CHECK: [[VAL_8:%.*]] = "tf.BatchMatMulV2"([[VAL_7]], [[VAL_4]]) <{adj_x = false, adj_y = true}> : (tensor, tensor<10x40xf32>) -> tensor // CHECK: [[VAL_9:%.*]] = "tf.Add"([[VAL_8]], [[VAL_1]]) : (tensor, tensor) -> tensor // CHECK: [[VAL_10:%.*]] = "tf.Add"([[VAL_8]], [[VAL_2]]) : (tensor, tensor) -> tensor // CHECK: [[VAL_11:%.*]] = "tf.Add"([[VAL_1]], [[VAL_2]]) : (tensor, tensor) -> tensor -// CHECK: [[VAL_12:%.*]] = "tf.Const"() {_output_shapes = ["tfshape$"], device = "/device:CPU:0", dtype = f32, value = dense<1.000000e+00> : tensor} : () -> tensor +// CHECK: [[VAL_12:%.*]] = "tf.Const"() <{value = dense<1.000000e+00> : tensor}> {_output_shapes = ["tfshape$"], device = "/device:CPU:0", dtype = f32} : () -> tensor // CHECK: return [[VAL_11]], [[VAL_10]], [[VAL_11]], [[VAL_11]], [[VAL_12]] : tensor, tensor, tensor, tensor, tensor // CHECK: } } @@ -745,13 +745,13 @@ func.func @dynamic_shape_non_fuse_standard_lstm(%arg0: tensor, %arg1: } // CHECK: func @dynamic_shape_non_fuse_standard_lstm(%[[VAL_0:.*]]: tensor, %[[VAL_1:.*]]: tensor, %[[VAL_2:.*]]: tensor, %[[VAL_3:.*]]: tensor<8x40xf32>, %[[VAL_4:.*]]: tensor<10x40xf32>, %[[VAL_5:.*]]: tensor<40xf32>) -> (tensor, tensor, tensor, tensor, tensor) attributes {tf._input_shapes = ["tfshape$dim { size: -1 } dim { size: 8 } dim { size: 8 }", "tfshape$dim { size: -1 } dim { size: 10 }", "tfshape$dim { size: -1 } dim { size: 10 }", "tfshape$unknown_rank: true", "tfshape$unknown_rank: true", "tfshape$unknown_rank: true"], tf.api_implements = "lstm_b4e9f0e7-ac55-42bc-8ef2-8496419a608c", tf.api_preferred_device = "CPU", tf.go_backwards = false, tf.time_major = true} { -// CHECK: %[[VAL_6:.*]] = "tf.BatchMatMulV2"(%[[VAL_0]], %[[VAL_3]]) {adj_x = false, adj_y = false} : (tensor, tensor<8x40xf32>) -> tensor +// CHECK: %[[VAL_6:.*]] = "tf.BatchMatMulV2"(%[[VAL_0]], %[[VAL_3]]) <{adj_x = false, adj_y = false}> : (tensor, tensor<8x40xf32>) -> tensor // CHECK: %[[VAL_7:.*]] = "tf.Add"(%[[VAL_6]], %[[VAL_5]]) : (tensor, tensor<40xf32>) -> tensor -// CHECK: %[[VAL_8:.*]] = "tf.BatchMatMulV2"(%[[VAL_7]], %[[VAL_4]]) {adj_x = false, adj_y = true} : (tensor, tensor<10x40xf32>) -> tensor +// CHECK: %[[VAL_8:.*]] = "tf.BatchMatMulV2"(%[[VAL_7]], %[[VAL_4]]) <{adj_x = false, adj_y = true}> : (tensor, tensor<10x40xf32>) -> tensor // CHECK: %[[VAL_9:.*]] = "tf.Add"(%[[VAL_8]], %[[VAL_1]]) : (tensor, tensor) -> tensor // CHECK: %[[VAL_10:.*]] = "tf.Add"(%[[VAL_8]], %[[VAL_2]]) : (tensor, tensor) -> tensor // CHECK: %[[VAL_11:.*]] = "tf.Add"(%[[VAL_1]], %[[VAL_2]]) : (tensor, tensor) -> tensor -// CHECK: %[[VAL_12:.*]] = "tf.Const"() {_output_shapes = ["tfshape$"], device = "/device:CPU:0", dtype = f32, value = dense<1.000000e+00> : tensor} : () -> tensor +// CHECK: %[[VAL_12:.*]] = "tf.Const"() <{value = dense<1.000000e+00> : tensor}> {_output_shapes = ["tfshape$"], device = "/device:CPU:0", dtype = f32} : () -> tensor // CHECK: return %[[VAL_11]], %[[VAL_10]], %[[VAL_11]], %[[VAL_11]], %[[VAL_12]] : tensor, tensor, tensor, tensor, tensor // CHECK: } } diff --git a/tensorflow/compiler/mlir/lite/tests/prepare-tf-fake-quant-4bit.mlir b/tensorflow/compiler/mlir/lite/tests/prepare-tf-fake-quant-4bit.mlir index 9a865f08464772..dca4c21766ee4a 100644 --- a/tensorflow/compiler/mlir/lite/tests/prepare-tf-fake-quant-4bit.mlir +++ b/tensorflow/compiler/mlir/lite/tests/prepare-tf-fake-quant-4bit.mlir @@ -40,7 +40,7 @@ func.func @fakeQuantForActivationNoDuplication(tensor<8xf32>) -> (tensor<8x!quan %1 = "tfl.quantize"(%0) {qtype = tensor<8x!quant.uniform>} : (tensor<8xf32>) -> tensor<8x!quant.uniform> func.return %1 : tensor<8x!quant.uniform> -// CHECK: %0 = "tf.FakeQuantWithMinMaxVars"(%arg0, %cst, %cst_0) {narrow_range = false, num_bits = 3 : i64} +// CHECK: %0 = "tf.FakeQuantWithMinMaxVars"(%arg0, %cst, %cst_0) <{narrow_range = false, num_bits = 3 : i64}> // CHECK: %1 = "tfl.quantize"(%0) {qtype = tensor<8x!quant.uniform>} // CHECK: return %1 } diff --git a/tensorflow/compiler/mlir/lite/tests/prepare-tf-fake-quant.mlir b/tensorflow/compiler/mlir/lite/tests/prepare-tf-fake-quant.mlir index 0c5bdac0a0b792..c65cecc188f468 100644 --- a/tensorflow/compiler/mlir/lite/tests/prepare-tf-fake-quant.mlir +++ b/tensorflow/compiler/mlir/lite/tests/prepare-tf-fake-quant.mlir @@ -39,7 +39,7 @@ func.func @fakeQuantForActivationNoDuplication(tensor<8xf32>) -> (tensor<8x!quan %1 = "tfl.quantize"(%0) {qtype = tensor<8x!quant.uniform>} : (tensor<8xf32>) -> tensor<8x!quant.uniform> func.return %1 : tensor<8x!quant.uniform> -// CHECK: %0 = "tf.FakeQuantWithMinMaxVars"(%arg0, %cst, %cst_0) {narrow_range = false, num_bits = 5 : i64} +// CHECK: %0 = "tf.FakeQuantWithMinMaxVars"(%arg0, %cst, %cst_0) <{narrow_range = false, num_bits = 5 : i64}> // CHECK: %1 = "tfl.quantize"(%0) {qtype = tensor<8x!quant.uniform>} // CHECK: return %1 } diff --git a/tensorflow/compiler/mlir/lite/tests/prepare-tf.mlir b/tensorflow/compiler/mlir/lite/tests/prepare-tf.mlir index 4f3914265b4da5..fff00820ce353b 100644 --- a/tensorflow/compiler/mlir/lite/tests/prepare-tf.mlir +++ b/tensorflow/compiler/mlir/lite/tests/prepare-tf.mlir @@ -30,7 +30,7 @@ func.func @conv(tensor<256x32x32x3xf32>, tensor<3x3x3x16xf32>, tensor<256x3x32x3 // CHECK: %5 = "tf.Pad"(%arg0, %[[CONSTANT1]]) : (tensor<256x32x32x3xf32>, tensor<4x2xi32>) -> tensor<*xf32> // CHECK: %6 = "tf.Transpose"(%arg1, %[[CONSTANT0]]) : (tensor<3x3x3x16xf32>, tensor<4xi32>) -> tensor<16x3x3x3xf32> // CHECK: %7 = "tfl.conv_2d"(%5, %6, %[[CONSTANT]]) {dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32} : (tensor<*xf32>, tensor<16x3x3x3xf32>, tensor<16xf32>) -> tensor<256x32x32x16xf32> -// CHECK: %8 = "tf.Conv2D"(%arg0, %arg1) {T = "tfdtype$DT_FLOAT", data_format = "NHWC", dilations = [1, 1, 1, 1], padding = "SAME", strides = [2, 1, 1, 1]} : (tensor<256x32x32x3xf32>, tensor<3x3x3x16xf32>) -> tensor<256x32x32x16xf32> +// CHECK: %8 = "tf.Conv2D"(%arg0, %arg1) <{data_format = "NHWC", dilations = [1, 1, 1, 1], padding = "SAME", strides = [2, 1, 1, 1]}> {T = "tfdtype$DT_FLOAT"} : (tensor<256x32x32x3xf32>, tensor<3x3x3x16xf32>) -> tensor<256x32x32x16xf32> } func.func @depthwiseConv2D(tensor<256x32x32x3xf32>, tensor<3x3x3x4xf32>, tensor<256x3x32x32xf32>) -> (tensor<256x30x30x12xf32>, tensor<256x12x30x30xf32>, tensor<256x30x30x12xf32>, tensor<256x30x30x12xf32>) { @@ -224,9 +224,9 @@ func.func @matmulNoTransposeAOrB(%arg0: tensor<1x1280xf32>, %arg1: tensor<1280x1 func.return %166 : tensor<1x1000xf32> // CHECK-LABEL: matmulNoTransposeAOrB - // CHECK: %[[RES:.*]] = "tf.Const"() {value = dense<[1, 0]> : tensor<2xi32>} : () -> tensor + // CHECK: %[[RES:.*]] = "tf.Const"() <{value = dense<[1, 0]> : tensor<2xi32>}> : () -> tensor // CHECK: %[[TRANS:.*]] = "tf.Transpose"(%arg1, %[[RES]]) : (tensor<1280x1000xf32>, tensor) -> tensor<*xf32> - // CHECK: %[[MM:.*]] = "tf.MatMul"(%arg0, %[[TRANS]]) {transpose_a = false, transpose_b = true} : (tensor<1x1280xf32>, tensor<*xf32>) -> tensor<1x1000xf32> + // CHECK: %[[MM:.*]] = "tf.MatMul"(%arg0, %[[TRANS]]) <{transpose_a = false, transpose_b = true}> : (tensor<1x1280xf32>, tensor<*xf32>) -> tensor<1x1000xf32> // CHECK: return %[[MM]] : tensor<1x1000xf32> } @@ -235,10 +235,10 @@ func.func @matmulNoTransposeB(%arg0: tensor<1x1280xf32>, %arg1: tensor<1280x1000 func.return %166 : tensor<1x1000xf32> // CHECK-LABEL: matmulNoTransposeB - // CHECK: %[[RES:.*]] = "tf.Const"() {value = dense<[1, 0]> : tensor<2xi32>} : () -> tensor + // CHECK: %[[RES:.*]] = "tf.Const"() <{value = dense<[1, 0]> : tensor<2xi32>}> : () -> tensor // CHECK: %[[TRANS1:.*]] = "tf.Transpose"(%arg0, %[[RES]]) : (tensor<1x1280xf32>, tensor) -> tensor<*xf32> // CHECK: %[[TRANS2:.*]] = "tf.Transpose"(%arg1, %[[RES]]) : (tensor<1280x1000xf32>, tensor) -> tensor<*xf32> - // CHECK: %[[MM:.*]] = "tf.MatMul"(%[[TRANS1]], %[[TRANS2]]) {transpose_a = false, transpose_b = true} : (tensor<*xf32>, tensor<*xf32>) -> tensor<1x1000xf32> + // CHECK: %[[MM:.*]] = "tf.MatMul"(%[[TRANS1]], %[[TRANS2]]) <{transpose_a = false, transpose_b = true}> : (tensor<*xf32>, tensor<*xf32>) -> tensor<1x1000xf32> // CHECK: return %[[MM]] : tensor<1x1000xf32> } @@ -284,7 +284,7 @@ func.func @StridedSliceEllipsisMaskBefore(%arg0: tensor<21x15x7xf32>) -> tensor< // CHECK-DAG: %[[CST:.*]] = arith.constant dense<0> : tensor<3xi32> // CHECK-DAG: %[[CST_0:.*]] = arith.constant dense<1> : tensor<3xi32> - // CHECK: %[[STRIDED_SLICE:.*]] = "tf.StridedSlice"(%arg0, %[[CST]], %[[CST]], %[[CST_0]]) {begin_mask = 3 : i64, ellipsis_mask = 0 : i64, end_mask = 3 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 0 : i64} : (tensor<21x15x7xf32>, tensor<3xi32>, tensor<3xi32>, tensor<3xi32>) -> tensor<21x15x2xf32> + // CHECK: %[[STRIDED_SLICE:.*]] = "tf.StridedSlice"(%arg0, %[[CST]], %[[CST]], %[[CST_0]]) <{begin_mask = 3 : i64, ellipsis_mask = 0 : i64, end_mask = 3 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 0 : i64}> : (tensor<21x15x7xf32>, tensor<3xi32>, tensor<3xi32>, tensor<3xi32>) -> tensor<21x15x2xf32> } // CHECK-LABEL: @StridedSliceEllipsisMaskBeforeWithBeginAndEndMask @@ -298,7 +298,7 @@ func.func @StridedSliceEllipsisMaskBeforeWithBeginAndEndMask(%arg0: tensor<4x5x4 // CHECK-DAG: %[[CST:.*]] = arith.constant dense<[0, 1, 0]> : tensor<3xi32> // CHECK-DAG: %[[CST_0:.*]] = arith.constant dense<0> : tensor<3xi32> // CHECK-DAG: %[[CST_1:.*]] = arith.constant dense<1> : tensor<3xi32> - // CHECK: %[[STRIDED_SLICE:.*]] = "tf.StridedSlice"(%arg0, %[[CST]], %[[CST_0]], %[[CST_1]]) {begin_mask = 7 : i64, ellipsis_mask = 0 : i64, end_mask = 5 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 0 : i64} : (tensor<4x5x4xf32>, tensor<3xi32>, tensor<3xi32>, tensor<3xi32>) -> tensor<4x4x4xf32> + // CHECK: %[[STRIDED_SLICE:.*]] = "tf.StridedSlice"(%arg0, %[[CST]], %[[CST_0]], %[[CST_1]]) <{begin_mask = 7 : i64, ellipsis_mask = 0 : i64, end_mask = 5 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 0 : i64}> : (tensor<4x5x4xf32>, tensor<3xi32>, tensor<3xi32>, tensor<3xi32>) -> tensor<4x4x4xf32> } // CHECK-LABEL: @StridedSliceEllipsisMaskAfter @@ -310,7 +310,7 @@ func.func @StridedSliceEllipsisMaskAfter(%arg0: tensor<21x15x7xf32>) -> tensor<5 // CHECK-DAG: %[[CST:.*]] = arith.constant dense<0> : tensor<3xi32> // CHECK-DAG: %[[CST_0:.*]] = arith.constant dense<1> : tensor<3xi32> - // CHECK: %[[STRIDED_SLICE:.*]] = "tf.StridedSlice"(%arg0, %[[CST]], %[[CST]], %[[CST_0]]) {begin_mask = 6 : i64, ellipsis_mask = 0 : i64, end_mask = 6 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 0 : i64} : (tensor<21x15x7xf32>, tensor<3xi32>, tensor<3xi32>, tensor<3xi32>) -> tensor<5x15x7xf32> + // CHECK: %[[STRIDED_SLICE:.*]] = "tf.StridedSlice"(%arg0, %[[CST]], %[[CST]], %[[CST_0]]) <{begin_mask = 6 : i64, ellipsis_mask = 0 : i64, end_mask = 6 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 0 : i64}> : (tensor<21x15x7xf32>, tensor<3xi32>, tensor<3xi32>, tensor<3xi32>) -> tensor<5x15x7xf32> } // CHECK-LABEL: @NoStridedSliceEllipsisMask @@ -322,7 +322,7 @@ func.func @NoStridedSliceEllipsisMask(%arg0: tensor<*xf32>) -> tensor<21x15x2xf3 // CHECK-DAG: %[[CST:.*]] = arith.constant dense<0> : tensor<2xi32> // CHECK-DAG: %[[CST_0:.*]] = arith.constant dense<1> : tensor<2xi32> - // CHECK: %[[STRIDED_SLICE:.*]] = "tf.StridedSlice"(%arg0, %[[CST]], %[[CST]], %[[CST_0]]) {begin_mask = 0 : i64, ellipsis_mask = 1 : i64, end_mask = 0 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 0 : i64} : (tensor<*xf32>, tensor<2xi32>, tensor<2xi32>, tensor<2xi32>) -> tensor<21x15x2xf32> + // CHECK: %[[STRIDED_SLICE:.*]] = "tf.StridedSlice"(%arg0, %[[CST]], %[[CST]], %[[CST_0]]) <{begin_mask = 0 : i64, ellipsis_mask = 1 : i64, end_mask = 0 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 0 : i64}> : (tensor<*xf32>, tensor<2xi32>, tensor<2xi32>, tensor<2xi32>) -> tensor<21x15x2xf32> } // CHECK-LABEL: @NoPadStridedSliceNonNewAxisMask @@ -334,7 +334,7 @@ func.func @NoPadStridedSliceNonNewAxisMask(%arg0: tensor<1x2x3x1xf32>) -> tensor // CHECK-DAG: %cst = arith.constant dense<0> : tensor<4xi32> // CHECK-DAG: %cst_0 = arith.constant dense<1> : tensor<4xi32> - // CHECK: %0 = "tf.StridedSlice"(%arg0, %cst, %cst, %cst_0) {begin_mask = 15 : i64, ellipsis_mask = 0 : i64, end_mask = 15 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 0 : i64} : (tensor<1x2x3x1xf32>, tensor<4xi32>, tensor<4xi32>, tensor<4xi32>) -> tensor<1x2x3x1xf32> + // CHECK: %0 = "tf.StridedSlice"(%arg0, %cst, %cst, %cst_0) <{begin_mask = 15 : i64, ellipsis_mask = 0 : i64, end_mask = 15 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 0 : i64}> : (tensor<1x2x3x1xf32>, tensor<4xi32>, tensor<4xi32>, tensor<4xi32>) -> tensor<1x2x3x1xf32> } // CHECK-LABEL: @PadStridedSliceNewAxisMask1 @@ -348,7 +348,7 @@ func.func @PadStridedSliceNewAxisMask1(%arg0: tensor<2x3xf32>) -> tensor<1x2x3x1 // CHECK-DAG: %[[CST1:.*]] = arith.constant dense<1> : tensor<4xi32> // CHECK-DAG: %[[cst_1:.*]] = arith.constant dense<[1, 2, 3, 1]> : tensor<4xi32> // CHECK: %0 = "tf.Reshape"(%arg0, %[[cst_1]]) : (tensor<2x3xf32>, tensor<4xi32>) -> tensor<1x2x3x1xf32> - // CHECK: %1 = "tf.StridedSlice"(%0, %[[CST0]], %[[CST0]], %[[CST1]]) {begin_mask = 15 : i64, ellipsis_mask = 0 : i64, end_mask = 15 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 0 : i64} : (tensor<1x2x3x1xf32>, tensor<4xi32>, tensor<4xi32>, tensor<4xi32>) -> tensor<1x2x3x1xf32> + // CHECK: %1 = "tf.StridedSlice"(%0, %[[CST0]], %[[CST0]], %[[CST1]]) <{begin_mask = 15 : i64, ellipsis_mask = 0 : i64, end_mask = 15 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 0 : i64}> : (tensor<1x2x3x1xf32>, tensor<4xi32>, tensor<4xi32>, tensor<4xi32>) -> tensor<1x2x3x1xf32> } // CHECK-LABEL: @PadStridedSliceNewAxisMask2 @@ -401,7 +401,7 @@ func.func @strided_slice_with_constant_attributes(%arg0: tensor<10x10x10xf32>, % // CHECK-DAG: [[BEGIN:%cst.*]] = arith.constant dense<[-1, 0, 0]> : tensor<3xi32> // CHECK-DAG: [[END:%cst.*]] = arith.constant dense<[0, 10, 10]> : tensor<3xi32> // CHECK-DAG: [[STRIDES:%cst.*]] = arith.constant dense<1> : tensor<3xi32> - // CHECK-NEXT: "tf.StridedSlice"(%arg0, [[BEGIN]], [[END]], [[STRIDES]]) {begin_mask = 6 : i64, ellipsis_mask = 0 : i64, end_mask = 6 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 1 : i64} : (tensor<10x10x10xf32>, tensor<3xi32>, tensor<3xi32>, tensor<3xi32>) -> tensor<10x10xf32> + // CHECK-NEXT: "tf.StridedSlice"(%arg0, [[BEGIN]], [[END]], [[STRIDES]]) <{begin_mask = 6 : i64, ellipsis_mask = 0 : i64, end_mask = 6 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 1 : i64}> : (tensor<10x10x10xf32>, tensor<3xi32>, tensor<3xi32>, tensor<3xi32>) -> tensor<10x10xf32> } // CHECK-LABEL: @StridedSliceEllipsisAndNewAxisMaskBothSet @@ -419,7 +419,7 @@ func.func @StridedSliceEllipsisAndNewAxisMaskBothSet(%arg0: tensor<6x7x8xf32>) - // CHECK-DAG: %[[STEP:.*]] = arith.constant dense<1> : tensor<5xi32> // CHECK-DAG: %[[NEW_DIMS:.*]] = arith.constant dense<[6, 1, 7, 8, 1]> : tensor<5xi32> // CHECK: %[[RESHAPE:.*]] = "tf.Reshape"(%arg0, %[[NEW_DIMS]]) : (tensor<6x7x8xf32>, tensor<5xi32>) -> tensor<6x1x7x8x1xf32> - // CHECK: %[[STRIDED_SLICE:.*]] = "tf.StridedSlice"(%[[RESHAPE]], %[[BEGIN]], %[[END]], %[[STEP]]) {begin_mask = 30 : i64, ellipsis_mask = 0 : i64, end_mask = 30 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 0 : i64} : (tensor<6x1x7x8x1xf32>, tensor<5xi32>, tensor<5xi32>, tensor<5xi32>) -> tensor<2x1x7x8x1xf32> + // CHECK: %[[STRIDED_SLICE:.*]] = "tf.StridedSlice"(%[[RESHAPE]], %[[BEGIN]], %[[END]], %[[STEP]]) <{begin_mask = 30 : i64, ellipsis_mask = 0 : i64, end_mask = 30 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 0 : i64}> : (tensor<6x1x7x8x1xf32>, tensor<5xi32>, tensor<5xi32>, tensor<5xi32>) -> tensor<2x1x7x8x1xf32> } // CHECK-LABEL: @StridedSliceShrinkAxisAndNewAxisMaskBothSet @@ -437,7 +437,7 @@ func.func @StridedSliceShrinkAxisAndNewAxisMaskBothSet(%arg0: tensor<6x7x8xf32>) // CHECK-DAG: %[[END:.*]] = arith.constant dense<[2, 3, 4, 5, 8]> : tensor<5xi32> // CHECK-DAG: %[[STEP:.*]] = arith.constant dense<1> : tensor<5xi32> // CHECK: %[[RESHAPE:.*]] = "tf.Reshape"(%arg0, %[[NEW_DIMS]]) : (tensor<6x7x8xf32>, tensor<5xi32>) -> tensor<6x1x7x1x8xf32> - // CHECK: %[[STRIDED_SLICE:.*]] = "tf.StridedSlice"(%[[RESHAPE]], %[[BEGIN]], %[[END]], %[[STEP]]) {begin_mask = 26 : i64, ellipsis_mask = 0 : i64, end_mask = 26 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 1 : i64} : (tensor<6x1x7x1x8xf32>, tensor<5xi32>, tensor<5xi32>, tensor<5xi32>) -> tensor<1x4x1x8xf32> + // CHECK: %[[STRIDED_SLICE:.*]] = "tf.StridedSlice"(%[[RESHAPE]], %[[BEGIN]], %[[END]], %[[STEP]]) <{begin_mask = 26 : i64, ellipsis_mask = 0 : i64, end_mask = 26 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 1 : i64}> : (tensor<6x1x7x1x8xf32>, tensor<5xi32>, tensor<5xi32>, tensor<5xi32>) -> tensor<1x4x1x8xf32> } func.func @broadcast_to_f32_low_dim(%arg0: tensor<3xf32>, %arg1: tensor<2xi32>) -> tensor<3x3xf32> { @@ -572,7 +572,7 @@ func.func @lower_rfft_to_rfft2d(%input: tensor<10x20x30xf32>, %fft_len: tensor<1 // CHECK: %[[EXP:.*]] = "tf.ExpandDims"(%arg0, %[[CST]]) : (tensor<10x20x30xf32>, tensor) -> tensor<10x20x1x30xf32> // CHECK: %[[CON:.*]] = "tf.ConcatV2"(%[[CST0]], %arg1, %[[CST1]]) : (tensor<1xi32>, tensor<1xi32>, tensor) -> tensor<2xi32> // CHECK: %[[RFF:.*]] = "tf.RFFT2D"(%[[EXP]], %[[CON]]) : (tensor<10x20x1x30xf32>, tensor<2xi32>) -> tensor<10x20x1x30xcomplex> -// CHECK: %[[SQE:.*]] = "tf.Squeeze"(%[[RFF]]) {squeeze_dims = [-2]} : (tensor<10x20x1x30xcomplex>) -> tensor<10x20x30xcomplex> +// CHECK: %[[SQE:.*]] = "tf.Squeeze"(%[[RFF]]) <{squeeze_dims = [-2]}> : (tensor<10x20x1x30xcomplex>) -> tensor<10x20x30xcomplex> } // CHECK-LABEL: xla_gather_to_strided_slice @@ -585,7 +585,7 @@ func.func @xla_gather_to_strided_slice(%arg0 : tensor<1x9x104x768xf32>) -> tenso // CHECK-DAG: %[[CST:.*]] = arith.constant dense<0> : tensor<4xi64> // CHECK-DAG: %[[CST0:.*]] = arith.constant dense<[1, 9, 23, 768]> : tensor<4xi64> // CHECK-DAG: %[[CST1:.*]] = arith.constant dense<1> : tensor<4xi64> -// CHECK: %[[V0:.*]] = "tf.StridedSlice"(%arg0, %[[CST]], %[[CST0]], %[[CST1]]) {begin_mask = 0 : i64, ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 0 : i64} : (tensor<1x9x104x768xf32>, tensor<4xi64>, tensor<4xi64>, tensor<4xi64>) -> tensor<*xf32> +// CHECK: %[[V0:.*]] = "tf.StridedSlice"(%arg0, %[[CST]], %[[CST0]], %[[CST1]]) <{begin_mask = 0 : i64, ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 0 : i64}> : (tensor<1x9x104x768xf32>, tensor<4xi64>, tensor<4xi64>, tensor<4xi64>) -> tensor<*xf32> // CHECK: return %[[V0]] : tensor<*xf32> } @@ -660,9 +660,9 @@ func.func @fused_batch_norm_v3_training(%arg0 : tensor<1x1x6x2xf32>, %arg1 : ten // CHECK-LABEL: fused_batch_norm_v3_training // CHECK-DAG: %[[CST:.*]] = arith.constant dense<[0, 1, 2]> : tensor<3xi32> // CHECK-DAG: %[[CST1:.*]] = arith.constant dense<1.000000e-03> : tensor - // CHECK: %[[MEAN:.*]] = "tf.Mean"(%arg0, %[[CST]]) {keep_dims = false} : (tensor<1x1x6x2xf32>, tensor<3xi32>) -> tensor<2xf32> + // CHECK: %[[MEAN:.*]] = "tf.Mean"(%arg0, %[[CST]]) <{keep_dims = false}> : (tensor<1x1x6x2xf32>, tensor<3xi32>) -> tensor<2xf32> // CHECK: %[[SQ:.*]] = "tf.SquaredDifference"(%arg0, %[[MEAN]]) : (tensor<1x1x6x2xf32>, tensor<2xf32>) -> tensor<1x1x6x2xf32> - // CHECK: %[[MEAN0:.*]] = "tf.Mean"(%[[SQ]], %[[CST]]) {keep_dims = false} : (tensor<1x1x6x2xf32>, tensor<3xi32>) -> tensor<2xf32> + // CHECK: %[[MEAN0:.*]] = "tf.Mean"(%[[SQ]], %[[CST]]) <{keep_dims = false}> : (tensor<1x1x6x2xf32>, tensor<3xi32>) -> tensor<2xf32> // CHECK: %[[ADD:.*]] = "tf.Add"(%[[MEAN0]], %[[CST1]]) : (tensor<2xf32>, tensor) -> tensor<2xf32> // CHECK: %[[RSQRT:.*]] = "tf.Rsqrt"(%[[ADD]]) : (tensor<2xf32>) -> tensor<2xf32> // CHECK: %[[MUL1:.*]] = "tf.Mul"(%arg1, %[[RSQRT]]) : (tensor<2xf32>, tensor<2xf32>) -> tensor<2xf32> @@ -689,10 +689,10 @@ func.func @add_v2_uint32(%arg0: tensor, %arg1: tensor) -> tensor // CHECK-LABEL: add_v2_uint32 - // CHECK: %[[CAST:.*]] = "tf.Cast"(%arg0) {Truncate = false} : (tensor) -> tensor - // CHECK: %[[CAST1:.*]] = "tf.Cast"(%arg1) {Truncate = false} : (tensor) -> tensor + // CHECK: %[[CAST:.*]] = "tf.Cast"(%arg0) <{Truncate = false}> : (tensor) -> tensor + // CHECK: %[[CAST1:.*]] = "tf.Cast"(%arg1) <{Truncate = false}> : (tensor) -> tensor // CHECK: %[[ADD:.*]] = "tf.AddV2"(%[[CAST]], %[[CAST1]]) : (tensor, tensor) -> tensor - // CHECK: %[[CAST2:.*]] = "tf.Cast"(%[[ADD]]) {Truncate = false} : (tensor) -> tensor + // CHECK: %[[CAST2:.*]] = "tf.Cast"(%[[ADD]]) <{Truncate = false}> : (tensor) -> tensor // CHECK: return %[[CAST2]] : tensor } @@ -713,12 +713,12 @@ func.func @QuantDequantTranspose(%arg0: tensor<2x3xf32>) -> (tensor<2x4xf32>) { func.return %6 : tensor<2x4xf32> // CHECK-LABEL: QuantDequantTranspose - // CHECK-DAG: %[[CST:.*]] = "tf.Const"() {value = dense<[1, 0]> : tensor<2xi32>} : () -> tensor + // CHECK-DAG: %[[CST:.*]] = "tf.Const"() <{value = dense<[1, 0]> : tensor<2xi32>}> : () -> tensor // CHECK-DAG: %[[CST_0:.*]] = arith.constant dense<1.00392163> : tensor<3x4xf32> // CHECK: %[[QUANT:.*]] = "tfl.quantize"(%[[CST_0]]) {qtype = tensor<3x4x!quant.uniform>} : (tensor<3x4xf32>) -> tensor<3x4x!quant.uniform> // CHECK: %[[DEQUANT:.*]] = "tfl.dequantize"(%[[QUANT]]) : (tensor<3x4x!quant.uniform>) -> tensor<3x4xf32> // CHECK: %[[TRANSPOSE:.*]] = "tf.Transpose"(%[[DEQUANT]], %[[CST]]) : (tensor<3x4xf32>, tensor) -> tensor<*xf32> - // CHECK: %[[MATMUL:.*]] = "tf.MatMul"(%arg0, %[[TRANSPOSE]]) {transpose_a = false, transpose_b = true} : (tensor<2x3xf32>, tensor<*xf32>) -> tensor<2x4xf32> + // CHECK: %[[MATMUL:.*]] = "tf.MatMul"(%arg0, %[[TRANSPOSE]]) <{transpose_a = false, transpose_b = true}> : (tensor<2x3xf32>, tensor<*xf32>) -> tensor<2x4xf32> // CHECK: return %[[MATMUL]] : tensor<2x4xf32> } diff --git a/tensorflow/compiler/mlir/lite/tests/raise-custom-ops.mlir b/tensorflow/compiler/mlir/lite/tests/raise-custom-ops.mlir index e72f421ad9eafb..477315d696783c 100644 --- a/tensorflow/compiler/mlir/lite/tests/raise-custom-ops.mlir +++ b/tensorflow/compiler/mlir/lite/tests/raise-custom-ops.mlir @@ -42,11 +42,11 @@ func.func @tf_executor_wrapper(%arg0: tensor<*xf32>) -> tensor<*xf32> attributes // CHECK: tf_executor.island wraps "tf.FakeQuantWithMinMaxVarsPerChannel" // WRAPPED-NEXT: tf_executor.graph { -// WRAPPED-NEXT: tf_executor.island wraps "tf.Const"() {device = "", value = dense<1.000000e+00> : tensor<186xf32>} : () -> tensor<186xf32> -// WRAPPED-NEXT: tf_executor.island wraps "tf.Const"() {device = "", value = dense<2.000000e+00> : tensor<186xf32>} : () -> tensor<186xf32> +// WRAPPED-NEXT: tf_executor.island wraps "tf.Const"() <{value = dense<1.000000e+00> : tensor<186xf32>}> {device = ""} : () -> tensor<186xf32> +// WRAPPED-NEXT: tf_executor.island wraps "tf.Const"() <{value = dense<2.000000e+00> : tensor<186xf32>}> {device = ""} : () -> tensor<186xf32> // WRAPPED-NEXT: tf_executor.island wraps "tfl.custom_tf" // WRAPPED-NEXT: ^bb0(%arg1: tensor<*xf32>, %arg2: tensor<186xf32>, %arg3: tensor<186xf32>): -// WRAPPED-NEXT: %[[fq:.*]] = "tf.FakeQuantWithMinMaxVarsPerChannel"(%arg1, %arg2, %arg3) {device = "", narrow_range = true, num_bits = 8 : i64} : (tensor<*xf32>, tensor<186xf32>, tensor<186xf32>) -> tensor<*xf32> +// WRAPPED-NEXT: %[[fq:.*]] = "tf.FakeQuantWithMinMaxVarsPerChannel"(%arg1, %arg2, %arg3) <{narrow_range = true, num_bits = 8 : i64}> {device = ""} : (tensor<*xf32>, tensor<186xf32>, tensor<186xf32>) -> tensor<*xf32> // WRAPPED-NEXT: "tfl.yield"(%[[fq]]) : (tensor<*xf32>) -> () // WRAPPED-NEXT: }) {device = "", narrow_range = true, num_bits = 8 : i64} : (tensor<*xf32>, tensor<186xf32>, tensor<186xf32>) -> tensor<*xf32> } diff --git a/tensorflow/compiler/mlir/lite/transforms/lift_tflite_flex_ops.cc b/tensorflow/compiler/mlir/lite/transforms/lift_tflite_flex_ops.cc index 9563f5d93226fb..f8d1de04649169 100644 --- a/tensorflow/compiler/mlir/lite/transforms/lift_tflite_flex_ops.cc +++ b/tensorflow/compiler/mlir/lite/transforms/lift_tflite_flex_ops.cc @@ -19,6 +19,7 @@ limitations under the License. #include #include +#include "absl/strings/match.h" #include "flatbuffers/flexbuffers.h" // from @flatbuffers #include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project #include "mlir/IR/Attributes.h" // from @llvm-project @@ -111,20 +112,6 @@ class LiftFlexCustomOp : public OpRewritePattern { Operation* tf_op = rewriter.create(op_state); rewriter.replaceOp(op, tf_op->getResults()); - if (isa(tf_op)) { - constexpr StringRef kFuncAttrName = "f"; - tf_op->setAttr( - kFuncAttrName, - tf_op->getAttr(kFuncAttrName).cast().getName()); - } - - if (isa(tf_op)) { - constexpr StringRef kFuncAttrName = "predicate"; - tf_op->setAttr( - kFuncAttrName, - tf_op->getAttr(kFuncAttrName).cast().getName()); - } - // Special type fixes for TF Resource Tensors that are casted to // Int32 tensor during MLIR->TFLite flatbuffer conversion. // TODO(b/146131919): correct handling of resource type @@ -237,6 +224,10 @@ class LiftFlexCustomOp : public OpRewritePattern { if (!mlir_attr.ok()) { return emitError(loc, mlir_attr.status().message()); } + if (absl::StrContains(op_name, "Dataset") && + mlir_attr->isa()) { + mlir_attr = mlir_attr->cast().getName(); + } attributes.push_back(builder.getNamedAttr(attr_name, *mlir_attr)); } return success(); diff --git a/tensorflow/compiler/mlir/quantization/stablehlo/BUILD b/tensorflow/compiler/mlir/quantization/stablehlo/BUILD index a3210f64239403..ea2d77a865dcec 100644 --- a/tensorflow/compiler/mlir/quantization/stablehlo/BUILD +++ b/tensorflow/compiler/mlir/quantization/stablehlo/BUILD @@ -32,8 +32,10 @@ cc_library( "passes/lift_quantizable_spots_as_functions.cc", "passes/lift_quantizable_spots_as_functions_fusion.inc", "passes/lift_quantizable_spots_as_functions_simple.inc", + "passes/post_quantize.cc", "passes/prepare_quantize.cc", "passes/quantize.cc", + "passes/quantize_composite_functions.cc", "passes/quantize_weight.cc", "passes/replace_stablehlo_ops_in_main_function_with_xla_call_module_ops.cc", "passes/restore_function_name.cc", @@ -51,6 +53,7 @@ cc_library( ":quantization_options_proto_cc", ":stablehlo_passes_inc_gen", ":stablehlo_type_utils", + ":uniform_quantized_types", "//tensorflow/compiler/mlir/lite:tensorflow_lite", "//tensorflow/compiler/mlir/lite/quantization:quantization_config", "//tensorflow/compiler/mlir/lite/quantization:quantization_lib", @@ -58,6 +61,7 @@ cc_library( "//tensorflow/compiler/mlir/quantization/tensorflow:pass_utils", "//tensorflow/compiler/mlir/quantization/tensorflow:passes", "//tensorflow/compiler/mlir/quantization/tensorflow:quantization_options_proto_cc", + "//tensorflow/compiler/mlir/quantization/tensorflow/cc:run_passes", "//tensorflow/compiler/mlir/quantization/tensorflow/ops:tf_op_quant_spec", "//tensorflow/compiler/mlir/quantization/tensorflow/utils:lift_as_function_call_utils", "//tensorflow/compiler/mlir/tensorflow", @@ -185,7 +189,6 @@ cc_library( ], deps = [ ":bridge_passes_inc_gen", - ":math_utils", ":tf_type_utils", "//tensorflow/compiler/mlir/tensorflow", "//tensorflow/compiler/mlir/tensorflow:mangling_util", @@ -210,7 +213,6 @@ cc_library( "@llvm-project//mlir:Pass", "@llvm-project//mlir:QuantOps", "@llvm-project//mlir:ShapeDialect", - "@llvm-project//mlir:SparseTensorDialect", "@llvm-project//mlir:Support", "@llvm-project//mlir:Transforms", "@local_xla//xla:xla_data_proto_cc", @@ -238,9 +240,12 @@ tf_cc_test( "//tensorflow/compiler/mlir/tensorflow:tensorflow_types", "//tensorflow/compiler/tf2xla:common", "//tensorflow/core:framework", + "//tensorflow/core/kernels:math", + "//tensorflow/core/kernels:nn", "//tensorflow/core/kernels/uniform_quant_ops:kernels", "//tensorflow/core/ops", "@com_google_absl//absl/log:check", + "@com_google_absl//absl/random", "@com_google_absl//absl/status", "@com_google_absl//absl/strings", "@com_google_absl//absl/types:span", diff --git a/tensorflow/compiler/mlir/quantization/stablehlo/passes/bridge/convert_mhlo_quant_to_int.cc b/tensorflow/compiler/mlir/quantization/stablehlo/passes/bridge/convert_mhlo_quant_to_int.cc index 505554bdad4bb0..16af4b212ecfe8 100644 --- a/tensorflow/compiler/mlir/quantization/stablehlo/passes/bridge/convert_mhlo_quant_to_int.cc +++ b/tensorflow/compiler/mlir/quantization/stablehlo/passes/bridge/convert_mhlo_quant_to_int.cc @@ -15,12 +15,10 @@ limitations under the License. #include #include -#include #include #include -#include -#include #include +#include #include "absl/algorithm/container.h" #include "llvm/ADT/ArrayRef.h" @@ -31,12 +29,11 @@ limitations under the License. #include "mlir/Dialect/Func/Transforms/FuncConversions.h" // from @llvm-project #include "mlir/Dialect/Quant/QuantOps.h" // from @llvm-project #include "mlir/Dialect/Quant/QuantTypes.h" // from @llvm-project -#include "mlir/Dialect/Shape/IR/Shape.h" // from @llvm-project -#include "mlir/Dialect/SparseTensor/IR/SparseTensor.h" // from @llvm-project #include "mlir/IR/Attributes.h" // from @llvm-project #include "mlir/IR/BuiltinAttributes.h" // from @llvm-project #include "mlir/IR/BuiltinTypeInterfaces.h" // from @llvm-project #include "mlir/IR/BuiltinTypes.h" // from @llvm-project +#include "mlir/IR/Location.h" // from @llvm-project #include "mlir/IR/MLIRContext.h" // from @llvm-project #include "mlir/IR/OpDefinition.h" // from @llvm-project #include "mlir/IR/OperationSupport.h" // from @llvm-project @@ -49,7 +46,6 @@ limitations under the License. #include "mlir/Transforms/DialectConversion.h" // from @llvm-project #include "stablehlo/dialect/ChloOps.h" // from @stablehlo #include "tensorflow/compiler/mlir/quantization/stablehlo/passes/bridge/passes.h" -#include "tensorflow/compiler/mlir/quantization/stablehlo/utils/math_utils.h" #include "tensorflow/compiler/mlir/tf2xla/transforms/xla_legalize_targets.h" #include "xla/mlir_hlo/mhlo/IR/hlo_ops.h" #include "xla/mlir_hlo/mhlo/transforms/rewriters.h" @@ -57,88 +53,193 @@ limitations under the License. namespace mlir::quant::stablehlo { namespace { -#define GEN_PASS_DEF_CONVERTMHLOQUANTTOINT -#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/bridge/passes.h.inc" - -// This helper function create ops to requantize `input` tensor and output to -// `res_int32` tensor. Clamping is omitted because for some ops clamping can be -// done later to avoid duplicate. -LogicalResult RequantizeWithoutClamping( - mlir::OpState op, Value input, TensorType int32_tensor_type, - quant::UniformQuantizedType input_quantized_type, - quant::UniformQuantizedType result_quantized_type, Value &res_int32, - ConversionPatternRewriter &rewriter) { +// This helper function create ops to requantize `input` tensor and returns the +// output tensor. Clamping is done if output integer bit-width < 32. +// +// Requantization is essentially dequantize --> quantize. +// +// Dequantize: (input - zp) * scale +// Quantize: input / scale + zp +// +// Hence, +// output = (input - input_zp) * input_scale / output_scale + output_zp +// +// This is simplified as: +// output = input * merged_scale + merged_zp +// where: +// merged_zp = output_zp - input_zp * merged_scale. +// merged_scale = input_scale / output_scale. +Value Requantize(mlir::OpState op, Value input, + UniformQuantizedType input_quantized_type, + UniformQuantizedType output_quantized_type, + TensorType output_tensor_type, + ConversionPatternRewriter &rewriter) { // Skip requantization when input and result have the same type. - if (input_quantized_type == result_quantized_type) { - res_int32 = rewriter.create(op->getLoc(), - int32_tensor_type, input); - return success(); + if (input_quantized_type == output_quantized_type) { + return rewriter.create(op->getLoc(), output_tensor_type, + input); } - // Convert input to int32 tensor. - res_int32 = - rewriter.create(op->getLoc(), int32_tensor_type, input); - // Undo the input zero point. - Value input_zero_point = rewriter.create( - op->getLoc(), rewriter.getI32IntegerAttr(static_cast( - input_quantized_type.getZeroPoint()))); - res_int32 = rewriter.create( - op->getLoc(), int32_tensor_type, res_int32, input_zero_point, nullptr); - - // Adjust the scale. - const double effective_scale = - input_quantized_type.getScale() / result_quantized_type.getScale(); - int32_t effective_quantized_fraction; - int32_t effective_shift; - if (failed(quant::stablehlo::QuantizeMultiplier( - effective_scale, effective_quantized_fraction, effective_shift))) { - op->emitError("Invalid effective quantization scale."); + double merged_scale_fp = + input_quantized_type.getScale() / output_quantized_type.getScale(); + Value merged_scale = rewriter.create( + op->getLoc(), + rewriter.getF32FloatAttr(static_cast(merged_scale_fp))); + + auto float_tensor_type = + input.getType().cast().clone(rewriter.getF32Type()); + Value output_float = + rewriter.create(op->getLoc(), float_tensor_type, input); + + output_float = rewriter.create( + op->getLoc(), float_tensor_type, output_float, merged_scale, nullptr); + + // Add merged_zp only when it is non-zero. + double merged_zp_fp = output_quantized_type.getZeroPoint() - + input_quantized_type.getZeroPoint() * merged_scale_fp; + if (merged_zp_fp != 0) { + Value merged_zp = rewriter.create( + op->getLoc(), + rewriter.getF32FloatAttr(static_cast(merged_zp_fp))); + output_float = rewriter.create( + op->getLoc(), float_tensor_type, output_float, merged_zp, nullptr); + } + + // Clamp output if the output integer bit-width <32. + if (output_tensor_type.getElementType().cast().getWidth() < 32) { + Value quantization_min = rewriter.create( + op->getLoc(), rewriter.getF32FloatAttr(static_cast( + output_quantized_type.getStorageTypeMin()))); + Value quantization_max = rewriter.create( + op->getLoc(), rewriter.getF32FloatAttr(static_cast( + output_quantized_type.getStorageTypeMax()))); + // Clamp results by [quantization_min, quantization_max]. + output_float = rewriter.create( + op->getLoc(), float_tensor_type, quantization_min, output_float, + quantization_max); + } + + output_float = rewriter.create( + op->getLoc(), float_tensor_type, output_float); + return rewriter.create(op->getLoc(), output_tensor_type, + output_float); +} + +using QuantType = + std::variant; +FailureOr GetQuantType(Type type) { + if (auto quant_type = + getElementTypeOrSelf(type).dyn_cast()) { + return QuantType(quant_type); + } else if (auto quant_type = getElementTypeOrSelf(type) + .dyn_cast()) { + return QuantType(quant_type); + } else { return failure(); } - Value multiplier = rewriter.create( - op->getLoc(), rewriter.getI32IntegerAttr( - static_cast(effective_quantized_fraction))); - // The effective_quantized_fraction value has been quantized by multiplying - // (1 << 15). So, we have to shift it back by (15 - effective_shift) to get - // the desired outcome. - Value total_shift = rewriter.create( - op->getLoc(), - rewriter.getI32IntegerAttr(static_cast(15 - effective_shift))); - - // Apply the effective scale with rounding. - Value half = rewriter.create( - op->getLoc(), rewriter.getI32IntegerAttr( - static_cast(1 << (14 - effective_shift)))); - res_int32 = rewriter.create( - op->getLoc(), int32_tensor_type, res_int32, multiplier, nullptr); - res_int32 = rewriter.create( - op->getLoc(), int32_tensor_type, res_int32, half, nullptr); - res_int32 = rewriter.create( - op->getLoc(), int32_tensor_type, res_int32, total_shift, nullptr); - - // Apply the output zero point. - Value output_zero_point = rewriter.create( - op->getLoc(), rewriter.getI32IntegerAttr(static_cast( - result_quantized_type.getZeroPoint()))); - res_int32 = rewriter.create( - op->getLoc(), int32_tensor_type, res_int32, output_zero_point, nullptr); +} - return success(); +// Extract scale and zero point info from input quant type info. +void GetQuantizationParams(OpBuilder &builder, Location loc, + QuantType quant_type, Value &scales, + Value &zero_points, bool output_zero_point_in_fp, + DenseIntElementsAttr &broadcast_dims) { + // Get scales/zero points for per-tensor and per-axis quantization cases. + if (auto *quant_per_tensor_type = + std::get_if(&quant_type)) { + scales = builder.create( + loc, builder.getF32FloatAttr(quant_per_tensor_type->getScale())); + if (output_zero_point_in_fp) { + zero_points = builder.create( + loc, builder.getF32FloatAttr( + static_cast(quant_per_tensor_type->getZeroPoint()))); + } else { + zero_points = builder.create( + loc, builder.getI32IntegerAttr(static_cast( + quant_per_tensor_type->getZeroPoint()))); + } + } else { + auto &quant_per_channel_type = + std::get(quant_type); + llvm::SmallVector scales_vec; + for (auto scale : quant_per_channel_type.getScales()) + scales_vec.push_back(scale); + scales = builder.create( + loc, DenseFPElementsAttr::get( + RankedTensorType::get( + {static_cast( + quant_per_channel_type.getScales().size())}, + builder.getF32Type()), + scales_vec)); + if (output_zero_point_in_fp) { + llvm::SmallVector zero_points_vec; + for (auto zero_point : quant_per_channel_type.getZeroPoints()) + zero_points_vec.push_back(zero_point); + zero_points = builder.create( + loc, DenseFPElementsAttr::get( + RankedTensorType::get( + {static_cast( + quant_per_channel_type.getZeroPoints().size())}, + builder.getF32Type()), + zero_points_vec)); + } else { + llvm::SmallVector zero_points_vec; + for (auto zero_point : quant_per_channel_type.getZeroPoints()) + zero_points_vec.push_back(zero_point); + zero_points = builder.create( + loc, DenseIntElementsAttr::get( + RankedTensorType::get( + {static_cast( + quant_per_channel_type.getZeroPoints().size())}, + builder.getI32Type()), + zero_points_vec)); + } + broadcast_dims = DenseIntElementsAttr::get( + RankedTensorType::get({1}, builder.getI64Type()), + {static_cast(quant_per_channel_type.getQuantizedDimension())}); + } } -class ConvertMHLOQuantToInt - : public impl::ConvertMHLOQuantToIntBase { - public: - ConvertMHLOQuantToInt() = default; - ConvertMHLOQuantToInt(const ConvertMHLOQuantToInt &) {} +// Extract storage min/max from input quant type info. +void GetQuantizationStorageInfo(OpBuilder &builder, Location loc, + QuantType quant_type, Value &storage_min, + Value &storage_max) { + if (auto *quant_per_tensor_type = + std::get_if(&quant_type)) { + storage_min = builder.create( + loc, builder.getF32FloatAttr(static_cast( + quant_per_tensor_type->getStorageTypeMin()))); + storage_max = builder.create( + loc, builder.getF32FloatAttr(static_cast( + quant_per_tensor_type->getStorageTypeMax()))); + } else { + auto &quant_per_channel_type = + std::get(quant_type); + storage_min = builder.create( + loc, builder.getF32FloatAttr(static_cast( + quant_per_channel_type.getStorageTypeMin()))); + storage_max = builder.create( + loc, builder.getF32FloatAttr(static_cast( + quant_per_channel_type.getStorageTypeMax()))); + } +} - explicit ConvertMHLOQuantToInt(bool legalize_chlo) { - legalize_chlo_ = legalize_chlo; +// Get storage type of a UQ type. Return original type if it is no UQ type. +Type GetQuantStorageType(Type type) { + if (auto shaped = type.dyn_cast()) { + return shaped.clone(GetQuantStorageType(shaped.getElementType())); } - // Performs conversion of MHLO quant ops to primitive ops. - void runOnOperation() override; -}; + if (auto element_type = + getElementTypeOrSelf(type).dyn_cast()) { + return element_type.getStorageType(); + } else if (auto element_type = getElementTypeOrSelf(type) + .dyn_cast()) { + return element_type.getStorageType(); + } else { + return type; + } +} class ConvertUniformQuantizeOp : public OpConversionPattern { @@ -148,124 +249,66 @@ class ConvertUniformQuantizeOp LogicalResult matchAndRewrite( mhlo::UniformQuantizeOp op, mhlo::UniformQuantizeOpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override { - auto quantized_type = getElementTypeOrSelf(op.getResult().getType()) - .dyn_cast(); - // Currently for activation, PTQ supports per-tensor quantization only, and - // UniformQuantize op is only for activation. - if (!quantized_type) { - return rewriter.notifyMatchFailure( - op, "Legalization supports only per-tensor quantization."); - } auto input_element_type = getElementTypeOrSelf(op.getOperand().getType()); if (input_element_type.isF32()) { - return matchAndRewriteQuantize(op, adaptor, rewriter, quantized_type); - } else if (input_element_type.isa()) { - return matchAndRewriteRequantize(op, adaptor, rewriter, quantized_type); + auto quant_type = GetQuantType(op.getResult().getType()); + if (succeeded(quant_type)) { + return matchAndRewriteQuantize(op, adaptor, rewriter, *quant_type); + } + } else if (input_element_type.isa()) { + return matchAndRewriteRequantize(op, adaptor, rewriter); } return rewriter.notifyMatchFailure(op, "Unsupported input element type."); } - LogicalResult matchAndRewriteQuantize( - mhlo::UniformQuantizeOp op, mhlo::UniformQuantizeOpAdaptor adaptor, - ConversionPatternRewriter &rewriter, - const quant::UniformQuantizedType &quantized_type) const { - Value scale = rewriter.create( - op->getLoc(), rewriter.getF32FloatAttr(quantized_type.getScale())); - Value zero_point = rewriter.create( - op->getLoc(), rewriter.getF32FloatAttr( - static_cast(quantized_type.getZeroPoint()))); - Value quantization_min = rewriter.create( - op->getLoc(), rewriter.getF32FloatAttr(static_cast( - quantized_type.getStorageTypeMin()))); - Value quantization_max = rewriter.create( - op->getLoc(), rewriter.getF32FloatAttr(static_cast( - quantized_type.getStorageTypeMax()))); + LogicalResult matchAndRewriteQuantize(mhlo::UniformQuantizeOp op, + mhlo::UniformQuantizeOpAdaptor adaptor, + ConversionPatternRewriter &rewriter, + QuantType quant_type) const { + Value scales, zero_points; + DenseIntElementsAttr broadcast_dims; + GetQuantizationParams(rewriter, op->getLoc(), quant_type, scales, + zero_points, /*output_zero_point_in_fp=*/true, + broadcast_dims); + + Value quantization_min, quantization_max; + GetQuantizationStorageInfo(rewriter, op->getLoc(), quant_type, + quantization_min, quantization_max); auto res_float_tensor_type = op.getOperand().getType().clone(rewriter.getF32Type()); Value res_float = rewriter.create( - op->getLoc(), res_float_tensor_type, adaptor.getOperand(), scale, - nullptr); + op->getLoc(), res_float_tensor_type, adaptor.getOperand(), scales, + broadcast_dims); res_float = rewriter.create( - op->getLoc(), res_float_tensor_type, res_float, zero_point, nullptr); + op->getLoc(), res_float_tensor_type, res_float, zero_points, + broadcast_dims); res_float = rewriter.create( op->getLoc(), res_float_tensor_type, quantization_min, res_float, quantization_max); res_float = rewriter.create( op->getLoc(), res_float_tensor_type, res_float); - auto res_final_tensor_type = - res_float_tensor_type.clone(quantized_type.getStorageType()); + auto res_final_tensor_type = res_float_tensor_type.clone( + GetQuantStorageType(op.getResult().getType().getElementType())); rewriter.replaceOpWithNewOp(op, res_final_tensor_type, res_float); return success(); } - // Requantization is essentially dequantize --> quantize. - // - // Dequantize: (input - zp) * scale - // Quantize: input / scale + zp - // - // Hence, - // result = (input - input_zp) * input_scale / output_scale + output_zp - // - // This is simplified as: - // result = input * merged_scale + merged_zp - // where: - // merged_zp = output_zp - input_zp * merged_scale. - // merged_scale = input_scale / output_scale. LogicalResult matchAndRewriteRequantize( mhlo::UniformQuantizeOp op, mhlo::UniformQuantizeOpAdaptor adaptor, - ConversionPatternRewriter &rewriter, - const quant::UniformQuantizedType &output_quantized_type) const { + ConversionPatternRewriter &rewriter) const { auto input_quantized_type = getElementTypeOrSelf(op.getOperand().getType()) - .cast(); - auto result_quantized_type = getElementTypeOrSelf(op.getResult().getType()) - .cast(); - - double merged_scale_fp = - input_quantized_type.getScale() / result_quantized_type.getScale(); - Value merged_scale = rewriter.create( - op->getLoc(), - rewriter.getF32FloatAttr(static_cast(merged_scale_fp))); - - auto res_float_tensor_type = - op.getOperand().getType().clone(rewriter.getF32Type()); - Value res_float = rewriter.create( - op->getLoc(), res_float_tensor_type, adaptor.getOperand()); - - res_float = rewriter.create( - op->getLoc(), res_float_tensor_type, res_float, merged_scale, nullptr); - - // Add merged_zp only when it is non-zero. - double merged_zp_fp = result_quantized_type.getZeroPoint() - - input_quantized_type.getZeroPoint() * merged_scale_fp; - if (merged_zp_fp != 0) { - Value merged_zp = rewriter.create( - op->getLoc(), - rewriter.getF32FloatAttr(static_cast(merged_zp_fp))); - res_float = rewriter.create( - op->getLoc(), res_float_tensor_type, res_float, merged_zp, nullptr); - } - - Value quantization_min = rewriter.create( - op->getLoc(), rewriter.getF32FloatAttr(static_cast( - output_quantized_type.getStorageTypeMin()))); - Value quantization_max = rewriter.create( - op->getLoc(), rewriter.getF32FloatAttr(static_cast( - output_quantized_type.getStorageTypeMax()))); - - // Clamp results by [quantization_min, quantization_max]. - res_float = rewriter.create( - op->getLoc(), res_float_tensor_type, quantization_min, res_float, - quantization_max); - res_float = rewriter.create( - op->getLoc(), res_float_tensor_type, res_float); - - auto res_final_tensor_type = - res_float_tensor_type.clone(output_quantized_type.getStorageType()); - rewriter.replaceOpWithNewOp(op, res_final_tensor_type, - res_float); + .cast(); + auto output_quantized_type = getElementTypeOrSelf(op.getResult().getType()) + .cast(); + rewriter.replaceOp( + op, Requantize(op, adaptor.getOperand(), input_quantized_type, + output_quantized_type, + op.getResult().getType().cast().clone( + output_quantized_type.getStorageType()), + rewriter)); return success(); } }; @@ -278,19 +321,15 @@ class ConvertUniformDequantizeOp LogicalResult matchAndRewrite( mhlo::UniformDequantizeOp op, mhlo::UniformDequantizeOpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override { - auto element_type = getElementTypeOrSelf(op.getOperand().getType()) - .dyn_cast(); - // Currently for activation, PTQ supports per-tensor quantization only, and - // UniformQuantize op is only for activation. - if (!element_type) { - return rewriter.notifyMatchFailure( - op, "Legalization supports only per-tensor quantization."); + auto quant_type = GetQuantType(op.getOperand().getType()); + if (failed(quant_type)) { + return failure(); } - Value scale = rewriter.create( - op->getLoc(), rewriter.getF32FloatAttr(element_type.getScale())); - Value zero_point = rewriter.create( - op->getLoc(), rewriter.getI32IntegerAttr( - static_cast(element_type.getZeroPoint()))); + Value scales, zero_points; + DenseIntElementsAttr broadcast_dims; + GetQuantizationParams(rewriter, op->getLoc(), *quant_type, scales, + zero_points, + /*output_zero_point_in_fp=*/false, broadcast_dims); Value input = adaptor.getOperand(); // TODO: b/260280919 - Consider avoiding conversion to int32. @@ -299,13 +338,14 @@ class ConvertUniformDequantizeOp Value res_int32 = rewriter.create( op->getLoc(), res_int32_tensor_type, input); res_int32 = rewriter.create( - op->getLoc(), res_int32_tensor_type, res_int32, zero_point, nullptr); + op->getLoc(), res_int32_tensor_type, res_int32, zero_points, + broadcast_dims); auto res_float_tensor_type = res_int32.getType().cast().clone(rewriter.getF32Type()); Value res_float = rewriter.create( op->getLoc(), res_float_tensor_type, res_int32); res_float = rewriter.replaceOpWithNewOp( - op, res_float_tensor_type, res_float, scale, nullptr); + op, res_float_tensor_type, res_float, scales, broadcast_dims); return success(); } }; @@ -317,18 +357,14 @@ class ConvertUniformQuantizedAddOp : public OpConversionPattern { LogicalResult matchAndRewrite( mhlo::AddOp op, mhlo::AddOpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override { - auto lhs_element_type = op.getLhs() - .getType() - .getElementType() - .dyn_cast(); - auto rhs_element_type = op.getRhs() - .getType() - .getElementType() - .dyn_cast(); + auto lhs_element_type = + op.getLhs().getType().getElementType().dyn_cast(); + auto rhs_element_type = + op.getRhs().getType().getElementType().dyn_cast(); auto result_element_type = op.getResult() .getType() .getElementType() - .dyn_cast(); + .dyn_cast(); // We only handle cases where lhs, rhs and results all have quantized // element type. @@ -347,20 +383,14 @@ class ConvertUniformQuantizedAddOp : public OpConversionPattern { // be the same as the result. // TODO: b/260280919 - Consider avoiding conversion to int32. Value lhs = adaptor.getLhs(); - Value lhs_int32_tensor; - if (failed(RequantizeWithoutClamping(op, lhs, res_int32_tensor_type, - lhs_element_type, result_element_type, - lhs_int32_tensor, rewriter))) { - return failure(); - } + Value lhs_int32_tensor = + Requantize(op, lhs, lhs_element_type, result_element_type, + res_int32_tensor_type, rewriter); Value rhs = adaptor.getRhs(); - Value rhs_int32_tensor; - if (failed(RequantizeWithoutClamping(op, rhs, res_int32_tensor_type, - rhs_element_type, result_element_type, - rhs_int32_tensor, rewriter))) { - return failure(); - } + Value rhs_int32_tensor = + Requantize(op, rhs, rhs_element_type, result_element_type, + res_int32_tensor_type, rewriter); Value zero_point = rewriter.create( op->getLoc(), rewriter.getI32IntegerAttr(static_cast( @@ -437,9 +467,9 @@ LogicalResult matchAndRewriteDotLikeHybridOp( // result = hybridOp(lhs, dequant(rhs)) Value lhs_float32_tensor = adaptor.getLhs(); Value rhs = adaptor.getRhs(); - quant::UniformQuantizedType rhs_element_type = + UniformQuantizedType rhs_element_type = getElementTypeOrSelf(op.getRhs().getType()) - .template cast(); + .template cast(); auto res_float32_tensor_type = op.getResult().getType().template cast(); auto rhs_float32_tensor_type = @@ -481,7 +511,7 @@ Value CreateZeroPointPartialOffset(OpBuilder &builder, Location loc, // Calculate the output tensor shape. This is input tensor dims minus // contracting dims. - auto ranked_tensor = tensor.getType().dyn_cast(); + auto ranked_tensor = tensor.getType().cast(); llvm::SmallVector output_dims; for (int64_t i = 0; i < ranked_tensor.getRank(); ++i) { if (absl::c_count(reduction_dims, i) == 0) { @@ -492,7 +522,7 @@ Value CreateZeroPointPartialOffset(OpBuilder &builder, Location loc, // Convert input tensor to output type since mhlo::Reduce only supports same // element type for input/output. tensor = builder.create( - loc, tensor.getType().dyn_cast().clone(output_element_type), + loc, tensor.getType().cast().clone(output_element_type), tensor); auto reducer_tensor_type = RankedTensorType::get({}, output_element_type); @@ -592,7 +622,7 @@ Value BroadcastZpContribution(OpBuilder &builder, Location loc, // zero-point-offset tensor to the final output tensor, and then do the // broadcast. auto zp_contribution_rank = - zp_contribution.getType().dyn_cast().getRank(); + zp_contribution.getType().cast().getRank(); llvm::SmallVector broadcast_dims; broadcast_dims.resize(zp_contribution_rank, 0); // Result tensor will have batching dims first, then LHS result dims, then @@ -615,7 +645,7 @@ Value BroadcastZpContribution(OpBuilder &builder, Location loc, } // Use broadcast_in_dim or dyanmic_broadcast_in_dim based on input shape // dynamism. - if (zp_contribution.getType().dyn_cast().hasStaticShape()) { + if (zp_contribution.getType().cast().hasStaticShape()) { zp_contribution = builder.create( loc, output_tensor_type, zp_contribution, DenseIntElementsAttr::get( @@ -742,13 +772,13 @@ Value CreateDotLikeKernel(OpBuilder &builder, Location loc, DenseIntElementsAttr::get( RankedTensorType::get({}, builder.getI8Type()), {static_cast(getElementTypeOrSelf(op.getLhs().getType()) - .dyn_cast() + .cast() .getZeroPoint())})); // Convert Padding attributes from mhlo::Convolution to mhlo::Pad. Note that // Padding is applied for spatial dimensions [1...rank-1) only for // mhlo::Convolution. But mhlo::Pad require those for all dimensions. Hence // we add 0 to the beginning and end of the padding vectors. - int64_t rank = lhs.getType().dyn_cast().getRank(); + int64_t rank = lhs.getType().cast().getRank(); llvm::SmallVector padding_low(rank, 0), padding_high(rank, 0), padding_interior(rank, 0); for (int64_t i = 1; i < rank - 1; ++i) { @@ -786,15 +816,12 @@ LogicalResult matchAndRewriteDotLikeOp(DotLikeOp op, DotLikeOpAdaptor adaptor, ConversionPatternRewriter &rewriter) { // Lower Dot/DotGeneral UQ ops to DotGeneral int. // Assumes that operands and results are uq types. - auto lhs_element_quant_type = - getElementTypeOrSelf(op.getLhs().getType()) - .template dyn_cast(); - auto rhs_element_quant_type = - getElementTypeOrSelf(op.getRhs().getType()) - .template dyn_cast(); - auto res_element_quant_type = - getElementTypeOrSelf(op.getResult()) - .template dyn_cast(); + auto lhs_element_quant_type = getElementTypeOrSelf(op.getLhs().getType()) + .template dyn_cast(); + auto rhs_element_quant_type = getElementTypeOrSelf(op.getRhs().getType()) + .template dyn_cast(); + auto res_element_quant_type = getElementTypeOrSelf(op.getResult()) + .template dyn_cast(); Value lhs = adaptor.getLhs(); Value rhs = adaptor.getRhs(); auto res_int32_tensor_type = @@ -837,8 +864,7 @@ LogicalResult matchAndRewriteDotLikeOp(DotLikeOp op, DotLikeOpAdaptor adaptor, // Skip zp_offset if it is 0. if (zp_offset) { auto zp_offset_float32_tensor_type = - zp_offset.getType().dyn_cast().clone( - rewriter.getF32Type()); + zp_offset.getType().cast().clone(rewriter.getF32Type()); zp_offset = rewriter.create( op->getLoc(), zp_offset_float32_tensor_type, zp_offset); zp_offset = rewriter.create( @@ -867,15 +893,12 @@ template FailureOr IsDotLikeOpHybrid(DotLikeOp op) { // Checks whether a dot-like op is hybrid by looking at input/output types. // Returns failure() when the type is not supported. - auto lhs_element_quant_type = - getElementTypeOrSelf(op.getLhs().getType()) - .template dyn_cast(); - auto rhs_element_quant_type = - getElementTypeOrSelf(op.getRhs().getType()) - .template dyn_cast(); - auto res_element_quant_type = - getElementTypeOrSelf(op.getResult()) - .template dyn_cast(); + auto lhs_element_quant_type = getElementTypeOrSelf(op.getLhs().getType()) + .template dyn_cast(); + auto rhs_element_quant_type = getElementTypeOrSelf(op.getRhs().getType()) + .template dyn_cast(); + auto res_element_quant_type = getElementTypeOrSelf(op.getResult()) + .template dyn_cast(); if (lhs_element_quant_type && rhs_element_quant_type && res_element_quant_type) { return false; @@ -996,8 +1019,7 @@ bool IsConvNDHWC(const mhlo::ConvDimensionNumbersAttr &dims) { FailureOr VerifyConvolutionOp(mhlo::ConvolutionOp op) { // RHS (weight) must have zero zp. auto rhs_element_quant_type = - getElementTypeOrSelf(op.getRhs().getType()) - .template dyn_cast(); + getElementTypeOrSelf(op.getRhs().getType()).cast(); if (rhs_element_quant_type.getZeroPoint() != 0) { op->emitError("RHS UQ type must have zero zp."); return failure(); @@ -1074,15 +1096,15 @@ class ConvertGenericOp : public ConversionPattern { // Check that all operands and result uq types are the same. llvm::SmallVector uq_types; for (auto result_type : op->getResultTypes()) { - auto type = getElementTypeOrSelf(result_type) - .dyn_cast(); + auto type = + getElementTypeOrSelf(result_type).dyn_cast(); if (type) { uq_types.push_back(type); } } for (auto operand : op->getOperands()) { auto type = getElementTypeOrSelf(operand.getType()) - .dyn_cast(); + .dyn_cast(); if (type) { uq_types.push_back(type); } @@ -1097,15 +1119,7 @@ class ConvertGenericOp : public ConversionPattern { // type otherwise. llvm::SmallVector new_result_types; for (auto result_type : op->getResultTypes()) { - if (getElementTypeOrSelf(result_type) - .isa()) { - new_result_types.push_back(result_type.cast().clone( - getElementTypeOrSelf(result_type) - .cast() - .getStorageType())); - } else { - new_result_types.push_back(result_type); - } + new_result_types.push_back(GetQuantStorageType(result_type)); } OperationState state(op->getLoc(), op->getName().getStringRef(), operands, @@ -1120,73 +1134,78 @@ class ConvertGenericOp : public ConversionPattern { class UQTypeConverter : public TypeConverter { public: UQTypeConverter() { - addConversion([](Type type) -> Type { - auto to_legal_type = [](Type type) { - if (auto uq_type = dyn_cast(type)) { - return uq_type.getStorageType(); - } - return type; - }; - if (auto shaped = type.dyn_cast()) { - return shaped.clone(to_legal_type(shaped.getElementType())); - } else { - return to_legal_type(type); - } - }); + addConversion([](Type type) -> Type { return GetQuantStorageType(type); }); } }; -// Performs conversion of MHLO quant ops to primitive ops. -void ConvertMHLOQuantToInt::runOnOperation() { - Operation *op = getOperation(); - MLIRContext *context = op->getContext(); - RewritePatternSet patterns(context); - - // Populate MHLO quant ops conversion patterns. - patterns.add(context); - - // uq->int convert patterns for func.func and func.return. - UQTypeConverter converter; - populateFunctionOpInterfaceTypeConversionPattern(patterns, - converter); - populateReturnOpTypeConversionPattern(patterns, converter); - - ConversionTarget target(*op->getContext()); - auto is_legal = [&converter](Operation *op) { return converter.isLegal(op); }; - target.addDynamicallyLegalDialect(is_legal); - target.addDynamicallyLegalDialect(is_legal); - target.addDynamicallyLegalDialect( - [&converter](Operation *op) { - if (auto func = dyn_cast(op)) { - return converter.isSignatureLegal(func.getFunctionType()); - } - return converter.isLegal(op); - }); - - LogicalResult result = - applyPartialConversion(op, target, std::move(patterns)); - if (failed(result)) { - signalPassFailure(); +#define GEN_PASS_DEF_CONVERTMHLOQUANTTOINT +#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/bridge/passes.h.inc" + +class ConvertMHLOQuantToInt + : public impl::ConvertMHLOQuantToIntBase { + public: + ConvertMHLOQuantToInt() = default; + ConvertMHLOQuantToInt(const ConvertMHLOQuantToInt &) {} + + explicit ConvertMHLOQuantToInt(bool legalize_chlo) { + legalize_chlo_ = legalize_chlo; } - // Legalize CHLO if needed. - if (!legalize_chlo_) return; - RewritePatternSet patterns_2(context); + // Performs conversion of MHLO quant ops to primitive ops. + void runOnOperation() override { + Operation *op = getOperation(); + MLIRContext *context = op->getContext(); + RewritePatternSet patterns(context); + + // Populate MHLO quant ops conversion patterns. + patterns.add( + context); + + // uq->int convert patterns for func.func and func.return. + UQTypeConverter converter; + populateFunctionOpInterfaceTypeConversionPattern(patterns, + converter); + populateReturnOpTypeConversionPattern(patterns, converter); + + ConversionTarget target(*op->getContext()); + auto is_legal = [&converter](Operation *op) { + return converter.isLegal(op); + }; + target.addDynamicallyLegalDialect(is_legal); + target.addDynamicallyLegalDialect(is_legal); + target.addDynamicallyLegalDialect( + [&converter](Operation *op) { + if (auto func = dyn_cast(op)) { + return converter.isSignatureLegal(func.getFunctionType()); + } + return converter.isLegal(op); + }); + + LogicalResult result = + applyPartialConversion(op, target, std::move(patterns)); + if (failed(result)) { + signalPassFailure(); + } + + // Legalize CHLO if needed. + if (!legalize_chlo_) return; + RewritePatternSet patterns_2(context); - chlo::populateDecomposeChloPatterns(context, &patterns_2); - chlo::populateChloBroadcastingPatterns(context, &patterns_2); + chlo::populateDecomposeChloPatterns(context, &patterns_2); + chlo::populateChloBroadcastingPatterns(context, &patterns_2); - ConversionTarget target_2 = - mhlo::GetDefaultLegalConversionTargets(*op->getContext(), legalize_chlo_); + ConversionTarget target_2 = mhlo::GetDefaultLegalConversionTargets( + *op->getContext(), legalize_chlo_); - result = applyPartialConversion(op, target_2, std::move(patterns_2)); - if (failed(result)) { - signalPassFailure(); + result = applyPartialConversion(op, target_2, std::move(patterns_2)); + if (failed(result)) { + signalPassFailure(); + } } -} +}; } // namespace diff --git a/tensorflow/compiler/mlir/quantization/stablehlo/passes/bridge/convert_tf_quant_to_mhlo_int_test.cc b/tensorflow/compiler/mlir/quantization/stablehlo/passes/bridge/convert_tf_quant_to_mhlo_int_test.cc index 9c2b88abddb8a5..dc686d8bf0fb35 100644 --- a/tensorflow/compiler/mlir/quantization/stablehlo/passes/bridge/convert_tf_quant_to_mhlo_int_test.cc +++ b/tensorflow/compiler/mlir/quantization/stablehlo/passes/bridge/convert_tf_quant_to_mhlo_int_test.cc @@ -13,6 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +#include #include #include #include @@ -21,6 +22,7 @@ limitations under the License. #include #include "absl/log/check.h" +#include "absl/random/random.h" #include "absl/status/status.h" #include "absl/strings/string_view.h" #include "absl/types/span.h" @@ -29,6 +31,8 @@ limitations under the License. #include "mlir/Dialect/Quant/QuantOps.h" // from @llvm-project #include "mlir/IR/BuiltinOps.h" // from @llvm-project #include "mlir/IR/DialectRegistry.h" // from @llvm-project +#include "mlir/IR/OwningOpRef.h" // from @llvm-project +#include "mlir/IR/Value.h" // from @llvm-project #include "mlir/Parser/Parser.h" // from @llvm-project #include "mlir/Pass/PassManager.h" // from @llvm-project #include "mlir/Support/LogicalResult.h" // from @llvm-project @@ -65,6 +69,7 @@ class ConvertTfQuantToMhloIntTest : public ::testing::Test { dialects.insert(); ctx_ = std::make_unique(dialects); + ctx_->loadAllAvailableDialects(); // Create a CPU client with 1 device. TF_ASSERT_OK_AND_ASSIGN( @@ -74,11 +79,10 @@ class ConvertTfQuantToMhloIntTest : public ::testing::Test { CHECK(device_); } - // Evaluate return value of a function using TF kernel. - // This assumes that the module op has only 1 function and it has TF ops only. - absl::StatusOr> EvaluateTfFunction( + absl::StatusOr> ReplaceFuncArgsByConstant( absl::string_view program, - absl::Span arguments) { + absl::Span arguments, + bool use_mhlo_const = false) { auto module_op = parseSourceString(program, ctx_.get()); CHECK(module_op); auto func_op = llvm::dyn_cast( @@ -90,9 +94,9 @@ class ConvertTfQuantToMhloIntTest : public ::testing::Test { return absl::InternalError("Input argument has wrong size"); } - // Convert input xla::Literal arguments to tf.Const, this allows using + // Convert input xla::Literal arguments to constants, this allows using // constant folding to evaluate function return value. - mlir::OpBuilder builder(func_op); + mlir::OpBuilder builder(ctx_.get()); for (int i = 0; i < arguments.size(); ++i) { const xla::Literal* const xla_literal = arguments[i]; tensorflow::TensorShape shape; @@ -109,15 +113,38 @@ class ConvertTfQuantToMhloIntTest : public ::testing::Test { xla_literal->element_count()); TF_ASSIGN_OR_RETURN(auto attrs, tensorflow::ConvertTensor(tensor, &builder)); - auto cst = builder.create(func_op->getLoc(), attrs); + builder.setInsertionPoint( + &func_op.getFunctionBody().getBlocks().front().front()); + // Use mhlo.Constant when it is consumed by the lowering passes since they + // can't lower tf.Const. + Value cst; + if (use_mhlo_const) { + cst = builder.create(func_op->getLoc(), attrs); + } else { + cst = builder.create(func_op->getLoc(), attrs); + } func_op.getArgument(i).replaceAllUsesWith(cst); } + return module_op; + } + // Evaluate return value of a function using TF kernel. + // This assumes that the module op has only 1 function and it has TF ops only. + absl::StatusOr> EvaluateTfFunction( + absl::string_view program, + absl::Span arguments) { + TF_ASSIGN_OR_RETURN(auto module_op, + ReplaceFuncArgsByConstant(program, arguments)); // Constant fold the func.Return op's producer op to evaluate the return // value. The evaluation will use TF kernels. // This assumes that func.Return is the last op in the function and it // returns only 1 value. - auto& return_op = func_op.getFunctionBody().getBlocks().back().back(); + auto& return_op = llvm::dyn_cast( + *module_op->getBodyRegion().getOps().begin()) + .getFunctionBody() + .getBlocks() + .back() + .back(); if (!llvm::isa(return_op) || return_op.getNumOperands() != 1) { return absl::InternalError( @@ -150,17 +177,19 @@ class ConvertTfQuantToMhloIntTest : public ::testing::Test { } absl::StatusOr> CompileProgram( - absl::string_view program) { - // Parse the program. - auto module_op = parseSourceString(program, ctx_.get()); - CHECK(module_op); + absl::string_view program, + absl::Span arguments) { + // Replace args by mhlo.constant since the lowering passes can't lower + // tf.Const. + TF_ASSIGN_OR_RETURN( + auto module_op, + ReplaceFuncArgsByConstant(program, arguments, /*use_mhlo_const=*/true)); + // Run the Convert TF Quant Types, TF Quant -> MHLO Quant and MHLO Quant -> // MHLO int passes. PassManager pm(module_op->getContext()); pm.addNestedPass(CreateConvertTFQuantTypesPass()); - pm.addNestedPass(CreateConvertTFQuantOpsToMHLOPass()); - pm.addNestedPass( - stablehlo::createConvertMHLOQuantToIntPass(false)); + AddQuantizationLoweringPasses(pm); CHECK(succeeded(pm.run(module_op.get()))); // Compile the program. return pjrt_client_->Compile(*module_op, xla::CompileOptions{}); @@ -190,206 +219,366 @@ class ConvertTfQuantToMhloIntTest : public ::testing::Test { void ExecuteAndCompareResultsWithTfKernel( absl::string_view program, absl::Span arguments, - float error_tolerance = 0.1) { - TF_ASSERT_OK_AND_ASSIGN(auto executable, this->CompileProgram(program)); + std::optional tf_program = std::nullopt, + double error_tolerance = 0.1) { + // Expected result is calculated by evaluating using TF kernels. In some + // cases, TF kernel behaves differently from lowered graph (e.g. Hybrid + // ops). So we optionally use a different graph to calculate the expected + // result. + TF_ASSERT_OK_AND_ASSIGN( + auto expected, + this->EvaluateTfFunction( + (tf_program.has_value() ? *tf_program : program), arguments)); + TF_ASSERT_OK_AND_ASSIGN(auto executable, + this->CompileProgram(program, arguments)); TF_ASSERT_OK_AND_ASSIGN( - auto result_literal, + auto result, this->ExecuteProgramAndReturnSingleResult(executable.get(), arguments)); - TF_ASSERT_OK_AND_ASSIGN(auto expected, - this->EvaluateTfFunction(program, arguments)); - EXPECT_TRUE(xla::LiteralTestUtil::Near(*expected, *result_literal, + // Convert to double for comparison. This is needed for comparing integers + // since it LiteralTestUtil asserts different integers even if it is within + // error_spec. + TF_ASSERT_OK_AND_ASSIGN(auto expected_double, expected->Convert(xla::F64)) + TF_ASSERT_OK_AND_ASSIGN(auto result_double, result->Convert(xla::F64)) + EXPECT_TRUE(xla::LiteralTestUtil::Near(expected_double, result_double, xla::ErrorSpec(error_tolerance))); } + absl::StatusOr CreateRandomF32Literal( + absl::Span dims, float min = -100, float max = 100) { + TF_ASSIGN_OR_RETURN(auto shape, + xla::ShapeUtil::MakeValidatedShape(xla::F32, dims)); + return xla::LiteralUtil::CreateLiteralWithGenerator( + shape, [this, min, max](absl::Span dims) -> float { + return absl::Uniform(bitgen_, min, max); + }); + } + + absl::StatusOr CreateRandomI8Literal( + absl::Span dims, int8_t min = -128, int8_t max = 127) { + TF_ASSIGN_OR_RETURN(auto shape, + xla::ShapeUtil::MakeValidatedShape(xla::S8, dims)); + return xla::LiteralUtil::CreateLiteralWithGenerator( + shape, [this, min, max](absl::Span dims) -> int8_t { + return absl::Uniform(bitgen_, min, max); + }); + } + + absl::StatusOr CreateRandomI32Literal( + absl::Span dims, int32_t min = -128, int32_t max = 127) { + TF_ASSIGN_OR_RETURN(auto shape, + xla::ShapeUtil::MakeValidatedShape(xla::S32, dims)); + return xla::LiteralUtil::CreateLiteralWithGenerator( + shape, [this, min, max](absl::Span dims) -> int32_t { + return absl::Uniform(bitgen_, min, max); + }); + } + std::unique_ptr ctx_; std::unique_ptr pjrt_client_; xla::PjRtDevice* device_; + absl::BitGen bitgen_; }; TEST_F(ConvertTfQuantToMhloIntTest, UniformQuantizeAndDequantize) { constexpr absl::string_view kProgram = R"mlir( -func.func @main(%arg0: tensor<4xf32>) -> tensor<4xf32> { - %scale = "tf.Const"() { value = dense<10.0> : tensor } : () - -> tensor +func.func @main(%arg0: tensor<10xf32>) -> tensor<10xf32> { + %scale = "tf.Const"() { value = dense<0.347> : tensor } : () -> tensor %zp = "tf.Const"() { value = dense<3> : tensor } : () -> tensor %0 = "tf.UniformQuantize"(%arg0, %scale, %zp) { quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64 - } : (tensor<4xf32>, tensor, tensor) -> tensor<4x!tf_type.qint8> + } : (tensor<10xf32>, tensor, tensor) -> tensor<10x!tf_type.qint8> %1 = "tf.UniformDequantize"(%0, %scale, %zp) { quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64 - } : (tensor<4x!tf_type.qint8>, tensor, tensor) -> tensor<4xf32> - return %1 : tensor<4xf32> + } : (tensor<10x!tf_type.qint8>, tensor, tensor) -> tensor<10xf32> + return %1 : tensor<10xf32> })mlir"; - auto arg0 = - xla::LiteralUtil::CreateR1({100.0f, 20000.0f, -2409.0f, -25.1f}); - ExecuteAndCompareResultsWithTfKernel(kProgram, {&arg0}); + TF_ASSERT_OK_AND_ASSIGN(auto arg0, CreateRandomF32Literal({10})); + // error_tolerance is set to be slightly > scale because different rounding + // implementations for UniformQuantize in TF kernel and the lowering passes + // may cause +/-1 differences. + ExecuteAndCompareResultsWithTfKernel( + kProgram, {&arg0}, /*tf_program=*/std::nullopt, /*error_tolerance=*/0.35); +} + +TEST_F(ConvertTfQuantToMhloIntTest, UniformQuantizePerChannel) { + constexpr absl::string_view kProgram = R"mlir( +func.func @main( + %arg0: tensor<10x10xf32>, %scale: tensor<10xf32>, %zp: tensor<10xi32> + ) -> tensor<10x10xi8> { + %0 = "tf.UniformQuantize"(%arg0, %scale, %zp) { + quantization_axis = 1 : i64, + quantization_min_val = -128 : i64, + quantization_max_val = 127 : i64 + } : (tensor<10x10xf32>, tensor<10xf32>, tensor<10xi32>) -> tensor<10x10x!tf_type.qint8> + %1 = "tf.Cast"(%0) {} : (tensor<10x10x!tf_type.qint8>) -> tensor<10x10xi8> + return %1 : tensor<10x10xi8> +})mlir"; + TF_ASSERT_OK_AND_ASSIGN(auto arg0, CreateRandomF32Literal({10, 10})); + TF_ASSERT_OK_AND_ASSIGN( + auto scale, CreateRandomF32Literal({10}, /*min=*/0.0001, /*max=*/2)); + TF_ASSERT_OK_AND_ASSIGN(auto zp, CreateRandomI32Literal({10})); + // Different rounding implementations for UniformQuantize in TF kernel and the + // lowering passes may cause +/-1 differences. + ExecuteAndCompareResultsWithTfKernel(kProgram, {&arg0, &scale, &zp}, + /*tf_program=*/std::nullopt, + /*error_tolerance=*/1.0); +} + +TEST_F(ConvertTfQuantToMhloIntTest, UniformDequantizePerChannel) { + constexpr absl::string_view kProgram = R"mlir( +func.func @main( + %arg0: tensor<10x10xi8>, %scale: tensor<10xf32>, %zp: tensor<10xi32> + ) -> tensor<10x10xf32> { + %0 = "tf.Cast"(%arg0) {} : (tensor<10x10xi8>) -> tensor<10x10x!tf_type.qint8> + %1 = "tf.UniformDequantize"(%0, %scale, %zp) { + quantization_axis = 1 : i64, + quantization_min_val = -128 : i64, + quantization_max_val = 127 : i64 + } : (tensor<10x10x!tf_type.qint8>, tensor<10xf32>, tensor<10xi32>) -> tensor<10x10xf32> + return %1 : tensor<10x10xf32> +})mlir"; + TF_ASSERT_OK_AND_ASSIGN(auto arg0, CreateRandomI8Literal({10, 10})); + TF_ASSERT_OK_AND_ASSIGN( + auto scale, CreateRandomF32Literal({10}, /*min=*/0.0001, /*max=*/2)); + TF_ASSERT_OK_AND_ASSIGN(auto zp, CreateRandomI32Literal({10})); + ExecuteAndCompareResultsWithTfKernel(kProgram, {&arg0, &scale, &zp}); } TEST_F(ConvertTfQuantToMhloIntTest, UniformQuantizeConvolution) { constexpr absl::string_view kProgram = R"mlir( -func.func @main(%input: tensor<1x2x2x1xf32>, %filter: tensor<2x1x1x1xf32>) -> tensor<1x2x2x1xf32> { - %input_scale = "tf.Const"() { value = dense<7.3> : tensor } : () - -> tensor - %input_zp = "tf.Const"() { value = dense<-45> : tensor } : () -> tensor - %filter_scale = "tf.Const"() { value = dense<0.047> : tensor } : () - -> tensor - %filter_zp = "tf.Const"() { value = dense<0> : tensor } : () -> tensor - %accum_scale = "tf.Const"() { value = dense<0.3431> : tensor } : () - -> tensor - %accum_zp = "tf.Const"() { value = dense<0> : tensor } : () -> tensor - %quant_input = "tf.UniformQuantize"(%input, %input_scale, %input_zp) { - Tin = "tfdtype$DT_FLOAT", Tout = "tfdtype$DT_QINT8", - attr_map = "", quantization_axis = -1 : i64, quantization_max_val = 127 : i64, - quantization_min_val = -128 : i64 - } : (tensor<1x2x2x1xf32>, tensor, tensor) -> tensor<1x2x2x1x!tf_type.qint8> - %quant_filter = "tf.UniformQuantize"(%filter, %filter_scale, %filter_zp) { - Tin = "tfdtype$DT_FLOAT", Tout = "tfdtype$DT_QINT8", - attr_map = "", quantization_axis = -1 : i64, - quantization_max_val = 127 : i64, quantization_min_val = -128 : i64 - } : (tensor<2x1x1x1xf32>, tensor, tensor) -> tensor<2x1x1x1x!tf_type.qint8> - %0 = "tf.UniformQuantizedConvolution"( - %quant_input, %quant_filter, %input_scale, %input_zp, - %filter_scale, %filter_zp, %accum_scale, %accum_zp - ) { - Tin = "tfdtype$DT_QINT8", Tout = "tfdtype$DT_QINT32", - attr_map = "", batch_group_count = 1 : i64, - dimension_numbers = "\10\03\1A\02\01\02 \02(\032\02\00\01@\03J\02\01\02", - explicit_padding = [], feature_group_count = 1 : i64, lhs_dilation = [1, 1], - lhs_quantization_axis = -1 : i64, lhs_quantization_max_val = 127 : i64, - lhs_quantization_min_val = -128 : i64, output_quantization_axis = -1 : i64, - output_quantization_max_val = 2147483647 : i64, - output_quantization_min_val = -2147483648 : i64, padding = "SAME", - rhs_dilation = [1, 1], rhs_quantization_axis = -1 : i64, - rhs_quantization_max_val = 127 : i64, rhs_quantization_min_val = -128 : i64, - window_strides = [1, 1] - } : (tensor<1x2x2x1x!tf_type.qint8>, tensor<2x1x1x1x!tf_type.qint8>, - tensor, tensor, tensor, tensor, tensor, tensor - ) -> tensor<1x2x2x1x!tf_type.qint32> - %output = "tf.UniformDequantize"(%0, %accum_scale, %accum_zp) { - quantization_axis = -1 : i64, quantization_min_val = -128 : i64, - quantization_max_val = 127 : i64 - } : (tensor<1x2x2x1x!tf_type.qint32>, tensor, tensor) -> tensor<1x2x2x1xf32> - return %output : tensor<1x2x2x1xf32> +func.func @main(%input: tensor<1x9x9x9xi8>, %filter: tensor<3x3x9x10xi8>) -> tensor<1x9x9x10xi32> { + %input_scale = "tf.Const"() { value = dense<2.0> : tensor } : () -> tensor + %input_zp = "tf.Const"() { value = dense<-10> : tensor } : () -> tensor + %filter_scale = "tf.Const"() { value = dense<0.5> : tensor } : () -> tensor + %filter_zp = "tf.Const"() { value = dense<0> : tensor } : () -> tensor + %accum_scale = "tf.Const"() { value = dense<1.0> : tensor } : () -> tensor + %accum_zp = "tf.Const"() { value = dense<0> : tensor } : () -> tensor + %quant_input = "tf.Cast"(%input) {} : (tensor<1x9x9x9xi8>) -> + tensor<1x9x9x9x!tf_type.qint8> + %quant_filter = "tf.Cast"(%filter) {} : (tensor<3x3x9x10xi8>) -> + tensor<3x3x9x10x!tf_type.qint8> + %0 = "tf.UniformQuantizedConvolution"( + %quant_input, %quant_filter, %input_scale, %input_zp, + %filter_scale, %filter_zp, %accum_scale, %accum_zp + ) { + Tin = "tfdtype$DT_QINT8", Tout = "tfdtype$DT_QINT32", + attr_map = "", batch_group_count = 1 : i64, + dimension_numbers = "\10\03\1A\02\01\02 \02(\032\02\00\01@\03J\02\01\02", + explicit_padding = [], feature_group_count = 1 : i64, lhs_dilation = [1, 1], + lhs_quantization_axis = -1 : i64, lhs_quantization_max_val = 127 : i64, + lhs_quantization_min_val = -128 : i64, output_quantization_axis = -1 : i64, + output_quantization_max_val = 2147483647 : i64, + output_quantization_min_val = -2147483648 : i64, padding = "SAME", + rhs_dilation = [1, 1], rhs_quantization_axis = -1 : i64, + rhs_quantization_max_val = 127 : i64, rhs_quantization_min_val = -128 : i64, + window_strides = [1, 1] + } : (tensor<1x9x9x9x!tf_type.qint8>, tensor<3x3x9x10x!tf_type.qint8>, + tensor, tensor, tensor, tensor, tensor, tensor + ) -> tensor<1x9x9x10x!tf_type.qint32> + %output = "tf.Cast"(%0) {} : (tensor<1x9x9x10x!tf_type.qint32>) -> tensor<1x9x9x10xi32> + return %output : tensor<1x9x9x10xi32> })mlir"; - auto input = xla::LiteralUtil::CreateR4( - {{{{14.f}, {-100.f}}, {{-200.f}, {350.f}}}}); - auto filter = xla::LiteralUtil::CreateR4({{{{4.1f}}}, {{{-2.f}}}}); + TF_ASSERT_OK_AND_ASSIGN(auto input, CreateRandomI8Literal({1, 9, 9, 9})); + TF_ASSERT_OK_AND_ASSIGN(auto filter, CreateRandomI8Literal({3, 3, 9, 10})); ExecuteAndCompareResultsWithTfKernel(kProgram, {&input, &filter}); } TEST_F(ConvertTfQuantToMhloIntTest, UniformQuantizeConvolutionHybrid) { + constexpr absl::string_view kTfProgram = R"mlir( +func.func @main(%input: tensor<2x10x10x10xf32>, %filter: tensor<3x3x10x20xi8>) -> tensor<2x10x10x20xf32> { + %filter_scale = "tf.Const"() { value = dense<0.047> : tensor } : () -> tensor + %filter_zp = "tf.Const"() { value = dense<0> : tensor } : () -> tensor + %quant_filter = "tf.Cast"(%filter) {} : (tensor<3x3x10x20xi8>) -> + tensor<3x3x10x20x!tf_type.qint8> + %filter_new = "tf.UniformDequantize"(%quant_filter, %filter_scale, %filter_zp) { + quantization_axis = -1 : i64, quantization_min_val = -128 : i64, + quantization_max_val = 127 : i64 + } : ( + tensor<3x3x10x20x!tf_type.qint8>, tensor, tensor + ) -> tensor<3x3x10x20xf32> + %0 = "tf.Conv2D"(%input, %filter_new) { + Tin = "tfdtype$DT_FLOAT", Tout = "tfdtype$DT_FLOAT", + attr_map = "", batch_group_count = 1 : i64, + explicit_padding = [], feature_group_count = 1 : i64, lhs_dilation = [1, 1], + padding = "SAME", rhs_dilation = [1, 1], strides = [1, 1, 1, 1] + } : (tensor<2x10x10x10xf32>, tensor<3x3x10x20xf32>) -> tensor<2x10x10x20xf32> + return %0 : tensor<2x10x10x20xf32> +})mlir"; constexpr absl::string_view kProgram = R"mlir( -func.func @main(%input: tensor<1x2x2x1xf32>, %filter: tensor<2x1x1x1xf32>) -> tensor<1x2x2x1xf32> { - %filter_scale = "tf.Const"() { value = dense<0.047> : tensor } : () - -> tensor - %filter_zp = "tf.Const"() { value = dense<0> : tensor } : () -> tensor - %quant_filter = "tf.UniformQuantize"(%filter, %filter_scale, %filter_zp) { - Tin = "tfdtype$DT_FLOAT", Tout = "tfdtype$DT_QINT8", - attr_map = "", quantization_axis = -1 : i64, - quantization_max_val = 127 : i64, quantization_min_val = -128 : i64 - } : (tensor<2x1x1x1xf32>, tensor, tensor) -> tensor<2x1x1x1x!tf_type.qint8> - %0 = "tf.UniformQuantizedConvolutionHybrid"( - %input, %quant_filter, %filter_scale, %filter_zp - ) { - Tin = "tfdtype$DT_QINT8", Tout = "tfdtype$DT_FLOAT", - attr_map = "", batch_group_count = 1 : i64, - dimension_numbers = "\10\03\1A\02\01\02 \02(\032\02\00\01@\03J\02\01\02", - explicit_padding = [], feature_group_count = 1 : i64, lhs_dilation = [1, 1], - padding = "SAME", rhs_dilation = [1, 1], rhs_quantization_axis = -1 : i64, - rhs_quantization_max_val = 127 : i64, rhs_quantization_min_val = -128 : i64, - window_strides = [1, 1] - } : (tensor<1x2x2x1xf32>, tensor<2x1x1x1x!tf_type.qint8>, - tensor, tensor) -> tensor<1x2x2x1xf32> - return %0 : tensor<1x2x2x1xf32> +func.func @main(%input: tensor<2x10x10x10xf32>, %filter: tensor<3x3x10x20xi8>) -> tensor<2x10x10x20xf32> { + %filter_scale = "tf.Const"() { value = dense<0.047> : tensor } : () -> tensor + %filter_zp = "tf.Const"() { value = dense<0> : tensor } : () -> tensor + %quant_filter = "tf.Cast"(%filter) {} : (tensor<3x3x10x20xi8>) -> tensor<3x3x10x20x!tf_type.qint8> + %0 = "tf.UniformQuantizedConvolutionHybrid"( + %input, %quant_filter, %filter_scale, %filter_zp + ) { + Tin = "tfdtype$DT_QINT8", Tout = "tfdtype$DT_FLOAT", + attr_map = "", batch_group_count = 1 : i64, + dimension_numbers = "\10\03\1A\02\01\02 \02(\032\02\00\01@\03J\02\01\02", + explicit_padding = [], feature_group_count = 1 : i64, lhs_dilation = [1, 1], + padding = "SAME", rhs_dilation = [1, 1], rhs_quantization_axis = -1 : i64, + rhs_quantization_max_val = 127 : i64, rhs_quantization_min_val = -128 : i64, + window_strides = [1, 1] + } : (tensor<2x10x10x10xf32>, tensor<3x3x10x20x!tf_type.qint8>, + tensor, tensor) -> tensor<2x10x10x20xf32> + return %0 : tensor<2x10x10x20xf32> })mlir"; - auto input = xla::LiteralUtil::CreateR4( - {{{{14.f}, {-100.f}}, {{-200.f}, {350.f}}}}); - auto filter = xla::LiteralUtil::CreateR4({{{{4.1f}}}, {{{-2.f}}}}); - // The large tolerance here is expected because - // tf.UniformQuantizedConvolutionHybrid does DRQ. But StableHLO hybrid ops - // does weight-only. - ExecuteAndCompareResultsWithTfKernel(kProgram, {&input, &filter}, - /*error_tolerance=*/5.0); + TF_ASSERT_OK_AND_ASSIGN(auto input, CreateRandomF32Literal({2, 10, 10, 10})); + TF_ASSERT_OK_AND_ASSIGN(auto filter, CreateRandomI8Literal({3, 3, 10, 20})); + // TF kernels for UniformQuantizedConvolutionHybrid does DRQ. But StableHLO + // hybrid ops does weight-only. So we use a different TF graph for evaluating + // expected weight-only quantized results. + ExecuteAndCompareResultsWithTfKernel(kProgram, {&input, &filter}, kTfProgram); } TEST_F(ConvertTfQuantToMhloIntTest, UniformQuantizeDot) { constexpr absl::string_view kProgram = R"mlir( -func.func @main(%input: tensor<1x2xf32>, %filter: tensor<2x3xf32>) -> tensor<1x3xf32> { - %input_scale = "tf.Const"() { value = dense<0.588> : tensor } : () - -> tensor - %input_zp = "tf.Const"() { value = dense<42> : tensor } : () -> tensor - %filter_scale = "tf.Const"() { value = dense<0.0235> : tensor } : () - -> tensor - %filter_zp = "tf.Const"() { value = dense<0> : tensor } : () -> tensor - %accum_scale = "tf.Const"() { value = dense<0.013818> : tensor } : () - -> tensor - %accum_zp = "tf.Const"() { value = dense<0> : tensor } : () -> tensor - %quant_input = "tf.UniformQuantize"(%input, %input_scale, %input_zp) { - Tin = "tfdtype$DT_FLOAT", Tout = "tfdtype$DT_QINT8", attr_map = "", - quantization_axis = -1 : i64, quantization_max_val = 127 : i64, - quantization_min_val = -128 : i64 - } : (tensor<1x2xf32>, tensor, tensor) -> tensor<1x2x!tf_type.qint8> - %quant_filter = "tf.UniformQuantize"(%filter, %filter_scale, %filter_zp) { - Tin = "tfdtype$DT_FLOAT", Tout = "tfdtype$DT_QINT8", attr_map = "", - quantization_axis = -1 : i64, quantization_max_val = 127 : i64, - quantization_min_val = -128 : i64 - } : (tensor<2x3xf32>, tensor, tensor) -> tensor<2x3x!tf_type.qint8> - %0 = "tf.UniformQuantizedDot"( - %quant_input, %quant_filter, %input_scale, %input_zp, %filter_scale, - %filter_zp, %accum_scale, %accum_zp - ) { - Tin = "tfdtype$DT_QINT8", Tout = "tfdtype$DT_QINT32", attr_map = "", - device = "", lhs_quantization_axis = -1 : i64, - lhs_quantization_max_val = 127 : i64, lhs_quantization_min_val = -128 : i64, - output_quantization_axis = -1 : i64, output_quantization_max_val = 2147483647 : i64, - output_quantization_min_val = -2147483648 : i64, rhs_quantization_axis = -1 : i64, - rhs_quantization_max_val = 127 : i64, rhs_quantization_min_val = -128 : i64 - } : ( - tensor<1x2x!tf_type.qint8>, tensor<2x3x!tf_type.qint8>, tensor, - tensor, tensor, tensor, tensor, tensor - ) -> tensor<1x3x!tf_type.qint32> - %output = "tf.UniformDequantize"(%0, %accum_scale, %accum_zp) { - quantization_axis = -1 : i64, quantization_min_val = -128 : i64, - quantization_max_val = 127 : i64 - } : (tensor<1x3x!tf_type.qint32>, tensor, tensor) -> tensor<1x3xf32> - return %output : tensor<1x3xf32> +func.func @main(%input: tensor<8x9xi8>, %filter: tensor<9x10xi8>) -> tensor<8x10xi32> { + %input_scale = "tf.Const"() { value = dense<0.588> : tensor } : () -> tensor + %input_zp = "tf.Const"() { value = dense<42> : tensor } : () -> tensor + %filter_scale = "tf.Const"() { value = dense<0.0235> : tensor } : () -> tensor + %filter_zp = "tf.Const"() { value = dense<0> : tensor } : () -> tensor + %accum_scale = "tf.Const"() { value = dense<0.013818> : tensor } : () -> tensor + %accum_zp = "tf.Const"() { value = dense<0> : tensor } : () -> tensor + %quant_input = "tf.Cast"(%input) {} : (tensor<8x9xi8>) -> tensor<8x9x!tf_type.qint8> + %quant_filter = "tf.Cast"(%filter) {} : (tensor<9x10xi8>) -> tensor<9x10x!tf_type.qint8> + %0 = "tf.UniformQuantizedDot"( + %quant_input, %quant_filter, %input_scale, %input_zp, %filter_scale, + %filter_zp, %accum_scale, %accum_zp + ) { + Tin = "tfdtype$DT_QINT8", Tout = "tfdtype$DT_QINT32", attr_map = "", + device = "", lhs_quantization_axis = -1 : i64, + lhs_quantization_max_val = 127 : i64, + lhs_quantization_min_val = -128 : i64, + output_quantization_axis = -1 : i64, + output_quantization_max_val = 2147483647 : i64, + output_quantization_min_val = -2147483648 : i64, + rhs_quantization_axis = -1 : i64, + rhs_quantization_max_val = 127 : i64, + rhs_quantization_min_val = -128 : i64 + } : ( + tensor<8x9x!tf_type.qint8>, tensor<9x10x!tf_type.qint8>, tensor, + tensor, tensor, tensor, tensor, tensor + ) -> tensor<8x10x!tf_type.qint32> + %output = "tf.Cast"(%0) {} : (tensor<8x10x!tf_type.qint32>) -> tensor<8x10xi32> + return %output : tensor<8x10xi32> })mlir"; - auto input = xla::LiteralUtil::CreateR2({{50.f, -100.f}}); - auto filter = - xla::LiteralUtil::CreateR2({{1.f, 2.f, 3.f}, {-1.f, -3.f, 1.f}}); + TF_ASSERT_OK_AND_ASSIGN(auto input, CreateRandomI8Literal({8, 9})); + TF_ASSERT_OK_AND_ASSIGN(auto filter, CreateRandomI8Literal({9, 10})); ExecuteAndCompareResultsWithTfKernel(kProgram, {&input, &filter}); } TEST_F(ConvertTfQuantToMhloIntTest, UniformQuantizeDotHybrid) { + constexpr absl::string_view kTfProgram = R"mlir( +func.func @main(%input: tensor<8x9xf32>, %filter: tensor<9x10xi8>) -> tensor<8x10xf32> { + %filter_scale = "tf.Const"() { value = dense<0.0235> : tensor } : () -> tensor + %filter_zp = "tf.Const"() { value = dense<0> : tensor } : () -> tensor + %quant_filter = "tf.Cast"(%filter) {} : (tensor<9x10xi8>) -> tensor<9x10x!tf_type.qint8> + %filter_new = "tf.UniformDequantize"(%quant_filter, %filter_scale, %filter_zp) { + quantization_axis = -1 : i64, quantization_min_val = -128 : i64, + quantization_max_val = 127 : i64 + } : (tensor<9x10x!tf_type.qint8>, tensor, tensor) -> tensor<9x10xf32> + %0 = "tf.MatMul"(%input, %filter_new) { + } : (tensor<8x9xf32>, tensor<9x10xf32>) -> tensor<8x10xf32> + return %0 : tensor<8x10xf32> +})mlir"; constexpr absl::string_view kProgram = R"mlir( -func.func @main(%input: tensor<1x2xf32>, %filter: tensor<2x3xf32>) -> tensor<1x3xf32> { - %filter_scale = "tf.Const"() { value = dense<0.0235> : tensor } : () - -> tensor - %filter_zp = "tf.Const"() { value = dense<0> : tensor } : () -> tensor - %quant_filter = "tf.UniformQuantize"(%filter, %filter_scale, %filter_zp) { - Tin = "tfdtype$DT_FLOAT", Tout = "tfdtype$DT_QINT8", attr_map = "", - quantization_axis = -1 : i64, quantization_max_val = 127 : i64, - quantization_min_val = -128 : i64 - } : (tensor<2x3xf32>, tensor, tensor) -> tensor<2x3x!tf_type.qint8> - %0 = "tf.UniformQuantizedDotHybrid"( - %input, %quant_filter, %filter_scale, %filter_zp - ) { - Tin = "tfdtype$DT_QINT8", Tout = "tfdtype$DT_FLOAT", attr_map = "", - device = "", rhs_quantization_axis = -1 : i64, - rhs_quantization_max_val = 127 : i64, rhs_quantization_min_val = -128 : i64 - } : (tensor<1x2xf32>, tensor<2x3x!tf_type.qint8>, tensor, tensor) -> tensor<1x3xf32> - return %0 : tensor<1x3xf32> +func.func @main(%input: tensor<8x9xf32>, %filter: tensor<9x10xi8>) -> tensor<8x10xf32> { + %filter_scale = "tf.Const"() { value = dense<0.0235> : tensor } : () + -> tensor + %filter_zp = "tf.Const"() { value = dense<0> : tensor } : () -> tensor + %quant_filter = "tf.Cast"(%filter) {} : (tensor<9x10xi8>) -> tensor<9x10x!tf_type.qint8> + %0 = "tf.UniformQuantizedDotHybrid"( + %input, %quant_filter, %filter_scale, %filter_zp + ) { + Tin = "tfdtype$DT_QINT8", Tout = "tfdtype$DT_FLOAT", attr_map = "", + device = "", rhs_quantization_axis = -1 : i64, + rhs_quantization_max_val = 127 : i64, rhs_quantization_min_val = -128 : i64 + } : (tensor<8x9xf32>, tensor<9x10x!tf_type.qint8>, tensor, tensor) -> tensor<8x10xf32> + return %0 : tensor<8x10xf32> })mlir"; - auto input = xla::LiteralUtil::CreateR2({{50.f, -100.f}}); - auto filter = - xla::LiteralUtil::CreateR2({{1.f, 2.f, 3.f}, {-1.f, -3.f, 1.f}}); - ExecuteAndCompareResultsWithTfKernel(kProgram, {&input, &filter}); + TF_ASSERT_OK_AND_ASSIGN(auto input, CreateRandomF32Literal({8, 9})); + TF_ASSERT_OK_AND_ASSIGN(auto filter, CreateRandomI8Literal({9, 10})); + // TF kernels for UniformQuantizedDotHybrid does DRQ. But StableHLO hybrid ops + // does weight-only. So we use a different TF graph for evaluating expected + // weight-only quantized results. + ExecuteAndCompareResultsWithTfKernel(kProgram, {&input, &filter}, kTfProgram); +} + +TEST_F(ConvertTfQuantToMhloIntTest, UniformRequantize) { + constexpr absl::string_view kProgram = R"mlir( +func.func @main(%input: tensor<10xi8>) -> tensor<10xi8> { + %input_scale = "tf.Const"() { value = dense<0.2235> : tensor } : () -> tensor + %input_zp = "tf.Const"() { value = dense<-2> : tensor } : () -> tensor + %output_scale = "tf.Const"() { value = dense<0.11> : tensor } : () -> tensor + %output_zp = "tf.Const"() { value = dense<3> : tensor } : () -> tensor + %0 = "tf.Cast"(%input) {} : (tensor<10xi8>) -> tensor<10x!tf_type.qint8> + %1 = "tf.UniformRequantize"( + %0, %input_scale, %input_zp, %output_scale, %output_zp + ) { + Tin = "tfdtype$DT_QINT8", Tout = "tfdtype$DT_QINT8", attr_map = "", + device = "", input_quantization_axis = -1, + input_quantization_max_val = 127 : i64, + input_quantization_min_val = -128 : i64, + output_quantization_axis = -1 : i64, + output_quantization_max_val = 127 : i64, + output_quantization_min_val = -128 : i64 + } : ( + tensor<10x!tf_type.qint8>, tensor, tensor, tensor, + tensor + ) -> tensor<10x!tf_type.qint8> + %2 = "tf.Cast"(%1) {} : (tensor<10x!tf_type.qint8>) -> tensor<10xi8> + return %2 : tensor<10xi8> +})mlir"; + TF_ASSERT_OK_AND_ASSIGN(auto input, CreateRandomI8Literal({10})); + ExecuteAndCompareResultsWithTfKernel(kProgram, {&input}); +} + +TEST_F(ConvertTfQuantToMhloIntTest, UniformQuantizeAdd) { + constexpr absl::string_view kProgram = R"mlir( +func.func @main(%lhs: tensor<10x10xi32>, %rhs: tensor<10x10xi32>) -> tensor<10x10xi32> { + %lhs_scale = "tf.Const"() { value = dense<0.518> : tensor } : () -> tensor + %lhs_zp = "tf.Const"() { value = dense<42> : tensor } : () -> tensor + %rhs_scale = "tf.Const"() { value = dense<0.0239> : tensor } : () -> tensor + %rhs_zp = "tf.Const"() { value = dense<0> : tensor } : () -> tensor + %accum_scale = "tf.Const"() { value = dense<0.013> : tensor } : () -> tensor + %accum_zp = "tf.Const"() { value = dense<0> : tensor } : () -> tensor + %quant_lhs = "tf.Cast"(%lhs) {} : (tensor<10x10xi32>) -> tensor<10x10x!tf_type.qint32> + %quant_rhs = "tf.Cast"(%rhs) {} : (tensor<10x10xi32>) -> tensor<10x10x!tf_type.qint32> + %0 = "tf.UniformQuantizedAdd"( + %quant_lhs, %quant_rhs, %lhs_scale, %lhs_zp, %rhs_scale, + %rhs_zp, %accum_scale, %accum_zp + ) { + Tin = "tfdtype$DT_QINT32", Tout = "tfdtype$DT_QINT32", attr_map = "", + device = "", lhs_quantization_axis = -1 : i64, + lhs_quantization_max_val = 2147483647 : i64, + lhs_quantization_min_val = -2147483648 : i64, + output_quantization_axis = -1 : i64, + output_quantization_max_val = 2147483647 : i64, + output_quantization_min_val = -2147483648 : i64, + rhs_quantization_axis = -1 : i64, + rhs_quantization_max_val = 2147483647 : i64, + rhs_quantization_min_val = -2147483648 : i64 + } : ( + tensor<10x10x!tf_type.qint32>, tensor<10x10x!tf_type.qint32>, tensor, + tensor, tensor, tensor, tensor, tensor + ) -> tensor<10x10x!tf_type.qint32> + %1 = "tf.Cast"(%0) {} : (tensor<10x10x!tf_type.qint32>) -> tensor<10x10xi32> + return %1 : tensor<10x10xi32> +})mlir"; + TF_ASSERT_OK_AND_ASSIGN(auto lhs, CreateRandomI32Literal({10, 10})); + TF_ASSERT_OK_AND_ASSIGN(auto rhs, CreateRandomI32Literal({10, 10})); + // error_tolerance is set to be 1 because different rounding implementations + // in TF kernel and the lowering passes may cause +/-1 differences. + ExecuteAndCompareResultsWithTfKernel(kProgram, {&lhs, &rhs}, + /*tf_program=*/std::nullopt, + /*error_tolerance=*/1.0); } } // namespace diff --git a/tensorflow/compiler/mlir/quantization/stablehlo/passes/lift_quantizable_spots_as_functions_fusion.td b/tensorflow/compiler/mlir/quantization/stablehlo/passes/lift_quantizable_spots_as_functions_fusion.td index a449b4d59018d4..116037d9130df2 100644 --- a/tensorflow/compiler/mlir/quantization/stablehlo/passes/lift_quantizable_spots_as_functions_fusion.td +++ b/tensorflow/compiler/mlir/quantization/stablehlo/passes/lift_quantizable_spots_as_functions_fusion.td @@ -13,9 +13,10 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -include "mlir/IR/OpBase.td" -include "mlir/Dialect/Func/IR/FuncOps.td" include "mlir/Dialect/Arith/IR/ArithOps.td" +include "mlir/Dialect/Func/IR/FuncOps.td" +include "mlir/Dialect/Shape/IR/ShapeOps.td" +include "mlir/IR/OpBase.td" include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td" include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td" include "stablehlo/dialect/StablehloOps.td" @@ -47,7 +48,7 @@ def LiftConvWithBias : Pat< (NamedAttr<"feature_group_count"> $feature_group_count), (NamedAttr<"batch_group_count"> $batch_group_count), (NamedAttr<"precision_config"> (DefaultOrNullAttr $precision_config)))), - [(IsNotInLiftedFunc $res)], [], (addBenefit 5)>; + [(IsNotInLiftedFunc $res), (IsStableHLOConstantOp $bias)], [], (addBenefit 5)>; def LiftDotGeneralWithBias : Pat< (StableHLO_AddOp:$res @@ -60,7 +61,44 @@ def LiftDotGeneralWithBias : Pat< (NamedAttributeList (NamedAttr<"dot_dimension_numbers"> $dot_dimension_numbers), (NamedAttr<"precision_config"> (DefaultOrNullAttr $precision_config)))), - [(IsNotInLiftedFunc $res)], [], (addBenefit 5)>; + [(IsNotInLiftedFunc $res), (IsStableHLOConstantOp $bias)], [], (addBenefit 5)>; + +def LiftConvWithBiasDynamic : Pat< + (StableHLO_AddOp:$res + (StableHLO_ConvolutionOp $lhs, $rhs, $window_strides, $padding, + $lhs_dilation, $rhs_dilation, $window_reversal, $dimension_numbers, + $feature_group_count, $batch_group_count, $precision_config), + (StableHLO_DynamicBroadcastInDimOp + $bias, + (Shape_ShapeOfOp $conv), $_, $_, $_)), + (LiftAsTFXlaCallModule<"composite_conv_with_bias_dynamic_fn"> + (ArgumentList $lhs, $rhs, $bias), + (ResultList $res), + (NamedAttributeList + (NamedAttr<"window_strides"> (DefaultOrNullAttr $window_strides)), + (NamedAttr<"padding"> (DefaultOrNullAttr $padding)), + (NamedAttr<"lhs_dilation"> (DefaultOrNullAttr $lhs_dilation)), + (NamedAttr<"rhs_dilation"> (DefaultOrNullAttr $rhs_dilation)), + (NamedAttr<"window_reversal"> (DefaultOrNullAttr $window_reversal)), + (NamedAttr<"dimension_numbers"> $dimension_numbers), + (NamedAttr<"feature_group_count"> $feature_group_count), + (NamedAttr<"batch_group_count"> $batch_group_count), + (NamedAttr<"precision_config"> (DefaultOrNullAttr $precision_config)))), + [(IsNotInLiftedFunc $res), (IsStableHLOConstantOp $bias)], [], (addBenefit 10)>; + +def LiftDotGeneralWithBiasDynamic : Pat< + (StableHLO_AddOp:$res + (StableHLO_DotGeneralOp $lhs, $rhs, $dot_dimension_numbers, $precision_config), + (StableHLO_DynamicBroadcastInDimOp + $bias, + (Shape_ShapeOfOp $dot_general), $_, $_, $_)), + (LiftAsTFXlaCallModule<"composite_dot_general_with_bias_dynamic_fn"> + (ArgumentList $lhs, $rhs, $bias), + (ResultList $res), + (NamedAttributeList + (NamedAttr<"dot_dimension_numbers"> $dot_dimension_numbers), + (NamedAttr<"precision_config"> (DefaultOrNullAttr $precision_config)))), + [(IsNotInLiftedFunc $res), (IsStableHLOConstantOp $bias)], [], (addBenefit 10)>; //===----------------------------------------------------------------------===// // Pattern rules for lifting ops with activation as functions @@ -101,6 +139,45 @@ def LiftDotGeneralWithRelu : Pat< [(IsNotInLiftedFunc $res), (FloatValueEquals<"0"> $cst)], [], (addBenefit 10)>; +def LiftConvWithReluDynamic : Pat< + (StableHLO_MaxOp:$res + (StableHLO_ConvolutionOp $lhs, $rhs, $window_strides, $padding, + $lhs_dilation, $rhs_dilation, $window_reversal, $dimension_numbers, + $feature_group_count, $batch_group_count, $precision_config), + (StableHLO_DynamicBroadcastInDimOp + (StableHLO_ConstantOp $cst), + (Shape_ShapeOfOp $conv), $_, $_, $_)), + (LiftAsTFXlaCallModule<"composite_conv_with_relu_dynamic_fn"> + (ArgumentList $lhs, $rhs), + (ResultList $res), + (NamedAttributeList + (NamedAttr<"window_strides"> (DefaultOrNullAttr $window_strides)), + (NamedAttr<"padding"> (DefaultOrNullAttr $padding)), + (NamedAttr<"lhs_dilation"> (DefaultOrNullAttr $lhs_dilation)), + (NamedAttr<"rhs_dilation"> (DefaultOrNullAttr $rhs_dilation)), + (NamedAttr<"window_reversal"> (DefaultOrNullAttr $window_reversal)), + (NamedAttr<"dimension_numbers"> $dimension_numbers), + (NamedAttr<"feature_group_count"> $feature_group_count), + (NamedAttr<"batch_group_count"> $batch_group_count), + (NamedAttr<"precision_config"> (DefaultOrNullAttr $precision_config)))), + [(IsNotInLiftedFunc $res), + (FloatValueEquals<"0"> $cst)], [], (addBenefit 15)>; + +def LiftDotGeneralWithReluDynamic : Pat< + (StableHLO_MaxOp:$res + (StableHLO_DotGeneralOp $lhs, $rhs, $dot_dimension_numbers, $precision_config), + (StableHLO_DynamicBroadcastInDimOp + (StableHLO_ConstantOp $cst), + (Shape_ShapeOfOp $dot_general), $_, $_, $_)), + (LiftAsTFXlaCallModule<"composite_dot_general_with_relu_dynamic_fn"> + (ArgumentList $lhs, $rhs), + (ResultList $res), + (NamedAttributeList + (NamedAttr<"dot_dimension_numbers"> $dot_dimension_numbers), + (NamedAttr<"precision_config"> (DefaultOrNullAttr $precision_config)))), + [(IsNotInLiftedFunc $res), + (FloatValueEquals<"0"> $cst)], [], (addBenefit 15)>; + def LiftConvWithRelu6 : Pat< (StableHLO_ClampOp:$res (StableHLO_ConstantOp $cst_0), @@ -163,7 +240,7 @@ def LiftConvWithBiasAndRelu : Pat< (NamedAttr<"batch_group_count"> $batch_group_count), (NamedAttr<"precision_config"> (DefaultOrNullAttr $precision_config)))), [(IsNotInLiftedFunc $res), - (FloatValueEquals<"0"> $cst)], [], (addBenefit 10)>; + (FloatValueEquals<"0"> $cst), (IsStableHLOConstantOp $bias)], [], (addBenefit 10)>; def LiftDotGeneralWithBiasAndRelu : Pat< (StableHLO_MaxOp:$res @@ -179,7 +256,55 @@ def LiftDotGeneralWithBiasAndRelu : Pat< (NamedAttr<"dot_dimension_numbers"> $dot_dimension_numbers), (NamedAttr<"precision_config"> (DefaultOrNullAttr $precision_config)))), [(IsNotInLiftedFunc $res), - (FloatValueEquals<"0"> $cst)], [], (addBenefit 10)>; + (FloatValueEquals<"0"> $cst), (IsStableHLOConstantOp $bias)], [], (addBenefit 10)>; + +def LiftConvWithBiasAndReluDynamic : Pat< + (StableHLO_MaxOp:$res + (StableHLO_AddOp + (StableHLO_ConvolutionOp $lhs, $rhs, $window_strides, $padding, + $lhs_dilation, $rhs_dilation, $window_reversal, $dimension_numbers, + $feature_group_count, $batch_group_count, $precision_config), + (StableHLO_DynamicBroadcastInDimOp + $bias, + (Shape_ShapeOfOp $conv), $_, $_, $_)), + (StableHLO_DynamicBroadcastInDimOp + (StableHLO_ConstantOp $cst), + (Shape_ShapeOfOp $add), $_, $_, $_)), + (LiftAsTFXlaCallModule<"composite_conv_with_bias_and_relu_dynamic_fn"> + (ArgumentList $lhs, $rhs, $bias), + (ResultList $res), + (NamedAttributeList + (NamedAttr<"window_strides"> (DefaultOrNullAttr $window_strides)), + (NamedAttr<"padding"> (DefaultOrNullAttr $padding)), + (NamedAttr<"lhs_dilation"> (DefaultOrNullAttr $lhs_dilation)), + (NamedAttr<"rhs_dilation"> (DefaultOrNullAttr $rhs_dilation)), + (NamedAttr<"window_reversal"> (DefaultOrNullAttr $window_reversal)), + (NamedAttr<"dimension_numbers"> $dimension_numbers), + (NamedAttr<"feature_group_count"> $feature_group_count), + (NamedAttr<"batch_group_count"> $batch_group_count), + (NamedAttr<"precision_config"> (DefaultOrNullAttr $precision_config)))), + [(IsNotInLiftedFunc $res), + (FloatValueEquals<"0"> $cst), (IsStableHLOConstantOp $bias)], [], (addBenefit 15)>; + +def LiftDotGeneralWithBiasAndReluDynamic : Pat< + (StableHLO_MaxOp:$res + (StableHLO_AddOp + (StableHLO_DotGeneralOp $lhs, $rhs, $dot_dimension_numbers, $precision_config), + (StableHLO_DynamicBroadcastInDimOp + $bias, + (Shape_ShapeOfOp $dot_general), $_, $_, $_)), + (StableHLO_DynamicBroadcastInDimOp + (StableHLO_ConstantOp $cst), + (Shape_ShapeOfOp $add), $_, $_, $_)), + (LiftAsTFXlaCallModule<"composite_dot_general_with_bias_and_relu_dynamic_fn"> + (ArgumentList $lhs, $rhs, $bias), + (ResultList $res), + (NamedAttributeList + (NamedAttr<"dot_dimension_numbers"> $dot_dimension_numbers), + (NamedAttr<"precision_config"> (DefaultOrNullAttr $precision_config)))), + [(IsNotInLiftedFunc $res), + (FloatValueEquals<"0"> $cst), (IsStableHLOConstantOp $bias)], [], (addBenefit 15)>; + def LiftConvWithBiasAndRelu6 : Pat< (StableHLO_ClampOp:$res @@ -203,7 +328,7 @@ def LiftConvWithBiasAndRelu6 : Pat< (NamedAttr<"feature_group_count"> $feature_group_count), (NamedAttr<"batch_group_count"> $batch_group_count), (NamedAttr<"precision_config"> (DefaultOrNullAttr $precision_config)))), - [(IsNotInLiftedFunc $res), (FloatValueEquals<"0"> $cst_0), (FloatValueEquals<"6"> $cst_1)], [], (addBenefit 10)>; + [(IsNotInLiftedFunc $res), (IsStableHLOConstantOp $bias), (FloatValueEquals<"0"> $cst_0), (FloatValueEquals<"6"> $cst_1)], [], (addBenefit 10)>; def LiftDotGeneralWithBiasAndRelu6 : Pat< (StableHLO_ClampOp:$res @@ -219,4 +344,4 @@ def LiftDotGeneralWithBiasAndRelu6 : Pat< (NamedAttributeList (NamedAttr<"dot_dimension_numbers"> $dot_dimension_numbers), (NamedAttr<"precision_config"> (DefaultOrNullAttr $precision_config)))), - [(IsNotInLiftedFunc $res), (FloatValueEquals<"0"> $cst_0), (FloatValueEquals<"6"> $cst_1)], [], (addBenefit 10)>; + [(IsNotInLiftedFunc $res), (IsStableHLOConstantOp $bias), (FloatValueEquals<"0"> $cst_0), (FloatValueEquals<"6"> $cst_1)], [], (addBenefit 10)>; diff --git a/tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h b/tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h index 8e9cfb57f32e63..0b05069b265989 100644 --- a/tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h +++ b/tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h @@ -21,10 +21,16 @@ limitations under the License. #include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project #include "mlir/IR/BuiltinOps.h" // from @llvm-project // IWYU pragma: keep #include "mlir/Pass/Pass.h" // from @llvm-project +#include "tensorflow/compiler/mlir/lite/quantization/quantization_config.h" #include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_options.pb.h" namespace mlir::quant::stablehlo { +// Creates a `QuantizePass` that quantizes ops according to surrounding qcast / +// dcast ops. +std::unique_ptr> CreateQuantizePass( + const quant::QuantizationSpecs& quantization_specs); + // Creates a pass that quantizes weight component of StableHLO graph. std::unique_ptr> CreateQuantizeWeightPass( const ::stablehlo::quantization::QuantizationComponentSpec& diff --git a/tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.td b/tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.td index 0992c042175f01..52dca7897ea05d 100644 --- a/tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.td +++ b/tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.td @@ -76,7 +76,30 @@ def QuantizePass : Pass<"stablehlo-quantize", "mlir::func::FuncOp"> { ]; } - def RestoreFunctionNamePass : Pass<"stablehlo-restore-function-name", "ModuleOp"> { let summary = "Restores function name from XlaCallModule op."; } + +def PostQuantizePass : Pass<"stablehlo-post-quantize", "mlir::func::FuncOp"> { + let summary = "Apply clean-up after quantization."; + let dependentDialects = [ + "mlir::stablehlo::StablehloDialect", + "mlir::quantfork::QuantizationForkDialect", + ]; +} + +def QuantizeCompositeFunctionsPass : Pass<"stablehlo-quantize-composite-functions", "ModuleOp"> { + let summary = "Quantize composite functions with QDQ input / outputs."; + let options = [ + Option<"mlir_dump_file_name_", "mlir-dump-file-name", + "std::optional", /*default=*/"std::nullopt", + "MLIR dump file name."> + ]; + let dependentDialects = [ + "mlir::arith::ArithDialect", + "mlir::stablehlo::StablehloDialect", + "mlir::quant::QuantizationDialect", + "mlir::quantfork::QuantizationForkDialect", + "TF::TensorFlowDialect", + ]; +} diff --git a/tensorflow/compiler/mlir/quantization/stablehlo/passes/post_quantize.cc b/tensorflow/compiler/mlir/quantization/stablehlo/passes/post_quantize.cc new file mode 100644 index 00000000000000..0416bbdbff73a0 --- /dev/null +++ b/tensorflow/compiler/mlir/quantization/stablehlo/passes/post_quantize.cc @@ -0,0 +1,158 @@ +/* Copyright 2023 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include + +#include "llvm/Support/Casting.h" +#include "mlir/Dialect/Quant/QuantTypes.h" // from @llvm-project +#include "mlir/IR/BuiltinAttributeInterfaces.h" // from @llvm-project +#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project +#include "mlir/IR/MLIRContext.h" // from @llvm-project +#include "mlir/IR/Matchers.h" // from @llvm-project +#include "mlir/IR/OpDefinition.h" // from @llvm-project +#include "mlir/IR/PatternMatch.h" // from @llvm-project +#include "mlir/Support/LogicalResult.h" // from @llvm-project +#include "mlir/Support/TypeID.h" // from @llvm-project +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" // from @llvm-project +#include "stablehlo/dialect/StablehloOps.h" // from @stablehlo +#include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.h" +#include "tensorflow/compiler/mlir/lite/quantization/quantization_utils.h" +#include "tensorflow/compiler/mlir/lite/transforms/passes.h" + +namespace mlir::quant::stablehlo { + +#define GEN_PASS_DEF_POSTQUANTIZEPASS +#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h.inc" + +namespace { + +// Applies clean-up patterns after quantization. +class PostQuantizePass : public impl::PostQuantizePassBase { + public: + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PostQuantizePass) + + explicit PostQuantizePass() = default; + + private: + void runOnOperation() override; +}; + +// TODO: b/305815328 - Consider preserving leading and trailing QDQs for +// ModifyIONodesPass in TFLite use cases. +// Removes the back-to-back quantize and dequantize ops with volatile attribute. +class RemoveVolatileQdqPattern + : public OpRewritePattern { + public: + explicit RemoveVolatileQdqPattern(MLIRContext* context) + : OpRewritePattern(context) {} + + LogicalResult matchAndRewrite(quantfork::DequantizeCastOp op, + PatternRewriter& rewriter) const override { + auto input_op = op.getArg().getDefiningOp(); + if (auto q = llvm::dyn_cast_or_null(input_op)) { + if (!q->getAttr(kVolatileOpAttrName)) return failure(); + + // If the quantize op is a requantize op, it is being used in other scale + // adjustments and should be kept. Instead, move dequantize op before the + // requantize op to remove the unnecessary requantize op. + if (auto qtype = + QuantizedType::getQuantizedElementType(q.getArg().getType())) { + rewriter.setInsertionPoint(op); + rewriter.replaceOpWithNewOp( + op, op.getResult().getType(), q.getArg()); + return success(); + } + + op.replaceAllUsesWith(q.getArg()); + return success(); + } + return failure(); + } +}; + +// Replaces constant and uniform_quantize ops with single quantized constant op. +class QuantizeConstPattern + : public OpRewritePattern { + public: + explicit QuantizeConstPattern(MLIRContext* context) + : OpRewritePattern(context) {} + + LogicalResult matchAndRewrite(mlir::stablehlo::UniformQuantizeOp op, + PatternRewriter& rewriter) const override { + DenseFPElementsAttr attr; + if (matchPattern(op.getOperand(), m_Constant(&attr))) { + auto qtype = op.getResult().getType(); + ElementsAttr quantized_attr = Quantize(attr, qtype); + if (quantized_attr) { + rewriter.replaceOpWithNewOp( + op, qtype, quantized_attr); + return success(); + } + } + return failure(); + } +}; + +// Replaces quantfork.dcast with stablehlo.uniform_dequantize. +class ConvertDequantizeCastToUniformDequantizePattern + : public OpRewritePattern { + public: + explicit ConvertDequantizeCastToUniformDequantizePattern(MLIRContext* context) + : OpRewritePattern(context) {} + LogicalResult matchAndRewrite(quantfork::DequantizeCastOp dq_op, + PatternRewriter& rewriter) const override { + rewriter.replaceOpWithNewOp( + dq_op, dq_op.getResult().getType(), dq_op.getArg()); + return success(); + } +}; + +// Replaces quantfork.qcast with stablehlo.uniform_quantize. +class ConvertQuantizeCastToUniformQuantizePattern + : public OpRewritePattern { + public: + explicit ConvertQuantizeCastToUniformQuantizePattern(MLIRContext* context) + : OpRewritePattern(context) {} + LogicalResult matchAndRewrite(quantfork::QuantizeCastOp q_op, + PatternRewriter& rewriter) const override { + rewriter.replaceOpWithNewOp( + q_op, q_op.getResult().getType(), q_op.getArg()); + return success(); + } +}; + +void PostQuantizePass::runOnOperation() { + RewritePatternSet patterns(&getContext()); + func::FuncOp func = getOperation(); + MLIRContext* ctx = func.getContext(); + // TODO: b/307463853 - Consider splitting passes for each pattern set. + patterns.add, + RemoveVolatileQdqPattern>(ctx); + if (failed(applyPatternsAndFoldGreedily(func, std::move(patterns)))) { + signalPassFailure(); + } + + RewritePatternSet patterns_2(&getContext()); + patterns_2 + .add( + ctx); + if (failed(applyPatternsAndFoldGreedily(func, std::move(patterns_2)))) { + signalPassFailure(); + } +} + +} // namespace +} // namespace mlir::quant::stablehlo diff --git a/tensorflow/compiler/mlir/quantization/stablehlo/passes/quantize.cc b/tensorflow/compiler/mlir/quantization/stablehlo/passes/quantize.cc index c5b7a3a3a3420d..16e7ad1cfd7010 100644 --- a/tensorflow/compiler/mlir/quantization/stablehlo/passes/quantize.cc +++ b/tensorflow/compiler/mlir/quantization/stablehlo/passes/quantize.cc @@ -13,6 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +#include #include #include @@ -137,4 +138,9 @@ void QuantizePass::runOnOperation() { } // namespace +std::unique_ptr> CreateQuantizePass( + const QuantizationSpecs& quantization_specs) { + return std::make_unique(quantization_specs); +} + } // namespace mlir::quant::stablehlo diff --git a/tensorflow/compiler/mlir/quantization/stablehlo/passes/quantize_composite_functions.cc b/tensorflow/compiler/mlir/quantization/stablehlo/passes/quantize_composite_functions.cc new file mode 100644 index 00000000000000..cf0c44f779a9ae --- /dev/null +++ b/tensorflow/compiler/mlir/quantization/stablehlo/passes/quantize_composite_functions.cc @@ -0,0 +1,358 @@ +/* Copyright 2023 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ +#include +#include +#include +#include + +#include "absl/algorithm/container.h" +#include "absl/status/status.h" +#include "llvm/ADT/STLExtras.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project +#include "mlir/Dialect/Quant/QuantOps.h" // from @llvm-project // IWYU pragma: keep +#include "mlir/Dialect/Quant/QuantTypes.h" // from @llvm-project +#include "mlir/IR/Block.h" // from @llvm-project +#include "mlir/IR/Builders.h" // from @llvm-project +#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project +#include "mlir/IR/BuiltinOps.h" // from @llvm-project +#include "mlir/IR/BuiltinTypes.h" // from @llvm-project +#include "mlir/IR/Location.h" // from @llvm-project +#include "mlir/IR/MLIRContext.h" // from @llvm-project +#include "mlir/IR/OperationSupport.h" // from @llvm-project +#include "mlir/IR/PatternMatch.h" // from @llvm-project +#include "mlir/IR/SymbolTable.h" // from @llvm-project +#include "mlir/IR/TypeUtilities.h" // from @llvm-project +#include "mlir/IR/Visitors.h" // from @llvm-project +#include "mlir/Pass/Pass.h" // from @llvm-project // IWYU pragma: keep +#include "mlir/Pass/PassRegistry.h" // from @llvm-project +#include "mlir/Support/LLVM.h" // from @llvm-project +#include "mlir/Support/LogicalResult.h" // from @llvm-project +#include "mlir/Support/TypeID.h" // from @llvm-project +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" // from @llvm-project +#include "stablehlo/dialect/StablehloOps.h" // from @stablehlo // IWYU pragma: keep +#include "tensorflow/compiler/mlir/lite/quantization/quantization_config.h" +#include "tensorflow/compiler/mlir/lite/quantization/quantization_utils.h" +#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h" +#include "tensorflow/compiler/mlir/quantization/stablehlo/uniform_quantized_types.h" +#include "tensorflow/compiler/mlir/quantization/tensorflow/cc/run_passes.h" +#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h" +#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h" + +namespace mlir::quant::stablehlo { + +#define GEN_PASS_DEF_QUANTIZECOMPOSITEFUNCTIONSPASS +#include "tensorflow/compiler/mlir/quantization/stablehlo/passes/passes.h.inc" + +namespace { + +using QuantMethod = tensorflow::quantization::QuantizationMethod::PresetMethod; +using ::mlir::stablehlo::DotGeneralOp; +using ::mlir::stablehlo::UniformQuantizeOp; +using ::tensorflow::quantization::RunPassesOnModuleOp; + +constexpr StringRef kCompositeFuncPrefix = "composite_"; +constexpr StringRef kQuantizedFuncPrefix = "quantized_"; +constexpr StringRef kEntryFuncAttrName = "_entry_function"; + +class QuantizeCompositeFunctionsPass + : public impl::QuantizeCompositeFunctionsPassBase< + QuantizeCompositeFunctionsPass> { + public: + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(QuantizeCompositeFunctionsPass) + + using impl::QuantizeCompositeFunctionsPassBase< + QuantizeCompositeFunctionsPass>::QuantizeCompositeFunctionsPassBase; + + private: + void runOnOperation() override; +}; + +// Returns true if `type` is a TensorType with quantized elements. +bool IsQuantizedTensorType(const Type type) { + return type.isa() && + type.cast().getElementType().isa(); +} + +// Checks if all inputs and outputs are quantized. +bool HasQuantizedOperandOrOutput(Operation* call_op) { + SmallVector arg_types; + for (const Value arg : call_op->getOperands()) { + arg_types.push_back(arg.getType()); + } + + SmallVector output_types; + for (const Value output : call_op->getResults()) { + output_types.push_back(output.getType()); + } + + return absl::c_all_of(arg_types, IsQuantizedTensorType) && + absl::c_all_of(output_types, IsQuantizedTensorType); +} + +// Get the corresponding quantized function name from the given function name. +// Example: "composite_dot_general_fn_1" => "quantized_dot_general_fn" +std::string GetQuantizedFunctionName(const StringRef func_name) { + return Twine(kQuantizedFuncPrefix) + .concat(func_name.rsplit(kCompositeFuncPrefix).second) + .str(); +} + +// Returns true if `xla_call_module_op` is quantized. To be considered +// quantized, it should meet three conditions: +// 1. At least one of the inputs or outputs should be a uniform quantized type. +// 2. `xla_call_module_op` should have the `kQuantTraitAttrName` attribute. +// 3. It should also have the `kEntryFuncAttrName` attribute, which points to +// the function that `xla_call_module_op` represents. +bool IsQuantizedXlaCallModuleOp(TF::XlaCallModuleOp xla_call_module_op) { + return HasQuantizedOperandOrOutput(xla_call_module_op) && + xla_call_module_op->hasAttr(kQuantTraitAttrName) && + xla_call_module_op->hasAttr(kEntryFuncAttrName); +} + +// Returns the entry function, i.e. the callee of `xla_call_module_op`. +func::FuncOp GetEntryFuncOp(TF::XlaCallModuleOp xla_call_module_op, + SymbolTable symbol_table) { + auto entry_function_symbol_ref = + xla_call_module_op->getAttrOfType(kEntryFuncAttrName); + + // Don't match if there are no DotGeneralOp. + // if (target_func_op.getOps().empty()) return {}; + return dyn_cast_or_null( + symbol_table.lookup(entry_function_symbol_ref.getValue())); +} + +// Replaces the function type of `entry_func_op` to a quantized one, matching +// the input and output types of `xla_call_module_op`. +void SetQuantizedFunctionType(PatternRewriter& rewriter, + func::FuncOp entry_func_op, + TF::XlaCallModuleOp xla_call_module_op) { + SmallVector arg_types; + SmallVector arg_locs; + for (const Value arg : xla_call_module_op.getArgs()) { + arg_types.push_back(arg.getType()); + arg_locs.push_back(arg.getLoc()); + } + + SmallVector output_types; + for (const Value output : xla_call_module_op.getOutput()) { + output_types.push_back(output.getType()); + } + + entry_func_op.setFunctionType( + rewriter.getFunctionType(arg_types, output_types)); + + // Replace argument types and locs. + Block& entry = entry_func_op->getRegion(0).front(); + for (auto [arg, arg_type, arg_loc] : + llvm::zip_equal(entry.getArguments(), arg_types, arg_locs)) { + arg.setType(arg_type); + arg.setLoc(arg_loc); + } +} + +// An interface representing patterns that quantizes an entry function's body. +// The entry function's signatures should have already been quantized at the +// point of rewriting. +class EntryFuncBodyQuantizationPattern { + public: + virtual ~EntryFuncBodyQuantizationPattern() = default; + + // Returns `success()` if `entry_func_op`'s body is eligible for rewriting. At + // this point `entry_func_op`'s signature has not been reset with quantized + // types. + virtual LogicalResult match(func::FuncOp entry_func_op) const = 0; + + // Rewrites the `entry_func_op`'s body. + virtual void rewrite(func::FuncOp entry_func_op, + PatternRewriter& rewriter) const = 0; +}; + +// Quantizes the entry function's body containing a `DotGeneralOp`. +class QuantizeDotGeneralOpPattern : public EntryFuncBodyQuantizationPattern { + public: + explicit QuantizeDotGeneralOpPattern(MLIRContext& ctx) : ctx_(&ctx) {} + + LogicalResult match(func::FuncOp entry_func_op) const override { + auto& operations = entry_func_op.getBody().front().getOperations(); + return success(operations.size() == 2 && + isa(operations.front())); + } + + void rewrite(func::FuncOp entry_func_op, + PatternRewriter& rewriter) const override { + // Update the output type of the dot_general op. + auto dot_general_op = *entry_func_op.getOps().begin(); + + const Type input_type = entry_func_op.getArgumentTypes()[0]; + const Type rhs_type = entry_func_op.getArgumentTypes()[1]; + const Type func_result_type = entry_func_op.getResultTypes()[0]; + + const double input_scale = getElementTypeOrSelf(input_type) + .cast() + .getScale(); + const double rhs_scale = + getElementTypeOrSelf(rhs_type).cast().getScale(); + + // Define the intermediate output type, which is an i32 quantized type. + // This is intermediate because the final output type of the entry_func_op + // should be an i8 quantized type. + const UniformQuantizedType output_quantized_element_type = + CreateI32F32UniformQuantizedType(dot_general_op->getLoc(), *ctx_, + input_scale * rhs_scale, + /*zero_point=*/0); + + Value dot_general_op_result = dot_general_op->getResult(0); + const auto dot_general_op_result_type = + dot_general_op_result.getType().cast(); + const ArrayRef shape = dot_general_op_result_type.getShape(); + + const TensorType new_dot_general_op_result_type = + dot_general_op_result_type.cloneWith(shape, + output_quantized_element_type); + dot_general_op_result.setType(new_dot_general_op_result_type); + + // Add i32 -> i8 requantization. + rewriter.setInsertionPointAfter(dot_general_op); + auto uniform_quant_op = rewriter.create( + dot_general_op->getLoc(), func_result_type, + dot_general_op->getResults()); + + auto return_op = + cast(entry_func_op.getBody().front().getTerminator()); + return_op.setOperand(0, uniform_quant_op); + } + + private: + MLIRContext* ctx_ = nullptr; +}; + +// Converts `entry_func_op` to be quantized according to the respective +// inputs and outputs of `xla_call_module_op` that are possibly quantized. It +// signature (type) is reset to match that of `xla_call_module_op`. +// `entry_func_body_quantization_pattern` rewrites the function's body, based on +// the new signature. +void QuantizeEntryFuncOp( + MLIRContext& ctx, PatternRewriter& rewriter, + TF::XlaCallModuleOp xla_call_module_op, func::FuncOp entry_func_op, + const EntryFuncBodyQuantizationPattern& body_rewrite_pattern) { + SetQuantizedFunctionType(rewriter, entry_func_op, xla_call_module_op); + + body_rewrite_pattern.rewrite(entry_func_op, rewriter); + + // Rename the function to be clear that the function has been quantized. + const std::string quantized_function_name = + GetQuantizedFunctionName(entry_func_op.getSymName()); + entry_func_op.setSymName(quantized_function_name); +} + +// Replaces a quantized `xla_call_module_op` with a `func::CallOp`. The callee +// is expected to remain unquantized (thus having a signature mismatch), and it +// is also quantized accordingly. +void ReplaceQuantizedXlaCallModuleOpWithQuantizedCallOp( + MLIRContext& ctx, PatternRewriter& rewriter, + TF::XlaCallModuleOp xla_call_module_op, + const EntryFuncBodyQuantizationPattern& body_rewrite_pattern) { + auto module_op = xla_call_module_op->getParentOfType(); + SymbolTable symbol_table(module_op); + + func::FuncOp entry_func_op = GetEntryFuncOp(xla_call_module_op, symbol_table); + QuantizeEntryFuncOp(ctx, rewriter, xla_call_module_op, entry_func_op, + body_rewrite_pattern); + + // Replace the XlaCallModuleOp with a new CallOp. + rewriter.setInsertionPoint(xla_call_module_op); + rewriter.replaceOpWithNewOp(xla_call_module_op, entry_func_op, + xla_call_module_op.getArgs()); +} + +// Pattern that mainly does two things: +// +// 1. Replaces quantized `TF::XlaCallModuleOp` with a `func::CallOp`. +// 2. Quantizes the callee function. +// +// The inputs of this pattern assumes an invalid IR, where even if a +// `TF::XlaCallModuleOp` is quantized the callee remains unquantized. Step (2) +// not only replaces the input and output tensor types into quantized ones, but +// also rewrites the body with a quantized equivalent. +// +// `FuncBodyRewritePatternT` defines how a function body is quantized and +// rewritten. +template >> +class XlaCallModuleOpToCallOp : public OpRewritePattern { + public: + explicit XlaCallModuleOpToCallOp(MLIRContext& ctx) + : OpRewritePattern(&ctx) {} + + LogicalResult match(TF::XlaCallModuleOp op) const override { + auto module_op = op->getParentOfType(); + SymbolTable symbol_table(module_op); + + // Ignore unquantized ops. + if (!IsQuantizedXlaCallModuleOp(op)) return failure(); + + func::FuncOp entry_func_op = GetEntryFuncOp(op, symbol_table); + if (!entry_func_op) { + op->emitError("Failed to find a valid entry function."); + return failure(); + } + + return FuncBodyRewritePatternT(*getContext()).match(entry_func_op); + } + + void rewrite(TF::XlaCallModuleOp xla_call_module_op, + PatternRewriter& rewriter) const override { + ReplaceQuantizedXlaCallModuleOpWithQuantizedCallOp( + *rewriter.getContext(), rewriter, xla_call_module_op, + FuncBodyRewritePatternT(*getContext())); + } +}; + +void QuantizeCompositeFunctionsPass::runOnOperation() { + MLIRContext& ctx = getContext(); + + QuantizationSpecs quant_specs; + quant_specs.inference_type = tensorflow::DT_QINT8; + + PassManager pm(&ctx); + // Intermediate output from QuantizePass will have quantized ops + // (XlaCallModuleOps) with quantized input and output types, which are not + // allowed in the TF dialect. + pm.enableVerifier(false); + + pm.addNestedPass(CreatePrepareQuantizePass()); + pm.addNestedPass(CreateQuantizePass(quant_specs)); + pm.addNestedPass(createPostQuantizePass()); + + ModuleOp module_op = getOperation(); + if (const absl::Status pm_run_status = + RunPassesOnModuleOp(mlir_dump_file_name_, pm, module_op); + !pm_run_status.ok()) { + signalPassFailure(); + } + + // TODO - b/307839649: Move this as a separate pass. + RewritePatternSet patterns(&ctx); + patterns.add>(ctx); + + if (failed(applyPatternsAndFoldGreedily(module_op, std::move(patterns)))) { + signalPassFailure(); + } +} + +} // namespace + +} // namespace mlir::quant::stablehlo diff --git a/tensorflow/compiler/mlir/quantization/stablehlo/passes/replace_stablehlo_ops_in_main_function_with_xla_call_module_ops.cc b/tensorflow/compiler/mlir/quantization/stablehlo/passes/replace_stablehlo_ops_in_main_function_with_xla_call_module_ops.cc index 8cc75fbc261a20..5bf8ba7ec07657 100644 --- a/tensorflow/compiler/mlir/quantization/stablehlo/passes/replace_stablehlo_ops_in_main_function_with_xla_call_module_ops.cc +++ b/tensorflow/compiler/mlir/quantization/stablehlo/passes/replace_stablehlo_ops_in_main_function_with_xla_call_module_ops.cc @@ -78,7 +78,43 @@ std::string CreateStablehloFunctionName(const int id) { return Twine("_stablehlo_main_").concat(std::to_string(id)).str(); } -// Follows the structure of Live-variable analysis. +// Follows the structure of Live-variable analysis. It is a form of +// CFG (Control Flow Graph) analysis, often used in compilers. +// +// A variable is live if it holds a value that may be used in the future. +// It is live-in at node n if it is live on any of the node's in-edges. +// It is live-out at node n if it is live on any of the node's out-edges. +// def[n] refers to values that are defined at node n. +// use[n] refers to values that are used at node n. +// +// Given a node n, variables' liveliness is defined like the following: +// live_in[n] = use[n] U (live_out[n] - def[n]) +// live_out[n] = U {live_in[s] | s ε succ[n]} +// +// Consider a sequence of op: +// +// ``` +// node 1: %0 = stablehlo.constant +// node 2: %1 = stablehlo.constant +// node 3: %2 = stablehlo.add %0, %1 +// node 4: %3 = stablehlo.multiply %2, %1 +// node 5: return %3 +// ``` +// +// In Backward Liveliness analysis, the liveliness for each node above becomes: +// live_in[5] = use[5] U (live_out[5] - def[5]) +// = {%3} U ({∅} - {∅}) = {%3} +// live_in[4] = use[4] U (live_out[4] - def[4]) +// = {%1, %2} U ({%3} - {%3}) = {%1, %2} +// live_in[3] = use[3] U (live_out[3] - def[3]) +// = {%0, %1} U ({%1, %2} - {%2}) = {%0, %1} +// live_in[2] = use[2] U (live_out[2] - def[2]) +// = {∅} U ({%0, %1} - {%1}) = {%0} +// live_in[1] = use[1] U (live_out[1] - def[1]) +// = {∅} U ({%0} - {%0}) = {∅} +// +// This analogy is used throughout this pass to ensure only live edges form +// proper subgraphs. class LiveOuts { public: LiveOuts() = default; @@ -100,10 +136,10 @@ class LiveOuts { void snapshot_previous_state() { prev_liveouts_ = liveouts_; } // Return the current live values. - DenseSet& get() { return liveouts_; } + const DenseSet& get() const { return liveouts_; } // Return the previous live values. - DenseSet& get_previous() { return prev_liveouts_; } + const DenseSet& get_previous() const { return prev_liveouts_; } private: DenseSet liveouts_; @@ -212,6 +248,38 @@ void ReplaceStablehloOpsWithXlaCallModuleOp( } } +// Contains the actual logic for updating states and replacing StableHLO ops +// with tf.XlaCallModuleOps. +void UpdateStatesAndReplaceStablehloOps( + const DenseSet& operands, const DenseSet& defined_values, + const LiveOuts& liveouts, ModuleOp module_op, + ArrayRef reverse_subgraph, const int stablehlo_func_id, + func::FuncOp main_func, const bool is_last_subgraph = false) { + DenseSet inputs = operands; + for (Value defined_value : defined_values) { + inputs.erase(defined_value); + } + + DenseSet outputs = liveouts.get_previous(); + for (Value live_value : liveouts.get()) { + outputs.erase(live_value); + } + + if (is_last_subgraph) { + // Additionally remove arguments from the outputs, as it provides liveness + // throughout (functions as an invisible op above the very first op that + // returns the arguments). + for (const BlockArgument arg : main_func.getArguments()) { + outputs.erase(arg); + } + } + + ReplaceStablehloOpsWithXlaCallModuleOp( + SmallVector(inputs.begin(), inputs.end()), + SmallVector(outputs.begin(), outputs.end()), reverse_subgraph, + stablehlo_func_id, module_op); +} + // Replaces the StableHLO ops in the main function block with // tf.XlaCallModuleOps as separate subgraphs. Wires them back to the main // function block to be compatible with SavedModel structure. @@ -241,20 +309,14 @@ void ReplaceStablehloOpsInMainFunctionWithXlaCallModuleOps( DenseSet operands; DenseSet defined_values; - int stablehlo_func_id = 0; + int stablehlo_func_id = -1; for (Operation* op : reverse_main_func_block_ops) { if (!IsStablehloOp(op)) { // Create an XlaCallModuleOp if reverse_subgraph isn't empty. if (!reverse_subgraph.empty()) { - DenseSet outputs = liveouts.get_previous(); - for (Value live_value : liveouts.get()) { - outputs.erase(live_value); - } - - ReplaceStablehloOpsWithXlaCallModuleOp( - SmallVector(operands.begin(), operands.end()), - SmallVector(outputs.begin(), outputs.end()), - reverse_subgraph, stablehlo_func_id++, module_op); + UpdateStatesAndReplaceStablehloOps(operands, defined_values, liveouts, + module_op, reverse_subgraph, + ++stablehlo_func_id, main_func); // Reset states and start a new subgraph. reverse_subgraph.clear(); @@ -273,25 +335,16 @@ void ReplaceStablehloOpsInMainFunctionWithXlaCallModuleOps( } reverse_subgraph.push_back(op); + + defined_values.insert(op->getResults().begin(), op->getResults().end()); + operands.insert(op->getOperands().begin(), op->getOperands().end()); } // Create the last subgraph if it isn't empty. if (!reverse_subgraph.empty()) { - DenseSet outputs = liveouts.get_previous(); - for (Value live_value : liveouts.get()) { - outputs.erase(live_value); - } - // Additionally remove arguments from the outputs, as it provides liveness - // throughout (functions as an invisible op above the very first op that - // returns the arguments). - for (const BlockArgument arg : main_func.getArguments()) { - outputs.erase(arg); - } - - ReplaceStablehloOpsWithXlaCallModuleOp( - SmallVector(operands.begin(), operands.end()), - SmallVector(outputs.begin(), outputs.end()), reverse_subgraph, - stablehlo_func_id++, module_op); + UpdateStatesAndReplaceStablehloOps( + operands, defined_values, liveouts, module_op, reverse_subgraph, + ++stablehlo_func_id, main_func, /*is_last_subgraph=*/true); } } diff --git a/tensorflow/compiler/mlir/quantization/stablehlo/passes/utils.td b/tensorflow/compiler/mlir/quantization/stablehlo/passes/utils.td index 8c6ab88b3368e9..744637d58d8760 100644 --- a/tensorflow/compiler/mlir/quantization/stablehlo/passes/utils.td +++ b/tensorflow/compiler/mlir/quantization/stablehlo/passes/utils.td @@ -24,3 +24,7 @@ class FloatValueEquals : Constraint; + +// Returns true if the given op is a StableHLO constant op. +def IsStableHLOConstantOp : Constraint($0.getDefiningOp())">>; + diff --git a/tensorflow/compiler/mlir/quantization/stablehlo/tests/bridge/convert-mhlo-quant-to-int.mlir b/tensorflow/compiler/mlir/quantization/stablehlo/tests/bridge/convert-mhlo-quant-to-int.mlir index 9f9c11454c5c2a..cb8ad65a5cdd48 100644 --- a/tensorflow/compiler/mlir/quantization/stablehlo/tests/bridge/convert-mhlo-quant-to-int.mlir +++ b/tensorflow/compiler/mlir/quantization/stablehlo/tests/bridge/convert-mhlo-quant-to-int.mlir @@ -106,6 +106,53 @@ func.func @uniform_quantize_and_dequantize_sparse_tensor_encoding(%arg0: tensor< // ----- +// CHECK-LABEL: func @quantize_per_channel +func.func @quantize_per_channel(%arg0: tensor<26x26x3x2xf32> + ) -> tensor<26x26x3x2x!quant.uniform> { + // CHECK-DAG: %[[SCALES:.*]] = mhlo.constant dense<[1.100000e+00, 1.100000e-01]> + // CHECK-DAG: %[[ZPS:.*]] = mhlo.constant dense<[-1.000000e+01, 2.000000e+00]> + // CHECK-DAG: %[[QMIN:.*]] = mhlo.constant dense<-2.14748365E+9> : tensor + // CHECK-DAG: %[[QMAX:.*]] = mhlo.constant dense<2.14748365E+9> : tensor + // CHECK: %[[DIVIDE:.*]] = chlo.broadcast_divide %arg0, %[[SCALES]] + // CHECK-SAME: {broadcast_dimensions = dense<3> : tensor<1xi64>} + // CHECK-SAME: (tensor<26x26x3x2xf32>, tensor<2xf32>) -> tensor<26x26x3x2xf32> + // CHECK: %[[ADD:.*]] = chlo.broadcast_add %[[DIVIDE]], %[[ZPS]] + // CHECK-SAME: {broadcast_dimensions = dense<3> : tensor<1xi64>} + // CHECK-SAME: (tensor<26x26x3x2xf32>, tensor<2xf32>) -> tensor<26x26x3x2xf32> + // CHECK: %[[CLAMP:.*]] = mhlo.clamp %[[QMIN]], %[[ADD]], %[[QMAX]] + // CHECK: %[[ROUND:.*]] = mhlo.round_nearest_even %[[CLAMP]] + // CHECK: %[[RESULT:.*]] = mhlo.convert %[[ROUND]] + // CHECK-SAME: (tensor<26x26x3x2xf32>) -> tensor<26x26x3x2xi32> + %0 = mhlo.uniform_quantize %arg0 : (tensor<26x26x3x2xf32> + ) -> tensor<26x26x3x2x!quant.uniform> + return %0 : tensor<26x26x3x2x!quant.uniform> +} + +// ----- + +// CHECK-LABEL: func @dequantize_per_channel +func.func @dequantize_per_channel( + %arg0: tensor<26x26x3x2x!quant.uniform> + ) -> tensor<26x26x3x2xf32> { + // CHECK-DAG: %[[SCALES:.*]] = mhlo.constant dense<[1.100000e+00, 1.100000e-01]> + // CHECK-DAG: %[[ZPS:.*]] = mhlo.constant dense<[-10, 2]> : tensor<2xi32> + // CHECK: %[[SUBTRACT:.*]] = chlo.broadcast_subtract + // CHECK-SAME: %[[INPUT:.*]], %[[ZPS]] + // CHECK-SAME: {broadcast_dimensions = dense<3> : tensor<1xi64>} + // CHECK-SAME: (tensor<26x26x3x2xi32>, tensor<2xi32>) -> tensor<26x26x3x2xi32> + // CHECK: %[[FLOAT:.*]] = mhlo.convert %[[SUBTRACT]] + // CHECK: %[[RESULT:.*]] = chlo.broadcast_multiply + // CHECK-SAME: %[[FLOAT]], %[[SCALES]] + // CHECK-SAME: {broadcast_dimensions = dense<3> : tensor<1xi64>} + // CHECK-SAME: (tensor<26x26x3x2xf32>, tensor<2xf32>) -> tensor<26x26x3x2xf32> + %0 = mhlo.uniform_dequantize %arg0 : ( + tensor<26x26x3x2x!quant.uniform> + ) -> tensor<26x26x3x2xf32> + return %0 : tensor<26x26x3x2xf32> +} + +// ----- + // CHECK-LABEL: func @add func.func @add( %arg0: tensor>, @@ -173,17 +220,11 @@ func.func @add_different_lhs_type( %arg0: tensor>, %arg1: tensor> ) -> tensor> { - // CHECK: %[[VAL1:.*]] = mhlo.convert %[[LHS:.*]] : (tensor) -> tensor - // CHECK-DAG: %[[INPUT_ZPS:.*]] = mhlo.constant dense<3> : tensor - // CHECK: %[[VAL2:.*]] = chlo.broadcast_subtract %[[VAL1]], %[[INPUT_ZPS]] : (tensor, tensor) -> tensor - // CHECK-DAG: %[[MULTIPLIER:.*]] = mhlo.constant dense<16384> : tensor - // CHECK-DAG: %[[TOTAL_SHIFT:.*]] = mhlo.constant dense<13> : tensor - // CHECK-DAG: %[[HALF:.*]] = mhlo.constant dense<4096> : tensor - // CHECK: %[[VAL3:.*]] = chlo.broadcast_multiply %[[VAL2]], %[[MULTIPLIER]] : (tensor, tensor) -> tensor - // CHECK: %[[VAL4:.*]] = chlo.broadcast_add %[[VAL3]], %[[HALF]] : (tensor, tensor) -> tensor - // CHECK: %[[VAL5:.*]] = chlo.broadcast_shift_right_arithmetic %[[VAL4]], %[[TOTAL_SHIFT]] : (tensor, tensor) -> tensor - // CHECK-DAG: %[[OUTPUT_ZPS:.*]] = mhlo.constant dense<1> : tensor - // CHECK: %[[LHS_32_REQ:.*]] = chlo.broadcast_add %[[VAL5]], %[[OUTPUT_ZPS]] : (tensor, tensor) -> tensor + // CHECK-DAG: %[[COMBINED_SCALE:.*]] = mhlo.constant dense<2.000000e+00> : tensor + // CHECK-DAG: %[[LHS:.*]] = mhlo.convert %arg0 : (tensor) -> tensor + // CHECK-DAG: %[[MUL:.*]] = chlo.broadcast_multiply %[[LHS]], %[[COMBINED_SCALE]] : (tensor, tensor) -> tensor + // CHECK-DAG: %[[COMBINED_ZP:.*]] = mhlo.constant dense<-5.000000e+00> + // CHECK: %[[LHS_32:.*]] = chlo.broadcast_add %[[MUL]], %[[COMBINED_ZP]] : (tensor, tensor) -> tensor // CHECK-DAG: %[[RHS_32:.*]] = mhlo.convert %[[RHS:.*]] : (tensor) -> tensor // CHECK-DAG: %[[RES_ZPS:.*]] = mhlo.constant dense<1> : tensor @@ -207,18 +248,11 @@ func.func @add_different_rhs_type( %arg0: tensor>, %arg1: tensor> ) -> tensor> { - // CHECK: %[[VAL0:.*]] = mhlo.convert %[[LHS:.*]] : (tensor) -> tensor - // CHECK: %[[VAL1:.*]] = mhlo.convert %[[RHS:.*]] : (tensor) -> tensor - // CHECK-DAG: %[[INPUT_ZPS:.*]] = mhlo.constant dense<3> : tensor - // CHECK: %[[VAL2:.*]] = chlo.broadcast_subtract %[[VAL1]], %[[INPUT_ZPS]] : (tensor, tensor) -> tensor - // CHECK-DAG: %[[MULTIPLIER:.*]] = mhlo.constant dense<16384> : tensor - // CHECK-DAG: %[[TOTAL_SHIFT:.*]] = mhlo.constant dense<13> : tensor - // CHECK-DAG: %[[HALF:.*]] = mhlo.constant dense<4096> : tensor - // CHECK: %[[VAL3:.*]] = chlo.broadcast_multiply %[[VAL2]], %[[MULTIPLIER]] : (tensor, tensor) -> tensor - // CHECK: %[[VAL4:.*]] = chlo.broadcast_add %[[VAL3]], %[[HALF]] : (tensor, tensor) -> tensor - // CHECK: %[[VAL5:.*]] = chlo.broadcast_shift_right_arithmetic %[[VAL4]], %[[TOTAL_SHIFT]] : (tensor, tensor) -> tensor - // CHECK-DAG: %[[OUTPUT_ZPS:.*]] = mhlo.constant dense<1> : tensor - // CHECK: %[[RHS_32_REQ:.*]] = chlo.broadcast_add %[[VAL5]], %[[OUTPUT_ZPS]] : (tensor, tensor) -> tensor + // CHECK-DAG: %[[COMBINED_SCALE:.*]] = mhlo.constant dense<2.000000e+00> : tensor + // CHECK-DAG: %[[RHS:.*]] = mhlo.convert %arg1 : (tensor) -> tensor + // CHECK-DAG: %[[MUL:.*]] = chlo.broadcast_multiply %[[RHS]], %[[COMBINED_SCALE]] : (tensor, tensor) -> tensor + // CHECK-DAG: %[[COMBINED_ZP:.*]] = mhlo.constant dense<-5.000000e+00> + // CHECK: %[[RHS_32:.*]] = chlo.broadcast_add %[[MUL]], %[[COMBINED_ZP]] : (tensor, tensor) -> tensor // CHECK-DAG: %[[RES_ZPS:.*]] = mhlo.constant dense<1> : tensor // CHECK-DAG: %[[VAL7:.*]] = chlo.broadcast_add %[[LHS_32:.*]], %[[RHS_32_REQ:.*]] : (tensor, tensor) -> tensor @@ -239,29 +273,17 @@ func.func @add_different_res_type( %arg0: tensor>, %arg1: tensor> ) -> tensor> { - // CHECK: %[[VAL1:.*]] = mhlo.convert %[[LHS:.*]] : (tensor) -> tensor - // CHECK-DAG: %[[INPUT_ZPS:.*]] = mhlo.constant dense<3> : tensor - // CHECK: %[[VAL2:.*]] = chlo.broadcast_subtract %[[VAL1]], %[[INPUT_ZPS]] : (tensor, tensor) -> tensor - // CHECK-DAG: %[[MULTIPLIER:.*]] = mhlo.constant dense<16384> : tensor - // CHECK-DAG: %[[TOTAL_SHIFT:.*]] = mhlo.constant dense<13> : tensor - // CHECK-DAG: %[[HALF:.*]] = mhlo.constant dense<4096> : tensor - // CHECK: %[[VAL3:.*]] = chlo.broadcast_multiply %[[VAL2]], %[[MULTIPLIER]] : (tensor, tensor) -> tensor - // CHECK: %[[VAL4:.*]] = chlo.broadcast_add %[[VAL3]], %[[HALF]] : (tensor, tensor) -> tensor - // CHECK: %[[VAL5:.*]] = chlo.broadcast_shift_right_arithmetic %[[VAL4]], %[[TOTAL_SHIFT]] : (tensor, tensor) -> tensor - // CHECK-DAG: %[[OUTPUT_ZPS:.*]] = mhlo.constant dense<1> : tensor - // CHECK: %[[LHS_32_REQ:.*]] = chlo.broadcast_add %[[VAL5]], %[[OUTPUT_ZPS]] : (tensor, tensor) -> tensor - - // CHECK: %[[VAL6:.*]] = mhlo.convert %[[RHS:.*]] : (tensor) -> tensor - // CHECK-DAG: %[[INPUT_ZPS:.*]] = mhlo.constant dense<3> : tensor - // CHECK: %[[VAL7:.*]] = chlo.broadcast_subtract %[[VAL6]], %[[INPUT_ZPS]] : (tensor, tensor) -> tensor - // CHECK-DAG: %[[MULTIPLIER:.*]] = mhlo.constant dense<16384> : tensor - // CHECK-DAG: %[[TOTAL_SHIFT:.*]] = mhlo.constant dense<13> : tensor - // CHECK-DAG: %[[HALF:.*]] = mhlo.constant dense<4096> : tensor - // CHECK: %[[VAL8:.*]] = chlo.broadcast_multiply %[[VAL7]], %[[MULTIPLIER]] : (tensor, tensor) -> tensor - // CHECK: %[[VAL9:.*]] = chlo.broadcast_add %[[VAL8]], %[[HALF]] : (tensor, tensor) -> tensor - // CHECK: %[[VAL10:.*]] = chlo.broadcast_shift_right_arithmetic %[[VAL9]], %[[TOTAL_SHIFT]] : (tensor, tensor) -> tensor - // CHECK-DAG: %[[OUTPUT_ZPS:.*]] = mhlo.constant dense<1> : tensor - // CHECK: %[[RHS_32_REQ:.*]] = chlo.broadcast_add %[[VAL10]], %[[OUTPUT_ZPS]] : (tensor, tensor) -> tensor + // CHECK-DAG: %[[COMBINED_SCALE:.*]] = mhlo.constant dense<2.000000e+00> : tensor + // CHECK-DAG: %[[LHS:.*]] = mhlo.convert %arg0 : (tensor) -> tensor + // CHECK-DAG: %[[MUL:.*]] = chlo.broadcast_multiply %[[LHS]], %[[COMBINED_SCALE]] : (tensor, tensor) -> tensor + // CHECK-DAG: %[[COMBINED_ZP:.*]] = mhlo.constant dense<-5.000000e+00> + // CHECK: %[[LHS_32_REQ:.*]] = chlo.broadcast_add %[[MUL]], %[[COMBINED_ZP]] : (tensor, tensor) -> tensor + + // CHECK-DAG: %[[COMBINED_SCALE:.*]] = mhlo.constant dense<2.000000e+00> : tensor + // CHECK-DAG: %[[RHS:.*]] = mhlo.convert %arg1 : (tensor) -> tensor + // CHECK-DAG: %[[MUL:.*]] = chlo.broadcast_multiply %[[RHS]], %[[COMBINED_SCALE]] : (tensor, tensor) -> tensor + // CHECK-DAG: %[[COMBINED_ZP:.*]] = mhlo.constant dense<-5.000000e+00> + // CHECK: %[[RHS_32_REQ:.*]] = chlo.broadcast_add %[[MUL]], %[[COMBINED_ZP]] : (tensor, tensor) -> tensor // CHECK-DAG: %[[RES_ZPS:.*]] = mhlo.constant dense<1> : tensor // CHECK-DAG: %[[VAL11:.*]] = chlo.broadcast_add %[[LHS_32_REQ:.*]], %[[RHS_32_REQ:.*]] : (tensor, tensor) -> tensor @@ -1413,6 +1435,17 @@ func.func @mhlo_constant_uniform_quantized() -> tensor<1x!quant.uniform> } +// ----- + +// CHECK-LABEL: func @mhlo_constant_uniform_quantized_per_channel +func.func @mhlo_constant_uniform_quantized_per_channel() -> () { + // CHECK: mhlo.constant dense<[9, 4]> : tensor<2xi8> + %0 = mhlo.constant() {value = dense<[9, 4]> : tensor<2xi8>} : () + -> tensor<2x!quant.uniform> + return +} + + // ----- // CHECK-LABEL: func @mhlo_constant_int diff --git a/tensorflow/compiler/mlir/quantization/stablehlo/tests/bridge/convert-tf-quant-types.mlir b/tensorflow/compiler/mlir/quantization/stablehlo/tests/bridge/convert-tf-quant-types.mlir index c83f95ce90cdcc..73555e6913b91d 100644 --- a/tensorflow/compiler/mlir/quantization/stablehlo/tests/bridge/convert-tf-quant-types.mlir +++ b/tensorflow/compiler/mlir/quantization/stablehlo/tests/bridge/convert-tf-quant-types.mlir @@ -11,17 +11,17 @@ func.func @relu_qint8(%arg0: tensor<1x!tf_type.qint8>) -> tensor<1x!tf_type.qint // CHECK-LABEL: func @if_qint8(%arg0: tensor, %arg1: tensor<1xi8>, %arg2: tensor<1xi8>) -> tensor<1xi8> func.func @if_qint8(%arg0: tensor, %arg1: tensor<1x!tf_type.qint8>, %arg2: tensor<1x!tf_type.qint8>) -> tensor<1x!tf_type.qint8> { - // CHECK-NEXT: %0 = "tf.IfRegion"(%arg0) ({ + // CHECK-NEXT: %0 = "tf.IfRegion"(%arg0) <{is_stateless = false}> ({ // CHECK-NEXT: "tf.Yield"(%arg1) : (tensor<1xi8>) -> () // CHECK-NEXT: }, { // CHECK-NEXT: "tf.Yield"(%arg2) : (tensor<1xi8>) -> () - // CHECK-NEXT: }) {is_stateless = false} : (tensor) -> tensor<1xi8> + // CHECK-NEXT: }) : (tensor) -> tensor<1xi8> // CHECK-NEXT: return %0 : tensor<1xi8> - %0 = "tf.IfRegion"(%arg0) ({ + %0 = "tf.IfRegion"(%arg0) <{is_stateless = false}> ({ "tf.Yield"(%arg1) : (tensor<1x!tf_type.qint8>) -> () }, { "tf.Yield"(%arg2) : (tensor<1x!tf_type.qint8>) -> () - }) {is_stateless = false} : (tensor) -> tensor<1x!tf_type.qint8> + }) : (tensor) -> tensor<1x!tf_type.qint8> func.return %0 : tensor<1x!tf_type.qint8> } @@ -74,7 +74,7 @@ func.func @uniform_quantize(%arg0: tensor<1xf32>) -> tensor<1x!tf_type.qint8> %zps = "tf.Const"() { value = dense<3> : tensor } : () -> tensor // CHECK: %[[qint:.*]] = "tf.UniformQuantize" - // CHECK: %[[int:.*]] = "tf.Cast"(%[[qint]]) {Truncate = false} : (tensor<1x!tf_type.qint8>) -> tensor<1xi8> + // CHECK: %[[int:.*]] = "tf.Cast"(%[[qint]]) <{Truncate = false}> : (tensor<1x!tf_type.qint8>) -> tensor<1xi8> %0 = "tf.UniformQuantize"(%arg0, %scales, %zps) { quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64 } : (tensor<1xf32>, tensor, tensor) -> tensor<1x!tf_type.qint8> @@ -92,7 +92,7 @@ func.func @uniform_quantize_no_return(%arg0: tensor<1xf32>) -> () %zps = "tf.Const"() { value = dense<3> : tensor } : () -> tensor // CHECK: %[[qint:.*]] = "tf.UniformQuantize" - // CHECK: %[[int:.*]] = "tf.Cast"(%[[qint]]) {Truncate = false} : (tensor<1x!tf_type.qint8>) -> tensor<1xi8> + // CHECK: %[[int:.*]] = "tf.Cast"(%[[qint]]) <{Truncate = false}> : (tensor<1x!tf_type.qint8>) -> tensor<1xi8> %0 = "tf.UniformQuantize"(%arg0, %scales, %zps) { quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64 } : (tensor<1xf32>, tensor, tensor) -> tensor<1x!tf_type.qint8> @@ -109,7 +109,7 @@ func.func @uniform_dequantize(%arg0: tensor<1x!tf_type.qint8>) -> tensor<1xf32> %scales = "tf.Const"() { value = dense<1.0> : tensor } : () -> tensor %zps = "tf.Const"() { value = dense<3> : tensor } : () -> tensor - // CHECK: %[[x:.*]] = "tf.Cast"(%arg0) {Truncate = false} : (tensor<1xi8>) -> tensor<1x!tf_type.qint8> + // CHECK: %[[x:.*]] = "tf.Cast"(%arg0) <{Truncate = false}> : (tensor<1xi8>) -> tensor<1x!tf_type.qint8> // CHECK: %[[y:.*]] = "tf.UniformDequantize"(%[[x]] %0 = "tf.UniformDequantize"(%arg0, %scales, %zps) { quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64 @@ -132,8 +132,8 @@ func.func @uniform_quantize_dequantize(%arg0: tensor<1xf32>) -> tensor<1xf32> quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64 } : (tensor<1xf32>, tensor, tensor) -> tensor<1x!tf_type.qint8> - // CHECK: %[[int:.*]] = "tf.Cast"(%[[qint0]]) {Truncate = false} : (tensor<1x!tf_type.qint8>) -> tensor<1xi8> - // CHECK: %[[qint1:.*]] = "tf.Cast"(%[[int]]) {Truncate = false} : (tensor<1xi8>) -> tensor<1x!tf_type.qint8> + // CHECK: %[[int:.*]] = "tf.Cast"(%[[qint0]]) <{Truncate = false}> : (tensor<1x!tf_type.qint8>) -> tensor<1xi8> + // CHECK: %[[qint1:.*]] = "tf.Cast"(%[[int]]) <{Truncate = false}> : (tensor<1xi8>) -> tensor<1x!tf_type.qint8> // CHECK: %[[res:.*]] = "tf.UniformDequantize"(%[[qint1]] %1 = "tf.UniformDequantize"(%0, %scales, %zps) { quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64 @@ -155,10 +155,10 @@ func.func @uniform_quantized_add(%arg0: tensor<2x!tf_type.qint32>, %arg1: tensor %output_scales = "tf.Const"() { value = dense<2.0> : tensor } : () -> tensor %output_zps = "tf.Const"() { value = dense<4> : tensor } : () -> tensor - // CHECK: %[[lhs:.*]] = "tf.Cast"(%arg0) {Truncate = false} : (tensor<2xi32>) -> tensor<2x!tf_type.qint32> - // CHECK: %[[rhs:.*]] = "tf.Cast"(%arg1) {Truncate = false} : (tensor<2xi32>) -> tensor<2x!tf_type.qint32> + // CHECK: %[[lhs:.*]] = "tf.Cast"(%arg0) <{Truncate = false}> : (tensor<2xi32>) -> tensor<2x!tf_type.qint32> + // CHECK: %[[rhs:.*]] = "tf.Cast"(%arg1) <{Truncate = false}> : (tensor<2xi32>) -> tensor<2x!tf_type.qint32> // CHECK: %[[res_qint:.*]] = "tf.UniformQuantizedAdd"(%[[lhs]], %[[rhs]] - // CHECK: %[[res_int:.*]] = "tf.Cast"(%[[res_qint]]) {Truncate = false} : (tensor<2x!tf_type.qint32>) -> tensor<2xi32> + // CHECK: %[[res_int:.*]] = "tf.Cast"(%[[res_qint]]) <{Truncate = false}> : (tensor<2x!tf_type.qint32>) -> tensor<2xi32> // CHECK: return %[[res_int]] : tensor<2xi32> %1 = "tf.UniformQuantizedAdd"( %arg0, %arg1, @@ -190,13 +190,13 @@ func.func @while_region_qint(%arg0: tensor<2x2xf32>) -> (tensor<2x?xf32>, tensor %zps4 = "tf.Const"() { value = dense<4> : tensor } : () -> tensor // CHECK: %[[qint_0:.*]] = "tf.UniformQuantize" - // CHECK: %[[int_0:.*]] = "tf.Cast"(%[[qint_0]]) {Truncate = false} : (tensor<2x2x!tf_type.qint8>) -> tensor<2x2xi8> + // CHECK: %[[int_0:.*]] = "tf.Cast"(%[[qint_0]]) <{Truncate = false}> : (tensor<2x2x!tf_type.qint8>) -> tensor<2x2xi8> %0 = "tf.UniformQuantize"(%arg0, %scales, %zps2) { quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64 } : (tensor<2x2xf32>, tensor, tensor) -> tensor<2x2x!tf_type.qint8> // CHECK: %[[qint_1:.*]] = "tf.UniformQuantize" - // CHECK: %[[int_1:.*]] = "tf.Cast"(%[[qint_1]]) {Truncate = false} : (tensor<2x2x!tf_type.qint8>) -> tensor<2x2xi8> + // CHECK: %[[int_1:.*]] = "tf.Cast"(%[[qint_1]]) <{Truncate = false}> : (tensor<2x2x!tf_type.qint8>) -> tensor<2x2xi8> %1 = "tf.UniformQuantize"(%arg0, %scales, %zps4) { quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64 } : (tensor<2x2xf32>, tensor, tensor) -> tensor<2x2x!tf_type.qint8> @@ -212,11 +212,11 @@ func.func @while_region_qint(%arg0: tensor<2x2xf32>) -> (tensor<2x?xf32>, tensor "tf.Yield"(%id, %barg1) : (tensor<2x?x!tf_type.qint8>, tensor) -> () }) {is_stateless = false} : (tensor<2x2x!tf_type.qint8>, tensor<2x2x!tf_type.qint8>) -> (tensor<2x?x!tf_type.qint8>, tensor) - // CHECK: %[[out_qint_0:.*]] = "tf.Cast"(%[[while_result]]#0) {Truncate = false} : (tensor<2x?xi8>) -> tensor<2x?x!tf_type.qint8> + // CHECK: %[[out_qint_0:.*]] = "tf.Cast"(%[[while_result]]#0) <{Truncate = false}> : (tensor<2x?xi8>) -> tensor<2x?x!tf_type.qint8> // CHECK: %[[out_f_0:.*]] = "tf.UniformDequantize"(%[[out_qint_0]] %3 = "tf.UniformDequantize"(%2#0, %scales, %zps2) {quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64} : (tensor<2x?x!tf_type.qint8>, tensor, tensor) -> tensor<2x?xf32> - // CHECK: %[[out_qint_1:.*]] = "tf.Cast"(%[[while_result]]#1) {Truncate = false} : (tensor) -> tensor + // CHECK: %[[out_qint_1:.*]] = "tf.Cast"(%[[while_result]]#1) <{Truncate = false}> : (tensor) -> tensor // CHECK: %[[out_f_1:.*]] = "tf.UniformDequantize"(%[[out_qint_1]] %4 = "tf.UniformDequantize"(%2#1, %scales, %zps4) {quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64} : (tensor, tensor, tensor) -> tensor @@ -234,7 +234,7 @@ func.func @concat_uniform_quantize(%arg0: tensor<3x3xf32>, %arg1: tensor<3x3xf32 // CHECK: %[[input:.*]] = "tf.ConcatV2"(%arg0, %arg1 // CHECK: %[[output_qint:.*]] = "tf.UniformQuantize"(%[[input]] - // CHECK: %[[output:.*]] = "tf.Cast"(%[[output_qint]]) {Truncate = false} : (tensor<6x3x!tf_type.qint8>) -> tensor<6x3xi8> + // CHECK: %[[output:.*]] = "tf.Cast"(%[[output_qint]]) <{Truncate = false}> : (tensor<6x3x!tf_type.qint8>) -> tensor<6x3xi8> // CHECK: return %[[output]] : tensor<6x3xi8> %0 = "tf.ConcatV2"(%arg0, %arg1, %axis) : (tensor<3x3xf32>, tensor<3x3xf32>, tensor) -> tensor<6x3xf32> %1 = "tf.UniformQuantize"(%0, %scales, %zps) { @@ -252,7 +252,7 @@ func.func @concat_uniform_dequantize(%arg0: tensor<3x3x!tf_type.qint8>, %arg1: t %zps = "tf.Const"() { value = dense<3> : tensor } : () -> tensor // CHECK: %[[input:.*]] = "tf.ConcatV2"(%arg0, %arg1, %[[VAL:.*]]) : (tensor<3x3xi8>, tensor<3x3xi8>, tensor) -> tensor<6x3xi8> - // CHECK: %[[input_qint:.*]] = "tf.Cast"(%[[input]]) {Truncate = false} : (tensor<6x3xi8>) -> tensor<6x3x!tf_type.qint8> + // CHECK: %[[input_qint:.*]] = "tf.Cast"(%[[input]]) <{Truncate = false}> : (tensor<6x3xi8>) -> tensor<6x3x!tf_type.qint8> // CHECK: %[[output:.*]] = "tf.UniformDequantize"(%[[input_qint]] // CHECK: return %[[output]] : tensor<6x3xf32> %0 = "tf.ConcatV2"(%arg0, %arg1, %axis) : (tensor<3x3x!tf_type.qint8>, tensor<3x3x!tf_type.qint8>, tensor) -> tensor<6x3x!tf_type.qint8> @@ -266,7 +266,7 @@ func.func @concat_uniform_dequantize(%arg0: tensor<3x3x!tf_type.qint8>, %arg1: t // CHECK-LABEL: func @tf_const_qint32 func.func @tf_const_qint32() -> tensor<1x!tf_type.qint32> { - // CHECK: %[[result:.*]] = "tf.Const"() {value = dense<127> : tensor<1xi32>} : () -> tensor<1xi32> + // CHECK: %[[result:.*]] = "tf.Const"() <{value = dense<127> : tensor<1xi32>}> : () -> tensor<1xi32> %0 = "tf.Const"() { value = #tf_type : tensor<1x!tf_type.qint32> } : () -> tensor<1x!tf_type.qint32> // CHECK: return %[[result]] : tensor<1xi32> func.return %0 : tensor<1x!tf_type.qint32> @@ -276,7 +276,7 @@ func.func @tf_const_qint32() -> tensor<1x!tf_type.qint32> { // CHECK-LABEL: func @tf_const_qint8 func.func @tf_const_qint8() -> tensor<2x!tf_type.qint8> { - // CHECK: %[[result:.*]] = "tf.Const"() {value = dense<[127, 18]> : tensor<2xi8>} : () -> tensor<2xi8> + // CHECK: %[[result:.*]] = "tf.Const"() <{value = dense<[127, 18]> : tensor<2xi8>}> : () -> tensor<2xi8> %0 = "tf.Const"() { value = #tf_type : tensor<2x!tf_type.qint8> } : () -> tensor<2x!tf_type.qint8> // CHECK: return %[[result]] : tensor<2xi8> func.return %0 : tensor<2x!tf_type.qint8> @@ -295,7 +295,7 @@ func.func @tf_const_invalid_proto() -> tensor<2x!tf_type.qint32> { // CHECK-LABEL: func @cast_op_qint32_int32 func.func @cast_op_qint32_int32(%arg0: tensor<1x!tf_type.qint32>) -> tensor<1xi32> { - // CHECK: "tf.Cast"(%arg0) {Truncate = false} : (tensor<1xi32>) -> tensor<1xi32> + // CHECK: "tf.Cast"(%arg0) <{Truncate = false}> : (tensor<1xi32>) -> tensor<1xi32> %0 = "tf.Cast"(%arg0) {Truncate = false} : (tensor<1x!tf_type.qint32>) -> tensor<1xi32> func.return %0: tensor<1xi32> } @@ -304,7 +304,7 @@ func.func @cast_op_qint32_int32(%arg0: tensor<1x!tf_type.qint32>) -> tensor<1xi3 // CHECK-LABEL: func @cast_op_int32_qint32 func.func @cast_op_int32_qint32(%arg0: tensor<1xi32>) -> tensor<1x!tf_type.qint32> { - // CHECK: "tf.Cast"(%arg0) {Truncate = false} : (tensor<1xi32>) -> tensor<1xi32> + // CHECK: "tf.Cast"(%arg0) <{Truncate = false}> : (tensor<1xi32>) -> tensor<1xi32> %0 = "tf.Cast"(%arg0) {Truncate = false} : (tensor<1xi32>) -> tensor<1x!tf_type.qint32> func.return %0: tensor<1x!tf_type.qint32> } @@ -313,7 +313,7 @@ func.func @cast_op_int32_qint32(%arg0: tensor<1xi32>) -> tensor<1x!tf_type.qint3 // CHECK-LABEL: func @cast_op_qint8_int8 func.func @cast_op_qint8_int8(%arg0: tensor<1x!tf_type.qint8>) -> tensor<1xi8> { - // CHECK: "tf.Cast"(%arg0) {Truncate = false} : (tensor<1xi8>) -> tensor<1xi8> + // CHECK: "tf.Cast"(%arg0) <{Truncate = false}> : (tensor<1xi8>) -> tensor<1xi8> %0 = "tf.Cast"(%arg0) {Truncate = false} : (tensor<1x!tf_type.qint8>) -> tensor<1xi8> func.return %0: tensor<1xi8> } @@ -322,7 +322,7 @@ func.func @cast_op_qint8_int8(%arg0: tensor<1x!tf_type.qint8>) -> tensor<1xi8> { // CHECK-LABEL: func @cast_op_int8_qint8 func.func @cast_op_int8_qint8(%arg0: tensor<1xi8>) -> tensor<1x!tf_type.qint8> { - // CHECK: "tf.Cast"(%arg0) {Truncate = false} : (tensor<1xi8>) -> tensor<1xi8> + // CHECK: "tf.Cast"(%arg0) <{Truncate = false}> : (tensor<1xi8>) -> tensor<1xi8> %0 = "tf.Cast"(%arg0) {Truncate = false} : (tensor<1xi8>) -> tensor<1x!tf_type.qint8> func.return %0: tensor<1x!tf_type.qint8> } @@ -331,7 +331,7 @@ func.func @cast_op_int8_qint8(%arg0: tensor<1xi8>) -> tensor<1x!tf_type.qint8> { // CHECK-LABEL: func @cast_op_qint32_int8 func.func @cast_op_qint32_int8(%arg0: tensor<1x!tf_type.qint32>) -> tensor<1xi8> { - // CHECK: "tf.Cast"(%arg0) {Truncate = false} : (tensor<1xi32>) -> tensor<1xi8> + // CHECK: "tf.Cast"(%arg0) <{Truncate = false}> : (tensor<1xi32>) -> tensor<1xi8> %0 = "tf.Cast"(%arg0) {Truncate = false} : (tensor<1x!tf_type.qint32>) -> tensor<1xi8> func.return %0: tensor<1xi8> } @@ -340,7 +340,7 @@ func.func @cast_op_qint32_int8(%arg0: tensor<1x!tf_type.qint32>) -> tensor<1xi8> // CHECK-LABEL: func @cast_op_int8_qint32 func.func @cast_op_int8_qint32(%arg0: tensor<1xi8>) -> tensor<1x!tf_type.qint32> { - // CHECK: "tf.Cast"(%arg0) {Truncate = false} : (tensor<1xi8>) -> tensor<1xi32> + // CHECK: "tf.Cast"(%arg0) <{Truncate = false}> : (tensor<1xi8>) -> tensor<1xi32> %0 = "tf.Cast"(%arg0) {Truncate = false} : (tensor<1xi8>) -> tensor<1x!tf_type.qint32> func.return %0: tensor<1x!tf_type.qint32> } @@ -353,10 +353,10 @@ func.func @cast_uniform_dequantize(%arg0: tensor<1x!tf_type.qint32>) -> tensor<1 %scales = "tf.Const"() { value = dense<1.0> : tensor } : () -> tensor %zps = "tf.Const"() { value = dense<3> : tensor } : () -> tensor - // CHECK: %[[x:.*]] = "tf.Cast"(%arg0) {Truncate = false} : (tensor<1xi32>) -> tensor<1xi8> + // CHECK: %[[x:.*]] = "tf.Cast"(%arg0) <{Truncate = false}> : (tensor<1xi32>) -> tensor<1xi8> %0 = "tf.Cast"(%arg0) {Truncate = false} : (tensor<1x!tf_type.qint32>) -> tensor<1x!tf_type.qint8> - // CHECK: %[[y:.*]] = "tf.Cast"(%[[x]]) {Truncate = false} : (tensor<1xi8>) -> tensor<1x!tf_type.qint8> + // CHECK: %[[y:.*]] = "tf.Cast"(%[[x]]) <{Truncate = false}> : (tensor<1xi8>) -> tensor<1x!tf_type.qint8> // CHECK: %[[z:.*]] = "tf.UniformDequantize"(%[[y]] %1 = "tf.UniformDequantize"(%0, %scales, %zps) { quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64 @@ -379,8 +379,8 @@ func.func @uniform_quantize_cast(%arg0: tensor<1xf32>) -> tensor<1x!tf_type.qint quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64 } : (tensor<1xf32>, tensor, tensor) -> tensor<1x!tf_type.qint8> - // CHECK: %1 = "tf.Cast"(%0) {Truncate = false} : (tensor<1x!tf_type.qint8>) -> tensor<1xi8> - // CHECK: %2 = "tf.Cast"(%1) {Truncate = false} : (tensor<1xi8>) -> tensor<1xi32> + // CHECK: %1 = "tf.Cast"(%0) <{Truncate = false}> : (tensor<1x!tf_type.qint8>) -> tensor<1xi8> + // CHECK: %2 = "tf.Cast"(%1) <{Truncate = false}> : (tensor<1xi8>) -> tensor<1xi32> %1 = "tf.Cast"(%0) {Truncate = false} : (tensor<1x!tf_type.qint8>) -> tensor<1x!tf_type.qint32> // CHECK: return %2 : tensor<1xi32> @@ -398,15 +398,15 @@ func.func @uniform_quantize_cast_dequantize(%arg0: tensor<1xf32>) -> tensor<1xf3 %zps1 = "tf.Const"() { value = dense<2> : tensor } : () -> tensor // CHECK: %[[qint_1:.*]] = "tf.UniformQuantize" - // CHECK: %[[int_1:.*]] = "tf.Cast"(%[[qint_1]]) {Truncate = false} : (tensor<1x!tf_type.qint8>) -> tensor<1xi8> + // CHECK: %[[int_1:.*]] = "tf.Cast"(%[[qint_1]]) <{Truncate = false}> : (tensor<1x!tf_type.qint8>) -> tensor<1xi8> %0 = "tf.UniformQuantize"(%arg0, %scales, %zps) { quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64 } : (tensor<1xf32>, tensor, tensor) -> tensor<1x!tf_type.qint8> - // CHECK: %[[int_2:.*]] = "tf.Cast"(%[[int_1]]) {Truncate = false} : (tensor<1xi8>) -> tensor<1xi32> + // CHECK: %[[int_2:.*]] = "tf.Cast"(%[[int_1]]) <{Truncate = false}> : (tensor<1xi8>) -> tensor<1xi32> %1 = "tf.Cast"(%0) {Truncate = false} : (tensor<1x!tf_type.qint8>) -> tensor<1x!tf_type.qint32> - // CHECK: %[[qint_2:.*]] = "tf.Cast"(%[[int_2]]) {Truncate = false} : (tensor<1xi32>) -> tensor<1x!tf_type.qint32> + // CHECK: %[[qint_2:.*]] = "tf.Cast"(%[[int_2]]) <{Truncate = false}> : (tensor<1xi32>) -> tensor<1x!tf_type.qint32> // CHECK: %[[int_3:.*]] = "tf.UniformDequantize"(%[[qint_2]] %2 = "tf.UniformDequantize"(%1, %scales1, %zps1) { quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64 @@ -423,10 +423,10 @@ func.func @uniform_quantize_clip_min_cast(%arg0: tensor<1x2x2x1x!tf_type.qint32> %scale = "tf.Const"() { value = dense<1.0> : tensor } : () -> tensor %zp = "tf.Const"() { value = dense<3> : tensor } : () -> tensor - // CHECK-DAG: %[[MIN_QINT:.*]] = "tf.Cast"(%arg1) {Truncate = false} : (tensor) -> tensor + // CHECK-DAG: %[[MIN_QINT:.*]] = "tf.Cast"(%arg1) <{Truncate = false}> : (tensor) -> tensor %q_min = "tf.Cast"(%arg1) {Truncate = false} : (tensor) -> tensor - // CHECK-DAG: %[[INPUT_QINT:.*]] = "tf.Cast"(%arg0) {Truncate = false} : (tensor<1x2x2x1xi32>) -> tensor<1x2x2x1x!tf_type.qint32> + // CHECK-DAG: %[[INPUT_QINT:.*]] = "tf.Cast"(%arg0) <{Truncate = false}> : (tensor<1x2x2x1xi32>) -> tensor<1x2x2x1x!tf_type.qint32> // CHECK: "tf.UniformQuantizedClipByValue"(%[[INPUT_QINT]], %[[MIN_QINT]], %[[MIN_QINT]] %output = "tf.UniformQuantizedClipByValue"(%arg0, %q_min, %q_min, %scale, %zp) {quantization_axis = -1 : i64, quantization_max_val = 2147483647 : i64, quantization_min_val = -2147483648 : i64} : @@ -441,11 +441,11 @@ func.func @uniform_quantize_clip_input_cast(%arg0: tensor<1x2x2x1xi32>, %arg1: t %scale = "tf.Const"() { value = dense<1.0> : tensor } : () -> tensor %zp = "tf.Const"() { value = dense<3> : tensor } : () -> tensor - // CHECK-DAG: %[[INPUT_QINT:.*]] = "tf.Cast"(%arg0) {Truncate = false} : (tensor<1x2x2x1xi32>) -> tensor<1x2x2x1x!tf_type.qint32> + // CHECK-DAG: %[[INPUT_QINT:.*]] = "tf.Cast"(%arg0) <{Truncate = false}> : (tensor<1x2x2x1xi32>) -> tensor<1x2x2x1x!tf_type.qint32> %q_input = "tf.Cast"(%arg0) {Truncate = false} : (tensor<1x2x2x1xi32>) -> tensor<1x2x2x1x!tf_type.qint32> - // CHECK-DAG: %[[MIN_QINT:.*]] = "tf.Cast"(%arg1) {Truncate = false} : (tensor) -> tensor - // CHECK-DAG: %[[MAX_QINT:.*]] = "tf.Cast"(%arg1) {Truncate = false} : (tensor) -> tensor + // CHECK-DAG: %[[MIN_QINT:.*]] = "tf.Cast"(%arg1) <{Truncate = false}> : (tensor) -> tensor + // CHECK-DAG: %[[MAX_QINT:.*]] = "tf.Cast"(%arg1) <{Truncate = false}> : (tensor) -> tensor // CHECK: "tf.UniformQuantizedClipByValue"(%[[INPUT_QINT]], %[[MIN_QINT]], %[[MAX_QINT]] %output = "tf.UniformQuantizedClipByValue"(%q_input, %arg1, %arg1, %scale, %zp) {quantization_axis = -1 : i64, quantization_max_val = 2147483647 : i64, quantization_min_val = -2147483648 : i64} : @@ -460,15 +460,15 @@ func.func @uniform_quantize_clip_output_cast(%arg0: tensor<1x2x2x1x!tf_type.qint %scale = "tf.Const"() { value = dense<1.0> : tensor } : () -> tensor %zp = "tf.Const"() { value = dense<3> : tensor } : () -> tensor - // CHECK-DAG: %[[INPUT_QINT:.*]] = "tf.Cast"(%arg0) {Truncate = false} : (tensor<1x2x2x1xi32>) -> tensor<1x2x2x1x!tf_type.qint32> - // CHECK-DAG: %[[MIN_QINT:.*]] = "tf.Cast"(%arg1) {Truncate = false} : (tensor) -> tensor - // CHECK-DAG: %[[MAX_QINT:.*]] = "tf.Cast"(%arg1) {Truncate = false} : (tensor) -> tensor + // CHECK-DAG: %[[INPUT_QINT:.*]] = "tf.Cast"(%arg0) <{Truncate = false}> : (tensor<1x2x2x1xi32>) -> tensor<1x2x2x1x!tf_type.qint32> + // CHECK-DAG: %[[MIN_QINT:.*]] = "tf.Cast"(%arg1) <{Truncate = false}> : (tensor) -> tensor + // CHECK-DAG: %[[MAX_QINT:.*]] = "tf.Cast"(%arg1) <{Truncate = false}> : (tensor) -> tensor // CHECK: %[[OUTPUT_QINT:.*]] = "tf.UniformQuantizedClipByValue"(%[[INPUT_QINT]], %[[MIN_QINT]], %[[MAX_QINT]] %q_output = "tf.UniformQuantizedClipByValue"(%arg0, %arg1, %arg1, %scale, %zp) {quantization_axis = -1 : i64, quantization_max_val = 2147483647 : i64, quantization_min_val = -2147483648 : i64} : (tensor<1x2x2x1x!tf_type.qint32>, tensor, tensor, tensor, tensor) -> tensor<1x2x2x1x!tf_type.qint32> - // CHECK: %[[OUTPUT:.*]] = "tf.Cast"(%[[OUTPUT_QINT]]) {Truncate = false} : (tensor<1x2x2x1x!tf_type.qint32>) -> tensor<1x2x2x1xi32> + // CHECK: %[[OUTPUT:.*]] = "tf.Cast"(%[[OUTPUT_QINT]]) <{Truncate = false}> : (tensor<1x2x2x1x!tf_type.qint32>) -> tensor<1x2x2x1xi32> %output = "tf.Cast"(%q_output) {Truncate = false} : (tensor<1x2x2x1x!tf_type.qint32>) -> tensor<1x2x2x1xi32> return %output : tensor<1x2x2x1xi32> @@ -481,19 +481,19 @@ func.func @uniform_quantize_clip_output_cast_multiple_uses(%arg0: tensor<1x2x2x1 %scale = "tf.Const"() { value = dense<1.0> : tensor } : () -> tensor %zp = "tf.Const"() { value = dense<3> : tensor } : () -> tensor - // CHECK-DAG: %[[INPUT_QINT:.*]] = "tf.Cast"(%arg0) {Truncate = false} : (tensor<1x2x2x1xi32>) -> tensor<1x2x2x1x!tf_type.qint32> - // CHECK-DAG: %[[MIN_QINT:.*]] = "tf.Cast"(%arg1) {Truncate = false} : (tensor) -> tensor - // CHECK-DAG: %[[MAX_QINT:.*]] = "tf.Cast"(%arg1) {Truncate = false} : (tensor) -> tensor + // CHECK-DAG: %[[INPUT_QINT:.*]] = "tf.Cast"(%arg0) <{Truncate = false}> : (tensor<1x2x2x1xi32>) -> tensor<1x2x2x1x!tf_type.qint32> + // CHECK-DAG: %[[MIN_QINT:.*]] = "tf.Cast"(%arg1) <{Truncate = false}> : (tensor) -> tensor + // CHECK-DAG: %[[MAX_QINT:.*]] = "tf.Cast"(%arg1) <{Truncate = false}> : (tensor) -> tensor // CHECK: %[[OUTPUT_QINT:.*]] = "tf.UniformQuantizedClipByValue"(%[[INPUT_QINT]], %[[MIN_QINT]], %[[MAX_QINT]] %q_output = "tf.UniformQuantizedClipByValue"(%arg0, %arg1, %arg1, %scale, %zp) {quantization_axis = -1 : i64, quantization_max_val = 2147483647 : i64, quantization_min_val = -2147483648 : i64} : (tensor<1x2x2x1x!tf_type.qint32>, tensor, tensor, tensor, tensor) -> tensor<1x2x2x1x!tf_type.qint32> - // CHECK-DAG: %[[OUTPUT_1:.*]] = "tf.Cast"(%[[OUTPUT_QINT]]) {Truncate = false} : (tensor<1x2x2x1x!tf_type.qint32>) -> tensor<1x2x2x1xi32> + // CHECK-DAG: %[[OUTPUT_1:.*]] = "tf.Cast"(%[[OUTPUT_QINT]]) <{Truncate = false}> : (tensor<1x2x2x1x!tf_type.qint32>) -> tensor<1x2x2x1xi32> %output = "tf.Cast"(%q_output) {Truncate = false} : (tensor<1x2x2x1x!tf_type.qint32>) -> tensor<1x2x2x1xi32> - // CHECK-DAG: %[[OUTPUT_2:.*]] = "tf.Cast"(%[[OUTPUT_QINT]]) {Truncate = false} : (tensor<1x2x2x1x!tf_type.qint32>) -> tensor<1x2x2x1xi32> - // CHECK-DAG: %[[OUTPUT_QINT_1:.*]] = "tf.Cast"(%[[OUTPUT_1]]) {Truncate = false} : (tensor<1x2x2x1xi32>) -> tensor<1x2x2x1x!tf_type.qint32> + // CHECK-DAG: %[[OUTPUT_2:.*]] = "tf.Cast"(%[[OUTPUT_QINT]]) <{Truncate = false}> : (tensor<1x2x2x1x!tf_type.qint32>) -> tensor<1x2x2x1xi32> + // CHECK-DAG: %[[OUTPUT_QINT_1:.*]] = "tf.Cast"(%[[OUTPUT_1]]) <{Truncate = false}> : (tensor<1x2x2x1xi32>) -> tensor<1x2x2x1x!tf_type.qint32> // CHECK: "tf.UniformDequantize"(%[[OUTPUT_QINT_1:.*]] %dq = "tf.UniformDequantize"(%q_output, %scale, %zp) { quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64 diff --git a/tensorflow/compiler/mlir/quantization/stablehlo/tests/lift_quantizable_spots_as_functions.mlir b/tensorflow/compiler/mlir/quantization/stablehlo/tests/lift_quantizable_spots_as_functions.mlir index 3e588f5aeef1bc..f085a38f952c38 100644 --- a/tensorflow/compiler/mlir/quantization/stablehlo/tests/lift_quantizable_spots_as_functions.mlir +++ b/tensorflow/compiler/mlir/quantization/stablehlo/tests/lift_quantizable_spots_as_functions.mlir @@ -1,13 +1,14 @@ // RUN: stablehlo-quant-opt %s -split-input-file -stablehlo-lift-quantizable-spots-as-functions | FileCheck %s // CHECK-LABEL: @conv_fn( -// CHECK-SAME: %[[ARG_0:.*]]: tensor<1x3x3x4xf32>, -// CHECK-SAME: %[[ARG_1:.*]]: tensor<3x3x4x4xf32>) -func.func @conv_fn(%arg0: tensor<1x3x3x4xf32>, %arg1: tensor<3x3x4x4xf32>) -> tensor<1x3x3x4xf32> { - %0 = stablehlo.convolution(%arg0, %arg1) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {pad = [[1, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x3x3x4xf32>, tensor<3x3x4x4xf32>) -> tensor<1x3x3x4xf32> - func.return %0: tensor<1x3x3x4xf32> +// CHECK-SAME: %[[ARG_0:.*]]: tensor<1x3x3x4xf32> +func.func @conv_fn(%arg0: tensor<1x3x3x4xf32>) -> tensor<1x3x3x4xf32> { + %0 = stablehlo.constant dense<2.000000e+00> : tensor<3x3x4x4xf32> + %1 = stablehlo.convolution(%arg0, %0) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {pad = [[1, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x3x3x4xf32>, tensor<3x3x4x4xf32>) -> tensor<1x3x3x4xf32> + func.return %1: tensor<1x3x3x4xf32> } -// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %arg1) +// CHECK: %[[CONST:.*]] = stablehlo.constant dense<2.000000e+00> +// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %[[CONST]]) // CHECK: return %[[XLA_CALL_MODULE:.*]] : tensor<1x3x3x4xf32> // CHECK: } @@ -19,13 +20,14 @@ func.func @conv_fn(%arg0: tensor<1x3x3x4xf32>, %arg1: tensor<3x3x4x4xf32>) -> te // ----- // CHECK-LABEL: @dot_general_fn( -// CHECK-SAME: %[[ARG_0:.*]]: tensor<1x1x167xf32>, -// CHECK-SAME: %[[ARG_1:.*]]: tensor<167x64xf32> -func.func @dot_general_fn(%arg0: tensor<1x1x167xf32>, %arg1: tensor<167x64xf32>) -> tensor<1x1x64xf32> { - %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [2] x [0], precision = [DEFAULT, DEFAULT] : (tensor<1x1x167xf32>, tensor<167x64xf32>) -> tensor<1x1x64xf32> - return %0 : tensor<1x1x64xf32> +// CHECK-SAME: %[[ARG_0:.*]]: tensor<1x1x167xf32> +func.func @dot_general_fn(%arg0: tensor<1x1x167xf32>) -> tensor<1x1x64xf32> { + %0 = stablehlo.constant dense<2.000000e+00> : tensor<167x64xf32> + %1 = stablehlo.dot_general %arg0, %0, contracting_dims = [2] x [0], precision = [DEFAULT, DEFAULT] : (tensor<1x1x167xf32>, tensor<167x64xf32>) -> tensor<1x1x64xf32> + return %1 : tensor<1x1x64xf32> } -// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %arg1) +// CHECK: %[[CONST:.*]] = stablehlo.constant dense<2.000000e+00> +// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %[[CONST]]) // CHECK: return %[[XLA_CALL_MODULE:.*]] : tensor<1x1x64xf32> // CHECK: } @@ -37,15 +39,17 @@ func.func @dot_general_fn(%arg0: tensor<1x1x167xf32>, %arg1: tensor<167x64xf32>) // ----- // CHECK-LABEL: @conv_with_bias_fn( -// CHECK-SAME: %[[ARG_0:.*]]: tensor<1x3x3x4xf32>, -// CHECK-SAME: %[[ARG_1:.*]]: tensor<3x3x4x4xf32>, -// CHECK-SAME: %[[ARG_2:.*]]: tensor<1x3x3x4xf32>) -func.func @conv_with_bias_fn(%arg0: tensor<1x3x3x4xf32>, %arg1: tensor<3x3x4x4xf32>, %arg2: tensor<1x3x3x4xf32>) -> tensor<1x3x3x4xf32> { - %0 = stablehlo.convolution(%arg0, %arg1) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {pad = [[1, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x3x3x4xf32>, tensor<3x3x4x4xf32>) -> tensor<1x3x3x4xf32> - %1 = stablehlo.add %0, %arg2 : tensor<1x3x3x4xf32> - func.return %1: tensor<1x3x3x4xf32> +// CHECK-SAME: %[[ARG_0:.*]]: tensor<1x3x3x4xf32> +func.func @conv_with_bias_fn(%arg0: tensor<1x3x3x4xf32>) -> tensor<1x3x3x4xf32> { + %0 = stablehlo.constant dense<2.000000e+00> : tensor<3x3x4x4xf32> + %1 = stablehlo.constant dense<2.000000e+00> : tensor<1x3x3x4xf32> + %2 = stablehlo.convolution(%arg0, %0) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {pad = [[1, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x3x3x4xf32>, tensor<3x3x4x4xf32>) -> tensor<1x3x3x4xf32> + %3 = stablehlo.add %2, %1 : tensor<1x3x3x4xf32> + func.return %3: tensor<1x3x3x4xf32> } -// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %arg1, %arg2) +// CHECK: %[[CONST_0:.*]] = stablehlo.constant dense<2.000000e+00> +// CHECK: %[[CONST_1:.*]] = stablehlo.constant dense<2.000000e+00> +// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %[[CONST_0]], %[[CONST_1]]) // CHECK: return %[[XLA_CALL_MODULE:.*]] : tensor<1x3x3x4xf32> // CHECK: } @@ -58,15 +62,17 @@ func.func @conv_with_bias_fn(%arg0: tensor<1x3x3x4xf32>, %arg1: tensor<3x3x4x4xf // ----- // CHECK-LABEL: @dot_general_with_bias_fn( -// CHECK-SAME: %[[ARG_0:.*]]: tensor<1x1x167xf32>, -// CHECK-SAME: %[[ARG_1:.*]]: tensor<167x64xf32> -// CHECK-SAME: %[[ARG_2:.*]]: tensor<1x1x64xf32>) -func.func @dot_general_with_bias_fn(%arg0: tensor<1x1x167xf32>, %arg1: tensor<167x64xf32>, %arg2: tensor<1x1x64xf32>) -> tensor<1x1x64xf32> { - %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [2] x [0], precision = [DEFAULT, DEFAULT] : (tensor<1x1x167xf32>, tensor<167x64xf32>) -> tensor<1x1x64xf32> - %1 = stablehlo.add %0, %arg2 : tensor<1x1x64xf32> - func.return %1: tensor<1x1x64xf32> +// CHECK-SAME: %[[ARG_0:.*]]: tensor<1x1x167xf32> +func.func @dot_general_with_bias_fn(%arg0: tensor<1x1x167xf32>) -> tensor<1x1x64xf32> { + %0 = stablehlo.constant dense<2.000000e+00> : tensor<167x64xf32> + %1 = stablehlo.constant dense<2.000000e+00> : tensor<1x1x64xf32> + %2 = stablehlo.dot_general %arg0, %0, contracting_dims = [2] x [0], precision = [DEFAULT, DEFAULT] : (tensor<1x1x167xf32>, tensor<167x64xf32>) -> tensor<1x1x64xf32> + %3 = stablehlo.add %2, %1 : tensor<1x1x64xf32> + func.return %3: tensor<1x1x64xf32> } -// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %arg1, %arg2) +// CHECK: %[[CONST_0:.*]] = stablehlo.constant dense<2.000000e+00> +// CHECK: %[[CONST_1:.*]] = stablehlo.constant dense<2.000000e+00> +// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %[[CONST_0]], %[[CONST_1]]) // CHECK: return %[[XLA_CALL_MODULE:.*]] : tensor<1x1x64xf32> // CHECK: } @@ -78,16 +84,71 @@ func.func @dot_general_with_bias_fn(%arg0: tensor<1x1x167xf32>, %arg1: tensor<16 // ----- +// CHECK-LABEL: @conv_with_bias_dynamic_fn( +// CHECK-SAME: %[[ARG_0:.*]]: tensor +func.func @conv_with_bias_dynamic_fn(%arg0: tensor) -> tensor { + %0 = stablehlo.constant dense<2.000000e+00> : tensor<3x3x1x16xf32> + %1 = stablehlo.constant dense<2.000000e+00> : tensor<16xf32> + %2 = stablehlo.convolution(%arg0, %0) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {stride = [1, 1], pad = [[1, 1], [1, 1]], rhs_dilate = [1, 1]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64, precision_config = [#stablehlo, #stablehlo]} : (tensor, tensor<3x3x1x16xf32>) -> tensor + %3 = shape.shape_of %2 : tensor -> tensor<4xindex> + %4 = stablehlo.dynamic_broadcast_in_dim %1, %3, dims = [3] : (tensor<16xf32>, tensor<4xindex>) -> tensor + %5 = stablehlo.add %2, %4 : tensor + func.return %5: tensor +} +// CHECK: %[[CONST_0:.*]] = stablehlo.constant dense<2.000000e+00> +// CHECK: %[[CONST_1:.*]] = stablehlo.constant dense<2.000000e+00> +// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %[[CONST_0]], %[[CONST_1]]) +// CHECK: return %[[XLA_CALL_MODULE:.*]] : tensor +// CHECK: } + +// CHECK-LABEL: private @composite_conv_with_bias_dynamic_fn_1 +// CHECK: %[[CONV:.*]] = stablehlo.convolution(%arg0, %arg1) +// CHECK: %[[SHAPE_OF:.*]] = shape.shape_of %[[CONV]] +// CHECK: %[[DYNAMIC_BROADCAST_IN_DIM:.*]] = stablehlo.dynamic_broadcast_in_dim %arg2, %[[SHAPE_OF]] +// CHECK: %[[ADD:.*]] = stablehlo.add %[[CONV]], %[[DYNAMIC_BROADCAST_IN_DIM]] +// CHECK: return %[[ADD]] : tensor +// CHECK: } + +// ----- + +// CHECK-LABEL: @dot_general_with_bias_dynamic_fn( +// CHECK-SAME: %[[ARG_0:.*]]: tensor +func.func @dot_general_with_bias_dynamic_fn(%arg0: tensor) -> tensor { + %0 = stablehlo.constant dense<2.000000e+00> : tensor<12544x10xf32> + %1 = stablehlo.constant dense<2.000000e+00> : tensor<10xf32> + %2 = stablehlo.dot_general %arg0, %0, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor, tensor<12544x10xf32>) -> tensor + %3 = shape.shape_of %2 : tensor -> tensor<2xindex> + %4 = stablehlo.dynamic_broadcast_in_dim %1, %3, dims = [1] : (tensor<10xf32>, tensor<2xindex>) -> tensor + %5 = stablehlo.add %2, %4 : tensor + func.return %5: tensor +} +// CHECK: %[[CONST_0:.*]] = stablehlo.constant dense<2.000000e+00> +// CHECK: %[[CONST_1:.*]] = stablehlo.constant dense<2.000000e+00> +// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %[[CONST_0]], %[[CONST_1]]) +// CHECK: return %[[XLA_CALL_MODULE:.*]] : tensor +// CHECK: } + +// CHECK-LABEL: private @composite_dot_general_with_bias_dynamic_fn_1 +// CHECK: %[[DOT_GENERAL:.*]] = stablehlo.dot_general %arg0, %arg1 +// CHECK: %[[SHAPE_OF_0:.*]] = shape.shape_of %[[DOT_GENERAL]] +// CHECK: %[[DYNAMIC_BROADCAST_IN_DIM_0:.*]] = stablehlo.dynamic_broadcast_in_dim %arg2, %[[SHAPE_OF_0]] +// CHECK: %[[ADD:.*]] = stablehlo.add %[[DOT_GENERAL]], %[[DYNAMIC_BROADCAST_IN_DIM_0]] +// CHECK: return %[[ADD]] : tensor +// CHECK: } + +// ----- + // CHECK-LABEL: @conv_with_relu_fn( -// CHECK-SAME: %[[ARG_0:.*]]: tensor<1x3x3x4xf32>, -// CHECK-SAME: %[[ARG_1:.*]]: tensor<3x3x4x4xf32>) -func.func @conv_with_relu_fn(%arg0: tensor<1x3x3x4xf32>, %arg1: tensor<3x3x4x4xf32>) -> tensor<1x3x3x4xf32> { - %0 = stablehlo.constant dense<0.000000e+00> : tensor<1x3x3x4xf32> - %1 = stablehlo.convolution(%arg0, %arg1) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {pad = [[1, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x3x3x4xf32>, tensor<3x3x4x4xf32>) -> tensor<1x3x3x4xf32> - %2 = stablehlo.maximum %1, %0 : tensor<1x3x3x4xf32> - func.return %2: tensor<1x3x3x4xf32> +// CHECK-SAME: %[[ARG_0:.*]]: tensor<1x3x3x4xf32> +func.func @conv_with_relu_fn(%arg0: tensor<1x3x3x4xf32>) -> tensor<1x3x3x4xf32> { + %0 = stablehlo.constant dense<2.000000e+00> : tensor<3x3x4x4xf32> + %1 = stablehlo.constant dense<0.000000e+00> : tensor<1x3x3x4xf32> + %2 = stablehlo.convolution(%arg0, %0) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {pad = [[1, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x3x3x4xf32>, tensor<3x3x4x4xf32>) -> tensor<1x3x3x4xf32> + %3 = stablehlo.maximum %2, %1 : tensor<1x3x3x4xf32> + func.return %3: tensor<1x3x3x4xf32> } -// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %arg1) +// CHECK: %[[CONST:.*]] = stablehlo.constant dense<2.000000e+00> +// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %[[CONST]]) // CHECK: return %[[XLA_CALL_MODULE:.*]] : tensor<1x3x3x4xf32> // CHECK: } @@ -102,14 +163,15 @@ func.func @conv_with_relu_fn(%arg0: tensor<1x3x3x4xf32>, %arg1: tensor<3x3x4x4xf // CHECK-LABEL: @dot_general_with_relu_fn( // CHECK-SAME: %[[ARG_0:.*]]: tensor<1x1x167xf32>, -// CHECK-SAME: %[[ARG_1:.*]]: tensor<167x64xf32> func.func @dot_general_with_relu_fn(%arg0: tensor<1x1x167xf32>, %arg1: tensor<167x64xf32>) -> tensor<1x1x64xf32> { - %0 = stablehlo.constant dense<0.000000e+00> : tensor<1x1x64xf32> - %1 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [2] x [0], precision = [DEFAULT, DEFAULT] : (tensor<1x1x167xf32>, tensor<167x64xf32>) -> tensor<1x1x64xf32> - %2 = stablehlo.maximum %1, %0 : tensor<1x1x64xf32> - return %2 : tensor<1x1x64xf32> + %0 = stablehlo.constant dense<2.000000e+00> : tensor<167x64xf32> + %1 = stablehlo.constant dense<0.000000e+00> : tensor<1x1x64xf32> + %2 = stablehlo.dot_general %arg0, %0, contracting_dims = [2] x [0], precision = [DEFAULT, DEFAULT] : (tensor<1x1x167xf32>, tensor<167x64xf32>) -> tensor<1x1x64xf32> + %3 = stablehlo.maximum %2, %1 : tensor<1x1x64xf32> + return %3 : tensor<1x1x64xf32> } -// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %arg1) +// CHECK: %[[CONST:.*]] = stablehlo.constant dense<2.000000e+00> +// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %[[CONST]]) // CHECK: return %[[XLA_CALL_MODULE:.*]] : tensor<1x1x64xf32> // CHECK: } @@ -122,36 +184,95 @@ func.func @dot_general_with_relu_fn(%arg0: tensor<1x1x167xf32>, %arg1: tensor<16 // ----- +// CHECK-LABEL: @conv_with_relu_dynamic_fn( +// CHECK-SAME: %[[ARG_0:.*]]: tensor +func.func @conv_with_relu_dynamic_fn(%arg0: tensor) -> tensor { + %0 = stablehlo.constant dense<2.000000e+00> : tensor<3x3x1x16xf32> + %1 = stablehlo.constant dense<0.000000e+00> : tensor + %2 = stablehlo.convolution(%arg0, %0) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {stride = [1, 1], pad = [[1, 1], [1, 1]], rhs_dilate = [1, 1]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64, precision_config = [#stablehlo, #stablehlo]} : (tensor, tensor<3x3x1x16xf32>) -> tensor + %3 = shape.shape_of %2 : tensor -> tensor<4xindex> + %4 = stablehlo.dynamic_broadcast_in_dim %1, %3, dims = [] : (tensor, tensor<4xindex>) -> tensor + %5 = stablehlo.maximum %2, %4 : tensor + func.return %5: tensor +} +// CHECK: %[[CONST:.*]] = stablehlo.constant dense<2.000000e+00> +// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %[[CONST]]) +// CHECK: return %[[XLA_CALL_MODULE:.*]] : tensor +// CHECK: } + +// CHECK-LABEL: private @composite_conv_with_relu_dynamic_fn_1 +// CHECK: %[[CONST:.*]] = stablehlo.constant dense<0.000000e+00> +// CHECK: %[[CONV:.*]] = stablehlo.convolution(%arg0, %arg1) +// CHECK: %[[SHAPE_OF:.*]] = shape.shape_of %[[CONV]] +// CHECK: %[[DYNAMIC_BROADCAST_IN_DIM:.*]] = stablehlo.dynamic_broadcast_in_dim %[[CONST]], %[[SHAPE_OF]] +// CHECK: %[[MAX:.*]] = stablehlo.maximum %[[CONV]], %[[DYNAMIC_BROADCAST_IN_DIM]] +// CHECK: return %[[MAX]] : tensor +// CHECK: } + +// ----- + +// CHECK-LABEL: @dot_general_with_relu_dynamic_fn( +// CHECK-SAME: %[[ARG_0:.*]]: tensor +func.func @dot_general_with_relu_dynamic_fn(%arg0: tensor) -> tensor { + %0 = stablehlo.constant dense<2.000000e+00> : tensor<12544x10xf32> + %1 = stablehlo.constant dense<0.000000e+00> : tensor + %2 = stablehlo.dot_general %arg0, %0, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor, tensor<12544x10xf32>) -> tensor + %3 = shape.shape_of %2 : tensor -> tensor<2xindex> + %4 = stablehlo.dynamic_broadcast_in_dim %1, %3, dims = [] : (tensor, tensor<2xindex>) -> tensor + %5 = stablehlo.maximum %2, %4 : tensor + func.return %5: tensor +} +// CHECK: %[[CONST:.*]] = stablehlo.constant dense<2.000000e+00> +// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %[[CONST]]) +// CHECK: return %[[XLA_CALL_MODULE:.*]] : tensor +// CHECK: } + +// CHECK-LABEL: private @composite_dot_general_with_relu_dynamic_fn_1 +// CHECK: %[[CONST:.*]] = stablehlo.constant dense<0.000000e+00> +// CHECK: %[[DOT_GENERAL:.*]] = stablehlo.dot_general %arg0, %arg1 +// CHECK: %[[SHAPE_OF:.*]] = shape.shape_of %[[DOT_GENERAL]] +// CHECK: %[[DYNAMIC_BROADCAST_IN_DIM:.*]] = stablehlo.dynamic_broadcast_in_dim %[[CONST]], %[[SHAPE_OF]] +// CHECK: %[[MAX:.*]] = stablehlo.maximum %[[DOT_GENERAL]], %[[DYNAMIC_BROADCAST_IN_DIM]] +// CHECK: return %[[MAX]] : tensor +// CHECK: } + +// ----- + // The pattern should not match when the const value for relu is not 0. // CHECK-LABEL: @conv_with_relu_wrong_const_fn( -// CHECK-SAME: %[[ARG_0:.*]]: tensor<1x3x3x4xf32>, -// CHECK-SAME: %[[ARG_1:.*]]: tensor<3x3x4x4xf32>) +// CHECK-SAME: %[[ARG_0:.*]]: tensor<1x3x3x4xf32> func.func @conv_with_relu_wrong_const_fn(%arg0: tensor<1x3x3x4xf32>, %arg1: tensor<3x3x4x4xf32>) -> tensor<1x3x3x4xf32> { - %0 = stablehlo.constant dense<2.000000e+00> : tensor<1x3x3x4xf32> - %1 = stablehlo.convolution(%arg0, %arg1) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {pad = [[1, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x3x3x4xf32>, tensor<3x3x4x4xf32>) -> tensor<1x3x3x4xf32> - %2 = stablehlo.maximum %1, %0 : tensor<1x3x3x4xf32> - func.return %2: tensor<1x3x3x4xf32> + %0 = stablehlo.constant dense<2.000000e+00> : tensor<3x3x4x4xf32> + %1 = stablehlo.constant dense<2.000000e+00> : tensor<1x3x3x4xf32> + %2 = stablehlo.convolution(%arg0, %0) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {pad = [[1, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x3x3x4xf32>, tensor<3x3x4x4xf32>) -> tensor<1x3x3x4xf32> + %3 = stablehlo.maximum %2, %1 : tensor<1x3x3x4xf32> + func.return %3: tensor<1x3x3x4xf32> } -// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %arg1) -// CHECK: return %[[XLA_CALL_MODULE:.*]] : tensor<1x3x3x4xf32> +// CHECK: %[[CONST_0:.*]] = stablehlo.constant dense<2.000000e+00> +// CHECK: %[[CONST_1:.*]] = stablehlo.constant dense<2.000000e+00> +// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %[[CONST_0]]) +// CHECK: %[[MAX:.*]] = stablehlo.maximum %[[XLA_CALL_MODULE]], %[[CONST_1]] +// CHECK: return %[[MAX]] : tensor<1x3x3x4xf32> // CHECK: } // CHECK-LABEL: private @composite_conv_fn_1 +// CHECK-NOT: private @composite_conv_with_relu_fn_1 // ----- // CHECK-LABEL: @conv_with_relu6_fn( -// CHECK-SAME: %[[ARG_0:.*]]: tensor<1x3x3x4xf32>, -// CHECK-SAME: %[[ARG_1:.*]]: tensor<3x3x4x4xf32>) -func.func @conv_with_relu6_fn(%arg0: tensor<1x3x3x4xf32>, %arg1: tensor<3x3x4x4xf32>) -> tensor<1x3x3x4xf32> { - %0 = stablehlo.constant dense<0.000000e+00> : tensor<1x3x3x4xf32> - %1 = stablehlo.constant dense<6.000000e+00> : tensor<1x3x3x4xf32> - %2 = stablehlo.convolution(%arg0, %arg1) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {pad = [[1, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x3x3x4xf32>, tensor<3x3x4x4xf32>) -> tensor<1x3x3x4xf32> - %3 = stablehlo.clamp %0, %2, %1 : tensor<1x3x3x4xf32> - func.return %3: tensor<1x3x3x4xf32> +// CHECK-SAME: %[[ARG_0:.*]]: tensor<1x3x3x4xf32> +func.func @conv_with_relu6_fn(%arg0: tensor<1x3x3x4xf32>) -> tensor<1x3x3x4xf32> { + %0 = stablehlo.constant dense<2.000000e+00> : tensor<3x3x4x4xf32> + %1 = stablehlo.constant dense<0.000000e+00> : tensor<1x3x3x4xf32> + %2 = stablehlo.constant dense<6.000000e+00> : tensor<1x3x3x4xf32> + %3 = stablehlo.convolution(%arg0, %0) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {pad = [[1, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x3x3x4xf32>, tensor<3x3x4x4xf32>) -> tensor<1x3x3x4xf32> + %4 = stablehlo.clamp %1, %3, %2 : tensor<1x3x3x4xf32> + func.return %4: tensor<1x3x3x4xf32> } -// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %arg1) +// CHECK: %[[CONST:.*]] = stablehlo.constant dense<2.000000e+00> +// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %[[CONST]]) // CHECK: return %[[XLA_CALL_MODULE:.*]] : tensor<1x3x3x4xf32> // CHECK: } @@ -166,16 +287,17 @@ func.func @conv_with_relu6_fn(%arg0: tensor<1x3x3x4xf32>, %arg1: tensor<3x3x4x4x // ----- // CHECK-LABEL: @dot_general_with_relu6_fn( -// CHECK-SAME: %[[ARG_0:.*]]: tensor<1x1x167xf32>, -// CHECK-SAME: %[[ARG_1:.*]]: tensor<167x64xf32> -func.func @dot_general_with_relu6_fn(%arg0: tensor<1x1x167xf32>, %arg1: tensor<167x64xf32>) -> tensor<1x1x64xf32> { - %0 = stablehlo.constant dense<0.000000e+00> : tensor<1x1x64xf32> - %1 = stablehlo.constant dense<6.000000e+00> : tensor<1x1x64xf32> - %2 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [2] x [0], precision = [DEFAULT, DEFAULT] : (tensor<1x1x167xf32>, tensor<167x64xf32>) -> tensor<1x1x64xf32> - %3 = stablehlo.clamp %0, %2, %1 : tensor<1x1x64xf32> - return %3 : tensor<1x1x64xf32> +// CHECK-SAME: %[[ARG_0:.*]]: tensor<1x1x167xf32> +func.func @dot_general_with_relu6_fn(%arg0: tensor<1x1x167xf32>) -> tensor<1x1x64xf32> { + %0 = stablehlo.constant dense<2.000000e+00> : tensor<167x64xf32> + %1 = stablehlo.constant dense<0.000000e+00> : tensor<1x1x64xf32> + %2 = stablehlo.constant dense<6.000000e+00> : tensor<1x1x64xf32> + %3 = stablehlo.dot_general %arg0, %0, contracting_dims = [2] x [0], precision = [DEFAULT, DEFAULT] : (tensor<1x1x167xf32>, tensor<167x64xf32>) -> tensor<1x1x64xf32> + %4 = stablehlo.clamp %1, %3, %2 : tensor<1x1x64xf32> + return %4 : tensor<1x1x64xf32> } -// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %arg1) +// CHECK: %[[CONST:.*]] = stablehlo.constant dense<2.000000e+00> +// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %[[CONST]]) // CHECK: return %[[XLA_CALL_MODULE:.*]] : tensor<1x1x64xf32> // CHECK: } @@ -190,17 +312,19 @@ func.func @dot_general_with_relu6_fn(%arg0: tensor<1x1x167xf32>, %arg1: tensor<1 // ----- // CHECK-LABEL: @conv_with_bias_and_relu_fn( -// CHECK-SAME: %[[ARG_0:.*]]: tensor<1x3x3x4xf32>, -// CHECK-SAME: %[[ARG_1:.*]]: tensor<3x3x4x4xf32>, -// CHECK-SAME: %[[ARG_2:.*]]: tensor<1x3x3x4xf32>) -func.func @conv_with_bias_and_relu_fn(%arg0: tensor<1x3x3x4xf32>, %arg1: tensor<3x3x4x4xf32>, %arg2: tensor<1x3x3x4xf32>) -> tensor<1x3x3x4xf32> { - %0 = stablehlo.constant dense<0.000000e+00> : tensor<1x3x3x4xf32> - %1 = stablehlo.convolution(%arg0, %arg1) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {pad = [[1, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x3x3x4xf32>, tensor<3x3x4x4xf32>) -> tensor<1x3x3x4xf32> - %2 = stablehlo.add %1, %arg2 : tensor<1x3x3x4xf32> - %3 = stablehlo.maximum %2, %0 : tensor<1x3x3x4xf32> - func.return %3: tensor<1x3x3x4xf32> +// CHECK-SAME: %[[ARG_0:.*]]: tensor<1x3x3x4xf32> +func.func @conv_with_bias_and_relu_fn(%arg0: tensor<1x3x3x4xf32>) -> tensor<1x3x3x4xf32> { + %0 = stablehlo.constant dense<2.000000e+00> : tensor<3x3x4x4xf32> + %1 = stablehlo.constant dense<2.000000e+00> : tensor<1x3x3x4xf32> + %2 = stablehlo.constant dense<0.000000e+00> : tensor<1x3x3x4xf32> + %3 = stablehlo.convolution(%arg0, %0) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {pad = [[1, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x3x3x4xf32>, tensor<3x3x4x4xf32>) -> tensor<1x3x3x4xf32> + %4 = stablehlo.add %3, %1 : tensor<1x3x3x4xf32> + %5 = stablehlo.maximum %4, %2 : tensor<1x3x3x4xf32> + func.return %5: tensor<1x3x3x4xf32> } -// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %arg1, %arg2) +// CHECK: %[[CONST_0:.*]] = stablehlo.constant dense<2.000000e+00> +// CHECK: %[[CONST_1:.*]] = stablehlo.constant dense<2.000000e+00> +// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %[[CONST_0]], %[[CONST_1]]) // CHECK: return %[[XLA_CALL_MODULE:.*]] : tensor<1x3x3x4xf32> // CHECK: } @@ -215,17 +339,19 @@ func.func @conv_with_bias_and_relu_fn(%arg0: tensor<1x3x3x4xf32>, %arg1: tensor< // ----- // CHECK-LABEL: @dot_general_with_bias_and_relu_fn( -// CHECK-SAME: %[[ARG_0:.*]]: tensor<1x1x167xf32>, -// CHECK-SAME: %[[ARG_1:.*]]: tensor<167x64xf32> -// CHECK-SAME: %[[ARG_2:.*]]: tensor<1x1x64xf32>) -func.func @dot_general_with_bias_and_relu_fn(%arg0: tensor<1x1x167xf32>, %arg1: tensor<167x64xf32>, %arg2: tensor<1x1x64xf32>) -> tensor<1x1x64xf32> { - %0 = stablehlo.constant dense<0.000000e+00> : tensor<1x1x64xf32> - %1 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [2] x [0], precision = [DEFAULT, DEFAULT] : (tensor<1x1x167xf32>, tensor<167x64xf32>) -> tensor<1x1x64xf32> - %2 = stablehlo.add %1, %arg2 : tensor<1x1x64xf32> - %3 = stablehlo.maximum %2, %0 : tensor<1x1x64xf32> - func.return %3: tensor<1x1x64xf32> +// CHECK-SAME: %[[ARG_0:.*]]: tensor<1x1x167xf32> +func.func @dot_general_with_bias_and_relu_fn(%arg0: tensor<1x1x167xf32>) -> tensor<1x1x64xf32> { + %0 = stablehlo.constant dense<2.000000e+00> : tensor<167x64xf32> + %1 = stablehlo.constant dense<2.000000e+00> : tensor<1x1x64xf32> + %2 = stablehlo.constant dense<0.000000e+00> : tensor<1x1x64xf32> + %3 = stablehlo.dot_general %arg0, %0, contracting_dims = [2] x [0], precision = [DEFAULT, DEFAULT] : (tensor<1x1x167xf32>, tensor<167x64xf32>) -> tensor<1x1x64xf32> + %4 = stablehlo.add %3, %1 : tensor<1x1x64xf32> + %5 = stablehlo.maximum %4, %2 : tensor<1x1x64xf32> + func.return %5: tensor<1x1x64xf32> } -// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %arg1, %arg2) +// CHECK: %[[CONST_0:.*]] = stablehlo.constant dense<2.000000e+00> +// CHECK: %[[CONST_1:.*]] = stablehlo.constant dense<2.000000e+00> +// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %[[CONST_0]], %[[CONST_1]]) // CHECK: return %[[XLA_CALL_MODULE:.*]] : tensor<1x1x64xf32> // CHECK: } @@ -239,19 +365,91 @@ func.func @dot_general_with_bias_and_relu_fn(%arg0: tensor<1x1x167xf32>, %arg1: // ----- +// CHECK-LABEL: @conv_with_bias_and_relu_dynamic_fn( +// CHECK-SAME: %[[ARG_0:.*]]: tensor +func.func @conv_with_bias_and_relu_dynamic_fn(%arg0: tensor) -> tensor { + %0 = stablehlo.constant dense<2.000000e+00> : tensor<3x3x1x16xf32> + %1 = stablehlo.constant dense<2.000000e+00> : tensor<16xf32> + %2 = stablehlo.constant dense<0.000000e+00> : tensor + %3 = stablehlo.convolution(%arg0, %0) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {stride = [1, 1], pad = [[1, 1], [1, 1]], rhs_dilate = [1, 1]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64, precision_config = [#stablehlo, #stablehlo]} : (tensor, tensor<3x3x1x16xf32>) -> tensor + %4 = shape.shape_of %3 : tensor -> tensor<4xindex> + %5 = stablehlo.dynamic_broadcast_in_dim %1, %4, dims = [3] : (tensor<16xf32>, tensor<4xindex>) -> tensor + %6 = stablehlo.add %3, %5 : tensor + %7 = shape.shape_of %6 : tensor -> tensor<4xindex> + %8 = stablehlo.dynamic_broadcast_in_dim %2, %7, dims = [] : (tensor, tensor<4xindex>) -> tensor + %9 = stablehlo.maximum %6, %8 : tensor + func.return %9: tensor +} +// CHECK: %[[CONST_0:.*]] = stablehlo.constant dense<2.000000e+00> +// CHECK: %[[CONST_1:.*]] = stablehlo.constant dense<2.000000e+00> +// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %[[CONST_0]], %[[CONST_1]]) +// CHECK: return %[[XLA_CALL_MODULE:.*]] : tensor +// CHECK: } + +// CHECK-LABEL: private @composite_conv_with_bias_and_relu_dynamic_fn_1 +// CHECK: %[[CONST:.*]] = stablehlo.constant dense<0.000000e+00> +// CHECK: %[[CONV:.*]] = stablehlo.convolution(%arg0, %arg1) +// CHECK: %[[SHAPE_OF_0:.*]] = shape.shape_of %[[CONV]] +// CHECK: %[[DYNAMIC_BROADCAST_IN_DIM_0:.*]] = stablehlo.dynamic_broadcast_in_dim %arg2, %[[SHAPE_OF_0]] +// CHECK: %[[ADD:.*]] = stablehlo.add %[[CONV]], %[[DYNAMIC_BROADCAST_IN_DIM_0]] +// CHECK: %[[SHAPE_OF_1:.*]] = shape.shape_of %[[ADD]] +// CHECK: %[[DYNAMIC_BROADCAST_IN_DIM_1:.*]] = stablehlo.dynamic_broadcast_in_dim %[[CONST]], %[[SHAPE_OF_1]] +// CHECK: %[[MAX:.*]] = stablehlo.maximum %[[ADD]], %[[DYNAMIC_BROADCAST_IN_DIM_1]] +// CHECK: return %[[MAX]] : tensor +// CHECK: } + +// ----- + +// CHECK-LABEL: @dot_general_with_bias_and_relu_dynamic_fn( +// CHECK-SAME: %[[ARG_0:.*]]: tensor +func.func @dot_general_with_bias_and_relu_dynamic_fn(%arg0: tensor) -> tensor { + %0 = stablehlo.constant dense<2.000000e+00> : tensor<12544x10xf32> + %1 = stablehlo.constant dense<2.000000e+00> : tensor<10xf32> + %2 = stablehlo.constant dense<0.000000e+00> : tensor + %3 = stablehlo.dot_general %arg0, %0, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor, tensor<12544x10xf32>) -> tensor + %4 = shape.shape_of %3 : tensor -> tensor<2xindex> + %5 = stablehlo.dynamic_broadcast_in_dim %1, %4, dims = [1] : (tensor<10xf32>, tensor<2xindex>) -> tensor + %6 = stablehlo.add %3, %5 : tensor + %7 = shape.shape_of %6 : tensor -> tensor<2xindex> + %8 = stablehlo.dynamic_broadcast_in_dim %2, %7, dims = [] : (tensor, tensor<2xindex>) -> tensor + %9 = stablehlo.maximum %6, %8 : tensor + func.return %9: tensor +} +// CHECK: %[[CONST_0:.*]] = stablehlo.constant dense<2.000000e+00> +// CHECK: %[[CONST_1:.*]] = stablehlo.constant dense<2.000000e+00> +// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %[[CONST_0]], %[[CONST_1]]) +// CHECK: return %[[XLA_CALL_MODULE:.*]] : tensor +// CHECK: } + +// CHECK-LABEL: private @composite_dot_general_with_bias_and_relu_dynamic_fn_1 +// CHECK: %[[CONST:.*]] = stablehlo.constant dense<0.000000e+00> +// CHECK: %[[DOT_GENERAL:.*]] = stablehlo.dot_general %arg0, %arg1 +// CHECK: %[[SHAPE_OF_0:.*]] = shape.shape_of %[[DOT_GENERAL]] +// CHECK: %[[DYNAMIC_BROADCAST_IN_DIM_0:.*]] = stablehlo.dynamic_broadcast_in_dim %arg2, %[[SHAPE_OF_0]] +// CHECK: %[[ADD:.*]] = stablehlo.add %[[DOT_GENERAL]], %[[DYNAMIC_BROADCAST_IN_DIM_0]] +// CHECK: %[[SHAPE_OF_1:.*]] = shape.shape_of %[[ADD]] +// CHECK: %[[DYNAMIC_BROADCAST_IN_DIM_1:.*]] = stablehlo.dynamic_broadcast_in_dim %[[CONST]], %[[SHAPE_OF_1]] +// CHECK: %[[MAX:.*]] = stablehlo.maximum %[[ADD]], %[[DYNAMIC_BROADCAST_IN_DIM_1]] +// CHECK: return %[[MAX]] : tensor +// CHECK: } + +// ----- + // CHECK-LABEL: @conv_with_bias_and_relu6_fn( -// CHECK-SAME: %[[ARG_0:.*]]: tensor<1x3x3x4xf32>, -// CHECK-SAME: %[[ARG_1:.*]]: tensor<3x3x4x4xf32>, -// CHECK-SAME: %[[ARG_2:.*]]: tensor<1x3x3x4xf32>) -func.func @conv_with_bias_and_relu6_fn(%arg0: tensor<1x3x3x4xf32>, %arg1: tensor<3x3x4x4xf32>, %arg2: tensor<1x3x3x4xf32>) -> tensor<1x3x3x4xf32> { - %0 = stablehlo.constant dense<0.000000e+00> : tensor<1x3x3x4xf32> - %1 = stablehlo.constant dense<6.000000e+00> : tensor<1x3x3x4xf32> - %2 = stablehlo.convolution(%arg0, %arg1) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {pad = [[1, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x3x3x4xf32>, tensor<3x3x4x4xf32>) -> tensor<1x3x3x4xf32> - %3 = stablehlo.add %2, %arg2 : tensor<1x3x3x4xf32> - %4 = stablehlo.clamp %0, %3, %1 : tensor<1x3x3x4xf32> - func.return %4: tensor<1x3x3x4xf32> +// CHECK-SAME: %[[ARG_0:.*]]: tensor<1x3x3x4xf32> +func.func @conv_with_bias_and_relu6_fn(%arg0: tensor<1x3x3x4xf32>) -> tensor<1x3x3x4xf32> { + %0 = stablehlo.constant dense<2.000000e+00> : tensor<3x3x4x4xf32> + %1 = stablehlo.constant dense<2.000000e+00> : tensor<1x3x3x4xf32> + %2 = stablehlo.constant dense<0.000000e+00> : tensor<1x3x3x4xf32> + %3 = stablehlo.constant dense<6.000000e+00> : tensor<1x3x3x4xf32> + %4 = stablehlo.convolution(%arg0, %0) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {pad = [[1, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x3x3x4xf32>, tensor<3x3x4x4xf32>) -> tensor<1x3x3x4xf32> + %5 = stablehlo.add %4, %1 : tensor<1x3x3x4xf32> + %6 = stablehlo.clamp %2, %5, %3 : tensor<1x3x3x4xf32> + func.return %6: tensor<1x3x3x4xf32> } -// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %arg1, %arg2) +// CHECK: %[[CONST_0:.*]] = stablehlo.constant dense<2.000000e+00> +// CHECK: %[[CONST_1:.*]] = stablehlo.constant dense<2.000000e+00> +// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %[[CONST_0]], %[[CONST_1]]) // CHECK: return %[[XLA_CALL_MODULE:.*]] : tensor<1x3x3x4xf32> // CHECK: } @@ -267,18 +465,20 @@ func.func @conv_with_bias_and_relu6_fn(%arg0: tensor<1x3x3x4xf32>, %arg1: tensor // ----- // CHECK-LABEL: @dot_general_with_bias_and_relu6_fn( -// CHECK-SAME: %[[ARG_0:.*]]: tensor<1x1x167xf32>, -// CHECK-SAME: %[[ARG_1:.*]]: tensor<167x64xf32> -// CHECK-SAME: %[[ARG_2:.*]]: tensor<1x1x64xf32>) -func.func @dot_general_with_bias_and_relu6_fn(%arg0: tensor<1x1x167xf32>, %arg1: tensor<167x64xf32>, %arg2: tensor<1x1x64xf32>) -> tensor<1x1x64xf32> { - %0 = stablehlo.constant dense<0.000000e+00> : tensor<1x1x64xf32> - %1 = stablehlo.constant dense<6.000000e+00> : tensor<1x1x64xf32> - %2 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [2] x [0], precision = [DEFAULT, DEFAULT] : (tensor<1x1x167xf32>, tensor<167x64xf32>) -> tensor<1x1x64xf32> - %3 = stablehlo.add %2, %arg2 : tensor<1x1x64xf32> - %4 = stablehlo.clamp %0, %3, %1 : tensor<1x1x64xf32> - func.return %4: tensor<1x1x64xf32> +// CHECK-SAME: %[[ARG_0:.*]]: tensor<1x1x167xf32> +func.func @dot_general_with_bias_and_relu6_fn(%arg0: tensor<1x1x167xf32>) -> tensor<1x1x64xf32> { + %0 = stablehlo.constant dense<2.000000e+00> : tensor<167x64xf32> + %1 = stablehlo.constant dense<2.000000e+00> : tensor<1x1x64xf32> + %2 = stablehlo.constant dense<0.000000e+00> : tensor<1x1x64xf32> + %3 = stablehlo.constant dense<6.000000e+00> : tensor<1x1x64xf32> + %4 = stablehlo.dot_general %arg0, %0, contracting_dims = [2] x [0], precision = [DEFAULT, DEFAULT] : (tensor<1x1x167xf32>, tensor<167x64xf32>) -> tensor<1x1x64xf32> + %5 = stablehlo.add %4, %1 : tensor<1x1x64xf32> + %6 = stablehlo.clamp %2, %5, %3 : tensor<1x1x64xf32> + func.return %6: tensor<1x1x64xf32> } -// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %arg1, %arg2) +// CHECK: %[[CONST_0:.*]] = stablehlo.constant dense<2.000000e+00> +// CHECK: %[[CONST_1:.*]] = stablehlo.constant dense<2.000000e+00> +// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%arg0, %[[CONST_0]], %[[CONST_1]]) // CHECK: return %[[XLA_CALL_MODULE:.*]] : tensor<1x1x64xf32> // CHECK: } diff --git a/tensorflow/compiler/mlir/quantization/stablehlo/tests/post_quantize.mlir b/tensorflow/compiler/mlir/quantization/stablehlo/tests/post_quantize.mlir new file mode 100644 index 00000000000000..ae2f57081e40f7 --- /dev/null +++ b/tensorflow/compiler/mlir/quantization/stablehlo/tests/post_quantize.mlir @@ -0,0 +1,72 @@ +// RUN: stablehlo-quant-opt %s -split-input-file -stablehlo-post-quantize | FileCheck %s + +// CHECK-LABEL: @remove_volatile_qdq +func.func @remove_volatile_qdq() -> tensor<3x2xf32> { + // CHECK: %[[CST:.*]] = stablehlo.constant + // CHECK-NOT: "quantfork.qcast" + // CHECK-NOT: "quantfork.dcast" + // CHECK: return %[[CST]] + %cst = stablehlo.constant dense<[[-0.960978984, -0.390246302], [-0.790828585, -0.601039409], [-1.0280807, -1.02731466]]> : tensor<3x2xf32> + %q = "quantfork.qcast"(%cst) {volatile} : (tensor<3x2xf32>) -> tensor<3x2x!quant.uniform> + %dq = "quantfork.dcast"(%q) : (tensor<3x2x!quant.uniform>) -> tensor<3x2xf32> + func.return %dq : tensor<3x2xf32> +} + +// ----- + +// CHECK-LABEL: @remove_volatile_qdq_with_requantization +// CHECK-SAME: %[[ARG0:.*]]: tensor<3x2xf32> +func.func @remove_volatile_qdq_with_requantization(%arg0: tensor<3x2xf32>) -> tensor<3x2xf32> { + // CHECK: %[[Q1:.*]] = stablehlo.uniform_quantize %[[ARG0]] + // CHECK: %[[Q2:.*]] = stablehlo.uniform_quantize %[[Q1]] + // CHECK: %[[ABS:.*]] = stablehlo.abs %[[Q2]] + // CHECK: %[[DQ:.*]] = stablehlo.uniform_dequantize %[[ABS]] + // CHECK: %[[ADD:.*]] = stablehlo.add %[[ARG0]], %[[DQ]] + // CHECK: return %[[ADD]] + %q1 = "quantfork.qcast"(%arg0) {volatile} : (tensor<3x2xf32>) -> tensor<3x2x!quant.uniform> + %q2 = "quantfork.qcast"(%q1) {volatile} : (tensor<3x2x!quant.uniform>) -> tensor<3x2x!quant.uniform> + %dq1 = "quantfork.dcast"(%q2) : (tensor<3x2x!quant.uniform>) -> tensor<3x2xf32> + %abs = stablehlo.abs %q2 : (tensor<3x2x!quant.uniform>) -> tensor<3x2x!quant.uniform> + %dq2 = "quantfork.dcast"(%abs) : (tensor<3x2x!quant.uniform>) -> tensor<3x2xf32> + %add = stablehlo.add %dq1, %dq2 : (tensor<3x2xf32>, tensor<3x2xf32>) -> tensor<3x2xf32> + func.return %add : tensor<3x2xf32> +} + +// ----- + +// CHECK-LABEL: @quantize_constant +// CHECK-SAME: %[[ARG0:.*]]: tensor<1x3xf32> +func.func @quantize_constant(%arg0: tensor<1x3xf32>) -> tensor<1x2xf32> { + // CHECK-DAG: %[[QCST:.*]] = stablehlo.constant() {value = dense<-78> : tensor<3x2xi8>} : () -> tensor<3x2x!quant.uniform:f32, 5.000000e-03>> + // CHECK-DAG: %[[Q1:.*]] = stablehlo.uniform_quantize %[[ARG0]] + // CHECK-NOT: "quantfork.qcast" + // CHECK: %[[DOT:.*]] = stablehlo.dot %[[Q1]], %[[QCST]] + // CHECK: %[[DQ:.*]] = stablehlo.uniform_dequantize %[[DOT]] + // CHECK: return %[[DQ]] + %cst = stablehlo.constant dense<-0.390246302> : tensor<3x2xf32> + %q1 = "quantfork.qcast"(%arg0) {volatile} : (tensor<1x3xf32>) -> tensor<1x3x!quant.uniform> + %q2 = "quantfork.qcast"(%cst) {volatile} : (tensor<3x2xf32>) -> tensor<3x2x!quant.uniform:f32, 5.000000e-03>> + %dot = stablehlo.dot %q1, %q2 : (tensor<1x3x!quant.uniform>, tensor<3x2x!quant.uniform:f32, 5.000000e-03>>) -> tensor<1x2x!quant.uniform> + %dq = "quantfork.dcast"(%dot) : (tensor<1x2x!quant.uniform>) -> tensor<1x2xf32> + func.return %dq : tensor<1x2xf32> +} + +// ----- + +// CHECK-LABEL: @convert_quantfork_qdq_to_stablehlo_uniform_qdq +// CHECK-SAME: %[[ARG0:.*]]: tensor<1x3xf32> +// CHECK-SAME: %[[ARG1:.*]]: tensor<3x2xf32> +func.func @convert_quantfork_qdq_to_stablehlo_uniform_qdq(%arg0: tensor<1x3xf32>, %arg1: tensor<3x2xf32>) -> tensor<1x2xf32> { + // CHECK: %[[Q1:.*]] = stablehlo.uniform_quantize %[[ARG0]] + // CHECK-NOT: "quantfork.qcast" + // CHECK: %[[Q2:.*]] = stablehlo.uniform_quantize %[[ARG1]] + // CHECK-NOT: "quantfork.qcast" + // CHECK: %[[DOT:.*]] = stablehlo.dot %[[Q1]], %[[Q2]] + // CHECK: %[[DQ:.*]] = stablehlo.uniform_dequantize %[[DOT]] + // CHECK: return %[[DQ]] + %q1 = "quantfork.qcast"(%arg0) {volatile} : (tensor<1x3xf32>) -> tensor<1x3x!quant.uniform> + %q2 = "quantfork.qcast"(%arg1) {volatile} : (tensor<3x2xf32>) -> tensor<3x2x!quant.uniform:f32, 5.000000e-03>> + %dot = stablehlo.dot %q1, %q2 : (tensor<1x3x!quant.uniform>, tensor<3x2x!quant.uniform:f32, 5.000000e-03>>) -> tensor<1x2x!quant.uniform> + %dq = "quantfork.dcast"(%dot) : (tensor<1x2x!quant.uniform>) -> tensor<1x2xf32> + func.return %dq : tensor<1x2xf32> +} diff --git a/tensorflow/compiler/mlir/quantization/stablehlo/tests/prepare_quantize.mlir b/tensorflow/compiler/mlir/quantization/stablehlo/tests/prepare_quantize.mlir index 612b8d6a519959..8f38f889f28e33 100644 --- a/tensorflow/compiler/mlir/quantization/stablehlo/tests/prepare_quantize.mlir +++ b/tensorflow/compiler/mlir/quantization/stablehlo/tests/prepare_quantize.mlir @@ -3,6 +3,7 @@ // ----- // CHECK-LABEL: func @dot +// CHECK-SAME: (%[[ARG_0:.*]]: tensor) -> tensor func.func @dot(%arg0: tensor) -> tensor { // CHECK: %[[cst:.*]] = stablehlo.constant // CHECK: %[[q1:.*]] = "quantfork.qcast"(%[[cst]]) @@ -10,7 +11,7 @@ func.func @dot(%arg0: tensor) -> tensor { // CHECK: %[[dq1:.*]] = "quantfork.dcast"(%[[q1]]) // CHECK-SAME: quant.uniform %cst = stablehlo.constant dense<[[-0.960978984, -0.390246302], [-0.790828585, -0.601039409], [-1.0280807, -1.02731466]]> : tensor<3x2xf32> - // CHECK: %[[q2:.*]] = "quantfork.qcast"(%arg0) + // CHECK: %[[q2:.*]] = "quantfork.qcast"(%[[ARG_0]]) // CHECK-SAME: quant.uniform // CHECK: %[[dq2:.*]] = "quantfork.dcast"(%[[q2]]) // CHECK-SAME: quant.uniform @@ -29,8 +30,9 @@ func.func @dot(%arg0: tensor) -> tensor { // ----- // CHECK-LABEL: func @duplicate_stats +// CHECK-SAME: (%[[ARG_0:.*]]: tensor<2x3xf32>) -> tensor<2x3xf32> func.func @duplicate_stats(%arg0: tensor<2x3xf32>) -> tensor<2x3xf32> { - // CHECK: %[[q1:.*]] = "quantfork.qcast"(%arg0) + // CHECK: %[[q1:.*]] = "quantfork.qcast"(%[[ARG_0]]) // CHECK: %[[dq1:.*]] = "quantfork.dcast"(%[[q1]]) // CHECK: %[[q2:.*]] = "quantfork.qcast"(%[[dq1]]) // CHECK: %[[dq2:.*]] = "quantfork.dcast"(%[[q2]]) @@ -44,6 +46,7 @@ func.func @duplicate_stats(%arg0: tensor<2x3xf32>) -> tensor<2x3xf32> { // ----- // CHECK-LABEL: func @dot_redundant_stats +// CHECK-SAME: (%[[ARG_0:.*]]: tensor) -> tensor func.func @dot_redundant_stats(%arg0: tensor) -> tensor { // CHECK: %[[cst:.*]] = stablehlo.constant // CHECK: %[[q1:.*]] = "quantfork.qcast"(%[[cst]]) @@ -51,7 +54,7 @@ func.func @dot_redundant_stats(%arg0: tensor) -> tensor { // CHECK: %[[dq1:.*]] = "quantfork.dcast"(%[[q1]]) // CHECK-SAME: quant.uniform %cst = stablehlo.constant dense<[[-0.960978984, -0.390246302], [-0.790828585, -0.601039409], [-1.0280807, -1.02731466]]> : tensor<3x2xf32> - // CHECK: %[[q2:.*]] = "quantfork.qcast"(%arg0) + // CHECK: %[[q2:.*]] = "quantfork.qcast"(%[[ARG_0]]) // CHECK-SAME: quant.uniform // CHECK: %[[dq2:.*]] = "quantfork.dcast"(%[[q2]]) // CHECK-SAME: quant.uniform @@ -87,10 +90,11 @@ func.func @convert_same_scale_propagate(%arg0: tensor<2x3xf32>) -> tensor<2x3xf3 // ----- // CHECK-LABEL: func @merge_consecutive_qcast +// CHECK-SAME: (%[[ARG_0:.*]]: tensor<*xf32>, %[[ARG_1:.*]]: tensor<*xf32>, %[[ARG_2:.*]]: tensor<*xf32>) -> (tensor<*xf32>, tensor<*xf32>) func.func @merge_consecutive_qcast(%arg0: tensor<*xf32>, %arg1: tensor<*xf32>, %arg2: tensor<*xf32>) -> (tensor<*xf32>, tensor<*xf32>) { - // CHECK: "quantfork.qcast"(%arg1) + // CHECK: "quantfork.qcast"(%[[ARG_1]]) // CHECK-SAME: -> tensor<*x!quant.uniform> - // CHECK: "quantfork.qcast"(%arg1) + // CHECK: "quantfork.qcast"(%[[ARG_1]]) // CHECK-SAME: -> tensor<*x!quant.uniform> %0 = "quantfork.stats"(%arg0) {layerStats = dense<[-0.83811146, 2.4960899]> : tensor<2xf32>} : (tensor<*xf32>) -> tensor<*xf32> %1 = "quantfork.stats"(%arg1) {layerStats = dense<[-0.835039615, 1.000000e+00]> : tensor<2xf32>} : (tensor<*xf32>) -> tensor<*xf32> diff --git a/tensorflow/compiler/mlir/quantization/stablehlo/tests/prepare_quantize_int4.mlir b/tensorflow/compiler/mlir/quantization/stablehlo/tests/prepare_quantize_int4.mlir index 0ddb01cd0e0be4..ca467d1180ea0a 100644 --- a/tensorflow/compiler/mlir/quantization/stablehlo/tests/prepare_quantize_int4.mlir +++ b/tensorflow/compiler/mlir/quantization/stablehlo/tests/prepare_quantize_int4.mlir @@ -1,6 +1,7 @@ // RUN: stablehlo-quant-opt %s -split-input-file -stablehlo-prepare-quantize=bit-width=4 -verify-diagnostics | FileCheck %s // CHECK-LABEL: func @dot_int4 +// CHECK-SAME: (%[[ARG_0:.*]]: tensor) -> tensor func.func @dot_int4(%arg0: tensor) -> tensor { // CHECK: %[[cst:.*]] = stablehlo.constant // CHECK: %[[q1:.*]] = "quantfork.qcast"(%[[cst]]) @@ -8,7 +9,7 @@ func.func @dot_int4(%arg0: tensor) -> tensor { // CHECK: %[[dq1:.*]] = "quantfork.dcast"(%[[q1]]) // CHECK-SAME: quant.uniform %cst = stablehlo.constant dense<[[-0.960978984, -0.390246302], [-0.790828585, -0.601039409], [-1.0280807, -1.02731466]]> : tensor<3x2xf32> - // CHECK: %[[q2:.*]] = "quantfork.qcast"(%arg0) + // CHECK: %[[q2:.*]] = "quantfork.qcast"(%[[ARG_0]]) // CHECK-SAME: quant.uniform // CHECK: %[[dq2:.*]] = "quantfork.dcast"(%[[q2]]) // CHECK-SAME: quant.uniform diff --git a/tensorflow/compiler/mlir/quantization/stablehlo/tests/quantize.mlir b/tensorflow/compiler/mlir/quantization/stablehlo/tests/quantize.mlir index eccd931d1b5008..d1bfea7a236448 100644 --- a/tensorflow/compiler/mlir/quantization/stablehlo/tests/quantize.mlir +++ b/tensorflow/compiler/mlir/quantization/stablehlo/tests/quantize.mlir @@ -20,7 +20,7 @@ func.func private @quantize_simple_xla_call_module(%arg0: tensor<1x4xf32>) -> te // CHECK: %[[CONST_0:.*]] = "stablehlo.constant"() {value = dense<1.000000e+00> : tensor<4x3xf32>} : () -> tensor<4x3xf32> // CHECK-DAG: %[[QCAST_0:.*]] = "quantfork.qcast"(%[[CONST_0]]) {volatile} : (tensor<4x3xf32>) -> tensor<4x3x!quant.uniform:f32, 5.000000e-03>> // CHECK-DAG: %[[QCAST_1:.*]] = "quantfork.qcast"(%[[ARG_0]]) {volatile} : (tensor<1x4xf32>) -> tensor<1x4x!quant.uniform> -// CHECK: %[[XLACALLMODULE_0:.*]] = "tf.XlaCallModule"(%[[QCAST_1]], %[[QCAST_0]]) {{{.*}}} : (tensor<1x4x!quant.uniform>, tensor<4x3x!quant.uniform:f32, 5.000000e-03>>) -> tensor<1x3x!quant.uniform> +// CHECK: %[[XLACALLMODULE_0:.*]] = "tf.XlaCallModule"(%[[QCAST_1]], %[[QCAST_0]]) <{{{.*}}}> {{{.*}}} : (tensor<1x4x!quant.uniform>, tensor<4x3x!quant.uniform:f32, 5.000000e-03>>) -> tensor<1x3x!quant.uniform> // CHECK: %[[DCAST_0:.*]] = "quantfork.dcast"(%[[XLACALLMODULE_0]]) : (tensor<1x3x!quant.uniform>) -> tensor<1x3xf32> // CHECK: "func.return"(%[[DCAST_0]]) : (tensor<1x3xf32>) -> () @@ -37,6 +37,6 @@ func.func private @quantize_simple_xla_call_module_no_operand() -> tensor<1x3xf3 // Tests that the output of the tf.XlaCallModule op has been replaced by // a quantized type, and the corresponding quantfork.qcast ops that turned // the float output to a quantized type is removed. -// CHECK: %[[XLACALLMODULE_0:.*]] = "tf.XlaCallModule"() {{{.*}}} : () -> tensor<1x3x!quant.uniform> +// CHECK: %[[XLACALLMODULE_0:.*]] = "tf.XlaCallModule"() <{{{.*}}}> {{{.*}}} : () -> tensor<1x3x!quant.uniform> // CHECK: %[[DCAST_0:.*]] = "quantfork.dcast"(%[[XLACALLMODULE_0]]) : (tensor<1x3x!quant.uniform>) -> tensor<1x3xf32> // CHECK: "func.return"(%[[DCAST_0]]) : (tensor<1x3xf32>) -> () diff --git a/tensorflow/compiler/mlir/quantization/stablehlo/tests/quantize_composite_functions.mlir b/tensorflow/compiler/mlir/quantization/stablehlo/tests/quantize_composite_functions.mlir new file mode 100644 index 00000000000000..97ea1f30be81ba --- /dev/null +++ b/tensorflow/compiler/mlir/quantization/stablehlo/tests/quantize_composite_functions.mlir @@ -0,0 +1,114 @@ +// RUN: stablehlo-quant-opt %s -split-input-file -verify-diagnostics \ +// RUN: -stablehlo-quantize-composite-functions | FileCheck %s + +module attributes {tf_saved_model.semantics} { +// The following pattern does not converge because of a bug in QuantizePass. +// TODO - b/305469508: Fix the QuantizePass to avoid this warning. +// expected-warning @+1 {{Failed to converge pattern at QuantizePass.}} + func.func private @quantize_dot_general(%arg0: tensor<1x3xf32>) -> tensor<1x3xf32> attributes {tf._original_func_name = "main_0"} { + %cst = "tf.Const"() {value = dense<3.00000000e-1> : tensor<3x3xf32>} : () -> tensor<3x3xf32> + %0 = "quantfork.stats"(%arg0) {layerStats = dense<[6.00000000e-6, 9.00000000e-1]> : tensor<2xf32>} : (tensor<1x3xf32>) -> tensor<1x3xf32> + %1 = "tf.XlaCallModule"(%0, %cst) {Sout = [#tf_type.shape<1x3>], _entry_function = @composite_dot_general_fn, _original_entry_function = "composite_dot_general_fn", _stablehlo_module_attrs = {}, _tfl_quant_trait = "fully_quantizable", device = "", dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = [], version = 5 : i64} : (tensor<1x3xf32>, tensor<3x3xf32>) -> tensor<1x3xf32> + %2 = "quantfork.stats"(%1) {layerStats = dense<[5.00000000e-6, 7.00000000e-1]> : tensor<2xf32>} : (tensor<1x3xf32>) -> tensor<1x3xf32> + return %2 : tensor<1x3xf32> + } +// Checks that the quantized XlaCallModule has been replaced by a CallOp, which +// calls the quantized entry function. + +// CHECK-LABEL: func.func private @quantize_dot_general +// CHECK-SAME: (%[[ARG_1:.*]]: tensor<1x3xf32>) -> tensor<1x3xf32> attributes {tf._original_func_name = "main_0"} +// CHECK: %[[CONST_0:.*]] = stablehlo.constant() {value = dense<{{.*}}> : tensor<3x3xi8>} : () -> tensor<3x3x!quant.uniform:f32, {{.*}}> +// CHECK: %[[UNIFORM_QUANTIZE_0:.*]] = stablehlo.uniform_quantize %[[ARG_1]] : (tensor<1x3xf32>) -> tensor<1x3x!quant.uniform> +// CHECK: %[[CALL_0:.*]] = call @quantized_dot_general_fn(%[[UNIFORM_QUANTIZE_0]], %[[CONST_0]]) : (tensor<1x3x!quant.uniform>, tensor<3x3x!quant.uniform:f32, {{.*}}>) -> tensor<1x3x!quant.uniform> +// CHECK: %[[UNIFORM_DEQUANTIZE_0:.*]] = stablehlo.uniform_dequantize %[[CALL_0]] : (tensor<1x3x!quant.uniform) -> tensor<1x3xf32> +// CHECK: return %[[UNIFORM_DEQUANTIZE_0]] : tensor<1x3xf32> + + func.func private @composite_dot_general_fn(%arg0: tensor<1x3xf32>, %arg1: tensor<3x3xf32>) -> tensor<1x3xf32> attributes {_from_xla_call_module} { + %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<1x3xf32>, tensor<3x3xf32>) -> tensor<1x3xf32> + return %0 : tensor<1x3xf32> + } +// Checks that the entry function is quantized for dot_general. Quantized +// dot_general outputs an i32 quantized tensor, followed by requantization to +// i8 quantized tensor. + +// CHECK: func.func private @quantized_dot_general_fn(%[[ARG_2:.*]]: tensor<1x3x!quant.uniform>, %[[ARG_3:.*]]: tensor<3x3x!quant.uniform:f32, {{.*}}>>) -> tensor<1x3x!quant.uniform> attributes {_from_xla_call_module} +// CHECK: %[[DOT_GENERAL_0:.*]] = stablehlo.dot_general %[[ARG_2]], %[[ARG_3]], contracting_dims = [1] x [0] : (tensor<1x3x!quant.uniform>, tensor<3x3x!quant.uniform:f32, {{.*}}>) -> tensor<1x3x!quant.uniform> +// CHECK: %[[UNIFORM_QUANTIZE_1:.*]] = stablehlo.uniform_quantize %[[DOT_GENERAL_0]] : (tensor<1x3x!quant.uniform>) -> tensor<1x3x!quant.uniform> +// CHECK: return %[[UNIFORM_QUANTIZE_1]] : tensor<1x3x!quant.uniform> +} + +// ----- + +// Tests error when there are no corresponding entry function to quantize +// (@composite_dot_general_fn). + +module attributes {tf_saved_model.semantics} { +// The following pattern does not converge because of a bug in QuantizePass. +// TODO - b/305469508: Fix the QuantizePass to avoid this warning. +// expected-warning @+1 {{Failed to converge pattern at QuantizePass.}} + func.func private @error_when_no_entry_function(%arg0: tensor<1x3xf32>) -> tensor<1x3xf32> attributes {tf._original_func_name = "main_0"} { + %cst = "tf.Const"() {value = dense<3.00000000e-1> : tensor<3x3xf32>} : () -> tensor<3x3xf32> + %0 = "quantfork.stats"(%arg0) {layerStats = dense<[6.00000000e-6, 9.00000000e-1]> : tensor<2xf32>} : (tensor<1x3xf32>) -> tensor<1x3xf32> +// expected-error @+2 {{Failed to find a valid entry function}} +// expected-error @+1 {{'tf.XlaCallModule' op operand #0 must be variadic of tensor of tf.dtype values}} + %1 = "tf.XlaCallModule"(%0, %cst) {Sout = [#tf_type.shape<1x3>], _entry_function = @composite_dot_general_fn, _original_entry_function = "composite_dot_general_fn", _stablehlo_module_attrs = {}, _tfl_quant_trait = "fully_quantizable", device = "", dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = [], version = 5 : i64} : (tensor<1x3xf32>, tensor<3x3xf32>) -> tensor<1x3xf32> + %2 = "quantfork.stats"(%1) {layerStats = dense<[5.00000000e-6, 7.00000000e-1]> : tensor<2xf32>} : (tensor<1x3xf32>) -> tensor<1x3xf32> + return %2 : tensor<1x3xf32> + } +} + +// ----- + +// Tests that XlaCallModule op is not quantized without the quantfork.stats ops. + +module attributes {tf_saved_model.semantics} { + func.func private @not_quantized_without_stats(%arg0: tensor<1x3xf32>) -> tensor<1x3xf32> attributes {tf._original_func_name = "main_0"} { + %cst = "tf.Const"() {value = dense<3.00000000e-1> : tensor<3x3xf32>} : () -> tensor<3x3xf32> + %1 = "tf.XlaCallModule"(%arg0, %cst) {Sout = [#tf_type.shape<1x3>], _entry_function = @composite_dot_general_fn, _original_entry_function = "composite_dot_general_fn", _stablehlo_module_attrs = {}, _tfl_quant_trait = "fully_quantizable", device = "", dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = [], version = 5 : i64} : (tensor<1x3xf32>, tensor<3x3xf32>) -> tensor<1x3xf32> + return %1 : tensor<1x3xf32> + } +// Check that "tf.Const" is converted to stablehlo.constant. XlaCallModule is +// not quantized. + +// CHECK-LABEL: func.func private @not_quantized_without_stats +// CHECK-SAME: (%[[ARG_1:.*]]: tensor<1x3xf32>) -> tensor<1x3xf32> attributes {tf._original_func_name = "main_0"} +// CHECK: %[[CONST_0:.*]] = stablehlo.constant dense<3.000000e-01> : tensor<3x3xf32> +// CHECK: %[[XLA_CALL_MODULE_0:.*]] = "tf.XlaCallModule"(%[[ARG_1]], %[[CONST_0]]) <{{{.*}}}> {{{.*_entry_function = @composite_dot_general_fn.*}}} : (tensor<1x3xf32>, tensor<3x3xf32>) -> tensor<1x3xf32> +// CHECK: return %[[XLA_CALL_MODULE_0]] + + func.func private @composite_dot_general_fn(%arg0: tensor<1x3xf32>, %arg1: tensor<3x3xf32>) -> tensor<1x3xf32> attributes {_from_xla_call_module} { + %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<1x3xf32>, tensor<3x3xf32>) -> tensor<1x3xf32> + return %0 : tensor<1x3xf32> + } +// Check that the composite_dot_general_fn is untouched. + +// CHECK: func.func private @composite_dot_general_fn(%[[ARG_2:.*]]: tensor<1x3xf32>, %[[ARG_3:.*]]: tensor<3x3xf32>) -> tensor<1x3xf32> attributes {_from_xla_call_module} +// CHECK: %[[DOT_GENERAL:.*]] = stablehlo.dot_general %[[ARG_2]], %[[ARG_3]] +// CHECK: return %[[DOT_GENERAL]] +} + +// ----- + +// Tests that a fusion pattern for dot_general is not yet supported. Further op +// coverage will be provided in the future. +// TODO - b/307620428: Increase op coverage to cover this test case. + +module attributes {tf_saved_model.semantics} { +// The following pattern does not converge because of a bug in QuantizePass. +// TODO - b/305469508: Fix the QuantizePass to avoid this warning. +// expected-warning @+1 {{Failed to converge pattern at QuantizePass.}} + func.func private @dot_general_fn_fusion_not_quantized(%arg0: tensor<1x3xf32>) -> tensor<1x3xf32> attributes {tf._original_func_name = "main_0"} { + %cst = "tf.Const"() {value = dense<3.00000000e-1> : tensor<3x3xf32>} : () -> tensor<3x3xf32> + %0 = "quantfork.stats"(%arg0) {layerStats = dense<[6.00000000e-6, 9.00000000e-1]> : tensor<2xf32>} : (tensor<1x3xf32>) -> tensor<1x3xf32> +// expected-error @+1 {{'tf.XlaCallModule' op operand #0 must be variadic of tensor of tf.dtype values}} + %1 = "tf.XlaCallModule"(%0, %cst) {Sout = [#tf_type.shape<1x3>], _entry_function = @composite_dot_general_fn, _original_entry_function = "composite_dot_general_fn", _stablehlo_module_attrs = {}, _tfl_quant_trait = "fully_quantizable", device = "", dim_args_spec = [], disabled_checks = [], has_token_input_output = false, module = "", platforms = [], version = 5 : i64} : (tensor<1x3xf32>, tensor<3x3xf32>) -> tensor<1x3xf32> + %2 = "quantfork.stats"(%1) {layerStats = dense<[5.00000000e-6, 7.00000000e-1]> : tensor<2xf32>} : (tensor<1x3xf32>) -> tensor<1x3xf32> + return %2 : tensor<1x3xf32> + } + + func.func private @composite_dot_general_fn(%arg0: tensor<1x3xf32>, %arg1: tensor<3x3xf32>, %arg2: tensor<1x3xf32>) -> tensor<1x3xf32> attributes {_from_xla_call_module} { + %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<1x3xf32>, tensor<3x3xf32>) -> tensor<1x3xf32> + %1 = stablehlo.add %0, %arg2 : tensor<1x3xf32> + return %1 : tensor<1x3xf32> + } +} diff --git a/tensorflow/compiler/mlir/quantization/stablehlo/tests/replace_stablehlo_ops_in_main_function_with_xla_call_module_ops.mlir b/tensorflow/compiler/mlir/quantization/stablehlo/tests/replace_stablehlo_ops_in_main_function_with_xla_call_module_ops.mlir index 57423f5f03a7ec..3d04c72dec7f7e 100644 --- a/tensorflow/compiler/mlir/quantization/stablehlo/tests/replace_stablehlo_ops_in_main_function_with_xla_call_module_ops.mlir +++ b/tensorflow/compiler/mlir/quantization/stablehlo/tests/replace_stablehlo_ops_in_main_function_with_xla_call_module_ops.mlir @@ -30,13 +30,13 @@ module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, p return %9 : tensor<1x64xf32> } - // CHECK: %[[STABLEHLO_SUBGRAPH_TO_XLA_CALL_MODULE_OP_0:.*]] = "tf.XlaCallModule"() {Sout = [#tf_type.shape<{{.*}}>, #tf_type.shape<{{.*}}>], _entry_function = @_stablehlo_main_1 - // CHECK: %[[CUSTOM_AGGREGATOR_0:.*]] = "tf.CustomAggregator"(%arg0) {calibration_method = 1 : i32, id = "0", initial_num_bins = 0 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32} : (tensor<1x1024xf32>) -> tensor<1x1024xf32> - // CHECK: %[[XLA_CALL_MODULE_0:.*]] = "tf.XlaCallModule"(%[[CUSTOM_AGGREGATOR_0]], %[[STABLEHLO_SUBGRAPH_TO_XLA_CALL_MODULE_OP_0:.*]], %[[STABLEHLO_SUBGRAPH_TO_XLA_CALL_MODULE_OP_0:.*]]) {Sout = [#tf_type.shape<1x3>], _entry_function = @composite_dot_general_fn_1, _original_entry_function = "composite_dot_general_fn_1" + // CHECK: %[[STABLEHLO_SUBGRAPH_TO_XLA_CALL_MODULE_OP_0:.*]] = "tf.XlaCallModule"() <{Sout = [#tf_type.shape<{{.*}}>, #tf_type.shape<{{.*}}>], {{.*}}}> {_entry_function = @_stablehlo_main_1 + // CHECK: %[[CUSTOM_AGGREGATOR_0:.*]] = "tf.CustomAggregator"(%arg0) <{id = "0"}> {calibration_method = 1 : i32, initial_num_bins = 0 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32} : (tensor<1x1024xf32>) -> tensor<1x1024xf32> + // CHECK: %[[XLA_CALL_MODULE_0:.*]] = "tf.XlaCallModule"(%[[CUSTOM_AGGREGATOR_0]], %[[STABLEHLO_SUBGRAPH_TO_XLA_CALL_MODULE_OP_0:.*]], %[[STABLEHLO_SUBGRAPH_TO_XLA_CALL_MODULE_OP_0:.*]]) <{Sout = [#tf_type.shape<1x3>], {{.*}}}> {_entry_function = @composite_dot_general_fn_1, _original_entry_function = "composite_dot_general_fn_1" // CHECK: %[[CUSTOM_AGGREGATOR_1:.*]] = "tf.CustomAggregator"(%[[XLA_CALL_MODULE_0]]) - // CHECK: %[[STABLEHLO_SUBGRAPH_TO_XLA_CALL_MODULE_OP_1:.*]] = "tf.XlaCallModule"() {Sout = [#tf_type.shape<{{.*}}>, #tf_type.shape<{{.*}}>], _entry_function = @_stablehlo_main_0 - // CHECK: %[[CUSTOM_AGGREGATOR_2:.*]] = "tf.CustomAggregator"(%[[CUSTOM_AGGREGATOR_1]]) {calibration_method = 1 : i32, id = "0", initial_num_bins = 0 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32} : (tensor<1x3xf32>) -> tensor<1x3xf32> - // CHECK: %[[XLA_CALL_MODULE_1:.*]] = "tf.XlaCallModule"(%[[CUSTOM_AGGREGATOR_2]], %[[STABLEHLO_SUBGRAPH_TO_XLA_CALL_MODULE_OP_1:.*]], %[[STABLEHLO_SUBGRAPH_TO_XLA_CALL_MODULE_OP_1:.*]]) {Sout = [#tf_type.shape<1x3>], _entry_function = @composite_dot_general_with_relu_fn_1, _original_entry_function = "composite_dot_general_with_relu_fn_1" + // CHECK: %[[STABLEHLO_SUBGRAPH_TO_XLA_CALL_MODULE_OP_1:.*]] = "tf.XlaCallModule"() <{Sout = [#tf_type.shape<{{.*}}>, #tf_type.shape<{{.*}}>], {{.*}}}> {_entry_function = @_stablehlo_main_0 + // CHECK: %[[CUSTOM_AGGREGATOR_2:.*]] = "tf.CustomAggregator"(%[[CUSTOM_AGGREGATOR_1]]) <{id = "0"}> {calibration_method = 1 : i32, initial_num_bins = 0 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32} : (tensor<1x3xf32>) -> tensor<1x3xf32> + // CHECK: %[[XLA_CALL_MODULE_1:.*]] = "tf.XlaCallModule"(%[[CUSTOM_AGGREGATOR_2]], %[[STABLEHLO_SUBGRAPH_TO_XLA_CALL_MODULE_OP_1:.*]], %[[STABLEHLO_SUBGRAPH_TO_XLA_CALL_MODULE_OP_1:.*]]) <{Sout = [#tf_type.shape<1x3>], {{.*}}}> {_entry_function = @composite_dot_general_with_relu_fn_1, _original_entry_function = "composite_dot_general_with_relu_fn_1" // CHECK: %[[CUSTOM_AGGREGATOR_3:.*]] = "tf.CustomAggregator"(%[[XLA_CALL_MODULE_1:.*]]) // CHECK: return %[[CUSTOM_AGGREGATOR_3]] : tensor<1x64xf32> // CHECK: } @@ -60,6 +60,39 @@ module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, p } } + +// ----- + +module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 1654 : i32}, tf_saved_model.semantics} { + + // CHECK: func private @_stablehlo_main_0 + // CHECK: %[[CONSTANT_0:.*]] = stablehlo.constant dense<0.134728625> : tensor<1x3xf32> + // CHECK: %[[CONSTANT_1:.*]] = stablehlo.constant dense<-1.280000e+02> : tensor<1x1024xf32> + // CHECK: %[[CONSTANT_2:.*]] = stablehlo.constant dense<0.003921567> : tensor<1x1024xf32> + // CHECK: %[[DIVIDE:.*]] = stablehlo.divide %arg0, %[[CONSTANT_2]] + // CHECK: %[[ADD:.*]] = stablehlo.add %[[DIVIDE]], %[[CONSTANT_1]] + // CHECK return %[[ADD]] + // CHECK: } + + // CHECK: @serving_default + func.func @serving_default(%arg0: tensor<1x1024xf32> {tf_saved_model.index_path = ["input_tensor"]}) -> (tensor<1x1024xf32> {tf_saved_model.index_path = ["output"]}) attributes {tf.entry_function = {control_outputs = "", inputs = "serving_default_input_tensor:0", outputs = "PartitionedCall:0"}, tf_saved_model.exported_names = ["serving_default"]} { + %0 = stablehlo.constant dense<0.134728625> : tensor<1x3xf32> + %1 = stablehlo.constant dense<-1.280000e+02> : tensor<1x1024xf32> + %2 = stablehlo.constant dense<0.003921567> : tensor<1x1024xf32> + %3 = stablehlo.divide %arg0, %2 : tensor<1x1024xf32> + %4 = stablehlo.add %3, %1 : tensor<1x1024xf32> + %5 = "tf.Identity"(%4) {device = ""} : (tensor<1x1024xf32>) -> tensor<1x1024xf32> + return %5 : tensor<1x1024xf32> + } + + // CHECK: %[[STABLEHLO_SUBGRAPH_TO_XLA_CALL_MODULE_OP:.*]] = "tf.XlaCallModule"(%arg0) <{Sout = [#tf_type.shape<1x1024>] + // CHECK-SAME: _entry_function = @_stablehlo_main_0 + // CHECK: %[[IDENTITY:.*]] = "tf.Identity"(%[[STABLEHLO_SUBGRAPH_TO_XLA_CALL_MODULE_OP]]) + // CHECK: return %[[IDENTITY]] + // CHECK } + +} + // ----- module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 1629 : i32}, tf_saved_model.semantics} { @@ -77,9 +110,9 @@ module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, p return %3 : tensor<1x3xf32> } - // CHECK: %[[STABLEHLO_SUBGRAPH_TO_XLA_CALL_MODULE_OP:.*]] = "tf.XlaCallModule"() {Sout = [#tf_type.shape<1024x3>], _entry_function = @_stablehlo_main_ - // CHECK: %[[CUSTOM_AGGREGATOR_0:.*]] = "tf.CustomAggregator"(%arg0) {calibration_method = 1 : i32, id = "0", initial_num_bins = 0 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32} : (tensor<1x1024xf32>) -> tensor<1x1024xf32> - // CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%[[CUSTOM_AGGREGATOR:.*]], %[[STABLEHLO_SUBGRAPH_TO_XLA_CALL_MODULE_OP:.*]]) {Sout = [#tf_type.shape<1x3>], _entry_function = @composite_dot_general_fn_1, _original_entry_function = "composite_dot_general_fn_1" + // CHECK: %[[STABLEHLO_SUBGRAPH_TO_XLA_CALL_MODULE_OP:.*]] = "tf.XlaCallModule"() <{Sout = [#tf_type.shape<1024x3>], {{.*}}}> {_entry_function = @_stablehlo_main_ + // CHECK: %[[CUSTOM_AGGREGATOR_0:.*]] = "tf.CustomAggregator"(%arg0) <{id = "0"}> {calibration_method = 1 : i32, initial_num_bins = 0 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32} : (tensor<1x1024xf32>) -> tensor<1x1024xf32> + // CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%[[CUSTOM_AGGREGATOR:.*]], %[[STABLEHLO_SUBGRAPH_TO_XLA_CALL_MODULE_OP:.*]]) <{Sout = [#tf_type.shape<1x3>], {{.*}}}> {_entry_function = @composite_dot_general_fn_1, _original_entry_function = "composite_dot_general_fn_1" // CHECK: %[[CUSTOM_AGGREGATOR_1:.*]] = "tf.CustomAggregator"(%[[XLA_CALL_MODULE:.*]]) // CHECK: return %[[CUSTOM_AGGREGATOR_1]] // CHECK: } @@ -109,8 +142,8 @@ module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, p } // CHECK: %[[CONSTANT:.*]] = stablehlo.constant dense<1.000000e+03> : tensor<1024x3xf32> - // CHECK: %[[CUSTOM_AGGREGATOR_0:.*]] = "tf.CustomAggregator"(%arg0) {calibration_method = 1 : i32, id = "0", initial_num_bins = 0 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32} : (tensor<1x1024xf32>) -> tensor<1x1024xf32> - // CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%[[CUSTOM_AGGREGATOR:.*]], %[[XLA_CALL_MODULE_EXTRACTED_FROM_SUBGRAPH:.*]]) {Sout = [#tf_type.shape<1x3>], _entry_function = @composite_dot_general_fn_1, _original_entry_function = "composite_dot_general_fn_1" + // CHECK: %[[CUSTOM_AGGREGATOR_0:.*]] = "tf.CustomAggregator"(%arg0) <{id = "0"}> {calibration_method = 1 : i32, initial_num_bins = 0 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32} : (tensor<1x1024xf32>) -> tensor<1x1024xf32> + // CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%[[CUSTOM_AGGREGATOR:.*]], %[[XLA_CALL_MODULE_EXTRACTED_FROM_SUBGRAPH:.*]]) <{Sout = [#tf_type.shape<1x3>], {{.*}}}> {_entry_function = @composite_dot_general_fn_1, _original_entry_function = "composite_dot_general_fn_1" // CHECK: %[[CUSTOM_AGGREGATOR_1:.*]] = "tf.CustomAggregator"(%[[XLA_CALL_MODULE:.*]]) // CHECK: return %[[CUSTOM_AGGREGATOR_1]] // CHECK: } diff --git a/tensorflow/compiler/mlir/quantization/stablehlo/uniform_quantized_types.cc b/tensorflow/compiler/mlir/quantization/stablehlo/uniform_quantized_types.cc index 197d5dd237e58d..bfd9de9ca60d25 100644 --- a/tensorflow/compiler/mlir/quantization/stablehlo/uniform_quantized_types.cc +++ b/tensorflow/compiler/mlir/quantization/stablehlo/uniform_quantized_types.cc @@ -37,6 +37,17 @@ UniformQuantizedType CreateI8F32UniformQuantizedType(const Location loc, /*storageTypeMin=*/llvm::minIntN(8), /*storageTypeMax=*/llvm::maxIntN(8)); } +UniformQuantizedType CreateI32F32UniformQuantizedType( + const Location loc, MLIRContext& context, const float scale, + const int32_t zero_point) { + return UniformQuantizedType::getChecked( + loc, /*flags=*/QuantizationFlags::Signed, + /*storageType=*/IntegerType::get(&context, /*width=*/32), + /*expressedType=*/FloatType::getF32(&context), scale, zero_point, + /*storageTypeMin=*/llvm::minIntN(32), + /*storageTypeMax=*/llvm::maxIntN(32)); +} + UniformQuantizedPerAxisType CreateI8F32UniformQuantizedPerAxisType( const Location loc, MLIRContext& context, const ArrayRef scales, const ArrayRef zero_points, const int quantization_dimension) { diff --git a/tensorflow/compiler/mlir/quantization/stablehlo/uniform_quantized_types.h b/tensorflow/compiler/mlir/quantization/stablehlo/uniform_quantized_types.h index 84f5ae274e047a..68774b2ecb876b 100644 --- a/tensorflow/compiler/mlir/quantization/stablehlo/uniform_quantized_types.h +++ b/tensorflow/compiler/mlir/quantization/stablehlo/uniform_quantized_types.h @@ -35,6 +35,16 @@ UniformQuantizedType CreateI8F32UniformQuantizedType(Location loc, float scale, int8_t zero_point); +// Creates a `UniformQuantizedType` with the given `scale` and `zero_point` +// values. The produced type has f32 as its expressed type and i32 as its +// storage type. The available values use the full range of the storage value. +// Assumes asymmetric quantization, meaning the zero point values can be +// non-zero values. +UniformQuantizedType CreateI32F32UniformQuantizedType(Location loc, + MLIRContext& context, + float scale, + int32_t zero_point); + // Creates a `UniformQuantizedPerAxisType` with the given `scales` and // `zero_points` values. The produced type has f32 as its expressed type and // i8 as its storage type. The available values use the full range of the diff --git a/tensorflow/compiler/mlir/quantization/stablehlo/uniform_quantized_types_test.cc b/tensorflow/compiler/mlir/quantization/stablehlo/uniform_quantized_types_test.cc index eb40dcb6afb617..0888bfa8d22908 100644 --- a/tensorflow/compiler/mlir/quantization/stablehlo/uniform_quantized_types_test.cc +++ b/tensorflow/compiler/mlir/quantization/stablehlo/uniform_quantized_types_test.cc @@ -15,6 +15,7 @@ limitations under the License. #include "tensorflow/compiler/mlir/quantization/stablehlo/uniform_quantized_types.h" #include +#include #include #include @@ -81,6 +82,60 @@ TEST_F(CreateI8F32UniformQuantizedTypeTest, HasScaleAndZeroPointProperlySet) { EXPECT_EQ(quantized_type.getZeroPoint(), 99); } +class CreateI32F32UniformQuantizedTypeTest : public ::testing::Test { + protected: + CreateI32F32UniformQuantizedTypeTest() : ctx_() { + ctx_.loadDialect(); + } + + MLIRContext ctx_; +}; + +TEST_F(CreateI32F32UniformQuantizedTypeTest, HasI32StorageType) { + const UniformQuantizedType quantized_type = + CreateI32F32UniformQuantizedType(UnknownLoc::get(&ctx_), ctx_, + /*scale=*/1.0, /*zero_point=*/0); + + EXPECT_TRUE(quantized_type.getStorageType().isSignlessInteger(32)); +} + +TEST_F(CreateI32F32UniformQuantizedTypeTest, HasF32ExpressedType) { + const UniformQuantizedType quantized_type = + CreateI32F32UniformQuantizedType(UnknownLoc::get(&ctx_), ctx_, + /*scale=*/1.0, /*zero_point=*/0); + + EXPECT_TRUE(quantized_type.getExpressedType().isF32()); +} + +TEST_F(CreateI32F32UniformQuantizedTypeTest, IsSigned) { + const UniformQuantizedType quantized_type = + CreateI32F32UniformQuantizedType(UnknownLoc::get(&ctx_), ctx_, + /*scale=*/1.0, /*zero_point=*/0); + + EXPECT_TRUE(quantized_type.isSigned()); +} + +TEST_F(CreateI32F32UniformQuantizedTypeTest, + SotrageTypeMinMaxEqualToI32MinMax) { + const UniformQuantizedType quantized_type = + CreateI32F32UniformQuantizedType(UnknownLoc::get(&ctx_), ctx_, + /*scale=*/1.0, /*zero_point=*/0); + + EXPECT_EQ(quantized_type.getStorageTypeMin(), + std::numeric_limits::min()); + EXPECT_EQ(quantized_type.getStorageTypeMax(), + std::numeric_limits::max()); +} + +TEST_F(CreateI32F32UniformQuantizedTypeTest, HasScaleAndZeroPointProperlySet) { + const UniformQuantizedType quantized_type = + CreateI32F32UniformQuantizedType(UnknownLoc::get(&ctx_), ctx_, + /*scale=*/8.0, /*zero_point=*/1111); + + EXPECT_EQ(quantized_type.getScale(), 8.0); + EXPECT_EQ(quantized_type.getZeroPoint(), 1111); +} + class CreateI8F32UniformQuantizedPerAxisTypeTest : public ::testing::Test { protected: CreateI8F32UniformQuantizedPerAxisTypeTest() : ctx_() { diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/calibrator/BUILD b/tensorflow/compiler/mlir/quantization/tensorflow/calibrator/BUILD index aec612c95b7b62..286e0e8799121f 100644 --- a/tensorflow/compiler/mlir/quantization/tensorflow/calibrator/BUILD +++ b/tensorflow/compiler/mlir/quantization/tensorflow/calibrator/BUILD @@ -117,6 +117,15 @@ cc_library( ], ) +cc_library( + name = "id_assigner", + hdrs = ["id_assigner.h"], + compatible_with = get_compatible_with_portable(), + deps = [ + "//tensorflow/compiler/mlir/quantization/tensorflow:exported_model_proto_cc", + ], +) + pytype_strict_library( name = "calibration_algorithm", srcs = ["calibration_algorithm.py"], diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/calibrator/id_assigner.h b/tensorflow/compiler/mlir/quantization/tensorflow/calibrator/id_assigner.h new file mode 100644 index 00000000000000..ae75d9f579b090 --- /dev/null +++ b/tensorflow/compiler/mlir/quantization/tensorflow/calibrator/id_assigner.h @@ -0,0 +1,37 @@ +/* Copyright 2023 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ +#ifndef TENSORFLOW_COMPILER_MLIR_QUANTIZATION_TENSORFLOW_CALIBRATOR_ID_ASSIGNER_H_ +#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_TENSORFLOW_CALIBRATOR_ID_ASSIGNER_H_ + +#include "tensorflow/compiler/mlir/quantization/tensorflow/exported_model.pb.h" + +namespace tensorflow::quantization { + +// An interface that assigns UUIDs to CustomAggregator ops. +class CustomAggregatorIdAssigner { + public: + virtual ~CustomAggregatorIdAssigner() = default; + + // Assigns UUIDs to each CustomAggregator op found in each GraphDef in + // `exported_model`. The UUIDs are set to the `id` attributes. The UUIDs will + // be used during calibration step to identify the collected quantization + // statistics for each CustsomAggregator op. + virtual ExportedModel AssignIds( + const ExportedModel& exported_model) const = 0; +}; + +} // namespace tensorflow::quantization + +#endif // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_TENSORFLOW_CALIBRATOR_ID_ASSIGNER_H_ diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/ops/tf_op_quant_spec.cc b/tensorflow/compiler/mlir/quantization/tensorflow/ops/tf_op_quant_spec.cc index d415e755e9c6f6..22f95fa5369215 100644 --- a/tensorflow/compiler/mlir/quantization/tensorflow/ops/tf_op_quant_spec.cc +++ b/tensorflow/compiler/mlir/quantization/tensorflow/ops/tf_op_quant_spec.cc @@ -45,8 +45,8 @@ bool IsOpWithQuantizableTrait(Operation* op) { // Supported quantizable ops. return isa(op); + TF::ResourceGatherOp, TF::DepthwiseConv2dNativeOp, TF::Conv3DOp, + TF::BatchMatMulV2Op, TF::EinsumOp>(op); } bool IsOpWithInt8TypeOperand(Operation* op) { diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/passes/quantize_composite_functions.cc b/tensorflow/compiler/mlir/quantization/tensorflow/passes/quantize_composite_functions.cc index d6ad4605863116..be72660f10091c 100644 --- a/tensorflow/compiler/mlir/quantization/tensorflow/passes/quantize_composite_functions.cc +++ b/tensorflow/compiler/mlir/quantization/tensorflow/passes/quantize_composite_functions.cc @@ -144,7 +144,7 @@ class QuantizeCompositeFunctionsPass "drq", "Post-training dynamic-range quantizaiton"), clEnumValN(tensorflow::quantization::QuantizationMethod:: METHOD_STATIC_RANGE_WEIGHT_ONLY_INT8, - "weight_only", "Post-training weight-only quantizaiton"))}; + "weight_only", "Post-training weight-only quantization"))}; Option target_opset_{ *this, "target-opset", llvm::cl::init(OpSet::TF), @@ -1283,9 +1283,7 @@ void QuantizeCompositeFunctionsPass::runOnOperation() { ctx, target_opset_); patterns_2.add(ctx, target_opset_); - if (target_opset_ == OpSet::XLA && enable_per_channel_quantization_ && - quantization_method_ == tensorflow::quantization::QuantizationMethod:: - METHOD_STATIC_RANGE_WEIGHT_ONLY_INT8) { + if (target_opset_ == OpSet::XLA && enable_per_channel_quantization_) { patterns_2.add(ctx); } diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/python/BUILD b/tensorflow/compiler/mlir/quantization/tensorflow/python/BUILD index 5b1815fc4fd8f8..ac1b64848b5bc1 100644 --- a/tensorflow/compiler/mlir/quantization/tensorflow/python/BUILD +++ b/tensorflow/compiler/mlir/quantization/tensorflow/python/BUILD @@ -108,8 +108,8 @@ tf_python_pybind_extension( "//tensorflow/compiler/mlir/quantization/tensorflow:quantization_options_proto_cc", "//tensorflow/compiler/mlir/quantization/tensorflow/calibrator:calibration_statistics_proto_cc", "//tensorflow/compiler/mlir/quantization/tensorflow/calibrator:calibrator_singleton", + "//tensorflow/compiler/mlir/quantization/tensorflow/calibrator:id_assigner", "//tensorflow/python/lib/core:pybind11_lib", - "@com_google_absl//absl/base:core_headers", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", @@ -178,7 +178,6 @@ pytype_strict_library( "//tensorflow/python/framework:ops", "//tensorflow/python/framework:tensor_conversion", "//tensorflow/python/lib/io:file_io", - "//tensorflow/python/platform:tf_logging", "//tensorflow/python/saved_model:load", "//tensorflow/python/saved_model:loader", "//tensorflow/python/saved_model:signature_constants", diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/python/integration_test/quantize_model_test.py b/tensorflow/compiler/mlir/quantization/tensorflow/python/integration_test/quantize_model_test.py index d7156be21fb843..c96e3a776c79af 100644 --- a/tensorflow/compiler/mlir/quantization/tensorflow/python/integration_test/quantize_model_test.py +++ b/tensorflow/compiler/mlir/quantization/tensorflow/python/integration_test/quantize_model_test.py @@ -1548,7 +1548,7 @@ def gen_data() -> repr_dataset.RepresentativeDataset: 'dilations': [1, 2, 2, 1], }, { - 'testcase_name': 'with_bias_and_relu6_to_xla', + 'testcase_name': 'with_bias_and_relu6_to_xla_per_tensor', 'activation_fn': nn_ops.relu6, 'has_bias': True, 'has_batch_norm': False, @@ -1556,6 +1556,15 @@ def gen_data() -> repr_dataset.RepresentativeDataset: 'input_shape_dynamic': False, 'enable_per_channel_quantization': False, }, + { + 'testcase_name': 'with_bias_and_relu6_to_xla_per_channel', + 'activation_fn': nn_ops.relu6, + 'has_bias': True, + 'has_batch_norm': False, + 'target_opset': quant_opts_pb2.XLA, + 'input_shape_dynamic': False, + 'enable_per_channel_quantization': True, + }, { 'testcase_name': 'dilation_with_bias_and_relu6_to_xla', 'activation_fn': nn_ops.relu6, @@ -1566,6 +1575,16 @@ def gen_data() -> repr_dataset.RepresentativeDataset: 'enable_per_channel_quantization': False, 'dilations': [1, 2, 2, 1], }, + { + 'testcase_name': 'dilation_with_bias_and_relu6_to_xla_per_channel', + 'activation_fn': nn_ops.relu6, + 'has_bias': True, + 'has_batch_norm': False, + 'target_opset': quant_opts_pb2.XLA, + 'input_shape_dynamic': False, + 'enable_per_channel_quantization': True, + 'dilations': [1, 2, 2, 1], + }, { 'testcase_name': 'with_bias_and_bn_and_relu6_to_xla', 'activation_fn': nn_ops.relu6, @@ -1575,6 +1594,15 @@ def gen_data() -> repr_dataset.RepresentativeDataset: 'input_shape_dynamic': False, 'enable_per_channel_quantization': False, }, + { + 'testcase_name': 'with_bias_and_bn_and_relu6_to_xla_per_channel', + 'activation_fn': nn_ops.relu6, + 'has_bias': True, + 'has_batch_norm': True, + 'target_opset': quant_opts_pb2.XLA, + 'input_shape_dynamic': False, + 'enable_per_channel_quantization': True, + }, { 'testcase_name': 'dilation_with_bias_and_bn_and_relu6_to_xla', 'activation_fn': nn_ops.relu6, @@ -1585,6 +1613,18 @@ def gen_data() -> repr_dataset.RepresentativeDataset: 'enable_per_channel_quantization': False, 'dilations': [1, 2, 2, 1], }, + { + 'testcase_name': ( + 'dilation_with_bias_and_bn_and_relu6_to_xla_per_channel' + ), + 'activation_fn': nn_ops.relu6, + 'has_bias': True, + 'has_batch_norm': True, + 'target_opset': quant_opts_pb2.XLA, + 'input_shape_dynamic': False, + 'enable_per_channel_quantization': True, + 'dilations': [1, 2, 2, 1], + }, { 'testcase_name': 'with_bias_and_relu6_to_xla_dynamic', 'activation_fn': nn_ops.relu6, @@ -1594,6 +1634,15 @@ def gen_data() -> repr_dataset.RepresentativeDataset: 'input_shape_dynamic': True, 'enable_per_channel_quantization': False, }, + { + 'testcase_name': 'with_bias_and_relu6_to_xla_dynamic_per_channel', + 'activation_fn': nn_ops.relu6, + 'has_bias': True, + 'has_batch_norm': False, + 'target_opset': quant_opts_pb2.XLA, + 'input_shape_dynamic': True, + 'enable_per_channel_quantization': True, + }, { 'testcase_name': 'dilation_with_bias_and_relu6_to_xla_dynamic', 'activation_fn': nn_ops.relu6, @@ -1604,6 +1653,18 @@ def gen_data() -> repr_dataset.RepresentativeDataset: 'enable_per_channel_quantization': False, 'dilations': [1, 2, 2, 1], }, + { + 'testcase_name': ( + 'dilation_with_bias_and_relu6_to_xla_dynamic_per_channel' + ), + 'activation_fn': nn_ops.relu6, + 'has_bias': True, + 'has_batch_norm': False, + 'target_opset': quant_opts_pb2.XLA, + 'input_shape_dynamic': True, + 'enable_per_channel_quantization': True, + 'dilations': [1, 2, 2, 1], + }, { 'testcase_name': 'with_bias_and_bn_and_relu6_to_xla_dynamic', 'activation_fn': nn_ops.relu6, @@ -1613,6 +1674,17 @@ def gen_data() -> repr_dataset.RepresentativeDataset: 'input_shape_dynamic': True, 'enable_per_channel_quantization': False, }, + { + 'testcase_name': ( + 'with_bias_and_bn_and_relu6_to_xla_dynamic_per_channel' + ), + 'activation_fn': nn_ops.relu6, + 'has_bias': True, + 'has_batch_norm': True, + 'target_opset': quant_opts_pb2.XLA, + 'input_shape_dynamic': True, + 'enable_per_channel_quantization': True, + }, { 'testcase_name': 'dilation_with_bias_and_bn_and_relu6_to_xla_dynamic', 'activation_fn': nn_ops.relu6, @@ -1623,6 +1695,18 @@ def gen_data() -> repr_dataset.RepresentativeDataset: 'enable_per_channel_quantization': False, 'dilations': [1, 2, 2, 1], }, + { + 'testcase_name': ( + 'dilation_with_bias_and_bn_and_relu6_to_xla_dynamic_per_channel' + ), + 'activation_fn': nn_ops.relu6, + 'has_bias': True, + 'has_batch_norm': True, + 'target_opset': quant_opts_pb2.XLA, + 'input_shape_dynamic': True, + 'enable_per_channel_quantization': True, + 'dilations': [1, 2, 2, 1], + }, { 'testcase_name': 'with_bias_and_relu6_to_uq', 'activation_fn': nn_ops.relu6, @@ -1787,6 +1871,28 @@ def data_gen() -> repr_dataset.RepresentativeDataset: if target_opset == quant_opts_pb2.XLA: self.assertTrue(self._contains_op(output_graphdef, 'XlaConvV2')) + if enable_per_channel_quantization: + per_channel_size_attr = attr_value_pb2.AttrValue( + list=attr_value_pb2.AttrValue.ListValue( + shape=[ + tensor_shape_pb2.TensorShapeProto( + dim=[ + tensor_shape_pb2.TensorShapeProto.Dim( + size=filter_shape[-1] + ) + ] + ) + ] + ) + ) + self.assertTrue( + self._contains_op( + output_graphdef, + 'Const', + '_output_shapes', + per_channel_size_attr, + ) + ) elif target_opset == quant_opts_pb2.UNIFORM_QUANTIZED: self.assertTrue( self._contains_op(output_graphdef, 'UniformQuantizedConvolution') @@ -2050,6 +2156,15 @@ def data_gen() -> repr_dataset.RepresentativeDataset: 'input_shape_dynamic': False, 'enable_per_channel_quantization': False, }, + { + 'testcase_name': 'with_bias_and_relu6_to_xla_per_channel', + 'activation_fn': nn_ops.relu6, + 'has_bias': True, + 'has_batch_norm': False, + 'target_opset': quant_opts_pb2.XLA, + 'input_shape_dynamic': False, + 'enable_per_channel_quantization': True, + }, { 'testcase_name': 'with_bias_and_bn_and_relu6_to_xla', 'activation_fn': nn_ops.relu6, @@ -2059,6 +2174,15 @@ def data_gen() -> repr_dataset.RepresentativeDataset: 'input_shape_dynamic': False, 'enable_per_channel_quantization': False, }, + { + 'testcase_name': 'with_bias_and_bn_and_relu6_to_xla_per_channel', + 'activation_fn': nn_ops.relu6, + 'has_bias': True, + 'has_batch_norm': True, + 'target_opset': quant_opts_pb2.XLA, + 'input_shape_dynamic': False, + 'enable_per_channel_quantization': True, + }, { 'testcase_name': 'with_bias_and_relu6_to_xla_dynamic', 'activation_fn': nn_ops.relu6, @@ -2068,6 +2192,15 @@ def data_gen() -> repr_dataset.RepresentativeDataset: 'input_shape_dynamic': True, 'enable_per_channel_quantization': False, }, + { + 'testcase_name': 'with_bias_and_relu6_to_xla_dynamic_per_channel', + 'activation_fn': nn_ops.relu6, + 'has_bias': True, + 'has_batch_norm': False, + 'target_opset': quant_opts_pb2.XLA, + 'input_shape_dynamic': True, + 'enable_per_channel_quantization': True, + }, { 'testcase_name': 'with_bias_and_bn_and_relu6_to_xla_dynamic', 'activation_fn': nn_ops.relu6, @@ -2077,6 +2210,17 @@ def data_gen() -> repr_dataset.RepresentativeDataset: 'input_shape_dynamic': True, 'enable_per_channel_quantization': False, }, + { + 'testcase_name': ( + 'with_bias_and_bn_and_relu6_to_xla_dynamic_per_channel' + ), + 'activation_fn': nn_ops.relu6, + 'has_bias': True, + 'has_batch_norm': True, + 'target_opset': quant_opts_pb2.XLA, + 'input_shape_dynamic': True, + 'enable_per_channel_quantization': True, + }, { 'testcase_name': 'with_bias_and_relu6_to_uq', 'activation_fn': nn_ops.relu6, @@ -2172,6 +2316,28 @@ def data_gen() -> repr_dataset.RepresentativeDataset: self.assertTrue( self._contains_op(output_graphdef, 'DepthwiseConv2dNative') ) + if enable_per_channel_quantization: + per_channel_size_attr = attr_value_pb2.AttrValue( + list=attr_value_pb2.AttrValue.ListValue( + shape=[ + tensor_shape_pb2.TensorShapeProto( + dim=[ + tensor_shape_pb2.TensorShapeProto.Dim( + size=filter_shape[-1] * filter_shape[2] + ) + ] + ) + ] + ) + ) + self.assertTrue( + self._contains_op( + output_graphdef, + 'Const', + '_output_shapes', + per_channel_size_attr, + ) + ) elif target_opset == quant_opts_pb2.UNIFORM_QUANTIZED: self.assertTrue( self._contains_op(output_graphdef, 'UniformQuantizedConvolution') @@ -2337,11 +2503,6 @@ def data_gen() -> repr_dataset.RepresentativeDataset: else: self.assertAllClose(new_outputs, expected_outputs, atol=0.13) - # NOTE: Isolated the most basic configuration from `test_matmul_ptq_model` - # for StableHLO PTQ prototype testing while integrating. Please note this - # test is for intermediate testing purposes as the migration is not complete. - # TODO: b/298581932 - Add the full test case for STABLEHLO opset once - # migration is complete. @test_util.run_in_graph_and_eager_modes def test_matmul_ptq_model_stablehlo(self): activation_fn = None @@ -2353,7 +2514,7 @@ def test_matmul_ptq_model_stablehlo(self): input_shape = (*lhs_batch_size, 1, 1024) filter_shape = (*rhs_batch_size, 1024, 3) static_input_shape = [dim if dim is not None else 2 for dim in input_shape] - self._create_matmul_model( + model = self._create_matmul_model( input_shape, filter_shape, self._input_saved_model_path, @@ -2362,38 +2523,47 @@ def test_matmul_ptq_model_stablehlo(self): ) rng = np.random.default_rng(seed=1234) + input_data = ops.convert_to_tensor( + rng.uniform(low=0.0, high=1.0, size=static_input_shape).astype( + np.float32 + ) + ) + expected_outputs = model.matmul(input_data) + def data_gen() -> repr_dataset.RepresentativeDataset: - for _ in range(5): + for _ in range(100): yield { 'input_tensor': rng.uniform( low=0.0, high=1.0, size=static_input_shape ).astype(np.float32) } - tags = {tag_constants.SERVING} - quantization_options = quant_opts_pb2.QuantizationOptions( quantization_method=quant_opts_pb2.QuantizationMethod( preset_method=_PresetMethod.METHOD_STATIC_RANGE_INT8 ), - tags=tags, + tags={tag_constants.SERVING}, signature_keys=['serving_default'], op_set=target_opset, ) - # TODO: b/299545836 - Remove exception handling below after migrating - # StableHLO export passes. - with self.assertRaisesRegex( # pylint: disable=g-error-prone-assert-raises - Exception, - "Failed to convert MLIR to GraphDef. op node 'quantfork.stats' was not" - ' a TF op', - ): - converted_model = quantize_model.quantize( - self._input_saved_model_path, - self._output_saved_model_path, - quantization_options, - representative_dataset=data_gen(), - ) - self.assertIsNotNone(converted_model) + converted_model = quantize_model.quantize( + self._input_saved_model_path, + self._output_saved_model_path, + quantization_options, + representative_dataset=data_gen(), + ) + + self.assertIsNotNone(converted_model) + self.assertCountEqual( + converted_model.signatures._signatures.keys(), {'serving_default'} + ) + + new_outputs = converted_model.signatures['serving_default']( + input_tensor=ops.convert_to_tensor(input_data) + ) + # Tests that the quantized graph outputs similar values. The rtol value is + # arbitrary. + self.assertAllClose(new_outputs, expected_outputs, rtol=0.02) @parameterized.named_parameters( { diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/python/pywrap_quantize_model.cc b/tensorflow/compiler/mlir/quantization/tensorflow/python/pywrap_quantize_model.cc index 97d36d4984f3e3..c08aaf78376f8b 100644 --- a/tensorflow/compiler/mlir/quantization/tensorflow/python/pywrap_quantize_model.cc +++ b/tensorflow/compiler/mlir/quantization/tensorflow/python/pywrap_quantize_model.cc @@ -34,6 +34,7 @@ limitations under the License. #include "pybind11_protobuf/native_proto_caster.h" // from @pybind11_protobuf #include "tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics.pb.h" #include "tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibrator_singleton.h" +#include "tensorflow/compiler/mlir/quantization/tensorflow/calibrator/id_assigner.h" #include "tensorflow/compiler/mlir/quantization/tensorflow/exported_model.pb.h" #include "tensorflow/compiler/mlir/quantization/tensorflow/python/quantize_model.h" #include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h" @@ -43,6 +44,7 @@ namespace { using ::tensorflow::calibrator::CalibrationStatistics; using ::tensorflow::calibrator::CalibratorSingleton; +using ::tensorflow::quantization::CustomAggregatorIdAssigner; using ::tensorflow::quantization::ExportedModel; using ::tensorflow::quantization::QuantizationOptions; using ::tensorflow::quantization::QuantizePtqDynamicRange; @@ -86,8 +88,8 @@ CalibrationStatistics GetStatisticsFromCalibrator(const absl::string_view id) { namespace pybind11 { namespace detail { -// Converts `ExportedModel` (c++) to `bytes` (python). The resulting `bytes` -// object is a serialization of `ExportedModel`. +// Handles `ExportedModel` (c++) <-> `bytes` (python) conversion. The `bytes` +// object in the python layer is a serialization of `ExportedModel`. // // See https://pybind11.readthedocs.io/en/stable/advanced/cast/custom.html for // further details on how custom type conversions work for pybind11. @@ -96,6 +98,21 @@ struct type_caster { public: PYBIND11_TYPE_CASTER(ExportedModel, const_name("ExportedModel")); + // Loads an `ExportedModel` instance from a python `bytes` object (`src`). + bool load(handle src, const bool convert) { + auto caster = make_caster(); + // Make sure the user passed a valid python string. + if (!caster.load(src, convert)) { + return false; + } + + const absl::string_view exported_model_serialized = + cast_op(std::move(caster)); + + // NOLINTNEXTLINE: Explicit std::string conversion required for OSS. + return value.ParseFromString(std::string(exported_model_serialized)); + } + // Constructs a `bytes` object after serializing `src`. static handle cast(ExportedModel&& src, return_value_policy policy, handle parent) { @@ -103,6 +120,14 @@ struct type_caster { // destruction of py::bytes and returns a raw python object handle. return py::bytes(Serialize(src)).release(); } + + // Constructs a `bytes` object after serializing `src`. + static handle cast(const ExportedModel& src, return_value_policy policy, + handle parent) { + // release() prevents the reference count from decreasing upon the + // destruction of py::bytes and returns a raw python object handle. + return py::bytes(Serialize(src)).release(); + } }; // Python -> cpp conversion for `QuantizationOptions`. Accepts a serialized @@ -130,10 +155,30 @@ struct type_caster { } // namespace detail } // namespace pybind11 +namespace { + +// A "trampoline" class that redirects virtual function calls to the python +// implementation. +// +// Reference: +// https://pybind11.readthedocs.io/en/stable/advanced/classes.html#overriding-virtual-functions-in-python +class CustomAggregatorIdAssignerTrampoline : public CustomAggregatorIdAssigner { + public: + using CustomAggregatorIdAssigner::CustomAggregatorIdAssigner; + + ExportedModel AssignIds(const ExportedModel& exported_model) const override { + PYBIND11_OVERRIDE_PURE(ExportedModel, CustomAggregatorIdAssigner, + assign_ids, exported_model); + } +}; + +} // namespace + PYBIND11_MODULE(pywrap_quantize_model, m) { // Supports absl::StatusOr type conversions. pybind11::google::ImportStatusModule(); pybind11_protobuf::ImportNativeProtoCasters(); + // Calibrator related functions. m.def( "clear_calibrator", @@ -156,6 +201,14 @@ PYBIND11_MODULE(pywrap_quantize_model, m) { Returns the proto CalibrationStatistics given id from calibrator. )pbdoc"); + // Exports `CustomAggregatorIdAssigner` class. A pure virtual member function + // `AssignIds` is mapped to `assign_ids` in python, which is expected to be + // inherited and overridden. + py::class_( + m, "CustomAggregatorIdAssigner") + .def(py::init<>()) + .def("assign_ids", &CustomAggregatorIdAssigner::AssignIds); + // Quantization functions. m.def( "quantize_qat_model", @@ -218,11 +271,17 @@ PYBIND11_MODULE(pywrap_quantize_model, m) { const std::vector& signature_keys, const std::unordered_set& tags, const QuantizationOptions& quant_opts, - const absl::flat_hash_map& function_aliases) + const absl::flat_hash_map& function_aliases, + const CustomAggregatorIdAssigner& custom_aggregator_id_assigner) -> absl::StatusOr { - return QuantizePtqModelPreCalibration(saved_model_path, signature_keys, - tags, quant_opts, - function_aliases); + const absl::StatusOr exported_model = + QuantizePtqModelPreCalibration(saved_model_path, signature_keys, + tags, quant_opts, function_aliases); + if (!exported_model.ok()) { + return exported_model.status(); + } + + return custom_aggregator_id_assigner.AssignIds(*exported_model); }, R"pbdoc( Returns serialized ExportedModel that contains the model's GraphDef and @@ -230,6 +289,10 @@ PYBIND11_MODULE(pywrap_quantize_model, m) { user should pass a serialized `QuantizationOptions` for the `quant_opts` argument. + The argument `custom_aggregator_id_assigner` is an instance of + `CustomAggregatorIdAssigner` whose virtual function `assign_ids` is + implemented in python. + Raises `StatusNotOk` exception if when the run was unsuccessful. )pbdoc"); diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/python/quantize_model.py b/tensorflow/compiler/mlir/quantization/tensorflow/python/quantize_model.py index 43dbd283a67fdf..b83461ba0ed542 100644 --- a/tensorflow/compiler/mlir/quantization/tensorflow/python/quantize_model.py +++ b/tensorflow/compiler/mlir/quantization/tensorflow/python/quantize_model.py @@ -740,6 +740,38 @@ def _get_saver_def_or_none( return None +class CustomAggregatorIdAssigner( + pywrap_quantize_model.CustomAggregatorIdAssigner +): + """Python impl. of `pywrap_quantize_model.CustomAggregatorIdAssigner`. + + The interface is defined in the C++ layer, exposing a pure virtual function + `assign_ids`. + """ + + def assign_ids(self, exported_model_serialized: bytes) -> bytes: + """Assigns UUIDs to each CustomAggregator op find in the graph def. + + Args: + exported_model_serialized: Serialized `ExportedModel` instance. + + Returns: + Serialized `ExportedModel` whose CustomAggregator ops are assigned UUIDs + to their `id` attributes. + """ + exported_model = exported_model_pb2.ExportedModel.FromString( + exported_model_serialized + ) + + graph_def = exported_model.graph_def + for function_def in graph_def.library.function: + for node_def in function_def.node_def: + if node_def.op == 'CustomAggregator': + node_def.attr['id'].s = uuid.uuid4().hex.encode('ascii') + + return exported_model.SerializeToString() + + def _run_static_range_ptq( src_saved_model_path: str, dst_saved_model_path: str, @@ -780,6 +812,7 @@ def _run_static_range_ptq( set(quant_opts.tags), quant_opts.SerializeToString(), dict(function_aliases), + CustomAggregatorIdAssigner(), ) ) @@ -788,11 +821,6 @@ def _run_static_range_ptq( ) graph_def = exported_model.graph_def - for function_def in graph_def.library.function: - for node_def in function_def.node_def: - if node_def.op == 'CustomAggregator': - node_def.attr['id'].s = uuid.uuid4().hex.encode('ascii') - pre_calib_output_model_path = tempfile.mkdtemp() save_model.save_model_v1( graph_def, @@ -1376,7 +1404,7 @@ def _populate_quantization_options_default_values( quantization_options.min_num_elements_for_weights = ( _DYNAMIC_RANGE_DEFAULT_MIN_NUM_ELEMENTS_FOR_WEIGHTS ) - logging.warn( + logging.warning( ( 'QuantizationOptions.min_num_elements_for_weights is not set (0).' ' Setting to the default value: %d.' @@ -1384,15 +1412,23 @@ def _populate_quantization_options_default_values( _DYNAMIC_RANGE_DEFAULT_MIN_NUM_ELEMENTS_FOR_WEIGHTS, ) - # TODO(b/281595329): Implement static range quantization per-channel support + # TODO: b/307900054 - Set the per-channel quantization by default. if quantization_options.enable_per_channel_quantization and not ( - quantization_options.op_set == quant_opts_pb2.OpSet.UNIFORM_QUANTIZED - or quantization_options.quantization_method.preset_method - == _PresetMethod.METHOD_STATIC_RANGE_WEIGHT_ONLY_INT8 + ( + quantization_options.op_set == quant_opts_pb2.OpSet.UNIFORM_QUANTIZED + or quantization_options.quantization_method.preset_method + == _PresetMethod.METHOD_STATIC_RANGE_WEIGHT_ONLY_INT8 + ) + or ( + quantization_options.op_set == quant_opts_pb2.OpSet.XLA + and quantization_options.quantization_method.preset_method + == _PresetMethod.METHOD_STATIC_RANGE_INT8 + ) ): raise ValueError( - 'Currently, per-channel quantization is supported for Uniform ' - 'Quantized opset and Weight-only.' + 'Currently, per-channel quantization is supported for Uniform Quantized' + ' opset, weight only quantization, or XLA opset with static range' + ' quantization.' ) if ( diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.proto b/tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.proto index 01cf3401f9a37a..4384a13c8e05eb 100644 --- a/tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.proto +++ b/tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.proto @@ -289,9 +289,10 @@ message QuantizationOptions { // If not set, it defaults to `true`. optional bool freeze_all_variables = 9; - // Enables chnanel-wise quantizaiton. By default, channel-wise quantization is + // Enables channel-wise quantization. By default, channel-wise quantization is // not applied regardless of the op support. Currently, it is supported for - // Uniform Quantized opset only. + // XLA opset for SRQ on weight tensors (not activation), + // and Uniform Quantized opset . bool enable_per_channel_quantization = 10; // Enables two inputs of an operation to be both tensors. diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/quantize_passes.cc b/tensorflow/compiler/mlir/quantization/tensorflow/quantize_passes.cc index 6fb3c91b15e129..0b9cdc09ca5b93 100644 --- a/tensorflow/compiler/mlir/quantization/tensorflow/quantize_passes.cc +++ b/tensorflow/compiler/mlir/quantization/tensorflow/quantize_passes.cc @@ -52,7 +52,7 @@ void AddStablehloQuantToIntPasses(mlir::PassManager &pm) { void AddStaticRangeQuantizationPass( mlir::PassManager &pm, const QuantizationOptions &quantization_options, std::optional mlir_dump_file_prefix) { - // TODO: b/299545840 - Include QuantizeCompositeFunctionsPass as in bug. + pm.addPass(mlir::quant::stablehlo::createQuantizeCompositeFunctionsPass()); } void AddConvertTpuToCpuModelPasses(mlir::PassManager &pm) { diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/tests/add_dump_tensor_op.mlir b/tensorflow/compiler/mlir/quantization/tensorflow/tests/add_dump_tensor_op.mlir index 4ee9b746fff9f6..bd3546299c1967 100644 --- a/tensorflow/compiler/mlir/quantization/tensorflow/tests/add_dump_tensor_op.mlir +++ b/tensorflow/compiler/mlir/quantization/tensorflow/tests/add_dump_tensor_op.mlir @@ -23,21 +23,21 @@ module { } // WholeModel-LABEL: func @conv -// WholeModel-DAG: %[[w:.*]] = "tf.Const"() {value = dense<{{\[\[\[\[}}1.600000e-01, 1.000000e-01 -// WholeModel-DAG: %[[b:.*]] = "tf.Const"() {value = dense<[-2.000000e+00, 3.000000e+00 -// WholeModel-DAG: %[[output0:.*]] = "tf.PartitionedCall"(%arg0, %[[w]], %[[b]]) {config = "", config_proto = "", executor_type = "", f = @composite_conv2d_with_bias_and_relu6_fn_2} -// WholeModel-DAG: %[[output1:.*]] = "tf.PartitionedCall"(%arg0, %[[w]], %[[b]]) {_tfl_quant_trait = "fully_quantizable", config = "", config_proto = "", executor_type = "", f = @composite_conv2d_with_bias_and_relu6_fn_1} -// WholeModel-DAG: "tf.DumpTensor"(%[[output1]]) {enabled = false, file_name = "unquantized_tensor_data.pb", func_name = "conv", log_dir_path = "/tmp/dumps/composite_conv2d_with_bias_and_relu6_fn_1", node_name = "Conv2D_1"} : (tensor<*xf32>) -> () +// WholeModel-DAG: %[[w:.*]] = "tf.Const"() <{value = dense<{{\[\[\[\[}}1.600000e-01, 1.000000e-01 +// WholeModel-DAG: %[[b:.*]] = "tf.Const"() <{value = dense<[-2.000000e+00, 3.000000e+00 +// WholeModel-DAG: %[[output0:.*]] = "tf.PartitionedCall"(%arg0, %[[w]], %[[b]]) <{config = "", config_proto = "", executor_type = "", f = @composite_conv2d_with_bias_and_relu6_fn_2}> +// WholeModel-DAG: %[[output1:.*]] = "tf.PartitionedCall"(%arg0, %[[w]], %[[b]]) <{config = "", config_proto = "", executor_type = "", f = @composite_conv2d_with_bias_and_relu6_fn_1}> {_tfl_quant_trait = "fully_quantizable"} +// WholeModel-DAG: "tf.DumpTensor"(%[[output1]]) <{enabled = false, file_name = "unquantized_tensor_data.pb", func_name = "conv", log_dir_path = "/tmp/dumps/composite_conv2d_with_bias_and_relu6_fn_1", node_name = "Conv2D_1"}> : (tensor<*xf32>) -> () // WholeModel-DAG: return %[[output0]], %[[output1]] // PerLayer-LABEL: func @conv -// PerLayer-DAG: %[[w:.*]] = "tf.Const"() {value = dense<{{\[\[\[\[}}1.600000e-01, 1.000000e-01 -// PerLayer-DAG: %[[b:.*]] = "tf.Const"() {value = dense<[-2.000000e+00, 3.000000e+00 -// PerLayer-DAG: %[[output0:.*]] = "tf.PartitionedCall"(%arg0, %[[w]], %[[b]]) {config = "", config_proto = "", executor_type = "", f = @composite_conv2d_with_bias_and_relu6_fn_2} -// PerLayer-DAG: %[[output1_quantized:.*]] = "tf.PartitionedCall"(%arg0, %[[w]], %[[b]]) {_tfl_quant_trait = "fully_quantizable", config = "", config_proto = "", executor_type = "", f = @composite_conv2d_with_bias_and_relu6_fn_1} -// PerLayer-DAG: %[[output1_unquantized:.*]] = "tf.PartitionedCall"(%arg0, %cst, %cst_0) {config = "", config_proto = "", executor_type = "", f = @composite_conv2d_with_bias_and_relu6_fn_1_0} -// PerLayer-DAG: "tf.DumpTensor"(%[[output1_quantized]]) {enabled = false, file_name = "quantized_tensor_data.pb", func_name = "conv", log_dir_path = "/tmp/dumps/composite_conv2d_with_bias_and_relu6_fn_1", node_name = "Conv2D_1"} : (tensor<*xf32>) -> () -// PerLayer-DAG: "tf.DumpTensor"(%[[output1_unquantized]]) {enabled = false, file_name = "unquantized_tensor_data.pb", func_name = "conv", log_dir_path = "/tmp/dumps/composite_conv2d_with_bias_and_relu6_fn_1", node_name = "Conv2D_1"} : (tensor<*xf32>) -> () +// PerLayer-DAG: %[[w:.*]] = "tf.Const"() <{value = dense<{{\[\[\[\[}}1.600000e-01, 1.000000e-01 +// PerLayer-DAG: %[[b:.*]] = "tf.Const"() <{value = dense<[-2.000000e+00, 3.000000e+00 +// PerLayer-DAG: %[[output0:.*]] = "tf.PartitionedCall"(%arg0, %[[w]], %[[b]]) <{config = "", config_proto = "", executor_type = "", f = @composite_conv2d_with_bias_and_relu6_fn_2}> +// PerLayer-DAG: %[[output1_quantized:.*]] = "tf.PartitionedCall"(%arg0, %[[w]], %[[b]]) <{config = "", config_proto = "", executor_type = "", f = @composite_conv2d_with_bias_and_relu6_fn_1}> {_tfl_quant_trait = "fully_quantizable"} +// PerLayer-DAG: %[[output1_unquantized:.*]] = "tf.PartitionedCall"(%arg0, %cst, %cst_0) <{config = "", config_proto = "", executor_type = "", f = @composite_conv2d_with_bias_and_relu6_fn_1_0}> +// PerLayer-DAG: "tf.DumpTensor"(%[[output1_quantized]]) <{enabled = false, file_name = "quantized_tensor_data.pb", func_name = "conv", log_dir_path = "/tmp/dumps/composite_conv2d_with_bias_and_relu6_fn_1", node_name = "Conv2D_1"}> : (tensor<*xf32>) -> () +// PerLayer-DAG: "tf.DumpTensor"(%[[output1_unquantized]]) <{enabled = false, file_name = "unquantized_tensor_data.pb", func_name = "conv", log_dir_path = "/tmp/dumps/composite_conv2d_with_bias_and_relu6_fn_1", node_name = "Conv2D_1"}> : (tensor<*xf32>) -> () // PerLayer-DAG: return %[[output0]], %[[output1_quantized]] } @@ -69,28 +69,28 @@ module { } // WholeModel-LABEL: func @multiple_conv2d -// WholeModel-DAG: %[[b0:.*]] = "tf.Const"() {value = dense<0.000000e+00> -// WholeModel-DAG: %[[b1:.*]] = "tf.Const"() {value = dense<1.000000e+00> -// WholeModel-DAG: %[[w0:.*]] = "tf.Const"() {value = dense<{{\[\[\[\[}}0.193340182, 0.285152316 -// WholeModel-DAG: %[[w1:.*]] = "tf.Const"() {value = dense<{{\[\[\[\[}}-0.174680978, -0.367524445 -// WholeModel-DAG: %[[output0:.*]] = "tf.PartitionedCall"(%arg0, %[[w0]], %[[b0]]) {_tfl_quant_trait = "fully_quantizable", config = "", config_proto = "", executor_type = "", f = @composite_conv2d_with_bias_and_relu6_fn_2} -// WholeModel-DAG: "tf.DumpTensor"(%[[output0]]) {enabled = false, file_name = "unquantized_tensor_data.pb", func_name = "multiple_conv2d", log_dir_path = "/tmp/dumps/composite_conv2d_with_bias_and_relu6_fn_2", node_name = "Conv2D"} -// WholeModel-DAG: %[[output1:.*]] = "tf.PartitionedCall"(%[[output0]], %[[w1]], %[[b1]]) {_tfl_quant_trait = "fully_quantizable", config = "", config_proto = "", executor_type = "", f = @composite_conv2d_with_bias_and_relu6_fn_1} -// WholeModel-DAG: "tf.DumpTensor"(%[[output1]]) {enabled = false, file_name = "unquantized_tensor_data.pb", func_name = "multiple_conv2d", log_dir_path = "/tmp/dumps/composite_conv2d_with_bias_and_relu6_fn_1", node_name = "Conv2D_1"} +// WholeModel-DAG: %[[b0:.*]] = "tf.Const"() <{value = dense<0.000000e+00> +// WholeModel-DAG: %[[b1:.*]] = "tf.Const"() <{value = dense<1.000000e+00> +// WholeModel-DAG: %[[w0:.*]] = "tf.Const"() <{value = dense<{{\[\[\[\[}}0.193340182, 0.285152316 +// WholeModel-DAG: %[[w1:.*]] = "tf.Const"() <{value = dense<{{\[\[\[\[}}-0.174680978, -0.367524445 +// WholeModel-DAG: %[[output0:.*]] = "tf.PartitionedCall"(%arg0, %[[w0]], %[[b0]]) <{config = "", config_proto = "", executor_type = "", f = @composite_conv2d_with_bias_and_relu6_fn_2}> {_tfl_quant_trait = "fully_quantizable"} +// WholeModel-DAG: "tf.DumpTensor"(%[[output0]]) <{enabled = false, file_name = "unquantized_tensor_data.pb", func_name = "multiple_conv2d", log_dir_path = "/tmp/dumps/composite_conv2d_with_bias_and_relu6_fn_2", node_name = "Conv2D"}> +// WholeModel-DAG: %[[output1:.*]] = "tf.PartitionedCall"(%[[output0]], %[[w1]], %[[b1]]) <{config = "", config_proto = "", executor_type = "", f = @composite_conv2d_with_bias_and_relu6_fn_1}> {_tfl_quant_trait = "fully_quantizable"} +// WholeModel-DAG: "tf.DumpTensor"(%[[output1]]) <{enabled = false, file_name = "unquantized_tensor_data.pb", func_name = "multiple_conv2d", log_dir_path = "/tmp/dumps/composite_conv2d_with_bias_and_relu6_fn_1", node_name = "Conv2D_1"}> // WholeModel-DAG: return %[[output1]] // PerLayer-LABEL: func @multiple_conv2d -// PerLayer-DAG: %[[b0:.*]] = "tf.Const"() {value = dense<0.000000e+00> -// PerLayer-DAG: %[[b1:.*]] = "tf.Const"() {value = dense<1.000000e+00> -// PerLayer-DAG: %[[w0:.*]] = "tf.Const"() {value = dense<{{\[\[\[\[}}0.193340182, 0.285152316 -// PerLayer-DAG: %[[w1:.*]] = "tf.Const"() {value = dense<{{\[\[\[\[}}-0.174680978, -0.367524445 -// PerLayer-DAG: %[[output0_quantized:.*]] = "tf.PartitionedCall"(%arg0, %[[w0]], %[[b0]]) {_tfl_quant_trait = "fully_quantizable", config = "", config_proto = "", executor_type = "", f = @composite_conv2d_with_bias_and_relu6_fn_2} -// PerLayer-DAG: %[[output0_unquantized:.*]] = "tf.PartitionedCall"(%arg0, %[[w0]], %[[b0]]) {config = "", config_proto = "", executor_type = "", f = @composite_conv2d_with_bias_and_relu6_fn_2_0} -// PerLayer-DAG: "tf.DumpTensor"(%[[output0_quantized]]) {enabled = false, file_name = "quantized_tensor_data.pb", func_name = "multiple_conv2d", log_dir_path = "/tmp/dumps/composite_conv2d_with_bias_and_relu6_fn_2", node_name = "Conv2D"} -// PerLayer-DAG: "tf.DumpTensor"(%[[output0_unquantized]]) {enabled = false, file_name = "unquantized_tensor_data.pb", func_name = "multiple_conv2d", log_dir_path = "/tmp/dumps/composite_conv2d_with_bias_and_relu6_fn_2", node_name = "Conv2D"} -// PerLayer-DAG: %[[output1_quantized:.*]] = "tf.PartitionedCall"(%[[output0_quantized]], %[[w1]], %[[b1]]) {_tfl_quant_trait = "fully_quantizable", config = "", config_proto = "", executor_type = "", f = @composite_conv2d_with_bias_and_relu6_fn_1} -// PerLayer-DAG: %[[output1_unquantized:.*]] = "tf.PartitionedCall"(%[[output0_quantized]], %[[w1]], %[[b1]]) {config = "", config_proto = "", executor_type = "", f = @composite_conv2d_with_bias_and_relu6_fn_1_0} -// PerLayer-DAG: "tf.DumpTensor"(%[[output1_quantized]]) {enabled = false, file_name = "quantized_tensor_data.pb", func_name = "multiple_conv2d", log_dir_path = "/tmp/dumps/composite_conv2d_with_bias_and_relu6_fn_1", node_name = "Conv2D_1"} -// PerLayer-DAG: "tf.DumpTensor"(%[[output1_unquantized]]) {enabled = false, file_name = "unquantized_tensor_data.pb", func_name = "multiple_conv2d", log_dir_path = "/tmp/dumps/composite_conv2d_with_bias_and_relu6_fn_1", node_name = "Conv2D_1"} +// PerLayer-DAG: %[[b0:.*]] = "tf.Const"() <{value = dense<0.000000e+00> +// PerLayer-DAG: %[[b1:.*]] = "tf.Const"() <{value = dense<1.000000e+00> +// PerLayer-DAG: %[[w0:.*]] = "tf.Const"() <{value = dense<{{\[\[\[\[}}0.193340182, 0.285152316 +// PerLayer-DAG: %[[w1:.*]] = "tf.Const"() <{value = dense<{{\[\[\[\[}}-0.174680978, -0.367524445 +// PerLayer-DAG: %[[output0_quantized:.*]] = "tf.PartitionedCall"(%arg0, %[[w0]], %[[b0]]) <{config = "", config_proto = "", executor_type = "", f = @composite_conv2d_with_bias_and_relu6_fn_2}> {_tfl_quant_trait = "fully_quantizable"} +// PerLayer-DAG: %[[output0_unquantized:.*]] = "tf.PartitionedCall"(%arg0, %[[w0]], %[[b0]]) <{config = "", config_proto = "", executor_type = "", f = @composite_conv2d_with_bias_and_relu6_fn_2_0}> +// PerLayer-DAG: "tf.DumpTensor"(%[[output0_quantized]]) <{enabled = false, file_name = "quantized_tensor_data.pb", func_name = "multiple_conv2d", log_dir_path = "/tmp/dumps/composite_conv2d_with_bias_and_relu6_fn_2", node_name = "Conv2D"}> +// PerLayer-DAG: "tf.DumpTensor"(%[[output0_unquantized]]) <{enabled = false, file_name = "unquantized_tensor_data.pb", func_name = "multiple_conv2d", log_dir_path = "/tmp/dumps/composite_conv2d_with_bias_and_relu6_fn_2", node_name = "Conv2D"}> +// PerLayer-DAG: %[[output1_quantized:.*]] = "tf.PartitionedCall"(%[[output0_quantized]], %[[w1]], %[[b1]]) <{config = "", config_proto = "", executor_type = "", f = @composite_conv2d_with_bias_and_relu6_fn_1}> {_tfl_quant_trait = "fully_quantizable"} +// PerLayer-DAG: %[[output1_unquantized:.*]] = "tf.PartitionedCall"(%[[output0_quantized]], %[[w1]], %[[b1]]) <{config = "", config_proto = "", executor_type = "", f = @composite_conv2d_with_bias_and_relu6_fn_1_0}> +// PerLayer-DAG: "tf.DumpTensor"(%[[output1_quantized]]) <{enabled = false, file_name = "quantized_tensor_data.pb", func_name = "multiple_conv2d", log_dir_path = "/tmp/dumps/composite_conv2d_with_bias_and_relu6_fn_1", node_name = "Conv2D_1"}> +// PerLayer-DAG: "tf.DumpTensor"(%[[output1_unquantized]]) <{enabled = false, file_name = "unquantized_tensor_data.pb", func_name = "multiple_conv2d", log_dir_path = "/tmp/dumps/composite_conv2d_with_bias_and_relu6_fn_1", node_name = "Conv2D_1"}> // PerLayer-DAG: return %[[output1_quantized]] } diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/tests/cast_bf16_ops_to_f32.mlir b/tensorflow/compiler/mlir/quantization/tensorflow/tests/cast_bf16_ops_to_f32.mlir index 4fc6cbf3f97b7d..deaafb3e2e99a9 100644 --- a/tensorflow/compiler/mlir/quantization/tensorflow/tests/cast_bf16_ops_to_f32.mlir +++ b/tensorflow/compiler/mlir/quantization/tensorflow/tests/cast_bf16_ops_to_f32.mlir @@ -10,8 +10,8 @@ func.func @cast_bf16_conv_to_fp32(%arg0: tensor<1x3x4x3xf32>) -> (tensor<1x3x2x2 } // CHECK: func @cast_bf16_conv_to_fp32 -// CHECK-DAG: %[[cst:.*]] = "tf.Const"() {device = "", value = dense_resource<__elided__> : tensor<2x3x3x2xbf16>} : () -> tensor<2x3x3x2xbf16> -// CHECK: %[[cast:.*]] = "tf.Cast"(%[[cst]]) {Truncate = false} : (tensor<2x3x3x2xbf16>) -> tensor<2x3x3x2xf32> +// CHECK-DAG: %[[cst:.*]] = "tf.Const"() <{value = dense_resource<__elided__> : tensor<2x3x3x2xbf16>}> {device = ""} : () -> tensor<2x3x3x2xbf16> +// CHECK: %[[cast:.*]] = "tf.Cast"(%[[cst]]) <{Truncate = false}> : (tensor<2x3x3x2xbf16>) -> tensor<2x3x3x2xf32> // CHECK: %[[conv:.*]] = "tf.Conv2D"(%arg0, %[[cast]]) // CHECK: %[[identity:.*]] = "tf.IdentityN"(%[[conv]]) {device = ""} : (tensor<1x3x2x2xf32>) -> tensor<1x3x2x2xf32> // CHECK: return %[[identity]] : tensor<1x3x2x2xf32> @@ -28,8 +28,8 @@ func.func @cast_bf16_conv_with_bias_to_fp32(%arg0: tensor<1x3x4x3xf32>) -> (tens } // CHECK: func @cast_bf16_conv_with_bias_to_fp32 -// CHECK-DAG: %[[cst:.*]] = "tf.Const"() {value = dense<1.000000e+00> : tensor<2x3x3x2xf32>} : () -> tensor<2x3x3x2xf32> -// CHECK-DAG: %[[cst_0:.*]] = "tf.Const"() {value = dense<1.000000e+00> : tensor<2xf32>} : () -> tensor<2xf32> +// CHECK-DAG: %[[cst:.*]] = "tf.Const"() <{value = dense<1.000000e+00> : tensor<2x3x3x2xf32>}> : () -> tensor<2x3x3x2xf32> +// CHECK-DAG: %[[cst_0:.*]] = "tf.Const"() <{value = dense<1.000000e+00> : tensor<2xf32>}> : () -> tensor<2xf32> // CHECK: %[[conv:.*]] = "tf.Conv2D"(%arg0, %[[cst]]) // CHECK: %[[bias_add:.*]] = "tf.BiasAdd"(%[[conv]], %[[cst_0]]) // CHECK: %[[identity:.*]] = "tf.IdentityN"(%[[bias_add]]) {device = ""} : (tensor<1x3x2x2xf32>) -> tensor<1x3x2x2xf32> @@ -46,7 +46,7 @@ func.func @cast_bf16_avg_pool_to_fp32(%arg0: tensor<1x3x4x3xf32>) -> (tensor<1x3 } // CHECK: func @cast_bf16_avg_pool_to_fp32 -// CHECK-DAG: %[[cst:.*]] = "tf.Const"() {value = dense<{{.*}}> : tensor<2x3x3x2xf32>} : () -> tensor<2x3x3x2xf32> +// CHECK-DAG: %[[cst:.*]] = "tf.Const"() <{value = dense<{{.*}}> : tensor<2x3x3x2xf32>}> : () -> tensor<2x3x3x2xf32> // CHECK: %[[conv:.*]] = "tf.Conv2D"(%arg0, %[[cst]]) // CHECK: %[[avg_pool:.*]] = "tf.AvgPool"(%[[conv]]) // CHECK: %[[identity:.*]] = "tf.IdentityN"(%[[avg_pool]]) {device = ""} : (tensor<1x3x2x2xf32>) -> tensor<1x3x2x2xf32> @@ -62,7 +62,7 @@ func.func @cast_bf16_matmul_to_fp32(%arg0: tensor<1x10xf32>) -> (tensor<1x2xf32> } // CHECK: func @cast_bf16_matmul_to_fp32 -// CHECK-DAG: %[[cst:.*]] = "tf.Const"() {value = dense<{{.*}}> : tensor<10x2xf32>} : () -> tensor<10x2xf32> +// CHECK-DAG: %[[cst:.*]] = "tf.Const"() <{value = dense<{{.*}}> : tensor<10x2xf32>}> : () -> tensor<10x2xf32> // CHECK: %[[matmul:.*]] = "tf.MatMul"(%arg0, %[[cst]]) // CHECK: %[[identity:.*]] = "tf.IdentityN"(%[[matmul]]) // CHECK: return %[[identity]] : tensor<1x2xf32> @@ -77,7 +77,7 @@ func.func @cast_bf16_depthwise_conv_to_fp32(%arg0: tensor<1x3x4x3xf32>) -> (tens } // CHECK: func @cast_bf16_depthwise_conv_to_fp32 -// CHECK-DAG: %[[cst:.*]] = "tf.Const"() {value = dense<{{.*}}> : tensor<2x3x3x2xf32>} : () -> tensor<2x3x3x2xf32> +// CHECK-DAG: %[[cst:.*]] = "tf.Const"() <{value = dense<{{.*}}> : tensor<2x3x3x2xf32>}> : () -> tensor<2x3x3x2xf32> // CHECK: %[[depthwise_conv:.*]] = "tf.DepthwiseConv2dNative"(%arg0, %[[cst]]) // CHECK: %[[identity:.*]] = "tf.IdentityN"(%[[depthwise_conv]]) {device = ""} : (tensor<1x2x2x6xf32>) -> tensor<1x2x2x6xf32> // CHECK: return %[[identity]] : tensor<1x2x2x6xf32> @@ -92,7 +92,7 @@ func.func @cast_bf16_batch_matmul_v2_to_fp32(%arg0: tensor<1x1x10xf32>) -> (tens } // CHECK: func @cast_bf16_batch_matmul_v2_to_fp32 -// CHECK-DAG: %[[cst:.*]] = "tf.Const"() {value = dense<{{.*}}> : tensor<10x2xf32>} : () -> tensor<10x2xf32> +// CHECK-DAG: %[[cst:.*]] = "tf.Const"() <{value = dense<{{.*}}> : tensor<10x2xf32>}> : () -> tensor<10x2xf32> // CHECK: %[[batch_matmul:.*]] = "tf.BatchMatMulV2"(%arg0, %[[cst]]) // CHECK: %[[identity:.*]] = "tf.IdentityN"(%[[batch_matmul]]) {device = ""} : (tensor<1x1x2xf32>) -> tensor<1x1x2xf32> // CHECK: return %[[identity]] : tensor<1x1x2xf32> @@ -108,7 +108,7 @@ func.func @cast_bf16_add_v2_to_fp32(%arg0: tensor<2xbf16>, %arg1: tensor<2xbf16> // CHECK: func @cast_bf16_add_v2_to_fp32(%[[ARG_0:.*]]: tensor<2xbf16>, %[[ARG_1:.*]]: tensor<2xbf16>) -> tensor<2xf32> // bfloat16 operands are cast to f32 operands. -// CHECK-DAG: %[[CAST_0:.*]] = "tf.Cast"(%[[ARG_0]]) {Truncate = false} : (tensor<2xbf16>) -> tensor<2xf32> -// CHECK-DAG: %[[CAST_1:.*]] = "tf.Cast"(%[[ARG_1]]) {Truncate = false} : (tensor<2xbf16>) -> tensor<2xf32> +// CHECK-DAG: %[[CAST_0:.*]] = "tf.Cast"(%[[ARG_0]]) <{Truncate = false}> : (tensor<2xbf16>) -> tensor<2xf32> +// CHECK-DAG: %[[CAST_1:.*]] = "tf.Cast"(%[[ARG_1]]) <{Truncate = false}> : (tensor<2xbf16>) -> tensor<2xf32> // CHECK: %[[ADD:.*]] = "tf.AddV2"(%[[CAST_0]], %[[CAST_1]]) : (tensor<2xf32>, tensor<2xf32>) -> tensor<2xf32> // CHECK: return %[[ADD]] : tensor<2xf32> diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/tests/convert_tf_xla_op_to_tf_op.mlir b/tensorflow/compiler/mlir/quantization/tensorflow/tests/convert_tf_xla_op_to_tf_op.mlir index d30c61f7df72dd..27a7bb62bbbd74 100644 --- a/tensorflow/compiler/mlir/quantization/tensorflow/tests/convert_tf_xla_op_to_tf_op.mlir +++ b/tensorflow/compiler/mlir/quantization/tensorflow/tests/convert_tf_xla_op_to_tf_op.mlir @@ -6,7 +6,7 @@ func.func @xla_dot_v2(%arg0: tensor, %arg1: tensor<3x4x5xf32>) -> (te } // CHECK: func @xla_dot_v2 -// CHECK: %[[einsum:.*]] = "tf.Einsum"(%arg0, %arg1) {equation = "abc,cde->abde"} : (tensor, tensor<3x4x5xf32>) -> tensor +// CHECK: %[[einsum:.*]] = "tf.Einsum"(%arg0, %arg1) <{equation = "abc,cde->abde"}> : (tensor, tensor<3x4x5xf32>) -> tensor // CHECK: return %[[einsum]] : tensor // ----- @@ -22,12 +22,12 @@ func.func @xla_gather(%arg0: tensor, %arg1: tensor<1xi32>, %arg2: tenso } // CHECK: func @xla_gather -// CHECK-DAG: %[[cst:.*]] = "tf.Const"() {value = dense<0> : tensor<2xi64>} : () -> tensor<2xi64> -// CHECK-DAG: %[[cst_0:.*]] = "tf.Const"() {value = dense<1> : tensor<1x1xi64>} : () -> tensor<1x1xi64> -// CHECK-DAG: %[[cst_1:.*]] = "tf.Const"() {value = dense<-1> : tensor<1xi64>} : () -> tensor<1xi64> -// CHECK: %[[arg1_i64:.*]] = "tf.Cast"(%arg1) {Truncate = false} : (tensor<1xi32>) -> tensor<1xi64> +// CHECK-DAG: %[[cst:.*]] = "tf.Const"() <{value = dense<0> : tensor<2xi64>}> : () -> tensor<2xi64> +// CHECK-DAG: %[[cst_0:.*]] = "tf.Const"() <{value = dense<1> : tensor<1x1xi64>}> : () -> tensor<1x1xi64> +// CHECK-DAG: %[[cst_1:.*]] = "tf.Const"() <{value = dense<-1> : tensor<1xi64>}> : () -> tensor<1xi64> +// CHECK: %[[arg1_i64:.*]] = "tf.Cast"(%arg1) <{Truncate = false}> : (tensor<1xi32>) -> tensor<1xi64> // CHECK: %[[tensor_scatter_update:.*]] = "tf.TensorScatterUpdate"(%[[cst]], %[[cst_0]], %[[arg1_i64]]) : (tensor<2xi64>, tensor<1x1xi64>, tensor<1xi64>) -> tensor<2xi64> -// CHECK: %[[arg2_i64:.*]] = "tf.Cast"(%arg2) {Truncate = false} : (tensor<2xi32>) -> tensor<2xi64> +// CHECK: %[[arg2_i64:.*]] = "tf.Cast"(%arg2) <{Truncate = false}> : (tensor<2xi32>) -> tensor<2xi64> // CHECK: %[[slice:.*]] = "tf.Slice"(%arg0, %[[tensor_scatter_update]], %[[arg2_i64]]) : (tensor, tensor<2xi64>, tensor<2xi64>) -> tensor<*xf32> // CHECK: %[[reshape:.*]] = "tf.Reshape"(%[[slice]], %[[cst_1]]) : (tensor<*xf32>, tensor<1xi64>) -> tensor<*xf32> // CHECK: return %[[reshape]] : tensor<*xf32> @@ -47,12 +47,12 @@ func.func @xla_gather_known_output_shape(%arg0: tensor<5xi32>, %arg1: tensor<1xi } // CHECK: func @xla_gather_known_output_shape -// CHECK-DAG: %[[cst:.*]] = "tf.Const"() {value = dense<0> : tensor<1xi64>} : () -> tensor<1xi64> -// CHECK-DAG: %[[cst_0:.*]] = "tf.Const"() {value = dense<0> : tensor<1x1xi64>} : () -> tensor<1x1xi64> -// CHECK-DAG: %[[cst_1:.*]] = "tf.Const"() {value = dense<> : tensor<0xi64>} : () -> tensor<0xi64> -// CHECK: %[[arg1_i64:.*]] = "tf.Cast"(%arg1) {Truncate = false} : (tensor<1xi32>) -> tensor<1xi64> +// CHECK-DAG: %[[cst:.*]] = "tf.Const"() <{value = dense<0> : tensor<1xi64>}> : () -> tensor<1xi64> +// CHECK-DAG: %[[cst_0:.*]] = "tf.Const"() <{value = dense<0> : tensor<1x1xi64>}> : () -> tensor<1x1xi64> +// CHECK-DAG: %[[cst_1:.*]] = "tf.Const"() <{value = dense<> : tensor<0xi64>}> : () -> tensor<0xi64> +// CHECK: %[[arg1_i64:.*]] = "tf.Cast"(%arg1) <{Truncate = false}> : (tensor<1xi32>) -> tensor<1xi64> // CHECK: %[[tensor_scatter_update:.*]] = "tf.TensorScatterUpdate"(%[[cst]], %[[cst_0]], %[[arg1_i64]]) : (tensor<1xi64>, tensor<1x1xi64>, tensor<1xi64>) -> tensor<1xi64> -// CHECK: %[[arg2_i64:.*]] = "tf.Cast"(%arg2) {Truncate = false} : (tensor<1xi32>) -> tensor<1xi64> +// CHECK: %[[arg2_i64:.*]] = "tf.Cast"(%arg2) <{Truncate = false}> : (tensor<1xi32>) -> tensor<1xi64> // CHECK: %[[slice:.*]] = "tf.Slice"(%arg0, %[[tensor_scatter_update]], %[[arg2_i64]]) : (tensor<5xi32>, tensor<1xi64>, tensor<1xi64>) -> tensor<1xi32> // CHECK: %[[reshape:.*]] = "tf.Reshape"(%[[slice]], %[[cst_1]]) : (tensor<1xi32>, tensor<0xi64>) -> tensor // CHECK: return %[[reshape]] : tensor diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/tests/convert_tpu_model_to_cpu.mlir b/tensorflow/compiler/mlir/quantization/tensorflow/tests/convert_tpu_model_to_cpu.mlir index 26809cd7e8d94f..ad13ebbfd59814 100644 --- a/tensorflow/compiler/mlir/quantization/tensorflow/tests/convert_tpu_model_to_cpu.mlir +++ b/tensorflow/compiler/mlir/quantization/tensorflow/tests/convert_tpu_model_to_cpu.mlir @@ -26,8 +26,8 @@ func.func private @tpu_func_0_optim0(%arg0: tensor<1x3x4x3xf32>) -> tensor<1x3x2 } // CHECK: func @tpu_conv(%[[ARG0:.*]]: tensor<1x3x4x3xf32>) -// CHECK-DAG: %[[cst:.*]] = "tf.Const"() {device = "", value = dense_resource<__elided__> : tensor<2x3x3x2xbf16>} : () -> tensor<2x3x3x2xbf16> -// CHECK: %[[cast:.*]] = "tf.Cast"(%[[cst]]) {Truncate = false} : (tensor<2x3x3x2xbf16>) -> tensor<2x3x3x2xf32> +// CHECK-DAG: %[[cst:.*]] = "tf.Const"() <{value = dense_resource<__elided__> : tensor<2x3x3x2xbf16>}> {device = ""} : () -> tensor<2x3x3x2xbf16> +// CHECK: %[[cast:.*]] = "tf.Cast"(%[[cst]]) <{Truncate = false}> : (tensor<2x3x3x2xbf16>) -> tensor<2x3x3x2xf32> // CHECK: %[[conv:.*]] = "tf.Conv2D"(%[[ARG0]], %[[cast]]) // CHECK: %[[identity:.*]] = "tf.IdentityN"(%[[conv]]) {device = ""} : (tensor<1x3x2x2xf32>) -> tensor<1x3x2x2xf32> // CHECK: return %[[identity]] : tensor<1x3x2x2xf32> diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/tests/fake_quant_e2e_flow.mlir b/tensorflow/compiler/mlir/quantization/tensorflow/tests/fake_quant_e2e_flow.mlir index b08c49f94f11ba..aaddb72fccd63d 100644 --- a/tensorflow/compiler/mlir/quantization/tensorflow/tests/fake_quant_e2e_flow.mlir +++ b/tensorflow/compiler/mlir/quantization/tensorflow/tests/fake_quant_e2e_flow.mlir @@ -13,18 +13,18 @@ func.func @fake_quant_conv(%arg0: tensor<1x3x4x3xf32>, %arg1: tensor<2x3x3x2xf32 // CHECK-LABEL: @fake_quant_conv // CHECK-SAME: %[[ARG0:.*]]: tensor<1x3 // CHECK-SAME: %[[ARG1:.*]]: tensor<2x3 -// CHECK-DAG: %[[CST:.*]] = "tf.Const"() {value = dense<0.00117647066> : tensor} : () -> tensor -// CHECK-DAG: %[[CST_0:.*]] = "tf.Const"() {value = dense<-43> : tensor} : () -> tensor -// CHECK-DAG: %[[CST_1:.*]] = "tf.Const"() {value = dense<0.0117647061> : tensor} : () -> tensor -// CHECK-DAG: %[[CST_2:.*]] = "tf.Const"() {value = dense<1.38408304E-5> : tensor} : () -> tensor -// CHECK-DAG: %[[CST_3:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor -// CHECK-DAG: %[[CST_4:.*]] = "tf.Const"() {value = dense<0.0027450982> : tensor} : () -> tensor -// CHECK-DAG: %[[CST_5:.*]] = "tf.Const"() {value = dense<-19> : tensor} : () -> tensor -// CHECK-DAG: %[[CST_6:.*]] = "tf.Const"() {value = dense<0> : tensor<2xi32>} : () -> tensor<2xi32> -// CHECK-NEXT: %[[V0:.*]] = "tf.PartitionedCall"(%[[ARG1]], %[[CST_1]], %[[CST_0]]) {config = "", config_proto = "", executor_type = "", f = @quantize_i8} : (tensor<2x3x3x2xf32>, tensor, tensor) -> tensor<2x3x3x2xi8> -// CHECK-NEXT: %[[V1:.*]] = "tf.PartitionedCall"(%[[ARG0]], %[[CST]], %[[CST_0]]) {config = "", config_proto = "", executor_type = "", f = @quantize_i8} : (tensor<1x3x4x3xf32>, tensor, tensor) -> tensor<1x3x4x3xi8> -// CHECK-NEXT: %[[V2:.*]] = "tf.PartitionedCall"(%[[V1]], %[[V0]], %[[CST_6]], %[[CST]], %[[CST_0]], %[[CST_1]], %[[CST_0]], %[[CST_2]], %[[CST_3]], %[[CST_4]], %[[CST_5]]) {config = "", config_proto = "", executor_type = "", f = @quantized_conv2d_with_bias_fn_0} : (tensor<1x3x4x3xi8>, tensor<2x3x3x2xi8>, tensor<2xi32>, tensor, tensor, tensor, tensor, tensor, tensor, tensor, tensor) -> tensor<*xi8> -// CHECK-NEXT: %[[V3:.*]] = "tf.PartitionedCall"(%[[V2]], %[[CST_4]], %[[CST_5]]) {config = "", config_proto = "", executor_type = "", f = @dequantize_i8} : (tensor<*xi8>, tensor, tensor) -> tensor<*xf32> +// CHECK-DAG: %[[CST:.*]] = "tf.Const"() <{value = dense<0.00117647066> : tensor}> : () -> tensor +// CHECK-DAG: %[[CST_0:.*]] = "tf.Const"() <{value = dense<-43> : tensor}> : () -> tensor +// CHECK-DAG: %[[CST_1:.*]] = "tf.Const"() <{value = dense<0.0117647061> : tensor}> : () -> tensor +// CHECK-DAG: %[[CST_2:.*]] = "tf.Const"() <{value = dense<1.38408304E-5> : tensor}> : () -> tensor +// CHECK-DAG: %[[CST_3:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor +// CHECK-DAG: %[[CST_4:.*]] = "tf.Const"() <{value = dense<0.0027450982> : tensor}> : () -> tensor +// CHECK-DAG: %[[CST_5:.*]] = "tf.Const"() <{value = dense<-19> : tensor}> : () -> tensor +// CHECK-DAG: %[[CST_6:.*]] = "tf.Const"() <{value = dense<0> : tensor<2xi32>}> : () -> tensor<2xi32> +// CHECK-NEXT: %[[V0:.*]] = "tf.PartitionedCall"(%[[ARG1]], %[[CST_1]], %[[CST_0]]) <{config = "", config_proto = "", executor_type = "", f = @quantize_i8}> : (tensor<2x3x3x2xf32>, tensor, tensor) -> tensor<2x3x3x2xi8> +// CHECK-NEXT: %[[V1:.*]] = "tf.PartitionedCall"(%[[ARG0]], %[[CST]], %[[CST_0]]) <{config = "", config_proto = "", executor_type = "", f = @quantize_i8}> : (tensor<1x3x4x3xf32>, tensor, tensor) -> tensor<1x3x4x3xi8> +// CHECK-NEXT: %[[V2:.*]] = "tf.PartitionedCall"(%[[V1]], %[[V0]], %[[CST_6]], %[[CST]], %[[CST_0]], %[[CST_1]], %[[CST_0]], %[[CST_2]], %[[CST_3]], %[[CST_4]], %[[CST_5]]) <{config = "", config_proto = "", executor_type = "", f = @quantized_conv2d_with_bias_fn_0}> : (tensor<1x3x4x3xi8>, tensor<2x3x3x2xi8>, tensor<2xi32>, tensor, tensor, tensor, tensor, tensor, tensor, tensor, tensor) -> tensor<*xi8> +// CHECK-NEXT: %[[V3:.*]] = "tf.PartitionedCall"(%[[V2]], %[[CST_4]], %[[CST_5]]) <{config = "", config_proto = "", executor_type = "", f = @dequantize_i8}> : (tensor<*xi8>, tensor, tensor) -> tensor<*xf32> // CHECK-NEXT: return %[[V3]] : tensor<*xf32> // CHECK: func private @quantize_i8( diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/tests/fake_quant_e2e_xla.mlir b/tensorflow/compiler/mlir/quantization/tensorflow/tests/fake_quant_e2e_xla.mlir index 5d9801e208593c..e5c5d8a4211d61 100644 --- a/tensorflow/compiler/mlir/quantization/tensorflow/tests/fake_quant_e2e_xla.mlir +++ b/tensorflow/compiler/mlir/quantization/tensorflow/tests/fake_quant_e2e_xla.mlir @@ -25,7 +25,7 @@ module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, p // CHECK: %[[pad:.*]] = "tf.PadV2"(%[[quant]] // CHECK: %[[xlaconv:.*]] = "tf.XlaConvV2"(%[[pad]] // CHECK: %[[sub:.*]] = "tf.Sub"(%[[xlaconv]] -// CHECK: %[[cast:.*]] = "tf.Cast"(%[[sub]]) {Truncate = false} : (tensor<1x3x2x2xi32>) -> tensor<1x3x2x2xf32> +// CHECK: %[[cast:.*]] = "tf.Cast"(%[[sub]]) <{Truncate = false}> : (tensor<1x3x2x2xi32>) -> tensor<1x3x2x2xf32> // CHECK: %[[dequant1:.*]] = "tf.Mul"(%[[cast]] // CHECK: %[[relu:.*]] = "tf.Relu"(%[[dequant1]] // CHECK: %[[clamped:.*]] = "tf.Minimum"(%[[relu]] @@ -35,12 +35,12 @@ module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, p // CHECK: %[[maximum2:.*]] = "tf.Maximum"(%[[add2]] // CHECK: %[[minimum2:.*]] = "tf.Minimum"(%[[maximum2]] // CHECK: %[[round2:.*]] = "tf.Round"(%[[minimum2]] -// CHECK: %[[quant2:.*]] = "tf.Cast"(%[[round2]]) {Truncate = false} : (tensor<1x3x2x2xf32>) -> tensor<1x3x2x2xi8> +// CHECK: %[[quant2:.*]] = "tf.Cast"(%[[round2]]) <{Truncate = false}> : (tensor<1x3x2x2xf32>) -> tensor<1x3x2x2xi8> // CHECK: %[[pad2:.*]] = "tf.PadV2"(%[[quant2]] // CHECK: %[[xlaconv2:.*]] = "tf.XlaConvV2"(%[[pad2]] // CHECK: %[[sub2:.*]] = "tf.Sub"(%[[xlaconv2]] -// CHECK: %[[cast2:.*]] = "tf.Cast"(%[[sub2]]) {Truncate = false} : (tensor<1x3x2x2xi32>) -> tensor<1x3x2x2xf32> +// CHECK: %[[cast2:.*]] = "tf.Cast"(%[[sub2]]) <{Truncate = false}> : (tensor<1x3x2x2xi32>) -> tensor<1x3x2x2xf32> // CHECK: %[[rescale2:.*]] = "tf.Mul"(%[[cast2]] // CHECK: %[[rescale2_maxclamped:.*]] = "tf.Maximum"(%[[rescale2]] // CHECK: %[[rescale2_minclamped:.*]] = "tf.Minimum"(%[[rescale2_maxclamped]] diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/tests/insert_custom_aggregation_ops.mlir b/tensorflow/compiler/mlir/quantization/tensorflow/tests/insert_custom_aggregation_ops.mlir index dd1c5a41bdf4d1..fa747357169f55 100644 --- a/tensorflow/compiler/mlir/quantization/tensorflow/tests/insert_custom_aggregation_ops.mlir +++ b/tensorflow/compiler/mlir/quantization/tensorflow/tests/insert_custom_aggregation_ops.mlir @@ -26,10 +26,10 @@ module { // CalibrationOptions(calibration_method=CALIBRATION_METHOD_MIN_MAX) // MIN-MAX-CHECK: func @add_custom_ops -// MIN-MAX-CHECK-NEXT: [[rhs:%.*]] = "tf.CustomAggregator"(%arg1) {calibration_method = 1 : i32, id = "", initial_num_bins = 0 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32} : (tensor<*xf32>) -> tensor<*xf32> -// MIN-MAX-CHECK-NEXT: [[lhs:%.*]] = "tf.CustomAggregator"(%arg0) {calibration_method = 1 : i32, id = "", initial_num_bins = 0 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32} : (tensor<*xf32>) -> tensor<*xf32> +// MIN-MAX-CHECK-NEXT: [[rhs:%.*]] = "tf.CustomAggregator"(%arg1) <{id = ""}> {calibration_method = 1 : i32, initial_num_bins = 0 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32} : (tensor<*xf32>) -> tensor<*xf32> +// MIN-MAX-CHECK-NEXT: [[lhs:%.*]] = "tf.CustomAggregator"(%arg0) <{id = ""}> {calibration_method = 1 : i32, initial_num_bins = 0 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32} : (tensor<*xf32>) -> tensor<*xf32> // MIN-MAX-CHECK-NEXT: [[add:%.*]] = "tf.AddV2"([[lhs]], [[rhs]]) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32> -// MIN-MAX-CHECK-NEXT: [[res:%.*]] = "tf.CustomAggregator"([[add]]) {calibration_method = 1 : i32, id = "", initial_num_bins = 0 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32} : (tensor<*xf32>) -> tensor<*xf32> +// MIN-MAX-CHECK-NEXT: [[res:%.*]] = "tf.CustomAggregator"([[add]]) <{id = ""}> {calibration_method = 1 : i32, initial_num_bins = 0 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32} : (tensor<*xf32>) -> tensor<*xf32> // MIN-MAX-CHECK-NEXT: return [[res]] : tensor<*xf32> // MIN-MAX-CHECK: func @no_custom_ops_on_non_f32_type @@ -44,10 +44,10 @@ module { // CalibrationOptions(calibration_method=CALIBRATION_METHOD_AVERAGE_MIN_MAX) // AVERAGE-MIN-MAX-CHECK: func @add_custom_ops -// AVERAGE-MIN-MAX-CHECK-NEXT: [[rhs:%.*]] = "tf.CustomAggregator"(%arg1) {calibration_method = 2 : i32, id = "", initial_num_bins = 0 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32} : (tensor<*xf32>) -> tensor<*xf32> -// AVERAGE-MIN-MAX-CHECK-NEXT: [[lhs:%.*]] = "tf.CustomAggregator"(%arg0) {calibration_method = 2 : i32, id = "", initial_num_bins = 0 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32} : (tensor<*xf32>) -> tensor<*xf32> +// AVERAGE-MIN-MAX-CHECK-NEXT: [[rhs:%.*]] = "tf.CustomAggregator"(%arg1) <{id = ""}> {calibration_method = 2 : i32, initial_num_bins = 0 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32} : (tensor<*xf32>) -> tensor<*xf32> +// AVERAGE-MIN-MAX-CHECK-NEXT: [[lhs:%.*]] = "tf.CustomAggregator"(%arg0) <{id = ""}> {calibration_method = 2 : i32, initial_num_bins = 0 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32} : (tensor<*xf32>) -> tensor<*xf32> // AVERAGE-MIN-MAX-CHECK-NEXT: [[add:%.*]] = "tf.AddV2"([[lhs]], [[rhs]]) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32> -// AVERAGE-MIN-MAX-CHECK-NEXT: [[res:%.*]] = "tf.CustomAggregator"([[add]]) {calibration_method = 2 : i32, id = "", initial_num_bins = 0 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32} : (tensor<*xf32>) -> tensor<*xf32> +// AVERAGE-MIN-MAX-CHECK-NEXT: [[res:%.*]] = "tf.CustomAggregator"([[add]]) <{id = ""}> {calibration_method = 2 : i32, initial_num_bins = 0 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32} : (tensor<*xf32>) -> tensor<*xf32> // AVERAGE-MIN-MAX-CHECK-NEXT: return [[res]] : tensor<*xf32> // AVERAGE-MIN-MAX-CHECK: func @no_custom_ops_on_non_f32_type @@ -65,10 +65,10 @@ module { // calibration_parameters=CalibrationParameters(initial_num_bins=256, min_percentile=0.001, max_percentile=99.999) // ) // HISTOGRAM-PERCENTILE-CHECK: func @add_custom_ops -// HISTOGRAM-PERCENTILE-CHECK-NEXT: [[rhs:%.*]] = "tf.CustomAggregator"(%arg1) {calibration_method = 3 : i32, id = "", initial_num_bins = 256 : i32, max_percentile = 9.999900e+01 : f32, min_percentile = 1.000000e-03 : f32} : (tensor<*xf32>) -> tensor<*xf32> -// HISTOGRAM-PERCENTILE-CHECK-NEXT: [[lhs:%.*]] = "tf.CustomAggregator"(%arg0) {calibration_method = 3 : i32, id = "", initial_num_bins = 256 : i32, max_percentile = 9.999900e+01 : f32, min_percentile = 1.000000e-03 : f32} : (tensor<*xf32>) -> tensor<*xf32> +// HISTOGRAM-PERCENTILE-CHECK-NEXT: [[rhs:%.*]] = "tf.CustomAggregator"(%arg1) <{id = ""}> {calibration_method = 3 : i32, initial_num_bins = 256 : i32, max_percentile = 9.999900e+01 : f32, min_percentile = 1.000000e-03 : f32} : (tensor<*xf32>) -> tensor<*xf32> +// HISTOGRAM-PERCENTILE-CHECK-NEXT: [[lhs:%.*]] = "tf.CustomAggregator"(%arg0) <{id = ""}> {calibration_method = 3 : i32, initial_num_bins = 256 : i32, max_percentile = 9.999900e+01 : f32, min_percentile = 1.000000e-03 : f32} : (tensor<*xf32>) -> tensor<*xf32> // HISTOGRAM-PERCENTILE-CHECK-NEXT: [[add:%.*]] = "tf.AddV2"([[lhs]], [[rhs]]) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32> -// HISTOGRAM-PERCENTILE-CHECK-NEXT: [[res:%.*]] = "tf.CustomAggregator"([[add]]) {calibration_method = 3 : i32, id = "", initial_num_bins = 256 : i32, max_percentile = 9.999900e+01 : f32, min_percentile = 1.000000e-03 : f32} : (tensor<*xf32>) -> tensor<*xf32> +// HISTOGRAM-PERCENTILE-CHECK-NEXT: [[res:%.*]] = "tf.CustomAggregator"([[add]]) <{id = ""}> {calibration_method = 3 : i32, initial_num_bins = 256 : i32, max_percentile = 9.999900e+01 : f32, min_percentile = 1.000000e-03 : f32} : (tensor<*xf32>) -> tensor<*xf32> // HISTOGRAM-PERCENTILE-CHECK-NEXT: return [[res]] : tensor<*xf32> // HISTOGRAM-PERCENTILE-CHECK: func @no_custom_ops_on_non_f32_type @@ -86,10 +86,10 @@ module { // calibration_parameters=CalibrationParameters(initial_num_bins=256) // ) // HISTOGRAM-MSE-BRUTEFORCE-CHECK: func @add_custom_ops -// HISTOGRAM-MSE-BRUTEFORCE-CHECK-NEXT: [[rhs:%.*]] = "tf.CustomAggregator"(%arg1) {calibration_method = 4 : i32, id = "", initial_num_bins = 256 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32} : (tensor<*xf32>) -> tensor<*xf32> -// HISTOGRAM-MSE-BRUTEFORCE-CHECK-NEXT: [[lhs:%.*]] = "tf.CustomAggregator"(%arg0) {calibration_method = 4 : i32, id = "", initial_num_bins = 256 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32} : (tensor<*xf32>) -> tensor<*xf32> +// HISTOGRAM-MSE-BRUTEFORCE-CHECK-NEXT: [[rhs:%.*]] = "tf.CustomAggregator"(%arg1) <{id = ""}> {calibration_method = 4 : i32, initial_num_bins = 256 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32} : (tensor<*xf32>) -> tensor<*xf32> +// HISTOGRAM-MSE-BRUTEFORCE-CHECK-NEXT: [[lhs:%.*]] = "tf.CustomAggregator"(%arg0) <{id = ""}> {calibration_method = 4 : i32, initial_num_bins = 256 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32} : (tensor<*xf32>) -> tensor<*xf32> // HISTOGRAM-MSE-BRUTEFORCE-CHECK-NEXT: [[add:%.*]] = "tf.AddV2"([[lhs]], [[rhs]]) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32> -// HISTOGRAM-MSE-BRUTEFORCE-CHECK-NEXT: [[res:%.*]] = "tf.CustomAggregator"([[add]]) {calibration_method = 4 : i32, id = "", initial_num_bins = 256 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32} : (tensor<*xf32>) -> tensor<*xf32> +// HISTOGRAM-MSE-BRUTEFORCE-CHECK-NEXT: [[res:%.*]] = "tf.CustomAggregator"([[add]]) <{id = ""}> {calibration_method = 4 : i32, initial_num_bins = 256 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32} : (tensor<*xf32>) -> tensor<*xf32> // HISTOGRAM-MSE-BRUTEFORCE-CHECK-NEXT: return [[res]] : tensor<*xf32> // HISTOGRAM-MSE-BRUTEFORCE-CHECK: func @no_custom_ops_on_non_f32_type @@ -107,10 +107,10 @@ module { // calibration_parameters=CalibrationParameters(initial_num_bins=256) // ) // HISTOGRAM-MSE-MAX-FREQUENCY-CHECK: func @add_custom_ops -// HISTOGRAM-MSE-MAX-FREQUENCY-CHECK-NEXT: [[rhs:%.*]] = "tf.CustomAggregator"(%arg1) {calibration_method = 5 : i32, id = "", initial_num_bins = 256 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32} : (tensor<*xf32>) -> tensor<*xf32> -// HISTOGRAM-MSE-MAX-FREQUENCY-CHECK-NEXT: [[lhs:%.*]] = "tf.CustomAggregator"(%arg0) {calibration_method = 5 : i32, id = "", initial_num_bins = 256 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32} : (tensor<*xf32>) -> tensor<*xf32> +// HISTOGRAM-MSE-MAX-FREQUENCY-CHECK-NEXT: [[rhs:%.*]] = "tf.CustomAggregator"(%arg1) <{id = ""}> {calibration_method = 5 : i32, initial_num_bins = 256 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32} : (tensor<*xf32>) -> tensor<*xf32> +// HISTOGRAM-MSE-MAX-FREQUENCY-CHECK-NEXT: [[lhs:%.*]] = "tf.CustomAggregator"(%arg0) <{id = ""}> {calibration_method = 5 : i32, initial_num_bins = 256 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32} : (tensor<*xf32>) -> tensor<*xf32> // HISTOGRAM-MSE-MAX-FREQUENCY-CHECK-NEXT: [[add:%.*]] = "tf.AddV2"([[lhs]], [[rhs]]) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32> -// HISTOGRAM-MSE-MAX-FREQUENCY-CHECK-NEXT: [[res:%.*]] = "tf.CustomAggregator"([[add]]) {calibration_method = 5 : i32, id = "", initial_num_bins = 256 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32} : (tensor<*xf32>) -> tensor<*xf32> +// HISTOGRAM-MSE-MAX-FREQUENCY-CHECK-NEXT: [[res:%.*]] = "tf.CustomAggregator"([[add]]) <{id = ""}> {calibration_method = 5 : i32, initial_num_bins = 256 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32} : (tensor<*xf32>) -> tensor<*xf32> // HISTOGRAM-MSE-MAX-FREQUENCY-CHECK-NEXT: return [[res]] : tensor<*xf32> // HISTOGRAM-MSE-MAX-FREQUENCY-CHECK: func @no_custom_ops_on_non_f32_type @@ -128,10 +128,10 @@ module { // calibration_parameters=CalibrationParameters(initial_num_bins=256) // ) // HISTOGRAM-MSE-SYMMETRIC-CHECK: func @add_custom_ops -// HISTOGRAM-MSE-SYMMETRIC-CHECK-NEXT: [[rhs:%.*]] = "tf.CustomAggregator"(%arg1) {calibration_method = 6 : i32, id = "", initial_num_bins = 256 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32} : (tensor<*xf32>) -> tensor<*xf32> -// HISTOGRAM-MSE-SYMMETRIC-CHECK-NEXT: [[lhs:%.*]] = "tf.CustomAggregator"(%arg0) {calibration_method = 6 : i32, id = "", initial_num_bins = 256 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32} : (tensor<*xf32>) -> tensor<*xf32> +// HISTOGRAM-MSE-SYMMETRIC-CHECK-NEXT: [[rhs:%.*]] = "tf.CustomAggregator"(%arg1) <{id = ""}> {calibration_method = 6 : i32, initial_num_bins = 256 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32} : (tensor<*xf32>) -> tensor<*xf32> +// HISTOGRAM-MSE-SYMMETRIC-CHECK-NEXT: [[lhs:%.*]] = "tf.CustomAggregator"(%arg0) <{id = ""}> {calibration_method = 6 : i32, initial_num_bins = 256 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32} : (tensor<*xf32>) -> tensor<*xf32> // HISTOGRAM-MSE-SYMMETRIC-CHECK-NEXT: [[add:%.*]] = "tf.AddV2"([[lhs]], [[rhs]]) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32> -// HISTOGRAM-MSE-SYMMETRIC-CHECK-NEXT: [[res:%.*]] = "tf.CustomAggregator"([[add]]) {calibration_method = 6 : i32, id = "", initial_num_bins = 256 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32} : (tensor<*xf32>) -> tensor<*xf32> +// HISTOGRAM-MSE-SYMMETRIC-CHECK-NEXT: [[res:%.*]] = "tf.CustomAggregator"([[add]]) <{id = ""}> {calibration_method = 6 : i32, initial_num_bins = 256 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32} : (tensor<*xf32>) -> tensor<*xf32> // HISTOGRAM-MSE-SYMMETRIC-CHECK-NEXT: return [[res]] : tensor<*xf32> // HISTOGRAM-MSE-SYMMETRIC-CHECK: func @no_custom_ops_on_non_f32_type diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/tests/insert_main_function.mlir b/tensorflow/compiler/mlir/quantization/tensorflow/tests/insert_main_function.mlir index 56208364ad8b3d..be98dd179bae11 100644 --- a/tensorflow/compiler/mlir/quantization/tensorflow/tests/insert_main_function.mlir +++ b/tensorflow/compiler/mlir/quantization/tensorflow/tests/insert_main_function.mlir @@ -25,7 +25,7 @@ module attributes {tf.versions = {producer = 930 : i32}, tf_saved_model.semantic func.return %1 : tensor<1xf32> } // CHECK: func private @mul2(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>) -> tensor<1xf32> attributes {tf.entry_function = {inputs = "mul2_y:0,mul2_x:0", outputs = "PartitionedCall_1:0"}} { -// CHECK: %[[CONST_0:.*]] = "tf.Const"() {value = dense<2.000000e+00> : tensor} : () -> tensor +// CHECK: %[[CONST_0:.*]] = "tf.Const"() <{value = dense<2.000000e+00> : tensor}> : () -> tensor // CHECK: %[[MUL_1:.*]] = "tf.Mul"(%arg1, %arg0) : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> // CHECK: %[[MUL_2:.*]] = "tf.Mul"(%[[MUL_1]], %[[CONST_0]]) : (tensor<1xf32>, tensor) -> tensor<1xf32> // CHECK: return %[[MUL_2]] : tensor<1xf32> @@ -33,8 +33,8 @@ module attributes {tf.versions = {producer = 930 : i32}, tf_saved_model.semantic // CHECK: func @main(%arg0: tensor<1xf32> {tf_saved_model.index_path = ["mul1_y:0"]}, %arg1: tensor<1xf32> {tf_saved_model.index_path = ["mul1_x:0"]}, %arg2: tensor<1xf32> {tf_saved_model.index_path = ["mul2_y:0"]}, %arg3: tensor<1xf32> {tf_saved_model.index_path = ["mul2_x:0"]}) -> (tensor<1xf32> {tf_saved_model.index_path = ["PartitionedCall:0"]}, tensor<1xf32> {tf_saved_model.index_path = ["PartitionedCall_1:0"]}) attributes {tf.entry_function = {inputs = "mul1_y:0,mul1_x:0,mul2_y:0,mul2_x:0", outputs = "PartitionedCall:0,PartitionedCall_1:0"}, tf_saved_model.exported_names = ["main"]} { // CHECK-NOT: f = @NoOp -// CHECK: %[[PARTITIONEDCALL_0:.*]] = "tf.PartitionedCall"(%arg0, %arg1) {config = "", config_proto = "", executor_type = "", f = @mul1} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> -// CHECK: %[[PARTITIONEDCALL_1:.*]] = "tf.PartitionedCall"(%arg2, %arg3) {config = "", config_proto = "", executor_type = "", f = @mul2} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> +// CHECK: %[[PARTITIONEDCALL_0:.*]] = "tf.PartitionedCall"(%arg0, %arg1) <{config = "", config_proto = "", executor_type = "", f = @mul1}> : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> +// CHECK: %[[PARTITIONEDCALL_1:.*]] = "tf.PartitionedCall"(%arg2, %arg3) <{config = "", config_proto = "", executor_type = "", f = @mul2}> : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> // CHECK-DAG: %[[IDENTITY_0:.*]] = "tf.Identity"(%[[PARTITIONEDCALL_0]]) // CHECK-DAG: %[[IDENTITY_1:.*]] = "tf.Identity"(%[[PARTITIONEDCALL_1]]) // CHECK: return %[[IDENTITY_0]], %[[IDENTITY_1]] : tensor<1xf32>, tensor<1xf32> @@ -82,10 +82,10 @@ module attributes {tf.versions = {producer = 1132 : i32}, tf_saved_model.semanti // CHECK-SAME: tf_saved_model.exported_names = ["main"] // Check that the function call to @add exists and not to @NoOp. -// CHECK: %[[CALL0:.*]] = "tf.PartitionedCall"(%[[ARG0]], %[[ARG1]]) { +// CHECK: %[[CALL0:.*]] = "tf.PartitionedCall"(%[[ARG0]], %[[ARG1]]) <{ // CHECK-NOT: f = @NoOp // CHECK-SAME: f = @add -// CHECK-SAME: } +// CHECK-SAME: }> // CHECK-SAME: : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32> // CHECK: %[[IDENTITY:.*]] = "tf.Identity"(%[[CALL0]]) // CHECK: return %[[IDENTITY]] : tensor<1xf32> @@ -111,7 +111,7 @@ module attributes {tf.versions = {producer = 930 : i32}, tf_saved_model.semantic // CHECK: func.func @main(%arg0: tensor<16xf32> {tf_saved_model.index_path = ["input:0"]}, %arg1: tensor {tf_saved_model.index_path = ["k:0"]}) // CHECK-SAME: -> (tensor {tf_saved_model.index_path = ["TopK:0"]}, tensor {tf_saved_model.index_path = ["TopK:1"]}) // CHECK-SAME: attributes {tf.entry_function = {inputs = "input:0,k:0", outputs = "TopK:0,TopK:1"}, tf_saved_model.exported_names = ["main"]} -// CHECK: %[[CALL0:.*]]:2 = "tf.PartitionedCall"(%arg0, %arg1) {config = "", config_proto = "", executor_type = "", f = @topk} +// CHECK: %[[CALL0:.*]]:2 = "tf.PartitionedCall"(%arg0, %arg1) <{config = "", config_proto = "", executor_type = "", f = @topk}> // Expects an IdentityN op to be created. // CHECK: %[[IDENTITY:.*]]:2 = "tf.IdentityN"(%[[CALL0]]#0, %[[CALL0]]#1) : (tensor, tensor) -> (tensor, tensor) // CHECK: return %[[IDENTITY]]#0, %[[IDENTITY]]#1 : tensor, tensor diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/tests/insert_restore_op.mlir b/tensorflow/compiler/mlir/quantization/tensorflow/tests/insert_restore_op.mlir index 385052ab9e5cb5..7f73eee6f056ca 100644 --- a/tensorflow/compiler/mlir/quantization/tensorflow/tests/insert_restore_op.mlir +++ b/tensorflow/compiler/mlir/quantization/tensorflow/tests/insert_restore_op.mlir @@ -32,10 +32,10 @@ module attributes {tf_saved_model.semantics} { // Test that RestoreV2 op is created with 1 resulting value. // CHECK: %[[RESTORE:.*]] = "tf.RestoreV2"(%[[ARG_0]], %[[CST_1]], %[[CST_2]]) : (tensor, tensor<1x!tf_type.string>, tensor<1x!tf_type.string>) -> tensor<2xf32> -// CHECK: "tf.AssignVariableOp"(%[[VAR_HANDLE]], %[[RESTORE]]) {validate_shape = false} : (tensor>>, tensor<2xf32>) -> () +// CHECK: "tf.AssignVariableOp"(%[[VAR_HANDLE]], %[[RESTORE]]) <{validate_shape = false}> : (tensor>>, tensor<2xf32>) -> () // Test that the loc is properly set to it's shared_name. -// CHECK-LOC: "tf.VarHandleOp"() {{{.*shared_name = "var_0".*}}} +// CHECK-LOC: "tf.VarHandleOp"() <{{{.*shared_name = "var_0".*}}}> // CHECK-LOC-SAME: loc("var_0") } @@ -66,19 +66,19 @@ module attributes {tf_saved_model.semantics} { // CHECK-DAG: %[[VAR_HANDLE_0:.*]] = "tf.VarHandleOp"() {{.*shared_name = "var_0".*}} : () -> tensor>> // CHECK-DAG: %[[VAR_HANDLE_1:.*]] = "tf.VarHandleOp"() {{.*shared_name = "var_1".*}} : () -> tensor>> -// CHECK-DAG: %[[CST_0:.*]] = "tf.Const"() {{{.*value = dense<\["var_0", "var_1"\]> : tensor<2x!tf_type.string>.*}}} -// CHECK-DAG: %[[CST_1:.*]] = "tf.Const"() {{{.*value = dense<""> : tensor<2x!tf_type.string>.*}}} +// CHECK-DAG: %[[CST_0:.*]] = "tf.Const"() <{{{.*value = dense<\["var_0", "var_1"\]> : tensor<2x!tf_type.string>.*}}}> +// CHECK-DAG: %[[CST_1:.*]] = "tf.Const"() <{{{.*value = dense<""> : tensor<2x!tf_type.string>.*}}}> // Test that RestoreV2 op is created with 2 resulting values. // CHECK: %[[RESTORE:.*]]:2 = "tf.RestoreV2"(%[[ARG_0]], %[[CST_0]], %[[CST_1]]) : (tensor, tensor<2x!tf_type.string>, tensor<2x!tf_type.string>) -> (tensor<2xf32>, tensor<4xi32>) -// CHECK: "tf.AssignVariableOp"(%[[VAR_HANDLE_0]], %[[RESTORE]]#0) {validate_shape = false} : (tensor>>, tensor<2xf32>) -> () -// CHECK: "tf.AssignVariableOp"(%[[VAR_HANDLE_1]], %[[RESTORE]]#1) {validate_shape = false} : (tensor>>, tensor<4xi32>) -> () +// CHECK: "tf.AssignVariableOp"(%[[VAR_HANDLE_0]], %[[RESTORE]]#0) <{validate_shape = false}> : (tensor>>, tensor<2xf32>) -> () +// CHECK: "tf.AssignVariableOp"(%[[VAR_HANDLE_1]], %[[RESTORE]]#1) <{validate_shape = false}> : (tensor>>, tensor<4xi32>) -> () // Test that the locs are properly set to their shared_names. -// CHECK-LOC: "tf.VarHandleOp"() {{{.*shared_name = "var_0".*}}} +// CHECK-LOC: "tf.VarHandleOp"() <{{{.*shared_name = "var_0".*}}}> // CHECK-LOC-SAME: loc("var_0") -// CHECK-LOC: "tf.VarHandleOp"() {{{.*shared_name = "var_1".*}}} +// CHECK-LOC: "tf.VarHandleOp"() <{{{.*shared_name = "var_1".*}}}> // CHECK-LOC-SAME: loc("var_1") } @@ -101,11 +101,11 @@ module attributes {tf_saved_model.semantics} { // Check that no function argument is created. // CHECK: func.func @init_func_init_op() -// CHECK-DAG: %[[VAR_HANDLE:.*]] = "tf.VarHandleOp"() {{{.*shared_name = "var_0".*}}} : () -> tensor>> -// CHECK-DAG: %[[CST:.*]] = "tf.Const"() {{{.*value = dense<1.000000e\+00> : tensor<2xf32>.*}}} +// CHECK-DAG: %[[VAR_HANDLE:.*]] = "tf.VarHandleOp"() <{{{.*shared_name = "var_0".*}}}> : () -> tensor>> +// CHECK-DAG: %[[CST:.*]] = "tf.Const"() <{{{.*value = dense<1.000000e\+00> : tensor<2xf32>.*}}}> // Make sure that "tf.RestoreV2" is not created. // CHECK-NOT: "tf.RestoreV2" -// CHECK: "tf.AssignVariableOp"(%[[VAR_HANDLE]], %[[CST]]) {validate_shape = false} : (tensor>>, tensor<2xf32>) -> () +// CHECK: "tf.AssignVariableOp"(%[[VAR_HANDLE]], %[[CST]]) <{validate_shape = false}> : (tensor>>, tensor<2xf32>) -> () // CHECK-LOC: @init_func_init_op // CHECK-LOC: return @@ -140,19 +140,19 @@ module attributes {tf_saved_model.semantics} { // CHECK-DAG: %[[VAR_HANDLE_0:.*]] = "tf.VarHandleOp"() {{.*shared_name = "var_0".*}} : () -> tensor>> // CHECK-DAG: %[[VAR_HANDLE_1:.*]] = "tf.VarHandleOp"() {{.*shared_name = "var_1".*}} : () -> tensor>> -// CHECK-DAG: %[[CST_0:.*]] = "tf.Const"() {{{.*value = dense<\["var_0", "var_1"\]> : tensor<2x!tf_type.string>.*}}} -// CHECK-DAG: %[[CST_1:.*]] = "tf.Const"() {{{.*value = dense<""> : tensor<2x!tf_type.string>.*}}} +// CHECK-DAG: %[[CST_0:.*]] = "tf.Const"() <{{{.*value = dense<\["var_0", "var_1"\]> : tensor<2x!tf_type.string>.*}}}> +// CHECK-DAG: %[[CST_1:.*]] = "tf.Const"() <{{{.*value = dense<""> : tensor<2x!tf_type.string>.*}}}> // Test that RestoreV2 op is created with 2 resulting values. // CHECK: %[[RESTORE:.*]]:2 = "tf.RestoreV2"(%[[ARG_0]], %[[CST_0]], %[[CST_1]]) : (tensor, tensor<2x!tf_type.string>, tensor<2x!tf_type.string>) -> (tensor<2xf32>, tensor<2xf32>) -// CHECK: "tf.AssignVariableOp"(%[[VAR_HANDLE_0]], %[[RESTORE]]#0) {validate_shape = false} : (tensor>>, tensor<2xf32>) -> () -// CHECK: "tf.AssignVariableOp"(%[[VAR_HANDLE_1]], %[[RESTORE]]#1) {validate_shape = false} : (tensor>>, tensor<2xf32>) -> () +// CHECK: "tf.AssignVariableOp"(%[[VAR_HANDLE_0]], %[[RESTORE]]#0) <{validate_shape = false}> : (tensor>>, tensor<2xf32>) -> () +// CHECK: "tf.AssignVariableOp"(%[[VAR_HANDLE_1]], %[[RESTORE]]#1) <{validate_shape = false}> : (tensor>>, tensor<2xf32>) -> () // Test that the locs are properly set to their shared_names. -// CHECK-LOC: "tf.VarHandleOp"() {{{.*shared_name = "var_0".*}}} +// CHECK-LOC: "tf.VarHandleOp"() <{{{.*shared_name = "var_0".*}}}> // CHECK-LOC-SAME: loc("var_0") -// CHECK-LOC: "tf.VarHandleOp"() {{{.*shared_name = "var_1".*}}} +// CHECK-LOC: "tf.VarHandleOp"() <{{{.*shared_name = "var_1".*}}}> // CHECK-LOC-SAME: loc("var_1") } diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/tests/insert_save_op.mlir b/tensorflow/compiler/mlir/quantization/tensorflow/tests/insert_save_op.mlir index 9483331b6fe31d..c142247cdbc235 100644 --- a/tensorflow/compiler/mlir/quantization/tensorflow/tests/insert_save_op.mlir +++ b/tensorflow/compiler/mlir/quantization/tensorflow/tests/insert_save_op.mlir @@ -30,8 +30,8 @@ module attributes {tf_saved_model.semantics} { // CHECK: %[[VAR_HANDLE:.*]] = "tf.VarHandleOp"() // CHECK-SAME: {{.*shared_name = "var_0".*}} // CHECK: %[[READ_VARIABLE:.*]] = "tf.ReadVariableOp"(%[[VAR_HANDLE]]) : (tensor>>) -> tensor<2xf32> -// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() {{{.*value = dense<"var_0"> : tensor<1x!tf_type.string>.*}}} -// CHECK-DAG: %[[CONST_1:.*]] = "tf.Const"() {{{.*value = dense<""> : tensor<1x!tf_type.string>.*}}} +// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() <{{{.*value = dense<"var_0"> : tensor<1x!tf_type.string>.*}}}> +// CHECK-DAG: %[[CONST_1:.*]] = "tf.Const"() <{{{.*value = dense<""> : tensor<1x!tf_type.string>.*}}}> // CHECK: "tf.SaveV2"(%[[ARG]], %[[CONST_0]], %[[CONST_1]], %[[READ_VARIABLE]]) // CHECK: return } @@ -73,8 +73,8 @@ module attributes {tf_saved_model.semantics} { // CHECK-DAG: %[[READ_VARIABLE_0:.*]] = "tf.ReadVariableOp"(%[[VAR_HANDLE_0]]) : (tensor>>) -> tensor<2xf32> // CHECK-DAG: %[[READ_VARIABLE_1:.*]] = "tf.ReadVariableOp"(%[[VAR_HANDLE_1]]) : (tensor>>) -> tensor<3xf32> -// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() {{{.*value = dense<\["var_0", "var_1"\]> : tensor<2x!tf_type.string>.*}}} -// CHECK-DAG: %[[CONST_1:.*]] = "tf.Const"() {{{.*value = dense<""> : tensor<2x!tf_type.string>.*}}} +// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() <{{{.*value = dense<\["var_0", "var_1"\]> : tensor<2x!tf_type.string>.*}}}> +// CHECK-DAG: %[[CONST_1:.*]] = "tf.Const"() <{{{.*value = dense<""> : tensor<2x!tf_type.string>.*}}}> // CHECK: "tf.SaveV2"(%[[ARG]], %[[CONST_0]], %[[CONST_1]], %[[READ_VARIABLE_0]], %[[READ_VARIABLE_1]]) // CHECK: return } diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/tests/issue_ids_of_custom_aggregation_ops.mlir b/tensorflow/compiler/mlir/quantization/tensorflow/tests/issue_ids_of_custom_aggregation_ops.mlir index 03fcbf4a2b0623..4aa1ae76b8a83d 100644 --- a/tensorflow/compiler/mlir/quantization/tensorflow/tests/issue_ids_of_custom_aggregation_ops.mlir +++ b/tensorflow/compiler/mlir/quantization/tensorflow/tests/issue_ids_of_custom_aggregation_ops.mlir @@ -10,8 +10,8 @@ func.func @issue_ids(%arg0: tensor<*xf32>, %arg1: tensor<*xf32>) -> tensor<*xf32 // CHECK: func @issue_ids -// CHECK-NEXT: [[rhs:%.*]] = "tf.CustomAggregator"(%arg1) {id = "0"} : (tensor<*xf32>) -> tensor<*xf32> -// CHECK-NEXT: [[lhs:%.*]] = "tf.CustomAggregator"(%arg0) {id = "1"} : (tensor<*xf32>) -> tensor<*xf32> +// CHECK-NEXT: [[rhs:%.*]] = "tf.CustomAggregator"(%arg1) <{id = "0"}> : (tensor<*xf32>) -> tensor<*xf32> +// CHECK-NEXT: [[lhs:%.*]] = "tf.CustomAggregator"(%arg0) <{id = "1"}> : (tensor<*xf32>) -> tensor<*xf32> // CHECK-NEXT: [[add:%.*]] = "tf.AddV2"([[lhs]], [[rhs]]) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32> -// CHECK-NEXT: [[res:%.*]] = "tf.CustomAggregator"([[add]]) {id = "2"} : (tensor<*xf32>) -> tensor<*xf32> +// CHECK-NEXT: [[res:%.*]] = "tf.CustomAggregator"([[add]]) <{id = "2"}> : (tensor<*xf32>) -> tensor<*xf32> // CHECK-NEXT: return [[res]] : tensor<*xf32> diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/tests/lift_quantizable_spots_as_functions.mlir b/tensorflow/compiler/mlir/quantization/tensorflow/tests/lift_quantizable_spots_as_functions.mlir index f61a9fbe9feb6f..6a7f9da6bc5563 100644 --- a/tensorflow/compiler/mlir/quantization/tensorflow/tests/lift_quantizable_spots_as_functions.mlir +++ b/tensorflow/compiler/mlir/quantization/tensorflow/tests/lift_quantizable_spots_as_functions.mlir @@ -26,10 +26,10 @@ func.func @float_conv(%arg0: tensor<1x3x4x3xf32>, %arg1: tensor<2x3x3x2xf32>) -> %7 = "tf.BiasAdd"(%6, %cst) {data_format = "NHWC", device = ""} : (tensor<*xf32>, tensor<2xf32>) -> tensor<*xf32> func.return %2, %5, %7 : tensor<*xf32>, tensor<*xf32>, tensor<*xf32> -// CHECK: %[[CONST_0:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<2xf32>} : () -> tensor<2xf32> +// CHECK: %[[CONST_0:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<2xf32>}> : () -> tensor<2xf32> // CHECK: %[[PARTITIONEDCALL_0:.*]] = "tf.PartitionedCall"(%arg0, %arg1, %[[CONST_0]]) -// CHECK-SAME: {_tfl_quant_trait = "fully_quantizable", -// CHECK-SAME: f = @composite_conv2d_with_bias_and_relu6_fn_1} +// CHECK-SAME: f = @composite_conv2d_with_bias_and_relu6_fn_1}> +// CHECK-SAME: {_tfl_quant_trait = "fully_quantizable" // CHECK: %[[PARTITIONEDCALL_1:.*]] = "tf.PartitionedCall"(%arg0, %arg1, %[[CONST_0]]) // CHECK-SAME: f = @composite_conv2d_with_bias_and_relu_fn_1} // CHECK: %[[PARTITIONEDCALL_2:.*]] = "tf.PartitionedCall"(%arg0, %arg1, %[[CONST_0]]) @@ -39,8 +39,8 @@ func.func @float_conv(%arg0: tensor<1x3x4x3xf32>, %arg1: tensor<2x3x3x2xf32>) -> // CHECK-LABEL: private @composite_conv2d_with_bias_and_relu6_fn_1 // CHECK-NEXT: %[[CONV2D_0:.*]] = "tf.Conv2D"(%arg0, %arg1) +// CHECK-SAME: data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 2, 1], use_cudnn_on_gpu = true // CHECK-SAME: attr_map = "0:strides,1:use_cudnn_on_gpu,2:padding,3:explicit_paddings,4:dilations" -// CHECK-SAME: data_format = "NHWC", device = "", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 2, 1], use_cudnn_on_gpu = true // CHECK-NEXT: %[[BIASADD_0:.*]] = "tf.BiasAdd"(%[[CONV2D_0]], %arg2) // CHECK-NEXT: %[[RELU6_0:.*]] = "tf.Relu6"(%[[BIASADD_0]]) // CHECK-NEXT: return %[[RELU6_0]] @@ -70,15 +70,15 @@ func.func @float_conv_strides_equals_to_dilations(%arg0: tensor<1x3x4x3xf32>, %a } // CHECK-LABEL: func @float_conv_strides_equals_to_dilations(%arg0: tensor<1x3x4x3xf32>, %arg1: tensor<2x3x3x2xf32>) -> tensor<*xf32> { -// CHECK: %[[CONST_0:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<2xf32>} : () -> tensor<2xf32> -// CHECK: %[[PARTITIONEDCALL_0:.*]] = "tf.PartitionedCall"(%arg0, %arg1, %[[CONST_0]]) {_tfl_quant_trait = "fully_quantizable", config = "", config_proto = "", executor_type = "", f = @composite_conv2d_with_bias_and_relu6_fn_1} : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>, tensor<2xf32>) -> tensor<*xf32> +// CHECK: %[[CONST_0:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<2xf32>}> : () -> tensor<2xf32> +// CHECK: %[[PARTITIONEDCALL_0:.*]] = "tf.PartitionedCall"(%arg0, %arg1, %[[CONST_0]]) <{config = "", config_proto = "", executor_type = "", f = @composite_conv2d_with_bias_and_relu6_fn_1}> {_tfl_quant_trait = "fully_quantizable"} : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>, tensor<2xf32>) -> tensor<*xf32> // CHECK: return %[[PARTITIONEDCALL_0]] : tensor<*xf32> // CHECK: } // CHECK-LABEL: func private @composite_conv2d_with_bias_and_relu6_fn_1(%arg0: tensor<1x3x4x3xf32>, %arg1: tensor<2x3x3x2xf32>, %arg2: tensor<2xf32>) -> tensor<*xf32> attributes {tf_quant.composite_function} { // CHECK-NEXT: %[[CONV2D_0:.*]] = "tf.Conv2D"(%arg0, %arg1) +// CHECK-SAME: data_format = "NHWC", dilations = [1, 1, 2, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 2, 1], use_cudnn_on_gpu = true // CHECK-SAME: attr_map = "0:strides,1:use_cudnn_on_gpu,2:padding,3:explicit_paddings,4:dilations" -// CHECK-SAME: data_format = "NHWC", device = "", dilations = [1, 1, 2, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 2, 1], use_cudnn_on_gpu = true // CHECK-NEXT: %[[BIASADD_0:.*]] = "tf.BiasAdd"(%[[CONV2D_0]], %arg2) // CHECK-NEXT: %[[RELU6_0:.*]] = "tf.Relu6"(%[[BIASADD_0]]) // CHECK-NEXT: return %[[RELU6_0]] @@ -111,14 +111,14 @@ func.func @float_depthwise_conv(%arg0: tensor<1x3x4x3xf32>, %arg1: tensor<2x3x3x %7 = "tf.BiasAdd"(%6, %cst) {data_format = "NHWC", device = ""} : (tensor<*xf32>, tensor<2xf32>) -> tensor<*xf32> func.return %2, %5, %7 : tensor<*xf32>, tensor<*xf32>, tensor<*xf32> -// CHECK: %[[CONST_0:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<2xf32>} : () -> tensor<2xf32> +// CHECK: %[[CONST_0:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<2xf32>}> : () -> tensor<2xf32> // CHECK: %[[PARTITIONEDCALL_0:.*]] = "tf.PartitionedCall"(%arg0, %arg1, %[[CONST_0]]) -// CHECK-SAME: _tfl_quant_trait = "fully_quantizable", -// CHECK-SAME: f = @composite_depthwise_conv2d_with_bias_and_relu6_fn_1} +// CHECK-SAME: f = @composite_depthwise_conv2d_with_bias_and_relu6_fn_1}> +// CHECK-SAME: _tfl_quant_trait = "fully_quantizable" // CHECK: %[[PARTITIONEDCALL_1:.*]] = "tf.PartitionedCall"(%arg0, %arg1, %[[CONST_0]]) -// CHECK-SAME: f = @composite_depthwise_conv2d_with_bias_and_relu_fn_1} +// CHECK-SAME: f = @composite_depthwise_conv2d_with_bias_and_relu_fn_1 // CHECK: %[[PARTITIONEDCALL_2:.*]] = "tf.PartitionedCall"(%arg0, %arg1, %[[CONST_0]]) -// CHECK-SAME: f = @composite_depthwise_conv2d_with_bias_fn_1} +// CHECK-SAME: f = @composite_depthwise_conv2d_with_bias_fn_1 // CHECK: return %[[PARTITIONEDCALL_0]], %[[PARTITIONEDCALL_1]], %[[PARTITIONEDCALL_2]] // CHECK: } @@ -161,14 +161,14 @@ func.func @float_matmul( %7 = "tf.BiasAdd"(%6, %cst) {data_format = "NHWC", device = ""} : (tensor<*xf32>, tensor<10xf32>) -> tensor<*xf32> func.return %2, %5, %7 : tensor<*xf32>, tensor<*xf32>, tensor<*xf32> -// CHECK: %[[CONST_0:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<10xf32>} +// CHECK: %[[CONST_0:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<10xf32>}> // CHECK: %[[PARTITIONEDCALL_0:.*]] = "tf.PartitionedCall"(%arg0, %arg1, %[[CONST_0]]) -// CHECK-SAME: {_tfl_quant_trait = "fully_quantizable", -// CHECK-SAME: f = @composite_matmul_with_bias_and_relu6_fn_1} +// CHECK-SAME: f = @composite_matmul_with_bias_and_relu6_fn_1}> +// CHECK-SAME: {_tfl_quant_trait = "fully_quantizable" // CHECK: %[[PARTITIONEDCALL_1:.*]] = "tf.PartitionedCall"(%arg0, %arg1, %[[CONST_0]]) -// CHECK-SAME: f = @composite_matmul_with_bias_and_relu_fn_1} +// CHECK-SAME: f = @composite_matmul_with_bias_and_relu_fn_1 // CHECK: %[[PARTITIONEDCALL_2:.*]] = "tf.PartitionedCall"(%arg0, %arg1, %[[CONST_0]]) -// CHECK-SAME: f = @composite_matmul_with_bias_fn_1} +// CHECK-SAME: f = @composite_matmul_with_bias_fn_1 // CHECK: return %[[PARTITIONEDCALL_0]], %[[PARTITIONEDCALL_1]], %[[PARTITIONEDCALL_2]] // CHECK: } @@ -207,10 +207,10 @@ func.func @float_matmul_with_reshape(%arg0: tensor<1x10xf32>, %arg1: tensor<10x1 func.return %3 : tensor<*xf32> -// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<10xf32>} -// CHECK-DAG: %[[SHAPE:.*]] = "tf.Const"() {value = dense<[-1, 10]> : tensor<2xi32>} +// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<10xf32>}> +// CHECK-DAG: %[[SHAPE:.*]] = "tf.Const"() <{value = dense<[-1, 10]> : tensor<2xi32>}> // CHECK: %[[PARTITIONEDCALL_0:.*]] = "tf.PartitionedCall"(%arg0, %arg1, %[[CONST_0]], %[[SHAPE]]) -// CHECK-SAME: f = @composite_matmul_with_reshape_and_bias_fn_1} +// CHECK-SAME: f = @composite_matmul_with_reshape_and_bias_fn_1 // CHECK: return %[[PARTITIONEDCALL_0]] // CHECK: } @@ -247,14 +247,14 @@ func.func @float_conv_no_bias(%arg0: tensor<1x3x4x3xf32>, %arg1: tensor<2x3x3x2x func.return %1, %4, %6 : tensor<*xf32>, tensor<*xf32>, tensor<*xf32> // CHECK: %[[PARTITIONEDCALL_0:.*]] = "tf.PartitionedCall"(%arg0, %arg1) -// CHECK-SAME: {_tfl_quant_trait = "fully_quantizable", -// CHECK-SAME: f = @composite_conv2d_with_relu6_fn_1} +// CHECK-SAME: f = @composite_conv2d_with_relu6_fn_1}> +// CHECK-SAME: {_tfl_quant_trait = "fully_quantizable" // CHECK: %[[PARTITIONEDCALL_1:.*]] = "tf.PartitionedCall"(%arg0, %arg1) -// CHECK-SAME: f = @composite_conv2d_with_relu_fn_1} +// CHECK-SAME: f = @composite_conv2d_with_relu_fn_1 // CHECK: %[[PARTITIONEDCALL_2:.*]] = "tf.PartitionedCall"(%arg0, %arg1) -// CHECK-SAME: f = @composite_conv2d_fn_1} +// CHECK-SAME: f = @composite_conv2d_fn_1 // CHECK: return %[[PARTITIONEDCALL_0]], %[[PARTITIONEDCALL_1]], %[[PARTITIONEDCALL_2]] // CHECK: } @@ -288,12 +288,12 @@ func.func @float_depthwise_conv_no_bias(%arg0: tensor<1x3x4x3xf32>, %arg1: tenso func.return %1, %4, %6 : tensor<*xf32>, tensor<*xf32>, tensor<*xf32> // CHECK: %[[PARTITIONEDCALL_0:.*]] = "tf.PartitionedCall"(%arg0, %arg1) -// CHECK-SAME: {_tfl_quant_trait = "fully_quantizable", -// CHECK-SAME: f = @composite_depthwise_conv2d_with_relu6_fn_1} +// CHECK-SAME: f = @composite_depthwise_conv2d_with_relu6_fn_1}> +// CHECK-SAME: {_tfl_quant_trait = "fully_quantizable" // CHECK: %[[PARTITIONEDCALL_1:.*]] = "tf.PartitionedCall"(%arg0, %arg1) -// CHECK-SAME: f = @composite_depthwise_conv2d_with_relu_fn_1} +// CHECK-SAME: f = @composite_depthwise_conv2d_with_relu_fn_1 // CHECK: %[[PARTITIONEDCALL_2:.*]] = "tf.PartitionedCall"(%arg0, %arg1) -// CHECK-SAME: f = @composite_depthwise_conv2d_fn_1} +// CHECK-SAME: f = @composite_depthwise_conv2d_fn_1 // CHECK: return %[[PARTITIONEDCALL_0]], %[[PARTITIONEDCALL_1]], %[[PARTITIONEDCALL_2]] // CHECK: } @@ -323,12 +323,12 @@ func.func @float_matmul_no_bias( func.return %1, %4, %6 : tensor<*xf32>, tensor<*xf32>, tensor<*xf32> // CHECK: %[[PARTITIONEDCALL_0:.*]] = "tf.PartitionedCall"(%arg0, %arg1) -// CHECK-SAME: {_tfl_quant_trait = "fully_quantizable", -// CHECK-SAME: f = @composite_matmul_with_relu6_fn_1} +// CHECK-SAME: f = @composite_matmul_with_relu6_fn_1}> +// CHECK-SAME: {_tfl_quant_trait = "fully_quantizable" // CHECK: %[[PARTITIONEDCALL_1:.*]] = "tf.PartitionedCall"(%arg0, %arg1) -// CHECK-SAME: f = @composite_matmul_with_relu_fn_1} +// CHECK-SAME: f = @composite_matmul_with_relu_fn_1 // CHECK: %[[PARTITIONEDCALL_2:.*]] = "tf.PartitionedCall"(%arg0, %arg1) -// CHECK-SAME: f = @composite_matmul_fn_1} +// CHECK-SAME: f = @composite_matmul_fn_1 // CHECK: return %[[PARTITIONEDCALL_0]], %[[PARTITIONEDCALL_1]], %[[PARTITIONEDCALL_2]] // CHECK: } @@ -361,14 +361,14 @@ func.func @conv3d_no_bias(%arg0: tensor<1x3x4x3x3xf32>) -> (tensor<1x3x2x3x2xf32 // CHECK-DAG: %[[CST:.*]] = "tf.Const"() {{.*}} : () -> tensor<2x3x3x3x2xf32> // CHECK: %[[PARTITIONEDCALL_0:.*]] = "tf.PartitionedCall"(%arg0, %[[CST]]) -// CHECK-SAME: {_tfl_quant_trait = "fully_quantizable", -// CHECK-SAME: f = @composite_conv3d_with_relu_fn_1} +// CHECK-SAME: f = @composite_conv3d_with_relu_fn_1}> +// CHECK-SAME: {_tfl_quant_trait = "fully_quantizable" // CHECK: %[[PARTITIONEDCALL_1:.*]] = "tf.PartitionedCall"(%arg0, %[[CST]]) -// CHECK-SAME: f = @composite_conv3d_with_relu6_fn_1} +// CHECK-SAME: f = @composite_conv3d_with_relu6_fn_1 // CHECK: %[[PARTITIONEDCALL_2:.*]] = "tf.PartitionedCall"(%arg0, %[[CST]]) -// CHECK-SAME: f = @composite_conv3d_fn_1} +// CHECK-SAME: f = @composite_conv3d_fn_1 // CHECK: return %[[PARTITIONEDCALL_0]], %[[PARTITIONEDCALL_1]], %[[PARTITIONEDCALL_2]] @@ -406,14 +406,14 @@ func.func @conv3d_with_bias(%arg0: tensor<1x3x4x3x3xf32>) -> (tensor<1x3x2x3x2xf // CHECK-DAG: %[[CST_1:.*]] = "tf.Const"() {{.*}} : () -> tensor<2xf32> // CHECK: %[[PARTITIONEDCALL_0:.*]] = "tf.PartitionedCall"(%arg0, %[[CST]], %[[CST_1]]) -// CHECK-SAME: {_tfl_quant_trait = "fully_quantizable", -// CHECK-SAME: f = @composite_conv3d_with_bias_and_relu_fn_1} +// CHECK-SAME: f = @composite_conv3d_with_bias_and_relu_fn_1}> +// CHECK-SAME: {_tfl_quant_trait = "fully_quantizable" // CHECK: %[[PARTITIONEDCALL_1:.*]] = "tf.PartitionedCall"(%arg0, %[[CST]], %[[CST_1]]) -// CHECK-SAME: f = @composite_conv3d_with_bias_and_relu6_fn_1} +// CHECK-SAME: f = @composite_conv3d_with_bias_and_relu6_fn_1 // CHECK: %[[PARTITIONEDCALL_2:.*]] = "tf.PartitionedCall"(%arg0, %[[CST]], %[[CST_1]]) -// CHECK-SAME: f = @composite_conv3d_with_bias_fn_1} +// CHECK-SAME: f = @composite_conv3d_with_bias_fn_1 // CHECK: return %[[PARTITIONEDCALL_0]], %[[PARTITIONEDCALL_1]], %[[PARTITIONEDCALL_2]] diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/tests/lift_quantizable_spots_as_functions_drq.mlir b/tensorflow/compiler/mlir/quantization/tensorflow/tests/lift_quantizable_spots_as_functions_drq.mlir index 35215eca54068e..305c6345566ab5 100644 --- a/tensorflow/compiler/mlir/quantization/tensorflow/tests/lift_quantizable_spots_as_functions_drq.mlir +++ b/tensorflow/compiler/mlir/quantization/tensorflow/tests/lift_quantizable_spots_as_functions_drq.mlir @@ -12,10 +12,10 @@ func.func @lift_float_matmul(%arg0: tensor<1x12x12x512xf32>) -> (tensor<*xf32>, } : (tensor<1x12x12x512xf32>, tensor<1x12x12x512xf32>) -> tensor<*xf32> func.return %out_1, %out_2 : tensor<*xf32>, tensor<*xf32> -// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<512x512xf32>} : () -> tensor<512x512xf32> +// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<512x512xf32>}> : () -> tensor<512x512xf32> // CHECK: %[[PARTITIONEDCALL:.*]] = "tf.PartitionedCall"(%arg0, %[[CONST]]) -// CHECK-SAME: {_tfl_quant_trait = "fully_quantizable", -// CHECK-SAME: f = @composite_matmul_fn_1} +// CHECK-SAME: f = @composite_matmul_fn_1}> +// CHECK-SAME: {_tfl_quant_trait = "fully_quantizable" // CHECK: %[[UNQUANTIZED_OUTPUT:.*]] = "tf.MatMul"(%arg0, %arg0) // CHECK: } @@ -45,23 +45,23 @@ func.func @lift_float_conv(%arg0: tensor<1x3x4x3xf32>) -> (tensor<*xf32>, tensor func.return %2, %4 : tensor<*xf32>, tensor<*xf32> -// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<2xf32>} : () -> tensor<2xf32> -// CHECK-DAG: %[[CONST_1:.*]] = "tf.Const"() {value = dense<3.000000e+00> : tensor<2x3x3x2xf32>} : () -> tensor<2x3x3x2xf32> +// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<2xf32>}> : () -> tensor<2xf32> +// CHECK-DAG: %[[CONST_1:.*]] = "tf.Const"() <{value = dense<3.000000e+00> : tensor<2x3x3x2xf32>}> : () -> tensor<2x3x3x2xf32> // CHECK: %[[PARTITIONEDCALL_0:.*]] = "tf.PartitionedCall"(%arg0, %[[CONST_1]]) -// CHECK-SAME: {_tfl_quant_trait = "fully_quantizable", -// CHECK-SAME: f = @composite_conv2d_fn_2} +// CHECK-SAME: f = @composite_conv2d_fn_2}> +// CHECK-SAME: {_tfl_quant_trait = "fully_quantizable" // CHECK: %[[BIASADD_0:.*]] = "tf.BiasAdd"(%[[PARTITIONEDCALL_0]], %[[CONST_0]]) // CHECK: %[[RELU6_0:.*]] = "tf.Relu6"(%[[BIASADD_0]]) // CHECK: %[[PARTITIONEDCALL_1:.*]] = "tf.PartitionedCall"(%arg0, %[[CONST_1]]) -// CHECK-SAME: f = @composite_conv2d_fn_1} +// CHECK-SAME: f = @composite_conv2d_fn_1 // CHECK: %[[BIASADD_1:.*]] = "tf.BiasAdd"(%[[PARTITIONEDCALL_1]], %[[CONST_0]]) // CHECK: return %[[RELU6_0]], %[[BIASADD_1]] // CHECK: } // CHECK-LABEL: private @composite_conv2d_fn_2 // CHECK-NEXT: %[[CONV2D_0:.*]] = "tf.Conv2D"(%arg0, %arg1) +// CHECK-SAME: data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 2, 1], use_cudnn_on_gpu = true // CHECK-SAME: attr_map = "0:strides,1:use_cudnn_on_gpu,2:padding,3:explicit_paddings,4:dilations" -// CHECK-SAME: data_format = "NHWC", device = "", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 2, 1], use_cudnn_on_gpu = true // CHECK-NEXT: return %[[CONV2D_0]] // CHECK-LABEL: private @composite_conv2d_fn_1 @@ -90,7 +90,7 @@ func.func @not_lift_float_conv_with_non_constant_weights(%arg0: tensor<1x3x4x3xf func.return %2, %4 : tensor<*xf32>, tensor<*xf32> -// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<2xf32>} : () -> tensor<2xf32> +// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<2xf32>}> : () -> tensor<2xf32> // CHECK-NOT: %[[PARTITIONEDCALL_0:.*]] = "tf.PartitionedCall"(%arg0, %arg1) // CHECK: %[[CONV2D_0:.*]] = "tf.Conv2D"(%arg0, %arg1) } @@ -115,15 +115,15 @@ func.func @lift_float_depthwise_conv(%arg0: tensor<1x3x4x3xf32>) -> (tensor<*xf3 %4 = "tf.BiasAdd"(%3, %cst) {data_format = "NHWC", device = ""} : (tensor<*xf32>, tensor<2xf32>) -> tensor<*xf32> func.return %2, %4 : tensor<*xf32>, tensor<*xf32> -// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<2xf32>} : () -> tensor<2xf32> -// CHECK-DAG: %[[CONST_1:.*]] = "tf.Const"() {value = dense<3.000000e+00> : tensor<2x3x3x1xf32>} : () -> tensor<2x3x3x1xf32> +// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<2xf32>}> : () -> tensor<2xf32> +// CHECK-DAG: %[[CONST_1:.*]] = "tf.Const"() <{value = dense<3.000000e+00> : tensor<2x3x3x1xf32>}> : () -> tensor<2x3x3x1xf32> // CHECK: %[[PARTITIONEDCALL_0:.*]] = "tf.PartitionedCall"(%arg0, %[[CONST_1]]) -// CHECK-SAME: _tfl_quant_trait = "fully_quantizable", -// CHECK-SAME: f = @composite_depthwise_conv2d_fn_2} +// CHECK-SAME: f = @composite_depthwise_conv2d_fn_2}> +// CHECK-SAME: {_tfl_quant_trait = "fully_quantizable" // CHECK: %[[BIASADD_0:.*]] = "tf.BiasAdd"(%[[PARTITIONEDCALL_0]], %[[CONST_0]]) // CHECK: %[[RELU6_0:.*]] = "tf.Relu6"(%[[BIASADD_0]]) // CHECK: %[[PARTITIONEDCALL_1:.*]] = "tf.PartitionedCall"(%arg0, %[[CONST_1]]) -// CHECK-SAME: f = @composite_depthwise_conv2d_fn_1} +// CHECK-SAME: f = @composite_depthwise_conv2d_fn_1 // CHECK: %[[BIASADD_0:.*]] = "tf.BiasAdd"(%[[PARTITIONEDCALL_1]], %[[CONST_0]]) // CHECK: return %[[RELU6_0]], %[[BIASADD_0]] // CHECK: } @@ -153,8 +153,8 @@ func.func @lift_float_conv3d(%arg0: tensor<1x3x4x3x3xf32>) -> (tensor<1x3x2x3x2x // CHECK-DAG: %[[CST:.*]] = "tf.Const"() {{.*}} : () -> tensor<2x3x3x3x2xf32> // CHECK: %[[PARTITIONEDCALL_0:.*]] = "tf.PartitionedCall"(%arg0, %[[CST]]) -// CHECK-NOT: {_tfl_quant_trait = "fully_quantizable", -// CHECK-SAME: f = @composite_conv3d_fn_1} +// CHECK-SAME: f = @composite_conv3d_fn_1}> +// CHECK-NOT: {_tfl_quant_trait = "fully_quantizable" // CHECK: %[[RELU:.*]] = "tf.Relu"(%[[PARTITIONEDCALL_0]]) // CHECK: return %[[RELU]] @@ -162,8 +162,8 @@ func.func @lift_float_conv3d(%arg0: tensor<1x3x4x3x3xf32>) -> (tensor<1x3x2x3x2x // WEIGHTONLY-DAG: %[[CST:.*]] = "tf.Const"() {{.*}} : () -> tensor<2x3x3x3x2xf32> // WEIGHTONLY: %[[PARTITIONEDCALL_0:.*]] = "tf.PartitionedCall"(%arg0, %[[CST]]) -// WEIGHTONLY: {_tfl_quant_trait = "fully_quantizable", -// WEIGHTONLY-SAME: f = @composite_conv3d_fn_1} +// WEIGHTONLY-SAME: f = @composite_conv3d_fn_1}> +// WEIGHTONLY: {_tfl_quant_trait = "fully_quantizable" // WEIGHTONLY: %[[RELU:.*]] = "tf.Relu"(%[[PARTITIONEDCALL_0]]) // WEIGHTONLY: return %[[RELU]] @@ -181,16 +181,16 @@ func.func @lift_float_batch_matmul(%arg0: tensor<4x4x3xf32>) -> (tensor<4x4x3xf3 // CHECK-DAG: %[[CST:.*]] = "tf.Const"() {{.*}} : () -> tensor<4x3x3xf32> // CHECK: %[[PARTITIONEDCALL_0:.*]] = "tf.PartitionedCall"(%arg0, %[[CST]]) -// CHECK-NOT: {_tfl_quant_trait = "fully_quantizable", -// CHECK-SAME: f = @composite_batch_matmul_fn_1} +// CHECK-SAME: f = @composite_batch_matmul_fn_1}> +// CHECK-NOT: {_tfl_quant_trait = "fully_quantizable" // CHECK: return %[[PARTITIONEDCALL_0]] // CHECK-LABEL: private @composite_batch_matmul_fn_1 // WEIGHTONLY-DAG: %[[CST:.*]] = "tf.Const"() {{.*}} : () -> tensor<4x3x3xf32> // WEIGHTONLY: %[[PARTITIONEDCALL_0:.*]] = "tf.PartitionedCall"(%arg0, %[[CST]]) -// WEIGHTONLY-SAME: {_tfl_quant_trait = "fully_quantizable", -// WEIGHTONLY-SAME: f = @composite_batch_matmul_fn_1} +// WEIGHTONLY-SAME: f = @composite_batch_matmul_fn_1}> +// WEIGHTONLY-SAME: {_tfl_quant_trait = "fully_quantizable" // WEIGHTONLY: return %[[PARTITIONEDCALL_0]] // WEIGHTONLY-LABEL: private @composite_batch_matmul_fn_1 @@ -209,15 +209,15 @@ func.func @lift_float_gather(%arg0: tensor<6xi64>) -> (tensor<6x32xf32>) { // CHECK-DAG: %[[CST:.*]] = "tf.Const"() {{.*}} : () -> tensor // CHECK-DAG: %[[CST_1:.*]] = "tf.Const"() {{.*}} : () -> tensor<128x32xf32> // CHECK: %[[PARTITIONEDCALL_0:.*]] = "tf.PartitionedCall"(%[[CST_1]], %arg0, %[[CST]]) -// CHECK-SAME: {_tfl_quant_trait = "fully_quantizable", -// CHECK-SAME: f = @composite_gather_fn_1} +// CHECK-SAME: f = @composite_gather_fn_1}> +// CHECK-SAME: {_tfl_quant_trait = "fully_quantizable" // CHECK: return %[[PARTITIONEDCALL_0]] // WEIGHTONLY-DAG: %[[CST:.*]] = "tf.Const"() {{.*}} : () -> tensor // WEIGHTONLY-DAG: %[[CST_1:.*]] = "tf.Const"() {{.*}} : () -> tensor<128x32xf32> // WEIGHTONLY: %[[PARTITIONEDCALL_0:.*]] = "tf.PartitionedCall"(%[[CST_1]], %arg0, %[[CST]]) -// WEIGHTONLY-SAME: {_tfl_quant_trait = "fully_quantizable", -// WEIGHTONLY-SAME: f = @composite_gather_fn_1} +// WEIGHTONLY-SAME: f = @composite_gather_fn_1}> +// WEIGHTONLY-SAME: {_tfl_quant_trait = "fully_quantizable" // WEIGHTONLY: return %[[PARTITIONEDCALL_0]] // WEIGHTONLY-LABEL: private @composite_gather_fn_1 diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/tests/lift_quantizable_spots_as_functions_drq_min_elements.mlir b/tensorflow/compiler/mlir/quantization/tensorflow/tests/lift_quantizable_spots_as_functions_drq_min_elements.mlir index b1104490025c6e..83d6b618c44489 100644 --- a/tensorflow/compiler/mlir/quantization/tensorflow/tests/lift_quantizable_spots_as_functions_drq_min_elements.mlir +++ b/tensorflow/compiler/mlir/quantization/tensorflow/tests/lift_quantizable_spots_as_functions_drq_min_elements.mlir @@ -11,11 +11,11 @@ func.func @lift_float_matmul(%arg0: tensor<1x12x12x512xf32>) -> (tensor<*xf32>, } : (tensor<1x12x12x512xf32>, tensor<1x12x12x512xf32>) -> tensor<*xf32> func.return %out_1, %out_2 : tensor<*xf32>, tensor<*xf32> -// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<512x512xf32>} : () -> tensor<512x512xf32> +// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<512x512xf32>}> : () -> tensor<512x512xf32> // CHECK: %[[PARTITIONEDCALL:.*]] = "tf.PartitionedCall"(%arg0, %[[CONST]]) -// CHECK-NOT: {_tfl_quant_trait = "fully_quantizable", -// CHECK-SAME: {config = "", -// CHECK-SAME: f = @composite_matmul_fn_1} +// CHECK-SAME: <{config = "", +// CHECK-SAME: f = @composite_matmul_fn_1}> +// CHECK-NOT: {_tfl_quant_trait = "fully_quantizable" // CHECK: %[[UNQUANTIZED_OUTPUT:.*]] = "tf.MatMul"(%arg0, %arg0) // CHECK: } @@ -33,9 +33,9 @@ func.func @not_lift_float_conv(%arg0: tensor<1x3x4x512xf32>) -> (tensor<*xf32>) } : (tensor<1x3x4x512xf32>, tensor<2x3x512x512xf32>) -> tensor<*xf32> func.return %0 : tensor<*xf32> -// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() {value = dense<3.000000e+00> : tensor<2x3x512x512xf32>} : () -> tensor<2x3x512x512xf32> +// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() <{value = dense<3.000000e+00> : tensor<2x3x512x512xf32>}> : () -> tensor<2x3x512x512xf32> // CHECK: %[[PARTITIONEDCALL:.*]] = "tf.PartitionedCall"(%arg0, %[[CONST]]) -// CHECK-NOT: {_tfl_quant_trait = "fully_quantizable", -// CHECK-SAME: {config = "", -// CHECK-SAME: f = @composite_conv2d_fn_1} +// CHECK-SAME: <{config = "", +// CHECK-SAME: f = @composite_conv2d_fn_1}> +// CHECK-NOT: {_tfl_quant_trait = "fully_quantizable" } diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/tests/lift_quantizable_spots_as_functions_xla.mlir b/tensorflow/compiler/mlir/quantization/tensorflow/tests/lift_quantizable_spots_as_functions_xla.mlir index 1d80a199d2e3ba..38911e2b70eb32 100644 --- a/tensorflow/compiler/mlir/quantization/tensorflow/tests/lift_quantizable_spots_as_functions_xla.mlir +++ b/tensorflow/compiler/mlir/quantization/tensorflow/tests/lift_quantizable_spots_as_functions_xla.mlir @@ -14,15 +14,15 @@ module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, p // CHECK-LABEL: func @depthwise_conv // CHECK: "tf.PartitionedCall" +// CHECK-SAME: f = @composite_depthwise_conv2d_with_bias_and_relu6_fn_1 // Check that the `_tfl_quant_trait` attribute has been removed. // CHECK-NOT: _tfl_quant_trait = "fully_quantizable" -// CHECK-SAME: f = @composite_depthwise_conv2d_with_bias_and_relu6_fn_1 // CHECK-LABEL: private @composite_depthwise_conv2d_with_bias_and_relu6_fn_1 // CHECK: %[[DEPTHWISECONV2D_0:.*]] = "tf.DepthwiseConv2dNative"(%arg0, %arg1) +// CHECK-SAME: <{data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 2, 2, 1]}> // Check that the `attr_map` attribute has been removed. // CHECK-NOT: attr_map -// CHECK-SAME: {data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 2, 2, 1]} // ----- @@ -35,17 +35,17 @@ func.func @conv_with_non_constant_filter(%arg0: tensor<1x3x4x3xf32>, %arg1: tens } // CHECK-LABEL: func @conv_with_non_constant_filter -// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<2xf32>} : () -> tensor<2xf32> +// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<2xf32>}> : () -> tensor<2xf32> // CHECK: %[[PARTITIONEDCALL_0:.*]] = "tf.PartitionedCall"(%arg0, %arg1, %[[CONST_0]]) +// CHECK-SAME: f = @composite_conv2d_with_bias_and_relu6_fn_1 // Check that the `_tfl_quant_trait` attribute has been removed. // CHECK-NOT: _tfl_quant_trait = "fully_quantizable" -// CHECK-SAME: f = @composite_conv2d_with_bias_and_relu6_fn_1 // CHECK-LABEL: func private @composite_conv2d_with_bias_and_relu6_fn_1 // CHECK: %[[CONV2D_0:.*]] = "tf.Conv2D"(%arg0, %arg1) +// CHECK-SAME: data_format = "NHWC", dilations = [1, 1, 2, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 2, 1] // Check that the `attr_map` attribute has been removed. // CHECK-NOT: attr_map -// CHECK-SAME: data_format = "NHWC", dilations = [1, 1, 2, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 2, 1] // ----- @@ -59,18 +59,18 @@ func.func @conv_with_dynamic_channel_dim(%arg0: tensor<1x3x4x?xf32>) -> tensor<* } // CHECK-LABEL: func @conv_with_dynamic_channel_dim -// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<1xf32>} : () -> tensor<1xf32> +// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<1xf32>}> : () -> tensor<1xf32> // CHECK-DAG: %[[CONST_1:.*]] = "tf.Const"() {{.*}} : () -> tensor<2x3x3x1xf32> // CHECK: %[[PARTITIONEDCALL_0:.*]] = "tf.PartitionedCall"(%arg0, %[[CONST_1]], %[[CONST_0]]) +// CHECK-SAME: f = @composite_conv2d_with_bias_and_relu6_fn_1 // Check that the `_tfl_quant_trait` attribute has been removed. // CHECK-NOT: _tfl_quant_trait = "fully_quantizable" -// CHECK-SAME: f = @composite_conv2d_with_bias_and_relu6_fn_1 // CHECK-LABEL: func private @composite_conv2d_with_bias_and_relu6_fn_1 // CHECK: %[[CONV2D_0:.*]] = "tf.Conv2D"(%arg0, %arg1) +// CHECK-SAME: data_format = "NHWC", dilations = [1, 1, 2, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 2, 1] // Check that the `attr_map` attribute has been removed. // CHECK-NOT: attr_map -// CHECK-SAME: data_format = "NHWC", dilations = [1, 1, 2, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 2, 1] // ----- @@ -93,11 +93,11 @@ func.func @const_filter_with_q_dq(%arg0: tensor<1x3x4x3xf32>) -> (tensor<1x3x2x2 // CHECK-LABEL: func @const_filter_with_q_dq // CHECK-DAG: %[[WEIGHT:.*]] = "tf.Const"() {{.*}} : () -> tensor<2x3x3x2xf32> -// CHECK-DAG: %[[BIAS:.*]] = "tf.Const"() {device = "", value = dense<[1.000000e-01, 2.000000e-01]> : tensor<2xf32>} +// CHECK-DAG: %[[BIAS:.*]] = "tf.Const"() <{value = dense<[1.000000e-01, 2.000000e-01]> : tensor<2xf32>}> {device = ""} // CHECK: %[[Q_W:.*]] = "quantfork.qcast"(%[[WEIGHT]]) // CHECK: %[[DQ_W:.*]] = "quantfork.dcast"(%[[Q_W]]) // CHECK: %[[PARTITIONEDCALL_0:.*]] = "tf.PartitionedCall"({{.*}}, %[[DQ_W]], %[[BIAS]]) -// CHECK-SAME: _tfl_quant_trait = "fully_quantizable" // CHECK-SAME: f = @composite_conv2d_with_bias_and_relu_fn_1 +// CHECK-SAME: _tfl_quant_trait = "fully_quantizable" -// CHECK-LABEL: func private @composite_conv2d_with_bias_and_relu_fn_1 \ No newline at end of file +// CHECK-LABEL: func private @composite_conv2d_with_bias_and_relu_fn_1 diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/tests/lift_quantizable_spots_as_functions_xla_selective_quantization.mlir b/tensorflow/compiler/mlir/quantization/tensorflow/tests/lift_quantizable_spots_as_functions_xla_selective_quantization.mlir index 8dfd2815692fc6..a2a86a2c112664 100644 --- a/tensorflow/compiler/mlir/quantization/tensorflow/tests/lift_quantizable_spots_as_functions_xla_selective_quantization.mlir +++ b/tensorflow/compiler/mlir/quantization/tensorflow/tests/lift_quantizable_spots_as_functions_xla_selective_quantization.mlir @@ -16,9 +16,9 @@ module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, p // CHECK-LABEL: func @conv2d_unmatching_unit // CHECK: "tf.PartitionedCall" +// CHECK-SAME: f = @composite_conv2d_fn_1 // Check that the `_tfl_quant_trait` attribute exists since the unit is not in `unit_wise_quantization_specs`. // CHECK-SAME: _tfl_quant_trait = "fully_quantizable" -// CHECK-SAME: f = @composite_conv2d_fn_1 // CHECK-SAME: loc(callsite("Model/conv2d@conv2d_unmatching_unit"("Conv2D") at "QuantizationUnit({{.*}})")) // ----- @@ -36,9 +36,9 @@ module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, p // CHECK-LABEL: func @conv2d_disable_quantization // CHECK: "tf.PartitionedCall" +// CHECK-SAME: f = @composite_conv2d_fn_1 // Check that quantization is disabled for this unit. // CHECK-NOT: _tfl_quant_trait = "fully_quantizable" -// CHECK-SAME: f = @composite_conv2d_fn_1 // CHECK-SAME: loc(callsite("test_opt_out@conv2d_disable_quantization"("Conv2D") at "QuantizationUnit({{.*}})")) // ----- @@ -58,9 +58,9 @@ module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, p // CHECK-LABEL: func @conv2d_with_bias_disable_quantization // CHECK: "tf.PartitionedCall" +// CHECK-SAME: f = @composite_conv2d_with_bias_fn_1 // Check that quantization is disabled for this unit. // CHECK-NOT: _tfl_quant_trait = "fully_quantizable" -// CHECK-SAME: f = @composite_conv2d_with_bias_fn_1 // CHECK-SAME: loc(callsite("test_opt_out@conv2d_with_bias_disable_quantization"("Conv2D") at "QuantizationUnit({{.*}})")) // ----- @@ -80,9 +80,9 @@ module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, p // CHECK-LABEL: func @matmul_with_reshape_disable_quantization // CHECK: "tf.PartitionedCall" +// CHECK-SAME: f = @composite_matmul_with_reshape_and_bias_fn_1 // Check that quantization is disabled for this unit. // CHECK-NOT: _tfl_quant_trait = "fully_quantizable" -// CHECK-SAME: f = @composite_matmul_with_reshape_and_bias_fn_1 // CHECK-SAME: loc(callsite("test_opt_out@matmul_with_reshape_disable_quantization"("MatMul") at "QuantizationUnit({{.*}})")) // ----- @@ -105,8 +105,8 @@ module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, p // CHECK-LABEL: func @serving_default // CHECK: "tf.PartitionedCall" +// CHECK-SAME: f = @composite_conv2d_fn_1 // Check that quantization is disabled for this unit. // CHECK-NOT: _tfl_quant_trait = "fully_quantizable" -// CHECK-SAME: f = @composite_conv2d_fn_1 // CHECK-SAME: loc(callsite("test_opt_out@conv2d_with_inliner"("Conv2D") at "QuantizationUnit({{.*}})")) } diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/tests/merge_initializer_function_ops_to_main.mlir b/tensorflow/compiler/mlir/quantization/tensorflow/tests/merge_initializer_function_ops_to_main.mlir index 0ca4745b730aa3..1b0e3e286f1f8d 100644 --- a/tensorflow/compiler/mlir/quantization/tensorflow/tests/merge_initializer_function_ops_to_main.mlir +++ b/tensorflow/compiler/mlir/quantization/tensorflow/tests/merge_initializer_function_ops_to_main.mlir @@ -146,14 +146,14 @@ module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, p // CHECK-NEXT: %[[OUT:.*]], %[[CTL:.*]] = tf_executor.island wraps "tf.PartitionedCall"(%[[ARG]]) // CHECK-SAME: f = @serving_default // Checks that the contents of @NoOp are copied here. -// CHECK-DAG: %[[OUT_0:.*]], %[[CTL_0:.*]] = tf_executor.island wraps "tf.Const"() {{{.*value = dense<"test_1">.*}}} -// CHECK-DAG: %[[OUT_1:.*]], %[[CTL_1:.*]] = tf_executor.island wraps "tf.Const"() {{{.*value = dense<1>.*}}} +// CHECK-DAG: %[[OUT_0:.*]], %[[CTL_0:.*]] = tf_executor.island wraps "tf.Const"() <{{{.*value = dense<"test_1">.*}}}> +// CHECK-DAG: %[[OUT_1:.*]], %[[CTL_1:.*]] = tf_executor.island wraps "tf.Const"() <{{{.*value = dense<1>.*}}}> // CHECK-NEXT: %[[OUT_2:.*]], %[[CTL_2:.*]] = tf_executor.island wraps "tf.HashTableV2"() // CHECK-NEXT: %[[CTL_3:.*]] = tf_executor.island wraps "tf.LookupTableImportV2"(%[[OUT_2]], %[[OUT_0]], %[[OUT_1]]) -// CHECK-DAG: %[[OUT_3:.*]], %[[CTL_4:.*]] = tf_executor.island wraps "tf.Const"() {{{.*value = dense<"test_2">.*}}} -// CHECK-DAG: %[[OUT_4:.*]], %[[CTL_5:.*]] = tf_executor.island wraps "tf.Const"() {{{.*value = dense<2>.*}}} +// CHECK-DAG: %[[OUT_3:.*]], %[[CTL_4:.*]] = tf_executor.island wraps "tf.Const"() <{{{.*value = dense<"test_2">.*}}}> +// CHECK-DAG: %[[OUT_4:.*]], %[[CTL_5:.*]] = tf_executor.island wraps "tf.Const"() <{{{.*value = dense<2>.*}}}> // CHECK-NEXT: %[[OUT_5:.*]], %[[CTL_6:.*]] = tf_executor.island(%[[CTL_3]]) wraps "tf.HashTableV2"() // CHECK-NEXT: %[[CTL_7:.*]] = tf_executor.island wraps "tf.LookupTableImportV2"(%[[OUT_5]], %[[OUT_3]], %[[OUT_4]]) @@ -330,9 +330,9 @@ module attributes {tf_saved_model.semantics} { // CHECK-NEXT: tf_executor.graph // Checks that the ops from @init_func_restore_op are cloned. -// CHECK-DAG: %[[CONST_0:.*]], %[[CTL:.*]] = tf_executor.island wraps "tf.Const"() {{{.*value = dense<""> : tensor<1x!tf_type\.string>.*}}} -// CHECK-DAG: %[[CONST_1:.*]], %[[CTL_0:.*]] = tf_executor.island wraps "tf.Const"() {{{.*value = dense<"var_0"> : tensor<1x!tf_type\.string>.*}}} -// CHECK: %[[VAR_HANDLE:.*]], %[[CTL_1:.*]] = tf_executor.island wraps "tf.VarHandleOp"() {{{.*shared_name = "var_0".*}}} +// CHECK-DAG: %[[CONST_0:.*]], %[[CTL:.*]] = tf_executor.island wraps "tf.Const"() <{{{.*value = dense<""> : tensor<1x!tf_type\.string>.*}}}> +// CHECK-DAG: %[[CONST_1:.*]], %[[CTL_0:.*]] = tf_executor.island wraps "tf.Const"() <{{{.*value = dense<"var_0"> : tensor<1x!tf_type\.string>.*}}}> +// CHECK: %[[VAR_HANDLE:.*]], %[[CTL_1:.*]] = tf_executor.island wraps "tf.VarHandleOp"() <{{{.*shared_name = "var_0".*}}}> // CHECK: %[[RESTORE:.*]], %[[CTL_2:.*]] = tf_executor.island wraps "tf.RestoreV2"(%[[ARG]], %[[CONST_1]], %[[CONST_0]]) // CHECK: %[[CTL_3:.*]] = tf_executor.island wraps "tf.AssignVariableOp"(%[[VAR_HANDLE]], %[[RESTORE]]) // CHECK: %[[CTL_4:.*]] = tf_executor.island(%[[CTL_3]]) wraps "tf.NoOp"() diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/tests/merge_save_function_ops_to_main.mlir b/tensorflow/compiler/mlir/quantization/tensorflow/tests/merge_save_function_ops_to_main.mlir index 5341e15c798161..bc0b283b808831 100644 --- a/tensorflow/compiler/mlir/quantization/tensorflow/tests/merge_save_function_ops_to_main.mlir +++ b/tensorflow/compiler/mlir/quantization/tensorflow/tests/merge_save_function_ops_to_main.mlir @@ -32,10 +32,10 @@ module attributes {tf_saved_model.semantics} { // CHECK: func.func @main // CHECK-SAME: %[[ARG_0:.*]]: tensor {tf_saved_model.index_path = ["__tf_file_prefix"]} // CHECK: tf_executor.graph -// CHECK: %[[VAR_HANDLE:.*]], {{.*}} = tf_executor.island wraps "tf.VarHandleOp"() {{{.*shared_name = "var_0".*}}} +// CHECK: %[[VAR_HANDLE:.*]], {{.*}} = tf_executor.island wraps "tf.VarHandleOp"() <{{{.*shared_name = "var_0".*}}}> // CHECK: %[[READ_VARIABLE:.*]], {{.*}} = tf_executor.island wraps "tf.ReadVariableOp"(%[[VAR_HANDLE]]) -// CHECK-DAG: %[[CST_0:.*]], {{.*}} = tf_executor.island wraps "tf.Const"() {{{.*value = dense<"var_0"> : tensor<1x!tf_type\.string>.*}}} -// CHECK-DAG: %[[CST_1:.*]], {{.*}} = tf_executor.island wraps "tf.Const"() {{{.*value = dense<""> : tensor<1x!tf_type\.string>.*}}} +// CHECK-DAG: %[[CST_0:.*]], {{.*}} = tf_executor.island wraps "tf.Const"() <{{{.*value = dense<"var_0"> : tensor<1x!tf_type\.string>.*}}}> +// CHECK-DAG: %[[CST_1:.*]], {{.*}} = tf_executor.island wraps "tf.Const"() <{{{.*value = dense<""> : tensor<1x!tf_type\.string>.*}}}> // CHECK: %[[CTL_0:.*]] = tf_executor.island wraps "tf.SaveV2"(%[[ARG_0]], %[[CST_0]], %[[CST_1]], %[[READ_VARIABLE]]) : (tensor, tensor<1x!tf_type.string>, tensor<1x!tf_type.string>, tensor<2xf32>) -> () // Test that the Identity op has been created to fetch the file prefix @@ -150,10 +150,10 @@ module attributes {tf_saved_model.semantics} { // CHECK-SAME: %[[ARG_0:.*]]: tensor {tf_saved_model.index_path = ["__tf_file_prefix"]} // CHECK-SAME: tf.entry_function = {inputs = "__tf_file_prefix:0", outputs = ""} // CHECK: tf_executor.graph -// CHECK: %[[VAR_HANDLE:.*]], {{.*}} = tf_executor.island wraps "tf.VarHandleOp"() {{{.*shared_name = "var_0".*}}} +// CHECK: %[[VAR_HANDLE:.*]], {{.*}} = tf_executor.island wraps "tf.VarHandleOp"() <{{{.*shared_name = "var_0".*}}}> // CHECK: %[[READ_VARIABLE:.*]], {{.*}} = tf_executor.island wraps "tf.ReadVariableOp"(%[[VAR_HANDLE]]) -// CHECK-DAG: %[[CST_0:.*]], {{.*}} = tf_executor.island wraps "tf.Const"() {{{.*value = dense<"var_0"> : tensor<1x!tf_type\.string>.*}}} -// CHECK-DAG: %[[CST_1:.*]], {{.*}} = tf_executor.island wraps "tf.Const"() {{{.*value = dense<""> : tensor<1x!tf_type\.string>.*}}} +// CHECK-DAG: %[[CST_0:.*]], {{.*}} = tf_executor.island wraps "tf.Const"() <{{{.*value = dense<"var_0"> : tensor<1x!tf_type\.string>.*}}}> +// CHECK-DAG: %[[CST_1:.*]], {{.*}} = tf_executor.island wraps "tf.Const"() <{{{.*value = dense<""> : tensor<1x!tf_type\.string>.*}}}> // CHECK: %[[CTL_0:.*]] = tf_executor.island wraps "tf.SaveV2"(%[[ARG_0]], %[[CST_0]], %[[CST_1]], %[[READ_VARIABLE]]) : (tensor, tensor<1x!tf_type.string>, tensor<1x!tf_type.string>, tensor<2xf32>) -> () // Test that the Identity op has been created to fetch the file prefix diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/tests/optimize.mlir b/tensorflow/compiler/mlir/quantization/tensorflow/tests/optimize.mlir index 2a69627bb40c96..48a5c43a5bcaa2 100644 --- a/tensorflow/compiler/mlir/quantization/tensorflow/tests/optimize.mlir +++ b/tensorflow/compiler/mlir/quantization/tensorflow/tests/optimize.mlir @@ -56,11 +56,11 @@ func.func @remove_redundant_cast(%arg0: tensor<1x100x100x1xf32>) -> (tensor<1x96 // CHECK: %[[CLIPBYVALUE_0:.*]] = "tf.ClipByValue" // CHECK-SAME: (tensor<1x100x100x1xi32>, tensor, tensor) -> tensor<1x100x100x1xi32> -// CHECK: %[[CAST_1:.*]] = "tf.Cast"(%[[CLIPBYVALUE_0]]) {Truncate = false} : (tensor<1x100x100x1xi32>) -> tensor<1x100x100x1xf32> +// CHECK: %[[CAST_1:.*]] = "tf.Cast"(%[[CLIPBYVALUE_0]]) <{Truncate = false}> : (tensor<1x100x100x1xi32>) -> tensor<1x100x100x1xf32> // CHECK: %[[CLIPBYVALUE_1:.*]] = "tf.ClipByValue" // CHECK-SAME: (tensor<1x98x98x1xi32>, tensor, tensor) -> tensor<1x98x98x1xi32> -// CHECK: %[[CAST_3:.*]] = "tf.Cast"(%[[CLIPBYVALUE_1]]) {Truncate = false} : (tensor<1x98x98x1xi32>) -> tensor<1x98x98x1xf32> +// CHECK: %[[CAST_3:.*]] = "tf.Cast"(%[[CLIPBYVALUE_1]]) <{Truncate = false}> : (tensor<1x98x98x1xi32>) -> tensor<1x98x98x1xf32> // CHECK: %[[CLIPBYVALUE_2:.*]] = "tf.ClipByValue" // CHECK-SAME: (tensor<1x96x96x1xi32>, tensor, tensor) -> tensor<1x96x96x1xi32> @@ -76,7 +76,7 @@ func.func @consecutive_add_add(%arg0: tensor) -> (tensor) { // CHECK-LABEL: func.func @consecutive_add_add -// CHECK: %[[CST:.*]] = "tf.Const"() {value = dense<-30> : tensor} : () -> tensor +// CHECK: %[[CST:.*]] = "tf.Const"() <{value = dense<-30> : tensor}> : () -> tensor // CHECK: %[[ADD:.*]] = "tf.AddV2"(%arg0, %[[CST]]) : (tensor, tensor) -> tensor // CHECK: return %[[ADD]] : tensor } @@ -90,7 +90,7 @@ func.func @consecutive_add_sub(%arg0: tensor) -> (tensor) { // CHECK-LABEL: func.func @consecutive_add_sub -// CHECK: %[[CST:.*]] = "tf.Const"() {value = dense<6> : tensor} : () -> tensor +// CHECK: %[[CST:.*]] = "tf.Const"() <{value = dense<6> : tensor}> : () -> tensor // CHECK: %[[SUB:.*]] = "tf.Sub"(%arg0, %[[CST]]) : (tensor, tensor) -> tensor // CHECK: return %[[SUB]] : tensor } @@ -104,7 +104,7 @@ func.func @consecutive_sub_add(%arg0: tensor) -> (tensor) { // CHECK-LABEL: func.func @consecutive_sub_add -// CHECK: %[[CST:.*]] = "tf.Const"() {value = dense<6> : tensor} : () -> tensor +// CHECK: %[[CST:.*]] = "tf.Const"() <{value = dense<6> : tensor}> : () -> tensor // CHECK: %[[ADD:.*]] = "tf.AddV2"(%arg0, %[[CST]]) : (tensor, tensor) -> tensor // CHECK: return %[[ADD]] : tensor } @@ -118,7 +118,7 @@ func.func @consecutive_sub_sub(%arg0: tensor) -> (tensor) { // CHECK-LABEL: func.func @consecutive_sub_sub -// CHECK: %[[CST:.*]] = "tf.Const"() {value = dense<-30> : tensor} : () -> tensor +// CHECK: %[[CST:.*]] = "tf.Const"() <{value = dense<-30> : tensor}> : () -> tensor // CHECK: %[[SUB:.*]] = "tf.Sub"(%arg0, %[[CST]]) : (tensor, tensor) -> tensor // CHECK: return %[[SUB]] : tensor } diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/tests/prepare_lifting.mlir b/tensorflow/compiler/mlir/quantization/tensorflow/tests/prepare_lifting.mlir index c99fed3d43a6f6..1e771e2586a61e 100644 --- a/tensorflow/compiler/mlir/quantization/tensorflow/tests/prepare_lifting.mlir +++ b/tensorflow/compiler/mlir/quantization/tensorflow/tests/prepare_lifting.mlir @@ -7,8 +7,8 @@ func.func @decompose_batch_norm(%arg0: tensor<*xf32>) -> (tensor<*xf32>) { func.return %add : tensor<*xf32> } // CHECK: func @decompose_batch_norm -// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() {value = dense<2.49743462E-5> : tensor<2xf32>} : () -> tensor<2xf32> -// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() {value = dense<0.999950051> : tensor<2xf32>} : () -> tensor<2xf32> +// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() <{value = dense<2.49743462E-5> : tensor<2xf32>}> : () -> tensor<2xf32> +// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() <{value = dense<0.999950051> : tensor<2xf32>}> : () -> tensor<2xf32> // CHECK: %[[mul:.*]] = "tf.Mul"(%arg0, %[[CONST_0]]) : (tensor<*xf32>, tensor<2xf32>) -> tensor<*xf32> // CHECK: %[[add:.*]] = "tf.AddV2"(%[[mul]], %[[CONST]]) : (tensor<*xf32>, tensor<2xf32>) -> tensor<*xf32> // CHECK-NEXT: return %[[add]] : tensor<*xf32> @@ -22,9 +22,9 @@ func.func @not_decompose_batch_norm(%arg0: tensor<*xf32>) -> (tensor<*xf32>) { func.return %bn : tensor<*xf32> } // CHECK: func @not_decompose_batch_norm -// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() {value = dense<1.000000e+00> : tensor<2xf32>} : () -> tensor<2xf32> -// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() {value = dense<5.000000e-01> : tensor<2xf32>} : () -> tensor<2xf32> -// CHECK: %[[bn:.*]], %batch_mean, %batch_variance, %reserve_space_1, %reserve_space_2, %reserve_space_3 = "tf.FusedBatchNormV3"(%arg0, %[[CONST]], %[[CONST_0]], %[[CONST_0]], %[[CONST]]) {data_format = "NHWC", device = "", epsilon = 9.99999974E-5 : f32, exponential_avg_factor = 1.000000e+00 : f32, is_training = true} : (tensor<*xf32>, tensor<2xf32>, tensor<2xf32>, tensor<2xf32>, tensor<2xf32>) -> (tensor<*xf32>, tensor<*xf32>, tensor<*xf32>, tensor<*xf32>, tensor<*xf32>, tensor<*xf32>) +// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() <{value = dense<1.000000e+00> : tensor<2xf32>}> : () -> tensor<2xf32> +// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() <{value = dense<5.000000e-01> : tensor<2xf32>}> : () -> tensor<2xf32> +// CHECK: %[[bn:.*]], %batch_mean, %batch_variance, %reserve_space_1, %reserve_space_2, %reserve_space_3 = "tf.FusedBatchNormV3"(%arg0, %[[CONST]], %[[CONST_0]], %[[CONST_0]], %[[CONST]]) <{data_format = "NHWC", epsilon = 9.99999974E-5 : f32, exponential_avg_factor = 1.000000e+00 : f32, is_training = true}> {device = ""} : (tensor<*xf32>, tensor<2xf32>, tensor<2xf32>, tensor<2xf32>, tensor<2xf32>) -> (tensor<*xf32>, tensor<*xf32>, tensor<*xf32>, tensor<*xf32>, tensor<*xf32>, tensor<*xf32>) // CHECK-NEXT: return %[[bn]] : tensor<*xf32> // ----- @@ -37,10 +37,10 @@ func.func @convert_add_to_biasadd(%arg0: tensor<1x3x4x3xf32>) -> (tensor<1x3x2x2 func.return %1 : tensor<1x3x2x2xf32> } // CHECK: func @convert_add_to_biasadd -// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() {value = dense<1.000000e+00> : tensor<2x3x3x2xf32>} : () -> tensor<2x3x3x2xf32> -// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() {value = dense<5.000000e-01> : tensor<2xf32>} : () -> tensor<2xf32> -// CHECK: %[[CONV2D:.*]] = "tf.Conv2D"(%arg0, %[[CONST]]) {data_format = "NHWC", dilations = [1, 1, 2, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 2, 1], use_cudnn_on_gpu = true} : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x3x2x2xf32> -// CHECK-NEXT: %[[BIASADD:.*]] = "tf.BiasAdd"(%[[CONV2D]], %[[CONST_0]]) {data_format = "NHWC"} : (tensor<1x3x2x2xf32>, tensor<2xf32>) -> tensor<1x3x2x2xf32> +// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() <{value = dense<1.000000e+00> : tensor<2x3x3x2xf32>}> : () -> tensor<2x3x3x2xf32> +// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() <{value = dense<5.000000e-01> : tensor<2xf32>}> : () -> tensor<2xf32> +// CHECK: %[[CONV2D:.*]] = "tf.Conv2D"(%arg0, %[[CONST]]) <{data_format = "NHWC", dilations = [1, 1, 2, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 2, 1], use_cudnn_on_gpu = true}> : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x3x2x2xf32> +// CHECK-NEXT: %[[BIASADD:.*]] = "tf.BiasAdd"(%[[CONV2D]], %[[CONST_0]]) <{data_format = "NHWC"}> : (tensor<1x3x2x2xf32>, tensor<2xf32>) -> tensor<1x3x2x2xf32> // CHECK-NEXT: return %[[BIASADD]] : tensor<1x3x2x2xf32> // ----- @@ -53,9 +53,9 @@ func.func @not_convert_add_to_biasadd(%arg0: tensor<1x3x4x3xf32>) -> (tensor<1x3 func.return %1 : tensor<1x3x2x3xf32> } // CHECK: func @not_convert_add_to_biasadd -// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() {value = dense<1.000000e+00> : tensor<2x3x3x3xf32>} : () -> tensor<2x3x3x3xf32> -// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() {value = dense<5.000000e-01> : tensor<1x3x2x3xf32>} : () -> tensor<1x3x2x3xf32> -// CHECK-NEXT: %[[CONV2D:.*]] = "tf.Conv2D"(%arg0, %[[CONST]]) {data_format = "NHWC", dilations = [1, 1, 2, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 2, 1], use_cudnn_on_gpu = true} : (tensor<1x3x4x3xf32>, tensor<2x3x3x3xf32>) -> tensor<1x3x2x3xf32> +// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() <{value = dense<1.000000e+00> : tensor<2x3x3x3xf32>}> : () -> tensor<2x3x3x3xf32> +// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() <{value = dense<5.000000e-01> : tensor<1x3x2x3xf32>}> : () -> tensor<1x3x2x3xf32> +// CHECK-NEXT: %[[CONV2D:.*]] = "tf.Conv2D"(%arg0, %[[CONST]]) <{data_format = "NHWC", dilations = [1, 1, 2, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 2, 1], use_cudnn_on_gpu = true}> : (tensor<1x3x4x3xf32>, tensor<2x3x3x3xf32>) -> tensor<1x3x2x3xf32> // CHECK-NEXT: %[[ADD:.*]] = "tf.AddV2"(%[[CONV2D]], %[[CONST_0]]) : (tensor<1x3x2x3xf32>, tensor<1x3x2x3xf32>) -> tensor<1x3x2x3xf32> // CHECK-NEXT: return %[[ADD]] : tensor<1x3x2x3xf32> @@ -69,8 +69,8 @@ func.func @fuse_conv2d_and_mul(%arg0: tensor<1x3x4x3xf32>) -> (tensor<1x3x2x2xf3 func.return %1 : tensor<1x3x2x2xf32> } // CHECK: func @fuse_conv2d_and_mul -// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() {value = dense<8.000000e-01> : tensor<2x3x3x2xf32>} : () -> tensor<2x3x3x2xf32> -// CHECK-NEXT: %[[CONV2D:.*]] = "tf.Conv2D"(%arg0, %[[CONST]]) {data_format = "NHWC", dilations = [1, 1, 2, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 2, 1], use_cudnn_on_gpu = true} : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x3x2x2xf32> +// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() <{value = dense<8.000000e-01> : tensor<2x3x3x2xf32>}> : () -> tensor<2x3x3x2xf32> +// CHECK-NEXT: %[[CONV2D:.*]] = "tf.Conv2D"(%arg0, %[[CONST]]) <{data_format = "NHWC", dilations = [1, 1, 2, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 2, 1], use_cudnn_on_gpu = true}> : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x3x2x2xf32> // CHECK-NEXT: return %[[CONV2D]] : tensor<1x3x2x2xf32> // ----- @@ -83,9 +83,9 @@ func.func @not_fuse_conv2d_and_mul(%arg0: tensor<1x3x4x3xf32>) -> (tensor<1x3x2x func.return %1 : tensor<1x3x2x2xf32> } // CHECK: func @not_fuse_conv2d_and_mul -// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() {value = dense<2.000000e+00> : tensor<2x3x3x2xf32>} : () -> tensor<2x3x3x2xf32> -// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() {value = dense<4.000000e-01> : tensor<2x2xf32>} : () -> tensor<2x2xf32> -// CHECK-NEXT: %[[CONV2D:.*]] = "tf.Conv2D"(%arg0, %[[CONST]]) {data_format = "NHWC", dilations = [1, 1, 2, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 2, 1], use_cudnn_on_gpu = true} : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x3x2x2xf32> +// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() <{value = dense<2.000000e+00> : tensor<2x3x3x2xf32>}> : () -> tensor<2x3x3x2xf32> +// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() <{value = dense<4.000000e-01> : tensor<2x2xf32>}> : () -> tensor<2x2xf32> +// CHECK-NEXT: %[[CONV2D:.*]] = "tf.Conv2D"(%arg0, %[[CONST]]) <{data_format = "NHWC", dilations = [1, 1, 2, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 2, 1], use_cudnn_on_gpu = true}> : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x3x2x2xf32> // CHECK-NEXT: %[[ADD:.*]] = "tf.Mul"(%[[CONV2D]], %[[CONST_0]]) : (tensor<1x3x2x2xf32>, tensor<2x2xf32>) -> tensor<1x3x2x2xf32> // CHECK-NEXT: return %[[ADD]] : tensor<1x3x2x2xf32> @@ -101,10 +101,10 @@ func.func @fuse_conv2d_with_bias_and_mul(%arg0: tensor<1x3x4x3xf32>) -> (tensor< func.return %2 : tensor<1x3x2x2xf32> } // CHECK: func @fuse_conv2d_with_bias_and_mul -// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() {value = dense<1.000000e+00> : tensor<2x3x3x2xf32>} : () -> tensor<2x3x3x2xf32> -// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() {value = dense<2.000000e-01> : tensor<2xf32>} : () -> tensor<2xf32> -// CHECK-NEXT: %[[CONV2D:.*]] = "tf.Conv2D"(%arg0, %[[CONST]]) {data_format = "NHWC", dilations = [1, 1, 2, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 2, 1], use_cudnn_on_gpu = true} : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x3x2x2xf32> -// CHECK-NEXT: %[[BIASADD:.*]] = "tf.BiasAdd"(%[[CONV2D]], %[[CONST_0]]) {data_format = "NHWC"} : (tensor<1x3x2x2xf32>, tensor<2xf32>) -> tensor<1x3x2x2xf32> +// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() <{value = dense<1.000000e+00> : tensor<2x3x3x2xf32>}> : () -> tensor<2x3x3x2xf32> +// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() <{value = dense<2.000000e-01> : tensor<2xf32>}> : () -> tensor<2xf32> +// CHECK-NEXT: %[[CONV2D:.*]] = "tf.Conv2D"(%arg0, %[[CONST]]) <{data_format = "NHWC", dilations = [1, 1, 2, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 2, 1], use_cudnn_on_gpu = true}> : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x3x2x2xf32> +// CHECK-NEXT: %[[BIASADD:.*]] = "tf.BiasAdd"(%[[CONV2D]], %[[CONST_0]]) <{data_format = "NHWC"}> : (tensor<1x3x2x2xf32>, tensor<2xf32>) -> tensor<1x3x2x2xf32> // CHECK-NEXT: return %[[BIASADD]] : tensor<1x3x2x2xf32> // ----- @@ -119,11 +119,11 @@ func.func @not_fuse_conv2d_with_bias_and_mul(%arg0: tensor<1x3x4x3xf32>) -> (ten func.return %1, %2 : tensor<1x3x2x2xf32>, tensor<1x3x2x2xf32> } // CHECK: func @not_fuse_conv2d_with_bias_and_mul -// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() {value = dense<2.000000e+00> : tensor<2x3x3x2xf32>} : () -> tensor<2x3x3x2xf32> -// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() {value = dense<4.000000e-01> : tensor<2xf32>} : () -> tensor<2xf32> -// CHECK-DAG: %[[CONST_1:.*]] = "tf.Const"() {value = dense<8.000000e-01> : tensor<2xf32>} : () -> tensor<2xf32> -// CHECK-NEXT: %[[CONV2D:.*]] = "tf.Conv2D"(%arg0, %[[CONST]]) {data_format = "NHWC", dilations = [1, 1, 2, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 2, 1], use_cudnn_on_gpu = true} : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x3x2x2xf32> -// CHECK-NEXT: %[[BIASADD:.*]] = "tf.BiasAdd"(%[[CONV2D]], %[[CONST_0]]) {data_format = "NHWC"} : (tensor<1x3x2x2xf32>, tensor<2xf32>) -> tensor<1x3x2x2xf32> +// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() <{value = dense<2.000000e+00> : tensor<2x3x3x2xf32>}> : () -> tensor<2x3x3x2xf32> +// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() <{value = dense<4.000000e-01> : tensor<2xf32>}> : () -> tensor<2xf32> +// CHECK-DAG: %[[CONST_1:.*]] = "tf.Const"() <{value = dense<8.000000e-01> : tensor<2xf32>}> : () -> tensor<2xf32> +// CHECK-NEXT: %[[CONV2D:.*]] = "tf.Conv2D"(%arg0, %[[CONST]]) <{data_format = "NHWC", dilations = [1, 1, 2, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 2, 1], use_cudnn_on_gpu = true}> : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x3x2x2xf32> +// CHECK-NEXT: %[[BIASADD:.*]] = "tf.BiasAdd"(%[[CONV2D]], %[[CONST_0]]) <{data_format = "NHWC"}> : (tensor<1x3x2x2xf32>, tensor<2xf32>) -> tensor<1x3x2x2xf32> // CHECK-NEXT: %[[MUL:.*]] = "tf.Mul"(%[[CONV2D]], %[[CONST_1]]) : (tensor<1x3x2x2xf32>, tensor<2xf32>) -> tensor<1x3x2x2xf32> // CHECK-NEXT: return %[[BIASADD]], %[[MUL]] : tensor<1x3x2x2xf32>, tensor<1x3x2x2xf32> @@ -139,10 +139,10 @@ func.func @fuse_conv2d_with_bias_and_add(%arg0: tensor<1x3x4x3xf32>) -> (tensor< func.return %2 : tensor<1x3x2x2xf32> } // CHECK: func @fuse_conv2d_with_bias_and_add -// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() {value = dense<2.000000e+00> : tensor<2x3x3x2xf32>} : () -> tensor<2x3x3x2xf32> -// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() {value = dense<1.000000e+00> : tensor<2xf32>} : () -> tensor<2xf32> -// CHECK-NEXT: %[[CONV2D:.*]] = "tf.Conv2D"(%arg0, %[[CONST]]) {data_format = "NHWC", dilations = [1, 1, 2, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 2, 1], use_cudnn_on_gpu = true} : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x3x2x2xf32> -// CHECK-NEXT: %[[BIASADD:.*]] = "tf.BiasAdd"(%[[CONV2D]], %[[CONST_0]]) {data_format = "NHWC"} : (tensor<1x3x2x2xf32>, tensor<2xf32>) -> tensor<1x3x2x2xf32> +// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() <{value = dense<2.000000e+00> : tensor<2x3x3x2xf32>}> : () -> tensor<2x3x3x2xf32> +// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() <{value = dense<1.000000e+00> : tensor<2xf32>}> : () -> tensor<2xf32> +// CHECK-NEXT: %[[CONV2D:.*]] = "tf.Conv2D"(%arg0, %[[CONST]]) <{data_format = "NHWC", dilations = [1, 1, 2, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 2, 1], use_cudnn_on_gpu = true}> : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x3x2x2xf32> +// CHECK-NEXT: %[[BIASADD:.*]] = "tf.BiasAdd"(%[[CONV2D]], %[[CONST_0]]) <{data_format = "NHWC"}> : (tensor<1x3x2x2xf32>, tensor<2xf32>) -> tensor<1x3x2x2xf32> // CHECK-NEXT: return %[[BIASADD]] : tensor<1x3x2x2xf32> // ----- @@ -156,10 +156,10 @@ func.func @not_fuse_conv2d_with_bias_and_add(%arg0: tensor<1x3x4x3xf32>, %arg1: func.return %2 : tensor<1x3x2x2xf32> } // CHECK: func @not_fuse_conv2d_with_bias_and_add -// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() {value = dense<2.000000e+00> : tensor<2x3x3x2xf32>} : () -> tensor<2x3x3x2xf32> -// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() {value = dense<4.000000e-01> : tensor<2xf32>} : () -> tensor<2xf32> -// CHECK-NEXT: %[[CONV2D:.*]] = "tf.Conv2D"(%arg0, %[[CONST]]) {data_format = "NHWC", dilations = [1, 1, 2, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 2, 1], use_cudnn_on_gpu = true} : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x3x2x2xf32> -// CHECK-NEXT: %[[BIASADD:.*]] = "tf.BiasAdd"(%[[CONV2D]], %[[CONST_0]]) {data_format = "NHWC"} : (tensor<1x3x2x2xf32>, tensor<2xf32>) -> tensor<1x3x2x2xf32> +// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() <{value = dense<2.000000e+00> : tensor<2x3x3x2xf32>}> : () -> tensor<2x3x3x2xf32> +// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() <{value = dense<4.000000e-01> : tensor<2xf32>}> : () -> tensor<2xf32> +// CHECK-NEXT: %[[CONV2D:.*]] = "tf.Conv2D"(%arg0, %[[CONST]]) <{data_format = "NHWC", dilations = [1, 1, 2, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 2, 1], use_cudnn_on_gpu = true}> : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x3x2x2xf32> +// CHECK-NEXT: %[[BIASADD:.*]] = "tf.BiasAdd"(%[[CONV2D]], %[[CONST_0]]) <{data_format = "NHWC"}> : (tensor<1x3x2x2xf32>, tensor<2xf32>) -> tensor<1x3x2x2xf32> // CHECK-NEXT: %[[ADD:.*]] = "tf.AddV2"(%[[BIASADD]], %arg1) : (tensor<1x3x2x2xf32>, tensor<2xf32>) -> tensor<1x3x2x2xf32> // CHECK-NEXT: return %[[ADD]] : tensor<1x3x2x2xf32> @@ -173,10 +173,10 @@ func.func @match_depthwise_conv2d_and_add(%arg0: tensor<*xf32>) -> (tensor<*xf32 func.return %1 : tensor<*xf32> } // CHECK: func @match_depthwise_conv2d_and_add -// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() {value = dense<2.000000e+00> : tensor<2x3x3x1xf32>} : () -> tensor<2x3x3x1xf32> -// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() {value = dense<4.000000e-01> : tensor<3xf32>} : () -> tensor<3xf32> -// CHECK-NEXT: %[[DEPTHWISE_CONV2D:.*]] = "tf.DepthwiseConv2dNative"(%arg0, %[[CONST]]) {data_format = "NHWC", device = "", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 1, 1]} : (tensor<*xf32>, tensor<2x3x3x1xf32>) -> tensor -// CHECK-NEXT: %[[BIASADD:.*]] = "tf.BiasAdd"(%[[DEPTHWISE_CONV2D]], %[[CONST_0]]) {data_format = "NHWC"} : (tensor, tensor<3xf32>) -> tensor<*xf32> +// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() <{value = dense<2.000000e+00> : tensor<2x3x3x1xf32>}> : () -> tensor<2x3x3x1xf32> +// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() <{value = dense<4.000000e-01> : tensor<3xf32>}> : () -> tensor<3xf32> +// CHECK-NEXT: %[[DEPTHWISE_CONV2D:.*]] = "tf.DepthwiseConv2dNative"(%arg0, %[[CONST]]) <{data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 1, 1]}> {device = ""} : (tensor<*xf32>, tensor<2x3x3x1xf32>) -> tensor +// CHECK-NEXT: %[[BIASADD:.*]] = "tf.BiasAdd"(%[[DEPTHWISE_CONV2D]], %[[CONST_0]]) <{data_format = "NHWC"}> : (tensor, tensor<3xf32>) -> tensor<*xf32> // CHECK-NEXT: return %[[BIASADD]] : tensor<*xf32> // ----- @@ -189,8 +189,8 @@ func.func @match_depthwise_conv2d_and_mul(%arg0: tensor<*xf32>) -> (tensor } // CHECK: func @match_depthwise_conv2d_and_mul -// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() {value = dense<8.000000e-01> : tensor<2x3x3x1xf32>} : () -> tensor<2x3x3x1xf32> -// CHECK-NEXT: %[[DEPTHWISE_CONV2D:.*]] = "tf.DepthwiseConv2dNative"(%arg0, %[[CONST]]) {data_format = "NHWC", device = "", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 1, 1]} : (tensor<*xf32>, tensor<2x3x3x1xf32>) -> tensor +// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() <{value = dense<8.000000e-01> : tensor<2x3x3x1xf32>}> : () -> tensor<2x3x3x1xf32> +// CHECK-NEXT: %[[DEPTHWISE_CONV2D:.*]] = "tf.DepthwiseConv2dNative"(%arg0, %[[CONST]]) <{data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 1, 1]}> {device = ""} : (tensor<*xf32>, tensor<2x3x3x1xf32>) -> tensor // CHECK-NEXT: return %[[DEPTHWISE_CONV2D]] : tensor // ----- @@ -205,10 +205,10 @@ func.func @match_depthwise_conv2d_with_bias_and_add(%arg0: tensor<*xf32>) -> (te func.return %2 : tensor } // CHECK: func @match_depthwise_conv2d_with_bias_and_add -// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() {value = dense<2.000000e+00> : tensor<2x3x3x1xf32>} : () -> tensor<2x3x3x1xf32> -// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() {value = dense<8.000000e-01> : tensor<3xf32>} : () -> tensor<3xf32> -// CHECK-NEXT: %[[DEPTHWISE_CONV2D:.*]] = "tf.DepthwiseConv2dNative"(%arg0, %[[CONST]]) {data_format = "NHWC", device = "", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 1, 1]} : (tensor<*xf32>, tensor<2x3x3x1xf32>) -> tensor -// CHECK-NEXT: %[[BIASADD:.*]] = "tf.BiasAdd"(%[[DEPTHWISE_CONV2D]], %[[CONST_0]]) {data_format = "NHWC"} : (tensor, tensor<3xf32>) -> tensor +// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() <{value = dense<2.000000e+00> : tensor<2x3x3x1xf32>}> : () -> tensor<2x3x3x1xf32> +// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() <{value = dense<8.000000e-01> : tensor<3xf32>}> : () -> tensor<3xf32> +// CHECK-NEXT: %[[DEPTHWISE_CONV2D:.*]] = "tf.DepthwiseConv2dNative"(%arg0, %[[CONST]]) <{data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 1, 1]}> {device = ""} : (tensor<*xf32>, tensor<2x3x3x1xf32>) -> tensor +// CHECK-NEXT: %[[BIASADD:.*]] = "tf.BiasAdd"(%[[DEPTHWISE_CONV2D]], %[[CONST_0]]) <{data_format = "NHWC"}> : (tensor, tensor<3xf32>) -> tensor // CHECK-NEXT: return %[[BIASADD]] : tensor // ----- @@ -223,10 +223,10 @@ func.func @match_depthwise_conv2d_with_bias_and_mul(%arg0: tensor<*xf32>) -> (te func.return %2 : tensor } // CHECK: func @match_depthwise_conv2d_with_bias_and_mul -// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() {value = dense<1.000000e+00> : tensor<2x3x3x1xf32>} : () -> tensor<2x3x3x1xf32> -// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() {value = dense<2.000000e-01> : tensor<3xf32>} : () -> tensor<3xf32> -// CHECK-NEXT: %[[DEPTHWISE_CONV2D:.*]] = "tf.DepthwiseConv2dNative"(%arg0, %[[CONST]]) {data_format = "NHWC", device = "", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 1, 1]} : (tensor<*xf32>, tensor<2x3x3x1xf32>) -> tensor -// CHECK-NEXT: %[[BIASADD:.*]] = "tf.BiasAdd"(%[[DEPTHWISE_CONV2D]], %[[CONST_0]]) {data_format = "NHWC"} : (tensor, tensor<3xf32>) -> tensor +// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() <{value = dense<1.000000e+00> : tensor<2x3x3x1xf32>}> : () -> tensor<2x3x3x1xf32> +// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() <{value = dense<2.000000e-01> : tensor<3xf32>}> : () -> tensor<3xf32> +// CHECK-NEXT: %[[DEPTHWISE_CONV2D:.*]] = "tf.DepthwiseConv2dNative"(%arg0, %[[CONST]]) <{data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 1, 1]}> {device = ""} : (tensor<*xf32>, tensor<2x3x3x1xf32>) -> tensor +// CHECK-NEXT: %[[BIASADD:.*]] = "tf.BiasAdd"(%[[DEPTHWISE_CONV2D]], %[[CONST_0]]) <{data_format = "NHWC"}> : (tensor, tensor<3xf32>) -> tensor // CHECK-NEXT: return %[[BIASADD]] : tensor // ----- @@ -236,7 +236,7 @@ func.func @lower_einsum(%arg0: tensor<3x4x5xf32>, %arg1: tensor<3x5x6xf32>) -> t func.return %0 : tensor<3x4x6xf32> } // CHECK-LABEL: lower_einsum -// CHECK: "tf.BatchMatMulV2"(%arg0, %arg1) {adj_x = false, adj_y = false} : (tensor<3x4x5xf32>, tensor<3x5x6xf32>) -> tensor<3x4x6xf32> +// CHECK: "tf.BatchMatMulV2"(%arg0, %arg1) <{adj_x = false, adj_y = false}> : (tensor<3x4x5xf32>, tensor<3x5x6xf32>) -> tensor<3x4x6xf32> // ----- @@ -251,7 +251,7 @@ func.func @removing_identity_after_const(%arg0: tensor<*xf32>) -> (tensor<*xf32> func.return %2 : tensor<*xf32> } // CHECK: func @removing_identity_after_const -// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() {value = dense<1.000000e+00> : tensor<2x3x3x1xf32>} : () -> tensor<2x3x3x1xf32> +// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() <{value = dense<1.000000e+00> : tensor<2x3x3x1xf32>}> : () -> tensor<2x3x3x1xf32> // CHECK: %[[DEPTHWISE_CONV2D:.*]] = "tf.DepthwiseConv2dNative"(%arg0, %[[CONST]]) // ----- @@ -291,14 +291,14 @@ func.func @batch_norm_with_q_dq(%arg0: tensor<1x3x4x3xf32>) -> (tensor<1x3x2x2xf } // CHECK: func @batch_norm_with_q_dq -// CHECK-DAG: %[[cst:.*]] = "tf.Const"() {value = dense<0.707036077> : tensor<2x3x3x2xf32>} : () -> tensor<2x3x3x2xf32> -// CHECK-DAG: %[[cst_0:.*]] = "tf.Const"() {value = dense<-0.914072155> : tensor<2xf32>} : () -> tensor<2xf32> +// CHECK-DAG: %[[cst:.*]] = "tf.Const"() <{value = dense<0.707036077> : tensor<2x3x3x2xf32>}> : () -> tensor<2x3x3x2xf32> +// CHECK-DAG: %[[cst_0:.*]] = "tf.Const"() <{value = dense<-0.914072155> : tensor<2xf32>}> : () -> tensor<2xf32> // CHECK: %[[q_input:.*]] = "quantfork.qcast"(%arg0) : (tensor<1x3x4x3xf32>) -> tensor<1x3x4x3x!quant.uniform> // CHECK: %[[dq_input:.*]] = "quantfork.dcast"(%[[q_input]]) : (tensor<1x3x4x3x!quant.uniform>) -> tensor<1x3x4x3xf32> // CHECK: %[[q_weight:.*]] = "quantfork.qcast"(%[[cst]]) : (tensor<2x3x3x2xf32>) -> tensor<2x3x3x2x!quant.uniform:f32:3, {0.005567213212411235,0.005567213212411235}>> // CHECK: %[[dq_weight:.*]] = "quantfork.dcast"(%[[q_weight]]) : (tensor<2x3x3x2x!quant.uniform:f32:3, {0.005567213212411235,0.005567213212411235}>>) -> tensor<2x3x3x2xf32> // CHECK: %[[conv:.*]] = "tf.Conv2D"(%[[dq_input]], %[[dq_weight]]) -// CHECK: %[[bias:.*]] = "tf.BiasAdd"(%[[conv]], %[[cst_0]]) {data_format = "NHWC"} +// CHECK: %[[bias:.*]] = "tf.BiasAdd"(%[[conv]], %[[cst_0]]) <{data_format = "NHWC"}> // CHECK: %[[relu6:.*]] = "tf.Relu6"(%[[bias]]) // ----- @@ -334,8 +334,8 @@ func.func @conv2d_with_large_weight_and_mul(%arg0: tensor) -> (tens func.return %2 : tensor } // CHECK: func @conv2d_with_large_weight_and_mul -// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() {value = dense<1.250000e+00> : tensor<48x48x3x256xf32>} : () -> tensor<48x48x3x256xf32> -// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() {value = dense<2.000000e-01> : tensor<256xf32>} : () -> tensor<256xf32> +// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() <{value = dense<1.250000e+00> : tensor<48x48x3x256xf32>}> : () -> tensor<48x48x3x256xf32> +// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() <{value = dense<2.000000e-01> : tensor<256xf32>}> : () -> tensor<256xf32> // CHECK-NEXT: %[[CONV2D:.*]] = "tf.Conv2D"(%arg0, %[[CONST]]) // CHECK-NEXT: %[[BIASADD:.*]] = "tf.BiasAdd"(%[[CONV2D]], %[[CONST_0]]) // CHECK-NEXT: return %[[BIASADD]] @@ -354,8 +354,8 @@ func.func @depthwise_conv2d_with_large_weight_and_add(%arg0: tensor<*xf32>) -> ( func.return %2 : tensor } // CHECK: func @depthwise_conv2d_with_large_weight_and_add -// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() {value = dense<2.500000e+00> : tensor<48x48x3x256xf32>} : () -> tensor<48x48x3x256xf32> -// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() {value = dense<8.000000e-01> : tensor<3xf32>} : () -> tensor<3xf32> +// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() <{value = dense<2.500000e+00> : tensor<48x48x3x256xf32>}> : () -> tensor<48x48x3x256xf32> +// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() <{value = dense<8.000000e-01> : tensor<3xf32>}> : () -> tensor<3xf32> // CHECK-NEXT: %[[DEPTHWISE_CONV2D:.*]] = "tf.DepthwiseConv2dNative"(%arg0, %[[CONST]]) // CHECK-NEXT: %[[BIASADD:.*]] = "tf.BiasAdd"(%[[DEPTHWISE_CONV2D]], %[[CONST_0]]) // CHECK-NEXT: return %[[BIASADD]] diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/tests/prepare_quantize_drq.mlir b/tensorflow/compiler/mlir/quantization/tensorflow/tests/prepare_quantize_drq.mlir index 4da50d4ac91b31..0176867dd48865 100644 --- a/tensorflow/compiler/mlir/quantization/tensorflow/tests/prepare_quantize_drq.mlir +++ b/tensorflow/compiler/mlir/quantization/tensorflow/tests/prepare_quantize_drq.mlir @@ -15,11 +15,11 @@ module { // CHECK-DAG: %[[CONST:.*]] = arith.constant dense<0.000000e+00> : tensor<2x1024xf32> // CHECK: %0 = "quantfork.qcast"(%[[CONST]]) : (tensor<2x1024xf32>) -> tensor<2x1024x!quant.uniform:f32, 3.9370078740157481E-9>> // CHECK: %1 = "quantfork.dcast"(%0) : (tensor<2x1024x!quant.uniform:f32, 3.9370078740157481E-9>>) -> tensor<2x1024xf32> -// CHECK: %2 = "tf.PartitionedCall"(%arg0, %1) {_tfl_quant_trait = "fully_quantizable", config = "", config_proto = "", executor_type = "", f = @composite_matmul_fn} : (tensor<1x2x2x3xf32>, tensor<2x1024xf32>) -> tensor<*xf32> +// CHECK: %2 = "tf.PartitionedCall"(%arg0, %1) <{config = "", config_proto = "", executor_type = "", f = @composite_matmul_fn}> {_tfl_quant_trait = "fully_quantizable"} : (tensor<1x2x2x3xf32>, tensor<2x1024xf32>) -> tensor<*xf32> // CHECK: return %2 : tensor<*xf32> // CHECK-LABEL: func private @composite_matmul_fn -// CHECK: %0 = "tf.MatMul"(%arg0, %arg1) {attr_map = "0:transpose_a,1:transpose_a", device = "", transpose_a = false, transpose_b = false} : (tensor<1x2x2x3xf32>, tensor<2x1024xf32>) -> tensor<*xf32> +// CHECK: %0 = "tf.MatMul"(%arg0, %arg1) <{transpose_a = false, transpose_b = false}> {attr_map = "0:transpose_a,1:transpose_a", device = ""} : (tensor<1x2x2x3xf32>, tensor<2x1024xf32>) -> tensor<*xf32> // CHECK: return %0 : tensor<*xf32> } @@ -43,7 +43,7 @@ module { // CHECK-DAG: %[[CONST_1:.*]] = arith.constant dense<3.000000e+00> : tensor<2x3x3x512xf32> // CHECK: %0 = "quantfork.qcast"(%[[CONST_1]]) : (tensor<2x3x3x512xf32>) -> tensor<2x3x3x512x!quant.uniform:f32, 0.023622047244094488>> // CHECK: %1 = "quantfork.dcast"(%0) : (tensor<2x3x3x512x!quant.uniform:f32, 0.023622047244094488>>) -> tensor<2x3x3x512xf32> -// CHECK: %2 = "tf.PartitionedCall"(%arg0, %1) {_tfl_quant_trait = "fully_quantizable", config = "", config_proto = "", executor_type = "", f = @composite_conv2d_fn_1} : (tensor<1x3x4x3xf32>, tensor<2x3x3x512xf32>) -> tensor<*xf32> +// CHECK: %2 = "tf.PartitionedCall"(%arg0, %1) <{config = "", config_proto = "", executor_type = "", f = @composite_conv2d_fn_1}> {_tfl_quant_trait = "fully_quantizable"} : (tensor<1x3x4x3xf32>, tensor<2x3x3x512xf32>) -> tensor<*xf32> // CHECK: %3 = "tf.BiasAdd"(%2, %[[CONST_0]]) // CHECK: return %3 : tensor<*xf32> @@ -74,7 +74,7 @@ module { // CHECK-DAG: %[[CONST_1:.*]] = arith.constant dense<3.000000e+00> : tensor<2x3x1x1536xf32> // CHECK: %0 = "quantfork.qcast"(%[[CONST_1]]) : (tensor<2x3x1x1536xf32>) -> tensor<2x3x1x1536x!quant.uniform:f32, 0.023622047244094488>> // CHECK: %1 = "quantfork.dcast"(%0) : (tensor<2x3x1x1536x!quant.uniform:f32, 0.023622047244094488>>) -> tensor<2x3x1x1536xf32> -// CHECK: %2 = "tf.PartitionedCall"(%arg0, %1) {_tfl_quant_trait = "fully_quantizable", config = "", config_proto = "", executor_type = "", f = @composite_depthwise_conv2d_fn_0} : (tensor<1x3x4x512xf32>, tensor<2x3x1x1536xf32>) -> tensor<*xf32> +// CHECK: %2 = "tf.PartitionedCall"(%arg0, %1) <{config = "", config_proto = "", executor_type = "", f = @composite_depthwise_conv2d_fn_0}> {_tfl_quant_trait = "fully_quantizable"} : (tensor<1x3x4x512xf32>, tensor<2x3x1x1536xf32>) -> tensor<*xf32> // CHECK: %3 = "tf.BiasAdd"(%2, %[[CONST_0]]) // CHECK: return %3 : tensor<*xf32> @@ -85,6 +85,6 @@ module { // CHECK-LABEL: func private @composite_depthwise_conv2d_fn_0( // CHECK-SAME: %arg0: tensor<1x3x4x512xf32>, // CHECK-SAME: %arg1: tensor<2x3x1x1536xf32>) -// CHECK: %0 = "tf.DepthwiseConv2dNative"(%arg0, %arg1) {attr_map = "0:strides,1:padding,2:explicit_paddings,3:dilations", data_format = "NHWC", device = "", +// CHECK: %0 = "tf.DepthwiseConv2dNative"(%arg0, %arg1) <{data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 2, 1]}> {attr_map = "0:strides,1:padding,2:explicit_paddings,3:dilations", device = ""} // CHECK: return %0 : tensor<*xf32> } diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/tests/prepare_quantize_drq_per_channel.mlir b/tensorflow/compiler/mlir/quantization/tensorflow/tests/prepare_quantize_drq_per_channel.mlir index f2d80c0bf4e01a..927fc34fb564dd 100644 --- a/tensorflow/compiler/mlir/quantization/tensorflow/tests/prepare_quantize_drq_per_channel.mlir +++ b/tensorflow/compiler/mlir/quantization/tensorflow/tests/prepare_quantize_drq_per_channel.mlir @@ -15,11 +15,11 @@ module { // CHECK-DAG: %[[CONST:.*]] = arith.constant dense<0.000000e+00> : tensor<2x1024xf32> // CHECK: %0 = "quantfork.qcast"(%[[CONST]]) : (tensor<2x1024xf32>) -> tensor<2x1024x!quant.uniform:f32, 3.9370078740157481E-9>> // CHECK: %1 = "quantfork.dcast"(%0) : (tensor<2x1024x!quant.uniform:f32, 3.9370078740157481E-9>>) -> tensor<2x1024xf32> -// CHECK: %2 = "tf.PartitionedCall"(%arg0, %1) {_tfl_quant_trait = "fully_quantizable", config = "", config_proto = "", executor_type = "", f = @composite_matmul_fn} : (tensor<1x2x2x3xf32>, tensor<2x1024xf32>) -> tensor<*xf32> +// CHECK: %2 = "tf.PartitionedCall"(%arg0, %1) <{config = "", config_proto = "", executor_type = "", f = @composite_matmul_fn}> {_tfl_quant_trait = "fully_quantizable"} : (tensor<1x2x2x3xf32>, tensor<2x1024xf32>) -> tensor<*xf32> // CHECK: return %2 : tensor<*xf32> // CHECK-LABEL: func private @composite_matmul_fn -// CHECK: %0 = "tf.MatMul"(%arg0, %arg1) {attr_map = "0:transpose_a,1:transpose_a", device = "", transpose_a = false, transpose_b = false} : (tensor<1x2x2x3xf32>, tensor<2x1024xf32>) -> tensor<*xf32> +// CHECK: %0 = "tf.MatMul"(%arg0, %arg1) <{transpose_a = false, transpose_b = false}> {attr_map = "0:transpose_a,1:transpose_a", device = ""} : (tensor<1x2x2x3xf32>, tensor<2x1024xf32>) -> tensor<*xf32> // CHECK: return %0 : tensor<*xf32> } @@ -43,7 +43,7 @@ module { // CHECK-DAG: %[[CONST_1:.*]] = arith.constant dense<3.000000e+00> : tensor<2x3x512x2xf32> // CHECK: %0 = "quantfork.qcast"(%[[CONST_1]]) : (tensor<2x3x512x2xf32>) -> tensor<2x3x512x2x!quant.uniform:f32:3, {0.023622047244094488,0.023622047244094488}>> // CHECK: %1 = "quantfork.dcast"(%0) : (tensor<2x3x512x2x!quant.uniform:f32:3, {0.023622047244094488,0.023622047244094488}>>) -> tensor<2x3x512x2xf32> -// CHECK: %2 = "tf.PartitionedCall"(%arg0, %1) {_tfl_quant_trait = "fully_quantizable", config = "", config_proto = "", executor_type = "", f = @composite_conv2d_fn_1} : (tensor<1x3x4x512xf32>, tensor<2x3x512x2xf32>) -> tensor<*xf32> +// CHECK: %2 = "tf.PartitionedCall"(%arg0, %1) <{config = "", config_proto = "", executor_type = "", f = @composite_conv2d_fn_1}> {_tfl_quant_trait = "fully_quantizable"} : (tensor<1x3x4x512xf32>, tensor<2x3x512x2xf32>) -> tensor<*xf32> // CHECK: %3 = "tf.BiasAdd"(%2, %[[CONST_0]]) // CHECK: return %3 : tensor<*xf32> @@ -74,7 +74,7 @@ module { // CHECK-DAG: %[[CONST_1:.*]] = arith.constant dense<3.000000e+00> : tensor<2x3x1x1536xf32> // CHECK: %0 = "quantfork.qcast"(%[[CONST_1]]) : (tensor<2x3x1x1536xf32>) -> tensor<2x3x1x1536x!quant.uniform:f32:3, {0.023622047244094488, // CHECK: %1 = "quantfork.dcast"(%0) : (tensor<2x3x1x1536x!quant.uniform:f32:3, {0.023622047244094488, -// CHECK: %2 = "tf.PartitionedCall"(%arg0, %1) {_tfl_quant_trait = "fully_quantizable", config = "", config_proto = "", executor_type = "", f = @composite_depthwise_conv2d_fn_0} : (tensor<1x3x4x512xf32>, tensor<2x3x1x1536xf32>) -> tensor<*xf32> +// CHECK: %2 = "tf.PartitionedCall"(%arg0, %1) <{config = "", config_proto = "", executor_type = "", f = @composite_depthwise_conv2d_fn_0}> {_tfl_quant_trait = "fully_quantizable"} : (tensor<1x3x4x512xf32>, tensor<2x3x1x1536xf32>) -> tensor<*xf32> // CHECK: %3 = "tf.BiasAdd"(%2, %[[CONST_0]]) // CHECK: return %3 : tensor<*xf32> @@ -85,6 +85,6 @@ module { // CHECK-LABEL: func private @composite_depthwise_conv2d_fn_0( // CHECK-SAME: %arg0: tensor<1x3x4x512xf32>, // CHECK-SAME: %arg1: tensor<2x3x1x1536xf32>) -// CHECK: %0 = "tf.DepthwiseConv2dNative"(%arg0, %arg1) {attr_map = "0:strides,1:padding,2:explicit_paddings,3:dilations", data_format = "NHWC", device = "", +// CHECK: %0 = "tf.DepthwiseConv2dNative"(%arg0, %arg1) <{data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 2, 1]}> {attr_map = "0:strides,1:padding,2:explicit_paddings,3:dilations", device = ""} // CHECK: return %0 : tensor<*xf32> } diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/tests/preprocess_op.mlir b/tensorflow/compiler/mlir/quantization/tensorflow/tests/preprocess_op.mlir index 0ef69f6f6f791f..ae8a20d6b82fa0 100644 --- a/tensorflow/compiler/mlir/quantization/tensorflow/tests/preprocess_op.mlir +++ b/tensorflow/compiler/mlir/quantization/tensorflow/tests/preprocess_op.mlir @@ -22,8 +22,8 @@ module { // CHECK: %[[CONST_1:.*]] = arith.constant dense // CHECK-NOT: tensor<2x3x3x2xf32> // CHECK-SAME: tensor<2x3x1x6xf32> -// CHECK: %[[PARTITIONEDCALL_0:.*]] = "tf.PartitionedCall"(%arg0, %[[CONST_1:.*]]) {_tfl_quant_trait = "fully_quantizable", config = "", config_proto = "", executor_type = "", f = @composite_depthwise_conv2d_fn_0} : (tensor<1x3x4x3xf32>, tensor<2x3x1x6xf32>) -> tensor<*xf32> -// CHECK: %[[BIAS_0:.*]] = "tf.BiasAdd"(%[[PARTITIONEDCALL_0]], %[[CONST_0:.*]]) {data_format = "NHWC", device = ""} : (tensor<*xf32>, tensor<6xf32>) -> tensor<*xf32> +// CHECK: %[[PARTITIONEDCALL_0:.*]] = "tf.PartitionedCall"(%arg0, %[[CONST_1:.*]]) <{config = "", config_proto = "", executor_type = "", f = @composite_depthwise_conv2d_fn_0}> {_tfl_quant_trait = "fully_quantizable"} : (tensor<1x3x4x3xf32>, tensor<2x3x1x6xf32>) -> tensor<*xf32> +// CHECK: %[[BIAS_0:.*]] = "tf.BiasAdd"(%[[PARTITIONEDCALL_0]], %[[CONST_0:.*]]) <{data_format = "NHWC"}> {device = ""} : (tensor<*xf32>, tensor<6xf32>) -> tensor<*xf32> // CHECK: return %[[BIAS_0:.*]] : tensor<*xf32> // CHECK-LABEL: func private @composite_depthwise_conv2d_fn( @@ -33,7 +33,7 @@ module { // CHECK-LABEL: func private @composite_depthwise_conv2d_fn_0( // CHECK-SAME: %arg0: tensor<1x3x4x3xf32>, // CHECK-SAME: %arg1: tensor<2x3x1x6xf32>) -// CHECK: %0 = "tf.DepthwiseConv2dNative"(%arg0, %arg1) {attr_map = "0:strides,1:padding,2:explicit_paddings,3:dilations", data_format = "NHWC", device = "", +// CHECK: %0 = "tf.DepthwiseConv2dNative"(%arg0, %arg1) <{data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 2, 1]}> {attr_map = "0:strides,1:padding,2:explicit_paddings,3:dilations", device = ""} // CHECK: return %0 : tensor<*xf32> } diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/tests/preprocess_op_weight_only.mlir b/tensorflow/compiler/mlir/quantization/tensorflow/tests/preprocess_op_weight_only.mlir index 4f36784e67a6d4..e80db7ff049b5b 100644 --- a/tensorflow/compiler/mlir/quantization/tensorflow/tests/preprocess_op_weight_only.mlir +++ b/tensorflow/compiler/mlir/quantization/tensorflow/tests/preprocess_op_weight_only.mlir @@ -23,14 +23,14 @@ module { // PerTensor: %[[CONST_1:.*]] = arith.constant dense // PerTensor-NOT: tensor<2x3x1x6xf32> // PerTensor-SAME: tensor<2x3x3x2xf32> -// PerTensor: %[[PARTITIONEDCALL_0:.*]] = "tf.PartitionedCall"(%arg0, %[[CONST_1:.*]]) {_tfl_quant_trait = "fully_quantizable", config = "", config_proto = "", executor_type = "", f = @composite_depthwise_conv2d_fn} : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<*xf32> -// PerTensor: %[[BIAS_0:.*]] = "tf.BiasAdd"(%[[PARTITIONEDCALL_0]], %[[CONST_0:.*]]) {data_format = "NHWC", device = ""} : (tensor<*xf32>, tensor<6xf32>) -> tensor<*xf32> +// PerTensor: %[[PARTITIONEDCALL_0:.*]] = "tf.PartitionedCall"(%arg0, %[[CONST_1:.*]]) <{config = "", config_proto = "", executor_type = "", f = @composite_depthwise_conv2d_fn}> {_tfl_quant_trait = "fully_quantizable"} : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<*xf32> +// PerTensor: %[[BIAS_0:.*]] = "tf.BiasAdd"(%[[PARTITIONEDCALL_0]], %[[CONST_0:.*]]) <{data_format = "NHWC"}> {device = ""} : (tensor<*xf32>, tensor<6xf32>) -> tensor<*xf32> // PerTensor: return %[[BIAS_0:.*]] : tensor<*xf32> // PerTensor-LABEL: func private @composite_depthwise_conv2d_fn( // PerTensor-SAME: %arg0: tensor<1x3x4x3xf32>, // PerTensor-SAME: %arg1: tensor<2x3x3x2xf32>) -// PerTensor: %0 = "tf.DepthwiseConv2dNative"(%arg0, %arg1) {attr_map = "0:strides,1:padding,2:explicit_paddings,3:dilations", data_format = "NHWC", device = "", +// PerTensor: %0 = "tf.DepthwiseConv2dNative"(%arg0, %arg1) <{data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 2, 1]}> {attr_map = "0:strides,1:padding,2:explicit_paddings,3:dilations", device = ""} // PerTensor: return %0 : tensor<*xf32> // PerChannel-LABEL: func @depthwise_conv @@ -38,8 +38,8 @@ module { // PerChannel: %[[CONST_1:.*]] = arith.constant dense // PerChannel-NOT: tensor<2x3x3x2xf32> // PerChannel-SAME: tensor<2x3x1x6xf32> -// PerChannel: %[[PARTITIONEDCALL_0:.*]] = "tf.PartitionedCall"(%arg0, %[[CONST_1:.*]]) {_tfl_quant_trait = "fully_quantizable", config = "", config_proto = "", executor_type = "", f = @composite_depthwise_conv2d_fn_0} : (tensor<1x3x4x3xf32>, tensor<2x3x1x6xf32>) -> tensor<*xf32> -// PerChannel: %[[BIAS_0:.*]] = "tf.BiasAdd"(%[[PARTITIONEDCALL_0]], %[[CONST_0:.*]]) {data_format = "NHWC", device = ""} : (tensor<*xf32>, tensor<6xf32>) -> tensor<*xf32> +// PerChannel: %[[PARTITIONEDCALL_0:.*]] = "tf.PartitionedCall"(%arg0, %[[CONST_1:.*]]) <{config = "", config_proto = "", executor_type = "", f = @composite_depthwise_conv2d_fn_0}> {_tfl_quant_trait = "fully_quantizable"} : (tensor<1x3x4x3xf32>, tensor<2x3x1x6xf32>) -> tensor<*xf32> +// PerChannel: %[[BIAS_0:.*]] = "tf.BiasAdd"(%[[PARTITIONEDCALL_0]], %[[CONST_0:.*]]) <{data_format = "NHWC"}> {device = ""} : (tensor<*xf32>, tensor<6xf32>) -> tensor<*xf32> // PerChannel: return %[[BIAS_0:.*]] : tensor<*xf32> // PerChannel-LABEL: func private @composite_depthwise_conv2d_fn( @@ -49,7 +49,7 @@ module { // PerChannel-LABEL: func private @composite_depthwise_conv2d_fn_0( // PerChannel-SAME: %arg0: tensor<1x3x4x3xf32>, // PerChannel-SAME: %arg1: tensor<2x3x1x6xf32>) -// PerChannel: %0 = "tf.DepthwiseConv2dNative"(%arg0, %arg1) {attr_map = "0:strides,1:padding,2:explicit_paddings,3:dilations", data_format = "NHWC", device = "", +// PerChannel: %0 = "tf.DepthwiseConv2dNative"(%arg0, %arg1) <{data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 2, 1]}> {attr_map = "0:strides,1:padding,2:explicit_paddings,3:dilations", device = ""} // PerChannel: return %0 : tensor<*xf32> } diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/tests/propagate_quantize_type.mlir b/tensorflow/compiler/mlir/quantization/tensorflow/tests/propagate_quantize_type.mlir index 0c69477a0c8c42..6a737b7b20e1e4 100644 --- a/tensorflow/compiler/mlir/quantization/tensorflow/tests/propagate_quantize_type.mlir +++ b/tensorflow/compiler/mlir/quantization/tensorflow/tests/propagate_quantize_type.mlir @@ -12,8 +12,8 @@ module { } // CHECK-LABEL: func @not_propagate_matmul -// CHECK: %[[CASTED_W:.*]] = "tf.Cast"(%0) {Truncate = false} : (tensor<2x1024xi8>) -> tensor<2x1024xf32> -// CHECK: %2 = "tf.MatMul"(%arg0, %[[CASTED_W]]) {attr_map = "0:transpose_a,1:transpose_a", device = "", transpose_a = false, transpose_b = false} : (tensor<1x2x2x2xf32>, tensor<2x1024xf32>) -> tensor<*xf32> +// CHECK: %[[CASTED_W:.*]] = "tf.Cast"(%0) <{Truncate = false}> : (tensor<2x1024xi8>) -> tensor<2x1024xf32> +// CHECK: %2 = "tf.MatMul"(%arg0, %[[CASTED_W]]) <{transpose_a = false, transpose_b = false}> {attr_map = "0:transpose_a,1:transpose_a", device = ""} : (tensor<1x2x2x2xf32>, tensor<2x1024xf32>) -> tensor<*xf32> } // ----- @@ -37,8 +37,8 @@ module { // CHECK-LABEL: func @propagate_xladotv2_bf16 // CHECK: %[[IDENTITY:.*]] = "tf.Identity"(%cst) : (tensor<2x1024xi8>) -> tensor<2x1024xi8> -// CHECK: %[[MATMUL:.*]] = "tf.XlaDotV2"(%arg0, %[[IDENTITY]]) {device = "", dimension_numbers = "\12\01\00\0A\01\03", precision_config = ""} : (tensor<1x2x2x2xbf16>, tensor<2x1024xi8>) -> tensor<1x2x2x1024xbf16> -// CHECK: %[[DEQUANTIZED:.*]] = "tf.PartitionedCall"(%[[MATMUL]]) {config = "", config_proto = "", executor_type = "", f = @composite_dequantize_uniform} : (tensor<1x2x2x1024xbf16>) -> tensor<1x2x2x1024xbf16> +// CHECK: %[[MATMUL:.*]] = "tf.XlaDotV2"(%arg0, %[[IDENTITY]]) <{dimension_numbers = "\12\01\00\0A\01\03", precision_config = ""}> {device = ""} : (tensor<1x2x2x2xbf16>, tensor<2x1024xi8>) -> tensor<1x2x2x1024xbf16> +// CHECK: %[[DEQUANTIZED:.*]] = "tf.PartitionedCall"(%[[MATMUL]]) <{config = "", config_proto = "", executor_type = "", f = @composite_dequantize_uniform}> : (tensor<1x2x2x1024xbf16>) -> tensor<1x2x2x1024xbf16> } // ----- @@ -64,8 +64,8 @@ module { // CHECK-LABEL: func @not_propagate_last_op // CHECK: %[[IDENTITY:.*]] = "tf.Identity"(%cst_0) : (tensor<200x100x300xi8>) -> tensor<200x100x300xi8> -// CHECK: %[[DEQUANTIZED:.*]] = "tf.PartitionedCall"(%[[IDENTITY]]) {config = "", config_proto = "", executor_type = "", f = @composite_dequantize_uniform} : (tensor<200x100x300xi8>) -> tensor<200x100x300xf32> -// CHECK: %[[GATHER:.*]] = "tf.XlaGather"(%[[DEQUANTIZED]], %arg0, %cst) {dimension_numbers = "\0A\02\00\01\12\01\00\1A\02\00\01 \01", indices_are_sorted = true} : (tensor<200x100x300xf32>, tensor<10x2xi32>, tensor<3xi64>) -> tensor<1x300x10xf32> +// CHECK: %[[DEQUANTIZED:.*]] = "tf.PartitionedCall"(%[[IDENTITY]]) <{config = "", config_proto = "", executor_type = "", f = @composite_dequantize_uniform}> : (tensor<200x100x300xi8>) -> tensor<200x100x300xf32> +// CHECK: %[[GATHER:.*]] = "tf.XlaGather"(%[[DEQUANTIZED]], %arg0, %cst) <{dimension_numbers = "\0A\02\00\01\12\01\00\1A\02\00\01 \01", indices_are_sorted = true}> : (tensor<200x100x300xf32>, tensor<10x2xi32>, tensor<3xi64>) -> tensor<1x300x10xf32> // CHECK: return %[[GATHER]] : tensor<1x300x10xf32> // ----- @@ -91,7 +91,7 @@ module { // CHECK-LABEL: func @propagate_xlagather // CHECK: %[[IDENTITY:.*]] = "tf.Identity"(%cst_0) : (tensor<200x100x300xi8>) -> tensor<200x100x300xi8> -// CHECK: %[[GATHER:.*]] = "tf.XlaGather"(%[[IDENTITY]], %arg0, %cst) {dimension_numbers = "\0A\02\00\01\12\01\00\1A\02\00\01 \01", indices_are_sorted = true} : (tensor<200x100x300xi8>, tensor<10x2xi32>, tensor<3xi64>) -> tensor<1x300x10xi8> -// CHECK: %[[DEQUANTIZED:.*]] = "tf.PartitionedCall"(%[[GATHER]]) {config = "", config_proto = "", executor_type = "", f = @composite_dequantize_uniform} : (tensor<1x300x10xi8>) -> tensor<1x300x10xf32> +// CHECK: %[[GATHER:.*]] = "tf.XlaGather"(%[[IDENTITY]], %arg0, %cst) <{dimension_numbers = "\0A\02\00\01\12\01\00\1A\02\00\01 \01", indices_are_sorted = true}> : (tensor<200x100x300xi8>, tensor<10x2xi32>, tensor<3xi64>) -> tensor<1x300x10xi8> +// CHECK: %[[DEQUANTIZED:.*]] = "tf.PartitionedCall"(%[[GATHER]]) <{config = "", config_proto = "", executor_type = "", f = @composite_dequantize_uniform}> : (tensor<1x300x10xi8>) -> tensor<1x300x10xf32> // CHECK: %[[ORIGINAL_IDENTITY:.*]] = "tf.Identity"(%[[DEQUANTIZED]]) : (tensor<1x300x10xf32>) -> tensor<1x300x10xf32> // CHECK: return %[[ORIGINAL_IDENTITY]] : tensor<1x300x10xf32> diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/tests/quantize.mlir b/tensorflow/compiler/mlir/quantization/tensorflow/tests/quantize.mlir index d04ec262f6ff85..0f3c7024dba4b4 100644 --- a/tensorflow/compiler/mlir/quantization/tensorflow/tests/quantize.mlir +++ b/tensorflow/compiler/mlir/quantization/tensorflow/tests/quantize.mlir @@ -23,7 +23,7 @@ func.func private @conv(%input: tensor<1x3x4x3xf32> {tf._user_specified_name = " // CHECK-DAG: [[weight:%.+]] = "arith.constant"() <{value = dense_resource<__elided__> : tensor<2x3x3x2xf32>}> : () -> tensor<2x3x3x2x!quant.uniform> // CHECK: [[q_input:%.+]] = "quantfork.qcast"(%arg0) : (tensor<1x3x4x3xf32>) -> tensor<1x3x4x3x!quant.uniform> // CHECK-NEXT: [[q_bias:%.+]] = "quantfork.qcast"([[bias]]) : (tensor<2xf32>) -> tensor<2x!quant.uniform> -// CHECK-NEXT: [[conv:%.+]] = "tf.PartitionedCall"([[q_input]], [[weight]], [[q_bias]]) {_tfl_quant_trait = "fully_quantizable", config = "", config_proto = "", executor_type = "", f = @[[composite_fn:composite_conv2d_with_bias_and_relu6_fn.*]]} : (tensor<1x3x4x3x!quant.uniform>, tensor<2x3x3x2x!quant.uniform>, tensor<2x!quant.uniform>) -> tensor<*x!quant.uniform> +// CHECK-NEXT: [[conv:%.+]] = "tf.PartitionedCall"([[q_input]], [[weight]], [[q_bias]]) <{config = "", config_proto = "", executor_type = "", f = @[[composite_fn:composite_conv2d_with_bias_and_relu6_fn.*]]}> {_tfl_quant_trait = "fully_quantizable"} : (tensor<1x3x4x3x!quant.uniform>, tensor<2x3x3x2x!quant.uniform>, tensor<2x!quant.uniform>) -> tensor<*x!quant.uniform> // CHECK-NEXT: [[res:%.+]] = "quantfork.dcast"([[conv]]) : (tensor<*x!quant.uniform>) -> tensor<*xf32> // CHECK-NEXT: "func.return"([[res]]) : (tensor<*xf32>) -> () @@ -69,11 +69,11 @@ func.func @avgpool_test(%arg0: tensor<*xf32>) -> tensor<*xf32> { // CHECK: %[[q:.*]] = "quantfork.qcast"(%arg0) // CHECK: %[[sc1:.*]] = "quantfork.scast"(%[[q]]) : (tensor<*x!quant.uniform>) -// CHECK: %[[fcast:.*]] = "tf.Cast"(%[[sc1]]) {Truncate = false} : (tensor<*xi8>) -> tensor<*xf32> +// CHECK: %[[fcast:.*]] = "tf.Cast"(%[[sc1]]) <{Truncate = false}> : (tensor<*xi8>) -> tensor<*xf32> // CHECK: %[[avgpool_f32:.*]] = "tf.AvgPool"(%[[fcast]]) // CHECK-SAME: (tensor<*xf32>) -> tensor<*xf32> // CHECK: %[[round:.*]] = "tf.Round"(%[[avgpool_f32]]) -// CHECK: %[[icast:.*]] = "tf.Cast"(%[[round]]) {Truncate = false} : (tensor<*xf32>) -> tensor<*xi8> +// CHECK: %[[icast:.*]] = "tf.Cast"(%[[round]]) <{Truncate = false}> : (tensor<*xf32>) -> tensor<*xi8> // CHECK: %[[sc2:.*]] = "quantfork.scast"(%[[icast]]) // CHECK: %[[dq:.*]] = "quantfork.dcast"(%[[sc2]]) : (tensor<*x!quant.uniform>) // CHECK: return %[[dq]] diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/tests/quantize_composite_functions.mlir b/tensorflow/compiler/mlir/quantization/tensorflow/tests/quantize_composite_functions.mlir index eab1244fb62b3d..5b5addde086898 100644 --- a/tensorflow/compiler/mlir/quantization/tensorflow/tests/quantize_composite_functions.mlir +++ b/tensorflow/compiler/mlir/quantization/tensorflow/tests/quantize_composite_functions.mlir @@ -28,17 +28,17 @@ module { } // CHECK-LABEL: func @conv -// CHECK-DAG: %[[w_float:.*]] = "tf.Const"() {value = dense<{{\[\[\[\[}}1.600000e-01 -// CHECK-DAG: %[[b_float:.*]] = "tf.Const"() {value = dense<[-2.000000e+00, 3.000000e+00]> : tensor<2xf32> -// CHECK-DAG: %[[in_scale:.*]] = "tf.Const"() {value = dense<8.000000e-03> : tensor} : () -> tensor -// CHECK-DAG: %[[in_zp:.*]] = "tf.Const"() {value = dense<0> : tensor} -// CHECK-DAG: %[[w_scale:.*]] = "tf.Const"() {value = dense<[4.000000e-03 -// CHECK-DAG: %[[w_zp:.*]] = "tf.Const"() {value = dense<0> : tensor<2xi32>} -// CHECK-DAG: %[[b_scale:.*]] = "tf.Const"() {value = dense<[3.200000e-05, 4.000000e-05]> : tensor<2xf32>} -// CHECK-DAG: %[[out_scale:.*]] = "tf.Const"() {value = dense<5.000000e-02> : tensor} -// CHECK-DAG: %[[out_zp:.*]] = "tf.Const"() {value = dense<-1> : tensor} -// CHECK-DAG: %[[b_quant:.*]] = "tf.Const"() {value = dense<[-62500, 75000]> : tensor<2xi32>} -// CHECK-DAG: %[[w_quant:.*]] = "tf.Const"() {value = dense<{{\[\[\[\[}}40, 20] +// CHECK-DAG: %[[w_float:.*]] = "tf.Const"() <{value = dense<{{\[\[\[\[}}1.600000e-01 +// CHECK-DAG: %[[b_float:.*]] = "tf.Const"() <{value = dense<[-2.000000e+00, 3.000000e+00]> : tensor<2xf32> +// CHECK-DAG: %[[in_scale:.*]] = "tf.Const"() <{value = dense<8.000000e-03> : tensor}> : () -> tensor +// CHECK-DAG: %[[in_zp:.*]] = "tf.Const"() <{value = dense<0> : tensor}> +// CHECK-DAG: %[[w_scale:.*]] = "tf.Const"() <{value = dense<[4.000000e-03 +// CHECK-DAG: %[[w_zp:.*]] = "tf.Const"() <{value = dense<0> : tensor<2xi32>}> +// CHECK-DAG: %[[b_scale:.*]] = "tf.Const"() <{value = dense<[3.200000e-05, 4.000000e-05]> : tensor<2xf32>} +// CHECK-DAG: %[[out_scale:.*]] = "tf.Const"() <{value = dense<5.000000e-02> : tensor}> +// CHECK-DAG: %[[out_zp:.*]] = "tf.Const"() <{value = dense<-1> : tensor}> +// CHECK-DAG: %[[b_quant:.*]] = "tf.Const"() <{value = dense<[-62500, 75000]> : tensor<2xi32>}> +// CHECK-DAG: %[[w_quant:.*]] = "tf.Const"() <{value = dense<{{\[\[\[\[}}40, 20] // CHECK-DAG: {{\[\[\[}}-87, -42] // CHECK: %[[quantize:.*]] = "tf.PartitionedCall"(%arg0, %[[in_scale]], %[[in_zp]]) @@ -58,7 +58,8 @@ module { // CHECK-LABEL: func private @composite_conv2d_with_bias_and_relu6_fn_1 // CHECK: %[[CONV2D_0:.*]] = "tf.Conv2D" -// CHECK-SAME: data_format = "NHWC", device = "", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "VALID", strides = [1, 1, 2, 1], use_cudnn_on_gpu = true +// CHECK-SAME: data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "VALID", strides = [1, 1, 2, 1], use_cudnn_on_gpu = true +// CHECK-SAME: device = "" // CHECK: %[[BIASADD_0:.*]] = "tf.BiasAdd" // CHECK: %[[RELU6_0:.*]] = "tf.Relu6" @@ -158,10 +159,10 @@ module { // CHECK: %[[conv_quant:.*]] = "tf.PartitionedCall"(%[[quantize]] // CHECK-SAME: f = @quantized_conv2d_with_bias_and_relu6_fn_0 // CHECK-SAME: (tensor<1x2x2x3xi8>, tensor<2x2x3x2xi8>, tensor<2xi32>, tensor, tensor, tensor<2xf32>, tensor<2xi32>, tensor<2xf32>, tensor<2xi32>, tensor, tensor) -> tensor<*xi8> -// CHECK: %[[cast_1:.*]] = "tf.Cast"(%[[conv_quant]]) {Truncate = false} : (tensor<*xi8>) -> tensor<*xf32> -// CHECK: %[[avgpool:.*]] = "tf.AvgPool"(%[[cast_1]]) {data_format = "NHWC", ksize = [1, 2, 2, 1], padding = "VALID", strides = [1, 1, 1, 1]} : (tensor<*xf32>) -> tensor<*xf32> +// CHECK: %[[cast_1:.*]] = "tf.Cast"(%[[conv_quant]]) <{Truncate = false}> : (tensor<*xi8>) -> tensor<*xf32> +// CHECK: %[[avgpool:.*]] = "tf.AvgPool"(%[[cast_1]]) <{data_format = "NHWC", ksize = [1, 2, 2, 1], padding = "VALID", strides = [1, 1, 1, 1]}> : (tensor<*xf32>) -> tensor<*xf32> // CHECK: %[[round:.*]] = "tf.Round"(%[[avgpool]]) : (tensor<*xf32>) -> tensor<*xf32> -// CHECK: %[[cast_2:.*]] = "tf.Cast"(%[[round]]) {Truncate = false} : (tensor<*xf32>) -> tensor<*xi8> +// CHECK: %[[cast_2:.*]] = "tf.Cast"(%[[round]]) <{Truncate = false}> : (tensor<*xf32>) -> tensor<*xi8> // CHECK: %[[dequantize:.*]] = "tf.PartitionedCall"(%[[cast_2]] // CHECK-SAME: f = @dequantize_i8 // CHECK: return %[[dequantize]] @@ -252,35 +253,35 @@ module { } // CHECK-LABE: @conv_with_dump -// CHECK-DAG: %[[w0_float:.*]] = "tf.Const"() {value = dense<{{\[\[\[\[}}-0.282878935, -0.211567819 -// CHECK-DAG: %[[b0_float:.*]] = "tf.Const"() {value = dense<[-0.0192535277, -5.998660e-03]> : tensor<2xf32>} : () -> tensor<2xf32> -// CHECK-DAG: %[[w1_float:.*]] = "tf.Const"() {value = dense<{{\[\[\[\[}}0.208403707, 0.478067577 -// CHECK-DAG: %[[b1_float:.*]] = "tf.Const"() {value = dense<[-0.0291469581, 0.0106381178]> : tensor<2xf32>} : () -> tensor<2xf32> -// CHECK-DAG: %[[w0_quantized:.*]] = "tf.Const"() {value = dense<{{\[\[\[\[}}-59, -44 -// CHECK-DAG: %[[b0_quantized:.*]] = "tf.Const"() {value = dense<[-1040, -324]> : tensor<2xi32>} : () -> tensor<2xi32> -// CHECK-DAG: %[[w1_quantized:.*]] = "tf.Const"() {value = dense<{{\[\[\[\[}}44, 100 -// CHECK-DAG: %[[b1_quantized:.*]] = "tf.Const"() {value = dense<[-4312, 1574]> : tensor<2xi32>} : () -> tensor<2xi32> -// CHECK-DAG: %[[in_scale:.*]] = "tf.Const"() {value = dense<0.00387597573> : tensor} : () -> tensor -// CHECK-DAG: %[[in_out_zp:.*]] = "tf.Const"() {value = dense<-128> : tensor} : () -> tensor -// CHECK-DAG: %[[w0_scale:.*]] = "tf.Const"() {value = dense<0.00477493973> : tensor} : () -> tensor -// CHECK-DAG: %[[w_b_zp:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor -// CHECK-DAG: %[[b0_scale:.*]] = "tf.Const"() {value = dense<1.85075514E-5> : tensor} : () -> tensor -// CHECK-DAG: %[[mid_scale:.*]] = "tf.Const"() {value = dense<0.00141507247> : tensor} : () -> tensor -// CHECK-DAG: %[[w1_scale:.*]] = "tf.Const"() {value = dense<0.00477652298> : tensor} : () -> tensor -// CHECK-DAG: %[[b1_scale:.*]] = "tf.Const"() {value = dense<6.75912588E-6> : tensor} : () -> tensor -// CHECK-DAG: %[[out_scale:.*]] = "tf.Const"() {value = dense<7.24974147E-4> : tensor} : () -> tensor -// CHECK-DAG: %[[arg_quantized:.*]] = "tf.PartitionedCall"(%arg0, %[[in_scale]], %[[in_out_zp]]) {config = "", config_proto = "", executor_type = "", f = @quantize_i8} -// CHECK-DAG: %[[conv0_quantized:.*]] = "tf.PartitionedCall"(%[[arg_quantized]], %[[w0_quantized]], %[[b0_quantized]], %[[in_scale]], %[[in_out_zp]], %[[w0_scale]], %[[w_b_zp]], %[[b0_scale]], %[[w_b_zp]], %[[mid_scale]], %[[in_out_zp]]) {config = "", config_proto = "", executor_type = "", f = @quantized_conv2d_with_bias_and_relu6_fn_1} -// CHECK-DAG: %[[conv0_dequantized:.*]] = "tf.PartitionedCall"(%[[conv0_quantized]], %[[mid_scale]], %[[in_out_zp]]) {config = "", config_proto = "", executor_type = "", f = @dequantize_i8} -// CHECK-DAG: %[[conv1_quantized:.*]] = "tf.PartitionedCall"(%[[conv0_quantized]], %[[w1_quantized]], %[[b1_quantized]], %[[mid_scale]], %[[in_out_zp]], %[[w1_scale]], %[[w_b_zp]], %[[b1_scale]], %[[w_b_zp]], %[[out_scale]], %[[in_out_zp]]) {config = "", config_proto = "", executor_type = "", f = @quantized_conv2d_with_bias_and_relu6_fn_0} -// CHECK-DAG: %[[conv1_dequantized_0:.*]] = "tf.PartitionedCall"(%[[conv1_quantized]], %[[out_scale]], %[[in_out_zp]]) {config = "", config_proto = "", executor_type = "", f = @dequantize_i8} -// CHECK-DAG: %[[conv1_dequantized_1:.*]] = "tf.PartitionedCall"(%[[conv1_quantized]], %[[out_scale]], %[[in_out_zp]]) {config = "", config_proto = "", executor_type = "", f = @dequantize_i8} +// CHECK-DAG: %[[w0_float:.*]] = "tf.Const"() <{value = dense<{{\[\[\[\[}}-0.282878935, -0.211567819 +// CHECK-DAG: %[[b0_float:.*]] = "tf.Const"() <{value = dense<[-0.0192535277, -5.998660e-03]> : tensor<2xf32>}> : () -> tensor<2xf32> +// CHECK-DAG: %[[w1_float:.*]] = "tf.Const"() <{value = dense<{{\[\[\[\[}}0.208403707, 0.478067577 +// CHECK-DAG: %[[b1_float:.*]] = "tf.Const"() <{value = dense<[-0.0291469581, 0.0106381178]> : tensor<2xf32>}> : () -> tensor<2xf32> +// CHECK-DAG: %[[w0_quantized:.*]] = "tf.Const"() <{value = dense<{{\[\[\[\[}}-59, -44 +// CHECK-DAG: %[[b0_quantized:.*]] = "tf.Const"() <{value = dense<[-1040, -324]> : tensor<2xi32>}> : () -> tensor<2xi32> +// CHECK-DAG: %[[w1_quantized:.*]] = "tf.Const"() <{value = dense<{{\[\[\[\[}}44, 100 +// CHECK-DAG: %[[b1_quantized:.*]] = "tf.Const"() <{value = dense<[-4312, 1574]> : tensor<2xi32>}> : () -> tensor<2xi32> +// CHECK-DAG: %[[in_scale:.*]] = "tf.Const"() <{value = dense<0.00387597573> : tensor}> : () -> tensor +// CHECK-DAG: %[[in_out_zp:.*]] = "tf.Const"() <{value = dense<-128> : tensor}> : () -> tensor +// CHECK-DAG: %[[w0_scale:.*]] = "tf.Const"() <{value = dense<0.00477493973> : tensor}> : () -> tensor +// CHECK-DAG: %[[w_b_zp:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor +// CHECK-DAG: %[[b0_scale:.*]] = "tf.Const"() <{value = dense<1.85075514E-5> : tensor}> : () -> tensor +// CHECK-DAG: %[[mid_scale:.*]] = "tf.Const"() <{value = dense<0.00141507247> : tensor}> : () -> tensor +// CHECK-DAG: %[[w1_scale:.*]] = "tf.Const"() <{value = dense<0.00477652298> : tensor}> : () -> tensor +// CHECK-DAG: %[[b1_scale:.*]] = "tf.Const"() <{value = dense<6.75912588E-6> : tensor}> : () -> tensor +// CHECK-DAG: %[[out_scale:.*]] = "tf.Const"() <{value = dense<7.24974147E-4> : tensor}> : () -> tensor +// CHECK-DAG: %[[arg_quantized:.*]] = "tf.PartitionedCall"(%arg0, %[[in_scale]], %[[in_out_zp]]) <{config = "", config_proto = "", executor_type = "", f = @quantize_i8}> +// CHECK-DAG: %[[conv0_quantized:.*]] = "tf.PartitionedCall"(%[[arg_quantized]], %[[w0_quantized]], %[[b0_quantized]], %[[in_scale]], %[[in_out_zp]], %[[w0_scale]], %[[w_b_zp]], %[[b0_scale]], %[[w_b_zp]], %[[mid_scale]], %[[in_out_zp]]) <{config = "", config_proto = "", executor_type = "", f = @quantized_conv2d_with_bias_and_relu6_fn_1}> +// CHECK-DAG: %[[conv0_dequantized:.*]] = "tf.PartitionedCall"(%[[conv0_quantized]], %[[mid_scale]], %[[in_out_zp]]) <{config = "", config_proto = "", executor_type = "", f = @dequantize_i8}> +// CHECK-DAG: %[[conv1_quantized:.*]] = "tf.PartitionedCall"(%[[conv0_quantized]], %[[w1_quantized]], %[[b1_quantized]], %[[mid_scale]], %[[in_out_zp]], %[[w1_scale]], %[[w_b_zp]], %[[b1_scale]], %[[w_b_zp]], %[[out_scale]], %[[in_out_zp]]) <{config = "", config_proto = "", executor_type = "", f = @quantized_conv2d_with_bias_and_relu6_fn_0}> +// CHECK-DAG: %[[conv1_dequantized_0:.*]] = "tf.PartitionedCall"(%[[conv1_quantized]], %[[out_scale]], %[[in_out_zp]]) <{config = "", config_proto = "", executor_type = "", f = @dequantize_i8}> +// CHECK-DAG: %[[conv1_dequantized_1:.*]] = "tf.PartitionedCall"(%[[conv1_quantized]], %[[out_scale]], %[[in_out_zp]]) <{config = "", config_proto = "", executor_type = "", f = @dequantize_i8}> // CHECK-DAG: %[[identity:.*]] = "tf.Identity"(%[[conv1_dequantized_1]]) -// CHECK-DAG: %[[conv0_float:.*]] = "tf.PartitionedCall"(%arg0, %[[w0_float]], %[[b0_float]]) {config = "", config_proto = "", device = "", executor_type = "", f = @composite_conv2d_with_bias_and_relu6_fn_2_00} -// CHECK-DAG: %[[conv1_float:.*]] = "tf.PartitionedCall"(%[[conv0_dequantized]], %[[w1_float]], %[[b1_float]]) {config = "", config_proto = "", device = "", executor_type = "", f = @composite_conv2d_with_bias_and_relu6_fn_1_00} -// CHECK-DAG: "tf.DumpTensor"(%[[conv0_dequantized]]) {device = "", enabled = true, file_name = "quantized_tensor_data.pb", func_name = "conv_with_dump", log_dir_path = "/tmp/dumps/composite_conv2d_with_bias_and_relu6_fn_2", node_name = "Conv2D"} -// CHECK-DAG: "tf.DumpTensor"(%[[conv0_float]]) {device = "", enabled = true, file_name = "unquantized_tensor_data.pb", func_name = "conv_with_dump", log_dir_path = "/tmp/dumps/composite_conv2d_with_bias_and_relu6_fn_2", node_name = "Conv2D"} -// CHECK-DAG: "tf.DumpTensor"(%[[conv1_dequantized_0]]) {device = "", enabled = true, file_name = "quantized_tensor_data.pb", func_name = "conv_with_dump", log_dir_path = "/tmp/dumps/composite_conv2d_with_bias_and_relu6_fn_1", node_name = "Conv2D_1"} -// CHECK-DAG: "tf.DumpTensor"(%[[conv1_float]]) {device = "", enabled = true, file_name = "unquantized_tensor_data.pb", func_name = "conv_with_dump", log_dir_path = "/tmp/dumps/composite_conv2d_with_bias_and_relu6_fn_1", node_name = "Conv2D_1"} +// CHECK-DAG: %[[conv0_float:.*]] = "tf.PartitionedCall"(%arg0, %[[w0_float]], %[[b0_float]]) <{config = "", config_proto = "", executor_type = "", f = @composite_conv2d_with_bias_and_relu6_fn_2_00}> {device = ""} +// CHECK-DAG: %[[conv1_float:.*]] = "tf.PartitionedCall"(%[[conv0_dequantized]], %[[w1_float]], %[[b1_float]]) <{config = "", config_proto = "", executor_type = "", f = @composite_conv2d_with_bias_and_relu6_fn_1_00}> {device = ""} +// CHECK-DAG: "tf.DumpTensor"(%[[conv0_dequantized]]) <{enabled = true, file_name = "quantized_tensor_data.pb", func_name = "conv_with_dump", log_dir_path = "/tmp/dumps/composite_conv2d_with_bias_and_relu6_fn_2", node_name = "Conv2D"}> {device = ""} +// CHECK-DAG: "tf.DumpTensor"(%[[conv0_float]]) <{enabled = true, file_name = "unquantized_tensor_data.pb", func_name = "conv_with_dump", log_dir_path = "/tmp/dumps/composite_conv2d_with_bias_and_relu6_fn_2", node_name = "Conv2D"}> {device = ""} +// CHECK-DAG: "tf.DumpTensor"(%[[conv1_dequantized_0]]) <{enabled = true, file_name = "quantized_tensor_data.pb", func_name = "conv_with_dump", log_dir_path = "/tmp/dumps/composite_conv2d_with_bias_and_relu6_fn_1", node_name = "Conv2D_1"}> {device = ""} +// CHECK-DAG: "tf.DumpTensor"(%[[conv1_float]]) <{enabled = true, file_name = "unquantized_tensor_data.pb", func_name = "conv_with_dump", log_dir_path = "/tmp/dumps/composite_conv2d_with_bias_and_relu6_fn_1", node_name = "Conv2D_1"}> {device = ""} // CHECK-DAG: return %[[identity]] } diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/tests/quantize_composite_functions_drq.mlir b/tensorflow/compiler/mlir/quantization/tensorflow/tests/quantize_composite_functions_drq.mlir index 0731e6bfe8e975..6ec99282a8e252 100644 --- a/tensorflow/compiler/mlir/quantization/tensorflow/tests/quantize_composite_functions_drq.mlir +++ b/tensorflow/compiler/mlir/quantization/tensorflow/tests/quantize_composite_functions_drq.mlir @@ -13,11 +13,11 @@ module { } // CHECK-LABEL: func @matmul -// CHECK-DAG: %[[q_w:.*]] = "tf.Const"() {value = #tf_type : tensor} : () -> tensor -// CHECK-DAG: %[[zp:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor -// CHECK: %0 = "tf.PartitionedCall"(%arg0, %[[q_w]], %[[scale]], %[[zp]]) {config = "", config_proto = "", executor_type = "", -// CHECK-SAME: f = @quantized_matmul_fn_0} : (tensor<2x12xf32>, tensor<12x2x!tf_type.qint8>, tensor, tensor) -> tensor<*xf32> +// CHECK-DAG: %[[q_w:.*]] = "tf.Const"() <{value = #tf_type : tensor}> : () -> tensor +// CHECK-DAG: %[[zp:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor +// CHECK: %0 = "tf.PartitionedCall"(%arg0, %[[q_w]], %[[scale]], %[[zp]]) <{config = "", config_proto = "", executor_type = "", +// CHECK-SAME: f = @quantized_matmul_fn_0}> : (tensor<2x12xf32>, tensor<12x2x!tf_type.qint8>, tensor, tensor) -> tensor<*xf32> // CHECK-LABEL: func private @quantized_matmul_fn_0 // CHECK: %0 = "tf.UniformQuantizedDotHybrid"(%arg0, %arg1, %arg2, %arg3) @@ -48,11 +48,11 @@ module { } // CHECK-LABEL: func @conv -// CHECK-DAG: %[[q_w:.*]] = "tf.Const"() {value = #tf_type : tensor} : () -> tensor -// CHECK-DAG: %[[w_zp:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor -// CHECK: %[[quantize_1:.*]] = "tf.PartitionedCall"(%arg0, %[[q_w]], %[[w_scale]], %[[w_zp]]) {config = "", config_proto = "", executor_type = "", f = @quantized_conv2d_fn_1} : (tensor<1x2x2x3xf32>, tensor<2x3x3x2x!tf_type.qint8>, tensor, tensor) -> tensor<*xf32> -// CHECK: %[[quantize_2:.*]] = "tf.PartitionedCall"(%arg0, %[[q_w]], %[[w_scale]], %[[w_zp]]) {config = "", config_proto = "", executor_type = "", f = @quantized_conv2d_fn_0} : (tensor<1x2x2x3xf32>, tensor<2x3x3x2x!tf_type.qint8>, tensor, tensor) -> tensor<*xf32> +// CHECK-DAG: %[[q_w:.*]] = "tf.Const"() <{value = #tf_type : tensor}> : () -> tensor +// CHECK-DAG: %[[w_zp:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor +// CHECK: %[[quantize_1:.*]] = "tf.PartitionedCall"(%arg0, %[[q_w]], %[[w_scale]], %[[w_zp]]) <{config = "", config_proto = "", executor_type = "", f = @quantized_conv2d_fn_1}> : (tensor<1x2x2x3xf32>, tensor<2x3x3x2x!tf_type.qint8>, tensor, tensor) -> tensor<*xf32> +// CHECK: %[[quantize_2:.*]] = "tf.PartitionedCall"(%arg0, %[[q_w]], %[[w_scale]], %[[w_zp]]) <{config = "", config_proto = "", executor_type = "", f = @quantized_conv2d_fn_0}> : (tensor<1x2x2x3xf32>, tensor<2x3x3x2x!tf_type.qint8>, tensor, tensor) -> tensor<*xf32> // CHECK: return %[[quantize_1]], %[[quantize_2]] // CHECK-LABEL: func private @quantized_conv2d_fn_0 @@ -102,17 +102,17 @@ module { } // CHECK-LABEL: func @depthwise_conv -// CHECK-DAG: %[[bias:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<3xf32>} : () -> tensor<3xf32> -// CHECK-DAG: %[[q_w1:.*]] = "tf.Const"() {value = #tf_type : tensor<3xf32>}> : () -> tensor<3xf32> +// CHECK-DAG: %[[q_w1:.*]] = "tf.Const"() <{value = #tf_type tensor<2x3x1x3x!tf_type.qint8> -// CHECK-DAG: %[[q_w2:.*]] = "tf.Const"() {value = #tf_type tensor<2x3x1x6x!tf_type.qint8> -// CHECK-DAG: %[[w_scale:.*]] = "tf.Const"() {value = dense<0.0236220472> : tensor} : () -> tensor -// CHECK-DAG: %[[w_zp:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor +// CHECK-DAG: %[[w_scale:.*]] = "tf.Const"() <{value = dense<0.0236220472> : tensor}> : () -> tensor +// CHECK-DAG: %[[w_zp:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor -// CHECK: %[[quantize_1:.*]] = "tf.PartitionedCall"(%arg0, %[[q_w1]], %[[w_scale]], %[[w_zp]]) {config = "", config_proto = "", executor_type = "", f = @quantized_depthwise_conv2d_fn_1} : (tensor<1x3x4x3xf32>, tensor<2x3x1x3x!tf_type.qint8>, tensor, tensor) -> tensor<*xf32> -// CHECK: %[[quantize_1_add:.*]] = "tf.BiasAdd"(%[[quantize_1]], %[[bias]]) {data_format = "NHWC", device = ""} : (tensor<*xf32>, tensor<3xf32>) -> tensor<*xf32> -// CHECK: %[[quantize_2:.*]] = "tf.PartitionedCall"(%arg0, %[[q_w2]], %[[w_scale]], %[[w_zp]]) {config = "", config_proto = "", executor_type = "", f = @quantized_depthwise_conv2d_fn_0} : (tensor<1x3x4x3xf32>, tensor<2x3x1x6x!tf_type.qint8>, tensor, tensor) -> tensor<*xf32> +// CHECK: %[[quantize_1:.*]] = "tf.PartitionedCall"(%arg0, %[[q_w1]], %[[w_scale]], %[[w_zp]]) <{config = "", config_proto = "", executor_type = "", f = @quantized_depthwise_conv2d_fn_1}> : (tensor<1x3x4x3xf32>, tensor<2x3x1x3x!tf_type.qint8>, tensor, tensor) -> tensor<*xf32> +// CHECK: %[[quantize_1_add:.*]] = "tf.BiasAdd"(%[[quantize_1]], %[[bias]]) <{data_format = "NHWC"}> {device = ""} : (tensor<*xf32>, tensor<3xf32>) -> tensor<*xf32> +// CHECK: %[[quantize_2:.*]] = "tf.PartitionedCall"(%arg0, %[[q_w2]], %[[w_scale]], %[[w_zp]]) <{config = "", config_proto = "", executor_type = "", f = @quantized_depthwise_conv2d_fn_0}> : (tensor<1x3x4x3xf32>, tensor<2x3x1x6x!tf_type.qint8>, tensor, tensor) -> tensor<*xf32> // CHECK: return %[[quantize_1_add]], %[[quantize_2]] // CHECK-LABEL: func private @quantized_depthwise_conv2d_fn_0 diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/tests/quantize_composite_functions_weight_only.mlir b/tensorflow/compiler/mlir/quantization/tensorflow/tests/quantize_composite_functions_weight_only.mlir index 8c0786178ee37e..ba8f21380b33e6 100644 --- a/tensorflow/compiler/mlir/quantization/tensorflow/tests/quantize_composite_functions_weight_only.mlir +++ b/tensorflow/compiler/mlir/quantization/tensorflow/tests/quantize_composite_functions_weight_only.mlir @@ -15,19 +15,19 @@ module { } // PerTensor-LABEL: func @matmul -// PerTensor-DAG: %[[q_w:.*]] = "tf.Const"() {value = dense<0> : tensor<12x2xi8>} : () -> tensor<12x2xi8> -// PerTensor-DAG: %[[scale:.*]] = "tf.Const"() {value = dense<3.93700805E-9> : tensor} : () -> tensor -// PerTensor-DAG: %[[zp:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor -// PerTensor: %[[out:.*]] = "tf.PartitionedCall"(%arg0, %[[q_w]], %[[scale]], %[[zp]]) {config = "", config_proto = "", executor_type = "", -// PerTensor-SAME: f = @quantized_matmul_fn_0} : (tensor<2x12xf32>, tensor<12x2xi8>, tensor, tensor) -> tensor<*xf32> +// PerTensor-DAG: %[[q_w:.*]] = "tf.Const"() <{value = dense<0> : tensor<12x2xi8>}> : () -> tensor<12x2xi8> +// PerTensor-DAG: %[[scale:.*]] = "tf.Const"() <{value = dense<3.93700805E-9> : tensor}> : () -> tensor +// PerTensor-DAG: %[[zp:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor +// PerTensor: %[[out:.*]] = "tf.PartitionedCall"(%arg0, %[[q_w]], %[[scale]], %[[zp]]) <{config = "", config_proto = "", executor_type = "", +// PerTensor-SAME: f = @quantized_matmul_fn_0}> : (tensor<2x12xf32>, tensor<12x2xi8>, tensor, tensor) -> tensor<*xf32> // PerTensor: return %[[out]] // PerChannel-LABEL: func @matmul -// PerChannel-DAG: %[[q_w:.*]] = "tf.Const"() {value = dense<0> : tensor<12x2xi8>} : () -> tensor<12x2xi8> -// PerChannel-DAG: %[[scale:.*]] = "tf.Const"() {value = dense<3.93700805E-9> : tensor} : () -> tensor -// PerChannel-DAG: %[[zp:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor -// PerChannel: %[[out:.*]] = "tf.PartitionedCall"(%arg0, %[[q_w]], %[[scale]], %[[zp]]) {config = "", config_proto = "", executor_type = "", -// PerChannel-SAME: f = @quantized_matmul_fn_0} : (tensor<2x12xf32>, tensor<12x2xi8>, tensor, tensor) -> tensor<*xf32> +// PerChannel-DAG: %[[q_w:.*]] = "tf.Const"() <{value = dense<0> : tensor<12x2xi8>}> : () -> tensor<12x2xi8> +// PerChannel-DAG: %[[scale:.*]] = "tf.Const"() <{value = dense<3.93700805E-9> : tensor}> : () -> tensor +// PerChannel-DAG: %[[zp:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor +// PerChannel: %[[out:.*]] = "tf.PartitionedCall"(%arg0, %[[q_w]], %[[scale]], %[[zp]]) <{config = "", config_proto = "", executor_type = "", +// PerChannel-SAME: f = @quantized_matmul_fn_0}> : (tensor<2x12xf32>, tensor<12x2xi8>, tensor, tensor) -> tensor<*xf32> // PerChannel: return %[[out]] // ----- @@ -51,23 +51,23 @@ module { } // PerTensor-LABEL: func @conv -// PerTensor-DAG: %[[q_w:.*]] = "tf.Const"() {value = dense<{{[0-9]+}}> : tensor<2x3x3x2xi8>} : () -> tensor<2x3x3x2xi8> -// PerTensor-DAG: %[[scale:.*]] = "tf.Const"() {value = dense<{{[0-9\.Ee\+\-]+}}> : tensor} : () -> tensor -// PerTensor-DAG: %[[zp:.*]] = "tf.Const"() {value = dense<{{[0-9]+}}> : tensor} : () -> tensor -// PerTensor: %[[out_1:.*]] = "tf.PartitionedCall"(%arg0, %[[q_w]], %[[scale]], %[[zp]]) {config = "", config_proto = "", executor_type = "", -// PerTensor-SAME: f = @quantized_conv2d_fn_1} : (tensor<1x2x2x3xf32>, tensor<2x3x3x2xi8>, tensor, tensor) -> tensor<*xf32> -// PerTensor: %[[out_2:.*]] = "tf.PartitionedCall"(%arg0, %[[q_w]], %[[scale]], %[[zp]]) {config = "", config_proto = "", executor_type = "", -// PerTensor-SAME: f = @quantized_conv2d_fn_0} : (tensor<1x2x2x3xf32>, tensor<2x3x3x2xi8>, tensor, tensor) -> tensor<*xf32> +// PerTensor-DAG: %[[q_w:.*]] = "tf.Const"() <{value = dense<{{[0-9]+}}> : tensor<2x3x3x2xi8>}> : () -> tensor<2x3x3x2xi8> +// PerTensor-DAG: %[[scale:.*]] = "tf.Const"() <{value = dense<{{[0-9\.Ee\+\-]+}}> : tensor}> : () -> tensor +// PerTensor-DAG: %[[zp:.*]] = "tf.Const"() <{value = dense<{{[0-9]+}}> : tensor}> : () -> tensor +// PerTensor: %[[out_1:.*]] = "tf.PartitionedCall"(%arg0, %[[q_w]], %[[scale]], %[[zp]]) <{config = "", config_proto = "", executor_type = "", +// PerTensor-SAME: f = @quantized_conv2d_fn_1}> : (tensor<1x2x2x3xf32>, tensor<2x3x3x2xi8>, tensor, tensor) -> tensor<*xf32> +// PerTensor: %[[out_2:.*]] = "tf.PartitionedCall"(%arg0, %[[q_w]], %[[scale]], %[[zp]]) <{config = "", config_proto = "", executor_type = "", +// PerTensor-SAME: f = @quantized_conv2d_fn_0}> : (tensor<1x2x2x3xf32>, tensor<2x3x3x2xi8>, tensor, tensor) -> tensor<*xf32> // PerTensor: return %[[out_1]], %[[out_2]] // PerChannel-LABEL: func @conv -// PerChannel-DAG: %[[q_w:.*]] = "tf.Const"() {value = dense<{{[0-9]+}}> : tensor<2x3x3x2xi8>} : () -> tensor<2x3x3x2xi8> -// PerChannel-DAG: %[[scale:.*]] = "tf.Const"() {value = dense<{{[0-9\.Ee\+\-]+}}> : tensor<2xf32>} : () -> tensor<2xf32> -// PerChannel-DAG: %[[zp:.*]] = "tf.Const"() {value = dense<{{[0-9]+}}> : tensor<2xi32>} : () -> tensor<2xi32> -// PerChannel: %[[out_1:.*]] = "tf.PartitionedCall"(%arg0, %[[q_w]], %[[scale]], %[[zp]]) {config = "", config_proto = "", executor_type = "", -// PerChannel-SAME: f = @quantized_conv2d_fn_1} : (tensor<1x2x2x3xf32>, tensor<2x3x3x2xi8>, tensor<2xf32>, tensor<2xi32>) -> tensor<*xf32> -// PerChannel: %[[out_2:.*]] = "tf.PartitionedCall"(%arg0, %[[q_w]], %[[scale]], %[[zp]]) {config = "", config_proto = "", executor_type = "", -// PerChannel-SAME: f = @quantized_conv2d_fn_0} : (tensor<1x2x2x3xf32>, tensor<2x3x3x2xi8>, tensor<2xf32>, tensor<2xi32>) -> tensor<*xf32> +// PerChannel-DAG: %[[q_w:.*]] = "tf.Const"() <{value = dense<{{[0-9]+}}> : tensor<2x3x3x2xi8>}> : () -> tensor<2x3x3x2xi8> +// PerChannel-DAG: %[[scale:.*]] = "tf.Const"() <{value = dense<{{[0-9\.Ee\+\-]+}}> : tensor<2xf32>}> : () -> tensor<2xf32> +// PerChannel-DAG: %[[zp:.*]] = "tf.Const"() <{value = dense<{{[0-9]+}}> : tensor<2xi32>}> : () -> tensor<2xi32> +// PerChannel: %[[out_1:.*]] = "tf.PartitionedCall"(%arg0, %[[q_w]], %[[scale]], %[[zp]]) <{config = "", config_proto = "", executor_type = "", +// PerChannel-SAME: f = @quantized_conv2d_fn_1}> : (tensor<1x2x2x3xf32>, tensor<2x3x3x2xi8>, tensor<2xf32>, tensor<2xi32>) -> tensor<*xf32> +// PerChannel: %[[out_2:.*]] = "tf.PartitionedCall"(%arg0, %[[q_w]], %[[scale]], %[[zp]]) <{config = "", config_proto = "", executor_type = "", +// PerChannel-SAME: f = @quantized_conv2d_fn_0}> : (tensor<1x2x2x3xf32>, tensor<2x3x3x2xi8>, tensor<2xf32>, tensor<2xi32>) -> tensor<*xf32> // PerChannel: return %[[out_1]], %[[out_2]] } @@ -98,30 +98,30 @@ module { } // PerTensor-LABEL: func @depthwise_conv -// PerTensor-DAG: %[[q_w1:.*]] = "tf.Const"() {value = dense<127> : tensor<2x3x3x1xi8>} -// PerTensor-DAG: %[[q_w2:.*]] = "tf.Const"() {value = dense<127> : tensor<2x3x3x2xi8>} : () -> tensor<2x3x3x2xi8> -// PerTensor-DAG: %[[scale:.*]] = "tf.Const"() {value = dense<0.0236220472> : tensor} : () -> tensor -// PerTensor-DAG: %[[zp:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor -// PerTensor-DAG: %[[bias:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<3xf32>} -// PerTensor: %[[out_1:.*]] = "tf.PartitionedCall"(%arg0, %[[q_w1]], %[[scale]], %[[zp]]) {config = "", config_proto = "", executor_type = "", -// PerTensor-SAME: f = @quantized_depthwise_conv2d_fn_1} : (tensor<1x3x4x3xf32>, tensor<2x3x3x1xi8>, tensor, tensor) -> tensor<*xf32> +// PerTensor-DAG: %[[q_w1:.*]] = "tf.Const"() <{value = dense<127> : tensor<2x3x3x1xi8>}> +// PerTensor-DAG: %[[q_w2:.*]] = "tf.Const"() <{value = dense<127> : tensor<2x3x3x2xi8>}> : () -> tensor<2x3x3x2xi8> +// PerTensor-DAG: %[[scale:.*]] = "tf.Const"() <{value = dense<0.0236220472> : tensor}> : () -> tensor +// PerTensor-DAG: %[[zp:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor +// PerTensor-DAG: %[[bias:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<3xf32>}> +// PerTensor: %[[out_1:.*]] = "tf.PartitionedCall"(%arg0, %[[q_w1]], %[[scale]], %[[zp]]) <{config = "", config_proto = "", executor_type = "", +// PerTensor-SAME: f = @quantized_depthwise_conv2d_fn_1}> : (tensor<1x3x4x3xf32>, tensor<2x3x3x1xi8>, tensor, tensor) -> tensor<*xf32> // PerTensor: %[[out_1_add:.*]] = "tf.BiasAdd"(%[[out_1]], %[[bias]]) -// PerTensor: %[[out_2:.*]] = "tf.PartitionedCall"(%arg0, %[[q_w2]], %[[scale]], %[[zp]]) {config = "", config_proto = "", executor_type = "", -// PerTensor-SAME: f = @quantized_depthwise_conv2d_fn_0} : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xi8>, tensor, tensor) -> tensor<*xf32> +// PerTensor: %[[out_2:.*]] = "tf.PartitionedCall"(%arg0, %[[q_w2]], %[[scale]], %[[zp]]) <{config = "", config_proto = "", executor_type = "", +// PerTensor-SAME: f = @quantized_depthwise_conv2d_fn_0}> : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xi8>, tensor, tensor) -> tensor<*xf32> // PerTensor: return %[[out_1_add]], %[[out_2]] // PerChannel-LABEL: func @depthwise_conv -// PerChannel-DAG: %[[bias1:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<3xf32>} : () -> tensor<3xf32> -// PerChannel-DAG: %[[q_w1:.*]] = "tf.Const"() {value = dense<{{[0-9]+}}> : tensor<2x3x3x1xi8>} : () -> tensor<2x3x3x1xi8> -// PerChannel-DAG: %[[q_w2:.*]] = "tf.Const"() {value = dense<{{[0-9]+}}> : tensor<2x3x3x2xi8>} : () -> tensor<2x3x3x2xi8> -// PerChannel-DAG: %[[scale1:.*]] = "tf.Const"() {value = dense<{{[0-9\.Ee\+\-]+}}> : tensor<3xf32>} : () -> tensor<3xf32> -// PerChannel-DAG: %[[scale2:.*]] = "tf.Const"() {value = dense<{{[0-9\.Ee\+\-]+}}> : tensor<6xf32>} : () -> tensor<6xf32> -// PerChannel-DAG: %[[zp1:.*]] = "tf.Const"() {value = dense<{{[0-9]+}}> : tensor<3xi32>} : () -> tensor<3xi32> -// PerChannel-DAG: %[[zp2:.*]] = "tf.Const"() {value = dense<{{[0-9]+}}> : tensor<6xi32>} : () -> tensor<6xi32> -// PerChannel: %[[out_1:.*]] = "tf.PartitionedCall"(%arg0, %[[q_w1]], %[[scale1]], %[[zp1]]) {config = "", config_proto = "", executor_type = "", -// PerChannel-SAME: f = @quantized_depthwise_conv2d_fn_1} : (tensor<1x3x4x3xf32>, tensor<2x3x3x1xi8>, tensor<3xf32>, tensor<3xi32>) -> tensor<*xf32> +// PerChannel-DAG: %[[bias1:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<3xf32>}> : () -> tensor<3xf32> +// PerChannel-DAG: %[[q_w1:.*]] = "tf.Const"() <{value = dense<{{[0-9]+}}> : tensor<2x3x3x1xi8>}> : () -> tensor<2x3x3x1xi8> +// PerChannel-DAG: %[[q_w2:.*]] = "tf.Const"() <{value = dense<{{[0-9]+}}> : tensor<2x3x3x2xi8>}> : () -> tensor<2x3x3x2xi8> +// PerChannel-DAG: %[[scale1:.*]] = "tf.Const"() <{value = dense<{{[0-9\.Ee\+\-]+}}> : tensor<3xf32>}> : () -> tensor<3xf32> +// PerChannel-DAG: %[[scale2:.*]] = "tf.Const"() <{value = dense<{{[0-9\.Ee\+\-]+}}> : tensor<6xf32>}> : () -> tensor<6xf32> +// PerChannel-DAG: %[[zp1:.*]] = "tf.Const"() <{value = dense<{{[0-9]+}}> : tensor<3xi32>}> : () -> tensor<3xi32> +// PerChannel-DAG: %[[zp2:.*]] = "tf.Const"() <{value = dense<{{[0-9]+}}> : tensor<6xi32>}> : () -> tensor<6xi32> +// PerChannel: %[[out_1:.*]] = "tf.PartitionedCall"(%arg0, %[[q_w1]], %[[scale1]], %[[zp1]]) <{config = "", config_proto = "", executor_type = "", +// PerChannel-SAME: f = @quantized_depthwise_conv2d_fn_1}> : (tensor<1x3x4x3xf32>, tensor<2x3x3x1xi8>, tensor<3xf32>, tensor<3xi32>) -> tensor<*xf32> // PerChannel: %[[out_1_add:.*]] = "tf.BiasAdd"(%[[out_1]], %[[bias1]]) -// PerChannel: %[[out_2:.*]] = "tf.PartitionedCall"(%arg0, %[[q_w2]], %[[scale2]], %[[zp2]]) {config = "", config_proto = "", executor_type = "", -// PerChannel-SAME: f = @quantized_depthwise_conv2d_fn_0} : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xi8>, tensor<6xf32>, tensor<6xi32>) -> tensor<*xf32> +// PerChannel: %[[out_2:.*]] = "tf.PartitionedCall"(%arg0, %[[q_w2]], %[[scale2]], %[[zp2]]) <{config = "", config_proto = "", executor_type = "", +// PerChannel-SAME: f = @quantized_depthwise_conv2d_fn_0}> : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xi8>, tensor<6xf32>, tensor<6xi32>) -> tensor<*xf32> // PerChannel: return %[[out_1_add]], %[[out_2]] } diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/tests/quantize_composite_functions_xla.mlir b/tensorflow/compiler/mlir/quantization/tensorflow/tests/quantize_composite_functions_xla.mlir index 38f0662e097063..38b41273f3016b 100644 --- a/tensorflow/compiler/mlir/quantization/tensorflow/tests/quantize_composite_functions_xla.mlir +++ b/tensorflow/compiler/mlir/quantization/tensorflow/tests/quantize_composite_functions_xla.mlir @@ -1,4 +1,5 @@ // RUN: tf-quant-opt %s -split-input-file -quant-insert-quantized-functions -quant-quantize-composite-functions='target-opset=XLA' | FileCheck %s +// RUN: tf-quant-opt %s -split-input-file -quant-insert-quantized-functions -quant-quantize-composite-functions='target-opset=XLA enable-per-channel-quantization=true' | FileCheck --check-prefix=PerChannel %s module { func.func @conv_with_single_layer(%arg0: tensor<1x2x2x3xf32>) -> (tensor<*xf32>) { @@ -32,7 +33,7 @@ module { // CHECK-LABEL: func private @quantized_conv2d_with_bias_and_relu6_float_output_fn_0 // CHECK-SAME: (%arg0: tensor<1x2x2x3xi8>, %arg1: tensor<2x2x3x2xi8>, %arg2: tensor<2xi32>, %arg3: tensor, %arg4: tensor, %arg5: tensor<2xf32>, %arg6: tensor<2xi32>, %arg7: tensor<2xf32>, %arg8: tensor<2xi32>, %arg9: tensor, %arg10: tensor) -> tensor<*xf32> // CHECK: %[[CONV2D_0:.*]] = "tf.Conv2D" -// CHECK-SAME: {dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "VALID", strides = [1, 1, 2, 1], use_cudnn_on_gpu = true} +// CHECK-SAME: <{dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "VALID", strides = [1, 1, 2, 1], use_cudnn_on_gpu = true}> // CHECK: -------- Quantization Summary -------- // CHECK: Number of quantized layers in the model @@ -123,7 +124,7 @@ module { // CHECK: %[[conv_quant:.*]] = "tf.PartitionedCall"(%[[quantize]] // CHECK-SAME: f = @quantized_conv2d_with_bias_and_relu6_fn_0 // CHECK-SAME: (tensor<1x2x2x3xi8>, tensor<2x2x3x2xi8>, tensor<2xi32>, tensor, tensor, tensor<2xf32>, tensor<2xi32>, tensor<2xf32>, tensor<2xi32>, tensor, tensor) -> tensor<*xi8> -// CHECK: %[[maxpool:.*]] = "tf.MaxPool"(%[[conv_quant]]) {data_format = "NHWC", ksize = [1, 2, 2, 1], padding = "VALID", strides = [1, 1, 1, 1]} : (tensor<*xi8>) -> tensor<*xi8> +// CHECK: %[[maxpool:.*]] = "tf.MaxPool"(%[[conv_quant]]) <{data_format = "NHWC", ksize = [1, 2, 2, 1], padding = "VALID", strides = [1, 1, 1, 1]}> : (tensor<*xi8>) -> tensor<*xi8> // CHECK: %[[dequantize:.*]] = "tf.PartitionedCall"(%[[maxpool]] // CHECK-SAME: f = @dequantize_i8 // CHECK: return %[[dequantize]] @@ -297,34 +298,94 @@ module { func.return %2 : tensor<*xf32> } -// CHECK-LABE: @conv_with_dump -// CHECK-DAG: %[[w0_float:.*]] = "tf.Const"() {value = dense<{{\[\[\[\[}}-0.282878935, -0.211567819 -// CHECK-DAG: %[[b0_float:.*]] = "tf.Const"() {value = dense<[-0.0192535277, -5.998660e-03]> : tensor<2xf32>} : () -> tensor<2xf32> -// CHECK-DAG: %[[w1_float:.*]] = "tf.Const"() {value = dense<{{\[\[\[\[}}0.208403707, 0.478067577 -// CHECK-DAG: %[[b1_float:.*]] = "tf.Const"() {value = dense<[-0.0291469581, 0.0106381178]> : tensor<2xf32>} : () -> tensor<2xf32> -// CHECK-DAG: %[[w0_quantized:.*]] = "tf.Const"() {value = dense<{{\[\[\[\[}}-59, -44 -// CHECK-DAG: %[[b0_quantized:.*]] = "tf.Const"() {value = dense<[-1040, -324]> : tensor<2xi32>} : () -> tensor<2xi32> -// CHECK-DAG: %[[w1_quantized:.*]] = "tf.Const"() {value = dense<{{\[\[\[\[}}44, 100 -// CHECK-DAG: %[[b1_quantized:.*]] = "tf.Const"() {value = dense<[-4312, 1574]> : tensor<2xi32>} : () -> tensor<2xi32> -// CHECK-DAG: %[[in_scale:.*]] = "tf.Const"() {value = dense<0.00387597573> : tensor} : () -> tensor -// CHECK-DAG: %[[in_out_zp:.*]] = "tf.Const"() {value = dense<-128> : tensor} : () -> tensor -// CHECK-DAG: %[[w0_scale:.*]] = "tf.Const"() {value = dense<0.00477493973> : tensor} : () -> tensor -// CHECK-DAG: %[[w_b_zp:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor -// CHECK-DAG: %[[b0_scale:.*]] = "tf.Const"() {value = dense<1.85075514E-5> : tensor} : () -> tensor -// CHECK-DAG: %[[mid_scale:.*]] = "tf.Const"() {value = dense<0.00141507247> : tensor} : () -> tensor -// CHECK-DAG: %[[w1_scale:.*]] = "tf.Const"() {value = dense<0.00477652298> : tensor} : () -> tensor -// CHECK-DAG: %[[b1_scale:.*]] = "tf.Const"() {value = dense<6.75912588E-6> : tensor} : () -> tensor -// CHECK-DAG: %[[out_scale:.*]] = "tf.Const"() {value = dense<7.24974147E-4> : tensor} : () -> tensor -// CHECK-DAG: %[[quantized:.*]] = "tf.PartitionedCall"(%arg0, %[[in_scale]], %[[in_out_zp]]) {config = "", config_proto = "", executor_type = "", f = @quantize_i8} -// CHECK-DAG: %[[conv0_dequantized:.*]] = "tf.PartitionedCall"(%[[quantized]], %[[w0_quantized]], %[[b0_quantized]], %[[in_scale]], %[[in_out_zp]], %[[w0_scale]], %[[w_b_zp]], %[[b0_scale]], %[[w_b_zp]], %[[mid_scale]], %[[in_out_zp]]) {config = "", config_proto = "", executor_type = "", f = @quantized_conv2d_with_bias_and_relu6_float_output_fn_1} -// CHECK-DAG: %[[conv0_quantized:.*]] = "tf.PartitionedCall"(%[[quantized]], %[[w0_quantized]], %[[b0_quantized]], %[[in_scale]], %[[in_out_zp]], %[[w0_scale]], %[[w_b_zp]], %[[b0_scale]], %[[w_b_zp]], %[[mid_scale]], %[[in_out_zp]]) {config = "", config_proto = "", executor_type = "", f = @quantized_conv2d_with_bias_and_relu6_fn_1} -// CHECK-DAG: %[[conv1_dequantized:.*]] = "tf.PartitionedCall"(%[[conv0_quantized]], %[[w1_quantized]], %[[b1_quantized]], %[[mid_scale]], %[[in_out_zp]], %[[w1_scale]], %[[w_b_zp]], %[[b1_scale]], %[[w_b_zp]], %[[out_scale]], %[[in_out_zp]]) {config = "", config_proto = "", executor_type = "", f = @quantized_conv2d_with_bias_and_relu6_float_output_fn_0} +// CHECK-LABEL: func @conv_with_dump +// CHECK-DAG: %[[w0_float:.*]] = "tf.Const"() <{value = dense<{{\[\[\[\[}}-0.282878935, -0.211567819 +// CHECK-DAG: %[[b0_float:.*]] = "tf.Const"() <{value = dense<[-0.0192535277, -5.998660e-03]> : tensor<2xf32>}> : () -> tensor<2xf32> +// CHECK-DAG: %[[w1_float:.*]] = "tf.Const"() <{value = dense<{{\[\[\[\[}}0.208403707, 0.478067577 +// CHECK-DAG: %[[b1_float:.*]] = "tf.Const"() <{value = dense<[-0.0291469581, 0.0106381178]> : tensor<2xf32>}> : () -> tensor<2xf32> +// CHECK-DAG: %[[w0_quantized:.*]] = "tf.Const"() <{value = dense<{{\[\[\[\[}}-59, -44 +// CHECK-DAG: %[[b0_quantized:.*]] = "tf.Const"() <{value = dense<[-1040, -324]> : tensor<2xi32>}> : () -> tensor<2xi32> +// CHECK-DAG: %[[w1_quantized:.*]] = "tf.Const"() <{value = dense<{{\[\[\[\[}}44, 100 +// CHECK-DAG: %[[b1_quantized:.*]] = "tf.Const"() <{value = dense<[-4312, 1574]> : tensor<2xi32>}> : () -> tensor<2xi32> +// CHECK-DAG: %[[in_scale:.*]] = "tf.Const"() <{value = dense<0.00387597573> : tensor}> : () -> tensor +// CHECK-DAG: %[[in_out_zp:.*]] = "tf.Const"() <{value = dense<-128> : tensor}> : () -> tensor +// CHECK-DAG: %[[w0_scale:.*]] = "tf.Const"() <{value = dense<0.00477493973> : tensor}> : () -> tensor +// CHECK-DAG: %[[w_b_zp:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor +// CHECK-DAG: %[[b0_scale:.*]] = "tf.Const"() <{value = dense<1.85075514E-5> : tensor}> : () -> tensor +// CHECK-DAG: %[[mid_scale:.*]] = "tf.Const"() <{value = dense<0.00141507247> : tensor}> : () -> tensor +// CHECK-DAG: %[[w1_scale:.*]] = "tf.Const"() <{value = dense<0.00477652298> : tensor}> : () -> tensor +// CHECK-DAG: %[[b1_scale:.*]] = "tf.Const"() <{value = dense<6.75912588E-6> : tensor}> : () -> tensor +// CHECK-DAG: %[[out_scale:.*]] = "tf.Const"() <{value = dense<7.24974147E-4> : tensor}> : () -> tensor +// CHECK-DAG: %[[quantized:.*]] = "tf.PartitionedCall"(%arg0, %[[in_scale]], %[[in_out_zp]]) <{config = "", config_proto = "", executor_type = "", f = @quantize_i8}> +// CHECK-DAG: %[[conv0_dequantized:.*]] = "tf.PartitionedCall"(%[[quantized]], %[[w0_quantized]], %[[b0_quantized]], %[[in_scale]], %[[in_out_zp]], %[[w0_scale]], %[[w_b_zp]], %[[b0_scale]], %[[w_b_zp]], %[[mid_scale]], %[[in_out_zp]]) <{config = "", config_proto = "", executor_type = "", f = @quantized_conv2d_with_bias_and_relu6_float_output_fn_1}> +// CHECK-DAG: %[[conv0_quantized:.*]] = "tf.PartitionedCall"(%[[quantized]], %[[w0_quantized]], %[[b0_quantized]], %[[in_scale]], %[[in_out_zp]], %[[w0_scale]], %[[w_b_zp]], %[[b0_scale]], %[[w_b_zp]], %[[mid_scale]], %[[in_out_zp]]) <{config = "", config_proto = "", executor_type = "", f = @quantized_conv2d_with_bias_and_relu6_fn_1}> +// CHECK-DAG: %[[conv1_dequantized:.*]] = "tf.PartitionedCall"(%[[conv0_quantized]], %[[w1_quantized]], %[[b1_quantized]], %[[mid_scale]], %[[in_out_zp]], %[[w1_scale]], %[[w_b_zp]], %[[b1_scale]], %[[w_b_zp]], %[[out_scale]], %[[in_out_zp]]) <{config = "", config_proto = "", executor_type = "", f = @quantized_conv2d_with_bias_and_relu6_float_output_fn_0}> // CHECK-DAG: %[[identity:.*]] = "tf.Identity"(%[[conv1_dequantized]]) -// CHECK-DAG: %[[conv0_float:.*]] = "tf.PartitionedCall"(%arg0, %[[w0_float]], %[[b0_float]]) {config = "", config_proto = "", device = "", executor_type = "", f = @composite_conv2d_with_bias_and_relu6_fn_2_00} -// CHECK-DAG: %[[conv1_float:.*]] = "tf.PartitionedCall"(%[[conv0_dequantized]], %[[w1_float]], %[[b1_float]]) {config = "", config_proto = "", device = "", executor_type = "", f = @composite_conv2d_with_bias_and_relu6_fn_1_00} -// CHECK-DAG: "tf.DumpTensor"(%[[conv0_dequantized]]) {device = "", enabled = true, file_name = "quantized_tensor_data.pb", func_name = "conv_with_dump", log_dir_path = "/tmp/dumps/composite_conv2d_with_bias_and_relu6_fn_2", node_name = "Conv2D"} -// CHECK-DAG: "tf.DumpTensor"(%[[conv0_float]]) {device = "", enabled = true, file_name = "unquantized_tensor_data.pb", func_name = "conv_with_dump", log_dir_path = "/tmp/dumps/composite_conv2d_with_bias_and_relu6_fn_2", node_name = "Conv2D"} -// CHECK-DAG: "tf.DumpTensor"(%[[conv1_dequantized]]) {device = "", enabled = true, file_name = "quantized_tensor_data.pb", func_name = "conv_with_dump", log_dir_path = "/tmp/dumps/composite_conv2d_with_bias_and_relu6_fn_1", node_name = "Conv2D_1"} -// CHECK-DAG: "tf.DumpTensor"(%[[conv1_float]]) {device = "", enabled = true, file_name = "unquantized_tensor_data.pb", func_name = "conv_with_dump", log_dir_path = "/tmp/dumps/composite_conv2d_with_bias_and_relu6_fn_1", node_name = "Conv2D_1"} +// CHECK-DAG: %[[conv0_float:.*]] = "tf.PartitionedCall"(%arg0, %[[w0_float]], %[[b0_float]]) <{config = "", config_proto = "", executor_type = "", f = @composite_conv2d_with_bias_and_relu6_fn_2_00}> {device = ""} +// CHECK-DAG: %[[conv1_float:.*]] = "tf.PartitionedCall"(%[[conv0_dequantized]], %[[w1_float]], %[[b1_float]]) <{config = "", config_proto = "", executor_type = "", f = @composite_conv2d_with_bias_and_relu6_fn_1_00}> {device = ""} +// CHECK-DAG: "tf.DumpTensor"(%[[conv0_dequantized]]) <{enabled = true, file_name = "quantized_tensor_data.pb", func_name = "conv_with_dump", log_dir_path = "/tmp/dumps/composite_conv2d_with_bias_and_relu6_fn_2", node_name = "Conv2D"}> {device = ""} +// CHECK-DAG: "tf.DumpTensor"(%[[conv0_float]]) <{enabled = true, file_name = "unquantized_tensor_data.pb", func_name = "conv_with_dump", log_dir_path = "/tmp/dumps/composite_conv2d_with_bias_and_relu6_fn_2", node_name = "Conv2D"}> {device = ""} +// CHECK-DAG: "tf.DumpTensor"(%[[conv1_dequantized]]) <{enabled = true, file_name = "quantized_tensor_data.pb", func_name = "conv_with_dump", log_dir_path = "/tmp/dumps/composite_conv2d_with_bias_and_relu6_fn_1", node_name = "Conv2D_1"}> {device = ""} +// CHECK-DAG: "tf.DumpTensor"(%[[conv1_float]]) <{enabled = true, file_name = "unquantized_tensor_data.pb", func_name = "conv_with_dump", log_dir_path = "/tmp/dumps/composite_conv2d_with_bias_and_relu6_fn_1", node_name = "Conv2D_1"}> {device = ""} // CHECK-DAG: return %[[identity]] + +// PerChannel-LABEL: func @conv_with_dump +// PerChannel-DAG: %[[PerChannel_w0_float:.*]] = "tf.Const"() <{value = dense<{{\[\[\[\[}}-0.282878935, -0.211567819 +// PerChannel-DAG: %[[b0_float:.*]] = "tf.Const"() <{value = dense<[-0.0192535277, -5.998660e-03]> : tensor<2xf32>}> : () -> tensor<2xf32> +// PerChannel-DAG: %[[w1_float:.*]] = "tf.Const"() <{value = dense<{{\[\[\[\[}}0.208403707, 0.478067577 +// PerChannel-DAG: %[[b1_float:.*]] = "tf.Const"() <{value = dense<[-0.0291469581, 0.0106381178]> : tensor<2xf32>}> : () -> tensor<2xf32> +// PerChannel-DAG: %[[w0_quantized:.*]] = "tf.Const"() <{value = dense<{{\[\[\[\[}}-59, -77 +// PerChannel-DAG: %[[b0_quantized:.*]] = "tf.Const"() <{value = dense<[-1040, -561]> : tensor<2xi32>}> : () -> tensor<2xi32> +// PerChannel-DAG: %[[w1_quantized:.*]] = "tf.Const"() <{value = dense<{{\[\[\[\[}}45, 100 +// PerChannel-DAG: %[[b1_quantized:.*]] = "tf.Const"() <{value = dense<[-4411, 1574]> : tensor<2xi32>}> : () -> tensor<2xi32> +// PerChannel-DAG: %[[in_scale:.*]] = "tf.Const"() <{value = dense<0.00387597573> : tensor}> : () -> tensor +// PerChannel-DAG: %[[in_out_zp:.*]] = "tf.Const"() <{value = dense<-128> : tensor}> : () -> tensor +// PerChannel-DAG: %[[w0_scale:.*]] = "tf.Const"() <{value = dense<[0.00477493973, 0.00275693159]> : tensor<2xf32>}> : () -> tensor<2xf32> +// PerChannel-DAG: %[[w_b_zp:.*]] = "tf.Const"() <{value = dense<0> : tensor<2xi32>}> : () -> tensor<2xi32> +// PerChannel-DAG: %[[b0_scale:.*]] = "tf.Const"() <{value = dense<[1.85075514E-5, 1.06858006E-5]> : tensor<2xf32>}> : () -> tensor<2xf32> +// PerChannel-DAG: %[[mid_scale:.*]] = "tf.Const"() <{value = dense<0.00141507247> : tensor}> : () -> tensor +// PerChannel-DAG: %[[w1_scale:.*]] = "tf.Const"() <{value = dense<[0.00467005931, 0.00477652298]> : tensor<2xf32>}> : () -> tensor<2xf32> +// PerChannel-DAG: %[[b1_scale:.*]] = "tf.Const"() <{value = dense<[6.60847217E-6, 6.75912588E-6]> : tensor<2xf32>}> : () -> tensor<2xf32> +// PerChannel-DAG: %[[out_scale:.*]] = "tf.Const"() <{value = dense<7.24974147E-4> : tensor}> : () -> tensor +} + +// ----- + +module { + func.func @conv_with_per_channel_and_tensor_weight(%arg0: tensor<1x3x4x3xf32>) -> tensor<1x3x4x2xf32> { + %cst = "tf.Const"() {device = "", value = dense<[7.11401462, 7.05456924]> : tensor<2xf32>} : () -> tensor<2xf32> + %cst_0 = "tf.Const"() {device = "", value = dense<[[[[-0.630731344, 0.277245104], [0.54962182, 0.927732646], [0.180364341, 1.90948534]], [[-0.764542698, -0.287541777], [-0.211145893, -1.59367061], [-0.708605706, 1.79999375]], [[-0.954062759, 0.197947085], [-0.614013135, -0.966769516], [0.612640202, -1.45540595]]], [[[-0.418223292, 0.234433219], [5.057390e-01, 1.86747122], [0.899269938, 0.145780042]], [[0.335351914, 1.02572429], [0.084816426, 1.79729116], [-0.664676845, 0.310017586]], [[-0.795477629, -7.709830e-01], [0.581315517, 0.740075528], [0.921566545, 1.85318887]]]]> : tensor<2x3x3x2xf32>} : () -> tensor<2x3x3x2xf32> + %0 = "quantfork.stats"(%arg0) {layerStats = dense<[4.6128589E-5, 0.999998927]> : tensor<2xf32>} : (tensor<1x3x4x3xf32>) -> tensor<1x3x4x3xf32> + %1 = "tf.PartitionedCall"(%0, %cst_0, %cst) {_tfl_quant_trait = "fully_quantizable", config = "", config_proto = "", device = "", executor_type = "", f = @composite_conv2d_with_bias_and_relu6_fn_1} : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>, tensor<2xf32>) -> tensor<1x3x4x2xf32> + %2 = "quantfork.stats"(%1) {layerStats = dense<[3.50919247, 6.000000e+00]> : tensor<2xf32>} : (tensor<1x3x4x2xf32>) -> tensor<1x3x4x2xf32> + %3 = "tf.Identity"(%2) {device = ""} : (tensor<1x3x4x2xf32>) -> tensor<1x3x4x2xf32> + %4 = "quantfork.stats"(%3) {layerStats = dense<[3.50919247, 6.000000e+00]> : tensor<2xf32>} : (tensor<1x3x4x2xf32>) -> tensor<1x3x4x2xf32> + func.return %4 : tensor<1x3x4x2xf32> + } + func.func private @composite_conv2d_with_bias_and_relu6_fn_1(%arg0: tensor<1x3x4x3xf32>, %arg1: tensor<2x3x3x2xf32>, %arg2: tensor<2xf32>) -> tensor<1x3x4x2xf32> attributes {tf._original_func_name = "composite_conv2d_with_bias_and_relu6_fn_1", tf.tf_quant.composite_function} { + %0 = "tf.Conv2D"(%arg0, %arg1) {attr_map = "0:strides,1:use_cudnn_on_gpu,2:padding,3:explicit_paddings,4:dilations", data_format = "NHWC", device = "", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true} : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x3x4x2xf32> + %1 = "tf.BiasAdd"(%0, %arg2) {data_format = "NHWC", device = ""} : (tensor<1x3x4x2xf32>, tensor<2xf32>) -> tensor<1x3x4x2xf32> + %2 = "tf.Relu6"(%1) {device = ""} : (tensor<1x3x4x2xf32>) -> tensor<1x3x4x2xf32> + func.return %2 : tensor<1x3x4x2xf32> + } + +// CHECK-LABEL: func @conv_with_per_channel_and_tensor_weight +// CHECK-DAG: %[[b0_quantized:.*]] = "tf.Const"() <{value = dense<[120654, 119646]> : tensor<2xi32>}> : () -> tensor<2xi32> +// CHECK-DAG: %[[w0_quantized:.*]] = "tf.Const"() <{value = dense<{{\[\[\[\[}}-42, 18 +// CHECK-DAG: %[[in_scale:.*]] = "tf.Const"() <{value = dense<0.0039215642> : tensor}> : () -> tensor +// CHECK-DAG: %[[in_out_zp:.*]] = "tf.Const"() <{value = dense<-128> : tensor}> : () -> tensor +// CHECK-DAG: %[[w0_scale:.*]] = "tf.Const"() <{value = dense<0.0150353173> : tensor}> : () -> tensor +// CHECK-DAG: %[[w_b_zp:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor +// CHECK-DAG: %[[b0_scale:.*]] = "tf.Const"() <{value = dense<5.89619667E-5> : tensor}> : () -> tensor +// CHECK-DAG: %[[mid_scale:.*]] = "tf.Const"() <{value = dense<0.0235294122> : tensor}> : () -> tensor + +// PerChannel-LABEL: func @conv_with_per_channel_and_tensor_weight +// PerChannel-DAG: %[[b0_quantized:.*]] = "tf.Const"() <{value = dense<[241481, 119646]> : tensor<2xi32>}> : () -> tensor<2xi32> +// PerChannel-DAG: %[[w0_quantized:.*]] = "tf.Const"() <{value = dense<{{\[\[\[\[}}-84, 18 +// PerChannel-DAG: %[[in_scale:.*]] = "tf.Const"() <{value = dense<0.0039215642> : tensor}> : () -> tensor +// PerChannel-DAG: %[[in_out_zp:.*]] = "tf.Const"() <{value = dense<-128> : tensor}> : () -> tensor +// PerChannel-DAG: %[[w0_scale:.*]] = "tf.Const"() <{value = dense<[0.0075123054, 0.0150353173]> : tensor<2xf32>}> : () -> tensor<2xf32> +// PerChannel-DAG: %[[w_b_zp:.*]] = "tf.Const"() <{value = dense<0> : tensor<2xi32>}> : () -> tensor<2xi32> +// PerChannel-DAG: %[[b0_scale:.*]] = "tf.Const"() <{value = dense<[2.94599886E-5, 5.89619667E-5]> : tensor<2xf32>}> : () -> tensor<2xf32> +// PerChannel-DAG: %[[mid_scale:.*]] = "tf.Const"() <{value = dense<0.0235294122> : tensor}> : () -> tensor } diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/tests/quantize_drq.mlir b/tensorflow/compiler/mlir/quantization/tensorflow/tests/quantize_drq.mlir index c500b3c72e8c86..e3bda3f5d09af9 100644 --- a/tensorflow/compiler/mlir/quantization/tensorflow/tests/quantize_drq.mlir +++ b/tensorflow/compiler/mlir/quantization/tensorflow/tests/quantize_drq.mlir @@ -15,7 +15,7 @@ module { // CHECK: %[[cst:.*]] = "arith.constant"() <{value = dense<0.000000e+00> : tensor<2x1024xf32>}> : () -> tensor<2x1024xf32> // CHECK: %[[q_cst:.*]] = "quantfork.qcast"(%[[cst]]) : (tensor<2x1024xf32>) -> tensor<2x1024x!quant.uniform:f32, 3.9370078740157481E-9>> -// CHECK: %[[out:.*]] = "tf.PartitionedCall"(%arg0, %[[q_cst]]) {_tfl_quant_trait = "fully_quantizable", config = "", config_proto = "", executor_type = "", f = @composite_matmul_fn} : (tensor<1x2x2x3xf32>, tensor<2x1024x!quant.uniform:f32, 3.9370078740157481E-9>>) -> tensor<*xf32> +// CHECK: %[[out:.*]] = "tf.PartitionedCall"(%arg0, %[[q_cst]]) <{config = "", config_proto = "", executor_type = "", f = @composite_matmul_fn}> {_tfl_quant_trait = "fully_quantizable"} : (tensor<1x2x2x3xf32>, tensor<2x1024x!quant.uniform:f32, 3.9370078740157481E-9>>) -> tensor<*xf32> // CHECK: "func.return"(%[[out]]) : (tensor<*xf32>) -> () } diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/tests/quantize_weights.mlir b/tensorflow/compiler/mlir/quantization/tensorflow/tests/quantize_weights.mlir index c41d43d0f662fb..7f7a5090439e28 100644 --- a/tensorflow/compiler/mlir/quantization/tensorflow/tests/quantize_weights.mlir +++ b/tensorflow/compiler/mlir/quantization/tensorflow/tests/quantize_weights.mlir @@ -8,7 +8,7 @@ module { } // CHECK-LABEL: func @not_quantize_const -// CHECK-DAG: %[[W:.*]] = "tf.Const"() {value = dense<2.000000e+00> : tensor<2x1024xf32> +// CHECK-DAG: %[[W:.*]] = "tf.Const"() <{value = dense<2.000000e+00> : tensor<2x1024xf32> // CHECK: return %[[W]] : tensor<2x1024xf32> } @@ -22,15 +22,15 @@ module { } // CHECK-LABEL: func @matmul -// CHECK-DAG: %[[W:.*]] = "tf.Const"() {value = dense<127> : tensor<2x1024xi8> +// CHECK-DAG: %[[W:.*]] = "tf.Const"() <{value = dense<127> : tensor<2x1024xi8> // CHECK: %[[PRESERVE_W:.*]] = "tf.Identity"(%[[W]]) : (tensor<2x1024xi8>) -> tensor<2x1024xi8> -// CHECK: %[[DEQUANTIZED:.*]] = "tf.PartitionedCall"(%[[PRESERVE_W]]) {config = "", config_proto = "", executor_type = "", f = @composite_dequantize_uniform} : (tensor<2x1024xi8>) -> tensor<2x1024xf32> -// CHECK: %[[MATMUL:.*]] = "tf.MatMul"(%arg0, %[[DEQUANTIZED]]) {attr_map = "0:transpose_a,1:transpose_a", device = "", transpose_a = false, transpose_b = false} : (tensor<1x2x2x2xf32>, tensor<2x1024xf32>) -> tensor<*xf32> +// CHECK: %[[DEQUANTIZED:.*]] = "tf.PartitionedCall"(%[[PRESERVE_W]]) <{config = "", config_proto = "", executor_type = "", f = @composite_dequantize_uniform}> : (tensor<2x1024xi8>) -> tensor<2x1024xf32> +// CHECK: %[[MATMUL:.*]] = "tf.MatMul"(%arg0, %[[DEQUANTIZED]]) <{transpose_a = false, transpose_b = false}> {attr_map = "0:transpose_a,1:transpose_a", device = ""} : (tensor<1x2x2x2xf32>, tensor<2x1024xf32>) -> tensor<*xf32> // CHECK: return %[[MATMUL]] : tensor<*xf32> // CHECK-LABEL: func.func private @composite_dequantize_uniform(%arg0: tensor<*xi8>) -> tensor<*xf32> -// CHECK-DAG: %[[SCALE:.*]] = "tf.Const"() {value = dense<0.0157480314> : tensor -// CHECK: %[[CASTED_W:.*]] = "tf.Cast"(%arg0) {Truncate = false} : (tensor<*xi8>) -> tensor<*xf32> +// CHECK-DAG: %[[SCALE:.*]] = "tf.Const"() <{value = dense<0.0157480314> : tensor +// CHECK: %[[CASTED_W:.*]] = "tf.Cast"(%arg0) <{Truncate = false}> : (tensor<*xi8>) -> tensor<*xf32> // CHECK: %[[DEQUANTIZED:.*]] = "tf.Mul"(%[[CASTED_W]], %[[SCALE]]) : (tensor<*xf32>, tensor) -> tensor<*xf32> // CHECK: return %[[DEQUANTIZED]] : tensor<*xf32> } @@ -48,7 +48,7 @@ module { // CHECK-LABEL: func @not_quantize_matmul_without_const // CHECK: %[[ORIGINAL_IDENTITY_1:.*]] = "tf.Identity"(%arg0) {device = ""} : (tensor<1x2x2x2xf32>) -> tensor<1x2x2x2xf32> // CHECK: %[[ORIGINAL_IDENTITY_2:.*]] = "tf.Identity"(%arg1) {device = ""} : (tensor<2x1024xf32>) -> tensor<2x1024xf32> -// CHECK: %[[MATMUL:.*]] = "tf.MatMul"(%[[ORIGINAL_IDENTITY_1]], %[[ORIGINAL_IDENTITY_2]]) {attr_map = "0:transpose_a,1:transpose_a", device = "", transpose_a = false, transpose_b = false} : (tensor<1x2x2x2xf32>, tensor<2x1024xf32>) -> tensor<*xf32> +// CHECK: %[[MATMUL:.*]] = "tf.MatMul"(%[[ORIGINAL_IDENTITY_1]], %[[ORIGINAL_IDENTITY_2]]) <{transpose_a = false, transpose_b = false}> {attr_map = "0:transpose_a,1:transpose_a", device = ""} : (tensor<1x2x2x2xf32>, tensor<2x1024xf32>) -> tensor<*xf32> // CHECK: return %[[MATMUL]] : tensor<*xf32> } @@ -63,14 +63,14 @@ module { } // CHECK-LABEL: func @quantize_xladotv2_bf16 -// CHECK-DAG: %[[W:.*]] = "tf.Const"() {value = dense<127> : tensor<2x1024xi8> +// CHECK-DAG: %[[W:.*]] = "tf.Const"() <{value = dense<127> : tensor<2x1024xi8> // CHECK: %[[IDENTITY:.*]] = "tf.Identity"(%[[W]]) : (tensor<2x1024xi8>) -> tensor<2x1024xi8> -// CHECK: %[[DEQUANTIZED:.*]] = "tf.PartitionedCall"(%[[IDENTITY]]) {config = "", config_proto = "", executor_type = "", f = @composite_dequantize_uniform} : (tensor<2x1024xi8>) -> tensor<2x1024xbf16> -// CHECK: %[[MATMUL:.*]] = "tf.XlaDotV2"(%arg0, %[[DEQUANTIZED]]) {device = "", dimension_numbers = "\12\01\00\0A\01\03", precision_config = ""} : (tensor<1x2x2x2xbf16>, tensor<2x1024xbf16>) -> tensor<1x2x2x1024xbf16> +// CHECK: %[[DEQUANTIZED:.*]] = "tf.PartitionedCall"(%[[IDENTITY]]) <{config = "", config_proto = "", executor_type = "", f = @composite_dequantize_uniform}> : (tensor<2x1024xi8>) -> tensor<2x1024xbf16> +// CHECK: %[[MATMUL:.*]] = "tf.XlaDotV2"(%arg0, %[[DEQUANTIZED]]) <{dimension_numbers = "\12\01\00\0A\01\03", precision_config = ""}> {device = ""} : (tensor<1x2x2x2xbf16>, tensor<2x1024xbf16>) -> tensor<1x2x2x1024xbf16> // CHECK: return %[[MATMUL]] : tensor<1x2x2x1024xbf16> // CHECK-LABEL: func.func private @composite_dequantize_uniform(%arg0: tensor<*xi8>) -> tensor<*xbf16> -// CHECK-DAG: %[[SCALE:.*]] = "tf.Const"() {value = dense<1.574710e-02> : tensor +// CHECK-DAG: %[[SCALE:.*]] = "tf.Const"() <{value = dense<1.574710e-02> : tensor } // ----- @@ -87,17 +87,17 @@ module { } // CHECK-LABEL: func @matmul_with_identity_and_reshape -// CHECK-DAG: %[[W:.*]] = "tf.Const"() {value = dense<127> : tensor<1024x2xi8> -// CHECK-DAG: %[[SHAPE:.*]] = "tf.Const"() {value = dense<[2, 1024]> : tensor<2xi32> +// CHECK-DAG: %[[W:.*]] = "tf.Const"() <{value = dense<127> : tensor<1024x2xi8> +// CHECK-DAG: %[[SHAPE:.*]] = "tf.Const"() <{value = dense<[2, 1024]> : tensor<2xi32> // CHECK: %[[PRESERVE_W:.*]] = "tf.Identity"(%[[W]]) : (tensor<1024x2xi8>) -> tensor<1024x2xi8> -// CHECK: %[[DEQUANTIZED:.*]] = "tf.PartitionedCall"(%[[PRESERVE_W]]) {config = "", config_proto = "", executor_type = "", f = @composite_dequantize_uniform} : (tensor<1024x2xi8>) -> tensor<1024x2xf32> +// CHECK: %[[DEQUANTIZED:.*]] = "tf.PartitionedCall"(%[[PRESERVE_W]]) <{config = "", config_proto = "", executor_type = "", f = @composite_dequantize_uniform}> : (tensor<1024x2xi8>) -> tensor<1024x2xf32> // CHECK: %[[ORIGINAL_IDENTITY:.*]] = "tf.Identity"(%[[DEQUANTIZED]]) {device = ""} : (tensor<1024x2xf32>) -> tensor<1024x2xf32> // CHECK: %[[RESHAPED_W:.*]] = "tf.Reshape"(%[[ORIGINAL_IDENTITY]], %[[SHAPE]]) : (tensor<1024x2xf32>, tensor<2xi32>) -> tensor<2x1024xf32> -// CHECK: %[[MATMUL:.*]] = "tf.MatMul"(%arg0, %[[RESHAPED_W]]) {attr_map = "0:transpose_a,1:transpose_a", device = "", transpose_a = false, transpose_b = false} : (tensor<1x2x2x2xf32>, tensor<2x1024xf32>) -> tensor<*xf32> +// CHECK: %[[MATMUL:.*]] = "tf.MatMul"(%arg0, %[[RESHAPED_W]]) <{transpose_a = false, transpose_b = false}> {attr_map = "0:transpose_a,1:transpose_a", device = ""} : (tensor<1x2x2x2xf32>, tensor<2x1024xf32>) -> tensor<*xf32> // CHECK: return %[[MATMUL]] : tensor<*xf32> // CHECK-LABEL: func.func private @composite_dequantize_uniform(%arg0: tensor<*xi8>) -> tensor<*xf32> -// CHECK-DAG: %[[SCALE:.*]] = "tf.Const"() {value = dense<0.0157480314> : tensor +// CHECK-DAG: %[[SCALE:.*]] = "tf.Const"() <{value = dense<0.0157480314> : tensor } // ----- @@ -113,16 +113,16 @@ module { } // CHECK-LABEL: func @conv2d -// CHECK-DAG: %[[W:.*]] = "tf.Const"() {value = dense<127> : tensor<2x3x3x512xi8> -// CHECK-DAG: %[[BIAS:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<2xf32> +// CHECK-DAG: %[[W:.*]] = "tf.Const"() <{value = dense<127> : tensor<2x3x3x512xi8> +// CHECK-DAG: %[[BIAS:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<2xf32> // CHECK: %[[PRESERVE_W:.*]] = "tf.Identity"(%[[W]]) : (tensor<2x3x3x512xi8>) -> tensor<2x3x3x512xi8> -// CHECK: %[[DEQUANTIZED:.*]] = "tf.PartitionedCall"(%[[PRESERVE_W]]) {config = "", config_proto = "", executor_type = "", f = @composite_dequantize_uniform} : (tensor<2x3x3x512xi8>) -> tensor<2x3x3x512xf32> -// CHECK: %[[CONV2D:.*]] = "tf.Conv2D"(%arg0, %[[DEQUANTIZED:.*]]) {attr_map = "0:strides,1:use_cudnn_on_gpu,2:padding,3:explicit_paddings,4:dilations", data_format = "NHWC", device = "", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 2, 1], use_cudnn_on_gpu = true} : (tensor<1x3x4x3xf32>, tensor<2x3x3x512xf32>) -> tensor<*xf32> -// CHECK: %[[BIASADD:.*]] = "tf.BiasAdd"(%[[CONV2D]], %[[BIAS]]) {data_format = "NHWC", device = ""} : (tensor<*xf32>, tensor<2xf32>) -> tensor<*xf32> +// CHECK: %[[DEQUANTIZED:.*]] = "tf.PartitionedCall"(%[[PRESERVE_W]]) <{config = "", config_proto = "", executor_type = "", f = @composite_dequantize_uniform}> : (tensor<2x3x3x512xi8>) -> tensor<2x3x3x512xf32> +// CHECK: %[[CONV2D:.*]] = "tf.Conv2D"(%arg0, %[[DEQUANTIZED:.*]]) <{data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 2, 1], use_cudnn_on_gpu = true}> {attr_map = "0:strides,1:use_cudnn_on_gpu,2:padding,3:explicit_paddings,4:dilations", device = ""} : (tensor<1x3x4x3xf32>, tensor<2x3x3x512xf32>) -> tensor<*xf32> +// CHECK: %[[BIASADD:.*]] = "tf.BiasAdd"(%[[CONV2D]], %[[BIAS]]) <{data_format = "NHWC"}> {device = ""} : (tensor<*xf32>, tensor<2xf32>) -> tensor<*xf32> // CHECK: return %[[BIASADD]] : tensor<*xf32> // CHECK-LABEL: func.func private @composite_dequantize_uniform(%arg0: tensor<*xi8>) -> tensor<*xf32> -// CHECK-DAG: %[[SCALE:.*]] = "tf.Const"() {value = dense<0.0236220472> : tensor +// CHECK-DAG: %[[SCALE:.*]] = "tf.Const"() <{value = dense<0.0236220472> : tensor } // ----- @@ -138,14 +138,14 @@ module { } // CHECK-LABEL: func @depthwise_conv -// CHECK-DAG: %[[W:.*]] = "tf.Const"() {value = dense<127> : tensor<2x3x3x512xi8> +// CHECK-DAG: %[[W:.*]] = "tf.Const"() <{value = dense<127> : tensor<2x3x3x512xi8> // CHECK: %[[PRESERVE_W:.*]] = "tf.Identity"(%[[W]]) : (tensor<2x3x3x512xi8>) -> tensor<2x3x3x512xi8> -// CHECK: %[[DEQUANTIZED:.*]] = "tf.PartitionedCall"(%[[PRESERVE_W]]) {config = "", config_proto = "", executor_type = "", f = @composite_dequantize_uniform} : (tensor<2x3x3x512xi8>) -> tensor<2x3x3x512xf32> -// CHECK: %[[DEPTHWISE_CONV2D:.*]] = "tf.DepthwiseConv2dNative"(%arg0, %[[DEQUANTIZED]]) {attr_map = "0:strides,1:padding,2:explicit_paddings,3:dilations", data_format = "NHWC", device = "", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 2, 1]} : (tensor<1x3x4x512xf32>, tensor<2x3x3x512xf32>) -> tensor<*xf32> +// CHECK: %[[DEQUANTIZED:.*]] = "tf.PartitionedCall"(%[[PRESERVE_W]]) <{config = "", config_proto = "", executor_type = "", f = @composite_dequantize_uniform}> : (tensor<2x3x3x512xi8>) -> tensor<2x3x3x512xf32> +// CHECK: %[[DEPTHWISE_CONV2D:.*]] = "tf.DepthwiseConv2dNative"(%arg0, %[[DEQUANTIZED]]) <{data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 2, 1]}> {attr_map = "0:strides,1:padding,2:explicit_paddings,3:dilations", device = ""} : (tensor<1x3x4x512xf32>, tensor<2x3x3x512xf32>) -> tensor<*xf32> // CHECK: return %[[DEPTHWISE_CONV2D]] : tensor<*xf32> // CHECK-LABEL: func.func private @composite_dequantize_uniform(%arg0: tensor<*xi8>) -> tensor<*xf32> -// CHECK-DAG: %[[SCALE:.*]] = "tf.Const"() {value = dense<0.00787401571> : tensor +// CHECK-DAG: %[[SCALE:.*]] = "tf.Const"() <{value = dense<0.00787401571> : tensor } // ----- @@ -160,16 +160,16 @@ module { } // CHECK-LABEL: func @quantize_sharded_weights_with_xladot -// CHECK-DAG: %[[W:.*]] = "tf.Const"() {value = dense<127> : tensor<512x512xi8> +// CHECK-DAG: %[[W:.*]] = "tf.Const"() <{value = dense<127> : tensor<512x512xi8> // CHECK: %[[PRESERVE_W:.*]] = "tf.Identity"(%[[W]]) : (tensor<512x512xi8>) -> tensor<512x512xi8> -// CHECK: %[[DEQUANTIZED:.*]] = "tf.PartitionedCall"(%[[PRESERVE_W]]) {config = "", config_proto = "", executor_type = "", f = @composite_dequantize_uniform} : (tensor<512x512xi8>) -> tensor<512x512xf32> -// CHECK: %[[SHARDED_W:.*]] = "tf.XlaSharding"(%[[DEQUANTIZED]]) {_XlaSharding = "\08\03\1A\03\01\04\02\22\08\00\04\01\05\02\06\03\070\01", device = "", sharding = "\08\03\1A\03\01\04\02\22\08\00\04\01\05\02\06\03\070\01", unspecified_dims = []} : (tensor<512x512xf32>) -> tensor<512x512xf32> -// CHECK: %[[XLADOT:.*]] = "tf.XlaDotV2"(%arg0, %[[SHARDED_W]]) {device = "", dimension_numbers = "\12\01\00\0A\01\03", precision_config = ""} : (tensor, tensor<512x512xf32>) -> tensor -// CHECK: %[[ORIGINAL_CAST:.*]] = "tf.Cast"(%[[XLADOT]]) {Truncate = false} : (tensor) -> tensor +// CHECK: %[[DEQUANTIZED:.*]] = "tf.PartitionedCall"(%[[PRESERVE_W]]) <{config = "", config_proto = "", executor_type = "", f = @composite_dequantize_uniform}> : (tensor<512x512xi8>) -> tensor<512x512xf32> +// CHECK: %[[SHARDED_W:.*]] = "tf.XlaSharding"(%[[DEQUANTIZED]]) <{_XlaSharding = "\08\03\1A\03\01\04\02\22\08\00\04\01\05\02\06\03\070\01", sharding = "\08\03\1A\03\01\04\02\22\08\00\04\01\05\02\06\03\070\01"}> {device = "", unspecified_dims = []} : (tensor<512x512xf32>) -> tensor<512x512xf32> +// CHECK: %[[XLADOT:.*]] = "tf.XlaDotV2"(%arg0, %[[SHARDED_W]]) <{dimension_numbers = "\12\01\00\0A\01\03", precision_config = ""}> {device = ""} : (tensor, tensor<512x512xf32>) -> tensor +// CHECK: %[[ORIGINAL_CAST:.*]] = "tf.Cast"(%[[XLADOT]]) <{Truncate = false}> : (tensor) -> tensor // CHECK: return %[[ORIGINAL_CAST]] : tensor // CHECK-LABEL: func.func private @composite_dequantize_uniform(%arg0: tensor<*xi8>) -> tensor<*xf32> -// CHECK-DAG: %[[SCALE:.*]] = "tf.Const"() {value = dense<0.0787401571> : tensor +// CHECK-DAG: %[[SCALE:.*]] = "tf.Const"() <{value = dense<0.0787401571> : tensor } // ----- @@ -184,16 +184,16 @@ module { } // CHECK-LABEL: func @quantize_sharded_weights_with_xladot_with_identity -// CHECK-DAG: %[[W:.*]] = "tf.Const"() {value = dense<127> : tensor<512x512xi8> +// CHECK-DAG: %[[W:.*]] = "tf.Const"() <{value = dense<127> : tensor<512x512xi8> // CHECK: %[[PRESERVE_W:.*]] = "tf.Identity"(%[[W]]) : (tensor<512x512xi8>) -> tensor<512x512xi8> -// CHECK: %[[DEQUANTIZED:.*]] = "tf.PartitionedCall"(%[[PRESERVE_W]]) {config = "", config_proto = "", executor_type = "", f = @composite_dequantize_uniform} : (tensor<512x512xi8>) -> tensor<512x512xf32> +// CHECK: %[[DEQUANTIZED:.*]] = "tf.PartitionedCall"(%[[PRESERVE_W]]) <{config = "", config_proto = "", executor_type = "", f = @composite_dequantize_uniform}> : (tensor<512x512xi8>) -> tensor<512x512xf32> // CHECK: %[[IDENTITY_W:.*]] = "tf.Identity"(%[[DEQUANTIZED]]) {device = ""} : (tensor<512x512xf32>) -> tensor<512x512xf32> -// CHECK: %[[SHARDED_W:.*]] = "tf.XlaSharding"(%[[IDENTITY_W]]) {_XlaSharding = "\08\03\1A\03\01\04\02\22\08\00\04\01\05\02\06\03\070\01", device = "", sharding = "\08\03\1A\03\01\04\02\22\08\00\04\01\05\02\06\03\070\01", unspecified_dims = []} : (tensor<512x512xf32>) -> tensor<512x512xf32> -// CHECK: %[[XLADOT:.*]] = "tf.XlaDotV2"(%arg0, %[[SHARDED_W]]) {device = "", dimension_numbers = "\12\01\00\0A\01\03", precision_config = ""} : (tensor, tensor<512x512xf32>) -> tensor +// CHECK: %[[SHARDED_W:.*]] = "tf.XlaSharding"(%[[IDENTITY_W]]) <{_XlaSharding = "\08\03\1A\03\01\04\02\22\08\00\04\01\05\02\06\03\070\01", sharding = "\08\03\1A\03\01\04\02\22\08\00\04\01\05\02\06\03\070\01"}> {device = "", unspecified_dims = []} : (tensor<512x512xf32>) -> tensor<512x512xf32> +// CHECK: %[[XLADOT:.*]] = "tf.XlaDotV2"(%arg0, %[[SHARDED_W]]) <{dimension_numbers = "\12\01\00\0A\01\03", precision_config = ""}> {device = ""} : (tensor, tensor<512x512xf32>) -> tensor // CHECK: return %[[XLADOT]] : tensor // CHECK-LABEL: func.func private @composite_dequantize_uniform(%arg0: tensor<*xi8>) -> tensor<*xf32> -// CHECK-DAG: %[[SCALE:.*]] = "tf.Const"() {value = dense<0.0787401571> : tensor +// CHECK-DAG: %[[SCALE:.*]] = "tf.Const"() <{value = dense<0.0787401571> : tensor } // ----- @@ -208,16 +208,16 @@ module { } // CHECK-LABEL: func @quantize_xlagather -// CHECK-DAG: %[[W:.*]] = "tf.Const"() {value = dense<127> : tensor<200x100x300xi8>} : () -> tensor<200x100x300xi8> -// CHECK-DAG: %[[IDX:.*]] = "tf.Const"() {value = dense<[1, 1, 300]> : tensor<3xi64> +// CHECK-DAG: %[[W:.*]] = "tf.Const"() <{value = dense<127> : tensor<200x100x300xi8>}> : () -> tensor<200x100x300xi8> +// CHECK-DAG: %[[IDX:.*]] = "tf.Const"() <{value = dense<[1, 1, 300]> : tensor<3xi64> // CHECK: %[[PRESERVE_W:.*]] = "tf.Identity"(%[[W]]) : (tensor<200x100x300xi8>) -> tensor<200x100x300xi8> -// CHECK: %[[DEQUANTIZED:.*]] = "tf.PartitionedCall"(%[[PRESERVE_W]]) {config = "", config_proto = "", executor_type = "", f = @composite_dequantize_uniform} : (tensor<200x100x300xi8>) -> tensor<200x100x300xf32> -// CHECK: %[[GATHER:.*]] = "tf.XlaGather"(%[[DEQUANTIZED]], %arg0, %[[IDX]]) {dimension_numbers = "\0A\02\00\01\12\01\00\1A\02\00\01 \01", indices_are_sorted = true} : (tensor<200x100x300xf32>, tensor<10x2xi32>, tensor<3xi64>) -> tensor<1x300x10xf32> +// CHECK: %[[DEQUANTIZED:.*]] = "tf.PartitionedCall"(%[[PRESERVE_W]]) <{config = "", config_proto = "", executor_type = "", f = @composite_dequantize_uniform}> : (tensor<200x100x300xi8>) -> tensor<200x100x300xf32> +// CHECK: %[[GATHER:.*]] = "tf.XlaGather"(%[[DEQUANTIZED]], %arg0, %[[IDX]]) <{dimension_numbers = "\0A\02\00\01\12\01\00\1A\02\00\01 \01", indices_are_sorted = true}> : (tensor<200x100x300xf32>, tensor<10x2xi32>, tensor<3xi64>) -> tensor<1x300x10xf32> // CHECK: %[[IDENTITY:.*]] = "tf.Identity"(%[[GATHER]]) {device = ""} : (tensor<1x300x10xf32>) -> tensor<1x300x10xf32> // CHECK: return %[[IDENTITY]] : tensor<1x300x10xf32> // CHECK-LABEL: func.func private @composite_dequantize_uniform(%arg0: tensor<*xi8>) -> tensor<*xf32> -// CHECK-DAG: %[[SCALE:.*]] = "tf.Const"() {value = dense<0.0787401571> : tensor} : () -> tensor +// CHECK-DAG: %[[SCALE:.*]] = "tf.Const"() <{value = dense<0.0787401571> : tensor}> : () -> tensor } // ----- @@ -236,17 +236,17 @@ module { } // CHECK-LABEL: func @partitioned_call -// CHECK-DAG: %[[W:.*]] = "tf.Const"() {value = dense<127> : tensor<2x1024xi8> +// CHECK-DAG: %[[W:.*]] = "tf.Const"() <{value = dense<127> : tensor<2x1024xi8> // CHECK: %[[PRESERVE_W:.*]] = "tf.Identity"(%[[W]]) : (tensor<2x1024xi8>) -> tensor<2x1024xi8> -// CHECK: %[[DEQUANTIZED:.*]] = "tf.PartitionedCall"(%[[PRESERVE_W]]) {config = "", config_proto = "", executor_type = "", f = @composite_dequantize_uniform} : (tensor<2x1024xi8>) -> tensor<2x1024xf32> -// CHECK: %[[OUTPUT:.*]] = "tf.PartitionedCall"(%arg0, %[[DEQUANTIZED]]) {_tfl_quant_trait = "fully_quantizable", config = "", config_proto = "", executor_type = "", f = @composite_matmul_fn} : (tensor<1x2x2x3xf32>, tensor<2x1024xf32>) -> tensor<*xf32> +// CHECK: %[[DEQUANTIZED:.*]] = "tf.PartitionedCall"(%[[PRESERVE_W]]) <{config = "", config_proto = "", executor_type = "", f = @composite_dequantize_uniform}> : (tensor<2x1024xi8>) -> tensor<2x1024xf32> +// CHECK: %[[OUTPUT:.*]] = "tf.PartitionedCall"(%arg0, %[[DEQUANTIZED]]) <{config = "", config_proto = "", executor_type = "", f = @composite_matmul_fn}> {_tfl_quant_trait = "fully_quantizable"} : (tensor<1x2x2x3xf32>, tensor<2x1024xf32>) -> tensor<*xf32> // CHECK: return %[[OUTPUT]] : tensor<*xf32> // CHECK-LABEL: func.func private @composite_dequantize_uniform(%arg0: tensor<*xi8>) -> tensor<*xf32> -// CHECK-DAG: %[[SCALE:.*]] = "tf.Const"() {value = dense<0.0314960629> : tensor +// CHECK-DAG: %[[SCALE:.*]] = "tf.Const"() <{value = dense<0.0314960629> : tensor // CHECK-LABEL: func private @composite_matmul_fn -// CHECK: %[[MATMUL:.*]] = "tf.MatMul"(%arg0, %arg1) {attr_map = "0:transpose_a,1:transpose_a", device = "", transpose_a = false, transpose_b = false} : (tensor<1x2x2x3xf32>, tensor<2x1024xf32>) -> tensor<*xf32> +// CHECK: %[[MATMUL:.*]] = "tf.MatMul"(%arg0, %arg1) <{transpose_a = false, transpose_b = false}> {attr_map = "0:transpose_a,1:transpose_a", device = ""} : (tensor<1x2x2x3xf32>, tensor<2x1024xf32>) -> tensor<*xf32> // CHECK: return %[[MATMUL]] : tensor<*xf32> } @@ -272,21 +272,21 @@ module { } // CHECK-LABEL: func @recursive_partitioned_call(%arg0: tensor<1x2x2x3xf32>) -> tensor<*xf32> -// CHECK-DAG: %[[W:.*]] = "tf.Const"() {value = dense<127> : tensor<2x1024xi8> +// CHECK-DAG: %[[W:.*]] = "tf.Const"() <{value = dense<127> : tensor<2x1024xi8> // CHECK: %[[PRESERVE_W:.*]] = "tf.Identity"(%[[W]]) : (tensor<2x1024xi8>) -> tensor<2x1024xi8> -// CHECK: %[[DEQUANTIZED:.*]] = "tf.PartitionedCall"(%[[PRESERVE_W]]) {config = "", config_proto = "", executor_type = "", f = @composite_dequantize_uniform} : (tensor<2x1024xi8>) -> tensor<2x1024xf32> -// CHECK: %[[OUTPUT:.*]] = "tf.PartitionedCall"(%arg0, %[[DEQUANTIZED]]) {config = "", config_proto = "", executor_type = "", f = @outer_fn} : (tensor<1x2x2x3xf32>, tensor<2x1024xf32>) -> tensor<*xf32> +// CHECK: %[[DEQUANTIZED:.*]] = "tf.PartitionedCall"(%[[PRESERVE_W]]) <{config = "", config_proto = "", executor_type = "", f = @composite_dequantize_uniform}> : (tensor<2x1024xi8>) -> tensor<2x1024xf32> +// CHECK: %[[OUTPUT:.*]] = "tf.PartitionedCall"(%arg0, %[[DEQUANTIZED]]) <{config = "", config_proto = "", executor_type = "", f = @outer_fn}> : (tensor<1x2x2x3xf32>, tensor<2x1024xf32>) -> tensor<*xf32> // CHECK: return %[[OUTPUT]] : tensor<*xf32> // CHECK-LABEL: func private @composite_dequantize_uniform(%arg0: tensor<*xi8>) -> tensor<*xf32> -// CHECK-DAG: %[[SCALE:.*]] = "tf.Const"() {value = dense<0.0314960629> : tensor +// CHECK-DAG: %[[SCALE:.*]] = "tf.Const"() <{value = dense<0.0314960629> : tensor // CHECK-LABEL: func private @outer_fn -// CHECK: %[[OUTER_OUTPUT:.*]] = "tf.PartitionedCall"(%arg0, %arg1) {config = "", config_proto = "", executor_type = "", f = @inner_fn} : (tensor<1x2x2x3xf32>, tensor<2x1024xf32>) -> tensor<*xf32> +// CHECK: %[[OUTER_OUTPUT:.*]] = "tf.PartitionedCall"(%arg0, %arg1) <{config = "", config_proto = "", executor_type = "", f = @inner_fn}> : (tensor<1x2x2x3xf32>, tensor<2x1024xf32>) -> tensor<*xf32> // CHECK: return %[[OUTER_OUTPUT]] : tensor<*xf32> // CHECK-LABEL: func private @inner_fn -// CHECK: %[[MATMUL:.*]] = "tf.MatMul"(%arg0, %arg1) {attr_map = "0:transpose_a,1:transpose_a", device = "", transpose_a = false, transpose_b = false} : (tensor<1x2x2x3xf32>, tensor<2x1024xf32>) -> tensor<*xf32> +// CHECK: %[[MATMUL:.*]] = "tf.MatMul"(%arg0, %arg1) <{transpose_a = false, transpose_b = false}> {attr_map = "0:transpose_a,1:transpose_a", device = ""} : (tensor<1x2x2x3xf32>, tensor<2x1024xf32>) -> tensor<*xf32> // CHECK: return %[[MATMUL]] : tensor<*xf32> // ----- @@ -302,17 +302,17 @@ module { } // CHECK-LABEL: func @matmul_multiuses -// CHECK-DAG: %[[W:.*]] = "tf.Const"() {value = dense<127> : tensor<2x1024xi8> +// CHECK-DAG: %[[W:.*]] = "tf.Const"() <{value = dense<127> : tensor<2x1024xi8> // CHECK: %[[PRESERVE_W:.*]] = "tf.Identity"(%[[W]]) : (tensor<2x1024xi8>) -> tensor<2x1024xi8> -// CHECK: %[[DEQUANTIZED:.*]] = "tf.PartitionedCall"(%[[PRESERVE_W]]) {config = "", config_proto = "", executor_type = "", f = @composite_dequantize_uniform} : (tensor<2x1024xi8>) -> tensor<2x1024xf32> -// CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%arg0, %[[DEQUANTIZED]]) {attr_map = "0:transpose_a,1:transpose_a", device = "", transpose_a = false, transpose_b = false} : (tensor<1x2x2x2xf32>, tensor<2x1024xf32>) -> tensor<*xf32> -// CHECK: %[[MATMUL_2:.*]] = "tf.MatMul"(%arg1, %[[DEQUANTIZED]]) {attr_map = "0:transpose_a,1:transpose_a", device = "", transpose_a = false, transpose_b = false} : (tensor<1x2x2x2xf32>, tensor<2x1024xf32>) -> tensor<*xf32> +// CHECK: %[[DEQUANTIZED:.*]] = "tf.PartitionedCall"(%[[PRESERVE_W]]) <{config = "", config_proto = "", executor_type = "", f = @composite_dequantize_uniform}> : (tensor<2x1024xi8>) -> tensor<2x1024xf32> +// CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%arg0, %[[DEQUANTIZED]]) <{transpose_a = false, transpose_b = false}> {attr_map = "0:transpose_a,1:transpose_a", device = ""} : (tensor<1x2x2x2xf32>, tensor<2x1024xf32>) -> tensor<*xf32> +// CHECK: %[[MATMUL_2:.*]] = "tf.MatMul"(%arg1, %[[DEQUANTIZED]]) <{transpose_a = false, transpose_b = false}> {attr_map = "0:transpose_a,1:transpose_a", device = ""} : (tensor<1x2x2x2xf32>, tensor<2x1024xf32>) -> tensor<*xf32> // CHECK: %[[ORIGINAL_IDENTITY:.*]] = "tf.Identity"(%[[DEQUANTIZED]]) {device = ""} : (tensor<2x1024xf32>) -> tensor<2x1024xf32> -// CHECK: %[[MATMUL_3:.*]] = "tf.MatMul"(%arg0, %[[ORIGINAL_IDENTITY]]) {attr_map = "0:transpose_a,1:transpose_a", device = "", transpose_a = false, transpose_b = false} : (tensor<1x2x2x2xf32>, tensor<2x1024xf32>) -> tensor<*xf32> +// CHECK: %[[MATMUL_3:.*]] = "tf.MatMul"(%arg0, %[[ORIGINAL_IDENTITY]]) <{transpose_a = false, transpose_b = false}> {attr_map = "0:transpose_a,1:transpose_a", device = ""} : (tensor<1x2x2x2xf32>, tensor<2x1024xf32>) -> tensor<*xf32> // CHECK: return %[[MATMUL_1]], %[[MATMUL_2]], %[[MATMUL_3]] : tensor<*xf32>, tensor<*xf32>, tensor<*xf32> // CHECK-LABEL: func.func private @composite_dequantize_uniform(%arg0: tensor<*xi8>) -> tensor<*xf32> -// CHECK-DAG: %[[SCALE:.*]] = "tf.Const"() {value = dense<0.0157480314> : tensor +// CHECK-DAG: %[[SCALE:.*]] = "tf.Const"() <{value = dense<0.0157480314> : tensor } // ----- @@ -327,15 +327,15 @@ module { } // CHECK-LABEL: func @matmul_multiuses -// CHECK-DAG: %[[W:.*]] = "tf.Const"() {value = dense<127> : tensor<2x1024xi8> +// CHECK-DAG: %[[W:.*]] = "tf.Const"() <{value = dense<127> : tensor<2x1024xi8> // CHECK: %[[PRESERVE_W:.*]] = "tf.Identity"(%[[W]]) : (tensor<2x1024xi8>) -> tensor<2x1024xi8> -// CHECK: %[[DEQUANTIZED:.*]] = "tf.PartitionedCall"(%[[PRESERVE_W]]) {config = "", config_proto = "", executor_type = "", f = @composite_dequantize_uniform} : (tensor<2x1024xi8>) -> tensor<2x1024xf32> -// CHECK: %[[MATMUL:.*]] = "tf.MatMul"(%arg0, %[[DEQUANTIZED]]) {attr_map = "0:transpose_a,1:transpose_a", device = "", transpose_a = false, transpose_b = false} : (tensor<1x2x2x2xf32>, tensor<2x1024xf32>) -> tensor<*xf32> +// CHECK: %[[DEQUANTIZED:.*]] = "tf.PartitionedCall"(%[[PRESERVE_W]]) <{config = "", config_proto = "", executor_type = "", f = @composite_dequantize_uniform}> : (tensor<2x1024xi8>) -> tensor<2x1024xf32> +// CHECK: %[[MATMUL:.*]] = "tf.MatMul"(%arg0, %[[DEQUANTIZED]]) <{transpose_a = false, transpose_b = false}> {attr_map = "0:transpose_a,1:transpose_a", device = ""} : (tensor<1x2x2x2xf32>, tensor<2x1024xf32>) -> tensor<*xf32> // CHECK: %[[ADD:.*]] = "tf.AddV2"(%arg1, %[[DEQUANTIZED]]) {device = ""} : (tensor<2x1024xf32>, tensor<2x1024xf32>) -> tensor<2x1024xf32> // CHECK: return %[[MATMUL]], %[[ADD]] : tensor<*xf32>, tensor<2x1024xf32> // CHECK-LABEL: func.func private @composite_dequantize_uniform(%arg0: tensor<*xi8>) -> tensor<*xf32> -// CHECK-DAG: %[[SCALE:.*]] = "tf.Const"() {value = dense<0.0157480314> : tensor +// CHECK-DAG: %[[SCALE:.*]] = "tf.Const"() <{value = dense<0.0157480314> : tensor } // ----- @@ -375,21 +375,21 @@ module { } // CHECK-LABEL: func @matmul_with_while -// CHECK-DAG: %[[W:.*]] = "tf.Const"() {value = dense<127> : tensor<1024x1024xi8> -// CHECK-DAG: %[[CNT:.*]] = "tf.Const"() {value = dense<1> : tensor} : () -> tensor +// CHECK-DAG: %[[W:.*]] = "tf.Const"() <{value = dense<127> : tensor<1024x1024xi8> +// CHECK-DAG: %[[CNT:.*]] = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor // CHECK: %[[PRESERVE_W:.*]] = "tf.Identity"(%[[W]]) : (tensor<1024x1024xi8>) -> tensor<1024x1024xi8> -// CHECK: %[[DEQUANTIZED:.*]] = "tf.PartitionedCall"(%[[PRESERVE_W]]) {config = "", config_proto = "", executor_type = "", f = @composite_dequantize_uniform} : (tensor<1024x1024xi8>) -> tensor<1024x1024xf32> -// CHECK: %[[WHILE:.*]] = "tf.While"(%[[CNT]], %[[CNT]], %[[CNT]], %arg0, %[[DEQUANTIZED]]) {T = [i32, i32, i32, f32, f32], _lower_using_switch_merge = true, _num_original_outputs = 5 : i64, _read_only_resource_inputs = [], body = @while_body, cond = @while_cond, device = "", is_stateless = true, output_shapes = [#tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<1x1024>, #tf_type.shape<1024x1024>], parallel_iterations = 10 : i64, shape_invariant} : (tensor, tensor, tensor, tensor<1x1024xf32>, tensor<1024x1024xf32>) -> (tensor, tensor, tensor, tensor<1x1024xf32>, tensor<1024x1024xf32>) +// CHECK: %[[DEQUANTIZED:.*]] = "tf.PartitionedCall"(%[[PRESERVE_W]]) <{config = "", config_proto = "", executor_type = "", f = @composite_dequantize_uniform}> : (tensor<1024x1024xi8>) -> tensor<1024x1024xf32> +// CHECK: %[[WHILE:.*]] = "tf.While"(%[[CNT]], %[[CNT]], %[[CNT]], %arg0, %[[DEQUANTIZED]]) <{body = @while_body, cond = @while_cond, is_stateless = true, parallel_iterations = 10 : i64, shape_invariant}> {T = [i32, i32, i32, f32, f32], _lower_using_switch_merge = true, _num_original_outputs = 5 : i64, _read_only_resource_inputs = [], device = "", output_shapes = [#tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<1x1024>, #tf_type.shape<1024x1024>]} : (tensor, tensor, tensor, tensor<1x1024xf32>, tensor<1024x1024xf32>) -> (tensor, tensor, tensor, tensor<1x1024xf32>, tensor<1024x1024xf32>) // CHECK: %[[IDENTITY:.*]] = "tf.Identity"(%[[WHILE:.*]]) {device = ""} : (tensor<1x1024xf32>) -> tensor<1x1024xf32> // CHECK: return %[[IDENTITY]] : tensor<1x1024xf32> // CHECK-LABEL: func.func private @composite_dequantize_uniform(%arg0: tensor<*xi8>) -> tensor<*xf32> -// CHECK-DAG: %[[SCALE:.*]] = "tf.Const"() {value = dense<0.00787401571> : tensor +// CHECK-DAG: %[[SCALE:.*]] = "tf.Const"() <{value = dense<0.00787401571> : tensor // CHECK-LABEL: func private @while_body(%arg0: tensor, %arg1: tensor, %arg2: tensor, %arg3: tensor<1x1024xf32>, %arg4: tensor<1024x1024xf32>) -> (tensor, tensor, tensor, tensor<1x1024xf32>, tensor<1024x1024xf32>) -// CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%arg3, %arg4) {device = "", transpose_a = false, transpose_b = false} : (tensor<1x1024xf32>, tensor<1024x1024xf32>) -> tensor<1x1024xf32> +// CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%arg3, %arg4) <{transpose_a = false, transpose_b = false}> {device = ""} : (tensor<1x1024xf32>, tensor<1024x1024xf32>) -> tensor<1x1024xf32> // CHECK: %[[IDENTITY:.*]] = "tf.Identity"(%arg4) {device = ""} : (tensor<1024x1024xf32>) -> tensor<1024x1024xf32> -// CHECK: %[[MATMUL_2:.*]] = "tf.MatMul"(%arg3, %[[IDENTITY]]) {device = "", transpose_a = false, transpose_b = false} : (tensor<1x1024xf32>, tensor<1024x1024xf32>) -> tensor<1x1024xf32> +// CHECK: %[[MATMUL_2:.*]] = "tf.MatMul"(%arg3, %[[IDENTITY]]) <{transpose_a = false, transpose_b = false}> {device = ""} : (tensor<1x1024xf32>, tensor<1024x1024xf32>) -> tensor<1x1024xf32> // CHECK: %[[ADD:.*]] = "tf.AddV2"(%[[MATMUL_1]], %[[MATMUL_2]]) {device = ""} : (tensor<1x1024xf32>, tensor<1x1024xf32>) -> tensor<1x1024xf32> // CHECK-LABEL: func private @while_cond(%arg0: tensor, %arg1: tensor, %arg2: tensor, %arg3: tensor<1x1024xf32>, %arg4: tensor<1024x1024xf32>) -> tensor @@ -401,13 +401,13 @@ module { func.func @matmul_with_while_bf16(%arg0: tensor<1x1024xbf16>) -> tensor<1x1024xbf16> { %cst = "tf.Const"() {value = dense<1> : tensor} : () -> tensor %cst_0 = "tf.Const"() {value = dense<1> : tensor} : () -> tensor - %cst_1 = "tf.Const"(){value = dense<1.0> : tensor<1024x1024xbf16>} : () -> tensor<1024x1024xbf16> + %cst_1 = "tf.Const"() {value = dense<1.0> : tensor<1024x1024xbf16>} : () -> tensor<1024x1024xbf16> %0:5 = "tf.While"(%cst_0, %cst, %cst_0, %arg0, %cst_1) {T = [i32, i32, i32, bf16, bf16],_lower_using_switch_merge = true, _num_original_outputs = 5 : i64, _read_only_resource_inputs = [], body = @while_body, cond = @while_cond, device = "", is_stateless = true, output_shapes = [#tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<1x1024>, #tf_type.shape<1024x1024>], parallel_iterations = 10 : i64, shape_invariant} : (tensor, tensor, tensor, tensor<1x1024xbf16>, tensor<1024x1024xbf16>) -> (tensor, tensor, tensor, tensor<1x1024xbf16>, tensor<1024x1024xbf16>) %1 = "tf.Identity"(%0#3) {device = ""} : (tensor<1x1024xbf16>) -> tensor<1x1024xbf16> func.return %1 : tensor<1x1024xbf16> } - func.func private @while_body(%arg0: tensor, %arg1: tensor, %arg2: tensor, %arg3: tensor<1x1024xbf16>, %arg4: tensor<1024x1024xbf16>) -> (tensor, tensor, tensor, tensor<1x1024xbf16>, tensor<1024x1024xbf16>) + func.func private @while_body(%arg0: tensor, %arg1: tensor, %arg2: tensor, %arg3: tensor<1x1024xbf16>, %arg4: tensor<1024x1024xbf16>) -> (tensor, tensor, tensor, tensor<1x1024xbf16>, tensor<1024x1024xbf16>) { %cst = "tf.Const"() {value = dense<1> : tensor} : () -> tensor %0 = "tf.AddV2"(%arg2, %cst) {device = ""} : (tensor, tensor) -> tensor @@ -432,20 +432,20 @@ module { } // CHECK-LABEL: func @matmul_with_while_bf16 -// CHECK-DAG: %[[W:.*]] = "tf.Const"() {value = dense<127> : tensor<1024x1024xi8> -// CHECK-DAG: %[[CNT:.*]] = "tf.Const"() {value = dense<1> : tensor} : () -> tensor +// CHECK-DAG: %[[W:.*]] = "tf.Const"() <{value = dense<127> : tensor<1024x1024xi8> +// CHECK-DAG: %[[CNT:.*]] = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor // CHECK: %[[IDENTITY:.*]] = "tf.Identity"(%[[W]]) : (tensor<1024x1024xi8>) -> tensor<1024x1024xi8> -// CHECK: %[[DEQUANTIZED:.*]] = "tf.PartitionedCall"(%[[IDENTITY]]) {config = "", config_proto = "", executor_type = "", f = @composite_dequantize_uniform} : (tensor<1024x1024xi8>) -> tensor<1024x1024xbf16> -// CHECK: %[[WHILE:.*]] = "tf.While"(%[[CNT]], %[[CNT]], %[[CNT]], %arg0, %[[DEQUANTIZED]]) {T = [i32, i32, i32, bf16, bf16], _lower_using_switch_merge = true, _num_original_outputs = 5 : i64, _read_only_resource_inputs = [], body = @while_body, cond = @while_cond, device = "", is_stateless = true, output_shapes = [#tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<1x1024>, #tf_type.shape<1024x1024>], parallel_iterations = 10 : i64, shape_invariant} : (tensor, tensor, tensor, tensor<1x1024xbf16>, tensor<1024x1024xbf16>) -> (tensor, tensor, tensor, tensor<1x1024xbf16>, tensor<1024x1024xbf16>) +// CHECK: %[[DEQUANTIZED:.*]] = "tf.PartitionedCall"(%[[IDENTITY]]) <{config = "", config_proto = "", executor_type = "", f = @composite_dequantize_uniform}> : (tensor<1024x1024xi8>) -> tensor<1024x1024xbf16> +// CHECK: %[[WHILE:.*]] = "tf.While"(%[[CNT]], %[[CNT]], %[[CNT]], %arg0, %[[DEQUANTIZED]]) <{body = @while_body, cond = @while_cond, is_stateless = true, parallel_iterations = 10 : i64, shape_invariant}> {T = [i32, i32, i32, bf16, bf16], _lower_using_switch_merge = true, _num_original_outputs = 5 : i64, _read_only_resource_inputs = [], device = "", output_shapes = [#tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<1x1024>, #tf_type.shape<1024x1024>]} : (tensor, tensor, tensor, tensor<1x1024xbf16>, tensor<1024x1024xbf16>) -> (tensor, tensor, tensor, tensor<1x1024xbf16>, tensor<1024x1024xbf16>) // CHECK: %[[ORIGIANL_IDENTITY:.*]] = "tf.Identity"(%[[WHILE:.*]]) {device = ""} : (tensor<1x1024xbf16>) -> tensor<1x1024xbf16> // CHECK-LABEL: func.func private @composite_dequantize_uniform(%arg0: tensor<*xi8>) -> tensor<*xbf16> -// CHECK-DAG: %[[SCALE:.*]] = "tf.Const"() {value = dense<7.873530e-03> : tensor +// CHECK-DAG: %[[SCALE:.*]] = "tf.Const"() <{value = dense<7.873530e-03> : tensor // CHECK-LABEL: func private @while_body(%arg0: tensor, %arg1: tensor, %arg2: tensor, %arg3: tensor<1x1024xbf16>, %arg4: tensor<1024x1024xbf16>) -> (tensor, tensor, tensor, tensor<1x1024xbf16>, tensor<1024x1024xbf16>) { -// CHECK: %[[MATMUL_1:.*]] = "tf.XlaDotV2"(%arg3, %arg4) {device = "", dimension_numbers = "\12\01\00\0A\01\03", precision_config = ""} : (tensor<1x1024xbf16>, tensor<1024x1024xbf16>) -> tensor<1x1024xbf16> +// CHECK: %[[MATMUL_1:.*]] = "tf.XlaDotV2"(%arg3, %arg4) <{dimension_numbers = "\12\01\00\0A\01\03", precision_config = ""}> {device = ""} : (tensor<1x1024xbf16>, tensor<1024x1024xbf16>) -> tensor<1x1024xbf16> // CHECK: %[[IDENTITY_2:.*]] = "tf.Identity"(%arg4) {device = ""} : (tensor<1024x1024xbf16>) -> tensor<1024x1024xbf16> -// CHECK: %[[MATMUL_2:.*]] = "tf.XlaDotV2"(%arg3, %[[IDENTITY_2]]) {device = "", dimension_numbers = "\12\01\00\0A\01\03", precision_config = ""} : (tensor<1x1024xbf16>, tensor<1024x1024xbf16>) -> tensor<1x1024xbf16> +// CHECK: %[[MATMUL_2:.*]] = "tf.XlaDotV2"(%arg3, %[[IDENTITY_2]]) <{dimension_numbers = "\12\01\00\0A\01\03", precision_config = ""}> {device = ""} : (tensor<1x1024xbf16>, tensor<1024x1024xbf16>) -> tensor<1x1024xbf16> // CHECK: %[[ADD:.*]] = "tf.AddV2"(%[[MATMUL_1]], %[[MATMUL_2]]) {device = ""} : (tensor<1x1024xbf16>, tensor<1x1024xbf16>) -> tensor<1x1024xbf16> // CHECK-LABEL: func private @while_cond(%arg0: tensor, %arg1: tensor, %arg2: tensor, %arg3: tensor<1x1024xbf16>, %arg4: tensor<1024x1024xbf16>) -> tensor { @@ -482,7 +482,7 @@ module { } // CHECK-LABEL: func @matmul_with_while_returning_mutated_value -// CHECK-DAG: %[[W:.*]] = "tf.Const"() {value = dense<1.000000e+00> : tensor<1024x1024xf32>} : () -> tensor<1024x1024xf32> +// CHECK-DAG: %[[W:.*]] = "tf.Const"() <{value = dense<1.000000e+00> : tensor<1024x1024xf32>}> : () -> tensor<1024x1024xf32> // ----- module { @@ -499,27 +499,27 @@ module { } // CHECK-LABEL: func @multiple_quantizable_ops_in_graph -// CHECK-DAG: %[[W_1:.*]] = "tf.Const"() {value = dense<127> : tensor<2x3x3x1024xi8>} : () -> tensor<2x3x3x1024xi8> -// CHECK-DAG: %[[W_2:.*]] = "tf.Const"() {value = dense<127> : tensor<3x3x1024x1xi8>} : () -> tensor<3x3x1024x1xi8> -// CHECK-DAG: %[[W_3:.*]] = "tf.Const"() {value = dense<127> : tensor<1024x3x4x3xi8>} : () -> tensor<1024x3x4x3xi8> -// CHECK-DAG: %[[AXIS:.*]] = "tf.Const"() {device = "", value = dense<0> : tensor} : () -> tensor +// CHECK-DAG: %[[W_1:.*]] = "tf.Const"() <{value = dense<127> : tensor<2x3x3x1024xi8>}> : () -> tensor<2x3x3x1024xi8> +// CHECK-DAG: %[[W_2:.*]] = "tf.Const"() <{value = dense<127> : tensor<3x3x1024x1xi8>}> : () -> tensor<3x3x1024x1xi8> +// CHECK-DAG: %[[W_3:.*]] = "tf.Const"() <{value = dense<127> : tensor<1024x3x4x3xi8>}> : () -> tensor<1024x3x4x3xi8> +// CHECK-DAG: %[[AXIS:.*]] = "tf.Const"() <{value = dense<0> : tensor}> {device = ""} : () -> tensor // CHECK: %[[IDENTITY_1:.*]] = "tf.Identity"(%[[W_1]]) : (tensor<2x3x3x1024xi8>) -> tensor<2x3x3x1024xi8> -// CHECK: %[[DEQUANTIZED_1:.*]] = "tf.PartitionedCall"(%[[IDENTITY_1]]) {config = "", config_proto = "", executor_type = "", f = @composite_dequantize_uniform__} : (tensor<2x3x3x1024xi8>) -> tensor<2x3x3x1024xf32> +// CHECK: %[[DEQUANTIZED_1:.*]] = "tf.PartitionedCall"(%[[IDENTITY_1]]) <{config = "", config_proto = "", executor_type = "", f = @composite_dequantize_uniform__}> : (tensor<2x3x3x1024xi8>) -> tensor<2x3x3x1024xf32> // CHECK: %[[IDENTITY_2:.*]] = "tf.Identity"(%[[W_2]]) : (tensor<3x3x1024x1xi8>) -> tensor<3x3x1024x1xi8> -// CHECK: %[[DEQUANTIZED_2:.*]] = "tf.PartitionedCall"(%[[IDENTITY_2]]) {config = "", config_proto = "", executor_type = "", f = @composite_dequantize_uniform_} : (tensor<3x3x1024x1xi8>) -> tensor<3x3x1024x1xf32> +// CHECK: %[[DEQUANTIZED_2:.*]] = "tf.PartitionedCall"(%[[IDENTITY_2]]) <{config = "", config_proto = "", executor_type = "", f = @composite_dequantize_uniform_}> : (tensor<3x3x1024x1xi8>) -> tensor<3x3x1024x1xf32> // CHECK: %[[IDENTITY_3:.*]] = "tf.Identity"(%[[W_3]]) : (tensor<1024x3x4x3xi8>) -> tensor<1024x3x4x3xi8> -// CHECK: %[[DEQUANTIZED_3:.*]] = "tf.PartitionedCall"(%[[IDENTITY_3]]) {config = "", config_proto = "", executor_type = "", f = @composite_dequantize_uniform} : (tensor<1024x3x4x3xi8>) -> tensor<1024x3x4x3xf32> -// CHECK: %[[GATHER:.*]] = "tf.GatherV2"(%[[DEQUANTIZED_3]], %arg0, %[[AXIS]]) {batch_dims = 0 : i64, device = ""} : (tensor<1024x3x4x3xf32>, tensor<1xi32>, tensor) -> tensor<1x3x4x3xf32> -// CHECK: %[[CONV_1:.*]] = "tf.Conv2D"(%[[GATHER]], %[[DEQUANTIZED_1]]) {data_format = "NHWC", device = "", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 2, 1], use_cudnn_on_gpu = true} : (tensor<1x3x4x3xf32>, tensor<2x3x3x1024xf32>) -> tensor<1x3x2x1024xf32> -// CHECK: %[[CONV_2:.*]] = "tf.Conv2D"(%[[CONV_1]], %[[DEQUANTIZED_2]]) {data_format = "NHWC", device = "", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 2, 1], use_cudnn_on_gpu = true} : (tensor<1x3x2x1024xf32>, tensor<3x3x1024x1xf32>) -> tensor<1x3x1x1xf32> +// CHECK: %[[DEQUANTIZED_3:.*]] = "tf.PartitionedCall"(%[[IDENTITY_3]]) <{config = "", config_proto = "", executor_type = "", f = @composite_dequantize_uniform}> : (tensor<1024x3x4x3xi8>) -> tensor<1024x3x4x3xf32> +// CHECK: %[[GATHER:.*]] = "tf.GatherV2"(%[[DEQUANTIZED_3]], %arg0, %[[AXIS]]) <{batch_dims = 0 : i64}> {device = ""} : (tensor<1024x3x4x3xf32>, tensor<1xi32>, tensor) -> tensor<1x3x4x3xf32> +// CHECK: %[[CONV_1:.*]] = "tf.Conv2D"(%[[GATHER]], %[[DEQUANTIZED_1]]) <{data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 2, 1], use_cudnn_on_gpu = true}> {device = ""} : (tensor<1x3x4x3xf32>, tensor<2x3x3x1024xf32>) -> tensor<1x3x2x1024xf32> +// CHECK: %[[CONV_2:.*]] = "tf.Conv2D"(%[[CONV_1]], %[[DEQUANTIZED_2]]) <{data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 2, 1], use_cudnn_on_gpu = true}> {device = ""} : (tensor<1x3x2x1024xf32>, tensor<3x3x1024x1xf32>) -> tensor<1x3x1x1xf32> // CHECK-LABEL: func private @composite_dequantize_uniform__ -// CHECK-DAG: %[[SCALE:.*]] = "tf.Const"() {value = dense<0.00866141729> : tensor} : () -> tensor +// CHECK-DAG: %[[SCALE:.*]] = "tf.Const"() <{value = dense<0.00866141729> : tensor}> : () -> tensor // CHECK-LABEL: func private @composite_dequantize_uniform_ -// CHECK-DAG: %[[SCALE:.*]] = "tf.Const"() {value = dense<0.00866141729> : tensor} : () -> tensor +// CHECK-DAG: %[[SCALE:.*]] = "tf.Const"() <{value = dense<0.00866141729> : tensor}> : () -> tensor // CHECK-LABEL: func private @composite_dequantize_uniform -// CHECK-DAG: %[[SCALE:.*]] = "tf.Const"() {value = dense<0.00866141729> : tensor} : () -> tensor +// CHECK-DAG: %[[SCALE:.*]] = "tf.Const"() <{value = dense<0.00866141729> : tensor}> : () -> tensor -} \ No newline at end of file +} diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/tests/quantize_xla.mlir b/tensorflow/compiler/mlir/quantization/tensorflow/tests/quantize_xla.mlir index 4356d084a56845..f24b6399774f08 100644 --- a/tensorflow/compiler/mlir/quantization/tensorflow/tests/quantize_xla.mlir +++ b/tensorflow/compiler/mlir/quantization/tensorflow/tests/quantize_xla.mlir @@ -23,7 +23,7 @@ func.func private @conv(%input: tensor<1x3x4x3xf32> {tf._user_specified_name = " // CHECK-DAG: [[weight:%.+]] = "arith.constant"() <{value = dense_resource<__elided__> : tensor<2x3x3x2xf32>}> : () -> tensor<2x3x3x2x!quant.uniform> // CHECK: [[q_input:%.+]] = "quantfork.qcast"(%arg0) : (tensor<1x3x4x3xf32>) -> tensor<1x3x4x3x!quant.uniform> // CHECK-NEXT: [[q_bias:%.+]] = "quantfork.qcast"([[bias]]) : (tensor<2xf32>) -> tensor<2x!quant.uniform> -// CHECK-NEXT: [[conv:%.+]] = "tf.PartitionedCall"([[q_input]], [[weight]], [[q_bias]]) {_tfl_quant_trait = "fully_quantizable", config = "", config_proto = "", executor_type = "", f = @[[composite_fn:composite_conv2d_with_bias_and_relu6_fn.*]]} : (tensor<1x3x4x3x!quant.uniform>, tensor<2x3x3x2x!quant.uniform>, tensor<2x!quant.uniform>) -> tensor<*x!quant.uniform> +// CHECK-NEXT: [[conv:%.+]] = "tf.PartitionedCall"([[q_input]], [[weight]], [[q_bias]]) <{config = "", config_proto = "", executor_type = "", f = @[[composite_fn:composite_conv2d_with_bias_and_relu6_fn.*]]}> {_tfl_quant_trait = "fully_quantizable"} : (tensor<1x3x4x3x!quant.uniform>, tensor<2x3x3x2x!quant.uniform>, tensor<2x!quant.uniform>) -> tensor<*x!quant.uniform> // CHECK-NEXT: [[res:%.+]] = "quantfork.dcast"([[conv]]) : (tensor<*x!quant.uniform>) -> tensor<*xf32> // CHECK-NEXT: "func.return"([[res]]) : (tensor<*xf32>) -> () @@ -127,10 +127,10 @@ func.func private @avgpool_after_conv(%input: tensor<1x3x4x3xf32> {tf._user_spec // CHECK-SAME: f = @composite_conv2d_with_bias_and_relu6_fn_1 // CHECK-SAME: (tensor<1x3x4x3x!quant.uniform>, tensor<2x3x3x2x!quant.uniform>, tensor<2x!quant.uniform>) -> tensor<*x!quant.uniform> // CHECK: %[[scast:.*]] = "quantfork.scast"(%[[conv]] -// CHECK: %[[fcast:.*]] = "tf.Cast"(%[[scast]]) {Truncate = false} : (tensor<*xi8>) -> tensor<*xf32> +// CHECK: %[[fcast:.*]] = "tf.Cast"(%[[scast]]) <{Truncate = false}> : (tensor<*xi8>) -> tensor<*xf32> // CHECK: %[[avgpool_f32:.*]] = "tf.AvgPool"(%[[fcast]]) // CHECK-SAME: (tensor<*xf32>) -> tensor<*xf32> // CHECK: %[[round:.*]] = "tf.Round"(%[[avgpool_f32]]) -// CHECK: %[[icast:.*]] = "tf.Cast"(%[[round]]) {Truncate = false} : (tensor<*xf32>) -> tensor<*xi8> +// CHECK: %[[icast:.*]] = "tf.Cast"(%[[round]]) <{Truncate = false}> : (tensor<*xf32>) -> tensor<*xi8> // CHECK: %[[reshape:.*]] = "tf.Reshape"(%[[icast]] // CHECK: %[[sc2:.*]] = "quantfork.scast"(%[[reshape]]) diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/tests/remove_var_init_by_const.mlir b/tensorflow/compiler/mlir/quantization/tensorflow/tests/remove_var_init_by_const.mlir index d5e18209291871..da78ef781b3468 100644 --- a/tensorflow/compiler/mlir/quantization/tensorflow/tests/remove_var_init_by_const.mlir +++ b/tensorflow/compiler/mlir/quantization/tensorflow/tests/remove_var_init_by_const.mlir @@ -79,7 +79,7 @@ module attributes {tf_saved_model.semantics} { // CHECK-NOT: "tf.AssignVariableOp" // CHECK: %[[CST:.*]] = "tf.Const"() // CHECK-NEXT: %[[IDENTITY:.*]] = "tf.Identity"(%[[CST]]) - // CHECK-NEXT: %[[VAR:.*]] = "tf.VarHandleOp"() {{{.*shared_name = "var_1".*}}} + // CHECK-NEXT: %[[VAR:.*]] = "tf.VarHandleOp"() <{{{.*shared_name = "var_1".*}}}> // CHECK-NEXT: "tf.AssignVariableOp"(%[[VAR]], %[[IDENTITY]]) } diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/tests/replace_cast_hacks_with_tf_xla_ops.mlir b/tensorflow/compiler/mlir/quantization/tensorflow/tests/replace_cast_hacks_with_tf_xla_ops.mlir index c7fae4cf232b9f..04677b6e468828 100644 --- a/tensorflow/compiler/mlir/quantization/tensorflow/tests/replace_cast_hacks_with_tf_xla_ops.mlir +++ b/tensorflow/compiler/mlir/quantization/tensorflow/tests/replace_cast_hacks_with_tf_xla_ops.mlir @@ -57,17 +57,17 @@ module attributes {} { } // CHECK-LABEL: func @conv_with_bias_and_relu -// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() {value = dense<[1, 2]> : tensor<2xi32>} : () -> tensor<2xi32> -// CHECK-DAG: %[[CONST_1:.*]] = "tf.Const"() {value = dense<1> : tensor<2xi32>} : () -> tensor<2xi32> -// CHECK-DAG: %[[CONST_2:.*]] = "tf.Const"() {value = dense<1> : tensor} : () -> tensor -// CHECK-DAG: %[[CONST_3:.*]] = "tf.Const"() {value = dense<0> : tensor<2x2xi32>} : () -> tensor<2x2xi32> -// CHECK-DAG: %[[CONST_4:.*]] = "tf.Const"() {value = dense<{{.*}}> : tensor<4x2xi32>} : () -> tensor<4x2xi32> +// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() <{value = dense<[1, 2]> : tensor<2xi32>}> : () -> tensor<2xi32> +// CHECK-DAG: %[[CONST_1:.*]] = "tf.Const"() <{value = dense<1> : tensor<2xi32>}> : () -> tensor<2xi32> +// CHECK-DAG: %[[CONST_2:.*]] = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor +// CHECK-DAG: %[[CONST_3:.*]] = "tf.Const"() <{value = dense<0> : tensor<2x2xi32>}> : () -> tensor<2x2xi32> +// CHECK-DAG: %[[CONST_4:.*]] = "tf.Const"() <{value = dense<{{.*}}> : tensor<4x2xi32>}> : () -> tensor<4x2xi32> // CHECK-DAG-SAME{LITERAL}: value = dense<[[0, 0], [0, 1], [0, 1], [0, 0]]> -// CHECK-DAG: %[[CONST_5:.*]] = "tf.Const"() {value = dense<-128> : tensor} : () -> tensor -// CHECK-DAG: %[[CONST_6:.*]] = "tf.Const"() {value = dense<{{.*}}> : tensor<2x3x3x2xi8>} : () -> tensor<2x3x3x2xi8> -// CHECK-DAG: %[[CONST_7:.*]] = "tf.Const"() {value = dense<{{.*}}> : tensor<1x1x1x2xi32>} : () -> tensor<1x1x1x2xi32> +// CHECK-DAG: %[[CONST_5:.*]] = "tf.Const"() <{value = dense<-128> : tensor}> : () -> tensor +// CHECK-DAG: %[[CONST_6:.*]] = "tf.Const"() <{value = dense<{{.*}}> : tensor<2x3x3x2xi8>}> : () -> tensor<2x3x3x2xi8> +// CHECK-DAG: %[[CONST_7:.*]] = "tf.Const"() <{value = dense<{{.*}}> : tensor<1x1x1x2xi32>}> : () -> tensor<1x1x1x2xi32> // CHECK-DAG-SAME{LITERAL}: value = dense<[[[[-22016, -23680]]]]> -// CHECK-DAG: %[[CONST_8:.*]] = "tf.Const"() {value = dense<[162, 160]> : tensor<2xi32>} : () -> tensor<2xi32> +// CHECK-DAG: %[[CONST_8:.*]] = "tf.Const"() <{value = dense<[162, 160]> : tensor<2xi32>}> : () -> tensor<2xi32> // CHECK: %[[PADV2_0:.*]] = "tf.PadV2"({{.*}}, %[[CONST_4]], %[[CONST_5]]) : (tensor<1x3x4x3xi8>, tensor<4x2xi32>, tensor) -> tensor<1x4x5x3xi8> // CHECK: %[[XLACONVV2_0:.*]] = "tf.XlaConvV2"(%[[PADV2_0]], %[[CONST_6]], %[[CONST_0]], %[[CONST_3]], %[[CONST_1]], %[[CONST_1]], %[[CONST_2]]) // CHECK-SAME: (tensor<1x4x5x3xi8>, tensor<2x3x3x2xi8>, tensor<2xi32>, tensor<2x2xi32>, tensor<2xi32>, tensor<2xi32>, tensor) -> tensor<1x3x2x2xi32> @@ -144,16 +144,16 @@ module attributes {} { } // CHECK-LABEL: func @depthwise_conv_with_bias_and_relu6 -// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() {value = dense<{{.*}}> : tensor<4x2xi32>} : () -> tensor<4x2xi32> -// CHECK-DAG: %[[CONST_1:.*]] = "tf.Const"() {value = dense<-128> : tensor} : () -> tensor -// CHECK-DAG: %[[CONST_2:.*]] = "tf.Const"() {value = dense<{{.*}}> : tensor<2x3x1x3xi8>} : () -> tensor<2x3x1x3xi8> -// CHECK-DAG: %[[CONST_3:.*]] = "tf.Const"() {value = dense<2> : tensor<2xi32>} : () -> tensor<2xi32> -// CHECK-DAG: %[[CONST_4:.*]] = "tf.Const"() {value = dense<0> : tensor<2x2xi32>} : () -> tensor<2x2xi32> -// CHECK-DAG: %[[CONST_5:.*]] = "tf.Const"() {value = dense<1> : tensor<2xi32>} : () -> tensor<2xi32> -// CHECK-DAG: %[[CONST_6:.*]] = "tf.Const"() {value = dense<3> : tensor} : () -> tensor -// CHECK-DAG: %[[CONST_7:.*]] = "tf.Const"() {value = dense<{{.*}}> : tensor<1x1x1x3xi32>} : () -> tensor<1x1x1x3xi32> +// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() <{value = dense<{{.*}}> : tensor<4x2xi32>}> : () -> tensor<4x2xi32> +// CHECK-DAG: %[[CONST_1:.*]] = "tf.Const"() <{value = dense<-128> : tensor}> : () -> tensor +// CHECK-DAG: %[[CONST_2:.*]] = "tf.Const"() <{value = dense<{{.*}}> : tensor<2x3x1x3xi8>}> : () -> tensor<2x3x1x3xi8> +// CHECK-DAG: %[[CONST_3:.*]] = "tf.Const"() <{value = dense<2> : tensor<2xi32>}> : () -> tensor<2xi32> +// CHECK-DAG: %[[CONST_4:.*]] = "tf.Const"() <{value = dense<0> : tensor<2x2xi32>}> : () -> tensor<2x2xi32> +// CHECK-DAG: %[[CONST_5:.*]] = "tf.Const"() <{value = dense<1> : tensor<2xi32>}> : () -> tensor<2xi32> +// CHECK-DAG: %[[CONST_6:.*]] = "tf.Const"() <{value = dense<3> : tensor}> : () -> tensor +// CHECK-DAG: %[[CONST_7:.*]] = "tf.Const"() <{value = dense<{{.*}}> : tensor<1x1x1x3xi32>}> : () -> tensor<1x1x1x3xi32> // CHECK-DAG-SAME{LITERAL}: value = dense<[[[[55040, -15104, -21376]]]]> -// CHECK-DAG: %[[CONST_8:.*]] = "tf.Const"() {value = dense<[129, 166, 221]> : tensor<3xi32>} : () -> tensor<3xi32> +// CHECK-DAG: %[[CONST_8:.*]] = "tf.Const"() <{value = dense<[129, 166, 221]> : tensor<3xi32>}> : () -> tensor<3xi32> // CHECK: %[[PADV2_0:.*]] = "tf.PadV2"({{.*}}, %[[CONST_0]], %[[CONST_1]]) : (tensor<1x3x4x3xi8>, tensor<4x2xi32>, tensor) -> tensor<1x4x5x3xi8> // CHECK: %[[XLACONVV2_0:.*]] = "tf.XlaConvV2"(%[[PADV2_0]], %[[CONST_2]], %[[CONST_3]], %[[CONST_4]], %[[CONST_5]], %[[CONST_5]], %[[CONST_6]]) // CHECK-SAME: (tensor<1x4x5x3xi8>, tensor<2x3x1x3xi8>, tensor<2xi32>, tensor<2x2xi32>, tensor<2xi32>, tensor<2xi32>, tensor) -> tensor<1x2x2x3xi32> @@ -198,14 +198,14 @@ module attributes {} { } // CHECK-LABEL: func @dynamic_shaped_conv2d_with_bias_and_relu6_inlined -// CHECK-DAG: %[[filter:.*]] = "tf.Const"() {device = "", value = dense<2> : tensor<2x3x3x2xi8>} : () -> tensor<2x3x3x2xi8> +// CHECK-DAG: %[[filter:.*]] = "tf.Const"() <{value = dense<2> : tensor<2x3x3x2xi8>}> {device = ""} : () -> tensor<2x3x3x2xi8> // CHECK-DAG: %[[input_shape:.*]] = "tf.Shape"({{.*}}) : (tensor) -> tensor<4xi32> -// CHECK-DAG: %[[input_dim_1:.*]] = "tf.StridedSlice"(%[[input_shape]], {{.*}}, {{.*}}, {{.*}}) {begin_mask = 0 : i64, ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 1 : i64} : (tensor<4xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor -// CHECK-DAG: %[[input_dim_2:.*]] = "tf.StridedSlice"(%[[input_shape]], {{.*}}, {{.*}}, {{.*}}) {begin_mask = 0 : i64, ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 1 : i64} : (tensor<4xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor +// CHECK-DAG: %[[input_dim_1:.*]] = "tf.StridedSlice"(%[[input_shape]], {{.*}}, {{.*}}, {{.*}}) <{begin_mask = 0 : i64, ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 1 : i64}> : (tensor<4xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor +// CHECK-DAG: %[[input_dim_2:.*]] = "tf.StridedSlice"(%[[input_shape]], {{.*}}, {{.*}}, {{.*}}) <{begin_mask = 0 : i64, ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 1 : i64}> : (tensor<4xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor // CHECK-DAG: %[[padding_rank_1:.*]] = "tf.Concat"({{.*}}, {{.*}}, {{.*}}, {{.*}}, {{.*}}, {{.*}}, {{.*}}, {{.*}}, {{.*}}) : (tensor, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<8xi32> // CHECK-DAG: %[[padding_rank_2:.*]] = "tf.Reshape"(%[[padding_rank_1]], {{.*}}) : (tensor<8xi32>, tensor<2xi64>) -> tensor<4x2xi32> // CHECK-DAG: %[[input_padded:.*]] = "tf.PadV2"(%{{.*}}, %[[padding_rank_2]], {{.*}}) : (tensor, tensor<4x2xi32>, tensor) -> tensor -// CHECK: %[[conv_output:.*]] = "tf.XlaConvV2"(%[[input_padded]], %[[filter]], {{.*}}, {{.*}}, {{.*}}, {{.*}}, {{.*}}) {batch_group_count = 1 : i64, dimension_numbers = "{{.*}}", precision_config = ""} : (tensor, tensor<2x3x3x2xi8>, tensor<2xi32>, tensor<2x2xi32>, tensor<2xi32>, tensor<2xi32>, tensor) -> tensor +// CHECK: %[[conv_output:.*]] = "tf.XlaConvV2"(%[[input_padded]], %[[filter]], {{.*}}, {{.*}}, {{.*}}, {{.*}}, {{.*}}) <{batch_group_count = 1 : i64, dimension_numbers = "{{.*}}", precision_config = ""}> : (tensor, tensor<2x3x3x2xi8>, tensor<2xi32>, tensor<2x2xi32>, tensor<2xi32>, tensor<2xi32>, tensor) -> tensor // CHECK: %[[conv_output_sub:.*]] = "tf.Sub"(%[[conv_output]], {{.*}}) : (tensor, tensor<1x1x1x2xi32>) -> tensor // CHECK: %[[conv_output_add:.*]] = "tf.AddV2"(%[[conv_output_sub]], {{.*}}) {device = ""} : (tensor, tensor<2xi32>) -> tensor } @@ -264,7 +264,7 @@ module attributes {tf_saved_model.semantics} { } // CHECK-LABEL: func @conv_with_filter_larger_than_1MB -// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() {value = dense<-264192> : tensor<1x1x1x512xi32>} : () -> tensor<1x1x1x512xi32> +// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() <{value = dense<-264192> : tensor<1x1x1x512xi32>}> : () -> tensor<1x1x1x512xi32> // CHECK: %[[PADV2_0:.*]] = "tf.PadV2" // CHECK: %[[XLACONVV2_0:.*]] = "tf.XlaConvV2"(%[[PADV2_0]] // CHECK: %[[SUB_0:.*]] = "tf.Sub"(%[[XLACONVV2_0]], %[[CONST]]) @@ -297,8 +297,8 @@ module attributes {tf_saved_model.semantics} { return %12 : tensor<1x3xf32> } // CHECK-LABEL: func @matmul_with_relu -// CHECK-DAG: %[[WEIGHT:.*]] = "tf.Const"() {device = "", value = dense<1> : tensor<1024x3xi8>} : () -> tensor<1024x3xi8> -// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() {value = dense<-131072> : tensor<1x3xi32>} : () -> tensor<1x3xi32> +// CHECK-DAG: %[[WEIGHT:.*]] = "tf.Const"() <{value = dense<1> : tensor<1024x3xi8>}> {device = ""} : () -> tensor<1024x3xi8> +// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() <{value = dense<-131072> : tensor<1x3xi32>}> : () -> tensor<1x3xi32> // CHECK: %[[MATMUL:.*]] = "tf.XlaDotV2"({{.*}}, %[[WEIGHT]]) // CHECK-SAME: (tensor<1x1024xi8>, tensor<1024x3xi8>) -> tensor<1x3xi32> // CHECK: %[[SUB:.*]] = "tf.Sub"(%[[MATMUL]], %[[CONST]]) : (tensor<1x3xi32>, tensor<1x3xi32>) -> tensor<1x3xi32> @@ -479,11 +479,11 @@ module attributes {tf_saved_model.semantics} { } // CHECK-LABEL: func @conv3d_with_static_shape -// CHECK-DAG: %[[WEIGHT:.*]] = "tf.Const"() {device = "", value = dense<1> : tensor<2x3x3x3x2xi8>} : () -> tensor<2x3x3x3x2xi8> +// CHECK-DAG: %[[WEIGHT:.*]] = "tf.Const"() <{value = dense<1> : tensor<2x3x3x3x2xi8>}> {device = ""} : () -> tensor<2x3x3x3x2xi8> // CHECK-DAG: %[[CONST:.*]] = "tf.Const"() {{.*}} : () -> tensor<5x2xi32> // CHECK-DAG-SAME{LITERAL}: value = dense<[[0, 0], [0, 1], [0, 1], [1, 1], [0, 0]]> : tensor<5x2xi32> -// CHECK-DAG: %[[CONST_1:.*]] = "tf.Const"() {value = dense<-43> : tensor} : () -> tensor -// CHECK-DAG: %[[CONST_2:.*]] = "tf.Const"() {value = dense<-2322> : tensor<1x1x1x1x2xi32>} : () -> tensor<1x1x1x1x2xi32> +// CHECK-DAG: %[[CONST_1:.*]] = "tf.Const"() <{value = dense<-43> : tensor}> : () -> tensor +// CHECK-DAG: %[[CONST_2:.*]] = "tf.Const"() <{value = dense<-2322> : tensor<1x1x1x1x2xi32>}> : () -> tensor<1x1x1x1x2xi32> // CHECK: %[[PAD:.*]] = "tf.PadV2"({{.*}}, %[[CONST]], %[[CONST_1]]) // CHECK: %[[CONV:.*]] = "tf.XlaConvV2"(%[[PAD]], %[[WEIGHT]] @@ -524,9 +524,9 @@ module attributes {tf_saved_model.semantics} { } // CHECK-LABEL: func @conv3d_with_dynamic_shape -// CHECK-DAG: %[[WEIGHT:.*]] = "tf.Const"() {device = "", value = dense<1> : tensor<2x3x3x3x2xi8>} : () -> tensor<2x3x3x3x2xi8> -// CHECK-DAG: %[[CONST_1:.*]] = "tf.Const"() {value = dense<-43> : tensor} : () -> tensor -// CHECK-DAG: %[[CONST_2:.*]] = "tf.Const"() {value = dense<-2322> : tensor<1x1x1x1x2xi32>} : () -> tensor<1x1x1x1x2xi32> +// CHECK-DAG: %[[WEIGHT:.*]] = "tf.Const"() <{value = dense<1> : tensor<2x3x3x3x2xi8>}> {device = ""} : () -> tensor<2x3x3x3x2xi8> +// CHECK-DAG: %[[CONST_1:.*]] = "tf.Const"() <{value = dense<-43> : tensor}> : () -> tensor +// CHECK-DAG: %[[CONST_2:.*]] = "tf.Const"() <{value = dense<-2322> : tensor<1x1x1x1x2xi32>}> : () -> tensor<1x1x1x1x2xi32> // CHECK: %[[CONCAT:.*]] = "tf.Concat"({{.*}}) // CHECK: %[[RESHAPE:.*]] = "tf.Reshape"(%[[CONCAT]], {{.*}}) : (tensor<10xi32>, tensor<2xi64>) -> tensor<5x2xi32> @@ -565,7 +565,7 @@ module attributes {tf_saved_model.semantics} { } // CHECK-LABEL: func @batch_matmul -// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() {value = dense<-131072> : tensor<20x30x1x3xi32>} : () -> tensor<20x30x1x3xi32> +// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() <{value = dense<-131072> : tensor<20x30x1x3xi32>}> : () -> tensor<20x30x1x3xi32> // CHECK: %[[CAST:.*]] = "tf.Cast" // CHECK: %[[XLADOTV2_0:.*]] = "tf.XlaDotV2"(%[[CAST]] // CHECK: %[[SUB_0:.*]] = "tf.Sub"(%[[XLADOTV2_0]], %[[CONST]]) : (tensor<20x30x64x3xi32>, tensor<20x30x1x3xi32>) -> tensor<20x30x64x3xi32> @@ -602,7 +602,7 @@ module attributes {tf_saved_model.semantics} { } // CHECK-LABEL: func @broadcasting_weight_batch_matmul -// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() {value = dense<[2, 1024, 3]> : tensor<3xi64>} : () -> tensor<3xi64> +// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() <{value = dense<[2, 1024, 3]> : tensor<3xi64>}> : () -> tensor<3xi64> // CHECK: %[[CAST:.*]] = "tf.Cast" // CHECK: %[[BROADCAST_TO:.*]] = "tf.BroadcastTo"({{.*}}, %[[CONST]]) : (tensor<1024x3xi8>, tensor<3xi64>) -> tensor<2x1024x3xi8> // CHECK: %[[XLADOTV2_0:.*]] = "tf.XlaDotV2"(%[[CAST]], %[[BROADCAST_TO]]) @@ -639,8 +639,8 @@ module attributes {tf_saved_model.semantics} { } // CHECK-LABEL: func @broadcasting_input_batch_matmul -// CHECK-DAG: %[[WEIGHT:.*]] = "tf.Const"() {device = "", value = {{.*}} : tensor<2x2x1024x3xi8>} : () -> tensor<2x2x1024x3xi8> -// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() {value = dense<[2, 2, 1, 1024]> : tensor<4xi64>} : () -> tensor<4xi64> +// CHECK-DAG: %[[WEIGHT:.*]] = "tf.Const"() <{value = {{.*}} : tensor<2x2x1024x3xi8>}> {device = ""} : () -> tensor<2x2x1024x3xi8> +// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() <{value = dense<[2, 2, 1, 1024]> : tensor<4xi64>}> : () -> tensor<4xi64> // CHECK: %[[CAST:.*]] = "tf.Cast" // CHECK: %[[BROADCAST_TO:.*]] = "tf.BroadcastTo"(%[[CAST]], %[[CONST]]) : (tensor<2x1x1024xi8>, tensor<4xi64>) -> tensor<2x2x1x1024xi8> // CHECK: %[[XLADOTV2_0:.*]] = "tf.XlaDotV2"(%[[BROADCAST_TO]], %[[WEIGHT]]) @@ -677,14 +677,14 @@ module attributes {tf_saved_model.semantics} { } // CHECK-LABEL: func @dynamic_shape_batch_matmul -// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() {value = dense<0> : tensor<1xi32>} : () -> tensor<1xi32> -// CHECK-DAG: %[[CONST_1:.*]] = "tf.Const"() {value = dense<2> : tensor<1xi32>} : () -> tensor<1xi32> -// CHECK-DAG: %[[CONST_2:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32> -// CHECK-DAG: %[[CONST_3:.*]] = "tf.Const"() {value = dense<[1024, 3]> : tensor<2xi64>} : () -> tensor<2xi64> -// CHECK-DAG: %[[CONST_4:.*]] = "tf.Const"() {value = dense<> : tensor<0xi64>} : () -> tensor<0xi64> -// CHECK-DAG: %[[CONST_5:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor -// CHECK-DAG: %[[WEIGHT:.*]] = "tf.Const"() {device = "", value = {{.*}} : tensor<1024x3xi8>} : () -> tensor<1024x3xi8> -// CHECK: %[[CAST:.*]] = "tf.Cast"({{.*}}) {Truncate = false, device = ""} : (tensor) -> tensor +// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() <{value = dense<0> : tensor<1xi32>}> : () -> tensor<1xi32> +// CHECK-DAG: %[[CONST_1:.*]] = "tf.Const"() <{value = dense<2> : tensor<1xi32>}> : () -> tensor<1xi32> +// CHECK-DAG: %[[CONST_2:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}> : () -> tensor<1xi32> +// CHECK-DAG: %[[CONST_3:.*]] = "tf.Const"() <{value = dense<[1024, 3]> : tensor<2xi64>}> : () -> tensor<2xi64> +// CHECK-DAG: %[[CONST_4:.*]] = "tf.Const"() <{value = dense<> : tensor<0xi64>}> : () -> tensor<0xi64> +// CHECK-DAG: %[[CONST_5:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor +// CHECK-DAG: %[[WEIGHT:.*]] = "tf.Const"() <{{{value = .* : tensor<1024x3xi8>}}}> {device = ""} : () -> tensor<1024x3xi8> +// CHECK: %[[CAST:.*]] = "tf.Cast"({{.*}}) <{Truncate = false}> {device = ""} : (tensor) -> tensor // CHECK: %[[SHAPE:.*]] = "tf.Shape"(%[[CAST]]) : (tensor) -> tensor<3xi64> // CHECK: %[[SLICE_1:.*]] = "tf.Slice"(%[[SHAPE]], %[[CONST]], %[[CONST_2]]) : (tensor<3xi64>, tensor<1xi32>, tensor<1xi32>) -> tensor<1xi64> // CHECK: %[[SLICE_2:.*]] = "tf.Slice"(%[[SHAPE]], %[[CONST_2]], %[[CONST_1]]) : (tensor<3xi64>, tensor<1xi32>, tensor<1xi32>) -> tensor<2xi64> diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/tests/replace_cast_hacks_with_tf_xla_ops_large_constants.mlir b/tensorflow/compiler/mlir/quantization/tensorflow/tests/replace_cast_hacks_with_tf_xla_ops_large_constants.mlir index 775ab82e10501a..3c0c36684916f3 100644 --- a/tensorflow/compiler/mlir/quantization/tensorflow/tests/replace_cast_hacks_with_tf_xla_ops_large_constants.mlir +++ b/tensorflow/compiler/mlir/quantization/tensorflow/tests/replace_cast_hacks_with_tf_xla_ops_large_constants.mlir @@ -56,7 +56,7 @@ module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, p } // CHECK-LABEL: func @conv_with_filter_larger_than_1GB -// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() {value = dense<-237772800> : tensor<1x1x1x512xi32>} : () -> tensor<1x1x1x512xi32> +// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() <{value = dense<-237772800> : tensor<1x1x1x512xi32>}> : () -> tensor<1x1x1x512xi32> // CHECK: %[[PADV2_0:.*]] = "tf.PadV2" // CHECK: %[[XLACONVV2_0:.*]] = "tf.XlaConvV2"(%[[PADV2_0]] // CHECK: %[[SUB_0:.*]] = "tf.Sub"(%[[XLACONVV2_0]], %[[CONST]]) diff --git a/tensorflow/compiler/mlir/quantization/tensorflow/tests/unfreeze_constants.mlir b/tensorflow/compiler/mlir/quantization/tensorflow/tests/unfreeze_constants.mlir index ddf33e312ab077..b7b4fa1cc39aed 100644 --- a/tensorflow/compiler/mlir/quantization/tensorflow/tests/unfreeze_constants.mlir +++ b/tensorflow/compiler/mlir/quantization/tensorflow/tests/unfreeze_constants.mlir @@ -15,7 +15,7 @@ module attributes {tf_saved_model.semantics} { // CHECK-SAME: tf_saved_model.initializer_type = "restore_op" // Check that variable is initialized by assigning the const value within the initializer function. -// CHECK-DAG: %[[CST_0:.*]] = "tf.Const"() {value = dense<1.000000e+00> : tensor<8xf32>} +// CHECK-DAG: %[[CST_0:.*]] = "tf.Const"() <{value = dense<1.000000e+00> : tensor<8xf32>}> // CHECK-DAG: %[[VAR_HANDLE_0:.*]] = "tf.VarHandleOp"() {{.*shared_name = "const_0".*}} // CHECK-DAG: "tf.AssignVariableOp"(%[[VAR_HANDLE_0]], %[[CST_0]]) @@ -44,11 +44,11 @@ module attributes {tf_saved_model.semantics} { // CHECK-SAME: tf_saved_model.exported_names = ["tf_saved_model.session_initializer_restore_op"] // CHECK-SAME: tf_saved_model.initializer_type = "restore_op" -// CHECK-DAG: %[[CST_0:.*]] = "tf.Const"() {{{.*value = dense<1.000000e\+00> : tensor<8xf32>.*}}} +// CHECK-DAG: %[[CST_0:.*]] = "tf.Const"() <{{{.*value = dense<1.000000e\+00> : tensor<8xf32>.*}}}> // CHECK-DAG: %[[VAR_HANDLE_0:.*]] = "tf.VarHandleOp"() {{.*shared_name = "const_0".*}} // CHECK-DAG: "tf.AssignVariableOp"(%[[VAR_HANDLE_0]], %[[CST_0]]) -// CHECK-DAG: %[[CST_1:.*]] = "tf.Const"() {{{.*value = dense<2.000000e\+00> : tensor<8xf32>.*}}} +// CHECK-DAG: %[[CST_1:.*]] = "tf.Const"() <{{{.*value = dense<2.000000e\+00> : tensor<8xf32>.*}}}> // CHECK-DAG: %[[VAR_HANDLE_1:.*]] = "tf.VarHandleOp"() {{.*shared_name = "const_1".*}} // CHECK-DAG: "tf.AssignVariableOp"(%[[VAR_HANDLE_1]], %[[CST_1]]) @@ -84,11 +84,11 @@ module attributes {tf_saved_model.semantics} { // CHECK-SAME: tf_saved_model.exported_names = ["tf_saved_model.session_initializer_init"] // CHECK-SAME: tf_saved_model.initializer_type = "restore_op" -// CHECK-DAG: %[[CST_0:.*]] = "tf.Const"() {value = dense<1.000000e+00> : tensor<8xf32>} +// CHECK-DAG: %[[CST_0:.*]] = "tf.Const"() <{value = dense<1.000000e+00> : tensor<8xf32>}> // CHECK-DAG: %[[VAR_HANDLE_0:.*]] = "tf.VarHandleOp"() // CHECK-DAG: "tf.AssignVariableOp"(%[[VAR_HANDLE_0]], %[[CST_0]]) -// CHECK-DAG: %[[CST_1:.*]] = "tf.Const"() {value = dense<2.000000e+00> : tensor<8xf32>} +// CHECK-DAG: %[[CST_1:.*]] = "tf.Const"() <{value = dense<2.000000e+00> : tensor<8xf32>}> // CHECK-DAG: %[[VAR_HANDLE_1:.*]] = "tf.VarHandleOp"() // CHECK-DAG: "tf.AssignVariableOp"(%[[VAR_HANDLE_1]], %[[CST_1]]) @@ -123,7 +123,7 @@ module attributes {tf_saved_model.semantics} { // CHECK-SAME: tf_saved_model.exported_names = ["tf_saved_model.session_initializer_restore_op"] // CHECK-SAME: tf_saved_model.initializer_type = "restore_op" -// CHECK-DAG: %[[CST_0:.*]] = "tf.Const"() {value = dense<3.000000e+00> : tensor<8xf32>} +// CHECK-DAG: %[[CST_0:.*]] = "tf.Const"() <{value = dense<3.000000e+00> : tensor<8xf32>}> // CHECK-DAG: %[[VAR_HANDLE_0:.*]] = "tf.VarHandleOp"() // CHECK-DAG: "tf.AssignVariableOp"(%[[VAR_HANDLE_0]], %[[CST_0]]) @@ -185,7 +185,7 @@ module attributes {tf_saved_model.semantics} { // Check that `tf.VarHandleOp` is only created for the constant that is larger // than the threshold (16 bytes for this test). -// CHECK-DAG: %[[CST_0:.*]] = "tf.Const"() {{{.*value = dense<5.000000e\+00> : tensor<8xf32>.*}}} +// CHECK-DAG: %[[CST_0:.*]] = "tf.Const"() <{{{.*value = dense<5.000000e\+00> : tensor<8xf32>.*}}}> // CHECK-DAG: %[[VAR_HANDLE_0:.*]] = "tf.VarHandleOp"() {{.*shared_name = "const_0".*}} // CHECK-DAG: "tf.AssignVariableOp"(%[[VAR_HANDLE_0]], %[[CST_0]]) @@ -199,8 +199,8 @@ module attributes {tf_saved_model.semantics} { // CHECK: @serving_default // CHECK-DAG: %[[VAR_HANDLE_2:.*]] = "tf.VarHandleOp"() {{.*shared_name = "const_0".*}} : () -> tensor>> // CHECK-DAG: %[[READ_VAR_0:.*]] = "tf.ReadVariableOp"(%[[VAR_HANDLE_2]]) : (tensor>>) -> tensor<8xf32> -// CHECK-DAG: %[[CST_1:.*]] = "tf.Const"() {{{.*value = dense<5.000000e\+00> : tensor<4xf32>.*}}} -// CHECK-DAG: %[[AXIS:.*]] = "tf.Const"() {{{.*value = dense<0> : tensor.*}}} +// CHECK-DAG: %[[CST_1:.*]] = "tf.Const"() <{{{.*value = dense<5.000000e\+00> : tensor<4xf32>.*}}}> +// CHECK-DAG: %[[AXIS:.*]] = "tf.Const"() <{{{.*value = dense<0> : tensor.*}}}> // CHECK-DAG: %[[CONCAT:.*]] = "tf.ConcatV2"(%[[READ_VAR_0]], %[[CST_1]], %[[AXIS]]) // CHECK: return %[[CONCAT]] : tensor<12xf32> } @@ -214,7 +214,7 @@ module attributes {tf_saved_model.semantics} { module attributes {tf_saved_model.semantics} { // CHECK: func.func @init_func_restore_op() -// CHECK-DAG: %[[CST_0:.*]] = "tf.Const"() {value = dense<1.000000e+00> : tensor<8xf32>} +// CHECK-DAG: %[[CST_0:.*]] = "tf.Const"() <{value = dense<1.000000e+00> : tensor<8xf32>}> // Check that the variable's shared_name contains the fused loc's items joined // by the delimiter "_" and suffixed with a number. // CHECK-DAG: %[[VAR_HANDLE_0:.*]] = "tf.VarHandleOp"() {{.*shared_name = "apple_banana_0".*}} @@ -247,7 +247,7 @@ module attributes {tf_saved_model.semantics} { %cst_2 = "tf.Const"() {value = dense<1.0> : tensor<1x5x5x1024xf32>} : () -> tensor<1x5x5x1024xf32> // Check that these constants are unfrozen. // CHECK: func private @__inference_main - // CHECK: %[[VAR_HANDLE_0:.*]] = "tf.VarHandleOp"() {container = "", shared_name = "const_0"} : () -> tensor>> + // CHECK: %[[VAR_HANDLE_0:.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "const_0"}> : () -> tensor>> // CHECK: %[[READ_VAR_0:.*]] = "tf.ReadVariableOp"(%0) : (tensor>>) -> tensor<1x5x5x1024xf32> %0:3 = "tf.While"(%cst_0, %cst_1, %arg0) {T = [i32, i32, f32], _lower_using_switch_merge = true, _num_original_outputs = 4 : i64, _read_only_resource_inputs = [], body = @while_body, cond = @while_cond, device = "", is_stateless = true, output_shapes = [#tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<1x5x5x1024>], parallel_iterations = 10 : i64, shape_invariant} : (tensor, tensor, tensor<1x5x5x1024xf32>) -> (tensor, tensor, tensor<1x5x5x1024xf32>) %1 = "tf.AddV2"(%0#2, %cst_2) {device = ""} : (tensor<1x5x5x1024xf32>, tensor<1x5x5x1024xf32>) -> tensor<1x5x5x1024xf32> @@ -260,7 +260,7 @@ module attributes {tf_saved_model.semantics} { %cst_0 = "tf.Const"() {value = dense<1.0> : tensor<1x5x5x1024xf32>} : () -> tensor<1x5x5x1024xf32> // Check that these constants are remained in constants. // CHECK: func private @while_body - // CHECK-DAG: %[[CST_0:.*]]= "tf.Const"() {value = dense<1.000000e+00> : tensor<1x5x5x1024xf32>} : () -> tensor<1x5x5x1024xf32> + // CHECK-DAG: %[[CST_0:.*]]= "tf.Const"() <{value = dense<1.000000e+00> : tensor<1x5x5x1024xf32>}> : () -> tensor<1x5x5x1024xf32> %0 = "tf.AddV2"(%arg0, %cst) {device = ""} : (tensor, tensor) -> tensor %1 = "tf.Identity"(%0) {device = ""} : (tensor) -> tensor %2 = "tf.Identity"(%arg1) {device = ""} : (tensor) -> tensor @@ -269,13 +269,13 @@ module attributes {tf_saved_model.semantics} { return %1, %2, %5 : tensor, tensor, tensor<1x5x5x1024xf32> } - func.func private @while_cond(%arg0: tensor {tf._user_specified_name = "while/loop_counter"}, %arg1: tensor {tf._user_specified_name = "while/maximum_iterations"}, %arg2: tensor<1x5x5x1024xf32>) -> tensor + func.func private @while_cond(%arg0: tensor {tf._user_specified_name = "while/loop_counter"}, %arg1: tensor {tf._user_specified_name = "while/maximum_iterations"}, %arg2: tensor<1x5x5x1024xf32>) -> tensor attributes {tf._construction_context = "kEagerRuntime", tf._input_shapes = [#tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<1x5x5x1024>], tf._original_func_name = "while_cond_60"} { %cst = "tf.Const"() {value = dense<[0, 1, 2, 3]> : tensor<4xi32>} : () -> tensor<4xi32> %cst_0 = "tf.Const"() {value = dense<5.0> : tensor} : () -> tensor // Check that these constants are remained in constants. // CHECK: func private @while_cond - // CHECK-DAG: %[[CST:.*]]= "tf.Const"() {value = dense<[0, 1, 2, 3]> : tensor<4xi32>} : () -> tensor<4xi32> + // CHECK-DAG: %[[CST:.*]]= "tf.Const"() <{value = dense<[0, 1, 2, 3]> : tensor<4xi32>}> : () -> tensor<4xi32> %0 = "tf.Sum"(%arg2, %cst) {device = "", keep_dims = false} : (tensor<1x5x5x1024xf32>, tensor<4xi32>) -> tensor %1 = "tf.Less"(%0, %cst_0) {device = ""} : (tensor, tensor) -> tensor %2 = "tf.Identity"(%1) {device = ""} : (tensor) -> tensor diff --git a/tensorflow/compiler/mlir/tensorflow/BUILD b/tensorflow/compiler/mlir/tensorflow/BUILD index 78d87e81c45525..01f225c9abd595 100644 --- a/tensorflow/compiler/mlir/tensorflow/BUILD +++ b/tensorflow/compiler/mlir/tensorflow/BUILD @@ -353,11 +353,14 @@ cc_library( ":attribute_utils", ":convert_type", ":dynamic_shape_utils", + ":tensorflow_all_ops_inc_gen", ":tensorflow_attributes", ":tensorflow_op_interfaces", ":tensorflow_op_interfaces_inc_gen", + ":tensorflow_remaining_ops_inc_gen", ":tensorflow_side_effects", ":tensorflow_structs", + ":tensorflow_tfrt_ops_inc_gen", ":tensorflow_traits", ":tensorflow_types", ":tf_arith_ops_folder", @@ -369,6 +372,8 @@ cc_library( "//tensorflow/compiler/mlir/tensorflow/transforms:tensorflow_canonicalize_inc_gen", "//tensorflow/core:framework", "//tensorflow/core:lib", + "@com_google_absl//absl/log:check", + "@com_google_absl//absl/strings", "@llvm-project//llvm:Support", "@llvm-project//mlir:ControlFlowInterfaces", "@llvm-project//mlir:DerivedAttributeOpInterface", diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_op_base.td b/tensorflow/compiler/mlir/tensorflow/ir/tf_op_base.td index c3fdc7a3c83685..a5bb0051cc8fe4 100644 --- a/tensorflow/compiler/mlir/tensorflow/ir/tf_op_base.td +++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_op_base.td @@ -48,7 +48,6 @@ TODO: Make invariants more structured so that we can reference them in ops. }]; let cppNamespace = "::mlir::TF"; - let usePropertiesForAttributes = 0; } //===----------------------------------------------------------------------===// diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h b/tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h index f9fa3d13732451..e81742e90cf56f 100644 --- a/tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h +++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h @@ -19,6 +19,7 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_MLIR_TENSORFLOW_IR_TF_OPS_H_ #define TENSORFLOW_COMPILER_MLIR_TENSORFLOW_IR_TF_OPS_H_ +#include "mlir/Bytecode/BytecodeOpInterface.h" // from @llvm-project // IWYU pragma: keep #include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project #include "mlir/Dialect/Traits.h" // from @llvm-project #include "mlir/IR/Attributes.h" // from @llvm-project diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td b/tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td index 5e4d546c4d47bc..e170bae81a8bcf 100644 --- a/tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td +++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td @@ -312,7 +312,7 @@ def TF_YieldOp : TF_Op<"Yield", [Terminator, Pure, NativeOpTrait<"ReturnLike", [], "", "">, - ParentOneOf<["CaseRegionOp", "IfRegionOp", "WhileRegionOp"]>, + ParentOneOf<["CaseRegionOp", "IfRegionOp", "WhileRegionOp", "GeneratorDatasetRegionOp"]>, DeclareOpInterfaceMethods, ]> { @@ -389,6 +389,57 @@ else_branch: A region that computes the outputs of the op if cond = false. let hasCanonicalizer = 1; } +def TF_GeneratorDatasetRegionOp : TF_Op<"GeneratorDatasetRegion", + [AttrSizedOperandSegments, + DeclareOpInterfaceMethods, + SingleBlockImplicitTerminator<"YieldOp">, + TF_GeneratorOpSideEffect, + ]> { + let summary = "Regional version of GeneratorDataset"; + + let description = [{ +Creates a dataset that invokes its 'next' region to generate elements. Conceptually, +within MLIR, we treat this op as if it fills a buffer with all the results right away, +and those results are then passed (through the variant tensor result) to +MakeIterator / IteratorGetNext. Note that the actual TF implementation differs: It +generates the next element just in time, during IteratorGetNext. + +init_extra_args: Additional arguments to pass to 'init'. +next_extra_args: Additional arguments to pass to 'next'. (Passed after the + normal arguments which are from the return values of 'init'.) +finalize_extra_args: Additional arguments to pass to 'finalize'. (Passed after + the normal arguments which are from the return values of 'init'.) + }]; + + let arguments = (ins + Variadic:$init_func_other_args, + Variadic:$next_func_other_args, + Variadic:$finalize_func_other_args, + + ConfinedAttr]>:$output_types, + ConfinedAttr]>:$output_shapes, + DefaultValuedOptionalAttr:$metadata + ); + + let results = (outs + TF_VariantTensor:$handle + ); + + let regions = (region SizedRegion<1>:$init, + SizedRegion<1>:$next, + SizedRegion<1>:$finalize + ); + + TF_DerivedOperandTypeListAttr Tinit_func_args = TF_DerivedOperandTypeListAttr<0>; + TF_DerivedOperandTypeListAttr Tnext_func_args = TF_DerivedOperandTypeListAttr<1>; + TF_DerivedOperandTypeListAttr Tfinalize_func_args = TF_DerivedOperandTypeListAttr<2>; +} + def TF_LegacyCallOp : TF_Op<"LegacyCall", [CallOpInterface, DeclareOpInterfaceMethods, Pure]> { @@ -455,9 +506,7 @@ def TF_ParseExampleOp : TF_Op<"ParseExample", Variadic:$dense_keys, Variadic>:$dense_defaults, - TF_ShapeAttrArray:$dense_shapes, - DenseI32ArrayAttr:$resultSegmentSizes, - DenseI32ArrayAttr:$operandSegmentSizes + TF_ShapeAttrArray:$dense_shapes ); let results = (outs @@ -491,8 +540,7 @@ def TF_ParseExampleV2Op : TF_Op<"ParseExampleV2", Variadic>:$dense_defaults, ConfinedAttr]>:$num_sparse, - TF_ShapeAttrArray:$dense_shapes, - DenseI32ArrayAttr:$resultSegmentSizes + TF_ShapeAttrArray:$dense_shapes ); let results = (outs diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_ops_a_m.cc b/tensorflow/compiler/mlir/tensorflow/ir/tf_ops_a_m.cc index 4f21118dc24b71..cee0e40f9cfeb5 100644 --- a/tensorflow/compiler/mlir/tensorflow/ir/tf_ops_a_m.cc +++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_ops_a_m.cc @@ -19,55 +19,60 @@ limitations under the License. #include #include #include +#include #include -#include #include -#include -#include #include #include #include #include +#include "absl/log/check.h" +#include "absl/strings/str_cat.h" #include "llvm/ADT/APFloat.h" #include "llvm/ADT/APInt.h" #include "llvm/ADT/ArrayRef.h" +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/DenseSet.h" #include "llvm/ADT/STLExtras.h" +#include "llvm/ADT/STLFunctionalExtras.h" #include "llvm/ADT/Sequence.h" #include "llvm/ADT/SmallVector.h" -#include "llvm/ADT/StringExtras.h" #include "llvm/ADT/StringRef.h" -#include "llvm/ADT/StringSwitch.h" #include "llvm/ADT/iterator_range.h" -#include "llvm/Support/Casting.h" +#include "llvm/Support/ErrorHandling.h" #include "llvm/Support/FormatVariadic.h" #include "llvm/Support/raw_ostream.h" #include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project #include "mlir/Dialect/Traits.h" // from @llvm-project #include "mlir/IR/Attributes.h" // from @llvm-project #include "mlir/IR/Builders.h" // from @llvm-project +#include "mlir/IR/BuiltinAttributeInterfaces.h" // from @llvm-project #include "mlir/IR/BuiltinAttributes.h" // from @llvm-project -#include "mlir/IR/BuiltinOps.h" // from @llvm-project #include "mlir/IR/BuiltinTypes.h" // from @llvm-project #include "mlir/IR/Diagnostics.h" // from @llvm-project -#include "mlir/IR/DialectImplementation.h" // from @llvm-project #include "mlir/IR/Location.h" // from @llvm-project #include "mlir/IR/MLIRContext.h" // from @llvm-project #include "mlir/IR/Matchers.h" // from @llvm-project #include "mlir/IR/OpDefinition.h" // from @llvm-project #include "mlir/IR/OpImplementation.h" // from @llvm-project +#include "mlir/IR/OperationSupport.h" // from @llvm-project #include "mlir/IR/PatternMatch.h" // from @llvm-project +#include "mlir/IR/Region.h" // from @llvm-project +#include "mlir/IR/SymbolTable.h" // from @llvm-project #include "mlir/IR/TypeUtilities.h" // from @llvm-project #include "mlir/IR/Types.h" // from @llvm-project #include "mlir/IR/Value.h" // from @llvm-project #include "mlir/IR/ValueRange.h" // from @llvm-project +#include "mlir/Interfaces/CallInterfaces.h" // from @llvm-project #include "mlir/Interfaces/ControlFlowInterfaces.h" // from @llvm-project +#include "mlir/Interfaces/InferTypeOpInterface.h" // from @llvm-project #include "mlir/Parser/Parser.h" // from @llvm-project #include "mlir/Support/LLVM.h" // from @llvm-project #include "mlir/Support/LogicalResult.h" // from @llvm-project #include "mlir/Transforms/InliningUtils.h" // from @llvm-project #include "tensorflow/compiler/mlir/tensorflow/ir/tf_arith_ops_folder.h" -#include "tensorflow/compiler/mlir/tensorflow/ir/tf_attributes.h" +#include "tensorflow/compiler/mlir/tensorflow/ir/tf_op_interfaces.h" #include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h" #include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops_canonicalization_helper.h" #include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops_device_helper.h" @@ -75,12 +80,14 @@ limitations under the License. #include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops_tensor_helper.h" #include "tensorflow/compiler/mlir/tensorflow/ir/tf_side_effects.h" #include "tensorflow/compiler/mlir/tensorflow/ir/tf_structs.h" +#include "tensorflow/compiler/mlir/tensorflow/ir/tf_traits.h" #include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h" #include "tensorflow/compiler/mlir/tensorflow/transforms/rewrite_util.h" #include "tensorflow/compiler/mlir/tensorflow/utils/attribute_utils.h" #include "tensorflow/compiler/mlir/tensorflow/utils/dynamic_shape_utils.h" #include "tensorflow/core/framework/kernel_shape_util.h" #include "tensorflow/core/platform/logging.h" +#include "tensorflow/core/platform/status.h" #include "tensorflow/core/util/padding.h" #include "tensorflow/core/util/tensor_format.h" @@ -1708,9 +1715,10 @@ void ConstOp::build(OpBuilder& builder, OperationState& result, Type type, LogicalResult ConstOp::inferReturnTypes( MLIRContext* context, std::optional location, ValueRange operands, - DictionaryAttr attributes, OpaqueProperties, RegionRange regions, + DictionaryAttr attributes, OpaqueProperties properties, RegionRange regions, SmallVectorImpl& inferredReturnTypes) { - auto value = attributes.get("value"); + ConstOpAdaptor adaptor(operands, attributes, properties, regions); + auto value = adaptor.getValue(); if (!value) return emitOptionalError(location, "missing attribute 'value'"); if (auto elem_attr = value.dyn_cast()) { inferredReturnTypes.assign({elem_attr.getType()}); @@ -1951,13 +1959,13 @@ static LogicalResult inferConvReturnTypeComponents( LogicalResult Conv2DOp::inferReturnTypeComponents( MLIRContext* context, std::optional location, - ValueShapeRange operands, DictionaryAttr attributes, OpaqueProperties, - RegionRange regions, + ValueShapeRange operands, DictionaryAttr attributes, + OpaqueProperties properties, RegionRange regions, SmallVectorImpl& inferredReturnShapes) { - Conv2DOpAdaptor op(operands.getValues(), attributes); + Conv2DOpAdaptor op(operands.getValues(), attributes, properties, regions); ArrayRef explicit_padding; ArrayAttr explicit_pad = - attributes.get("explicit_paddings").dyn_cast_or_null<::mlir::ArrayAttr>(); + op.getExplicitPaddings().dyn_cast_or_null<::mlir::ArrayAttr>(); if (!explicit_pad) { explicit_pad = ::mlir::Builder(context).getI64ArrayAttr({}); } @@ -2150,17 +2158,12 @@ StringRef Conv2DBackpropInputOp::GetOptimalLayout( LogicalResult Conv3DOp::inferReturnTypeComponents( MLIRContext* context, std::optional location, - ValueShapeRange operands, DictionaryAttr attributes, OpaqueProperties, - RegionRange regions, + ValueShapeRange operands, DictionaryAttr attributes, + OpaqueProperties properties, RegionRange regions, SmallVectorImpl& inferredReturnShapes) { - Conv3DOpAdaptor op(operands.getValues(), attributes); - ArrayRef explicit_padding; - ArrayAttr explicit_pad = - attributes.get("explicit_paddings").dyn_cast_or_null<::mlir::ArrayAttr>(); - if (!explicit_pad) { - explicit_pad = ::mlir::Builder(context).getI64ArrayAttr({}); - } - explicit_padding = explicit_pad.getValue(); + Conv3DOpAdaptor op(operands.getValues(), attributes, properties, regions); + ArrayAttr explicit_pad = ::mlir::Builder(context).getI64ArrayAttr({}); + ArrayRef explicit_padding = explicit_pad.getValue(); return inferConvReturnTypeComponents(location, op, explicit_padding, inferredReturnShapes); @@ -2968,6 +2971,70 @@ StringRef FusedBatchNormV3Op::GetOptimalLayout(const RuntimeDevices& devices) { return ::mlir::TF::GetOptimalLayout(devices, this); } +//===----------------------------------------------------------------------===// +// GeneratorDatasetRegionOp +//===----------------------------------------------------------------------===// + +bool GeneratorDatasetRegionOp::areTypesCompatible(Type t1, Type t2) { + return true; // Don't enforce type checking across control-flow edges. +} + +void GeneratorDatasetRegionOp::getRegionInvocationBounds( + ArrayRef operands, + SmallVectorImpl& invocationBounds) { + // We invoke `init` once, `finalize` once, and `next` any number of times. + invocationBounds.emplace_back(InvocationBounds(1, 1)); // init + invocationBounds.emplace_back(InvocationBounds::getUnknown()); // next + invocationBounds.emplace_back(InvocationBounds(1, 1)); // finalize +} + +OperandRange GeneratorDatasetRegionOp::getEntrySuccessorOperands( + RegionBranchPoint point) { + auto end = this->getOperation()->operand_end(); + if (point.isParent()) { + // The op itself doesn't branch back to itself. + return ::mlir::OperandRange(end, end); + } else if (point.getRegionOrNull() == &getInit()) { + return getInitFuncOtherArgs(); + } else if (point.getRegionOrNull() == &getNext()) { + return getNextFuncOtherArgs(); + } else /* finalize region */ { + return getFinalizeFuncOtherArgs(); + } +} + +void GeneratorDatasetRegionOp::getSuccessorRegions( + RegionBranchPoint point, SmallVectorImpl& regions) { + int n; + if (point.isParent()) { + // The op itself branches to `init` first. + regions.push_back( + RegionSuccessor(&getInit(), getInit().front().getArguments())); + } else if (point.getRegionOrNull() == &getInit()) { + // `init` branches to `next`, passing along the arguments given to `init`'s + // yield. Said arguments precede the "other args". + n = getInitFuncOtherArgs().size(); + regions.push_back(RegionSuccessor( + &getNext(), getNext().front().getArguments().drop_back(n))); + } else if (point.getRegionOrNull() == &getNext()) { + // `next` branches to itself, or to `finalize`, passing all arguments given + // to `next`s yield. + + // The number of values we're passing along. + int num = getNext().front().getTerminator()->getNumOperands(); + + // The number of extra values from the parent ops that should go to `next` + // and `finalize`. + regions.push_back(RegionSuccessor( + &getNext(), getNext().front().getArguments().slice(0, num))); + regions.push_back(RegionSuccessor( + &getFinalize(), getFinalize().front().getArguments().slice(0, num))); + } else { + // `finalize` branches back to the op itself, not passing any arguments. + regions.push_back(RegionSuccessor()); + } +} + //===----------------------------------------------------------------------===// // GatherV2Op //===----------------------------------------------------------------------===// diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_ops_n_z.cc b/tensorflow/compiler/mlir/tensorflow/ir/tf_ops_n_z.cc index 4ac541799afd4f..01cbbb9a46967c 100644 --- a/tensorflow/compiler/mlir/tensorflow/ir/tf_ops_n_z.cc +++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_ops_n_z.cc @@ -18,11 +18,12 @@ limitations under the License. #include #include +#include +#include +#include #include -#include #include #include -#include #include #include #include @@ -34,16 +35,15 @@ limitations under the License. #include "llvm/ADT/BitVector.h" #include "llvm/ADT/DenseMap.h" #include "llvm/ADT/STLExtras.h" +#include "llvm/ADT/STLFunctionalExtras.h" #include "llvm/ADT/Sequence.h" #include "llvm/ADT/SmallVector.h" -#include "llvm/ADT/StringExtras.h" #include "llvm/ADT/StringRef.h" -#include "llvm/ADT/StringSwitch.h" #include "llvm/ADT/Twine.h" #include "llvm/ADT/iterator_range.h" #include "llvm/Support/Casting.h" #include "llvm/Support/FormatVariadic.h" -#include "llvm/Support/raw_ostream.h" +#include "llvm/Support/MathExtras.h" #include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project #include "mlir/Dialect/Traits.h" // from @llvm-project #include "mlir/IR/Attributes.h" // from @llvm-project @@ -53,26 +53,28 @@ limitations under the License. #include "mlir/IR/BuiltinTypeInterfaces.h" // from @llvm-project #include "mlir/IR/BuiltinTypes.h" // from @llvm-project #include "mlir/IR/Diagnostics.h" // from @llvm-project -#include "mlir/IR/DialectImplementation.h" // from @llvm-project #include "mlir/IR/Location.h" // from @llvm-project #include "mlir/IR/MLIRContext.h" // from @llvm-project #include "mlir/IR/Matchers.h" // from @llvm-project #include "mlir/IR/OpDefinition.h" // from @llvm-project -#include "mlir/IR/OpImplementation.h" // from @llvm-project +#include "mlir/IR/OperationSupport.h" // from @llvm-project #include "mlir/IR/PatternMatch.h" // from @llvm-project +#include "mlir/IR/Region.h" // from @llvm-project +#include "mlir/IR/SymbolTable.h" // from @llvm-project #include "mlir/IR/TypeRange.h" // from @llvm-project #include "mlir/IR/TypeUtilities.h" // from @llvm-project #include "mlir/IR/Types.h" // from @llvm-project #include "mlir/IR/Value.h" // from @llvm-project +#include "mlir/IR/ValueRange.h" // from @llvm-project #include "mlir/Interfaces/CallInterfaces.h" // from @llvm-project #include "mlir/Interfaces/ControlFlowInterfaces.h" // from @llvm-project +#include "mlir/Interfaces/InferTypeOpInterface.h" // from @llvm-project #include "mlir/Interfaces/SideEffectInterfaces.h" // from @llvm-project #include "mlir/Parser/Parser.h" // from @llvm-project #include "mlir/Support/LLVM.h" // from @llvm-project #include "mlir/Support/LogicalResult.h" // from @llvm-project #include "mlir/Transforms/InliningUtils.h" // from @llvm-project #include "tensorflow/compiler/mlir/tensorflow/ir/tf_arith_ops_folder.h" -#include "tensorflow/compiler/mlir/tensorflow/ir/tf_attributes.h" #include "tensorflow/compiler/mlir/tensorflow/ir/tf_op_interfaces.h" #include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h" #include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops_canonicalization_helper.h" @@ -80,14 +82,12 @@ limitations under the License. #include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops_layout_helper.h" #include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops_tensor_helper.h" #include "tensorflow/compiler/mlir/tensorflow/ir/tf_side_effects.h" -#include "tensorflow/compiler/mlir/tensorflow/ir/tf_structs.h" +#include "tensorflow/compiler/mlir/tensorflow/ir/tf_traits.h" #include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h" #include "tensorflow/compiler/mlir/tensorflow/transforms/rewrite_util.h" #include "tensorflow/compiler/mlir/tensorflow/utils/attribute_utils.h" #include "tensorflow/compiler/mlir/tensorflow/utils/convert_type.h" #include "tensorflow/compiler/mlir/tensorflow/utils/dynamic_shape_utils.h" -#include "tensorflow/core/platform/logging.h" -#include "tensorflow/core/util/tensor_format.h" namespace mlir { namespace TF { @@ -4389,6 +4389,12 @@ MutableOperandRange YieldOp::getMutableSuccessorOperands( this->getOperation(), 1, this->getOperation()->getOperands().size() - 1); } + } else if (auto regionOp = llvm::dyn_cast( + this->getOperation()->getParentOp())) { + if (®ionOp.getFinalize() == this->getOperation()->getParentRegion()) { + // `finalize`'s returns get discarded. + return MutableOperandRange(this->getOperation(), 0, 0); + } } return MutableOperandRange(this->getOperation()); } diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tfrt_ops.cc b/tensorflow/compiler/mlir/tensorflow/ir/tfrt_ops.cc index 8cce823ae5233c..b6b8c41fdfaa6d 100644 --- a/tensorflow/compiler/mlir/tensorflow/ir/tfrt_ops.cc +++ b/tensorflow/compiler/mlir/tensorflow/ir/tfrt_ops.cc @@ -15,14 +15,19 @@ limitations under the License. #include "tensorflow/compiler/mlir/tensorflow/ir/tfrt_ops.h" +#include + +#include "llvm/ADT/DenseMap.h" #include "llvm/ADT/STLExtras.h" #include "llvm/ADT/SmallVector.h" +#include "llvm/ADT/StringRef.h" #include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project #include "mlir/IR/BuiltinAttributes.h" // from @llvm-project -#include "mlir/IR/BuiltinOps.h" // from @llvm-project +#include "mlir/IR/TypeUtilities.h" // from @llvm-project +#include "mlir/IR/Value.h" // from @llvm-project +#include "mlir/Support/LogicalResult.h" // from @llvm-project #include "tensorflow/compiler/mlir/tensorflow/ir/tf_op_interfaces.h" #include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h" -#include "tensorflow/core/framework/resource_handle.h" //===----------------------------------------------------------------------===// // _TfrtGetResourceOp @@ -84,6 +89,43 @@ mlir::LogicalResult PwStreamResultsOp::verify() { return mlir::success(); } +//===----------------------------------------------------------------------===// +// IfrtProgramCall +//===----------------------------------------------------------------------===// + +mlir::LogicalResult IfrtCallOp::verify() { + auto func = getOperation()->getParentOfType(); + if (func != nullptr && func->hasAttr("tfrt_ifrt_serving.program_id")) { + return emitOpError() << "cannot be nested inside an IFRT program"; + } + + for (mlir::Value arg : getArgs()) { + if (mlir::getElementTypeOrSelf(arg.getType()) + .isa()) { + return emitOpError() + << "does not support passing '!tf.resource' values as arguments"; + } + } + if (getArgs().size() != getArgNames().size()) { + return emitOpError() + << "expects every argument to have an 'arg_name' attribute"; + } + + for (mlir::Value result : getResults()) { + if (mlir::getElementTypeOrSelf(result.getType()) + .isa()) { + return emitOpError() + << "does not support returning '!tf.resource' values as results"; + } + } + if (getArgs().size() != getArgNames().size()) { + return emitOpError() + << "expects every result to have a 'result_name' attribute"; + } + + return mlir::success(); +} + } // namespace TF } // namespace mlir diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tfrt_ops.h b/tensorflow/compiler/mlir/tensorflow/ir/tfrt_ops.h index e4a41d41e0c90e..c6c3eb213d1778 100644 --- a/tensorflow/compiler/mlir/tensorflow/ir/tfrt_ops.h +++ b/tensorflow/compiler/mlir/tensorflow/ir/tfrt_ops.h @@ -16,6 +16,7 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_MLIR_TENSORFLOW_IR_TFRT_OPS_H_ #define TENSORFLOW_COMPILER_MLIR_TENSORFLOW_IR_TFRT_OPS_H_ +#include "mlir/Bytecode/BytecodeOpInterface.h" // from @llvm-project // IWYU pragma: keep #include "mlir/IR/BuiltinTypes.h" // from @llvm-project #include "mlir/Interfaces/DerivedAttributeOpInterface.h" // from @llvm-project #include "mlir/Interfaces/InferTypeOpInterface.h" // from @llvm-project diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tfrt_ops.td b/tensorflow/compiler/mlir/tensorflow/ir/tfrt_ops.td index a0e2935255e95a..0684e188db695d 100644 --- a/tensorflow/compiler/mlir/tensorflow/ir/tfrt_ops.td +++ b/tensorflow/compiler/mlir/tensorflow/ir/tfrt_ops.td @@ -63,6 +63,42 @@ def TF__TfrtGetResourceOp : TF_Op<"_TfrtGetResource", let hasVerifier = 1; } +def TF_IfrtCallOp : TF_Op<"IfrtCall", []> { + let summary = "Invokes a program via IFRT on a device"; + + let description = [{ + This op calls an IFRT program uniquely identified by the given program id. + + During lowering from a `tf_device.cluster_func` op to a `tf.IfrtCall` op, + the region owned by the former will be outlined to a function with a + `tfrt_ifrt_serving.program_id` attribute. After that, the runtime ensures + that the outlined function is compiled into an executable and is available + for lookup from `IfrtCall` TF ops. + + Ifrt program execution is encapsulated by `ServingExecutable` + abstraction, which takes named arguments and returns named results. Thus, + this op takes `arg_names` and `result_names` attributes to convert between + positional arguments/results of the `tf.IfrtCall` op and named ones for + invoking `ServingExecutable`. This op also takes `variable_names` attribute + to bind the variables (weights). + }]; + + let arguments = (ins + Variadic : $args, + I64Attr : $program_id, + StrArrayAttr : $arg_names, + StrArrayAttr : $variable_names, + StrArrayAttr : $result_names + ); + + let results = (outs Variadic : $results); + + TF_DerivedOperandTypeListAttr Tin = TF_DerivedOperandTypeListAttr<0>; + TF_DerivedResultTypeListAttr Tout = TF_DerivedResultTypeListAttr<0>; + + let hasVerifier = 1; +} + // TODO(chky): Consider adding this op to tensorflow core ops. def TF_PwStreamResultsOp : TF_Op<"PwStreamResults"> { let summary = "Streams results back to the controller"; diff --git a/tensorflow/compiler/mlir/tensorflow/tests/batchmatmul_to_einsum.mlir b/tensorflow/compiler/mlir/tensorflow/tests/batchmatmul_to_einsum.mlir index 30dd1200d0bab7..4b6a33cd657473 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/batchmatmul_to_einsum.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/batchmatmul_to_einsum.mlir @@ -2,42 +2,42 @@ func.func @test_batch_matmul_to_einsum(%arg0: tensor<1x2x3xf32>, %arg1: tensor<3x4xf32>) -> tensor<1x2x4xf32> { // CHECK-LABEL: test_batch_matmul_to_einsum - // CHECK: "tf.Einsum"(%arg0, %arg1) {equation = "...mk,...kn->...mn"} : (tensor<1x2x3xf32>, tensor<3x4xf32>) -> tensor<1x2x4xf32> + // CHECK: "tf.Einsum"(%arg0, %arg1) <{equation = "...mk,...kn->...mn"}> : (tensor<1x2x3xf32>, tensor<3x4xf32>) -> tensor<1x2x4xf32> %0 = "tf.BatchMatMul"(%arg0, %arg1) {adj_x = false, adj_y = false} : (tensor<1x2x3xf32>, tensor<3x4xf32>) -> tensor<1x2x4xf32> func.return %0: tensor<1x2x4xf32> } func.func @test_batch_matmul_broadcast_to_einsum(%arg0: tensor<2x2x4xf32>, %arg1: tensor<2x4x2xf32>) -> tensor<2x2x2xf32> { // CHECK-LABEL: test_batch_matmul_broadcast_to_einsum - // CHECK: "tf.Einsum"(%arg0, %arg1) {equation = "...mk,...kn->...mn"} : (tensor<2x2x4xf32>, tensor<2x4x2xf32>) -> tensor<2x2x2xf32> + // CHECK: "tf.Einsum"(%arg0, %arg1) <{equation = "...mk,...kn->...mn"}> : (tensor<2x2x4xf32>, tensor<2x4x2xf32>) -> tensor<2x2x2xf32> %0 = "tf.BatchMatMul"(%arg0, %arg1) {adj_x = false, adj_y = false} : (tensor<2x2x4xf32>, tensor<2x4x2xf32>) -> tensor<2x2x2xf32> func.return %0: tensor<2x2x2xf32> } func.func @test_batch_matmul_dynamic_shape_both_arg_to_einsum(%arg0: tensor<1x2x?xf32>, %arg1: tensor) -> tensor<1x2x4xf32> { // CHECK-LABEL: test_batch_matmul_dynamic_shape_both_arg_to_einsum - // CHECK: "tf.Einsum"(%arg0, %arg1) {equation = "...mk,...kn->...mn"} : (tensor<1x2x?xf32>, tensor) -> tensor<1x2x4xf32> + // CHECK: "tf.Einsum"(%arg0, %arg1) <{equation = "...mk,...kn->...mn"}> : (tensor<1x2x?xf32>, tensor) -> tensor<1x2x4xf32> %0 = "tf.BatchMatMul"(%arg0, %arg1) {adj_x = false, adj_y = false} : (tensor<1x2x?xf32>, tensor) -> tensor<1x2x4xf32> func.return %0: tensor<1x2x4xf32> } func.func @test_batch_matmul_dynamic_shape_one_arg_to_einsum(%arg0: tensor<1x2x?xf32>, %arg1: tensor<3x4xf32>) -> tensor<1x2x4xf32> { // CHECK-LABEL: test_batch_matmul_dynamic_shape_one_arg_to_einsum - // CHECK: "tf.Einsum"(%arg0, %arg1) {equation = "...mk,...kn->...mn"} : (tensor<1x2x?xf32>, tensor<3x4xf32>) -> tensor<1x2x4xf32> + // CHECK: "tf.Einsum"(%arg0, %arg1) <{equation = "...mk,...kn->...mn"}> : (tensor<1x2x?xf32>, tensor<3x4xf32>) -> tensor<1x2x4xf32> %0 = "tf.BatchMatMul"(%arg0, %arg1) {adj_x = false, adj_y = false} : (tensor<1x2x?xf32>, tensor<3x4xf32>) -> tensor<1x2x4xf32> func.return %0: tensor<1x2x4xf32> } func.func @test_batch_matmul_adj_to_einsum(%arg0: tensor<1x2x3xf32>, %arg1: tensor<4x3xf32>) -> tensor<1x2x4xf32> { // CHECK-LABEL: test_batch_matmul_adj_to_einsum - // CHECK: %[[RES_EINSUM:[0-9]*]] = "tf.Einsum"(%arg0, %arg1) {equation = "...mk,...nk->...mn"} : (tensor<1x2x3xf32>, tensor<4x3xf32>) -> tensor<1x2x4xf32> + // CHECK: %[[RES_EINSUM:[0-9]*]] = "tf.Einsum"(%arg0, %arg1) <{equation = "...mk,...nk->...mn"}> : (tensor<1x2x3xf32>, tensor<4x3xf32>) -> tensor<1x2x4xf32> // CHECK: return %[[RES_EINSUM]] : tensor<1x2x4xf32> %0 = "tf.BatchMatMul"(%arg0, %arg1) {adj_x = false, adj_y = true} : (tensor<1x2x3xf32>, tensor<4x3xf32>) -> tensor<1x2x4xf32> func.return %0: tensor<1x2x4xf32> } func.func @test_batch_matmulV2_adj_to_einsum(%arg0: tensor<1x3x2xf32>, %arg1: tensor<3x4xf32>) -> tensor<1x2x4xf32> { - // CHECK: %[[RES_EINSUM:[0-9]*]] = "tf.Einsum"(%arg0, %arg1) {equation = "...km,...kn->...mn"} : (tensor<1x3x2xf32>, tensor<3x4xf32>) -> tensor<1x2x4xf32> + // CHECK: %[[RES_EINSUM:[0-9]*]] = "tf.Einsum"(%arg0, %arg1) <{equation = "...km,...kn->...mn"}> : (tensor<1x3x2xf32>, tensor<3x4xf32>) -> tensor<1x2x4xf32> // CHECK: return %[[RES_EINSUM]] : tensor<1x2x4xf32> %0 = "tf.BatchMatMulV2"(%arg0, %arg1) {adj_x = true, adj_y = false} : (tensor<1x3x2xf32>, tensor<3x4xf32>) -> tensor<1x2x4xf32> func.return %0: tensor<1x2x4xf32> diff --git a/tensorflow/compiler/mlir/tensorflow/tests/breakup-islands.mlir b/tensorflow/compiler/mlir/tensorflow/tests/breakup-islands.mlir index 64b0bfb1a6202c..e94cb5f859ec34 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/breakup-islands.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/breakup-islands.mlir @@ -62,10 +62,10 @@ func.func @multiple_islands(%arg0: tensor<*xi32>, %arg1: tensor) -> (tensor // CHECK: %[[SUB1:.*]], %[[SUB1_control:.*]] = tf_executor.island(%[[ADD2_control]]) wraps "tf.Sub"(%arg0, %arg1) // CHECK: %[[MUL:.*]], %[[MUL_control:.*]] = tf_executor.island wraps "tf.Mul"(%[[SUB1]], %arg1) // CHECK: %[[SUB2:.*]], %[[SUB2_control:.*]] = tf_executor.island(%[[ADD2_control]], %[[MUL_control]]) wraps "tf.Sub"(%[[ADD1]], %[[SUB1]]) -// CHECK: %[[PRINT1:.*]], %[[PRINT1_control:.*]] = tf_executor.island wraps "tf.Print"(%[[SUB2]]) {message = "sub result"} +// CHECK: %[[PRINT1:.*]], %[[PRINT1_control:.*]] = tf_executor.island wraps "tf.Print"(%[[SUB2]]) <{message = "sub result"}> // CHECK: %[[ISLAND1:.*]] = tf_executor.island(%[[ADD2_control]], %[[MUL_control]]) wraps "tf.NoOp"() // CHECK: %[[ADD3:.*]], %[[ADD3_control:.*]] = tf_executor.island(%[[ISLAND1]], %[[ADD2_control]]) wraps "tf.Add"(%[[ADD2]], %[[ADD2]]) -// CHECK: %[[PRINT2:.*]], %[[PRINT2_control:.*]] = tf_executor.island wraps "tf.Print"(%[[ADD3]]) {message = "add result"} +// CHECK: %[[PRINT2:.*]], %[[PRINT2_control:.*]] = tf_executor.island wraps "tf.Print"(%[[ADD3]]) <{message = "add result"}> // CHECK: tf_executor.fetch %[[ADD2]], %[[MUL]], %[[PRINT1_control]], %[[PRINT2_control:.*]] : // CHECK: } // CHECK: return %[[GRAPH]]#0, %[[GRAPH]]#1 @@ -87,7 +87,7 @@ func.func @dangling_print(%arg0: tensor<*xi32>, %arg1: tensor) -> (tensor<* // CHECK: %[[GRAPH:.*]]:2 = tf_executor.graph { // CHECK: %[[ADD1:.*]], %[[ADD1_control:.*]] = tf_executor.island wraps "tf.Add"(%arg0, %arg1) // CHECK: %[[ADD2:.*]], %[[ADD2_control:.*]] = tf_executor.island wraps "tf.Add"(%[[ADD1_control:.*]], %arg1) -// CHECK: %[[PRINT:.*]], %[[PRINT_control:.*]] = tf_executor.island wraps "tf.Print"(%[[ADD2_control:.*]]) {message = "add result"} +// CHECK: %[[PRINT:.*]], %[[PRINT_control:.*]] = tf_executor.island wraps "tf.Print"(%[[ADD2_control:.*]]) <{message = "add result"}> // CHECK: tf_executor.fetch %[[ADD1]], %[[ADD2]], %[[PRINT_control]] : // CHECK: } // CHECK: return %[[GRAPH]]#0, %[[GRAPH]]#1 @@ -116,11 +116,11 @@ func.func @switch_and_merge(%arg0: tensor<*xi32>, %arg1: tensor) -> (tensor // CHECK: %[[GRAPH:.*]]:2 = tf_executor.graph { // CHECK: %[[ADD1:.*]], %[[ADD1_control:.*]] = tf_executor.island wraps "tf.Add"(%arg0, %arg1) // CHECK: %[[LESS:.*]], %[[LESS_control:.*]] = tf_executor.island wraps "tf.Less"(%arg1, %arg1) -// CHECK: %[[PRINT1:.*]], %[[PRINT1_control:.*]] = tf_executor.island wraps "tf.Print"(%[[ADD1]]) {message = "add result 1"} +// CHECK: %[[PRINT1:.*]], %[[PRINT1_control:.*]] = tf_executor.island wraps "tf.Print"(%[[ADD1]]) <{message = "add result 1"}> // CHECK: %[[ISLAND1:.*]] = tf_executor.island(%[[LESS_control]], %[[PRINT1_control]]) wraps "tf.NoOp"() // CHECK: %[[SWITCH_false:.*]], %[[SWITCH_true:.*]], {{.*}} = tf_executor.Switch %[[ADD1]], %[[LESS]], %[[ISLAND1]] // CHECK: %[[ADD2:.*]], %[[ADD2_control:.*]] = tf_executor.island wraps "tf.Add"(%[[SWITCH_false]], %arg1) -// CHECK: %[[PRINT2:.*]], %[[PRINT2_control:.*]] = tf_executor.island wraps "tf.Print"(%[[ADD2]]) {message = "add result 2"} +// CHECK: %[[PRINT2:.*]], %[[PRINT2_control:.*]] = tf_executor.island wraps "tf.Print"(%[[ADD2]]) <{message = "add result 2"}> // CHECK: %[[MERGE:.*]], %[[MERGE_index:.*]], %{{.*}} = tf_executor.Merge %[[ADD2]], %[[SWITCH_true]], %[[PRINT2_control]] // CHECK: tf_executor.fetch %[[MERGE]], %[[MERGE_index]] // CHECK: } @@ -141,7 +141,7 @@ func.func @control_flow_plumbing(%arg0: tensor<*xi32>, %arg1: tensor) -> te // CHECK-LABEL: func @control_flow_plumbing // CHECK: %[[GRAPH:.*]] = tf_executor.graph { -// CHECK: %[[PRINT:.*]], %[[PRINT_control:.*]] = tf_executor.island wraps "tf.Print"(%arg0) {message = "Random Print"} +// CHECK: %[[PRINT:.*]], %[[PRINT_control:.*]] = tf_executor.island wraps "tf.Print"(%arg0) <{message = "Random Print"}> // CHECK: %[[ADD1:.*]], %[[ADD1_control:.*]] = tf_executor.island(%[[PRINT_control]]) wraps "tf.Add"(%arg0, %arg1) // CHECK: %[[ADD2:.*]], %[[ADD2_control:.*]] = tf_executor.island wraps "tf.Add"(%[[ADD1]], %arg1) // CHECK: tf_executor.fetch %[[ADD2]] : tensor<*xi32> @@ -193,7 +193,7 @@ func.func @non_aliasing_reads_writes( // CHECK: %[[READ0:.*]], %[[READ0_CONTROL:.*]] = tf_executor.island wraps "tf.ReadVariableOp"(%arg0) // CHECK: %[[ASSIGN0_CONTROL:.*]] = tf_executor.island(%[[READ0_CONTROL]]) wraps "tf.AssignVariableOp"(%arg0, %arg2) // CHECK: %[[READ1:.*]], %[[READ1_CONTROL:.*]] = tf_executor.island wraps "tf.ReadVariableOp"(%arg1) -// CHECK: %[[VH0:.*]], %[[VH0_CONTROL:.*]] = tf_executor.island wraps "tf.VarHandleOp"() {container = "c", shared_name = "v0"} +// CHECK: %[[VH0:.*]], %[[VH0_CONTROL:.*]] = tf_executor.island wraps "tf.VarHandleOp"() <{container = "c", shared_name = "v0"}> // CHECK: %[[READ2:.*]], %[[READ2_CONTROL:.*]] = tf_executor.island wraps "tf.ReadVariableOp"(%[[VH0]]) // CHECK: %[[ASSIGN1_CONTROL:.*]] = tf_executor.island(%[[READ1_CONTROL]]) wraps "tf.AssignVariableOp"(%arg1, %[[READ0:.*]]) // CHECK: %[[ASSIGN2_CONTROL:.*]] = tf_executor.island(%[[ASSIGN0_CONTROL]]) wraps "tf.AssignVariableOp"(%arg0, %[[READ2]]) @@ -222,8 +222,8 @@ func.func @unknown_side_effecting_op(%arg0: tensor<32xf32>) -> () { // CHECK-LABEL: func @unknown_side_effecting_op // CHECK: tf_executor.graph { -// CHECK: %[[VH0:.*]], %[[VH0_CONTROL:.*]] = tf_executor.island wraps "tf.VarHandleOp"() {container = "c", shared_name = "v0"} -// CHECK: %[[VH1:.*]], %[[VH1_CONTROL:.*]] = tf_executor.island wraps "tf.VarHandleOp"() {container = "c", shared_name = "v1"} +// CHECK: %[[VH0:.*]], %[[VH0_CONTROL:.*]] = tf_executor.island wraps "tf.VarHandleOp"() <{container = "c", shared_name = "v0"}> +// CHECK: %[[VH1:.*]], %[[VH1_CONTROL:.*]] = tf_executor.island wraps "tf.VarHandleOp"() <{container = "c", shared_name = "v1"}> // CHECK: %[[READ0:.*]], %[[READ0_CONTROL:.*]] = tf_executor.island wraps "tf.ReadVariableOp"(%[[VH0]]) // CHECK: %[[ASSIGN0_CONTROL:.*]] = tf_executor.island wraps "tf.AssignVariableOp"(%[[VH1]], %arg0) // CHECK: %[[UNKNOWN_CONTROL:.*]] = tf_executor.island(%[[READ0_CONTROL]], %[[ASSIGN0_CONTROL]]) wraps "tf._UnknownSideEffectingOp_"() diff --git a/tensorflow/compiler/mlir/tensorflow/tests/canonicalize.mlir b/tensorflow/compiler/mlir/tensorflow/tests/canonicalize.mlir index 2adf88449d3853..612f01ce23ce8a 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/canonicalize.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/canonicalize.mlir @@ -19,29 +19,29 @@ func.func @tfAssertFalse(%arg0: tensor<1x1x6x2xf32>) { // CHECK-LABEL: testGatherToV2 // Ensures that axis param and batch_dims attr use their default values of 0. func.func @testGatherToV2(%params: tensor<4x3xf32>, %indices: tensor<1x2xi32>) -> tensor<2x3xf32> { - // CHECK: %[[AXIS:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor - // CHECK: "tf.GatherV2"(%arg0, %arg1, %[[AXIS]]) {batch_dims = 0 : i64, device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<4x3xf32>, tensor<1x2xi32>, tensor) -> tensor<2x3xf32> + // CHECK: %[[AXIS:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor + // CHECK: "tf.GatherV2"(%arg0, %arg1, %[[AXIS]]) <{batch_dims = 0 : i64}> {device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<4x3xf32>, tensor<1x2xi32>, tensor) -> tensor<2x3xf32> %0 = "tf.Gather"(%params, %indices) {device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<4x3xf32>, tensor<1x2xi32>) -> tensor<2x3xf32> func.return %0: tensor<2x3xf32> } // CHECK-LABEL: testBatchMatMulToV2 func.func @testBatchMatMulToV2(%arg0: tensor<2x3x5xf32>, %arg1: tensor<2x5x7xf32>) -> tensor<2x3x7xf32> { - // CHECK: "tf.BatchMatMulV2"(%arg0, %arg1) {adj_x = false, adj_y = false, device = "/job:localhost/replica:0/task:0/device:GPU:0"} - %0 = "tf.BatchMatMul"(%arg0, %arg1) {adj_x = false, adj_y = false, device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<2x3x5xf32>, tensor<2x5x7xf32>) -> tensor<2x3x7xf32> + // CHECK: "tf.BatchMatMulV2"(%arg0, %arg1) <{adj_x = false, adj_y = false}> {device = "/job:localhost/replica:0/task:0/device:GPU:0"} + %0 = "tf.BatchMatMul"(%arg0, %arg1) <{adj_x = false, adj_y = false}> {device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<2x3x5xf32>, tensor<2x5x7xf32>) -> tensor<2x3x7xf32> func.return %0: tensor<2x3x7xf32> } // CHECK-LABEL: testDynamicBatchMatMulToV2 func.func @testDynamicBatchMatMulToV2(%arg0: tensor<2x3x5xf32>, %arg1: tensor) -> tensor<2x3x7xf32> { - // CHECK: "tf.BatchMatMul"(%arg0, %arg1) {adj_x = false, adj_y = false, device = "/job:localhost/replica:0/task:0/device:GPU:0"} + // CHECK: "tf.BatchMatMul"(%arg0, %arg1) <{adj_x = false, adj_y = false}> {device = "/job:localhost/replica:0/task:0/device:GPU:0"} %0 = "tf.BatchMatMul"(%arg0, %arg1) {adj_x = false, adj_y = false, device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<2x3x5xf32>, tensor) -> tensor<2x3x7xf32> func.return %0: tensor<2x3x7xf32> } // CHECK-LABEL: testBatchMatMulToMatMul func.func @testBatchMatMulToMatMul(%arg0: tensor<2x3xf32>, %arg1: tensor<3x2xf32>) -> tensor<2x2xf32> { - // CHECK: %0 = "tf.MatMul"(%arg0, %arg1) {device = "/job:localhost/replica:0/task:0/device:GPU:0", transpose_a = false, transpose_b = false} : (tensor<2x3xf32>, tensor<3x2xf32>) -> tensor<2x2xf32> + // CHECK: %0 = "tf.MatMul"(%arg0, %arg1) <{transpose_a = false, transpose_b = false}> {device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<2x3xf32>, tensor<3x2xf32>) -> tensor<2x2xf32> %0 = "tf.BatchMatMul"(%arg0, %arg1) {adj_x = false, adj_y = false, device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<2x3xf32>, tensor<3x2xf32>) -> tensor<2x2xf32> // CHECK: return %0 func.return %0: tensor<2x2xf32> @@ -49,7 +49,7 @@ func.func @testBatchMatMulToMatMul(%arg0: tensor<2x3xf32>, %arg1: tensor<3x2xf32 // CHECK-LABEL: testBatchMatMulV2ToMatMul func.func @testBatchMatMulV2ToMatMul(%arg0: tensor<4x3xf32>, %arg1: tensor<4x5xf32>) -> tensor<3x5xf32> { - // CHECK: %0 = "tf.MatMul"(%arg0, %arg1) {device = "/job:localhost/replica:0/task:0/device:GPU:0", transpose_a = true, transpose_b = false} : (tensor<4x3xf32>, tensor<4x5xf32>) -> tensor<3x5xf32> + // CHECK: %0 = "tf.MatMul"(%arg0, %arg1) <{transpose_a = true, transpose_b = false}> {device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<4x3xf32>, tensor<4x5xf32>) -> tensor<3x5xf32> %0 = "tf.BatchMatMulV2"(%arg0, %arg1) {adj_x = true, adj_y = false, device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<4x3xf32>, tensor<4x5xf32>) -> tensor<3x5xf32> // CHECK: return %0 func.return %0: tensor<3x5xf32> @@ -58,7 +58,7 @@ func.func @testBatchMatMulV2ToMatMul(%arg0: tensor<4x3xf32>, %arg1: tensor<4x5xf // CHECK-LABEL: testBiasAddV1ToBiasAdd func.func @testBiasAddV1ToBiasAdd(%arg0: tensor<*xf32>, %arg1: tensor<128xf32>) -> tensor<*xf32> { - // CHECK: "tf.BiasAdd"(%arg0, %arg1) {data_format = "NHWC", device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<*xf32>, tensor<128xf32>) -> tensor<*xf32> + // CHECK: "tf.BiasAdd"(%arg0, %arg1) <{data_format = "NHWC"}> {device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<*xf32>, tensor<128xf32>) -> tensor<*xf32> %0 = "tf.BiasAddV1"(%arg0, %arg1) {device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<*xf32>, tensor<128xf32>) -> tensor<*xf32> func.return %0: tensor<*xf32> } @@ -124,8 +124,8 @@ func.func @testDifferentCastType(%arg0: tensor<8x16x32x64xf32>) -> (tensor<8x16x %1 = "tf.Cast"(%arg0) {Truncate = true} : (tensor<8x16x32x64xf32>) -> tensor<8x16x32x64xi32> func.return %0, %1: tensor<8x16x32x64xi32>, tensor<8x16x32x64xi32> - // CHECK: %0 = "tf.Cast"(%arg0) {Truncate = false} : (tensor<8x16x32x64xf32>) -> tensor<8x16x32x64xi32> - // CHECK: %1 = "tf.Cast"(%arg0) {Truncate = true} : (tensor<8x16x32x64xf32>) -> tensor<8x16x32x64xi32> + // CHECK: %0 = "tf.Cast"(%arg0) <{Truncate = false}> : (tensor<8x16x32x64xf32>) -> tensor<8x16x32x64xi32> + // CHECK: %1 = "tf.Cast"(%arg0) <{Truncate = true}> : (tensor<8x16x32x64xf32>) -> tensor<8x16x32x64xi32> // CHECK: return %0, %1 } @@ -135,8 +135,8 @@ func.func @testCompatibleCastType(%arg0: tensor) -> (tensor<10xf32>, tens %1 = "tf.Cast"(%arg0) {Truncate = true} : (tensor) -> tensor<10xf32> func.return %0, %1: tensor<10xf32>, tensor<10xf32> - // CHECK: %0 = "tf.Cast"(%arg0) {Truncate = false} : (tensor) -> tensor<10xf32> - // CHECK: %1 = "tf.Cast"(%arg0) {Truncate = true} : (tensor) -> tensor<10xf32> + // CHECK: %0 = "tf.Cast"(%arg0) <{Truncate = false}> : (tensor) -> tensor<10xf32> + // CHECK: %1 = "tf.Cast"(%arg0) <{Truncate = true}> : (tensor) -> tensor<10xf32> // CHECK: return %0, %1 } @@ -181,11 +181,11 @@ func.func @testConcatCwiseUnary(%arg0: tensor, %arg1: tensor, func.func @testConcatCwiseBinaryOnInnerDim(%arg0: tensor, %arg1: tensor, %arg2: tensor, %arg3: tensor) -> tensor { - // CHECK-DAG: %[[LHS_AXIS:.*]] = "tf.Const"() {value = dense<1> : tensor} + // CHECK-DAG: %[[LHS_AXIS:.*]] = "tf.Const"() <{value = dense<1> : tensor}> - // CHECK: %[[ADD_LHS_CONCAT:.*]] = "tf.Pack"(%arg2, %arg3) {axis = 0 : i64, device = "/job:localhost/replica:0/task:0/device:GPU:0"} + // CHECK: %[[ADD_LHS_CONCAT:.*]] = "tf.Pack"(%arg2, %arg3) <{axis = 0 : i64}> {device = "/job:localhost/replica:0/task:0/device:GPU:0"} // CHECK: %[[MUL_LHS_CONCAT:.*]] = "tf.ConcatV2"(%arg0, %arg1, %[[LHS_AXIS]]) {device = "/job:localhost/replica:0/task:0/device:GPU:0"} - // CHECK: %[[MUL_RHS_CONCAT:.*]] = "tf.Pack"(%arg2, %arg3) {axis = 0 : i64, device = "/job:localhost/replica:0/task:0/device:GPU:0"} + // CHECK: %[[MUL_RHS_CONCAT:.*]] = "tf.Pack"(%arg2, %arg3) <{axis = 0 : i64}> {device = "/job:localhost/replica:0/task:0/device:GPU:0"} // CHECK: %[[MUL:.*]] = "tf.Mul"(%[[MUL_LHS_CONCAT]], %[[MUL_RHS_CONCAT]]) {device = "/job:localhost/replica:0/task:0/device:GPU:0"} // CHECK-SAME: (tensor, tensor<2xf32>) -> tensor @@ -209,11 +209,11 @@ func.func @testConcatCwiseBinaryOnInnerDim(%arg0: tensor, func.func @testConcatCwiseBinaryPreserveAxisType(%arg0: tensor, %arg1: tensor, %arg2: tensor, %arg3: tensor) -> tensor { - // CHECK-DAG: %[[LHS_AXIS:.*]] = "tf.Const"() {value = dense<1> : tensor} + // CHECK-DAG: %[[LHS_AXIS:.*]] = "tf.Const"() <{value = dense<1> : tensor}> - // CHECK: %[[ADD_LHS_CONCAT:.*]] = "tf.Pack"(%arg2, %arg3) {axis = 0 : i64, device = "/job:localhost/replica:0/task:0/device:GPU:0"} + // CHECK: %[[ADD_LHS_CONCAT:.*]] = "tf.Pack"(%arg2, %arg3) <{axis = 0 : i64}> {device = "/job:localhost/replica:0/task:0/device:GPU:0"} // CHECK: %[[MUL_LHS_CONCAT:.*]] = "tf.ConcatV2"(%arg0, %arg1, %[[LHS_AXIS]]) {device = "/job:localhost/replica:0/task:0/device:GPU:0"} - // CHECK: %[[MUL_RHS_CONCAT:.*]] = "tf.Pack"(%arg2, %arg3) {axis = 0 : i64, device = "/job:localhost/replica:0/task:0/device:GPU:0"} + // CHECK: %[[MUL_RHS_CONCAT:.*]] = "tf.Pack"(%arg2, %arg3) <{axis = 0 : i64}> {device = "/job:localhost/replica:0/task:0/device:GPU:0"} // CHECK: %[[MUL:.*]] = "tf.Mul"(%[[MUL_LHS_CONCAT]], %[[MUL_RHS_CONCAT]]) // CHECK-SAME: {device = "/job:localhost/replica:0/task:0/device:GPU:0"} @@ -287,8 +287,8 @@ func.func @testConcatCwiseBinaryNegativeAxis(%arg0: tensor, // Synthesize binary ops when 1 of the 3 concat inputs is a non-binary op. // CHECK-LABEL: testConcatCwiseBinarySynthMulOp3Inputs func.func @testConcatCwiseBinarySynthMulOp3Inputs(%arg0: tensor, %arg1: tensor, %arg2: tensor) -> tensor { - // CHECK: %[[CONST:.*]] = "tf.Const"() {value = dense<1> : tensor} : () -> tensor - // CHECK-NEXT: %[[CONST0:.*]] = "tf.Const"() {value = dense<[2.000000e+00, 3.000000e+00, 1.000000e+00]> + // CHECK: %[[CONST:.*]] = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor + // CHECK-NEXT: %[[CONST0:.*]] = "tf.Const"() <{value = dense<[2.000000e+00, 3.000000e+00, 1.000000e+00]> // CHECK: %[[CONCAT:.*]] = "tf.ConcatV2"(%arg0, %arg1, %arg2, %[[CONST]]) {device = "/job:localhost/replica:0/task:0/device:GPU:0"} // CHECK: "tf.Mul"(%[[CONCAT]], %[[CONST0]]) {device = "/job:localhost/replica:0/task:0/device:GPU:0"} %axis = "tf.Const"() { value = dense<1> : tensor } : () -> tensor @@ -303,7 +303,7 @@ func.func @testConcatCwiseBinarySynthMulOp3Inputs(%arg0: tensor, %arg1: // Similar to the above, with tf.Sub as the binary op kind. func.func @testConcatCwiseBinarySynthSubOp3Inputs(%arg0: tensor, %arg1: tensor, %arg2: tensor) -> tensor { - // CHECK: %[[CONST:.*]] = "tf.Const"() {value = dense<[2.000000e+00, 3.000000e+00, 0.000000e+00]> + // CHECK: %[[CONST:.*]] = "tf.Const"() <{value = dense<[2.000000e+00, 3.000000e+00, 0.000000e+00]> // CHECK: %[[CONCAT:.*]] = "tf.ConcatV2"(%arg0, %arg1, %arg2, // CHECK: "tf.Sub"(%[[CONCAT]], %[[CONST]]) %axis = "tf.Const"() { value = dense<1> : tensor } : () -> tensor @@ -681,8 +681,8 @@ func.func @testTileMultiplesAllOnes(%arg0: tensor<2x3xf32>) -> tensor<2x3xf32> { // CHECK-LABEL: func @testStaticAndIdenticalTypeForEqualOp func.func @testStaticAndIdenticalTypeForEqualOp(%arg0: tensor<2xi32>, %arg1: tensor<2xi32>) -> tensor<2xi1> { // CHECK: "tf.Equal"(%arg0, %arg1) - // CHECK-SAME: device = "/job:localhost/replica:0/task:0/device:GPU:0" // CHECK-SAME: incompatible_shape_error = true + // CHECK-SAME: device = "/job:localhost/replica:0/task:0/device:GPU:0" %0 = "tf.Equal"(%arg0, %arg1) {incompatible_shape_error = false, device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi1> func.return %0: tensor<2xi1> } @@ -690,8 +690,8 @@ func.func @testStaticAndIdenticalTypeForEqualOp(%arg0: tensor<2xi32>, %arg1: ten // CHECK-LABEL: func @testStaticAndIdenticalTypeForNotEqualOp func.func @testStaticAndIdenticalTypeForNotEqualOp(%arg0: tensor<2xi32>, %arg1: tensor<2xi32>) -> tensor<2xi1> { // CHECK: "tf.NotEqual"(%arg0, %arg1) - // CHECK-SAME: device = "/job:localhost/replica:0/task:0/device:GPU:0" // CHECK-SAME: incompatible_shape_error = true + // CHECK-SAME: device = "/job:localhost/replica:0/task:0/device:GPU:0" %0 = "tf.NotEqual"(%arg0, %arg1) {incompatible_shape_error = false, device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi1> func.return %0: tensor<2xi1> } @@ -707,8 +707,8 @@ func.func @testUnknownBroadcastForNotEqualOp(%arg0: tensor, %arg1: tensor // CHECK-LABEL: func @testKnownGoodBroadcastForNotEqualOp func.func @testKnownGoodBroadcastForNotEqualOp(%arg0: tensor<1x?xi32>, %arg1: tensor) -> tensor { // CHECK: "tf.NotEqual"(%arg0, %arg1) - // CHECK-SAME: device = "/job:localhost/replica:0/task:0/device:GPU:0" // CHECK-SAME: incompatible_shape_error = true + // CHECK-SAME: device = "/job:localhost/replica:0/task:0/device:GPU:0" %0 = "tf.NotEqual"(%arg0, %arg1) {incompatible_shape_error = false, device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<1x?xi32>, tensor) -> tensor func.return %0: tensor } @@ -740,8 +740,8 @@ func.func @testUnrankedLHSForNotEqualOp(%arg0: tensor<*xi32>, %arg1: tensor // CHECK-LABEL: func @testScalarForNotEqualOp func.func @testScalarForNotEqualOp(%arg0: tensor, %arg1: tensor) -> tensor { // CHECK: "tf.NotEqual"(%arg0, %arg1) - // CHECK-SAME: device = "/job:localhost/replica:0/task:0/device:GPU:0" // CHECK-SAME: incompatible_shape_error = true + // CHECK-SAME: device = "/job:localhost/replica:0/task:0/device:GPU:0" %0 = "tf.NotEqual"(%arg0, %arg1) {incompatible_shape_error = false, device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor, tensor) -> tensor func.return %0: tensor } @@ -752,7 +752,7 @@ func.func @testLogicalNotOfEqual(%arg0: tensor<8x16xf32>, %arg1: tensor<8x16xf32 %1 = "tf.LogicalNot"(%0) {device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<8x16xi1>) -> tensor<8x16xi1> func.return %1: tensor<8x16xi1> - // CHECK: %[[NE:.*]] = "tf.NotEqual"(%arg0, %arg1) {device = "/job:localhost/replica:0/task:0/device:GPU:0", incompatible_shape_error = true} + // CHECK: %[[NE:.*]] = "tf.NotEqual"(%arg0, %arg1) <{incompatible_shape_error = true}> {device = "/job:localhost/replica:0/task:0/device:GPU:0"} // CHECK: return %[[NE]] } @@ -762,7 +762,7 @@ func.func @testLogicalNotOfNotEqual(%arg0: tensor<8x16xf32>, %arg1: tensor<8x16x %1 = "tf.LogicalNot"(%0) {device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<8x16xi1>) -> tensor<8x16xi1> func.return %1: tensor<8x16xi1> - // CHECK: %[[NE:.*]] = "tf.Equal"(%arg0, %arg1) {device = "/job:localhost/replica:0/task:0/device:GPU:0", incompatible_shape_error = true} + // CHECK: %[[NE:.*]] = "tf.Equal"(%arg0, %arg1) <{incompatible_shape_error = true}> {device = "/job:localhost/replica:0/task:0/device:GPU:0"} // CHECK: return %[[NE]] } @@ -811,7 +811,7 @@ func.func @testSizeFolding(%arg0: tensor<3x5x7xf32>) -> tensor { %0 = "tf.Size"(%arg0) : (tensor<3x5x7xf32>) -> tensor func.return %0: tensor -// CHECK: %[[CONST:.*]] = "tf.Const"() {value = dense<105> : tensor} : () -> tensor +// CHECK: %[[CONST:.*]] = "tf.Const"() <{value = dense<105> : tensor}> : () -> tensor // CHECK: return %[[CONST]] : tensor } @@ -873,7 +873,7 @@ func.func @testXdivyWithSqrtDivisor(%arg0: tensor<8x16xf32>, %arg1: tensor<8x16x // CHECK-LABEL: @identityTranspose func.func @identityTranspose(%arg0: tensor<2x3x4x5x6xf32>) -> tensor<2x3x4x5x6xf32> { - %0 = "tf.Const"() {value = dense<[0, 1, 2, 3, 4]> : tensor<5xi32>} : () -> tensor<5xi32> + %0 = "tf.Const"() <{value = dense<[0, 1, 2, 3, 4]> : tensor<5xi32>}> : () -> tensor<5xi32> %1 = "tf.Transpose"(%arg0, %0) : (tensor<2x3x4x5x6xf32>, tensor<5xi32>) -> tensor<2x3x4x5x6xf32> func.return %1 : tensor<2x3x4x5x6xf32> @@ -895,7 +895,7 @@ func.func @nonIdentityTranspose(%arg0: tensor<2x3x4x5x6xf32>) -> tensor<2x3x4x6x %1 = "tf.Transpose"(%arg0, %0) : (tensor<2x3x4x5x6xf32>, tensor<5xi32>) -> tensor<2x3x4x6x5xf32> func.return %1 : tensor<2x3x4x6x5xf32> - // CHECK: %[[CONST:.*]] = "tf.Const"() {value = dense<[0, 1, 2, 4, 3]> : tensor<5xi32>} : () -> tensor<5xi32> + // CHECK: %[[CONST:.*]] = "tf.Const"() <{value = dense<[0, 1, 2, 4, 3]> : tensor<5xi32>}> : () -> tensor<5xi32> // CHECK: %[[TRANS:.*]] = "tf.Transpose"(%arg0, %[[CONST]]) : (tensor<2x3x4x5x6xf32>, tensor<5xi32>) -> tensor<2x3x4x6x5xf32> // CHECK: return %[[TRANS]] } @@ -924,8 +924,8 @@ func.func @nonCancellableTransposeCrossRegion(%arg0: tensor<1x4x4x8xf32>) -> ten func.return %result : tensor<1x4x4x8xf32> - // CHECK-DAG: %[[CONST1:.*]] = "tf.Const"() {value = dense<[0, 3, 1, 2]> : tensor<4xi32>} - // CHECK-DAG: %[[CONST2:.*]] = "tf.Const"() {value = dense<[0, 2, 3, 1]> : tensor<4xi32>} + // CHECK-DAG: %[[CONST1:.*]] = "tf.Const"() <{value = dense<[0, 3, 1, 2]> : tensor<4xi32>}> + // CHECK-DAG: %[[CONST2:.*]] = "tf.Const"() <{value = dense<[0, 2, 3, 1]> : tensor<4xi32>}> // CHECK: %[[TRANS1:.*]] = "tf.Transpose"(%arg0, %[[CONST1]]) : (tensor<1x4x4x8xf32>, tensor<4xi32>) -> tensor<1x8x4x4xf32> // CHECK: %[[TRANS2:.*]] = "tf.Transpose"(%[[TRANS1]], %[[CONST2]]) : (tensor<1x8x4x4xf32>, tensor<4xi32>) -> tensor<1x4x4x8xf32> // CHECK: return %[[TRANS2]] @@ -951,8 +951,8 @@ func.func @nonCancellableTranspose(%arg0: tensor<1x4x4x8xf32>) -> tensor<4x1x4x8 func.return %3 : tensor<4x1x4x8xf32> - // CHECK-DAG: %[[CONST1:.*]] = "tf.Const"() {value = dense<[0, 3, 1, 2]> : tensor<4xi32>} - // CHECK-DAG: %[[CONST2:.*]] = "tf.Const"() {value = dense<[2, 0, 3, 1]> : tensor<4xi32>} + // CHECK-DAG: %[[CONST1:.*]] = "tf.Const"() <{value = dense<[0, 3, 1, 2]> : tensor<4xi32>}> + // CHECK-DAG: %[[CONST2:.*]] = "tf.Const"() <{value = dense<[2, 0, 3, 1]> : tensor<4xi32>}> // CHECK: %[[TRANS1:.*]] = "tf.Transpose"(%arg0, %[[CONST1]]) : (tensor<1x4x4x8xf32>, tensor<4xi32>) -> tensor<1x8x4x4xf32> // CHECK: %[[TRANS2:.*]] = "tf.Transpose"(%[[TRANS1]], %[[CONST2]]) : (tensor<1x8x4x4xf32>, tensor<4xi32>) -> tensor<4x1x4x8xf32> // CHECK: return %[[TRANS2]] @@ -969,8 +969,8 @@ func.func @addN(%arg0: tensor<*xf32>) -> tensor<*xf32> { func.func @addNWithZerosFloat(%arg0: tensor<2xf32>) -> (tensor<2xf32>, tensor<2xf32>, tensor<2xf32>, tensor<2xf32>) { %0 = "tf.Const"() {value = dense<1.000000e+00> : tensor<2xf32>} : () -> tensor<2xf32> %1 = "tf.Const"() {value = dense<0.000000e+00> : tensor<2xf32>} : () -> tensor<2xf32> - // CHECK-DAG: [[ZERO:%.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<2xf32>} - // CHECK-DAG: [[ONE:%.*]] = "tf.Const"() {value = dense<1.000000e+00> : tensor<2xf32>} + // CHECK-DAG: [[ZERO:%.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<2xf32>}> + // CHECK-DAG: [[ONE:%.*]] = "tf.Const"() <{value = dense<1.000000e+00> : tensor<2xf32>}> // CHECK: [[ADD_N:%.*]] = "tf.AddN"(%arg0, [[ZERO]], [[ONE]]) // CHECK: return %arg0, %arg0, [[ZERO]], [[ADD_N]] %2 = "tf.AddN"(%arg0, %1, %1) : (tensor<2xf32>, tensor<2xf32>, tensor<2xf32>) -> tensor<2xf32> @@ -984,8 +984,8 @@ func.func @addNWithZerosFloat(%arg0: tensor<2xf32>) -> (tensor<2xf32>, tensor<2x func.func @addNWithZerosInt(%arg0: tensor<2xi32>) -> (tensor<2xi32>, tensor<2xi32>, tensor<2xi32>, tensor<2xi32>) { %0 = "tf.Const"() {value = dense<1> : tensor<2xi32>} : () -> tensor<2xi32> %1 = "tf.Const"() {value = dense<0> : tensor<2xi32>} : () -> tensor<2xi32> - // CHECK-DAG: [[ZERO:%.*]] = "tf.Const"() {value = dense<0> : tensor<2xi32>} - // CHECK-DAG: [[ONE:%.*]] = "tf.Const"() {value = dense<1> : tensor<2xi32>} + // CHECK-DAG: [[ZERO:%.*]] = "tf.Const"() <{value = dense<0> : tensor<2xi32>}> + // CHECK-DAG: [[ONE:%.*]] = "tf.Const"() <{value = dense<1> : tensor<2xi32>}> // CHECK: [[ADD_N:%.*]] = "tf.AddN"(%arg0, [[ZERO]], [[ONE]]) // CHECK: return %arg0, %arg0, [[ZERO]], [[ADD_N]] %2 = "tf.AddN"(%arg0, %1, %1) : (tensor<2xi32>, tensor<2xi32>, tensor<2xi32>) -> tensor<2xi32> @@ -998,7 +998,7 @@ func.func @addNWithZerosInt(%arg0: tensor<2xi32>) -> (tensor<2xi32>, tensor<2xi3 // CHECK-LABEL: func @addNSkipFoldingIfBroadcasting func.func @addNSkipFoldingIfBroadcasting(%arg0: tensor<1xf32>) -> tensor<10xf32> { %0 = "tf.Const"() {value = dense<0.000000e+00> : tensor<10xf32>} : () -> tensor<10xf32> - // CHECK: [[ZERO:%.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<10xf32>} + // CHECK: [[ZERO:%.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<10xf32>}> // CHECK: [[ADD_N:%.*]] = "tf.AddN"(%arg0, [[ZERO]]) // CHECK: return [[ADD_N]] %1 = "tf.AddN"(%arg0, %0) : (tensor<1xf32>, tensor<10xf32>) -> tensor<10xf32> @@ -1014,8 +1014,8 @@ func.func @ToBool_0DScalarI1(%arg0: tensor) -> tensor { // CHECK-LABEL: func @ToBool_0DScalarInt func.func @ToBool_0DScalarInt(%arg0: tensor) -> tensor { - // CHECK: [[Zero:%.*]] = "tf.Const"() {value = dense<0> : tensor} - // CHECK: [[NE:%.*]] = "tf.NotEqual"(%arg0, [[Zero]]) {device = "/job:localhost/replica:0/task:0/device:GPU:0", incompatible_shape_error = true} + // CHECK: [[Zero:%.*]] = "tf.Const"() <{value = dense<0> : tensor}> + // CHECK: [[NE:%.*]] = "tf.NotEqual"(%arg0, [[Zero]]) <{incompatible_shape_error = true}> {device = "/job:localhost/replica:0/task:0/device:GPU:0"} // CHECK: return [[NE]] %0 = "tf.ToBool"(%arg0) {device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor) -> tensor func.return %0 : tensor @@ -1023,8 +1023,8 @@ func.func @ToBool_0DScalarInt(%arg0: tensor) -> tensor { // CHECK-LABEL: func @ToBool_0DScalarFloat func.func @ToBool_0DScalarFloat(%arg0: tensor) -> tensor { - // CHECK: [[Zero:%.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor} : () -> tensor - // CHECK: [[NE:%.*]] = "tf.NotEqual"(%arg0, [[Zero]]) {device = "/job:localhost/replica:0/task:0/device:GPU:0", incompatible_shape_error = true} + // CHECK: [[Zero:%.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor}> : () -> tensor + // CHECK: [[NE:%.*]] = "tf.NotEqual"(%arg0, [[Zero]]) <{incompatible_shape_error = true}> {device = "/job:localhost/replica:0/task:0/device:GPU:0"} // CHECK: return [[NE]] %0 = "tf.ToBool"(%arg0) {device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor) -> tensor func.return %0 : tensor @@ -1032,8 +1032,8 @@ func.func @ToBool_0DScalarFloat(%arg0: tensor) -> tensor { // CHECK-LABEL: func @ToBool_0DScalarString func.func @ToBool_0DScalarString(%arg0: tensor) -> tensor { - // CHECK: [[EmptyStr:%.*]] = "tf.Const"() {value = dense<""> : tensor} : () -> tensor - // CHECK: [[NE:%.*]] = "tf.NotEqual"(%arg0, [[EmptyStr]]) {device = "/job:localhost/replica:0/task:0/device:GPU:0", incompatible_shape_error = true} : (tensor, tensor) -> tensor + // CHECK: [[EmptyStr:%.*]] = "tf.Const"() <{value = dense<""> : tensor}> : () -> tensor + // CHECK: [[NE:%.*]] = "tf.NotEqual"(%arg0, [[EmptyStr]]) <{incompatible_shape_error = true}> {device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor, tensor) -> tensor // CHECK: return [[NE]] : tensor %0 = "tf.ToBool"(%arg0) {device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor) -> tensor func.return %0 : tensor @@ -1041,7 +1041,7 @@ func.func @ToBool_0DScalarString(%arg0: tensor) -> tensor { // CHECK-LABEL: func @ToBool_1DTensor func.func @ToBool_1DTensor(%arg0: tensor<1xf32>) -> tensor { - // CHECK: [[Const:%.*]] = "tf.Const"() {value = dense : tensor} : () -> tensor + // CHECK: [[Const:%.*]] = "tf.Const"() <{value = dense : tensor}> : () -> tensor // CHECK: return [[Const]] %0 = "tf.ToBool"(%arg0) : (tensor<1xf32>) -> tensor func.return %0 : tensor @@ -1049,7 +1049,7 @@ func.func @ToBool_1DTensor(%arg0: tensor<1xf32>) -> tensor { // CHECK-LABEL: func @ToBool_1DTensorZeroDim func.func @ToBool_1DTensorZeroDim(%arg0: tensor<0xf32>) -> tensor { - // CHECK: [[Const:%.*]] = "tf.Const"() {value = dense : tensor} : () -> tensor + // CHECK: [[Const:%.*]] = "tf.Const"() <{value = dense : tensor}> : () -> tensor // CHECK: return [[Const]] %0 = "tf.ToBool"(%arg0) : (tensor<0xf32>) -> tensor func.return %0 : tensor @@ -1057,7 +1057,7 @@ func.func @ToBool_1DTensorZeroDim(%arg0: tensor<0xf32>) -> tensor { // CHECK-LABEL: func @ToBool_2DTensor func.func @ToBool_2DTensor(%arg0: tensor<1x5xf32>) -> tensor { - // CHECK: [[Const:%.*]] = "tf.Const"() {value = dense : tensor} : () -> tensor + // CHECK: [[Const:%.*]] = "tf.Const"() <{value = dense : tensor}> : () -> tensor // CHECK: return [[Const]] %0 = "tf.ToBool"(%arg0) : (tensor<1x5xf32>) -> tensor func.return %0 : tensor @@ -1065,7 +1065,7 @@ func.func @ToBool_2DTensor(%arg0: tensor<1x5xf32>) -> tensor { // CHECK-LABEL: func @ToBool_2DTensorZeroDim func.func @ToBool_2DTensorZeroDim(%arg0: tensor<1x0xf32>) -> tensor { - // CHECK: [[Const:%.*]] = "tf.Const"() {value = dense : tensor} : () -> tensor + // CHECK: [[Const:%.*]] = "tf.Const"() <{value = dense : tensor}> : () -> tensor // CHECK: return [[Const]] %0 = "tf.ToBool"(%arg0) : (tensor<1x0xf32>) -> tensor func.return %0 : tensor @@ -1098,7 +1098,7 @@ func.func @testReadVariableOpOfCastMultiUse(%arg0: tensor, tensor) -> () func.return %1: tensor - // CHECK: %0 = "tf.Cast"(%arg0) {Truncate = false} : (tensor>>) -> tensor<*x!tf_type.resource> + // CHECK: %0 = "tf.Cast"(%arg0) <{Truncate = false}> : (tensor>>) -> tensor<*x!tf_type.resource> // CHECK: %1 = "tf.ReadVariableOp"(%0) : (tensor<*x!tf_type.resource>) -> tensor // CHECK: "tf.AssignVariableOp"(%0, %1) : (tensor<*x!tf_type.resource>, tensor) -> () // CHECK: return %1 @@ -1118,7 +1118,7 @@ func.func @testMultiReadVariableOpsOfCast(%arg0: tensor) -> tensor { - // CHECK:[[VAL0:%.+]] = "tf.Const"() {value = dense<3> : tensor} + // CHECK:[[VAL0:%.+]] = "tf.Const"() <{value = dense<3> : tensor}> %0 = "tf.Rank"(%arg0) : (tensor<4x3x2xf32>) -> tensor // CHECK: return [[VAL0]] @@ -1143,16 +1143,16 @@ func.func @testRankOfRankedTensorDynamicShapeOutput(%arg0 : tensor<4x3x2xf32>) - func.func @foldFill() -> (tensor<3x2x1xf32>, tensor<*xf32>, tensor<*xcomplex>) { %0 = "tf.Const"() {value = dense<[3, 2, 1]> : tensor<3xi32>} : () -> tensor<3xi32> %1 = "tf.Const"() {value = dense<23.0> : tensor} : () -> tensor - // CHECK-DAG: "tf.Const"() {value = dense<2.300000e+01> : tensor<3x2x1xf32>} + // CHECK-DAG: "tf.Const"() <{value = dense<2.300000e+01> : tensor<3x2x1xf32>}> %2 = "tf.Fill"(%0, %1) : (tensor<3xi32>, tensor) -> tensor<3x2x1xf32> - // CHECK-DAG: "tf.Const"() {value = dense<2.300000e+01> : tensor<3x2x1xf32>} + // CHECK-DAG: "tf.Const"() <{value = dense<2.300000e+01> : tensor<3x2x1xf32>}> %3 = "tf.Fill"(%0, %1) : (tensor<3xi32>, tensor) -> tensor<*xf32> %complex_cst = "tf.Const"() {value = dense<(0.000000e+00,1.000000e+00)> : tensor>} : () -> tensor> // Here, custom folder doesn't handle complex dtypes and it is folded through // the constant folding hook. // TODO(hinsu): Handle complex dtypes in the custom folder for FillOp. - // CHECK-DAG: "tf.Const"() {value = dense<(0.000000e+00,1.000000e+00)> : tensor<3x2x1xcomplex>} : () -> tensor<*xcomplex> + // CHECK-DAG: "tf.Const"() <{value = dense<(0.000000e+00,1.000000e+00)> : tensor<3x2x1xcomplex>}> : () -> tensor<*xcomplex> %4 = "tf.Fill"(%0, %complex_cst) : (tensor<3xi32>, tensor>) -> tensor<*xcomplex> func.return %2, %3, %4 : tensor<3x2x1xf32>, tensor<*xf32>, tensor<*xcomplex> @@ -1164,13 +1164,13 @@ func.func @foldIf(%arg0: tensor, %arg1: tensor, %arg2: tensor) -> %1 = "tf.Const"() {value = dense : tensor} : () -> tensor // CHECK: %0 = "tf.PartitionedCall"(%arg0, %arg1) - // CHECK-SAME: device = "noodle" // CHECK-SAME: f = @sub + // CHECK-SAME: device = "noodle" %2 = "tf.If"(%0, %arg0, %arg1) {then_branch = @add, else_branch = @sub, output_shapes = [#tf_type.shape<>], device = "noodle", is_stateless = true} : (tensor, tensor, tensor) -> tensor // CHECK: %1 = "tf.StatefulPartitionedCall"(%0, %arg1) + // CHECK-SAME: f = @add // CHECK-SAME: _underscore_attr = "something" // CHECK-SAME: device = "noodle" - // CHECK-SAME: f = @add %3 = "tf.If"(%1, %2, %arg1) {then_branch = @add, else_branch = @sub, output_shapes = [#tf_type.shape<>], device = "noodle", _underscore_attr = "something", is_stateless = false} : (tensor, tensor, tensor) -> tensor // CHECK: %2 = "tf.If" @@ -1233,13 +1233,13 @@ func.func @foldIfRegionMismatchedTypes(%arg0: tensor, %arg1: tensor, %arg1: tensor, %arg2: tensor) -> (tensor) { // CHECK: %[[PRED:.*]] = "tf._SomeOp"() : () -> tensor %pred = "tf._SomeOp"() : () -> tensor - // CHECK: %[[IF_OUTPUT:.*]] = "tf.IfRegion"(%[[PRED]]) ({ + // CHECK: %[[IF_OUTPUT:.*]] = "tf.IfRegion"(%[[PRED]]) <{is_stateless = true}> ({ // CHECK: %[[MUL:.*]] = "tf.Mul"(%[[ARG0]], %[[ARG1]]) // CHECK: "tf.Yield"(%[[MUL]]) : (tensor) // CHECK: }, { // CHECK: %[[SUB:.*]] = "tf.Sub"(%[[ARG0]], %[[ARG1]]) // CHECK: "tf.Yield"(%[[SUB]]) : (tensor) - // CHECK: }) {device = "/job:localhost/replica:0/task:0/device:GPU:0", is_stateless = true} : (tensor) -> tensor + // CHECK: }) {device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor) -> tensor %0:4 = "tf.IfRegion"(%pred) ({ %true_value = "tf.Mul"(%arg0, %arg1) : (tensor, tensor) -> tensor "tf.Yield"(%arg1, %arg2, %true_value, %arg2) : (tensor, tensor, tensor, tensor) -> () @@ -1260,7 +1260,7 @@ func.func @eliminatePassThroughIfRegion(%arg0: tensor, %arg1: tensor, func.func @eliminatePassThroughCaseRegion(%arg0: tensor, %arg1: tensor, %arg2: tensor) -> (tensor) { // CHECK: %[[INDEX:.*]] = "tf._SomeOp"() : () -> tensor %index = "tf._SomeOp"() : () -> tensor - // CHECK: %[[CASE_OUTPUT:.*]] = "tf.CaseRegion"(%[[INDEX]]) ({ + // CHECK: %[[CASE_OUTPUT:.*]] = "tf.CaseRegion"(%[[INDEX]]) <{is_stateless = true}> ({ // CHECK: %[[MUL:.*]] = "tf.Mul"(%[[ARG0]], %[[ARG1]]) // CHECK: "tf.Yield"(%[[MUL]]) : (tensor) // CHECK: }, { @@ -1269,7 +1269,7 @@ func.func @eliminatePassThroughCaseRegion(%arg0: tensor, %arg1: tensor // CHECK: }, { // CHECK: %[[ADD:.*]] = "tf.AddV2"(%[[ARG0]], %[[ARG1]]) // CHECK: "tf.Yield"(%[[ADD]]) : (tensor) - // CHECK: }) {device = "/job:localhost/replica:0/task:0/device:GPU:0", is_stateless = true} : (tensor) -> tensor + // CHECK: }) {device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor) -> tensor %0:3 = "tf.CaseRegion"(%index) ({ %mul = "tf.Mul"(%arg0, %arg1) : (tensor, tensor) -> tensor "tf.Yield"(%arg1, %mul, %arg2) : (tensor, tensor, tensor) -> () @@ -1293,13 +1293,13 @@ func.func @foldCase(%arg0: tensor, %arg1: tensor) -> (tensor) { %3 = arith.constant dense<0> : tensor // CHECK: PartitionedCall - // CHECK-SAME: device = "noodle" // CHECK-SAME: f = @add + // CHECK-SAME: device = "noodle" %4 = "tf.Case"(%2, %arg0, %arg1) {branches = [@sub, @add], output_shapes = [#tf_type.shape<>], device = "noodle", is_stateless = false} : (tensor, tensor, tensor) -> tensor // CHECK: PartitionedCall + // CHECK-SAME: f = @sub // CHECK-SAME: _cluster_launch = "not_ready" // CHECK-SAME: device = "noodle" - // CHECK-SAME: f = @sub %5 = "tf.Case"(%3, %4, %arg1) {branches = [@sub, @add], output_shapes = [#tf_type.shape<>], device= "noodle", _cluster_launch = "not_ready", is_stateless = false} : (tensor, tensor, tensor) -> tensor func.return %5 : tensor } @@ -1317,7 +1317,7 @@ func.func @sub(%arg0: tensor<*xf32>, %arg1: tensor<*xf32>) -> tensor<*xf32> { // CHECK-LABEL: testBatchToSpaceToBatchToSpaceND // CHECK-SAME: ([[INPUT:%.*]]: tensor, [[CROPS:%.*]]: tensor) func.func @testBatchToSpaceToBatchToSpaceND(%arg0: tensor, %arg1: tensor) -> tensor<*xf32> { - // CHECK: [[BLOCK_SHAPE:%.*]] = "tf.Const"() {value = dense<8> : tensor<2xi64>} + // CHECK: [[BLOCK_SHAPE:%.*]] = "tf.Const"() <{value = dense<8> : tensor<2xi64>}> // CHECK: [[BATCH_TO_SHAPE_ND:%.*]] = "tf.BatchToSpaceND"([[INPUT]], [[BLOCK_SHAPE]], [[CROPS]]) {device = "/job:localhost/replica:0/task:0/device:GPU:0"} %0 = "tf.BatchToSpace"(%arg0, %arg1) {block_size = 8 : i64, device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor, tensor) -> tensor<*xf32> // CHECK: return [[BATCH_TO_SHAPE_ND]] @@ -1612,8 +1612,8 @@ func.func private @testIfElse(tensor<*xf32>) -> tensor<*xf32> func.func @testIfDropOutputShapes(tensor, tensor<2xf32>) -> tensor<2xf32> { ^bb0(%arg0: tensor, %arg1: tensor<2xf32>): // CHECK: "tf.If" - // CHECK-SAME: device = "/job:localhost/replica:0/task:0/device:GPU:0" // CHECK-NOT: output_shapes + // CHECK-SAME: device = "/job:localhost/replica:0/task:0/device:GPU:0" %1 = "tf.If"(%arg0, %arg1) { then_branch = @testIfThen, else_branch = @testIfElse, is_stateless = false, output_shapes = [#tf_type.shape<>], device = "/job:localhost/replica:0/task:0/device:GPU:0" } : (tensor, tensor<2xf32>) -> tensor<2xf32> @@ -1647,10 +1647,10 @@ func.func @testSumFoldBypass(%arg0: tensor<4x?xf16>, %arg1: tensor<*xi64>) -> te // CHECK-LABEL: @testMatrixDiag func.func @testMatrixDiag(%diag: tensor<2x4xf32>) -> tensor<2x4x4xf32> { - // CHECK-DAG: %[[MINUS1:.*]] = "tf.Const"() {value = dense<-1> : tensor} : () -> tensor - // CHECK-DAG: %[[ZEROI:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor - // CHECK-DAG: %[[ZEROF:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor} : () -> tensor - // CHECK-DAG: "tf.MatrixDiagV3"(%arg0, %[[ZEROI]], %[[MINUS1]], %[[MINUS1]], %[[ZEROF]]) {align = "RIGHT_LEFT", device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<2x4xf32>, tensor, tensor, tensor, tensor) -> tensor<2x4x4xf32> + // CHECK-DAG: %[[MINUS1:.*]] = "tf.Const"() <{value = dense<-1> : tensor}> : () -> tensor + // CHECK-DAG: %[[ZEROI:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor + // CHECK-DAG: %[[ZEROF:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor}> : () -> tensor + // CHECK-DAG: "tf.MatrixDiagV3"(%arg0, %[[ZEROI]], %[[MINUS1]], %[[MINUS1]], %[[ZEROF]]) <{align = "RIGHT_LEFT"}> {device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<2x4xf32>, tensor, tensor, tensor, tensor) -> tensor<2x4x4xf32> %0 = "tf.MatrixDiag"(%diag) {device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<2x4xf32>) -> tensor<2x4x4xf32> func.return %0 : tensor<2x4x4xf32> } @@ -1660,9 +1660,9 @@ func.func @testMatrixSetDiag(%arg0: tensor<3x3xi64>, %arg1: tensor<3xi64>) -> te %0 = "tf.MatrixSetDiag"(%arg0, %arg1) {device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<3x3xi64>, tensor<3xi64>) -> tensor<3x3xi64> func.return %0 : tensor<3x3xi64> - // CHECK: %[[ZERO:.*]] = "tf.Const"() {value = dense<0> : tensor} + // CHECK: %[[ZERO:.*]] = "tf.Const"() <{value = dense<0> : tensor}> // CHECK: %[[RES:.*]] = "tf.MatrixSetDiagV3"(%arg0, %arg1, %[[ZERO]]) - // CHECK-SAME: {align = "RIGHT_LEFT", device = "/job:localhost/replica:0/task:0/device:GPU:0"} + // CHECK-SAME: <{align = "RIGHT_LEFT"}> {device = "/job:localhost/replica:0/task:0/device:GPU:0"} // CHECK-SAME: (tensor<3x3xi64>, tensor<3xi64>, tensor) -> tensor<3x3xi64> } @@ -1672,7 +1672,7 @@ func.func @testMatrixSetDiagV2(%arg0: tensor<3x3xi64>, %arg1: tensor<3xi64>, %ar func.return %0 : tensor<3x3xi64> // CHECK: %[[RES:.*]] = "tf.MatrixSetDiagV3"(%arg0, %arg1, %arg2) - // CHECK-SAME: {align = "LEFT_LEFT", device = "/job:localhost/replica:0/task:0/device:GPU:0"} + // CHECK-SAME: <{align = "LEFT_LEFT"}> {device = "/job:localhost/replica:0/task:0/device:GPU:0"} } // CHECK-LABEL: @testVariableToVariableV2 @@ -1680,7 +1680,7 @@ func.func @testVariableToVariableV2() { // CHECK-NOT: "tf.Variable" %0 = "tf.Const"() { value = dense<1> : tensor } : () -> tensor - // CHECK: "tf.VariableV2"() {container = "", device = "/job:localhost/replica:0/task:0/device:GPU:0", shape = #tf_type.shape<>, shared_name = "var"} + // CHECK: "tf.VariableV2"() <{container = "", shape = #tf_type.shape<>, shared_name = "var"}> {device = "/job:localhost/replica:0/task:0/device:GPU:0"} %1 = "tf.Variable"() {container = "", dtype = i32, shared_name = "var", shape = #tf_type.shape<>, device = "/job:localhost/replica:0/task:0/device:GPU:0"} : () -> tensor %2 = "tf.Assign"(%1, %0) : (tensor, tensor) -> (tensor) @@ -1691,7 +1691,7 @@ func.func @testVariableToVariableV2() { func.func @testUnpackAndCwiseUnary(%arg0: tensor) -> (tensor, tensor) { // CHECK: %[[NEG:.*]] = "tf.Neg"(%arg0) {device = ""} - // CHECK: %[[UNPACK:.*]]:2 = "tf.Unpack"(%[[NEG]]) {axis = 1 : i64, device = ""} + // CHECK: %[[UNPACK:.*]]:2 = "tf.Unpack"(%[[NEG]]) <{axis = 1 : i64}> {device = ""} %unpacked:2 = "tf.Unpack"(%arg0) {axis = 1 : i64, device = ""} : (tensor) -> (tensor, tensor) %0 = "tf.Neg"(%unpacked#0): (tensor) -> tensor @@ -1708,7 +1708,7 @@ func.func @testFoldStridedSliceShapeI32(%arg0: tensor) -> (tensor<2 %2 = "tf.Shape"(%arg0) : (tensor) -> tensor<4xi32> %3 = "tf.StridedSlice"(%2, %1, %0, %1) {begin_mask = 0 : i64, ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 0 : i64} : (tensor<4xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<2xi32> func.return %3 : tensor<2xi32> - // CHECK: %[[CST:.*]] = "tf.Const"() {value = dense<[1, 2]> : tensor<2xi32>} : () -> tensor<2xi32> + // CHECK: %[[CST:.*]] = "tf.Const"() <{value = dense<[1, 2]> : tensor<2xi32>}> : () -> tensor<2xi32> // CHECK: return %[[CST]] } @@ -1719,7 +1719,7 @@ func.func @testFoldStridedSliceShapeI64(%arg0: tensor) -> (tensor<2 %2 = "tf.Shape"(%arg0) : (tensor) -> tensor<4xi64> %3 = "tf.StridedSlice"(%2, %1, %0, %1) {begin_mask = 0 : i64, ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 0 : i64} : (tensor<4xi64>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<2xi64> func.return %3 : tensor<2xi64> - // CHECK: %[[CST:.*]] = "tf.Const"() {value = dense<[1, 2]> : tensor<2xi64>} : () -> tensor<2xi64> + // CHECK: %[[CST:.*]] = "tf.Const"() <{value = dense<[1, 2]> : tensor<2xi64>}> : () -> tensor<2xi64> // CHECK: return %[[CST]] } @@ -1730,7 +1730,7 @@ func.func @testFoldStridedSliceShapeDynamicOutput(%arg0: tensor) -> %2 = "tf.Shape"(%arg0) : (tensor) -> tensor<4xi32> %3 = "tf.StridedSlice"(%2, %1, %0, %1) {begin_mask = 0 : i64, ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 0 : i64} : (tensor<4xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor func.return %3 : tensor - // CHECK: %[[CST:.*]] = "tf.Const"() {value = dense<[1, 2]> : tensor<2xi32>} : () -> tensor + // CHECK: %[[CST:.*]] = "tf.Const"() <{value = dense<[1, 2]> : tensor<2xi32>}> : () -> tensor // CHECK: return %[[CST]] } @@ -1741,7 +1741,7 @@ func.func @testFoldStridedSliceShapeWithShrinkAxisMaskI32(%arg0: tensor) -> tensor<4xi32> %3 = "tf.StridedSlice"(%2, %1, %0, %1) {begin_mask = 0 : i64, ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 1 : i64} : (tensor<4xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor func.return %3 : tensor - // CHECK: %[[CST:.*]] = "tf.Const"() {value = dense<1> : tensor} : () -> tensor + // CHECK: %[[CST:.*]] = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor // CHECK: return %[[CST]] } @@ -1752,7 +1752,7 @@ func.func @testFoldStridedSliceShapeWithShrinkAxisMaskI64(%arg0: tensor) -> tensor<4xi64> %3 = "tf.StridedSlice"(%2, %1, %0, %1) {begin_mask = 0 : i64, ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 1 : i64} : (tensor<4xi64>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor func.return %3 : tensor - // CHECK: %[[CST:.*]] = "tf.Const"() {value = dense<1> : tensor} : () -> tensor + // CHECK: %[[CST:.*]] = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor // CHECK: return %[[CST]] } @@ -1763,7 +1763,7 @@ func.func @testFoldStridedSliceShapeWithShrinkAxisMaskUnrankedOutput(%arg0: tens %2 = "tf.Shape"(%arg0) : (tensor) -> tensor<4xi32> %3 = "tf.StridedSlice"(%2, %1, %0, %1) {begin_mask = 0 : i64, ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 1 : i64} : (tensor<4xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<*xi32> func.return %3 : tensor<*xi32> - // CHECK: %[[CST:.*]] = "tf.Const"() {value = dense<1> : tensor} : () -> tensor<*xi32> + // CHECK: %[[CST:.*]] = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor<*xi32> // CHECK: return %[[CST]] } @@ -1775,7 +1775,7 @@ func.func @testFoldStridedSliceShapeWithShrinkAxisMaskNegativeBegin1(%arg0: tens %3 = "tf.Shape"(%arg0) : (tensor) -> tensor<4xi32> %4 = "tf.StridedSlice"(%3, %0, %1, %2) {begin_mask = 0 : i64, ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 1 : i64} : (tensor<4xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor func.return %4 : tensor - // CHECK: %[[CST:.*]] = "tf.Const"() {value = dense<3> : tensor} : () -> tensor + // CHECK: %[[CST:.*]] = "tf.Const"() <{value = dense<3> : tensor}> : () -> tensor // CHECK: return %[[CST]] } @@ -1787,7 +1787,7 @@ func.func @testFoldStridedSliceShapeWithShrinkAxisMaskNegativeBegin2(%arg0: tens %3 = "tf.Shape"(%arg0) : (tensor) -> tensor<4xi32> %4 = "tf.StridedSlice"(%3, %0, %1, %2) {begin_mask = 0 : i64, ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 1 : i64} : (tensor<4xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor func.return %4 : tensor - // CHECK: %[[CST:.*]] = "tf.Const"() {value = dense<2> : tensor} : () -> tensor + // CHECK: %[[CST:.*]] = "tf.Const"() <{value = dense<2> : tensor}> : () -> tensor // CHECK: return %[[CST]] } @@ -1811,7 +1811,7 @@ func.func @testFoldStridedSliceShapeWithBeginMask(%arg0: tensor<1x2x3x?xf32>) -> %3 = "tf.Shape"(%arg0) : (tensor<1x2x3x?xf32>) -> tensor<4xi32> %4 = "tf.StridedSlice"(%3, %0, %1, %2) {begin_mask = 1 : i64, ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 0 : i64} : (tensor<4xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<2xi32> func.return %4 : tensor<2xi32> - // CHECK: %[[CST:.*]] = "tf.Const"() {value = dense<[1, 2]> : tensor<2xi32>} : () -> tensor<2xi32> + // CHECK: %[[CST:.*]] = "tf.Const"() <{value = dense<[1, 2]> : tensor<2xi32>}> : () -> tensor<2xi32> // CHECK: return %[[CST]] } @@ -1822,7 +1822,7 @@ func.func @testFoldStridedSliceShapeWithEndMask(%arg0: tensor) -> ( %2 = "tf.Shape"(%arg0) : (tensor) -> tensor<4xi32> %3 = "tf.StridedSlice"(%2, %1, %0, %1) {begin_mask = 0 : i64, ellipsis_mask = 0 : i64, end_mask = 1 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 0 : i64} : (tensor<4xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<3xi32> func.return %3 : tensor<3xi32> - // CHECK: %[[CST:.*]] = "tf.Const"() {value = dense<[1, 2, 3]> : tensor<3xi32>} : () -> tensor<3xi32> + // CHECK: %[[CST:.*]] = "tf.Const"() <{value = dense<[1, 2, 3]> : tensor<3xi32>}> : () -> tensor<3xi32> // CHECK: return %[[CST]] } @@ -1834,7 +1834,7 @@ func.func @testFoldStridedSliceShapeWithPositiveStrides(%arg0: tensor<1x2x3x4x?x %3 = "tf.Shape"(%arg0) : (tensor<1x2x3x4x?xf32>) -> tensor<5xi32> %4 = "tf.StridedSlice"(%3, %0, %1, %2) {begin_mask = 0 : i64, ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 0 : i64} : (tensor<5xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<2xi32> func.return %4 : tensor<2xi32> - // CHECK: %[[CST:.*]] = "tf.Const"() {value = dense<[2, 4]> : tensor<2xi32>} : () -> tensor<2xi32> + // CHECK: %[[CST:.*]] = "tf.Const"() <{value = dense<[2, 4]> : tensor<2xi32>}> : () -> tensor<2xi32> // CHECK: return %[[CST]] } @@ -1845,7 +1845,7 @@ func.func @testFoldStridedSliceShapeWithPositiveStridesOutOfBoundEnd(%arg0: tens %2 = "tf.Shape"(%arg0) : (tensor) -> tensor<4xi32> %3 = "tf.StridedSlice"(%2, %1, %0, %1) {begin_mask = 0 : i64, ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 0 : i64} : (tensor<4xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<3xi32> func.return %3 : tensor<3xi32> - // CHECK: %[[CST:.*]] = "tf.Const"() {value = dense<[1, 2, 3]> : tensor<3xi32>} : () -> tensor<3xi32> + // CHECK: %[[CST:.*]] = "tf.Const"() <{value = dense<[1, 2, 3]> : tensor<3xi32>}> : () -> tensor<3xi32> // CHECK: return %[[CST]] } @@ -1857,7 +1857,7 @@ func.func @testFoldStridedSliceShapeWithNegativeStrides(%arg0: tensor<1x2x3x?xf3 %3 = "tf.Shape"(%arg0) : (tensor<1x2x3x?xf32>) -> tensor<4xi32> %4 = "tf.StridedSlice"(%3, %0, %1, %2) {begin_mask = 0 : i64, ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 0 : i64} : (tensor<4xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32> func.return %4 : tensor<1xi32> - // CHECK: %[[CST:.*]] = "tf.Const"() {value = dense<3> : tensor<1xi32>} : () -> tensor<1xi32> + // CHECK: %[[CST:.*]] = "tf.Const"() <{value = dense<3> : tensor<1xi32>}> : () -> tensor<1xi32> // CHECK: return %[[CST]] } @@ -1869,7 +1869,7 @@ func.func @testFoldStridedSliceShapeWithNegativeStridesOutOfBoundBegin(%arg0: te %3 = "tf.Shape"(%arg0) : (tensor) -> tensor<4xi32> %4 = "tf.StridedSlice"(%3, %0, %1, %2) {begin_mask = 0 : i64, ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 0 : i64} : (tensor<4xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<2xi32> func.return %4 : tensor<2xi32> - // CHECK: %[[CST:.*]] = "tf.Const"() {value = dense<[3, 2]> : tensor<2xi32>} : () -> tensor<2xi32> + // CHECK: %[[CST:.*]] = "tf.Const"() <{value = dense<[3, 2]> : tensor<2xi32>}> : () -> tensor<2xi32> // CHECK: return %[[CST]] } @@ -1881,7 +1881,7 @@ func.func @testFoldStridedSliceShapeWithNegativeStridesBeginMask(%arg0: tensor) -> tensor<4xi32> %4 = "tf.StridedSlice"(%3, %0, %1, %2) {begin_mask = 1 : i64, ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 0 : i64} : (tensor<4xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<2xi32> func.return %4 : tensor<2xi32> - // CHECK: %[[CST:.*]] = "tf.Const"() {value = dense<[3, 2]> : tensor<2xi32>} : () -> tensor<2xi32> + // CHECK: %[[CST:.*]] = "tf.Const"() <{value = dense<[3, 2]> : tensor<2xi32>}> : () -> tensor<2xi32> // CHECK: return %[[CST]] } @@ -1893,7 +1893,7 @@ func.func @testFoldStridedSliceShapeWithNegativeStridesEndMask(%arg0: tensor<1x2 %3 = "tf.Shape"(%arg0) : (tensor<1x2x3x?xf32>) -> tensor<4xi32> %4 = "tf.StridedSlice"(%3, %0, %1, %2) {begin_mask = 0 : i64, ellipsis_mask = 0 : i64, end_mask = 1 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 0 : i64} : (tensor<4xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<3xi32> func.return %4 : tensor<3xi32> - // CHECK: %[[CST:.*]] = "tf.Const"() {value = dense<[3, 2, 1]> : tensor<3xi32>} : () -> tensor<3xi32> + // CHECK: %[[CST:.*]] = "tf.Const"() <{value = dense<[3, 2, 1]> : tensor<3xi32>}> : () -> tensor<3xi32> // CHECK: return %[[CST]] } @@ -1905,7 +1905,7 @@ func.func @testFoldStridedSliceShapeWithEmptySlice(%arg0: tensor) - %3 = "tf.Shape"(%arg0) : (tensor) -> tensor<4xi32> %4 = "tf.StridedSlice"(%3, %0, %1, %2) {begin_mask = 0 : i64, ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 0 : i64} : (tensor<4xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<0xi32> func.return %4 : tensor<0xi32> - // CHECK: %[[CST:.*]] = "tf.Const"() {value = dense<> : tensor<0xi32>} : () -> tensor<0xi32> + // CHECK: %[[CST:.*]] = "tf.Const"() <{value = dense<> : tensor<0xi32>}> : () -> tensor<0xi32> // CHECK: return %[[CST]] } @@ -1914,7 +1914,7 @@ func.func @testFoldEnsureShapeOp(%arg0: tensor<10x20xf32>) -> (tensor<10x20xf32> %0 = "tf.EnsureShape"(%arg0) {shape = #tf_type.shape<10x20>} : (tensor<10x20xf32>) -> tensor<10x20xf32> %1 = "tf.EnsureShape"(%arg0) {shape = #tf_type.shape} : (tensor<10x20xf32>) -> tensor<10x20xf32> // Failing case which should not be folded. - // CHECK: %[[NF:.*]] = "tf.EnsureShape"(%arg0) {shape = #tf_type.shape<20x10>} + // CHECK: %[[NF:.*]] = "tf.EnsureShape"(%arg0) <{shape = #tf_type.shape<20x10>}> %2 = "tf.EnsureShape"(%arg0) {shape = #tf_type.shape<20x10>} : (tensor<10x20xf32>) -> tensor<20x10xf32> // CHECK: return %arg0, %arg0, %[[NF]] func.return %0, %1, %2: tensor<10x20xf32>, tensor<10x20xf32>, tensor<20x10xf32> @@ -1924,7 +1924,7 @@ func.func @testFoldEnsureShapeOp(%arg0: tensor<10x20xf32>) -> (tensor<10x20xf32> func.func @testConvertPackToReshapeAxis0(%arg0: tensor<2x3xf32>) -> tensor<1x2x3xf32> { %0 = "tf.Pack"(%arg0) {axis = 0 : i64, _xla_outside_compilation = "1", device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<2x3xf32>) -> tensor<1x2x3xf32> func.return %0 : tensor<1x2x3xf32> - // CHECK: %[[SHAPE:.*]] = "tf.Const"() {value = dense<[1, 2, 3]> : tensor<3xi64>} : () -> tensor<3xi64> + // CHECK: %[[SHAPE:.*]] = "tf.Const"() <{value = dense<[1, 2, 3]> : tensor<3xi64>}> : () -> tensor<3xi64> // CHECK: %[[RESHAPE:.*]] = "tf.Reshape"(%arg0, %[[SHAPE]]) {_xla_outside_compilation = "1", device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<2x3xf32>, tensor<3xi64>) -> tensor<1x2x3xf32> // CHECK: return %[[RESHAPE]] : tensor<1x2x3xf32> } @@ -1933,7 +1933,7 @@ func.func @testConvertPackToReshapeAxis0(%arg0: tensor<2x3xf32>) -> tensor<1x2x3 func.func @testConvertPackToReshapeAxis1(%arg0: tensor<2x3xf32>) -> tensor<2x1x3xf32> { %0 = "tf.Pack"(%arg0) {axis = 1 : i64, device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<2x3xf32>) -> tensor<2x1x3xf32> func.return %0 : tensor<2x1x3xf32> - // CHECK: %[[SHAPE:.*]] = "tf.Const"() {value = dense<[2, 1, 3]> : tensor<3xi64>} : () -> tensor<3xi64> + // CHECK: %[[SHAPE:.*]] = "tf.Const"() <{value = dense<[2, 1, 3]> : tensor<3xi64>}> : () -> tensor<3xi64> // CHECK: %[[RESHAPE:.*]] = "tf.Reshape"(%arg0, %[[SHAPE]]) {device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<2x3xf32>, tensor<3xi64>) -> tensor<2x1x3xf32> // CHECK: return %[[RESHAPE]] : tensor<2x1x3xf32> } @@ -1942,7 +1942,7 @@ func.func @testConvertPackToReshapeAxis1(%arg0: tensor<2x3xf32>) -> tensor<2x1x3 func.func @testDontConvertPackToReshapeDynamicShape(%arg0: tensor<2x?xf32>) -> tensor<1x2x?xf32> { %0 = "tf.Pack"(%arg0) {axis = 0 : i64} : (tensor<2x?xf32>) -> tensor<1x2x?xf32> func.return %0 : tensor<1x2x?xf32> - // CHECK: %[[PACK:.*]] = "tf.Pack"(%arg0) {axis = 0 : i64} : (tensor<2x?xf32>) -> tensor<1x2x?xf32> + // CHECK: %[[PACK:.*]] = "tf.Pack"(%arg0) <{axis = 0 : i64}> : (tensor<2x?xf32>) -> tensor<1x2x?xf32> // CHECK: return %[[PACK]] : tensor<1x2x?xf32> } @@ -1950,7 +1950,7 @@ func.func @testDontConvertPackToReshapeDynamicShape(%arg0: tensor<2x?xf32>) -> t func.func @while_with_id_passthrough(%arg0: tensor<7xf32> {tf._user_specified_name = "x"}) -> tensor attributes {tf.entry_function = {control_outputs = "", inputs = "x", outputs = "identity_RetVal"}} { %0 = "tf.Const"() {value = dense<0> : tensor} : () -> tensor %1 = "tf.Const"() {value = dense<3> : tensor} : () -> tensor - // CHECK: %[[SHAPE:.*]] = "tf.Const"() {value = dense<7> : tensor<1xi32>} + // CHECK: %[[SHAPE:.*]] = "tf.Const"() <{value = dense<7> : tensor<1xi32>}> %2 = "tf.Const"() {value = dense<7> : tensor<1xi32>} : () -> tensor<1xi32> %3 = "tf.Const"() {value = dense : tensor} : () -> tensor %4 = "tf.Const"() {value = dense<1> : tensor} : () -> tensor @@ -1983,7 +1983,7 @@ func.func @while_with_id_passthrough(%arg0: tensor<7xf32> {tf._user_specified_na func.func @testConvertQuantizeAndDequantizeV2ToQuantizeAndDequantizeV4(%arg0 : tensor, %arg1 : tensor, %arg2 : tensor) -> tensor { %0 = "tf.QuantizeAndDequantizeV2"(%arg0, %arg1, %arg2) {device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor, tensor, tensor) -> tensor func.return %0 : tensor - // CHECK: %[[QUANT:.*]] = "tf.QuantizeAndDequantizeV4"(%arg0, %arg1, %arg2) {axis = -1 : i64, device = "/job:localhost/replica:0/task:0/device:GPU:0", narrow_range = false, num_bits = 8 : i64, range_given = false, round_mode = "HALF_TO_EVEN", signed_input = true} : (tensor, tensor, tensor) -> tensor + // CHECK: %[[QUANT:.*]] = "tf.QuantizeAndDequantizeV4"(%arg0, %arg1, %arg2) <{axis = -1 : i64, narrow_range = false, num_bits = 8 : i64, range_given = false, round_mode = "HALF_TO_EVEN", signed_input = true}> {device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor, tensor, tensor) -> tensor // CHECK: return %[[QUANT]] : tensor } @@ -1991,10 +1991,10 @@ func.func @testConvertQuantizeAndDequantizeV2ToQuantizeAndDequantizeV4(%arg0 : t func.func @testHashTableAndInitializeTableToV2(%arg0: tensor) { // CHECK: [[handle:%.*]] = "tf.HashTableV2"() // CHECK-SAME: container = "" - // CHECK-SAME: device = "" // CHECK-SAME: key_dtype = !tf_type.string // CHECK-SAME: shared_name = "table" // CHECK-SAME: value_dtype = i32 + // CHECK-SAME: device = "" // CHECK-SAME: () -> tensor %handle = "tf.HashTable"() {container = "", device = "", shared_name = "table", key_dtype = !tf_type.string, value_dtype = i32} : () -> tensor<*x!tf_type.stringref> @@ -2008,10 +2008,10 @@ func.func @testHashTableAndInitializeTableToV2(%arg0: tensor) { func.func @testHashTableAndLookupTableSizeToV2() -> tensor { // CHECK: [[handle:%.*]] = "tf.HashTableV2"() // CHECK-SAME: container = "" - // CHECK-SAME: device = "" // CHECK-SAME: key_dtype = !tf_type.string // CHECK-SAME: shared_name = "table" // CHECK-SAME: value_dtype = i32 + // CHECK-SAME: device = "" // CHECK-SAME: () -> tensor %handle = "tf.HashTable"() {container = "", device = "", shared_name = "table", key_dtype = !tf_type.string, value_dtype = i32} : () -> tensor<*x!tf_type.stringref> @@ -2025,10 +2025,10 @@ func.func @testHashTableAndLookupTableSizeToV2() -> tensor { func.func @testHashTableAndLookupTableFindToV2(%arg0: tensor, %arg1: tensor) -> tensor { // CHECK: [[handle:%.*]] = "tf.HashTableV2"() // CHECK-SAME: container = "" - // CHECK-SAME: device = "" // CHECK-SAME: key_dtype = !tf_type.string // CHECK-SAME: shared_name = "table" // CHECK-SAME: value_dtype = i32 + // CHECK-SAME: device = "" // CHECK-SAME: () -> tensor %handle = "tf.HashTable"() {container = "", device = "", shared_name = "table", key_dtype = !tf_type.string, value_dtype = i32} : () -> tensor<*x!tf_type.stringref> @@ -2041,9 +2041,9 @@ func.func @testHashTableAndLookupTableFindToV2(%arg0: tensor, % // CHECK-LABEL: testDivNoNanAndMulNoNanWithConstantY // CHECK-SAME: (%[[ARG0:.*]]: tensor<2xf32>) func.func @testDivNoNanAndMulNoNanWithConstantY(%arg0: tensor<2xf32>) -> (tensor<2xf32>, tensor<2xf32>, tensor<2xf32>) { - // CHECK: %[[CON1:.*]] = "tf.Const"() {value = dense<[1.000000e+00, 2.000000e+00]> : tensor<2xf32>} : () -> tensor<2xf32> - // CHECK-NEXT: %[[CON2:.*]] = "tf.Const"() {value = dense<[1.000000e+01, 0.000000e+00]> : tensor<2xf32>} : () -> tensor<2xf32> - // CHECK-NEXT: %[[CON3:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<2xf32>} : () -> tensor<2xf32> + // CHECK: %[[CON1:.*]] = "tf.Const"() <{value = dense<[1.000000e+00, 2.000000e+00]> : tensor<2xf32>}> : () -> tensor<2xf32> + // CHECK-NEXT: %[[CON2:.*]] = "tf.Const"() <{value = dense<[1.000000e+01, 0.000000e+00]> : tensor<2xf32>}> : () -> tensor<2xf32> + // CHECK-NEXT: %[[CON3:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<2xf32>}> : () -> tensor<2xf32> // CHECK-NEXT: %[[RES1:.*]] = "tf.Div"(%[[ARG0]], %[[CON1]]) {device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<2xf32>, tensor<2xf32>) -> tensor<2xf32> // CHECK-NEXT: %[[RES2:.*]] = "tf.MulNoNan"(%[[ARG0]], %[[CON2]]) : (tensor<2xf32>, tensor<2xf32>) -> tensor<2xf32> // CHECK-NEXT: return %[[RES1]], %[[RES2]], %[[CON3]] : tensor<2xf32>, tensor<2xf32>, tensor<2xf32> @@ -2060,9 +2060,9 @@ func.func @testDivNoNanAndMulNoNanWithConstantY(%arg0: tensor<2xf32>) -> (tensor // CHECK-LABEL: testComplexDivNoNanAndMulNoNanWithConstantY // CHECK-SAME: (%[[ARG0:.*]]: tensor<2xcomplex>) func.func @testComplexDivNoNanAndMulNoNanWithConstantY(%arg0: tensor<2xcomplex>) -> (tensor<2xcomplex>, tensor<2xcomplex>, tensor<2xcomplex>) { - // CHECK-NEXT: %[[COMP2:.*]] = "tf.Const"() {value = dense<[(0.000000e+00,0.000000e+00), (2.000000e+00,0.000000e+00)]> : tensor<2xcomplex>} : () -> tensor<2xcomplex> - // CHECK-NEXT: %[[COMP1:.*]] = "tf.Const"() {value = dense<[(1.000000e+00,3.000000e+00), (2.000000e+00,4.000000e+00)]> : tensor<2xcomplex>} : () -> tensor<2xcomplex> - // CHECK-NEXT: %[[COMP3:.*]] = "tf.Const"() {value = dense<(0.000000e+00,0.000000e+00)> : tensor<2xcomplex>} : () -> tensor<2xcomplex> + // CHECK-NEXT: %[[COMP2:.*]] = "tf.Const"() <{value = dense<[(0.000000e+00,0.000000e+00), (2.000000e+00,0.000000e+00)]> : tensor<2xcomplex>}> : () -> tensor<2xcomplex> + // CHECK-NEXT: %[[COMP1:.*]] = "tf.Const"() <{value = dense<[(1.000000e+00,3.000000e+00), (2.000000e+00,4.000000e+00)]> : tensor<2xcomplex>}> : () -> tensor<2xcomplex> + // CHECK-NEXT: %[[COMP3:.*]] = "tf.Const"() <{value = dense<(0.000000e+00,0.000000e+00)> : tensor<2xcomplex>}> : () -> tensor<2xcomplex> // CHECK-NEXT: %[[RES1:.*]] = "tf.Mul"(%[[ARG0]], %[[COMP1]]) {device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<2xcomplex>, tensor<2xcomplex>) -> tensor<2xcomplex> // CHECK-NEXT: %[[RES2:.*]] = "tf.DivNoNan"(%[[ARG0]], %[[COMP2]]) : (tensor<2xcomplex>, tensor<2xcomplex>) -> tensor<2xcomplex> // CHECK-NEXT: return %[[RES1]], %[[RES2]], %[[COMP3]] : tensor<2xcomplex>, tensor<2xcomplex>, tensor<2xcomplex> @@ -2104,7 +2104,7 @@ func.func @testDivNoNanAndMulNoNanWithNonConstantY(%arg0: tensor<2xf32>, %arg1: // CHECK-LABEL: testComplexDivNoNanOpWithNonConstantY // CHECK-SAME: (%[[ARG0:.*]]: tensor<2xcomplex>, %[[ARG1:.*]]: tensor<2xcomplex>, %[[ARG2:.*]]: tensor<2xf32>) func.func @testComplexDivNoNanOpWithNonConstantY(%arg0: tensor<2xcomplex>, %arg1: tensor<2xcomplex>, %arg2: tensor<2xf32>) -> (tensor<2xcomplex>, tensor<2xcomplex>, tensor<2xcomplex>) { - // CHECK: %[[CON1:.*]] = "tf.Const"() {value = dense<[1.000000e+00, 2.000000e+00]> : tensor<2xf32>} : () -> tensor<2xf32> + // CHECK: %[[CON1:.*]] = "tf.Const"() <{value = dense<[1.000000e+00, 2.000000e+00]> : tensor<2xf32>}> : () -> tensor<2xf32> // CHECK-NEXT: %[[NONCON2:.*]] = "tf.Sub"(%[[ARG0]], %[[ARG1]]) : (tensor<2xcomplex>, tensor<2xcomplex>) -> tensor<2xcomplex> // CHECK-NEXT: %[[NONCON3:.*]] = "tf.Complex"(%[[CON1]], %[[ARG2]]) : (tensor<2xf32>, tensor<2xf32>) -> tensor<2xcomplex> // CHECK-NEXT: %[[RES1:.*]] = "tf.MulNoNan"(%[[ARG0]], %[[ARG1]]) : (tensor<2xcomplex>, tensor<2xcomplex>) -> tensor<2xcomplex @@ -2128,7 +2128,7 @@ func.func @testXlaConvToV2(%lhs: tensor<8x4x16x16x16xf32>, %rhs: tensor<4x3x3x16 %rhs_dilation = "tf.Const"() {value = dense<1> : tensor<3xi32>} : () -> tensor<3xi32> %padding = "tf.Const"() {value = dense<0> : tensor<3x2xi32>} : () -> tensor<3x2xi32> %strides = "tf.Const"() {value = dense<[3, 1, 1]> : tensor<3xi32>} : () -> tensor<3xi32> - // CHECK: "tf.XlaConvV2"(%arg0, %arg1, %cst_3, %cst_2, %cst_0, %cst_1, %cst) {batch_group_count = 1 : i64, device = "/job:localhost/replica:0/task:0/device:GPU:0", dimension_numbers = "\18\03 \042\03\00\01\02@\04P\04Z\03\01\02\03b\03\01\02\03", precision_config = ""} : (tensor<8x4x16x16x16xf32>, tensor<4x3x3x16x16xf32>, tensor<3xi32>, tensor<3x2xi32>, tensor<3xi32>, tensor<3xi32>, tensor) -> tensor<8x4x14x14x16xf32> + // CHECK: "tf.XlaConvV2"(%arg0, %arg1, %cst_3, %cst_2, %cst_0, %cst_1, %cst) <{batch_group_count = 1 : i64, dimension_numbers = "\18\03 \042\03\00\01\02@\04P\04Z\03\01\02\03b\03\01\02\03", precision_config = ""}> {device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<8x4x16x16x16xf32>, tensor<4x3x3x16x16xf32>, tensor<3xi32>, tensor<3x2xi32>, tensor<3xi32>, tensor<3xi32>, tensor) -> tensor<8x4x14x14x16xf32> %0 = "tf.XlaConv"(%lhs, %rhs, %strides, %padding, %lhs_dilation, %rhs_dilation, %feature_group_count) {dimension_numbers = "\18\03 \042\03\00\01\02@\04P\04Z\03\01\02\03b\03\01\02\03", precision_config = "", device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<8x4x16x16x16xf32>, tensor<4x3x3x16x16xf32>, tensor<3xi32>, tensor<3x2xi32>, tensor<3xi32>, tensor<3xi32>, tensor) -> tensor<8x4x14x14x16xf32> func.return %0 : tensor<8x4x14x14x16xf32> } @@ -2136,7 +2136,7 @@ func.func @testXlaConvToV2(%lhs: tensor<8x4x16x16x16xf32>, %rhs: tensor<4x3x3x16 // CHECK-LABEL: testXlaReduceToXlaVariadicReduceV2 func.func @testXlaReduceToXlaVariadicReduceV2(%arg0: tensor<*xbf16>, %arg1: tensor<*xbf16>) -> tensor<*xbf16> { - // CHECK: "tf.XlaVariadicReduceV2"(%arg0, %arg1) {device = "/job:localhost/replica:0/task:0/device:GPU:0", dimensions_to_reduce = [], operandSegmentSizes = array, reducer = @sum1} : (tensor<*xbf16>, tensor<*xbf16>) -> tensor<*xbf16> + // CHECK: "tf.XlaVariadicReduceV2"(%arg0, %arg1) <{dimensions_to_reduce = [], operandSegmentSizes = array, reducer = @sum1}> {device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<*xbf16>, tensor<*xbf16>) -> tensor<*xbf16> %0 = "tf.XlaReduce"(%arg0, %arg1) {dimensions_to_reduce = [], reducer = @sum1, device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<*xbf16>, tensor<*xbf16>) -> tensor<*xbf16> func.return %0 : tensor<*xbf16> } @@ -2148,7 +2148,7 @@ func.func private @sum1(%arg0: tensor<*xbf16>, %arg1: tensor<*xbf16>) -> tensor< // CHECK-LABEL: testXlaVariadicReduceToV2 func.func @testXlaVariadicReduceToV2(%arg0: tensor<3x4xf32>, %arg1: tensor) -> tensor { - // CHECK: "tf.XlaVariadicReduceV2"(%arg0, %arg1) {device = "/job:localhost/replica:0/task:0/device:GPU:0", dimensions_to_reduce = [], operandSegmentSizes = array, reducer = @sum2} : (tensor<3x4xf32>, tensor) -> tensor + // CHECK: "tf.XlaVariadicReduceV2"(%arg0, %arg1) <{dimensions_to_reduce = [], operandSegmentSizes = array, reducer = @sum2}> {device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<3x4xf32>, tensor) -> tensor %0 = "tf.XlaVariadicReduce"(%arg0, %arg1) {dimensions_to_reduce = [], reducer = @sum2, device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<3x4xf32>, tensor) -> tensor func.return %0 : tensor } @@ -2229,8 +2229,8 @@ func.func @testTensorListGetItem(%arg0: tensor<1600x1x32xf32>, %arg1: tensor<2xi %1 = "tf.TensorListGetItem"(%0, %arg2, %arg1) {device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor>>, tensor, tensor<2xi32>) -> tensor<1x32xf32> func.return %1 : tensor<1x32xf32> - // CHECK: %[[CST:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor - // CHECK: %[[RES:.*]] = "tf.GatherV2"(%arg0, %arg2, %cst) {batch_dims = 0 : i64, device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<1600x1x32xf32>, tensor, tensor) -> tensor<1x32xf32> + // CHECK: %[[CST:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor + // CHECK: %[[RES:.*]] = "tf.GatherV2"(%arg0, %arg2, %cst) <{batch_dims = 0 : i64}> {device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<1600x1x32xf32>, tensor, tensor) -> tensor<1x32xf32> } // CHECK-LABEL: testTensorListGetItemMultipleUsers @@ -2240,9 +2240,9 @@ func.func @testTensorListGetItemMultipleUsers(%arg0: tensor<1600x1x32xf32>, %arg %2 = "tf.TensorListGetItem"(%0, %arg3, %arg1) {device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor>>, tensor, tensor<2xi32>) -> tensor<1x32xf32> func.return %1, %2 : tensor<1x32xf32>, tensor<1x32xf32> - // CHECK: %[[CST:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor - // CHECK: %[[RES0:.*]] = "tf.GatherV2"(%arg0, %arg2, %cst) {batch_dims = 0 : i64, device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<1600x1x32xf32>, tensor, tensor) -> tensor<1x32xf32> - // CHECK: %[[RES1:.*]] = "tf.GatherV2"(%arg0, %arg3, %cst) {batch_dims = 0 : i64, device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<1600x1x32xf32>, tensor, tensor) -> tensor<1x32xf32> + // CHECK: %[[CST:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor + // CHECK: %[[RES0:.*]] = "tf.GatherV2"(%arg0, %arg2, %cst) <{batch_dims = 0 : i64}> {device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<1600x1x32xf32>, tensor, tensor) -> tensor<1x32xf32> + // CHECK: %[[RES1:.*]] = "tf.GatherV2"(%arg0, %arg3, %cst) <{batch_dims = 0 : i64}> {device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<1600x1x32xf32>, tensor, tensor) -> tensor<1x32xf32> } // CHECK-LABEL: testUnaryIdempotent diff --git a/tensorflow/compiler/mlir/tensorflow/tests/canonicalize_compile_and_replicate_attributes.mlir b/tensorflow/compiler/mlir/tensorflow/tests/canonicalize_compile_and_replicate_attributes.mlir index 9f6dc7b017fb90..224395853c0901 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/canonicalize_compile_and_replicate_attributes.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/canonicalize_compile_and_replicate_attributes.mlir @@ -3,7 +3,7 @@ // CHECK-LABEL: func.func @convert_tpu_replicate func.func @convert_tpu_replicate() { tf_executor.graph { - // CHECK: tf_executor.island wraps "tf.TPUReplicateMetadata"() {_replication_info = "cluster", _xla_compile_device_type = "TPU", allow_soft_placement = false, computation_shape = [], device = "", device_assignment = [], host_compute_core = [], name = "TPUReplicateMetadata", num_cores_per_replica = 1 : i64, num_replicas = 1 : i64, step_marker_location = "STEP_MARK_AT_ENTRY", topology = "", use_spmd_for_xla_partitioning = false, use_tpu = true} : () -> () + // CHECK: tf_executor.island wraps "tf.TPUReplicateMetadata"() <{allow_soft_placement = false, computation_shape = [], device_assignment = [], host_compute_core = [], num_cores_per_replica = 1 : i64, num_replicas = 1 : i64, step_marker_location = "STEP_MARK_AT_ENTRY", topology = "", use_spmd_for_xla_partitioning = false, use_tpu = true}> {_replication_info = "cluster", _xla_compile_device_type = "TPU", device = "", name = "TPUReplicateMetadata"} : () -> () %control = tf_executor.island wraps "tf.TPUReplicateMetadata"() {_tpu_replicate = "cluster", allow_soft_placement = false, computation_shape = [], device = "", device_assignment = [], host_compute_core = [], name = "TPUReplicateMetadata", num_cores_per_replica = 1 : i64, num_replicas = 1 : i64, step_marker_location = "STEP_MARK_AT_ENTRY", topology = "", use_tpu = true, use_spmd_for_xla_partitioning = false} : () -> () %outputs_0, %control_0 = tf_executor.island wraps "tf.Placeholder"() {device = "", dtype = "tfdtype$DT_FLOAT", name = "y", shape = "tfshape$dim { }"} : () -> tensor<0xf32> %outputs_1, %control_1 = tf_executor.island wraps "tf.TPUReplicatedInput"(%outputs_0) {N = 1 : i64, T = "tfdtype$DT_FLOAT", device = "", name = "input1"} : (tensor<0xf32>) -> tensor<0xf32> @@ -21,7 +21,7 @@ func.func @convert_tpu_replicate() { // CHECK-LABEL: func.func @convert_xla_must_compile func.func @convert_xla_must_compile(%arg0: tensor) -> tensor { - // CHECK: "tf.StatefulPartitionedCall"(%arg0) {_xla_compile_device_type = "CPU", config = "", config_proto = "", device = "/device:CPU:0", executor_type = "", f = @stateful_pcall_func} : (tensor) -> tensor + // CHECK: "tf.StatefulPartitionedCall"(%arg0) <{config = "", config_proto = "", executor_type = "", f = @stateful_pcall_func}> {_xla_compile_device_type = "CPU", device = "/device:CPU:0"} : (tensor) -> tensor %0 = "tf.StatefulPartitionedCall"(%arg0) {_XlaMustCompile = true, config = "", config_proto = "", device = "/device:CPU:0", executor_type = "", f = @stateful_pcall_func} : (tensor) -> (tensor) func.return %0 : tensor } diff --git a/tensorflow/compiler/mlir/tensorflow/tests/cluster_tf_ops_pass.mlir b/tensorflow/compiler/mlir/tensorflow/tests/cluster_tf_ops_pass.mlir index 8d94a0778b839d..85f9c47058ecf5 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/cluster_tf_ops_pass.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/cluster_tf_ops_pass.mlir @@ -43,7 +43,7 @@ func.func @while_cond(%arg0: tensor {tf.device = "/job:localhost/replica:0/ // CHECK: func @while_body(%[[ARG_0:.*]]: tensor {tf.device = "/job:localhost/replica:0/task:0/device:CPU:0"}) // CHECK-NEXT: %[[RESULT_0:.*]] = "tf.Const"() // CHECK-NEXT: %[[RESULT_1:.*]] = "tf.AddV2"(%[[ARG_0]], %[[RESULT_0]]) -// CHECK-NEXT: %[[RESULT_2:.*]] = "tf.Const"() {value = dense<16> : tensor} : () -> tensor +// CHECK-NEXT: %[[RESULT_2:.*]] = "tf.Const"() <{value = dense<16> : tensor}> : () -> tensor // CHECK-NEXT: tf_device.send %[[RESULT_2]] "key-0" "/job:worker/replica:0/task:1/device:CPU:0" // CHECK-SAME: device = "/job:localhost/replica:0/task:0/device:CPU:0" // CHECK-NEXT: tf_device.remote_run "/job:worker/replica:0/task:1" @[[BODY_PARTITION_0:.*]]() : () -> () diff --git a/tensorflow/compiler/mlir/tensorflow/tests/constant-fold.mlir b/tensorflow/compiler/mlir/tensorflow/tests/constant-fold.mlir index 3231f05133c48f..be0ab858484acd 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/constant-fold.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/constant-fold.mlir @@ -9,7 +9,7 @@ func.func @testShape(tensor, tensor<1x32x32x16xf32>, tensor<*xf32>) -> (ten // Result shape need not be static. Folding harness uses TensorFlow constant // in that case. - // CHECK-DAG: "tf.Const"() {value = dense<[1, 32, 32, 16]> : tensor<4xi32>} : () -> tensor + // CHECK-DAG: "tf.Const"() <{value = dense<[1, 32, 32, 16]> : tensor<4xi32>}> : () -> tensor %1 = "tf.Shape"(%arg1) {T = "tfdtype$DT_FLOAT", output = "tfdtype$DT_INT32"} : (tensor<1x32x32x16xf32>) -> tensor // CHECK: "tf.Shape"(%arg2) {T = "tfdtype$DT_FLOAT", output = "tfdtype$DT_INT32"} : (tensor<*xf32>) -> tensor @@ -28,7 +28,7 @@ func.func @testPow(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> (tensor<4xf32> // CHECK-DAG: %[[RES_NO_FOLD:.*]] = "tf.Pow"(%arg0, %arg1) %0 = "tf.Pow"(%arg0, %arg1) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32> - // CHECK-DAG: %[[POW_ZERO:.*]] = "tf.Const"() {value = dense<1.000000e+00> : tensor<4xf32>} : () -> tensor<4xf32> + // CHECK-DAG: %[[POW_ZERO:.*]] = "tf.Const"() <{value = dense<1.000000e+00> : tensor<4xf32>}> : () -> tensor<4xf32> %1 = "tf.Pow"(%arg0, %cst_zero) : (tensor<4xf32>, tensor) -> tensor<4xf32> // CHECK-NOT: "tf.Pow" @@ -42,7 +42,7 @@ func.func @testPow(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> (tensor<4xf32> func.func @testEmpty32() -> (tensor<5xi32>) { %0 = "tf.Const"() { value = dense<5> : tensor } : () -> tensor - // CHECK: [[VAL:%.+]] = "tf.Const"() {value = dense<0> : tensor<5xi32>} + // CHECK: [[VAL:%.+]] = "tf.Const"() <{value = dense<0> : tensor<5xi32>}> // CHECK: return [[VAL]] %1 = "tf.Empty"(%0) : (tensor) -> (tensor<5xi32>) func.return %1 : tensor<5xi32> @@ -52,7 +52,7 @@ func.func @testEmpty32() -> (tensor<5xi32>) { func.func @testEmpty64() -> (tensor<5xi64>) { %0 = "tf.Const"() { value = dense<5> : tensor } : () -> tensor - // CHECK: [[VAL:%.+]] = "tf.Const"() {value = dense<0> : tensor<5xi64>} + // CHECK: [[VAL:%.+]] = "tf.Const"() <{value = dense<0> : tensor<5xi64>}> // CHECK: return [[VAL]] : tensor<5xi64> %1 = "tf.Empty"(%0) : (tensor) -> (tensor<5xi64>) func.return %1 : tensor<5xi64> @@ -62,7 +62,7 @@ func.func @testEmpty64() -> (tensor<5xi64>) { func.func @testEmptyFloat() -> (tensor<5xf64>) { %0 = "tf.Const"() { value = dense<5> : tensor } : () -> tensor - // CHECK: [[VAL:%.+]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<5xf64>} + // CHECK: [[VAL:%.+]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<5xf64>}> // CHECK: return [[VAL]] %1 = "tf.Empty"(%0) : (tensor) -> (tensor<5xf64>) func.return %1 : tensor<5xf64> @@ -72,7 +72,7 @@ func.func @testEmptyFloat() -> (tensor<5xf64>) { func.func @testEmptyf16() -> (tensor<5xf16>) { %0 = "tf.Const"() { value = dense<5> : tensor } : () -> tensor - // CHECK: [[VAL:%.+]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<5xf16>} + // CHECK: [[VAL:%.+]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<5xf16>}> // CHECK: return [[VAL]] %1 = "tf.Empty"(%0) : (tensor) -> (tensor<5xf16>) func.return %1 : tensor<5xf16> @@ -82,7 +82,7 @@ func.func @testEmptyf16() -> (tensor<5xf16>) { func.func @testEmptybf16() -> (tensor<5xbf16>) { %0 = "tf.Const"() { value = dense<5> : tensor } : () -> tensor - // CHECK: [[VAL:%.+]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<5xbf16>} + // CHECK: [[VAL:%.+]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<5xbf16>}> // CHECK: return [[VAL]] %1 = "tf.Empty"(%0) : (tensor) -> (tensor<5xbf16>) func.return %1 : tensor<5xbf16> @@ -91,8 +91,8 @@ func.func @testEmptybf16() -> (tensor<5xbf16>) { // CHECK-LABEL: func @testShapeN func.func @testShapeN(%arg0: tensor, %arg1: tensor<1x32x32x16xf32>) -> (tensor<0xi64>, tensor<4xi64>) { - // CHECK-DAG: %[[SHAPE0:.*]] = "tf.Const"() {value = dense<> : tensor<0xi64>} - // CHECK-DAG: %[[SHAPE1:.*]] = "tf.Const"() {value = dense<[1, 32, 32, 16]> : tensor<4xi64>} + // CHECK-DAG: %[[SHAPE0:.*]] = "tf.Const"() <{value = dense<> : tensor<0xi64>}> + // CHECK-DAG: %[[SHAPE1:.*]] = "tf.Const"() <{value = dense<[1, 32, 32, 16]> : tensor<4xi64>}> %0:2 = "tf.ShapeN"(%arg0, %arg1) : (tensor, tensor<1x32x32x16xf32>) -> (tensor<0xi64>, tensor<4xi64>) // CHECK: return %[[SHAPE0]], %[[SHAPE1]] @@ -101,8 +101,8 @@ func.func @testShapeN(%arg0: tensor, %arg1: tensor<1x32x32x16xf32>) -> (ten // CHECK-LABEL: func @testShapeNPartialStatic func.func @testShapeNPartialStatic(%arg0: tensor, %arg1: tensor<2x?x3xf32>, %arg2: tensor<1x32x32x16xf32>, %arg3: tensor<*xf32>) -> (tensor<0xi64>, tensor<3xi64>, tensor<4xi64>, tensor) { - // CHECK-DAG: %[[SHAPE0:.*]] = "tf.Const"() {value = dense<> : tensor<0xi64>} - // CHECK-DAG: %[[SHAPE2:.*]] = "tf.Const"() {value = dense<[1, 32, 32, 16]> : tensor<4xi64>} + // CHECK-DAG: %[[SHAPE0:.*]] = "tf.Const"() <{value = dense<> : tensor<0xi64>}> + // CHECK-DAG: %[[SHAPE2:.*]] = "tf.Const"() <{value = dense<[1, 32, 32, 16]> : tensor<4xi64>}> // CHECK: %[[SHAPE13:.*]]:2 = "tf.ShapeN"(%arg1, %arg3) : (tensor<2x?x3xf32>, tensor<*xf32>) -> (tensor<3xi64>, tensor) %0:4 = "tf.ShapeN"(%arg0, %arg1, %arg2, %arg3) : (tensor, tensor<2x?x3xf32>, tensor<1x32x32x16xf32>, tensor<*xf32>) -> (tensor<0xi64>, tensor<3xi64>, tensor<4xi64>, tensor) @@ -112,8 +112,8 @@ func.func @testShapeNPartialStatic(%arg0: tensor, %arg1: tensor<2x?x3xf32>, // CHECK-LABEL: func @testShapeNOneDynamic func.func @testShapeNOneDynamic(%arg0: tensor, %arg1: tensor<1x32x32x16xf32>, %arg2: tensor<*xf32>) -> (tensor<0xi64>, tensor<4xi64>, tensor) { - // CHECK-DAG: %[[SHAPE0:.*]] = "tf.Const"() {value = dense<> : tensor<0xi64>} - // CHECK-DAG: %[[SHAPE1:.*]] = "tf.Const"() {value = dense<[1, 32, 32, 16]> : tensor<4xi64>} + // CHECK-DAG: %[[SHAPE0:.*]] = "tf.Const"() <{value = dense<> : tensor<0xi64>}> + // CHECK-DAG: %[[SHAPE1:.*]] = "tf.Const"() <{value = dense<[1, 32, 32, 16]> : tensor<4xi64>}> // CHECK: %[[SHAPE2:.*]] = "tf.Shape"(%arg2) : (tensor<*xf32>) -> tensor %0:3 = "tf.ShapeN"(%arg0, %arg1, %arg2) : (tensor, tensor<1x32x32x16xf32>, tensor<*xf32>) -> (tensor<0xi64>, tensor<4xi64>, tensor) @@ -140,8 +140,8 @@ func.func @testLeakyRelu(%arg0 : tensor<16xf32>) -> (tensor<16xf32>, tensor %2 = "tf.LeakyRelu"(%arg0) {alpha = 3.0 : f32} : (tensor<16xf32>) -> tensor<16xf32> // CHECK-DAG: [[POS:%.*]] = "tf.Const{{.*}} dense<5.000000e+00> : tensor // CHECK-DAG: [[NEG:%.*]] = "tf.Const{{.*}} dense<-1.000000e+00> : tensor - // CHECK: [[NC1:%.*]] = "tf.LeakyRelu"(%arg0) {alpha = 2.000000e-01 : f32} : (tensor<16xf32>) -> tensor<16xf32> - // CHECK: [[NC2:%.*]] = "tf.LeakyRelu"(%arg0) {alpha = 3.000000e+00 : f32} : (tensor<16xf32>) -> tensor<16xf32> + // CHECK: [[NC1:%.*]] = "tf.LeakyRelu"(%arg0) <{alpha = 2.000000e-01 : f32}> : (tensor<16xf32>) -> tensor<16xf32> + // CHECK: [[NC2:%.*]] = "tf.LeakyRelu"(%arg0) <{alpha = 3.000000e+00 : f32}> : (tensor<16xf32>) -> tensor<16xf32> // CHECK: return [[NC1]], [[POS]], [[NEG]], [[NC2]] func.return %no, %0, %1, %2 : tensor<16xf32>, tensor, tensor, tensor<16xf32> } @@ -295,8 +295,8 @@ func.func @testUnimplementedOp() -> (tensor, tensor) { %3 = "tf.Minimum"(%0, %1) {random_attr = "hello"} : (tensor, tensor) -> tensor func.return %2, %3: tensor, tensor -// CHECK-DAG: %[[CST:.*]] = "tf.Const"() {value = dense<2> : tensor} : () -> tensor -// CHECK-DAG: %[[CST1:.*]] = "tf.Const"() {value = dense<1> : tensor} : () -> tensor +// CHECK-DAG: %[[CST:.*]] = "tf.Const"() <{value = dense<2> : tensor}> : () -> tensor +// CHECK-DAG: %[[CST1:.*]] = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor // CHECK-NEXT: return %[[CST]], %[[CST1]] } @@ -454,7 +454,7 @@ func.func @DontRemoveTrivialMul(%arg0: tensor<1x6x8x1xf32>) -> tensor<1x6x8x1xf3 %1 = "tf.Mul"(%arg0, %0) : (tensor<1x6x8x1xf32>, tensor) -> tensor<1x6x8x1xf32> func.return %1 : tensor<1x6x8x1xf32> // CHECK-LABEL: DontRemoveTrivialMul - // CHECK: %[[CONST:.*]] = "tf.Const"() {value = dense<2.000000e+00> : tensor} : () -> tensor + // CHECK: %[[CONST:.*]] = "tf.Const"() <{value = dense<2.000000e+00> : tensor}> : () -> tensor // CHECK: %[[RESULT:.*]] = "tf.Mul"(%arg0, %[[CONST]]) : (tensor<1x6x8x1xf32>, tensor) -> tensor<1x6x8x1xf32> // CHECK: return %[[RESULT]] : tensor<1x6x8x1xf32> } @@ -517,7 +517,7 @@ func.func @testBroadcastGradientArgsSameShape() -> (tensor<0xi32>, tensor<0xi32> %s1 = "tf.Const"() {value = dense<[1, 2]> : tensor<2xi32>} : () -> tensor<2xi32> %r0, %r1 = "tf.BroadcastGradientArgs"(%s0, %s1) {} : (tensor<2xi32>, tensor<2xi32>) -> (tensor<0xi32>, tensor<0xi32>) - // CHECK-DAG: %[[R:.*]] = "tf.Const"() {value = dense<> : tensor<0xi32>} : () -> tensor<0xi32> + // CHECK-DAG: %[[R:.*]] = "tf.Const"() <{value = dense<> : tensor<0xi32>}> : () -> tensor<0xi32> // CHECK-NOT: tf.BroadcastGradientArgs // CHECK: return %[[R]], %[[R]] @@ -529,8 +529,8 @@ func.func @testBroadcastGradientArgs1() -> (tensor<1xi32>, tensor<0xi32>) { %s0 = "tf.Const"() {value = dense<[4]> : tensor<1xi32>} : () -> tensor<1xi32> %s1 = "tf.Const"() {value = dense<[2, 4]> : tensor<2xi32>} : () -> tensor<2xi32> %r0, %r1 = "tf.BroadcastGradientArgs"(%s0, %s1) {} : (tensor<1xi32>, tensor<2xi32>) -> (tensor<1xi32>, tensor<0xi32>) - // CHECK-DAG: %[[R0:.*]] = "tf.Const"() {value = dense<0> : tensor<1xi32>} : () -> tensor<1xi32> - // CHECK-DAG: %[[R1:.*]] = "tf.Const"() {value = dense<> : tensor<0xi32>} : () -> tensor<0xi32> + // CHECK-DAG: %[[R0:.*]] = "tf.Const"() <{value = dense<0> : tensor<1xi32>}> : () -> tensor<1xi32> + // CHECK-DAG: %[[R1:.*]] = "tf.Const"() <{value = dense<> : tensor<0xi32>}> : () -> tensor<0xi32> // CHECK-NOT: tf.BroadcastGradientArgs // CHECK: return %[[R0]], %[[R1]] @@ -542,8 +542,8 @@ func.func @testBroadcastGradientArgs2() -> (tensor<1xi32>, tensor<3xi32>) { %s2 = "tf.Const"() {value = dense<[501, 1, 32, 1280]> : tensor<4xi32>} : () -> tensor<4xi32> %s3 = "tf.Const"() {value = dense<[ 1, 1, 1, 1280]> : tensor<4xi32>} : () -> tensor<4xi32> %r2, %r3 = "tf.BroadcastGradientArgs"(%s2, %s3) {} : (tensor<4xi32>, tensor<4xi32>) -> (tensor<1xi32>, tensor<3xi32>) - // CHECK-DAG: %[[R2:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32> - // CHECK-DAG: %[[R3:.*]] = "tf.Const"() {value = dense<[0, 1, 2]> : tensor<3xi32>} : () -> tensor<3xi32> + // CHECK-DAG: %[[R2:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}> : () -> tensor<1xi32> + // CHECK-DAG: %[[R3:.*]] = "tf.Const"() <{value = dense<[0, 1, 2]> : tensor<3xi32>}> : () -> tensor<3xi32> // CHECK-NOT: tf.BroadcastGradientArgs // CHECK: return %[[R2]], %[[R3]] @@ -555,7 +555,7 @@ func.func @testBroadcastGradientArgs3() -> (tensor<3xi32>, tensor<3xi32>) { %s4 = "tf.Const"() {value = dense<[]> : tensor<0xi32>} : () -> tensor<0xi32> %s5 = "tf.Const"() {value = dense<[1, 1, 1]> : tensor<3xi32>} : () -> tensor<3xi32> %r4, %r5 = "tf.BroadcastGradientArgs"(%s4, %s5) {} : (tensor<0xi32>, tensor<3xi32>) -> (tensor<3xi32>, tensor<3xi32>) - // CHECK: %[[R0:.*]] = "tf.Const"() {value = dense<[0, 1, 2]> : tensor<3xi32>} : () -> tensor<3xi32> + // CHECK: %[[R0:.*]] = "tf.Const"() <{value = dense<[0, 1, 2]> : tensor<3xi32>}> : () -> tensor<3xi32> // CHECK-NOT: tf.BroadcastGradientArgs // CHECK: return %[[R0]], %[[R0]] @@ -567,8 +567,8 @@ func.func @testBroadcastGradientArgs4() -> (tensor<2xi32>, tensor<3xi32>) { %s4 = "tf.Const"() {value = dense<[1, 2, 1]> : tensor<3xi32>} : () -> tensor<3xi32> %s5 = "tf.Const"() {value = dense<[]> : tensor<0xi32>} : () -> tensor<0xi32> %r4, %r5 = "tf.BroadcastGradientArgs"(%s4, %s5) {} : (tensor<3xi32>, tensor<0xi32>) -> (tensor<2xi32>, tensor<3xi32>) - // CHECK-DAG: %[[R0:.*]] = "tf.Const"() {value = dense<[0, 2]> : tensor<2xi32>} : () -> tensor<2xi32> - // CHECK-DAG: %[[R1:.*]] = "tf.Const"() {value = dense<[0, 1, 2]> : tensor<3xi32>} : () -> tensor<3xi32> + // CHECK-DAG: %[[R0:.*]] = "tf.Const"() <{value = dense<[0, 2]> : tensor<2xi32>}> : () -> tensor<2xi32> + // CHECK-DAG: %[[R1:.*]] = "tf.Const"() <{value = dense<[0, 1, 2]> : tensor<3xi32>}> : () -> tensor<3xi32> // CHECK-NOT: tf.BroadcastGradientArgs // CHECK: return %[[R0]], %[[R1]] @@ -580,7 +580,7 @@ func.func @testBroadcastGradientArgs5() -> (tensor<1xi32>, tensor<1xi32>) { %s4 = "tf.Const"() {value = dense<[]> : tensor<0xi32>} : () -> tensor<0xi32> %s5 = "tf.Const"() {value = dense<[1]> : tensor<1xi32>} : () -> tensor<1xi32> %r4, %r5 = "tf.BroadcastGradientArgs"(%s4, %s5) {} : (tensor<0xi32>, tensor<1xi32>) -> (tensor<1xi32>, tensor<1xi32>) - // CHECK: %[[R0:.*]] = "tf.Const"() {value = dense<0> : tensor<1xi32>} : () -> tensor<1xi32> + // CHECK: %[[R0:.*]] = "tf.Const"() <{value = dense<0> : tensor<1xi32>}> : () -> tensor<1xi32> // CHECK-NOT: tf.BroadcastGradientArgs // CHECK: return %[[R0]], %[[R0]] @@ -592,8 +592,8 @@ func.func @testBroadcastGradientArgs6() -> (tensor<1xi32>, tensor<0xi32>) { %s4 = "tf.Const"() {value = dense<[]> : tensor<0xi32>} : () -> tensor<0xi32> %s5 = "tf.Const"() {value = dense<[2]> : tensor<1xi32>} : () -> tensor<1xi32> %r4, %r5 = "tf.BroadcastGradientArgs"(%s4, %s5) {} : (tensor<0xi32>, tensor<1xi32>) -> (tensor<1xi32>, tensor<0xi32>) - // CHECK-DAG: %[[R0:.*]] = "tf.Const"() {value = dense<0> : tensor<1xi32>} : () -> tensor<1xi32> - // CHECK-DAG: %[[R1:.*]] = "tf.Const"() {value = dense<> : tensor<0xi32>} : () -> tensor<0xi32> + // CHECK-DAG: %[[R0:.*]] = "tf.Const"() <{value = dense<0> : tensor<1xi32>}> : () -> tensor<1xi32> + // CHECK-DAG: %[[R1:.*]] = "tf.Const"() <{value = dense<> : tensor<0xi32>}> : () -> tensor<0xi32> // CHECK-NOT: tf.BroadcastGradientArgs // CHECK: return %[[R0]], %[[R1]] @@ -605,8 +605,8 @@ func.func @testBroadcastGradientArgsHigherRank() -> (tensor<2xi32>, tensor<2xi32 %s0 = "tf.Const"() {value = dense<[1, 4, 1]> : tensor<3xi32>} : () -> tensor<3xi32> %s1 = "tf.Const"() {value = dense<[1, 4]> : tensor<2xi32>} : () -> tensor<2xi32> %r0, %r1 = "tf.BroadcastGradientArgs"(%s0, %s1) {} : (tensor<3xi32>, tensor<2xi32>) -> (tensor<2xi32>, tensor<2xi32>) - // CHECK-DAG: %[[R0:.*]] = "tf.Const"() {value = dense<[0, 2]> : tensor<2xi32>} : () -> tensor<2xi32> - // CHECK-DAG: %[[R1:.*]] = "tf.Const"() {value = dense<[0, 1]> : tensor<2xi32>} : () -> tensor<2xi32> + // CHECK-DAG: %[[R0:.*]] = "tf.Const"() <{value = dense<[0, 2]> : tensor<2xi32>}> : () -> tensor<2xi32> + // CHECK-DAG: %[[R1:.*]] = "tf.Const"() <{value = dense<[0, 1]> : tensor<2xi32>}> : () -> tensor<2xi32> // CHECK-NOT: tf.BroadcastGradientArgs // CHECK: return %[[R0]], %[[R1]] @@ -618,8 +618,8 @@ func.func @testBroadcastGradientArgsScalar() -> (tensor<2xi32>, tensor<0xi32>) { %s0 = "tf.Const"() {value = dense<> : tensor<0xi32>} : () -> tensor<0xi32> %s1 = "tf.Const"() {value = dense<[2, 4]> : tensor<2xi32>} : () -> tensor<2xi32> %r0, %r1 = "tf.BroadcastGradientArgs"(%s0, %s1) {} : (tensor<0xi32>, tensor<2xi32>) -> (tensor<2xi32>, tensor<0xi32>) - // CHECK-DAG: %[[R0:.*]] = "tf.Const"() {value = dense<[0, 1]> : tensor<2xi32>} : () -> tensor<2xi32> - // CHECK-DAG: %[[R1:.*]] = "tf.Const"() {value = dense<> : tensor<0xi32>} : () -> tensor<0xi32> + // CHECK-DAG: %[[R0:.*]] = "tf.Const"() <{value = dense<[0, 1]> : tensor<2xi32>}> : () -> tensor<2xi32> + // CHECK-DAG: %[[R1:.*]] = "tf.Const"() <{value = dense<> : tensor<0xi32>}> : () -> tensor<0xi32> // CHECK-NOT: tf.BroadcastGradientArgs // CHECK: return %[[R0]], %[[R1]] @@ -631,8 +631,8 @@ func.func @testBroadcastGradientArgI64() -> (tensor<2xi64>, tensor<0xi64>) { %s0 = "tf.Const"() {value = dense<> : tensor<0xi64>} : () -> tensor<0xi64> %s1 = "tf.Const"() {value = dense<[2, 4]> : tensor<2xi64>} : () -> tensor<2xi64> %r0, %r1 = "tf.BroadcastGradientArgs"(%s0, %s1) {} : (tensor<0xi64>, tensor<2xi64>) -> (tensor<2xi64>, tensor<0xi64>) - // CHECK-DAG: %[[R0:.*]] = "tf.Const"() {value = dense<[0, 1]> : tensor<2xi64>} : () -> tensor<2xi64> - // CHECK-DAG: %[[R1:.*]] = "tf.Const"() {value = dense<> : tensor<0xi64>} : () -> tensor<0xi64> + // CHECK-DAG: %[[R0:.*]] = "tf.Const"() <{value = dense<[0, 1]> : tensor<2xi64>}> : () -> tensor<2xi64> + // CHECK-DAG: %[[R1:.*]] = "tf.Const"() <{value = dense<> : tensor<0xi64>}> : () -> tensor<0xi64> // CHECK-NOT: tf.BroadcastGradientArgs // CHECK: return %[[R0]], %[[R1]] @@ -643,7 +643,7 @@ func.func @testBroadcastGradientArgI64() -> (tensor<2xi64>, tensor<0xi64>) { func.func @testEmptyResults(%arg0: tensor<0x2xf32>) -> tensor<0x2xf32> { %indices = "tf.Const"() {value = dense<> : tensor<0xi32>} : () -> tensor<0xi32> - // CHECK: "tf.Const"() {value = dense<> : tensor<0x2xf32>} : () -> tensor<0x2xf32> + // CHECK: "tf.Const"() <{value = dense<> : tensor<0x2xf32>}> : () -> tensor<0x2xf32> %0 = "tf.DynamicStitch"(%indices, %arg0) : (tensor<0xi32>, tensor<0x2xf32>) -> tensor<0x2xf32> func.return %0 : tensor<0x2xf32> } @@ -668,7 +668,7 @@ func.func @range_int() -> tensor { %cst_1 = arith.constant dense<4> : tensor %cst_2 = arith.constant dense<1> : tensor - // CHECK: %[[CST:.*]] = "tf.Const"() {value = dense<[0, 1, 2, 3]> : tensor<4xi32>} : () -> tensor + // CHECK: %[[CST:.*]] = "tf.Const"() <{value = dense<[0, 1, 2, 3]> : tensor<4xi32>}> : () -> tensor // CHECK: return %[[CST]] %0 = "tf.Range"(%cst, %cst_1, %cst_2) : (tensor, tensor, tensor) -> tensor func.return %0 : tensor @@ -680,7 +680,7 @@ func.func @range_uint() -> tensor { %cst_1 = arith.constant dense<4> : tensor %cst_2 = arith.constant dense<1> : tensor - // CHECK: %[[CST:.*]] = "tf.Const"() {value = dense<[0, 1, 2, 3]> : tensor<4xui32>} : () -> tensor + // CHECK: %[[CST:.*]] = "tf.Const"() <{value = dense<[0, 1, 2, 3]> : tensor<4xui32>}> : () -> tensor // CHECK: return %[[CST]] %0 = "tf.Range"(%cst, %cst_1, %cst_2) : (tensor, tensor, tensor) -> tensor func.return %0 : tensor @@ -692,7 +692,7 @@ func.func @range_float() -> tensor { %cst_1 = arith.constant dense<4.0> : tensor %cst_2 = arith.constant dense<1.0> : tensor - // CHECK: %[[CST:.*]] = "tf.Const"() {value = dense<[0.000000e+00, 1.000000e+00, 2.000000e+00, 3.000000e+00]> : tensor<4xf32>} : () -> tensor + // CHECK: %[[CST:.*]] = "tf.Const"() <{value = dense<[0.000000e+00, 1.000000e+00, 2.000000e+00, 3.000000e+00]> : tensor<4xf32>}> : () -> tensor // CHECK: return %[[CST]] %0 = "tf.Range"(%cst, %cst_1, %cst_2) : (tensor, tensor, tensor) -> tensor func.return %0 : tensor @@ -700,7 +700,7 @@ func.func @range_float() -> tensor { // CHECK-LABEL: func @testLogicalAndFoldsWithConstantFalse func.func @testLogicalAndFoldsWithConstantFalse(%arg0: tensor) -> (tensor) { - // CHECK: [[CST:%.+]] = "tf.Const"() {value = dense : tensor} : () -> tensor + // CHECK: [[CST:%.+]] = "tf.Const"() <{value = dense : tensor}> : () -> tensor %cst = arith.constant dense : tensor %0 = "tf.LogicalAnd"(%cst, %arg0) : (tensor, tensor) -> tensor @@ -711,7 +711,7 @@ func.func @testLogicalAndFoldsWithConstantFalse(%arg0: tensor) -> (tensor) -> (tensor) { - // CHECK: [[CST:%.+]] = "tf.Const"() {value = dense : tensor} : () -> tensor + // CHECK: [[CST:%.+]] = "tf.Const"() <{value = dense : tensor}> : () -> tensor %cst = arith.constant dense : tensor %0 = "tf.LogicalAnd"(%arg0, %cst) : (tensor, tensor) -> tensor diff --git a/tensorflow/compiler/mlir/tensorflow/tests/constant_op_device_assignment.mlir b/tensorflow/compiler/mlir/tensorflow/tests/constant_op_device_assignment.mlir index 5d890aa8898274..2e326474a82527 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/constant_op_device_assignment.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/constant_op_device_assignment.mlir @@ -2,8 +2,8 @@ // CHECK: func @replace_const_op_test func.func @replace_const_op_test() { - // CHECK-NEXT: %[[RESULT_0:.*]] = "tf.Const"() {device = "/job:worker/replica:0/task:0/device:CPU:1", value = dense<2.000000e+00> : tensor} - // CHECK-NEXT: %[[RESULT_1:.*]] = "tf.Const"() {device = "/job:worker/replica:0/task:0/device:CPU:0", value = dense<2.000000e+00> : tensor} + // CHECK-NEXT: %[[RESULT_0:.*]] = "tf.Const"() <{value = dense<2.000000e+00> : tensor}> {device = "/job:worker/replica:0/task:0/device:CPU:1"} + // CHECK-NEXT: %[[RESULT_1:.*]] = "tf.Const"() <{value = dense<2.000000e+00> : tensor}> {device = "/job:worker/replica:0/task:0/device:CPU:0"} // CHECK-NEXT: %[[RESULT_2:.*]] = "tf.AddV2"(%[[RESULT_1]], %[[RESULT_1]]) {device = "/job:worker/replica:0/task:0/device:CPU:0"} // CHECK-NEXT: %[[RESULT_3:.*]] = "tf.AddV2"(%[[RESULT_0]], %[[RESULT_0]]) {device = "/job:worker/replica:0/task:0/device:CPU:1"} %0 = "tf.Const"() {value = dense<2.000000e+00> : tensor} : () -> tensor @@ -14,7 +14,7 @@ func.func @replace_const_op_test() { // CHECK: func @no_change_test func.func @no_change_test() -> () { - // CHECK-NEXT: %[[RESULT_0:.*]] = "tf.Const"() {value = dense<1> : tensor} : () -> tensor + // CHECK-NEXT: %[[RESULT_0:.*]] = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor // CHECK-NEXT: %[[RESULT_1:.*]] = "tf.AddV2"(%[[RESULT_0]], %[[RESULT_0]]) : (tensor, tensor) -> tensor %0 = "tf.Const"() {value = dense<1> : tensor} : () -> tensor %1 = "tf.AddV2"(%0, %0) : (tensor, tensor) -> tensor diff --git a/tensorflow/compiler/mlir/tensorflow/tests/convert-tf-control-flow-to-scf.mlir b/tensorflow/compiler/mlir/tensorflow/tests/convert-tf-control-flow-to-scf.mlir index 77137819da6485..6368981ba13be5 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/convert-tf-control-flow-to-scf.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/convert-tf-control-flow-to-scf.mlir @@ -20,11 +20,11 @@ func.func @test_supported_lowering_of_tf_if_region1(%arg0: tensor, %arg1: te // CHECK-NEXT: %[[RES:.*]]:2 = scf.if %[[COND]] -> (tensor<*xf32>, tensor<4xf32>) { // CHECK-NEXT: %[[CALL:.*]] = func.call @test_if_then1(%[[ARG1]]) : (tensor<4xf32>) -> tensor<4xf32> // CHECK-NEXT: %[[ADD:.*]] = "tf.AddV2"(%[[CALL]], %[[CALL]]) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32> - // CHECK-NEXT: %[[CAST:.*]] = "tf.Cast"(%[[CALL]]) {Truncate = false} : (tensor<4xf32>) -> tensor<*xf32> + // CHECK-NEXT: %[[CAST:.*]] = "tf.Cast"(%[[CALL]]) <{Truncate = false}> : (tensor<4xf32>) -> tensor<*xf32> // CHECK-NEXT: scf.yield %[[CAST]], %[[ADD]] : tensor<*xf32>, tensor<4xf32> // CHECK-NEXT: } else { // CHECK-NEXT: %[[CALL_0:.*]] = func.call @test_if_else1(%[[ARG1]]) : (tensor<4xf32>) -> tensor<4xf32> - // CHECK-NEXT: %[[CAST_0:.*]] = "tf.Cast"(%[[CALL_0]]) {Truncate = false} : (tensor<4xf32>) -> tensor<*xf32> + // CHECK-NEXT: %[[CAST_0:.*]] = "tf.Cast"(%[[CALL_0]]) <{Truncate = false}> : (tensor<4xf32>) -> tensor<*xf32> // CHECK-NEXT: scf.yield %[[CAST_0]], %[[CALL_0]] : tensor<*xf32>, tensor<4xf32> // CHECK-NEXT: } // CHECK-NEXT: return %[[RES]]#0, %[[RES]]#1 : tensor<*xf32>, tensor<4xf32> @@ -72,7 +72,7 @@ func.func @test_supported_lowering_of_tf_while_region(%arg0: tensor, %arg1: }) {is_stateless = false} : (tensor, tensor<*xf32>) -> (tensor, tensor<*xf32>) func.return %0#0 : tensor - // CHECK-NEXT: %[[CST:.*]] = "tf.Const"() {value = dense<1.000000e+00> : tensor} : () -> tensor + // CHECK-NEXT: %[[CST:.*]] = "tf.Const"() <{value = dense<1.000000e+00> : tensor}> : () -> tensor // CHECK-NEXT: %[[RES:.*]]:2 = scf.while (%[[ARG3:.*]] = %[[ARG0]], %[[ARG4:.*]] = %[[ARG2]]) : (tensor, tensor<*xf32>) -> (tensor, tensor<*xf32>) { // CHECK-NEXT: %[[IDEN:.*]] = "tf.Identity"(%[[ARG3]]) : (tensor) -> tensor // CHECK-NEXT: %[[ADD:.*]] = "tf.Add"(%[[ARG1]], %[[ARG3]]) : (tensor, tensor) -> tensor diff --git a/tensorflow/compiler/mlir/tensorflow/tests/convert_control_to_data_outputs.mlir b/tensorflow/compiler/mlir/tensorflow/tests/convert_control_to_data_outputs.mlir index 4968054455f658..d473c1a7d7b67e 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/convert_control_to_data_outputs.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/convert_control_to_data_outputs.mlir @@ -49,7 +49,7 @@ func.func @simple_independent_chains(%arg0: !tf_res, %arg1: !tf_res, %arg2: tens tf_executor.graph { // CHECK: %[[A_CONTROL:.*]] = tf_executor.island wraps "tf.OpA"() : () -> () %control_A = tf_executor.island wraps "tf.OpA"() : () -> () - // CHECK: %[[CHAIN_CONSTANT:.*]], %{{.*}} = tf_executor.island wraps "tf.Const"() {value = dense<1> : tensor} : () -> tensor + // CHECK: %[[CHAIN_CONSTANT:.*]], %{{.*}} = tf_executor.island wraps "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor // CHECK: %[[WHILE_OUT:.*]]:6, %[[WHILE_CONTROL:.*]] = tf_executor.island(%[[A_CONTROL]]) wraps "tf.While"(%[[ARG_0]], %[[ARG_1]], %[[ARG_2]], %[[ARG_2]], %[[CHAIN_CONSTANT]], %[[CHAIN_CONSTANT]]) %while_out:4, %control_while = tf_executor.island(%control_A) wraps "tf.While"(%arg0, %arg1, %arg2, %arg2) {body = @simple_independent_chains_while_body, cond = @simple_independent_chains_while_cond, is_stateless = false} : (tensor>>, tensor>>, tensor, tensor) -> (tensor>>, tensor>>, tensor, tensor) // CHECK: %[[B_CONTROL:.*]] = tf_executor.island(%[[WHILE_CONTROL]]) wraps "tf.OpB"() : () -> () @@ -117,7 +117,7 @@ func.func @intersecting_chains_while_cond(%arg0: !tf_res, %arg1: !tf_res, %arg2: func.func @intersecting_chains(%arg0: !tf_res, %arg1: !tf_res, %arg2: tensor) { // CHECK: tf_executor.graph { tf_executor.graph { - // CHECK: %[[CHAIN_CONSTANT:.*]], %{{.*}} = tf_executor.island wraps "tf.Const"() {value = dense<1> : tensor} : () -> tensor + // CHECK: %[[CHAIN_CONSTANT:.*]], %{{.*}} = tf_executor.island wraps "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor // CHECK: %[[WHILE_OUT:.*]]:5, %[[WHILE_CONTROL:.*]] = tf_executor.island wraps "tf.While"(%[[ARG_0]], %[[ARG_1]], %[[ARG_2]], %[[ARG_2]], %[[CHAIN_CONSTANT]]) %while_out:4, %while_control = tf_executor.island wraps "tf.While"(%arg0, %arg1, %arg2, %arg2) {body = @intersecting_chains_while_body, cond = @intersecting_chains_while_cond, is_stateless = false} : (tensor>>, tensor>>, tensor, tensor) -> (tensor>>, tensor>>, tensor, tensor) // CHECK: tf_executor.fetch @@ -167,12 +167,12 @@ func.func @multiple_callers_while_cond(%arg0: !tf_res, %arg1: tensor) -> (t func.func @multiple_callers(%arg0: !tf_res, %arg1: tensor) { // CHECK: tf_executor.graph { tf_executor.graph { - // CHECK: %[[CHAIN_CONSTANT_0:.*]], %{{.*}} = tf_executor.island wraps "tf.Const"() {value = dense<1> : tensor} : () -> tensor + // CHECK: %[[CHAIN_CONSTANT_0:.*]], %{{.*}} = tf_executor.island wraps "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor // CHECK: %[[WHILE_OUT:.*]]:3, %[[WHILE_CONTROL:.*]] = tf_executor.island wraps "tf.While"(%[[ARG_0]], %[[ARG_1]], %[[CHAIN_CONSTANT_0]]) %while_0_out:2, %while_0_control = tf_executor.island wraps "tf.While"(%arg0, %arg1) {body = @multiple_callers_while_body, cond = @multiple_callers_while_cond, is_stateless = false} : (tensor>>, tensor) -> (tensor>>, tensor) // CHECK: %[[CONTROL_A:.*]] = tf_executor.island(%[[WHILE_CONTROL]]) wraps "tf.OpA"() : () -> () %control_A = tf_executor.island(%while_0_control) wraps "tf.OpA"() : () -> () - // CHECK: %[[CHAIN_CONSTANT_1:.*]], %{{.*}} = tf_executor.island wraps "tf.Const"() {value = dense<1> : tensor} : () -> tensor + // CHECK: %[[CHAIN_CONSTANT_1:.*]], %{{.*}} = tf_executor.island wraps "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor // CHECK: %[[WHILE_OUT:.*]]:3, %[[WHILE_CONTROL:.*]] = tf_executor.island(%[[CONTROL_A]]) wraps "tf.While"(%[[ARG_0]], %[[ARG_1]], %[[CHAIN_CONSTANT_1]]) %while_1_out:2, %while_1_control = tf_executor.island(%control_A) wraps "tf.While"(%arg0, %arg1) {body = @multiple_callers_while_body, cond = @multiple_callers_while_cond, is_stateless = false} : (tensor>>, tensor) -> (tensor>>, tensor) // CHECK: tf_executor.fetch @@ -223,7 +223,7 @@ func.func @nested_loop_while_body_outer(%arg0: !tf_res, %arg1: tensor) -> ( // CHECK: %[[GRAPH_OUT:.*]]:3 = tf_executor.graph { %graph:2 = tf_executor.graph { // CHECK: %{{.*}}, %[[CONTROL_CHAIN_0_SRC:.*]] = tf_executor.island wraps "tf.Identity"(%[[CHAIN_0]]) : (tensor) -> tensor - // CHECK: %[[CHAIN_CONSTANT:.*]], %{{.*}} = tf_executor.island wraps "tf.Const"() {value = dense<1> : tensor} : () -> tensor + // CHECK: %[[CHAIN_CONSTANT:.*]], %{{.*}} = tf_executor.island wraps "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor // CHECK: %[[WHILE_OUT:.*]]:3, %[[WHILE_CONTROL:.*]] = tf_executor.island(%[[CONTROL_CHAIN_0_SRC]]) wraps "tf.While"(%[[RES_0]], %[[ARG_1]], %[[CHAIN_CONSTANT]]) %while_out:2, %while_control = tf_executor.island() wraps "tf.While"(%arg0, %arg1) {body = @nested_loop_while_body_inner, cond = @nested_loop_while_cond_inner, is_stateless = false} : (tensor>>, tensor) -> (tensor>>, tensor) // CHECK: %[[CHAIN_0_SINK:.*]], %{{.*}} = tf_executor.island(%[[WHILE_CONTROL]]) wraps "tf.Identity"(%[[CHAIN_0]]) : (tensor) -> tensor @@ -252,7 +252,7 @@ func.func @nested_loop_while_cond_outer(%arg0: !tf_res, %arg1: tensor) -> ( func.func @nested_while(%arg0: !tf_res, %arg1: tensor) { // CHECK: tf_executor.graph { tf_executor.graph { - // CHECK: %[[CHAIN_CONSTANT:.*]], %{{.*}} = tf_executor.island wraps "tf.Const"() {value = dense<1> : tensor} : () -> tensor + // CHECK: %[[CHAIN_CONSTANT:.*]], %{{.*}} = tf_executor.island wraps "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor // CHECK: %[[WHILE_OUT:.*]]:3, %[[WHILE_CONTROL:.*]] = tf_executor.island wraps "tf.While"(%[[ARG_0]], %[[ARG_1]], %[[CHAIN_CONSTANT]]) %while_out:2, %while_control = tf_executor.island() wraps "tf.While"(%arg0, %arg1) {body = @nested_loop_while_body_outer, cond = @nested_loop_while_cond_outer, is_stateless = false} : (tensor>>, tensor) -> (tensor>>, tensor) // CHECK: tf_executor.fetch @@ -396,7 +396,7 @@ func.func @unique_resource_chain(%arg0: tensor, %arg1: tensor) { // CHECK-LABEL: func @unique_resource_chain // CHECK-SAME: %[[ARG_0:.*]]: tensor, %[[ARG_1:.*]]: tensor // CHECK: tf_executor.graph -// CHECK: %[[WHILE:.*]]:2, %[[WHILE_CONTROL:.*]] = tf_executor.island wraps "tf.While"(%[[ARG_0]], %[[ARG_1]]) {body = @unique_resource_chain_while_body, cond = @unique_resource_chain_while_cond, is_stateless = false} : (tensor, tensor) -> (tensor, tensor) +// CHECK: %[[WHILE:.*]]:2, %[[WHILE_CONTROL:.*]] = tf_executor.island wraps "tf.While"(%[[ARG_0]], %[[ARG_1]]) <{body = @unique_resource_chain_while_body, cond = @unique_resource_chain_while_cond, is_stateless = false}> : (tensor, tensor) -> (tensor, tensor) // CHECK: tf_executor.fetch // CHECK: } // CHECK: return @@ -417,12 +417,12 @@ func.func @unique_resource_chain_while_body(%arg0: tensor, %arg1: tensor, %[[ARG_1:.*]]: tensor // CHECK: %[[GRAPH:.*]]:2 = tf_executor.graph { -// CHECK: %[[THOUSAND:.*]], %{{.*}} = tf_executor.island wraps "tf.Const"() {value = dense<1000> : tensor} : () -> tensor -// CHECK: %[[STACK_HANDLE:.*]], %{{.*}} = tf_executor.island wraps "tf.StackV2"(%[[THOUSAND]]) {elem_type = f32} : (tensor) -> tensor>> +// CHECK: %[[THOUSAND:.*]], %{{.*}} = tf_executor.island wraps "tf.Const"() <{value = dense<1000> : tensor}> : () -> tensor +// CHECK: %[[STACK_HANDLE:.*]], %{{.*}} = tf_executor.island wraps "tf.StackV2"(%[[THOUSAND]]) <{elem_type = f32}> : (tensor) -> tensor>> // CHECK: %{{.*}}, %[[STACK_PUSH_CONTROL:.*]] = tf_executor.island wraps "tf.StackPushV2"(%[[STACK_HANDLE]], %[[ARG_1]]) : (tensor>>, tensor) -> tensor // CHECK: %[[ADD:.*]], %{{.*}} = tf_executor.island wraps "tf.Add"(%[[ARG_1]], %[[ARG_1]]) : (tensor, tensor) -> tensor // CHECK: %{{.*}}, %{{.*}} = tf_executor.island(%[[STACK_PUSH_CONTROL]]) wraps "tf.StackPushV2"(%[[STACK_HANDLE]], %[[ADD]]) : (tensor>>, tensor) -> tensor -// CHECK: %[[ONE:.*]], %{{.*}} = tf_executor.island wraps "tf.Const"() {value = dense<1> : tensor} : () -> tensor +// CHECK: %[[ONE:.*]], %{{.*}} = tf_executor.island wraps "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor // CHECK: %[[COUNTER:.*]], %{{.*}} = tf_executor.island wraps "tf.Add"(%[[ARG_0]], %[[ONE]]) : (tensor, tensor) -> tensor // CHECK: tf_executor.fetch %[[COUNTER]], %[[ARG_1]] : tensor, tensor // CHECK: } @@ -439,7 +439,7 @@ func.func @unique_resource_chain_while_cond(%arg0: tensor, %arg1: tensor, %[[ARG_1:.*]]: tensor // CHECK: %[[GRAPH:.*]] = tf_executor.graph -// CHECK: %[[CONST:.*]], %[[CONST_CONTROL:.*]] = tf_executor.island wraps "tf.Const"() {value = dense<1000> : tensor} : () -> tensor +// CHECK: %[[CONST:.*]], %[[CONST_CONTROL:.*]] = tf_executor.island wraps "tf.Const"() <{value = dense<1000> : tensor}> : () -> tensor // CHECK: %[[LESS:.*]], %[[LESS_CONTROL:.*]] = tf_executor.island wraps "tf.Less"(%[[CONST]], %[[ARG_0]]) : (tensor, tensor) -> tensor // CHECK: tf_executor.fetch %[[LESS]] : tensor // CHECK: } @@ -464,8 +464,8 @@ func.func @mixed_unique_resource_chain(%arg0: tensor, %arg1: tensor) { // CHECK-LABEL: func @mixed_unique_resource_chain // CHECK-SAME: %[[ARG_0:.*]]: tensor, %[[ARG_1:.*]]: tensor // CHECK: tf_executor.graph -// CHECK: %[[CHAIN_TOKEN:.*]], %{{.*}} = tf_executor.island wraps "tf.Const"() {value = dense<1> : tensor} : () -> tensor -// CHECK: %[[WHILE:.*]]:3, %[[WHILE_CONTROL:.*]] = tf_executor.island wraps "tf.While"(%[[ARG_0]], %[[ARG_1]], %[[CHAIN_TOKEN]]) {body = @mixed_unique_resource_chain_while_body, cond = @mixed_unique_resource_chain_while_cond, is_stateless = false} : (tensor, tensor, tensor) -> (tensor, tensor, tensor) +// CHECK: %[[CHAIN_TOKEN:.*]], %{{.*}} = tf_executor.island wraps "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor +// CHECK: %[[WHILE:.*]]:3, %[[WHILE_CONTROL:.*]] = tf_executor.island wraps "tf.While"(%[[ARG_0]], %[[ARG_1]], %[[CHAIN_TOKEN]]) <{body = @mixed_unique_resource_chain_while_body, cond = @mixed_unique_resource_chain_while_cond, is_stateless = false}> : (tensor, tensor, tensor) -> (tensor, tensor, tensor) // CHECK: tf_executor.fetch // CHECK: } // CHECK: return @@ -489,14 +489,14 @@ func.func @mixed_unique_resource_chain_while_body(%arg0: tensor, %arg1: ten // CHECK-SAME: %[[ARG_0:.*]]: tensor, %[[ARG_1:.*]]: tensor, %[[CHAIN_TOKEN:.*]]: tensor // CHECK: %[[GRAPH:.*]]:3 = tf_executor.graph // CHECK: %{{.*}}, %[[CHAIN_SRC:.*]] = tf_executor.island wraps "tf.Identity"(%[[CHAIN_TOKEN]]) : (tensor) -> tensor -// CHECK: %[[THOUSAND:.*]], %{{.*}} = tf_executor.island wraps "tf.Const"() {value = dense<1000> : tensor} : () -> tensor -// CHECK: %[[STACK_HANDLE:.*]], %{{.*}} = tf_executor.island wraps "tf.StackV2"(%[[THOUSAND]]) {elem_type = f32} : (tensor) -> tensor>> +// CHECK: %[[THOUSAND:.*]], %{{.*}} = tf_executor.island wraps "tf.Const"() <{value = dense<1000> : tensor}> : () -> tensor +// CHECK: %[[STACK_HANDLE:.*]], %{{.*}} = tf_executor.island wraps "tf.StackV2"(%[[THOUSAND]]) <{elem_type = f32}> : (tensor) -> tensor>> // CHECK: %{{.*}}, %[[STACK_PUSH_CONTROL:.*]] = tf_executor.island wraps "tf.StackPushV2"(%[[STACK_HANDLE]], %[[ARG_1]]) : (tensor>>, tensor) -> tensor // CHECK: %[[ADD:.*]], %{{.*}} = tf_executor.island wraps "tf.Add"(%[[ARG_1]], %[[ARG_1]]) : (tensor, tensor) -> tensor // CHECK: %{{.*}}, %{{.*}} = tf_executor.island(%[[STACK_PUSH_CONTROL]]) wraps "tf.StackPushV2"(%[[STACK_HANDLE]], %[[ADD]]) : (tensor>>, tensor) -> tensor -// CHECK: %[[ONE:.*]], %{{.*}} = tf_executor.island wraps "tf.Const"() {value = dense<1> : tensor} : () -> tensor +// CHECK: %[[ONE:.*]], %{{.*}} = tf_executor.island wraps "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor // CHECK: %[[COUNTER:.*]], %{{.*}} = tf_executor.island wraps "tf.Add"(%[[ARG_0]], %[[ONE]]) : (tensor, tensor) -> tensor -// CHECK: %[[VAR_HANDLE:.*]], %{{.*}} = tf_executor.island wraps "tf.VarHandleOp"() {container = "c", shared_name = "v0"} : () -> tensor>> +// CHECK: %[[VAR_HANDLE:.*]], %{{.*}} = tf_executor.island wraps "tf.VarHandleOp"() <{container = "c", shared_name = "v0"}> : () -> tensor>> // CHECK: %[[ASSIGN_CONTROL:.*]] = tf_executor.island(%[[CHAIN_SRC]]) wraps "tf.AssignVariableOp"(%[[VAR_HANDLE]], %[[ARG_1]]) : (tensor>>, tensor) -> () // CHECK: %[[CHAIN_SINK:.*]], %{{.*}} = tf_executor.island(%[[ASSIGN_CONTROL]]) wraps "tf.Identity"(%[[CHAIN_TOKEN]]) : (tensor) -> tensor // CHECK: tf_executor.fetch %[[COUNTER]], %[[ARG_1]], %[[CHAIN_SINK]] : tensor, tensor, tensor @@ -514,7 +514,7 @@ func.func @mixed_unique_resource_chain_while_cond(%arg0: tensor, %arg1: ten // CHECK-LABEL: func @mixed_unique_resource_chain_while_cond // CHECK-SAME: %[[ARG_0:.*]]: tensor, %[[ARG_1:.*]]: tensor, %[[CHAIN_TOKEN:.*]]: tensor // CHECK: %[[GRAPH:.*]] = tf_executor.graph -// CHECK: %[[CONST:.*]], %[[CONST_CONTROL:.*]] = tf_executor.island wraps "tf.Const"() {value = dense<1000> : tensor} : () -> tensor +// CHECK: %[[CONST:.*]], %[[CONST_CONTROL:.*]] = tf_executor.island wraps "tf.Const"() <{value = dense<1000> : tensor}> : () -> tensor // CHECK: %[[LESS:.*]], %[[LESS_CONTROL:.*]] = tf_executor.island wraps "tf.Less"(%[[CONST]], %[[ARG_0]]) : (tensor, tensor) -> tensor // CHECK: tf_executor.fetch %[[LESS]] : tensor // CHECK: } diff --git a/tensorflow/compiler/mlir/tensorflow/tests/convert_launch_func_to_tf_call.mlir b/tensorflow/compiler/mlir/tensorflow/tests/convert_launch_func_to_tf_call.mlir index ea44e0fccfb6ee..4c532cdabc0eaa 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/convert_launch_func_to_tf_call.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/convert_launch_func_to_tf_call.mlir @@ -11,8 +11,8 @@ func.func @single_launch_func(%arg0: tensor) -> tensor { %2 = "tf.A"(%arg0) : (tensor) -> tensor // CHECK: %[[CALL_OUTPUT:[0-9]*]] = "tf.PartitionedCall"(%[[A_OUTPUT]]) - // CHECK-SAME: device = "/device:test_device:0" // CHECK-SAME: f = @_func + // CHECK-SAME: device = "/device:test_device:0" %3 = "tf_device.launch_func"(%2) {device = "/device:test_device:0", func = @_func} : (tensor) -> tensor // CHECK: tf_executor.yield %[[CALL_OUTPUT]] @@ -40,13 +40,13 @@ func.func @multi_launch_func(%arg0: tensor) -> tensor { %2 = "tf.A"(%arg0) : (tensor) -> tensor // CHECK: %[[CALL_OUTPUT_0:[0-9]*]] = "tf.PartitionedCall"(%[[A_OUTPUT]]) - // CHECK-SAME: device = "/device:test_device:0" // CHECK-SAME: f = @_func + // CHECK-SAME: device = "/device:test_device:0" %3 = "tf_device.launch_func"(%2) {device = "/device:test_device:0", func = @_func} : (tensor) -> tensor // CHECK: %[[CALL_OUTPUT_1:[0-9]*]] = "tf.PartitionedCall"(%[[CALL_OUTPUT_0]]) - // CHECK-SAME: device = "/device:test_device:1" // CHECK-SAME: f = @_func + // CHECK-SAME: device = "/device:test_device:1" %4 = "tf_device.launch_func"(%3) {device = "/device:test_device:1", func = @_func} : (tensor) -> tensor // CHECK: tf_executor.yield %[[CALL_OUTPUT_1]] diff --git a/tensorflow/compiler/mlir/tensorflow/tests/decompose_resource_ops.mlir b/tensorflow/compiler/mlir/tensorflow/tests/decompose_resource_ops.mlir index 7af6ff18023782..300a766afb22ec 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/decompose_resource_ops.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/decompose_resource_ops.mlir @@ -9,7 +9,7 @@ // CHECK-LABEL: func @decomposition_outside_cluster func.func @decomposition_outside_cluster() { %0 = "tf.VarHandleOp"() {container = "c", shared_name = "v"} : () -> tensor<*x!tf_type.resource>> - // CHECK: %[[ONE:.*]] = "tf.Const"() {value = dense<1> : tensor} + // CHECK: %[[ONE:.*]] = "tf.Const"() <{value = dense<1> : tensor}> %1 = "tf.Const"() {value = dense<1> : tensor} : () -> tensor // CLUSTER-ONLY: "tf.AssignAddVariableOp" // ALWAYS-DECOMPOSE-NOT: "tf.AssignAddVariableOp" @@ -74,7 +74,7 @@ func.func @decompose_use_subtype() { "tf_device.cluster"() ({ %0 = "tf.VarHandleOp"() {container = "c", shared_name = "v"} : () -> tensor<*x!tf_type.resource>> - // CHECK: %[[ONE:.*]] = "tf.Const"() {value = dense<1> : tensor} + // CHECK: %[[ONE:.*]] = "tf.Const"() <{value = dense<1> : tensor}> // CHECK: %[[RES_READ_VAL:[0-9]*]] = "tf.ReadVariableOp" // CHECK-SAME: (tensor<*x!tf_type.resource>>) -> tensor<2x8xi32> // CHECK: "tf.AddV2"(%[[RES_READ_VAL]], %[[ONE]]) @@ -98,7 +98,7 @@ func.func @decompose_assign_add_variable_op() -> () { %0 = "tf.VarHandleOp"() {container = "c", shared_name = "v"} : () -> tensor>> - // CHECK: %[[ONE:.*]] = "tf.Const"() {value = dense<1> : tensor} + // CHECK: %[[ONE:.*]] = "tf.Const"() <{value = dense<1> : tensor}> // CHECK: %[[RES_READ_VAL:[0-9]*]] = "tf.ReadVariableOp" // CHECK: "tf.AddV2"(%[[RES_READ_VAL]], %[[ONE]]) // CHECK: "tf.AssignVariableOp" @@ -121,7 +121,7 @@ func.func @decompose_assign_sub_variable_op() -> () { %0 = "tf.VarHandleOp"() {container = "c", shared_name = "v"} : () -> tensor>> - // CHECK: %[[ONE:.*]] = "tf.Const"() {value = dense<1> : tensor} + // CHECK: %[[ONE:.*]] = "tf.Const"() <{value = dense<1> : tensor}> // CHECK: %[[RES_READ_VAL:[0-9]*]] = "tf.ReadVariableOp" // CHECK: "tf.Sub"(%[[RES_READ_VAL]], %[[ONE]]) // CHECK: "tf.AssignVariableOp" @@ -323,8 +323,8 @@ func.func @decompose_resource_apply_adagradv2(%arg0: tensor, %arg1: tensor< // CHECK: [[VAR_DELTA:%.*]] = "tf.Div"([[LR_MULTIPLY]], [[DIVISOR]]) : (tensor, tensor<*xf32>) -> tensor<*xf32> // CHECK: [[OLD_VAR:%.*]] = "tf.ReadVariableOp"([[VAR_HANDLE]]) : (tensor<*x!tf_type.resource>>) -> tensor<*xf32> // CHECK: [[NEW_VAR:%.*]] = "tf.Sub"(%9, %8) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32> - // CHECK: "tf.AssignVariableOp"([[VAR_HANDLE]], [[NEW_VAR]]) {validate_shape = false} : (tensor<*x!tf_type.resource>>, tensor<*xf32>) -> () - // CHECK: "tf.AssignVariableOp"([[ACC_HANDLE]], [[NEW_ACC]]) {validate_shape = false} : (tensor<*x!tf_type.resource>>, tensor<*xf32>) -> () + // CHECK: "tf.AssignVariableOp"([[VAR_HANDLE]], [[NEW_VAR]]) <{validate_shape = false}> : (tensor<*x!tf_type.resource>>, tensor<*xf32>) -> () + // CHECK: "tf.AssignVariableOp"([[ACC_HANDLE]], [[NEW_ACC]]) <{validate_shape = false}> : (tensor<*x!tf_type.resource>>, tensor<*xf32>) -> () %0 = "tf.VarHandleOp"() {container = "c", shared_name = "v"} : () -> tensor<*x!tf_type.resource>> %1 = "tf.VarHandleOp"() {container = "c", shared_name = "v"} : () -> tensor<*x!tf_type.resource>> @@ -341,8 +341,8 @@ func.func @decompose_resource_apply_adagradv2(%arg0: tensor, %arg1: tensor< func.func @decompose_resource_apply_adagrad(%arg0: tensor, %arg1: tensor) -> () { "tf_device.cluster"() ({ - // CHECK: %[[VAR_HANDLE:.*]] = "tf.VarHandleOp"() {container = "c", shared_name = "v"} : () -> tensor<*x!tf_type.resource>> - // CHECK: %[[ACCUM_HANDLE:.*]] = "tf.VarHandleOp"() {container = "c", shared_name = "v"} : () -> tensor<*x!tf_type.resource>> + // CHECK: %[[VAR_HANDLE:.*]] = "tf.VarHandleOp"() <{container = "c", shared_name = "v"}> : () -> tensor<*x!tf_type.resource>> + // CHECK: %[[ACCUM_HANDLE:.*]] = "tf.VarHandleOp"() <{container = "c", shared_name = "v"}> : () -> tensor<*x!tf_type.resource>> // CHECK: %[[ACCUM:.*]] = "tf.ReadVariableOp"(%[[ACCUM_HANDLE]]) : (tensor<*x!tf_type.resource>>) -> tensor<*xf32> // CHECK: %[[GRAD_SQUARE:.*]] = "tf.Mul"(%[[GRAD]], %[[GRAD]]) : (tensor, tensor) -> tensor // CHECK: %[[ACCUM_NEW:.*]] = "tf.AddV2"(%[[ACCUM]], %[[GRAD_SQUARE]]) : (tensor<*xf32>, tensor) -> tensor<*xf32> @@ -351,8 +351,8 @@ func.func @decompose_resource_apply_adagrad(%arg0: tensor, %arg1: tensor, tensor<*xf32>) -> tensor<*xf32> // CHECK: %[[VAR:.*]] = "tf.ReadVariableOp"(%[[VAR_HANDLE]]) : (tensor<*x!tf_type.resource>>) -> tensor<*xf32> // CHECK: %[[VAR_NEW:.*]] = "tf.Sub"(%[[VAR]], %[[DIV]]) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32> - // CHECK: "tf.AssignVariableOp"(%[[VAR_HANDLE]], %[[VAR_NEW]]) {validate_shape = false} : (tensor<*x!tf_type.resource>>, tensor<*xf32>) -> () - // CHECK: "tf.AssignVariableOp"(%[[ACCUM_HANDLE]], %[[ACCUM_NEW]]) {validate_shape = false} : (tensor<*x!tf_type.resource>>, tensor<*xf32>) -> () + // CHECK: "tf.AssignVariableOp"(%[[VAR_HANDLE]], %[[VAR_NEW]]) <{validate_shape = false}> : (tensor<*x!tf_type.resource>>, tensor<*xf32>) -> () + // CHECK: "tf.AssignVariableOp"(%[[ACCUM_HANDLE]], %[[ACCUM_NEW]]) <{validate_shape = false}> : (tensor<*x!tf_type.resource>>, tensor<*xf32>) -> () %0 = "tf.VarHandleOp"() {container = "c", shared_name = "v"} : () -> tensor<*x!tf_type.resource>> %1 = "tf.VarHandleOp"() {container = "c", shared_name = "v"} : () -> tensor<*x!tf_type.resource>> @@ -372,7 +372,7 @@ func.func @decompose_resource_apply_adagrad(%arg0: tensor, %arg1: tensor, %arg1: tensor, %arg2: tensor, %arg3: tensor, %arg4: tensor, %arg5: tensor, %arg6: tensor) -> () { "tf_device.cluster"() ({ - // CHECK: [[ONE:%.*]] = "tf.Const"() {value = dense<1.000000e+00> : tensor} + // CHECK: [[ONE:%.*]] = "tf.Const"() <{value = dense<1.000000e+00> : tensor}> // CHECK: [[VAR_HANDLE:%.*]] = "tf.VarHandleOp"() // CHECK: [[M_HANDLE:%.*]] = "tf.VarHandleOp"() // CHECK: [[V_HANDLE:%.*]] = "tf.VarHandleOp"() @@ -422,10 +422,10 @@ func.func @decompose_resource_apply_adam_non_nesterov(%arg0: tensor, %arg1: func.func @decompose_resource_apply_adam_nesterov(%arg0: tensor, %arg1: tensor, %arg2: tensor, %arg3: tensor, %arg4: tensor, %arg5: tensor, %arg6: tensor) -> () { "tf_device.cluster"() ({ - // CHECK: [[ONE:%.*]] = "tf.Const"() {value = dense<1.000000e+00> : tensor} - // CHECK: [[VAR_HANDLE:%.*]] = "tf.VarHandleOp"() {container = "c", shared_name = "v"} - // CHECK: [[M_HANDLE:%.*]] = "tf.VarHandleOp"() {container = "c", shared_name = "v"} - // CHECK: [[V_HANDLE:%.*]] = "tf.VarHandleOp"() {container = "c", shared_name = "v"} + // CHECK: [[ONE:%.*]] = "tf.Const"() <{value = dense<1.000000e+00> : tensor}> + // CHECK: [[VAR_HANDLE:%.*]] = "tf.VarHandleOp"() <{container = "c", shared_name = "v"}> + // CHECK: [[M_HANDLE:%.*]] = "tf.VarHandleOp"() <{container = "c", shared_name = "v"}> + // CHECK: [[V_HANDLE:%.*]] = "tf.VarHandleOp"() <{container = "c", shared_name = "v"}> // CHECK: [[VAL_82:%.*]] = "tf.Sub"([[ONE]], [[BETA2_POWER]]) // CHECK: [[VAL_83:%.*]] = "tf.Sqrt"([[VAL_82]]) // CHECK: [[VAL_84:%.*]] = "tf.Sub"([[ONE]], [[BETA1_POWER]]) @@ -452,9 +452,9 @@ func.func @decompose_resource_apply_adam_nesterov(%arg0: tensor, %arg1: ten // CHECK: [[VAL_105:%.*]] = "tf.Div"([[VAL_102]], [[VAL_104]]) // CHECK: [[OLD_VAR:%.*]] = "tf.ReadVariableOp"([[VAR_HANDLE]]) : (tensor<*x!tf_type.resource>>) -> tensor<*xf32> // CHECK: [[NEW_VAR:%.*]] = "tf.Sub"([[OLD_VAR]], [[VAL_105]]) - // CHECK: "tf.AssignVariableOp"([[VAR_HANDLE]], [[NEW_VAR]]) {validate_shape = false} : (tensor<*x!tf_type.resource>>, tensor<*xf32>) -> () - // CHECK: "tf.AssignVariableOp"([[M_HANDLE]], [[NEW_M]]) {validate_shape = false} : (tensor<*x!tf_type.resource>>, tensor<*xf32>) -> () - // CHECK: "tf.AssignVariableOp"([[V_HANDLE]], [[NEW_V]]) {validate_shape = false} : (tensor<*x!tf_type.resource>>, tensor<*xf32>) -> () + // CHECK: "tf.AssignVariableOp"([[VAR_HANDLE]], [[NEW_VAR]]) <{validate_shape = false}> : (tensor<*x!tf_type.resource>>, tensor<*xf32>) -> () + // CHECK: "tf.AssignVariableOp"([[M_HANDLE]], [[NEW_M]]) <{validate_shape = false}> : (tensor<*x!tf_type.resource>>, tensor<*xf32>) -> () + // CHECK: "tf.AssignVariableOp"([[V_HANDLE]], [[NEW_V]]) <{validate_shape = false}> : (tensor<*x!tf_type.resource>>, tensor<*xf32>) -> () %0 = "tf.VarHandleOp"() {container = "c", shared_name = "v"} : () -> tensor<*x!tf_type.resource>> %1 = "tf.VarHandleOp"() {container = "c", shared_name = "v"} : () -> tensor<*x!tf_type.resource>> @@ -474,7 +474,7 @@ func.func @decompose_resource_apply_adam_nesterov(%arg0: tensor, %arg1: ten func.func @decompose_adam_with_complex_inputs(%arg0: tensor>>>, %arg1: tensor>>>, %arg2: tensor>>>, %arg3: tensor>, %arg4: tensor>, %arg5: tensor>, %arg6: tensor>, %arg7: tensor>, %arg8: tensor>, %arg9: tensor<2xcomplex>) attributes {tf.entry_function = {control_outputs = "Adam/update_Variable_1/ResourceApplyAdam", inputs = "_arg0,_arg1,_arg2,_arg3,_arg4,_arg5,_arg6,_arg7,_arg8,_arg9", outputs = ""}} { "tf_device.cluster"() ({ - // CHECK: "tf.Const"() {value = dense<(1.000000e+00,0.000000e+00)> : tensor>} : () -> tensor> + // CHECK: "tf.Const"() <{value = dense<(1.000000e+00,0.000000e+00)> : tensor>}> : () -> tensor> // CHECK-NOT: tf.ResourceApplyAdam "tf.ResourceApplyAdam"(%arg0, %arg1, %arg2, %arg3, %arg4, %arg5, %arg6, %arg7, %arg8, %arg9) {_XlaHasReferenceVars = false, _xla_inferred_shapes = [], device = "/job:localhost/replica:0/task:0/device:TPU:0", use_locking = false, use_nesterov = false} : (tensor>>>, tensor>>>, tensor>>>, tensor>, tensor>, tensor>, tensor>, tensor>, tensor>, tensor<2xcomplex>) -> () @@ -489,13 +489,13 @@ func.func @decompose_adam_with_complex_inputs(%arg0: tensor func.func @decompose_resource_gather_op(%indices : tensor) -> tensor<*xi32> { %0 = "tf_device.cluster"() ({ - // CHECK: [[ZERO:%.+]] = "tf.Const"() {value = dense<0> : tensor} + // CHECK: [[ZERO:%.+]] = "tf.Const"() <{value = dense<0> : tensor}> // CHECK: [[VAR:%.+]] = "tf.VarHandleOp" %resource = "tf.VarHandleOp"() {container = "c", shared_name = "v"} : () -> tensor<*x!tf_type.resource>> // CHECK: [[READVAR:%.+]] = "tf.ReadVariableOp"([[VAR]]) - // CHECK: [[GATHER:%.+]] = "tf.GatherV2"([[READVAR]], [[INDEX]], [[ZERO]]) {batch_dims = 0 : i64} : (tensor<*xi32>, tensor, tensor) -> tensor<*xi32> + // CHECK: [[GATHER:%.+]] = "tf.GatherV2"([[READVAR]], [[INDEX]], [[ZERO]]) <{batch_dims = 0 : i64}> : (tensor<*xi32>, tensor, tensor) -> tensor<*xi32> // CHECK: return [[GATHER]] %1 = "tf.ResourceGather"(%resource, %indices) : (tensor<*x!tf_type.resource>>, tensor) -> (tensor<*xi32>) tf_device.return %1 : tensor<*xi32> @@ -512,7 +512,7 @@ func.func @decompose_resource_gather_op_subtype(%indices : tensor<5xi32>) -> ten %0 = "tf_device.cluster"() ({ %resource = "tf.VarHandleOp"() {container = "c", shared_name = "v"} : () -> tensor<*x!tf_type.resource>> - // CHECK: "tf.GatherV2"({{.+}}, {{.+}}, {{.+}}) {batch_dims = 1 : i64} : (tensor<2x8x16xi32>, tensor<5xi32>, tensor) -> tensor<2x5x16xi32> + // CHECK: "tf.GatherV2"({{.+}}, {{.+}}, {{.+}}) <{batch_dims = 1 : i64}> : (tensor<2x8x16xi32>, tensor<5xi32>, tensor) -> tensor<2x5x16xi32> %1 = "tf.ResourceGather"(%resource, %indices) {batch_dims = 1} : (tensor<*x!tf_type.resource>>, tensor<5xi32>) -> (tensor<2x5x16xi32>) tf_device.return %1 : tensor<2x5x16xi32> @@ -527,7 +527,7 @@ func.func @decompose_resource_gather_op_subtype(%indices : tensor<5xi32>) -> ten // CHECK-SAME: [[VAR:%.*]]: tensor, [[MG:%.*]]: tensor, [[MS:%.*]]: tensor, [[MOM:%.*]]: tensor, [[LR:%.*]]: tensor, [[RHO:%.*]]: tensor, [[MOMENTUM:%.*]]: tensor, [[EPSILON:%.*]]: tensor, [[GRAD:%.*]]: tensor func.func @decompose_resource_apply_centered_RMS_prop(%arg0: tensor, %arg1: tensor, %arg2: tensor, %arg3: tensor, %arg4: tensor, %arg5: tensor, %arg6: tensor, %arg7: tensor, %arg8: tensor) -> () { "tf_device.cluster"() ({ - // CHECK: [[ONE:%.*]] = "tf.Const"() {value = dense<1.000000e+00> : tensor} + // CHECK: [[ONE:%.*]] = "tf.Const"() <{value = dense<1.000000e+00> : tensor}> // CHECK: [[VAR_HANDLE:%.*]] = "tf.VarHandleOp" // CHECK: [[MG_HANDLE:%.*]] = "tf.VarHandleOp" // CHECK: [[MS_HANDLE:%.*]] = "tf.VarHandleOp" @@ -578,14 +578,14 @@ func.func @decompose_resource_apply_centered_RMS_prop(%arg0: tensor, %arg1: // CHECK-SAME: %[[LR:.*]]: tensor, %[[RHO:.*]]: tensor, %[[MOMENTUM:.*]]: tensor, %[[EPSILON:.*]]: tensor, %[[GRAD:.*]]: tensor) func.func @decompose_resource_apply_RMS_prop(%arg0: tensor<*x!tf_type.resource>, %arg1: tensor<*x!tf_type.resource>, %arg2: tensor<*x!tf_type.resource>, %arg3: tensor, %arg4: tensor, %arg5: tensor, %arg6: tensor, %arg7: tensor) -> () { "tf_device.cluster"() ({ - // CHECK: %[[ONE:.*]] = "tf.Const"() {value = dense<1.000000e+00> : tensor} : () -> tensor + // CHECK: %[[ONE:.*]] = "tf.Const"() <{value = dense<1.000000e+00> : tensor}> : () -> tensor // CHECK: %[[MS:.*]] = "tf.ReadVariableOp"(%[[MS_HANDLE]]) : (tensor<*x!tf_type.resource>) -> tensor<*xf32> // CHECK: %[[MS_RHO:.*]] = "tf.Mul"(%[[MS]], %[[RHO]]) : (tensor<*xf32>, tensor) -> tensor<*xf32> // CHECK: %[[GRAD_SQUARE:.*]] = "tf.Square"(%[[GRAD]]) : (tensor) -> tensor // CHECK: %[[ONE_RHO:.*]] = "tf.Sub"(%[[ONE]], %[[RHO]]) : (tensor, tensor) -> tensor // CHECK: %[[MUL:.*]] = "tf.Mul"(%[[GRAD_SQUARE]], %[[ONE_RHO]]) : (tensor, tensor) -> tensor // CHECK: %[[MS_NEW:.*]] = "tf.AddV2"(%[[MS_RHO]], %[[MUL]]) : (tensor<*xf32>, tensor) -> tensor<*xf32> - // CHECK: "tf.AssignVariableOp"(%[[MS_HANDLE]], %[[MS_NEW]]) {validate_shape = false} : (tensor<*x!tf_type.resource>, tensor<*xf32>) -> () + // CHECK: "tf.AssignVariableOp"(%[[MS_HANDLE]], %[[MS_NEW]]) <{validate_shape = false}> : (tensor<*x!tf_type.resource>, tensor<*xf32>) -> () // CHECK: %[[MOM:.*]] = "tf.ReadVariableOp"(%[[MOM_HANDLE]]) : (tensor<*x!tf_type.resource>) -> tensor<*xf32> // CHECK: %[[MOMENTUM_MOM:.*]] = "tf.Mul"(%[[MOMENTUM]], %[[MOM]]) : (tensor, tensor<*xf32>) -> tensor<*xf32> // CHECK: %[[LR_GRAD:.*]] = "tf.Mul"(%[[LR]], %[[GRAD]]) : (tensor, tensor) -> tensor @@ -593,10 +593,10 @@ func.func @decompose_resource_apply_RMS_prop(%arg0: tensor<*x!tf_type.resource>, // CHECK: %[[SQRT:.*]] = "tf.Sqrt"(%[[ADD]]) : (tensor<*xf32>) -> tensor<*xf32> // CHECK: %[[DIV:.*]] = "tf.Div"(%[[LR_GRAD]], %[[SQRT]]) : (tensor, tensor<*xf32>) -> tensor<*xf32> // CHECK: %[[MOM_NEW:.*]] = "tf.AddV2"(%[[MOMENTUM_MOM]], %[[DIV]]) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32> - // CHECK: "tf.AssignVariableOp"(%[[MOM_HANDLE]], %[[MOM_NEW]]) {validate_shape = false} : (tensor<*x!tf_type.resource>, tensor<*xf32>) -> () + // CHECK: "tf.AssignVariableOp"(%[[MOM_HANDLE]], %[[MOM_NEW]]) <{validate_shape = false}> : (tensor<*x!tf_type.resource>, tensor<*xf32>) -> () // CHECK: %[[VAR:.*]] = "tf.ReadVariableOp"(%[[VAR_HANDLE]]) : (tensor<*x!tf_type.resource>) -> tensor<*xf32> // CHECK: %[[VAR_NEW:.*]] = "tf.Sub"(%[[VAR]], %[[MOM_NEW]]) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32> - // CHECK: "tf.AssignVariableOp"(%[[VAR_HANDLE]], %[[VAR_NEW]]) {validate_shape = false} : (tensor<*x!tf_type.resource>, tensor<*xf32>) -> () + // CHECK: "tf.AssignVariableOp"(%[[VAR_HANDLE]], %[[VAR_NEW]]) <{validate_shape = false}> : (tensor<*x!tf_type.resource>, tensor<*xf32>) -> () "tf.ResourceApplyRMSProp"(%arg0, %arg1, %arg2, %arg3, %arg4, %arg5, %arg6, %arg7) {use_locking = false} : (tensor<*x!tf_type.resource>, tensor<*x!tf_type.resource>, tensor<*x!tf_type.resource>, tensor, tensor, tensor, tensor, tensor) -> () tf_device.return }) : () -> () @@ -609,7 +609,7 @@ func.func @decompose_resource_apply_RMS_prop(%arg0: tensor<*x!tf_type.resource>, // CHECK-LABEL: @decompose_resource_scatter_add_op // CHECK-SAME: ([[INDEX:%.+]]: tensor<2x?xi32>, [[UPDATE:%.+]]: tensor) func.func @decompose_resource_scatter_add_op(%indices : tensor<2x?xi32>, %updates: tensor) { - // CHECK: [[CST:%.+]] = "tf.Const"() {value = dense<-1> : tensor} : () -> tensor + // CHECK: [[CST:%.+]] = "tf.Const"() <{value = dense<-1> : tensor}> : () -> tensor "tf_device.cluster"() ({ // CHECK: [[VAR:%.+]] = "tf.VarHandleOp" %resource = "tf.VarHandleOp"() {container = "c", shared_name = "v"} : () -> tensor<*x!tf_type.resource>> @@ -630,7 +630,7 @@ func.func @decompose_resource_scatter_add_op(%indices : tensor<2x?xi32>, %update // CHECK-LABEL: @decompose_resource_scatter_add_op_1d_indices // CHECK-SAME: ([[INDEX:%.+]]: tensor, [[UPDATE:%.+]]: tensor) func.func @decompose_resource_scatter_add_op_1d_indices(%indices : tensor, %updates: tensor) { - // CHECK: [[CST:%.+]] = "tf.Const"() {value = dense<-1> : tensor} : () -> tensor + // CHECK: [[CST:%.+]] = "tf.Const"() <{value = dense<-1> : tensor}> : () -> tensor "tf_device.cluster"() ({ // CHECK: [[VAR:%.+]] = "tf.VarHandleOp" %resource = "tf.VarHandleOp"() {container = "c", shared_name = "v"} : () -> tensor<*x!tf_type.resource>> @@ -679,7 +679,7 @@ func.func @decompose_ResourceScatterAdd_with_unranked_updates(%resource : tensor // CHECK-LABEL: @decompose_resource_scatter_update_op // CHECK-SAME: ([[INDEX:%.+]]: tensor<2x?xi32>, [[UPDATE:%.+]]: tensor) func.func @decompose_resource_scatter_update_op(%indices : tensor<2x?xi32>, %updates: tensor) { - // CHECK: [[CST:%.+]] = "tf.Const"() {value = dense<-1> : tensor} : () -> tensor + // CHECK: [[CST:%.+]] = "tf.Const"() <{value = dense<-1> : tensor}> : () -> tensor "tf_device.cluster"() ({ // CHECK: [[VAR:%.+]] = "tf.VarHandleOp" %resource = "tf.VarHandleOp"() {container = "c", shared_name = "v"} : () -> tensor<*x!tf_type.resource>> @@ -782,10 +782,10 @@ func.func @decompose_resource_apply_proximal_adagrad_op(%lr: tensor, %l1: t %var = "tf.VarHandleOp"() {container = "c", shared_name = "var"} : () -> tensor<*x!tf_type.resource>> %accum = "tf.VarHandleOp"() {container = "c", shared_name = "accum"} : () -> tensor<*x!tf_type.resource>> - // CHECK-DAG: %[[ONE:.*]] = "tf.Const"() {value = dense<1.000000e+00> : tensor} : () -> tensor - // CHECK-DAG: %[[ZERO:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor} : () -> tensor - // CHECK-DAG: %[[VAR_HANDLE:.*]] = "tf.VarHandleOp"() {container = "c", shared_name = "var"} : () -> tensor<*x!tf_type.resource>> - // CHECK-DAG: %[[ACCUM_HANDLE:.*]] = "tf.VarHandleOp"() {container = "c", shared_name = "accum"} : () -> tensor<*x!tf_type.resource>> + // CHECK-DAG: %[[ONE:.*]] = "tf.Const"() <{value = dense<1.000000e+00> : tensor}> : () -> tensor + // CHECK-DAG: %[[ZERO:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor}> : () -> tensor + // CHECK-DAG: %[[VAR_HANDLE:.*]] = "tf.VarHandleOp"() <{container = "c", shared_name = "var"}> : () -> tensor<*x!tf_type.resource>> + // CHECK-DAG: %[[ACCUM_HANDLE:.*]] = "tf.VarHandleOp"() <{container = "c", shared_name = "accum"}> : () -> tensor<*x!tf_type.resource>> // CHECK-DAG: %[[GRAD_SQ:.*]] = "tf.Square"(%[[GRAD]]) : (tensor<4xf32>) -> tensor<4xf32> // CHECK-DAG: %[[ACCUM:.*]] = "tf.ReadVariableOp"(%[[ACCUM_HANDLE]]) : (tensor<*x!tf_type.resource>>) -> tensor<4xf32> // CHECK-DAG: %[[ACCUM_NEW:.*]] = "tf.AddV2"(%[[ACCUM]], %[[GRAD_SQ]]) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32> @@ -805,8 +805,8 @@ func.func @decompose_resource_apply_proximal_adagrad_op(%lr: tensor, %l1: t // CHECK-DAG: %[[SCALED_L2:.*]] = "tf.Mul"(%[[ADAGRAD_LR]], %[[L2]]) : (tensor<4xf32>, tensor) -> tensor<4xf32> // CHECK-DAG: %[[DENOMINATOR:.*]] = "tf.Add"(%[[ONE]], %[[SCALED_L2]]) : (tensor, tensor<4xf32>) -> tensor<4xf32> // CHECK-DAG: %[[VAR_NEW:.*]] = "tf.Div"(%[[NUMERATOR]], %[[DENOMINATOR]]) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32> - // CHECK-DAG: "tf.AssignVariableOp"(%[[VAR_HANDLE]], %[[VAR_NEW]]) {validate_shape = false} : (tensor<*x!tf_type.resource>>, tensor<4xf32>) -> () - // CHECK-DAG: "tf.AssignVariableOp"(%[[ACCUM_HANDLE]], %[[ACCUM_NEW]]) {validate_shape = false} : (tensor<*x!tf_type.resource>>, tensor<4xf32>) -> () + // CHECK-DAG: "tf.AssignVariableOp"(%[[VAR_HANDLE]], %[[VAR_NEW]]) <{validate_shape = false}> : (tensor<*x!tf_type.resource>>, tensor<4xf32>) -> () + // CHECK-DAG: "tf.AssignVariableOp"(%[[ACCUM_HANDLE]], %[[ACCUM_NEW]]) <{validate_shape = false}> : (tensor<*x!tf_type.resource>>, tensor<4xf32>) -> () "tf.ResourceApplyProximalAdagrad"(%var, %accum, %lr, %l1, %l2, %grad) {use_locking = false} : (tensor<*x!tf_type.resource>>, tensor<*x!tf_type.resource>>, tensor, tensor, tensor, tensor<4xf32>) -> () diff --git a/tensorflow/compiler/mlir/tensorflow/tests/einsum.mlir b/tensorflow/compiler/mlir/tensorflow/tests/einsum.mlir index e1071f3e899c00..6666a08dfde322 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/einsum.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/einsum.mlir @@ -6,7 +6,7 @@ func.func @unary_einsum_reduce_sum_transpose(%arg0: tensor<3x4x5x6xf32>) -> tens // CHECK-LABEL: unary_einsum_reduce_sum_transpose // CHECK-DAG: %[[cst:.*]] = arith.constant dense<3> : tensor<1xi32> // CHECK-DAG: %[[cst_1:.*]] = arith.constant dense<[0, 2, 1]> : tensor<3xi32> - // CHECK: %[[v0:.*]] = "tf.Sum"(%arg0, %[[cst]]) {keep_dims = false} : (tensor<3x4x5x6xf32>, tensor<1xi32>) -> tensor<3x4x5xf32> + // CHECK: %[[v0:.*]] = "tf.Sum"(%arg0, %[[cst]]) <{keep_dims = false}> : (tensor<3x4x5x6xf32>, tensor<1xi32>) -> tensor<3x4x5xf32> // CHECK: %[[v1:.*]] = "tf.Transpose"(%[[v0]], %[[cst_1]]) : (tensor<3x4x5xf32>, tensor<3xi32>) -> tensor<3x5x4xf32> // CHECK: return %[[v1]] : tensor<3x5x4xf32> } @@ -16,7 +16,7 @@ func.func @unary_einsum_reduce_sum_transpose1(%arg0: tensor<3x4x5x6xf32>) -> ten func.return %0 : tensor<3x4x5xf32> // CHECK-LABEL: unary_einsum_reduce_sum_transpose1 // CHECK-DAG: %[[cst:.*]] = arith.constant dense<3> : tensor<1xi32> - // CHECK: %[[v0:.*]] = "tf.Sum"(%arg0, %[[cst]]) {keep_dims = false} : (tensor<3x4x5x6xf32>, tensor<1xi32>) -> tensor<3x4x5xf32> + // CHECK: %[[v0:.*]] = "tf.Sum"(%arg0, %[[cst]]) <{keep_dims = false}> : (tensor<3x4x5x6xf32>, tensor<1xi32>) -> tensor<3x4x5xf32> // CHECK: return %[[v0]] : tensor<3x4x5xf32> } @@ -34,7 +34,7 @@ func.func @unary_einsum_reduce_sum(%arg0: tensor<4x5x6xf32>) -> tensor<4xf32> { func.return %0 : tensor<4xf32> // CHECK-LABEL: unary_einsum_reduce_sum // CHECK-DAG: %[[cst:.*]] = arith.constant dense<[1, 2]> : tensor<2xi32> - // CHECK: %[[v0:.*]] = "tf.Sum"(%arg0, %[[cst]]) {keep_dims = false} : (tensor<4x5x6xf32>, tensor<2xi32>) -> tensor<4xf32> + // CHECK: %[[v0:.*]] = "tf.Sum"(%arg0, %[[cst]]) <{keep_dims = false}> : (tensor<4x5x6xf32>, tensor<2xi32>) -> tensor<4xf32> // CHECK: return %[[v0]] } @@ -42,14 +42,14 @@ func.func @einsum_basic(%arg0: tensor<3x4x5xf32>, %arg1: tensor<3x5x6xf32>) -> t %0 = "tf.Einsum"(%arg0, %arg1) {T = "tfdtype$DT_FLOAT", equation = "ijk,ikm->ijm"}: (tensor<3x4x5xf32>, tensor<3x5x6xf32>) -> tensor<3x4x6xf32> func.return %0 : tensor<3x4x6xf32> // CHECK-LABEL: einsum_basic - // CHECK: "tf.BatchMatMulV2"(%arg0, %arg1) {adj_x = false, adj_y = false} : (tensor<3x4x5xf32>, tensor<3x5x6xf32>) -> tensor<3x4x6xf32> + // CHECK: "tf.BatchMatMulV2"(%arg0, %arg1) <{adj_x = false, adj_y = false}> : (tensor<3x4x5xf32>, tensor<3x5x6xf32>) -> tensor<3x4x6xf32> } func.func @einsum_matmul(%arg0: tensor<7x9xf32>, %arg1: tensor<9x5xf32>) -> tensor<7x5xf32> { %0 = "tf.Einsum"(%arg0, %arg1) {T = "tfdtype$DT_FLOAT", equation = "ae,ed->ad"}: (tensor<7x9xf32>, tensor<9x5xf32>) -> tensor<7x5xf32> func.return %0 : tensor<7x5xf32> // CHECK-LABEL: einsum_matmul - // CHECK: %[[v0:.*]] = "tf.BatchMatMulV2"(%arg0, %arg1) {adj_x = false, adj_y = false} : (tensor<7x9xf32>, tensor<9x5xf32>) -> tensor<7x5xf32> + // CHECK: %[[v0:.*]] = "tf.BatchMatMulV2"(%arg0, %arg1) <{adj_x = false, adj_y = false}> : (tensor<7x9xf32>, tensor<9x5xf32>) -> tensor<7x5xf32> // CHECK: return %[[v0]] : tensor<7x5xf32> } @@ -59,7 +59,7 @@ func.func @einsum_matmul_dynamic_size(%arg0: tensor<2x?x?x?xf32>, %arg1: tensor< // CHECK-LABEL: einsum_matmul_dynamic_size // CHECK-DAG: %[[cst:.*]] = arith.constant dense<[2, -1, 1, 1]> : tensor<4xi64> // CHECK: %[[v0:.*]] = "tf.Reshape"(%arg1, %cst) : (tensor<2x?xf32>, tensor<4xi64>) -> tensor<2x?x1x1xf32> - // CHECK: %[[v1:.*]] = "tf.BatchMatMulV2"(%arg0, %0) {adj_x = false, adj_y = false} : (tensor<2x?x?x?xf32>, tensor<2x?x1x1xf32>) -> tensor<2x?x?x1xf32> + // CHECK: %[[v1:.*]] = "tf.BatchMatMulV2"(%arg0, %0) <{adj_x = false, adj_y = false}> : (tensor<2x?x?x?xf32>, tensor<2x?x1x1xf32>) -> tensor<2x?x?x1xf32> // CHECK: return %[[v1]] : tensor<2x?x?x1xf32> } @@ -67,14 +67,14 @@ func.func @einsum_broadcast(%arg0: tensor<3x4x5xf32>, %arg1: tensor<5x6xf32>) -> %0 = "tf.Einsum"(%arg0, %arg1) {T = "tfdtype$DT_FLOAT", equation = "ijk,km->ijm"}: (tensor<3x4x5xf32>, tensor<5x6xf32>) -> tensor<3x4x6xf32> func.return %0 : tensor<3x4x6xf32> // CHECK-LABEL: einsum_broadcast - // CHECK: "tf.BatchMatMulV2"(%arg0, %arg1) {adj_x = false, adj_y = false} : (tensor<3x4x5xf32>, tensor<5x6xf32>) -> tensor<3x4x6xf32> + // CHECK: "tf.BatchMatMulV2"(%arg0, %arg1) <{adj_x = false, adj_y = false}> : (tensor<3x4x5xf32>, tensor<5x6xf32>) -> tensor<3x4x6xf32> } func.func @einsum_broadcast4(%arg0: tensor<3x4x5x6x7xf32>, %arg1: tensor<7x8xf32>) -> tensor<3x4x5x6x8xf32> { %0 = "tf.Einsum"(%arg0, %arg1) {T = "tfdtype$DT_FLOAT", equation = "abcdh,hg->abcdg"}: (tensor<3x4x5x6x7xf32>, tensor<7x8xf32>) -> tensor<3x4x5x6x8xf32> func.return %0 : tensor<3x4x5x6x8xf32> // CHECK-LABEL: einsum_broadcast4 - // CHECK: "tf.BatchMatMulV2"(%arg0, %arg1) {adj_x = false, adj_y = false} : (tensor<3x4x5x6x7xf32>, tensor<7x8xf32>) -> tensor<3x4x5x6x8xf32> + // CHECK: "tf.BatchMatMulV2"(%arg0, %arg1) <{adj_x = false, adj_y = false}> : (tensor<3x4x5x6x7xf32>, tensor<7x8xf32>) -> tensor<3x4x5x6x8xf32> } func.func @einsum_reducesum(%arg0: tensor<2x5x7xf32>, %arg1: tensor<5x2xf32>) -> tensor<5x7xf32> { @@ -86,7 +86,7 @@ func.func @einsum_reducesum(%arg0: tensor<2x5x7xf32>, %arg1: tensor<5x2xf32>) -> // CHECK-DAG: %[[cst_2:.*]] = arith.constant dense<[5, 7]> : tensor<2xi64> // CHECK: %[[v0:.*]] = "tf.Transpose"(%arg0, %[[cst]]) : (tensor<2x5x7xf32>, tensor<3xi32>) -> tensor<5x7x2xf32> // CHECK: %[[v1:.*]] = "tf.Reshape"(%arg1, %[[cst_1]]) : (tensor<5x2xf32>, tensor<3xi64>) -> tensor<5x2x1xf32> - // CHECK: %[[v2:.*]] = "tf.BatchMatMulV2"(%[[v0]], %[[v1]]) {adj_x = false, adj_y = false} : (tensor<5x7x2xf32>, tensor<5x2x1xf32>) -> tensor<5x7x1xf32> + // CHECK: %[[v2:.*]] = "tf.BatchMatMulV2"(%[[v0]], %[[v1]]) <{adj_x = false, adj_y = false}> : (tensor<5x7x2xf32>, tensor<5x2x1xf32>) -> tensor<5x7x1xf32> // CHECK: %[[v3:.*]] = "tf.Reshape"(%[[v2]], %[[cst_2]]) : (tensor<5x7x1xf32>, tensor<2xi64>) -> tensor<5x7xf32> // CHECK: return %[[v3:.*]] : tensor<5x7xf32> } @@ -99,7 +99,7 @@ func.func @einsum_transpose_matmul(%arg0: tensor<2x5x7xf32>, %arg1: tensor<5x3x2 // CHECK-DAG: %[[cst_0:.*]] = arith.constant dense<[0, 2, 1]> : tensor<3xi32> // CHECK: %[[v0:.*]] = "tf.Transpose"(%arg0, %[[cst]]) : (tensor<2x5x7xf32>, tensor<3xi32>) -> tensor<5x7x2xf32> // CHECK: %[[v1:.*]] = "tf.Transpose"(%arg1, %[[cst_0]]) : (tensor<5x3x2xf32>, tensor<3xi32>) -> tensor<5x2x3xf32> - // CHECK: %[[v2:.*]] = "tf.BatchMatMulV2"(%[[v0]], %[[v1]]) {adj_x = false, adj_y = false} : (tensor<5x7x2xf32>, tensor<5x2x3xf32>) -> tensor<5x7x3xf32> + // CHECK: %[[v2:.*]] = "tf.BatchMatMulV2"(%[[v0]], %[[v1]]) <{adj_x = false, adj_y = false}> : (tensor<5x7x2xf32>, tensor<5x2x3xf32>) -> tensor<5x7x3xf32> // CHECK: %[[v3:.*]] = "tf.Transpose"(%[[v2]], %[[cst_0]]) : (tensor<5x7x3xf32>, tensor<3xi32>) -> tensor<5x3x7xf32> } @@ -111,7 +111,7 @@ func.func @einsum_4D(%arg0: tensor<2x5x7x3xf32>, %arg1: tensor<2x4x7x3xf32>) -> // CHECK-DAG: %[[cst_1:.*]] = arith.constant dense<[0, 2, 3, 1]> : tensor<4xi32> // CHECK: %[[v0:.*]] = "tf.Transpose"(%arg0, %[[cst]]) : (tensor<2x5x7x3xf32>, tensor<4xi32>) -> tensor<2x7x5x3xf32> // CHECK: %[[v1:.*]] = "tf.Transpose"(%arg1, %[[cst_1]]) : (tensor<2x4x7x3xf32>, tensor<4xi32>) -> tensor<2x7x3x4xf32> - // CHECK: "tf.BatchMatMulV2"(%[[v0]], %[[v1]]) {adj_x = false, adj_y = false} : (tensor<2x7x5x3xf32>, tensor<2x7x3x4xf32>) -> tensor<2x7x5x4xf32> + // CHECK: "tf.BatchMatMulV2"(%[[v0]], %[[v1]]) <{adj_x = false, adj_y = false}> : (tensor<2x7x5x3xf32>, tensor<2x7x3x4xf32>) -> tensor<2x7x5x4xf32> } func.func @einsum_matrixdotprod(%arg0: tensor<2x5x7x3xf32>, %arg1: tensor<7x3x4xf32>) -> tensor<2x5x4xf32> { @@ -122,7 +122,7 @@ func.func @einsum_matrixdotprod(%arg0: tensor<2x5x7x3xf32>, %arg1: tensor<7x3x4x // CHECK-DAG: %[[cst_1:.*]] = arith.constant dense<[21, 4]> : tensor<2xi64> // CHECK: %[[v0:.*]] = "tf.Reshape"(%arg0, %[[cst]]) : (tensor<2x5x7x3xf32>, tensor<3xi64>) -> tensor<2x5x21xf32> // CHECK: %[[v1:.*]] = "tf.Reshape"(%arg1, %[[cst_1]]) : (tensor<7x3x4xf32>, tensor<2xi64>) -> tensor<21x4xf32> - // CHECK: "tf.BatchMatMulV2"(%[[v0]], %[[v1]]) {adj_x = false, adj_y = false} : (tensor<2x5x21xf32>, tensor<21x4xf32>) -> tensor<2x5x4xf32> + // CHECK: "tf.BatchMatMulV2"(%[[v0]], %[[v1]]) <{adj_x = false, adj_y = false}> : (tensor<2x5x21xf32>, tensor<21x4xf32>) -> tensor<2x5x4xf32> } func.func @einsum_reshapetail(%arg0: tensor<3x4x5xf32>, %arg1: tensor<5x6x2xf32>) -> tensor<3x4x6x2xf32> { @@ -132,7 +132,7 @@ func.func @einsum_reshapetail(%arg0: tensor<3x4x5xf32>, %arg1: tensor<5x6x2xf32> // CHECK-DAG: %[[cst:.*]] = arith.constant dense<[5, 12]> : tensor<2xi64> // CHECK-DAG: %[[cst_1:.*]] = arith.constant dense<[3, 4, 6, 2]> : tensor<4xi64> // CHECK: %[[v0:.*]] = "tf.Reshape"(%arg1, %[[cst]]) : (tensor<5x6x2xf32>, tensor<2xi64>) -> tensor<5x12xf32> - // CHECK: %[[v1:.*]] = "tf.BatchMatMulV2"(%arg0, %[[v0]]) {adj_x = false, adj_y = false} : (tensor<3x4x5xf32>, tensor<5x12xf32>) -> tensor<3x4x12xf32> + // CHECK: %[[v1:.*]] = "tf.BatchMatMulV2"(%arg0, %[[v0]]) <{adj_x = false, adj_y = false}> : (tensor<3x4x5xf32>, tensor<5x12xf32>) -> tensor<3x4x12xf32> // CHECK: %[[v2:.*]] = "tf.Reshape"(%[[v1]], %[[cst_1]]) : (tensor<3x4x12xf32>, tensor<4xi64>) -> tensor<3x4x6x2xf32> // CHECK: return %[[v2]] : tensor<3x4x6x2xf32> } @@ -144,7 +144,7 @@ func.func @einsum_reduceddim(%arg0: tensor<2x5x7xf32>, %arg1: tensor<2x5x7x3xf32 // CHECK-DAG: %[[cst:.*]] = arith.constant dense<[2, 5, 1, 7]> : tensor<4xi64> // CHECK-DAG: %[[cst_1:.*]] = arith.constant dense<[2, 5, 3]> : tensor<3xi64> // CHECK: %[[v0:.*]] = "tf.Reshape"(%arg0, %[[cst]]) : (tensor<2x5x7xf32>, tensor<4xi64>) -> tensor<2x5x1x7xf32> - // CHECK: %[[v1:.*]] = "tf.BatchMatMulV2"(%[[v0]], %arg1) {adj_x = false, adj_y = false} : (tensor<2x5x1x7xf32>, tensor<2x5x7x3xf32>) -> tensor<2x5x1x3xf32> + // CHECK: %[[v1:.*]] = "tf.BatchMatMulV2"(%[[v0]], %arg1) <{adj_x = false, adj_y = false}> : (tensor<2x5x1x7xf32>, tensor<2x5x7x3xf32>) -> tensor<2x5x1x3xf32> // CHECK: %[[v2:.*]] = "tf.Reshape"(%[[v1]], %[[cst_1]]) : (tensor<2x5x1x3xf32>, tensor<3xi64>) -> tensor<2x5x3xf32> // CHECK: return %[[v2]] : tensor<2x5x3xf32> } @@ -158,7 +158,7 @@ func.func @einsum_transposereduceddim(%arg0: tensor<2x5x7xf32>, %arg1: tensor<2x // CHECK-DAG: %[[cst_2:.*]] = arith.constant dense<[2, 5, 3]> : tensor<3xi64> // CHECK: %[[v0:.*]] = "tf.Transpose"(%arg1, %[[cst]]) : (tensor<2x5x3x7xf32>, tensor<4xi32>) -> tensor<2x5x7x3xf32> // CHECK: %[[v1:.*]] = "tf.Reshape"(%arg0, %[[cst_1]]) : (tensor<2x5x7xf32>, tensor<4xi64>) -> tensor<2x5x1x7xf32> - // CHECK: %[[v2:.*]] = "tf.BatchMatMulV2"(%[[v1]], %[[v0]]) {adj_x = false, adj_y = false} : (tensor<2x5x1x7xf32>, tensor<2x5x7x3xf32>) -> tensor<2x5x1x3xf32> + // CHECK: %[[v2:.*]] = "tf.BatchMatMulV2"(%[[v1]], %[[v0]]) <{adj_x = false, adj_y = false}> : (tensor<2x5x1x7xf32>, tensor<2x5x7x3xf32>) -> tensor<2x5x1x3xf32> // CHECK: %[[v3:.*]] = "tf.Reshape"(%[[v2]], %[[cst_2]]) : (tensor<2x5x1x3xf32>, tensor<3xi64>) -> tensor<2x5x3xf32> // CHECK: return %[[v3]] : tensor<2x5x3xf32> } @@ -169,7 +169,7 @@ func.func @einsum_fourdreducelast(%arg0: tensor<2x5x7x3xf32>, %arg1: tensor<2x3x // CHECK-LABEL: einsum_fourdreducelast // CHECK: %[[cst:.*]] = arith.constant dense<[0, 2, 1, 3]> : tensor<4xi32> // CHECK: %[[v0:.*]] = "tf.Transpose"(%arg1, %[[cst]]) : (tensor<2x3x5x13xf32>, tensor<4xi32>) -> tensor<2x5x3x13xf32> - // CHECK: %[[v1:.*]] = "tf.BatchMatMulV2"(%arg0, %[[v0]]) {adj_x = false, adj_y = false} : (tensor<2x5x7x3xf32>, tensor<2x5x3x13xf32>) -> tensor<2x5x7x13xf32> + // CHECK: %[[v1:.*]] = "tf.BatchMatMulV2"(%arg0, %[[v0]]) <{adj_x = false, adj_y = false}> : (tensor<2x5x7x3xf32>, tensor<2x5x3x13xf32>) -> tensor<2x5x7x13xf32> // CHECK: %[[v2:.*]] = "tf.Transpose"(%[[v1]], %[[cst]]) : (tensor<2x5x7x13xf32>, tensor<4xi32>) -> tensor<2x7x5x13xf32> // CHECK: return %[[v2]] : tensor<2x7x5x13xf32> } @@ -183,7 +183,7 @@ func.func @einsum_fourdtransposeall(%arg0: tensor<2x5x7x3xf32>, %arg1: tensor<2x // CHECK-DAG: %[[cst_2:.*]] = arith.constant dense<[0, 1, 3, 2]> : tensor<4xi32> // CHECK: %[[v0:.*]] = "tf.Transpose"(%arg0, %[[cst]]) : (tensor<2x5x7x3xf32>, tensor<4xi32>) -> tensor<2x7x5x3xf32> // CHECK: %[[v1:.*]] = "tf.Transpose"(%arg1, %[[cst_1]]) : (tensor<2x11x7x3xf32>, tensor<4xi32>) -> tensor<2x7x3x11xf32> - // CHECK: %[[v2:.*]] = "tf.BatchMatMulV2"(%[[v0]], %[[v1]]) {adj_x = false, adj_y = false} : (tensor<2x7x5x3xf32>, tensor<2x7x3x11xf32>) -> tensor<2x7x5x11xf32> + // CHECK: %[[v2:.*]] = "tf.BatchMatMulV2"(%[[v0]], %[[v1]]) <{adj_x = false, adj_y = false}> : (tensor<2x7x5x3xf32>, tensor<2x7x3x11xf32>) -> tensor<2x7x5x11xf32> // CHECK: %[[v3:.*]] = "tf.Transpose"(%[[v2]], %[[cst_2]]) : (tensor<2x7x5x11xf32>, tensor<4xi32>) -> tensor<2x7x11x5xf32> // CHECK: return %[[v3]] : tensor<2x7x11x5xf32> } @@ -196,7 +196,7 @@ func.func @einsum_4d_1(%arg0: tensor<3x4x5x6xf32>, %arg1: tensor<3x7x5x6xf32>) - // CHECK-DAG: %[[cst_1:.*]] = arith.constant dense<[0, 2, 3, 1]> : tensor<4xi32> // CHECK: %[[v0:.*]] = "tf.Transpose"(%arg0, %[[cst:.*]]) : (tensor<3x4x5x6xf32>, tensor<4xi32>) -> tensor<3x5x4x6xf32> // CHECK: %[[v1:.*]] = "tf.Transpose"(%arg1, %[[cst_1]]) : (tensor<3x7x5x6xf32>, tensor<4xi32>) -> tensor<3x5x6x7xf32> - // CHECK: %[[v2:.*]] = "tf.BatchMatMulV2"(%[[v0]], %[[v1]]) {adj_x = false, adj_y = false} : (tensor<3x5x4x6xf32>, tensor<3x5x6x7xf32>) -> tensor<3x5x4x7xf32> + // CHECK: %[[v2:.*]] = "tf.BatchMatMulV2"(%[[v0]], %[[v1]]) <{adj_x = false, adj_y = false}> : (tensor<3x5x4x6xf32>, tensor<3x5x6x7xf32>) -> tensor<3x5x4x7xf32> // CHECK: return %[[v2]] : tensor<3x5x4x7xf32> } @@ -204,7 +204,7 @@ func.func @einsum_no_match(%arg0: tensor<4x5x6xf32>, %arg1: tensor<5xf32>) -> te %0 = "tf.Einsum"(%arg0, %arg1) {T = "tfdtype$DT_FLOAT", equation = "ijk,j->i"}: (tensor<4x5x6xf32>, tensor<5xf32>) -> tensor<4xf32> func.return %0 : tensor<4xf32> // CHECK-LABEL: einsum_no_match -// CHECK: %[[v0:.*]] = "tf.Einsum"(%arg0, %arg1) {T = "tfdtype$DT_FLOAT", equation = "ijk,j->i"} : (tensor<4x5x6xf32>, tensor<5xf32>) -> tensor<4xf32> +// CHECK: %[[v0:.*]] = "tf.Einsum"(%arg0, %arg1) <{equation = "ijk,j->i"}> {T = "tfdtype$DT_FLOAT"} : (tensor<4x5x6xf32>, tensor<5xf32>) -> tensor<4xf32> // CHECK: return %[[v0]] } @@ -212,7 +212,7 @@ func.func @einsum_illegal_no_match(%arg0: tensor<4x5xf32>, %arg1: tensor<5xf32>) %0 = "tf.Einsum"(%arg0, %arg1) {T = "tfdtype$DT_FLOAT", equation = "ij,?zw->kq->i"}: (tensor<4x5xf32>, tensor<5xf32>) -> tensor<4xf32> func.return %0 : tensor<4xf32> // CHECK-LABEL: einsum_illegal_no_match -// CHECK: %[[v0:.*]] = "tf.Einsum"(%arg0, %arg1) {T = "tfdtype$DT_FLOAT", equation = "ij,?zw->kq->i"} : (tensor<4x5xf32>, tensor<5xf32>) -> tensor<4xf32> +// CHECK: %[[v0:.*]] = "tf.Einsum"(%arg0, %arg1) <{equation = "ij,?zw->kq->i"}> {T = "tfdtype$DT_FLOAT"} : (tensor<4x5xf32>, tensor<5xf32>) -> tensor<4xf32> // CHECK: return %[[v0]] } @@ -223,7 +223,7 @@ func.func @batch_multilhs_einsum(%arg0: tensor<2x1x1x11xf32>, %arg1: tensor<2x11 // CHECK-DAG: %[[cst:.*]] = arith.constant dense<[2, 1, 11]> : tensor<3xi64> // CHECK-DAG: %[[cst_1:.*]] = arith.constant dense<[2, 1, 1, 2]> : tensor<4xi64> // CHECK: %[[v0:.*]] = "tf.Reshape"(%arg0, %[[cst]]) : (tensor<2x1x1x11xf32>, tensor<3xi64>) -> tensor<2x1x11xf32> -// CHECK: %[[v1:.*]] = "tf.BatchMatMulV2"(%[[v0]], %arg1) {adj_x = false, adj_y = false} : (tensor<2x1x11xf32>, tensor<2x11x2xf32>) -> tensor<2x1x2xf32> +// CHECK: %[[v1:.*]] = "tf.BatchMatMulV2"(%[[v0]], %arg1) <{adj_x = false, adj_y = false}> : (tensor<2x1x11xf32>, tensor<2x11x2xf32>) -> tensor<2x1x2xf32> // CHECK: %[[v2:.*]] = "tf.Reshape"(%[[v1]], %[[cst_1]]) : (tensor<2x1x2xf32>, tensor<4xi64>) -> tensor<2x1x1x2xf32> // CHECK: return %[[v2]] : tensor<2x1x1x2xf32> } @@ -236,14 +236,14 @@ func.func @einsum_with_runtime_outputshape1(%arg0 : tensor, %arg1 : // CHECK-DAG: %[[cst_0:.*]] = arith.constant dense<[-1, 36, 1, 32]> : tensor<4xi64> // CHECK-DAG: %[[cst_1:.*]] = arith.constant dense<[0, 1]> : tensor<2xi32> // CHECK-DAG: %[[cst_2:.*]] = arith.constant dense<2> : tensor<1xi32> -// CHECK-DAG: %[[cst_3:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor +// CHECK-DAG: %[[cst_3:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK: %[[v0:.*]] = "tf.Transpose"(%arg1, %cst) : (tensor, tensor<4xi32>) -> tensor // CHECK: %[[v1:.*]] = "tf.Reshape"(%arg0, %cst_0) : (tensor, tensor<4xi64>) -> tensor -// CHECK: %[[v2:.*]] = "tf.BatchMatMulV2"(%1, %0) {adj_x = false, adj_y = false} : (tensor, tensor) -> tensor +// CHECK: %[[v2:.*]] = "tf.BatchMatMulV2"(%1, %0) <{adj_x = false, adj_y = false}> : (tensor, tensor) -> tensor // CHECK: %[[v3:.*]] = "tf.Shape"(%arg0) : (tensor) -> tensor<3xi32> // CHECK: %[[v4:.*]] = "tf.Shape"(%arg1) : (tensor) -> tensor<4xi32> -// CHECK: %[[v5:.*]] = "tf.Gather"(%3, %cst_1) {validate_indices = true} : (tensor<3xi32>, tensor<2xi32>) -> tensor<2xi32> -// CHECK: %[[v6:.*]] = "tf.Gather"(%4, %cst_2) {validate_indices = true} : (tensor<4xi32>, tensor<1xi32>) -> tensor<1xi32> +// CHECK: %[[v5:.*]] = "tf.Gather"(%3, %cst_1) <{validate_indices = true}> : (tensor<3xi32>, tensor<2xi32>) -> tensor<2xi32> +// CHECK: %[[v6:.*]] = "tf.Gather"(%4, %cst_2) <{validate_indices = true}> : (tensor<4xi32>, tensor<1xi32>) -> tensor<1xi32> // CHECK: %[[v7:.*]] = "tf.Concat"(%cst_3, %5, %6) : (tensor, tensor<2xi32>, tensor<1xi32>) -> tensor<3xi32> // CHECK: %[[v8:.*]] = "tf.Reshape"(%2, %7) : (tensor, tensor<3xi32>) -> tensor // CHECK: return %[[v8]] : tensor @@ -254,13 +254,13 @@ func.func @einsum_with_runtime_outputshape2(%arg0 : tensor, %arg1 func.return %0 : tensor // CHECK-LABEL: einsum_with_runtime_outputshape2 // CHECK-DAG: %[[cst:.*]] = arith.constant dense<1024> : tensor<2xi64> -// CHECK-DAG: %[[cst_0:.*]] = "tf.Const"() {value = dense<[8, 128]> : tensor<2xi32>} : () -> tensor<2xi32> +// CHECK-DAG: %[[cst_0:.*]] = "tf.Const"() <{value = dense<[8, 128]> : tensor<2xi32>}> : () -> tensor<2xi32> // CHECK-DAG: %[[cst_1:.*]] = arith.constant dense<[0, 1]> : tensor<2xi32> -// CHECK-DAG: %[[cst_2:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor +// CHECK-DAG: %[[cst_2:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK: %[[v0:.*]] = "tf.Reshape"(%arg1, %cst) : (tensor<1024x8x128xf32>, tensor<2xi64>) -> tensor<1024x1024xf32> -// CHECK: %[[v1:.*]] = "tf.BatchMatMulV2"(%arg0, %0) {adj_x = false, adj_y = false} : (tensor, tensor<1024x1024xf32>) -> tensor +// CHECK: %[[v1:.*]] = "tf.BatchMatMulV2"(%arg0, %0) <{adj_x = false, adj_y = false}> : (tensor, tensor<1024x1024xf32>) -> tensor // CHECK: %[[v2:.*]] = "tf.Shape"(%arg0) : (tensor) -> tensor<3xi32> -// CHECK: %[[v3:.*]] = "tf.Gather"(%2, %cst_1) {validate_indices = true} : (tensor<3xi32>, tensor<2xi32>) -> tensor<2xi32> +// CHECK: %[[v3:.*]] = "tf.Gather"(%2, %cst_1) <{validate_indices = true}> : (tensor<3xi32>, tensor<2xi32>) -> tensor<2xi32> // CHECK: %[[v4:.*]] = "tf.Concat"(%cst_2, %3, %cst_0) : (tensor, tensor<2xi32>, tensor<2xi32>) -> tensor<4xi32> // CHECK: %[[v5:.*]] = "tf.Reshape"(%1, %4) : (tensor, tensor<4xi32>) -> tensor // CHECK: return %[[v5]] : tensor @@ -275,7 +275,7 @@ func.func @einsum_with_runtime_shape1(%arg0 : tensor, %arg1 : tensor // CHECK: %[[v0:.*]] = "tf.Shape"(%arg0) : (tensor) -> tensor<3xi32> // CHECK: %[[v1:.*]] = "tf.UnsortedSegmentProd"(%0, %cst, %cst_0) : (tensor<3xi32>, tensor<3xi32>, tensor) -> tensor<4xi32> // CHECK: %[[v2:.*]] = "tf.Reshape"(%arg0, %1) : (tensor, tensor<4xi32>) -> tensor -// CHECK: %[[v3:.*]] = "tf.BatchMatMulV2"(%2, %arg1) {adj_x = false, adj_y = false} : (tensor, tensor) -> tensor +// CHECK: %[[v3:.*]] = "tf.BatchMatMulV2"(%2, %arg1) <{adj_x = false, adj_y = false}> : (tensor, tensor) -> tensor // CHECK: %[[v4:.*]] = "tf.Reshape"(%3, %cst_1) : (tensor, tensor<3xi64>) -> tensor // CHECK: return %[[v4]] : tensor } @@ -286,14 +286,14 @@ func.func @einsum_with_runtime_shape2(%arg0 : tensor, %arg1 : tens // CHECK-LABEL: einsum_with_runtime_shape2 // CHECK-DAG: %[[cst:.*]] = arith.constant dense<[1, 2, 0]> : tensor<3xi32> // CHECK-DAG: %[[cst_0:.*]] = arith.constant dense<[0, 1, 2, 2]> : tensor<4xi32> -// CHECK-DAG: %[[cst_1:.*]] = "tf.Const"() {value = dense<3> : tensor} : () -> tensor +// CHECK-DAG: %[[cst_1:.*]] = "tf.Const"() <{value = dense<3> : tensor}> : () -> tensor // CHECK-DAG: %[[cst_2:.*]] = arith.constant dense<[512, 8]> : tensor<2xi64> // CHECK: %[[v0:.*]] = "tf.Transpose"(%arg1, %cst) : (tensor<8x8x64xf32>, tensor<3xi32>) -> tensor<8x64x8xf32> // CHECK: %[[v1:.*]] = "tf.Shape"(%arg0) : (tensor) -> tensor<4xi32> // CHECK: %[[v2:.*]] = "tf.UnsortedSegmentProd"(%1, %cst_0, %cst_1) : (tensor<4xi32>, tensor<4xi32>, tensor) -> tensor<3xi32> // CHECK: %[[v3:.*]] = "tf.Reshape"(%arg0, %2) : (tensor, tensor<3xi32>) -> tensor // CHECK: %[[v4:.*]] = "tf.Reshape"(%0, %cst_2) : (tensor<8x64x8xf32>, tensor<2xi64>) -> tensor<512x8xf32> -// CHECK: %[[v5:.*]] = "tf.BatchMatMulV2"(%3, %4) {adj_x = false, adj_y = false} : (tensor, tensor<512x8xf32>) -> tensor +// CHECK: %[[v5:.*]] = "tf.BatchMatMulV2"(%3, %4) <{adj_x = false, adj_y = false}> : (tensor, tensor<512x8xf32>) -> tensor // CHECK: return %[[v5]] : tensor } @@ -305,7 +305,7 @@ func.func @einsum_no_reshape(%arg0 : tensor, %arg1 : tensor<1x?x8 // CHECK-DAG: %[[cst_0:.*]] = arith.constant dense<[0, 2, 3, 1]> : tensor<4xi32> // CHECK: %[[v0:.*]] = "tf.Transpose"(%arg0, %cst) : (tensor, tensor<4xi32>) -> tensor // CHECK: %[[v1:.*]] = "tf.Transpose"(%arg1, %cst_0) : (tensor<1x?x8x128xf32>, tensor<4xi32>) -> tensor<1x8x128x?xf32> -// CHECK: %[[v3:.*]] = "tf.BatchMatMulV2"(%0, %1) {adj_x = false, adj_y = false} : (tensor, tensor<1x8x128x?xf32>) -> tensor +// CHECK: %[[v3:.*]] = "tf.BatchMatMulV2"(%0, %1) <{adj_x = false, adj_y = false}> : (tensor, tensor<1x8x128x?xf32>) -> tensor // CHECK: %[[v4:.*]] = "tf.Transpose"(%2, %cst) : (tensor, tensor<4xi32>) -> tensor // CHECK: return %[[v4]] : tensor } @@ -314,7 +314,7 @@ func.func @einsum_ellipsis(%arg0: tensor<1x512x128xf32>, %arg1: tensor<128x256xf %0 = "tf.Einsum"(%arg0, %arg1) {device = "", equation = "...x,xy->...y"} : (tensor<1x512x128xf32>, tensor<128x256xf32>) -> tensor<1x512x256xf32> func.return %0 : tensor<1x512x256xf32> // CHECK-LABEL: einsum_ellipsis -// CHECK: "tf.BatchMatMulV2"(%arg0, %arg1) {adj_x = false, adj_y = false} : (tensor<1x512x128xf32>, tensor<128x256xf32>) -> tensor<1x512x256xf32> +// CHECK: "tf.BatchMatMulV2"(%arg0, %arg1) <{adj_x = false, adj_y = false}> : (tensor<1x512x128xf32>, tensor<128x256xf32>) -> tensor<1x512x256xf32> } func.func @einsum_ellipsis_in_both_sides(%arg0: tensor<1x11x19xf32>, %arg1: tensor<7x11x13x19xf32>) -> tensor<7x11x13xf32> { @@ -326,7 +326,7 @@ func.func @einsum_ellipsis_in_both_sides(%arg0: tensor<1x11x19xf32>, %arg1: tens // CHECK-DAG: %[[cst_2:.*]] = arith.constant dense<[7, 11, 13]> : tensor<3xi64> // CHECK: %[[v0:.*]] = "tf.Transpose"(%arg1, %[[cst]]) : (tensor<7x11x13x19xf32>, tensor<4xi32>) -> tensor<7x11x19x13xf32> // CHECK: %[[v1:.*]] = "tf.Reshape"(%arg0, %[[cst_1]]) : (tensor<1x11x19xf32>, tensor<4xi64>) -> tensor<1x11x1x19xf32> - // CHECK: %[[v2:.*]] = "tf.BatchMatMulV2"(%[[v1]], %[[v0]]) {adj_x = false, adj_y = false} : (tensor<1x11x1x19xf32>, tensor<7x11x19x13xf32>) -> tensor<7x11x1x13xf32> + // CHECK: %[[v2:.*]] = "tf.BatchMatMulV2"(%[[v1]], %[[v0]]) <{adj_x = false, adj_y = false}> : (tensor<1x11x1x19xf32>, tensor<7x11x19x13xf32>) -> tensor<7x11x1x13xf32> // CHECK: %[[v3:.*]] = "tf.Reshape"(%[[v2]], %[[cst_2]]) : (tensor<7x11x1x13xf32>, tensor<3xi64>) -> tensor<7x11x13xf32> // CHECK: return %[[v3]] : tensor<7x11x13xf32> } @@ -338,7 +338,7 @@ func.func @einsum_ellipsis_with_broadcast(%arg0: tensor<5x4x3xf32>, %arg1: tenso // CHECK-DAG: %[[cst:.*]] = arith.constant dense<[2, 0, 1]> : tensor<3xi32> // CHECK-DAG: %[[cst_1:.*]] = arith.constant dense<[1, 2, 0]> : tensor<3xi32> // CHECK: %[[v0:.*]] = "tf.Transpose"(%arg1, %[[cst]]) : (tensor<3x2x1xf32>, tensor<3xi32>) -> tensor<1x3x2xf32> - // CHECK: %[[v1:.*]] = "tf.BatchMatMulV2"(%arg0, %[[v0]]) {adj_x = false, adj_y = false} : (tensor<5x4x3xf32>, tensor<1x3x2xf32>) -> tensor<5x4x2xf32> + // CHECK: %[[v1:.*]] = "tf.BatchMatMulV2"(%arg0, %[[v0]]) <{adj_x = false, adj_y = false}> : (tensor<5x4x3xf32>, tensor<1x3x2xf32>) -> tensor<5x4x2xf32> // CHECK: %[[v2:.*]] = "tf.Transpose"(%[[v1]], %[[cst_1]]) : (tensor<5x4x2xf32>, tensor<3xi32>) -> tensor<4x2x5xf32> // CHECK: return %[[v2]] : tensor<4x2x5xf32> } diff --git a/tensorflow/compiler/mlir/tensorflow/tests/embedding_pipelining.mlir b/tensorflow/compiler/mlir/tensorflow/tests/embedding_pipelining.mlir index fdb64b900e7a8d..8aa1cf650550f4 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/embedding_pipelining.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/embedding_pipelining.mlir @@ -17,7 +17,7 @@ module { // CHECK: {{.*StatefulPartitionedCall.* f = @non_tpu.*}} // CHECK: {{.*StatefulPartitionedCall.* f = @start_step_1.*}} // CHECK: {{.*StatefulPartitionedCall.* f = @while_cond.*}} - // CHECK: {{.*tf.While.* body = @new_while_body.* cond = @new_while_cond.*}} + // CHECK: {{.*tf.While.* <{body = @new_while_body.* cond = @new_while_cond.*}} // CHECK: {{.*StatefulPartitionedCall.* f = @finish_step_nm2.*}} // CHECK: {{.*StatefulPartitionedCall.* f = @finish_step_nm1.*}} // CHECK: return @@ -73,7 +73,7 @@ module { // CHECK: {{.*StatefulPartitionedCall.* f = @non_tpu.*}} // CHECK: {{.*StatefulPartitionedCall.* f = @start_step_1.*}} // CHECK: {{.*StatefulPartitionedCall.* f = @while_cond.*}} - // CHECK: {{.*tf.While.* body = @new_while_body.* cond = @new_while_cond.*}} + // CHECK: {{.*tf.While.* <{body = @new_while_body.* cond = @new_while_cond.*}} // CHECK: {{.*StatefulPartitionedCall.* f = @finish_step_nm2.*}} // CHECK: {{.*StatefulPartitionedCall.* f = @finish_step_nm1.*}} // CHECK: return @@ -112,7 +112,7 @@ module { func.func private @while_body(%arg0: tensor) -> (tensor) { // The pipelining control flow and supporting functions stay the same as the training version above. // The order of these functions is also significant. - // CHECK: {{.*tf.While.* body = @new_while_body.* cond = @new_while_cond.* parallel_iterations = 3}} + // CHECK: {{.*tf.While.* <{body = @new_while_body.* cond = @new_while_cond.* parallel_iterations = 3}} // CHECK: return // metadata ops "tf.TPUReplicateMetadata"() {_has_manual_control_dependencies = true, _replication_info = "repl_info", num_replicas = 1 : i64} : () -> () diff --git a/tensorflow/compiler/mlir/tensorflow/tests/extract_outside_compilation.mlir b/tensorflow/compiler/mlir/tensorflow/tests/extract_outside_compilation.mlir index 488c98af16f55f..87acd459ed1350 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/extract_outside_compilation.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/extract_outside_compilation.mlir @@ -484,7 +484,7 @@ module attributes {tf.versions = {producer = 888 : i32}, tf.devices = ["/job:wor // CHECK: %[[RECV_OUTPUT_2:[0-9]*]] = "tf._XlaRecvAtHostV2"(%[[PROGRAM_OUTPUT]], %[[DEVICE_ORDINAL]]) // CHECK-SAME: key = "host_compute_channel_1_args" // CHECK: "tf.D"(%[[RECV_OUTPUT_2]]) - // CHECK: "tf_device.cluster" + // CHECK: "tf_device.cluster"() // CHECK: %[[A_OUTPUT:[0-9]*]] = "tf.A" // CHECK: "tf._XlaHostComputeMlir"(%[[A_OUTPUT]]) // CHECK-SAME: send_key = "host_compute_channel_0_args" @@ -581,14 +581,14 @@ module attributes {tf.versions = {producer = 888 : i32}, tf.devices = ["/job:wor // CHECK-NEXT: %[[PREDICATE_RECV_OUTPUT:[0-9]*]] = "tf._XlaRecvAtHostV2"(%[[PROGRAM_OUTPUT]], %[[DEVICE_ORDINAL]]) // CHECK-SAME: key = "if_predicate_channel_1" // CHECK-NEXT: tf.IfRegion"(%[[PREDICATE_RECV_OUTPUT]]) + // CHECK: _else_func_name = "test_else_name" + // CHECK-SAME: _then_func_name = "test_then_name" + // CHECK-SAME: is_stateless = false // CHECK-NEXT: %[[ARG_RECV_OUTPUT:[0-9]*]]:2 = "tf._XlaRecvAtHostV2"(%[[PROGRAM_OUTPUT]], %[[DEVICE_ORDINAL]]) // CHECK-SAME: key = "host_compute_channel_0_args" // CHECK: "tf.D"(%[[ARG_RECV_OUTPUT]]#0, %[[ARG_RECV_OUTPUT]]#1) // CHECK-NOT: "tf._XlaSendFromHostV2"(%[[PROGRAM_OUTPUT]], %[[DEVICE_ORDINAL]]) // CHECK: "tf.Yield"() : () -> () - // CHECK: _else_func_name = "test_else_name" - // CHECK-SAME: _then_func_name = "test_then_name" - // CHECK-SAME: is_stateless = false // CHECK: "tf_device.cluster" // CHECK: %[[A_OUTPUT:[0-9]*]] = "tf.A" // CHECK: %[[B_OUTPUT:[0-9]*]] = "tf.B" @@ -596,11 +596,11 @@ module attributes {tf.versions = {producer = 888 : i32}, tf.devices = ["/job:wor // CHECK: "tf._XlaHostComputeMlir" // CHECK-SAME: key = "if_predicate_channel_1" // CHECK-NEXT: tf.IfRegion"(%[[G_OUTPUT]]) + // CHECK: is_stateless = false // CHECK: "tf._XlaHostComputeMlir"(%[[B_OUTPUT]], %[[A_OUTPUT]]) // CHECK-SAME: recv_key = "host_compute_channel_0_retvals" // CHECK-SAME: send_key = "host_compute_channel_0_args" // CHECK-NEXT: "tf.Yield"() : () -> () - // CHECK: is_stateless = false %1:2 = tf_device.replicate([%0, %arg0] as %ri_0: tensor<2xi32>) {n = 2 : i32} { %2 = "tf_device.cluster"() ({ %3 = "tf.A"() : () -> (tensor<2xi32>) @@ -637,13 +637,13 @@ module attributes {tf.versions = {producer = 888 : i32}, tf.devices = ["/job:wor // CHECK-NEXT: %[[PREDICATE_RECV_OUTPUT:[0-9]*]] = "tf._XlaRecvAtHostV2"(%[[PROGRAM_OUTPUT]], %[[DEVICE_ORDINAL]]) // CHECK-SAME: key = "if_predicate_channel_1" // CHECK-NEXT: tf.IfRegion"(%[[PREDICATE_RECV_OUTPUT]]) + // CHECK: _else_func_name = "test_else_name" + // CHECK-SAME _then_func_name = "test_then_name" // CHECK-NEXT: %[[ARG_RECV_OUTPUT:[0-9]*]]:2 = "tf._XlaRecvAtHostV2"(%[[PROGRAM_OUTPUT]], %[[DEVICE_ORDINAL]]) // CHECK-SAME: key = "host_compute_channel_0_args" // CHECK: "tf.D"(%[[ARG_RECV_OUTPUT]]#0, %[[ARG_RECV_OUTPUT]]#1) // CHECK-NOT: "tf._XlaSendFromHostV2"(%[[PROGRAM_OUTPUT]], %[[DEVICE_ORDINAL]]) // CHECK: "tf.Yield"() : () -> () - // CHECK: _else_func_name = "test_else_name" - // CHECK-SAME _then_func_name = "test_then_name" // CHECK: "tf_device.cluster" // CHECK: %[[A_OUTPUT:[0-9]*]] = "tf.A" // CHECK: %[[B_OUTPUT:[0-9]*]] = "tf.B" @@ -1044,6 +1044,7 @@ module attributes {tf.versions = {producer = 888 : i32}, tf.devices = ["/job:wor // CHECK-DAG: %[[PROGRAM_OUTPUT:.+]] = "tf._XlaCompileMlirPlaceholderProgramKey" // CHECK-DAG: %[[DEVICE_ORDINAL:.+]] = "tf._TPUDeviceOrdinalPlaceholder" // CHECK-NEXT: tf.WhileRegion" + // CHECK: is_stateless = false // CHECK-NEXT: %[[COND_RECV_OUTPUT:[0-9]*]] = "tf._XlaRecvAtHostV2"(%[[PROGRAM_OUTPUT]], %[[DEVICE_ORDINAL]]) // CHECK-SAME: key = "while_condition_channel_0" // CHECK: "tf.Yield"(%[[COND_RECV_OUTPUT]]) @@ -1051,19 +1052,18 @@ module attributes {tf.versions = {producer = 888 : i32}, tf.devices = ["/job:wor // CHECK: %[[D_OUTPUT:[0-9]*]] = "tf.D" // CHECK: "tf._XlaSendFromHostV2"(%[[D_OUTPUT]], %[[PROGRAM_OUTPUT]], %[[DEVICE_ORDINAL]]) // CHECK-NEXT: "tf.Yield" - // CHECK: is_stateless = false // CHECK: "tf_device.cluster" // CHECK: %[[A_OUTPUT:[0-9]*]] = "tf.A" // CHECK: %[[B_OUTPUT:[0-9]*]] = "tf.B" // CHECK: %[[G_OUTPUT:[0-9]*]] = "tf.G" // CHECK-NEXT: tf.WhileRegion"(%[[B_OUTPUT]], %[[A_OUTPUT]]) + // CHECK: is_stateless = false // CHECK: %[[H_OUTPUT:[0-9]*]] = "tf.H" // CHECK-NEXT: "tf.XlaSendToHost"(%[[H_OUTPUT]]) // CHECK-NEXT: "tf.Yield"(%[[H_OUTPUT]]) // CHECK: %[[C_OUTPUT:[0-9]*]] = "tf.C" // CHECK-NEXT: %[[HOST_COMPUTE_OUTPUT:[0-9]*]] = "tf._XlaHostComputeMlir" // CHECK-NEXT: "tf.Yield"(%[[C_OUTPUT]], %[[HOST_COMPUTE_OUTPUT]]) - // CHECK: is_stateless = false %1:2 = tf_device.replicate([%0, %arg0] as %ri_0: tensor<2xi32>) {n = 2 : i32} { %2 = "tf_device.cluster"() ({ %3 = "tf.A"() : () -> (tensor<3xf32>) @@ -2118,34 +2118,34 @@ module attributes {tf.devices = {"/job:localhost/replica:0/task:0/device:CPU:0", // CHECK: "tf_device.launch" // CHECK: %[[PROGRAM0:.+]] = "tf._XlaCompileMlirPlaceholderProgramKey" // CHECK: %[[RECV0:.+]] = "tf._XlaRecvAtHost"(%[[PROGRAM0]]) - // CHECK-SAME: _xla_has_host_transfer = true // CHECK-SAME: device_ordinal = 0 // CHECK-SAME: key = "host_compute_channel_0_args" + // CHECK-SAME: _xla_has_host_transfer = true // CHECK: %[[B0:.+]] = "tf.OpB"(%[[RECV0]]) : (tensor<2x2xi64>) -> tensor<2x2xi64> // CHECK: "tf._XlaSendFromHost"(%[[B0]], %[[PROGRAM0]]) - // CHECK-SAME: _xla_has_host_transfer = true // CHECK-SAME: device_ordinal = 0 // CHECK-SAME: key = "host_compute_channel_0_retvals" + // CHECK-SAME: _xla_has_host_transfer = true // CHECK: }, { // CHECK: %[[PROGRAM1:.+]] = "tf._XlaCompileMlirPlaceholderProgramKey" // CHECK: %[[RECV1:.+]] = "tf._XlaRecvAtHost"(%[[PROGRAM1]]) - // CHECK-SAME: _xla_has_host_transfer = true // CHECK-SAME: device_ordinal = 1 // CHECK-SAME: key = "host_compute_channel_0_args" + // CHECK-SAME: _xla_has_host_transfer = true // CHECK: %[[B1:.+]] = "tf.OpB"(%[[RECV1]]) : (tensor<2x2xi64>) -> tensor<2x2xi64> // CHECK: "tf._XlaSendFromHost"(%[[B1]], %[[PROGRAM1]]) - // CHECK-SAME: _xla_has_host_transfer = true // CHECK-SAME: device_ordinal = 1 // CHECK-SAME: key = "host_compute_channel_0_retvals" + // CHECK-SAME: _xla_has_host_transfer = true // CHECK: }, { // CHECK: "tf_device.cluster" // CHECK: %[[A:.+]] = "tf.OpA" - // CHECK: %[[A_SHARD:.+]] = "tf.XlaSpmdFullToShardShape"(%[[A]]) {dim = -1 : i64, manual_sharding = "\08\03\1A\02\02\01\22\02\00\01", unspecified_dims = []} : (tensor<2x2xi64>) -> tensor<1x2xi64> + // CHECK: %[[A_SHARD:.+]] = "tf.XlaSpmdFullToShardShape"(%[[A]]) <{dim = -1 : i64, manual_sharding = "\08\03\1A\02\02\01\22\02\00\01", unspecified_dims = []}> : (tensor<2x2xi64>) -> tensor<1x2xi64> // CHECK: %[[B:.+]] = "tf._XlaHostComputeMlir"(%[[A_SHARD]]) // CHECK-SAME: manual_sharding = true // CHECK-SAME: recv_key = "host_compute_channel_0_retvals" // CHECK-SAME: send_key = "host_compute_channel_0_args" - // CHECK: %[[B_FULL:.+]] = "tf.XlaSpmdShardToFullShape"(%[[B]]) {dim = -1 : i64, full_shape = #tf_type.shape<2x2>, manual_sharding = "\08\03\1A\02\02\01\22\02\00\01", unspecified_dims = []} : (tensor<1x2xi64>) -> tensor<2x2xi64> + // CHECK: %[[B_FULL:.+]] = "tf.XlaSpmdShardToFullShape"(%[[B]]) <{dim = -1 : i64, full_shape = #tf_type.shape<2x2>, manual_sharding = "\08\03\1A\02\02\01\22\02\00\01", unspecified_dims = []}> : (tensor<1x2xi64>) -> tensor<2x2xi64> // CHECK: "tf.OpC"(%[[B_FULL]]) "tf_device.cluster"() ({ %0 = "tf.OpA"() {_XlaSharding = "\08\03\1A\02\02\01\22\02\00\01"} : () -> tensor<2x2xi64> @@ -2178,32 +2178,32 @@ module attributes {tf.devices = {"/job:localhost/replica:0/task:0/device:CPU:0", // CHECK: %[[DEVICE0:.+]] = "tf._TPUDeviceOrdinalPlaceholder" // CHECK-SAME: logical_core = 0 // CHECK: %[[RECV0:.+]] = "tf._XlaRecvAtHostV2"(%[[PROGRAM0]], %[[DEVICE0]]) - // CHECK-SAME: _xla_has_host_transfer = true // CHECK-SAME: key = "host_compute_channel_0_args" + // CHECK-SAME: _xla_has_host_transfer = true // CHECK: %[[B0:.+]] = "tf.OpB"(%[[RECV0]]) : (tensor<2x2xi64>) -> tensor<2x2xi64> // CHECK: "tf._XlaSendFromHostV2"(%[[B0]], %[[PROGRAM0]], %[[DEVICE0]]) - // CHECK-SAME: _xla_has_host_transfer = true // CHECK-SAME: key = "host_compute_channel_0_retvals" + // CHECK-SAME: _xla_has_host_transfer = true // CHECK: }, { // CHECK: %[[PROGRAM1:.+]] = "tf._XlaCompileMlirPlaceholderProgramKey" // CHECK: %[[DEVICE1:.+]] = "tf._TPUDeviceOrdinalPlaceholder" // CHECK-SAME: logical_core = 1 // CHECK: %[[RECV1:.+]] = "tf._XlaRecvAtHostV2"(%[[PROGRAM1]], %[[DEVICE1]]) - // CHECK-SAME: _xla_has_host_transfer = true // CHECK-SAME: key = "host_compute_channel_0_args" + // CHECK-SAME: _xla_has_host_transfer = true // CHECK: %[[B1:.+]] = "tf.OpB"(%[[RECV1]]) : (tensor<2x2xi64>) -> tensor<2x2xi64> // CHECK: "tf._XlaSendFromHostV2"(%[[B1]], %[[PROGRAM1]], %[[DEVICE1]]) - // CHECK-SAME: _xla_has_host_transfer = true // CHECK-SAME: key = "host_compute_channel_0_retvals" + // CHECK-SAME: _xla_has_host_transfer = true // CHECK: }, { // CHECK: "tf_device.cluster" // CHECK: %[[A:.+]] = "tf.OpA" - // CHECK: %[[A_SHARD:.+]] = "tf.XlaSpmdFullToShardShape"(%[[A]]) {dim = -1 : i64, manual_sharding = "\08\03\1A\02\02\01\22\02\00\01", unspecified_dims = []} : (tensor<2x2xi64>) -> tensor<1x2xi64> + // CHECK: %[[A_SHARD:.+]] = "tf.XlaSpmdFullToShardShape"(%[[A]]) <{dim = -1 : i64, manual_sharding = "\08\03\1A\02\02\01\22\02\00\01", unspecified_dims = []}> : (tensor<2x2xi64>) -> tensor<1x2xi64> // CHECK: %[[B:.+]] = "tf._XlaHostComputeMlir"(%[[A_SHARD]]) // CHECK-SAME: manual_sharding = true // CHECK-SAME: recv_key = "host_compute_channel_0_retvals" // CHECK-SAME: send_key = "host_compute_channel_0_args" - // CHECK: %[[B_FULL:.+]] = "tf.XlaSpmdShardToFullShape"(%[[B]]) {dim = -1 : i64, full_shape = #tf_type.shape<2x2>, manual_sharding = "\08\03\1A\02\02\01\22\02\00\01", unspecified_dims = []} : (tensor<1x2xi64>) -> tensor<2x2xi64> + // CHECK: %[[B_FULL:.+]] = "tf.XlaSpmdShardToFullShape"(%[[B]]) <{dim = -1 : i64, full_shape = #tf_type.shape<2x2>, manual_sharding = "\08\03\1A\02\02\01\22\02\00\01", unspecified_dims = []}> : (tensor<1x2xi64>) -> tensor<2x2xi64> // CHECK: "tf.OpC"(%[[B_FULL]]) tf_device.replicate() {n = 4 : i32} { "tf_device.cluster"() ({ diff --git a/tensorflow/compiler/mlir/tensorflow/tests/fold-broadcast.mlir b/tensorflow/compiler/mlir/tensorflow/tests/fold-broadcast.mlir index 9e7b5b23eb48dd..5535badb06a06c 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/fold-broadcast.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/fold-broadcast.mlir @@ -48,7 +48,7 @@ func.func @broadcast_eq(%arg0: tensor<5x7xf32>, %arg1: tensor<7xf32>) -> tensor< %0 = "tf.BroadcastTo"(%arg1, %cst) : (tensor<7xf32>, tensor<2xi32>) -> tensor<5x7xf32> %1 = "tf.Equal"(%arg0, %0) {incompatible_shape_error = true} : (tensor<5x7xf32>, tensor<5x7xf32>) -> tensor<5x7xi1> func.return %1 : tensor<5x7xi1> - // CHECK: %[[V0:.*]] = "tf.Equal"(%arg0, %arg1) {incompatible_shape_error = true} : (tensor<5x7xf32>, tensor<7xf32>) -> tensor<5x7xi1> + // CHECK: %[[V0:.*]] = "tf.Equal"(%arg0, %arg1) <{incompatible_shape_error = true}> : (tensor<5x7xf32>, tensor<7xf32>) -> tensor<5x7xi1> // CHECK: %[[V0]] : tensor<5x7xi1> } @@ -58,7 +58,7 @@ func.func @broadcast_neq(%arg0: tensor<5x7xf32>, %arg1: tensor<7xf32>) -> tensor %0 = "tf.BroadcastTo"(%arg1, %cst) : (tensor<7xf32>, tensor<2xi32>) -> tensor<5x7xf32> %1 = "tf.NotEqual"(%arg0, %0) {incompatible_shape_error = true} : (tensor<5x7xf32>, tensor<5x7xf32>) -> tensor<5x7xi1> func.return %1 : tensor<5x7xi1> - // CHECK: %[[V0:.*]] = "tf.NotEqual"(%arg0, %arg1) {incompatible_shape_error = true} : (tensor<5x7xf32>, tensor<7xf32>) -> tensor<5x7xi1> + // CHECK: %[[V0:.*]] = "tf.NotEqual"(%arg0, %arg1) <{incompatible_shape_error = true}> : (tensor<5x7xf32>, tensor<7xf32>) -> tensor<5x7xi1> // CHECK: %[[V0]] : tensor<5x7xi1> } @@ -79,7 +79,7 @@ func.func @broadcast_batch_matmul_v2_rhs(%arg0: tensor<17x17x17xf32>, %arg1: ten %0 = "tf.BroadcastTo"(%arg1, %cst) : (tensor<17x24xf32>, tensor<3xi64>) -> tensor<17x17x24xf32> %1 = "tf.BatchMatMulV2"(%arg0, %0) {adj_x = false, adj_y = false} : (tensor<17x17x17xf32>, tensor<17x17x24xf32>) -> tensor<17x17x24xf32> func.return %1 : tensor<17x17x24xf32> - // CHECK: %[[V0:.*]] = "tf.BatchMatMulV2"(%arg0, %arg1) {adj_x = false, adj_y = false} : (tensor<17x17x17xf32>, tensor<17x24xf32>) -> tensor<17x17x24xf32> + // CHECK: %[[V0:.*]] = "tf.BatchMatMulV2"(%arg0, %arg1) <{adj_x = false, adj_y = false}> : (tensor<17x17x17xf32>, tensor<17x24xf32>) -> tensor<17x17x24xf32> // CHECK: %[[V0]] : tensor<17x17x24xf32> } @@ -89,7 +89,7 @@ func.func @broadcast_batch_matmul_v2_lhs(%arg0: tensor<17x17xf32>, %arg1: tensor %0 = "tf.BroadcastTo"(%arg0, %cst) : (tensor<17x17xf32>, tensor<3xi64>) -> tensor<17x17x17xf32> %1 = "tf.BatchMatMulV2"(%0, %arg1) {adj_x = false, adj_y = false} : (tensor<17x17x17xf32>, tensor<17x17x24xf32>) -> tensor<17x17x24xf32> func.return %1 : tensor<17x17x24xf32> - // CHECK: %[[V0:.*]] = "tf.BatchMatMulV2"(%arg0, %arg1) {adj_x = false, adj_y = false} : (tensor<17x17xf32>, tensor<17x17x24xf32>) -> tensor<17x17x24xf32> + // CHECK: %[[V0:.*]] = "tf.BatchMatMulV2"(%arg0, %arg1) <{adj_x = false, adj_y = false}> : (tensor<17x17xf32>, tensor<17x17x24xf32>) -> tensor<17x17x24xf32> // CHECK: %[[V0]] : tensor<17x17x24xf32> } @@ -108,6 +108,6 @@ func.func @broadcast_splat_operand() -> tensor<5x5xi64> { %cst = arith.constant dense<5> : tensor<2xi64> %0 = "tf.BroadcastTo"(%cst, %cst) : (tensor<2xi64>, tensor<2xi64>) -> tensor<5x5xi64> func.return %0 : tensor<5x5xi64> - // CHECK: %[[V0:.*]] = "tf.Const"() {value = dense<5> : tensor<5x5xi64>} : () -> tensor<5x5xi64> + // CHECK: %[[V0:.*]] = "tf.Const"() <{value = dense<5> : tensor<5x5xi64>}> : () -> tensor<5x5xi64> // CHECK: %[[V0]] : tensor<5x5xi64> } diff --git a/tensorflow/compiler/mlir/tensorflow/tests/freeze_variables.mlir b/tensorflow/compiler/mlir/tensorflow/tests/freeze_variables.mlir index 7c18a389de94b5..a458a20d49e510 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/freeze_variables.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/freeze_variables.mlir @@ -464,7 +464,7 @@ module { func.return %arg0, %0 : tensor, tensor<0xf32> } // CHECK: func.func private @f_batch_callee(%[[ARG_0:.*]]: tensor) -> (tensor, tensor<0xf32>) - // CHECK-DAG: %[[CST_0:.*]] = "tf.Const"() {{{.*value = dense<> : tensor<0xf32>.*}}} : () -> tensor<0xf32> + // CHECK-DAG: %[[CST_0:.*]] = "tf.Const"() <{value = dense<> : tensor<0xf32>}> : () -> tensor<0xf32> // CHECK: return %[[ARG_0]], %[[CST_0]] : tensor, tensor<0xf32> func.func @f(%arg: tensor<1xf32>) -> (tensor<*xf32>, tensor<*xf32>) { @@ -474,6 +474,6 @@ module { } // CHECK: func.func @f(%[[ARG_1:.*]]: tensor<1xf32>) // Make sure that `operandSegmentSizes` attribute is also updated. - // CHECK-NEXT: %[[BATCH_FUNC:.*]]:2 = "tf.BatchFunction"(%[[ARG_1]]) {{{.*operandSegmentSizes = array.*}}} : (tensor<1xf32>) -> (tensor<*xf32>, tensor<*xf32>) + // CHECK-NEXT: %[[BATCH_FUNC:.*]]:2 = "tf.BatchFunction"(%[[ARG_1]]) <{{{.*operandSegmentSizes = array.*}}}> : (tensor<1xf32>) -> (tensor<*xf32>, tensor<*xf32>) // CHECK: return %[[BATCH_FUNC]]#0, %[[BATCH_FUNC]]#1 : tensor<*xf32>, tensor<*xf32> } diff --git a/tensorflow/compiler/mlir/tensorflow/tests/functional-control-flow-to-cfg.mlir b/tensorflow/compiler/mlir/tensorflow/tests/functional-control-flow-to-cfg.mlir index 4339cd725617fb..8deedc067f20fa 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/functional-control-flow-to-cfg.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/functional-control-flow-to-cfg.mlir @@ -68,14 +68,14 @@ func.func @testIfCasts(%arg0: tensor, %arg1: tensor // CHECK: cf.cond_br [[PRED]], ^bb1, ^bb2 // CHECK: ^bb1: -// CHECK: [[CAST0:%.+]] = "tf.Cast"(%arg1) {Truncate = false} : (tensor>>) -> tensor +// CHECK: [[CAST0:%.+]] = "tf.Cast"(%arg1) <{Truncate = false}> : (tensor>>) -> tensor // CHECK: [[THEN:%.+]] = call @testIfThen([[CAST0]]) : (tensor) -> tensor -// CHECK: [[CAST1:%.+]] = "tf.Cast"([[THEN]]) {Truncate = false} : (tensor) -> tensor>> +// CHECK: [[CAST1:%.+]] = "tf.Cast"([[THEN]]) <{Truncate = false}> : (tensor) -> tensor>> // CHECK: cf.br ^bb3([[CAST1]] : tensor>>) // CHECK: ^bb2: -// CHECK: [[CAST2:%.+]] = "tf.Cast"(%arg1) {Truncate = false} : (tensor>>) -> tensor +// CHECK: [[CAST2:%.+]] = "tf.Cast"(%arg1) <{Truncate = false}> : (tensor>>) -> tensor // CHECK: [[ELSE:%.+]] = call @testIfElse([[CAST2]]) : (tensor) -> tensor -// CHECK: [[CAST3:%.+]] = "tf.Cast"([[ELSE]]) {Truncate = false} : (tensor) -> tensor>> +// CHECK: [[CAST3:%.+]] = "tf.Cast"([[ELSE]]) <{Truncate = false}> : (tensor) -> tensor>> // CHECK: cf.br ^bb3([[CAST3]] : tensor>>) // CHECK: ^bb3([[BBARG0:%.+]]: tensor>>): // CHECK: return [[BBARG0]] : tensor>> @@ -201,20 +201,20 @@ func.func @testWhileCasts(%arg0: tensor>>) -> ( cond = @testWhileCond, body = @testWhileBody, is_stateless = false } : (tensor>>) -> (tensor>>) func.return %0 : tensor>> -// CHECK: [[CASTENTRY:%.+]] = "tf.Cast"(%arg0) {Truncate = false} : (tensor>>) -> tensor +// CHECK: [[CASTENTRY:%.+]] = "tf.Cast"(%arg0) <{Truncate = false}> : (tensor>>) -> tensor // CHECK: cf.br ^bb1([[CASTENTRY]] : tensor) // CHECK: ^bb1([[CONDARG0:%.+]]: tensor): // 2 preds: ^bb0, ^bb2 // CHECK: [[CONTINUE:%.+]] = call @testWhileCond([[CONDARG0]]) : (tensor) -> tensor // CHECK: [[TOBOOL:%.+]] = "tf.ToBool"([[CONTINUE]]) : (tensor) -> tensor // CHECK: [[PRED:%.+]] = tensor.extract [[TOBOOL]][] : tensor -// CHECK: [[CASTCONDARG0:%.+]] = "tf.Cast"([[CONDARG0]]) {Truncate = false} : (tensor) -> tensor>> +// CHECK: [[CASTCONDARG0:%.+]] = "tf.Cast"([[CONDARG0]]) <{Truncate = false}> : (tensor) -> tensor>> // CHECK: cf.cond_br [[PRED]], ^bb2([[CASTCONDARG0]] : tensor>>), ^bb3([[CASTCONDARG0]] : tensor>>) // CHECK: ^bb2([[BODYARG0:%.+]]: tensor>>): // pred: ^bb1 // CHECK: [[WHILERET:%.+]] = call @testWhileBody([[BODYARG0]]) : (tensor>>) -> tensor>> -// CHECK: [[CASTWHILERET:%.+]] = "tf.Cast"([[WHILERET]]) {Truncate = false} : (tensor>>) -> tensor +// CHECK: [[CASTWHILERET:%.+]] = "tf.Cast"([[WHILERET]]) <{Truncate = false}> : (tensor>>) -> tensor // CHECK: cf.br ^bb1([[CASTWHILERET]] : tensor) // CHECK: ^bb3([[EXITARG0:%.+]]: tensor>>): // pred: ^bb1 -// CHECK: [[CASTEXITARG0:%.+]] = "tf.Cast"([[EXITARG0]]) {Truncate = false} : (tensor>>) -> tensor>> +// CHECK: [[CASTEXITARG0:%.+]] = "tf.Cast"([[EXITARG0]]) <{Truncate = false}> : (tensor>>) -> tensor>> // CHECK: return [[CASTEXITARG0]] : tensor>> } diff --git a/tensorflow/compiler/mlir/tensorflow/tests/functional-control-flow-to-regions.mlir b/tensorflow/compiler/mlir/tensorflow/tests/functional-control-flow-to-regions.mlir index d426267f26ccd3..c5cf58971296fb 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/functional-control-flow-to-regions.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/functional-control-flow-to-regions.mlir @@ -14,17 +14,17 @@ func.func @testIf1Result(%arg0: tensor, %arg1: tensor<*xf32>) -> tensor<*xf3 } : (tensor, tensor<*xf32>) -> tensor<*xf32> // CHECK: "tf.IfRegion" + // CHECK-SAME: <{_else_func_name = "testIf1Else" + // CHECK-SAME: _then_func_name = "testIf1Then" + // CHECK-SAME: is_stateless = false // CHECK: [[Result0:%.*]] = func.call @testIf1Then // CHECK: "tf.Yield"([[Result0]]) // CHECK: [[Result1:%.*]] = func.call @testIf1Else // CHECK: "tf.Yield"([[Result1]]) // CHECK: _attr0 = 10 // CHECK-SAME: _attr1 = true - // CHECK-SAME: _else_func_name = "testIf1Else" - // CHECK-SAME: _then_func_name = "testIf1Then" // CHECK-NOT: attr2 = // CHECK-NOT: else_branch - // CHECK-SAME: is_stateless = false // CHECK-NOT: then_branch // CHECK-SAME: } func.return %0 : tensor<*xf32> @@ -179,6 +179,7 @@ func.func @testWhileResult(tensor<*xf32>) -> (tensor<*xf32>) { } : (tensor<*xf32>) -> (tensor<*xf32>) // CHECK: [[Result0:%.*]] = "tf.WhileRegion" + // CHECK-SAME: is_stateless = true // CHECK: ^bb0([[CARG0:%[^:]*]]: // CHECK: [[Result1:%.*]] = func.call @testWhileCond // CHECK: "tf.Yield"([[Result1]], [[CARG0]]) @@ -189,7 +190,6 @@ func.func @testWhileResult(tensor<*xf32>) -> (tensor<*xf32>) { // CHECK-NOT: attr2 = // CHECK-NOT: cond = // CHECK-NOT: body = - // CHECK-SAME: is_stateless = true // CHECK: return [[Result0]] func.return %1 : tensor<*xf32> } diff --git a/tensorflow/compiler/mlir/tensorflow/tests/fused_kernel_matcher.mlir b/tensorflow/compiler/mlir/tensorflow/tests/fused_kernel_matcher.mlir index 8458cb80df6983..38380ded0c93bb 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/fused_kernel_matcher.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/fused_kernel_matcher.mlir @@ -6,7 +6,7 @@ // CHECK-LABEL: conv2DBiasAdd_noActivation func.func @conv2DBiasAdd_noActivation(%arg0: tensor<128xf32>, %arg1: tensor<1x1x3x128xf32>, %arg2: tensor<8x32x32x3xf32>) -> (tensor<*xf32>) { - // CHECK: %[[VAL_0:.*]] = "tf._FusedConv2D"(%arg2, %arg1, %arg0) {TArgs = [f32], data_format = "NHWC", dilations = [1, 1, 1, 1], epsilon = 0.000000e+00 : f32, explicit_paddings = [], fused_ops = ["BiasAdd"], num_args = 1 : i64, operandSegmentSizes = array, padding = "SAME", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true} : (tensor<8x32x32x3xf32>, tensor<1x1x3x128xf32>, tensor<128xf32>) -> tensor<*xf32> + // CHECK: %[[VAL_0:.*]] = "tf._FusedConv2D"(%arg2, %arg1, %arg0) <{data_format = "NHWC", dilations = [1, 1, 1, 1], epsilon = 0.000000e+00 : f32, explicit_paddings = [], fused_ops = ["BiasAdd"], num_args = 1 : i64, operandSegmentSizes = array, padding = "SAME", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true}> {TArgs = [f32]} : (tensor<8x32x32x3xf32>, tensor<1x1x3x128xf32>, tensor<128xf32>) -> tensor<*xf32> // CHECK: %[[VAL_1:.*]] = "tf.Identity"(%[[VAL_0]]) : (tensor<*xf32>) -> tensor<*xf32> // CHECK: return %[[VAL_1]] %0 = "tf.Conv2D"(%arg2, %arg1) {data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true} : (tensor<8x32x32x3xf32>, tensor<1x1x3x128xf32>) -> tensor<*xf32> @@ -17,7 +17,7 @@ func.func @conv2DBiasAdd_noActivation(%arg0: tensor<128xf32>, %arg1: tensor<1x1x // CHECK-LABEL: conv2DBiasAdd_reluActivation func.func @conv2DBiasAdd_reluActivation(%arg0: tensor<128xf32>, %arg1: tensor<1x1x3x128xf32>, %arg2: tensor<8x32x32x3xf32>) -> (tensor<*xf32>) { - // CHECK: %[[VAL_0:.*]] = "tf._FusedConv2D"(%arg2, %arg1, %arg0) {TArgs = [f32], data_format = "NHWC", dilations = [1, 1, 1, 1], epsilon = 0.000000e+00 : f32, explicit_paddings = [], fused_ops = ["BiasAdd", "Relu"], num_args = 1 : i64, operandSegmentSizes = array, padding = "SAME", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true} : (tensor<8x32x32x3xf32>, tensor<1x1x3x128xf32>, tensor<128xf32>) -> tensor<*xf32> + // CHECK: %[[VAL_0:.*]] = "tf._FusedConv2D"(%arg2, %arg1, %arg0) <{data_format = "NHWC", dilations = [1, 1, 1, 1], epsilon = 0.000000e+00 : f32, explicit_paddings = [], fused_ops = ["BiasAdd", "Relu"], num_args = 1 : i64, operandSegmentSizes = array, padding = "SAME", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true}> {TArgs = [f32]} : (tensor<8x32x32x3xf32>, tensor<1x1x3x128xf32>, tensor<128xf32>) -> tensor<*xf32> // CHECK: %[[VAL_1:.*]] = "tf.Identity"(%[[VAL_0]]) : (tensor<*xf32>) -> tensor<*xf32> // CHECK: return %[[VAL_1]] %0 = "tf.Conv2D"(%arg2, %arg1) {data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true} : (tensor<8x32x32x3xf32>, tensor<1x1x3x128xf32>) -> tensor<*xf32> @@ -29,7 +29,7 @@ func.func @conv2DBiasAdd_reluActivation(%arg0: tensor<128xf32>, %arg1: tensor<1x // CHECK-LABEL: conv2DBiasAdd_relu6Activation func.func @conv2DBiasAdd_relu6Activation(%arg0: tensor<128xf32>, %arg1: tensor<1x1x3x128xf32>, %arg2: tensor<8x32x32x3xf32>) -> (tensor<*xf32>) { - // CHECK: %[[VAL_0:.*]] = "tf._FusedConv2D"(%arg2, %arg1, %arg0) {TArgs = [f32], data_format = "NHWC", dilations = [1, 1, 1, 1], epsilon = 0.000000e+00 : f32, explicit_paddings = [], fused_ops = ["BiasAdd", "Relu6"], num_args = 1 : i64, operandSegmentSizes = array, padding = "SAME", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true} : (tensor<8x32x32x3xf32>, tensor<1x1x3x128xf32>, tensor<128xf32>) -> tensor<*xf32> + // CHECK: %[[VAL_0:.*]] = "tf._FusedConv2D"(%arg2, %arg1, %arg0) <{data_format = "NHWC", dilations = [1, 1, 1, 1], epsilon = 0.000000e+00 : f32, explicit_paddings = [], fused_ops = ["BiasAdd", "Relu6"], num_args = 1 : i64, operandSegmentSizes = array, padding = "SAME", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true}> {TArgs = [f32]} : (tensor<8x32x32x3xf32>, tensor<1x1x3x128xf32>, tensor<128xf32>) -> tensor<*xf32> // CHECK: %[[VAL_1:.*]] = "tf.Identity"(%[[VAL_0]]) : (tensor<*xf32>) -> tensor<*xf32> // CHECK: return %[[VAL_1]] %0 = "tf.Conv2D"(%arg2, %arg1) {data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true} : (tensor<8x32x32x3xf32>, tensor<1x1x3x128xf32>) -> tensor<*xf32> @@ -41,7 +41,7 @@ func.func @conv2DBiasAdd_relu6Activation(%arg0: tensor<128xf32>, %arg1: tensor<1 // CHECK-LABEL: conv2DBiasAdd_eluActivation func.func @conv2DBiasAdd_eluActivation(%arg0: tensor<128xf32>, %arg1: tensor<1x1x3x128xf32>, %arg2: tensor<8x32x32x3xf32>) -> (tensor<*xf32>) { - // CHECK: %[[VAL_0:.*]] = "tf._FusedConv2D"(%arg2, %arg1, %arg0) {TArgs = [f32], data_format = "NHWC", dilations = [1, 1, 1, 1], epsilon = 0.000000e+00 : f32, explicit_paddings = [], fused_ops = ["BiasAdd", "Elu"], num_args = 1 : i64, operandSegmentSizes = array, padding = "SAME", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true} : (tensor<8x32x32x3xf32>, tensor<1x1x3x128xf32>, tensor<128xf32>) -> tensor<*xf32> + // CHECK: %[[VAL_0:.*]] = "tf._FusedConv2D"(%arg2, %arg1, %arg0) <{data_format = "NHWC", dilations = [1, 1, 1, 1], epsilon = 0.000000e+00 : f32, explicit_paddings = [], fused_ops = ["BiasAdd", "Elu"], num_args = 1 : i64, operandSegmentSizes = array, padding = "SAME", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true}> {TArgs = [f32]} : (tensor<8x32x32x3xf32>, tensor<1x1x3x128xf32>, tensor<128xf32>) -> tensor<*xf32> // CHECK: %[[VAL_1:.*]] = "tf.Identity"(%[[VAL_0]]) : (tensor<*xf32>) -> tensor<*xf32> // CHECK: return %[[VAL_1]] %0 = "tf.Conv2D"(%arg2, %arg1) {data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true} : (tensor<8x32x32x3xf32>, tensor<1x1x3x128xf32>) -> tensor<*xf32> @@ -64,7 +64,7 @@ func.func @conv2DBiasAdd_convMultipleUses(%arg0: tensor<128xf32>, %arg1: tensor< // CHECK-LABEL: conv2DBiasAdd_biasAddMultipleUse func.func @conv2DBiasAdd_biasAddMultipleUse(%arg0: tensor<128xf32>, %arg1: tensor<1x1x3x128xf32>, %arg2: tensor<8x32x32x3xf32>) -> (tensor<*xf32>, tensor<*xf32>) { - // CHECK-DAG: %[[VAL:.*]] = "tf._FusedConv2D"(%arg2, %arg1, %arg0) {TArgs = [f32], data_format = "NHWC", dilations = [1, 1, 1, 1], epsilon = 0.000000e+00 : f32, explicit_paddings = [], fused_ops = ["BiasAdd"], num_args = 1 : i64, operandSegmentSizes = array, padding = "SAME", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true} : (tensor<8x32x32x3xf32>, tensor<1x1x3x128xf32>, tensor<128xf32>) -> tensor<*xf32> + // CHECK-DAG: %[[VAL:.*]] = "tf._FusedConv2D"(%arg2, %arg1, %arg0) <{data_format = "NHWC", dilations = [1, 1, 1, 1], epsilon = 0.000000e+00 : f32, explicit_paddings = [], fused_ops = ["BiasAdd"], num_args = 1 : i64, operandSegmentSizes = array, padding = "SAME", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true}> {TArgs = [f32]} : (tensor<8x32x32x3xf32>, tensor<1x1x3x128xf32>, tensor<128xf32>) -> tensor<*xf32> // CHECK-DAG: %[[VAL_0:.*]] = "tf.Elu"(%[[VAL]]) : (tensor<*xf32>) -> tensor<*xf32> // CHECK-DAG: %[[VAL_1:.*]] = "tf.Identity"(%[[VAL_0]]) : (tensor<*xf32>) -> tensor<*xf32> // CHECK-DAG: %[[VAL_2:.*]] = "tf.Identity"(%[[VAL]]) : (tensor<*xf32>) -> tensor<*xf32> @@ -89,7 +89,7 @@ func.func @conv2D_noFusion(%arg0: tensor<128xf32>, %arg1: tensor<1x1x3x128xf32>, // CHECK-LABEL: conv2D_noFusion1 func.func @conv2D_noFusion1(%arg0: tensor<*xf32>, %arg1: tensor<1x1x3x128xf32>, %arg2: tensor<8x32x32x3xf32>) -> (tensor<*xf32>) { // CHECK-NOT: "tf._FusedConv2D" - %0 = "tf.Conv2D"(%arg2, %arg1) {data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true} : (tensor<8x32x32x3xf32>, tensor<1x1x3x128xf32>) -> tensor<*xf32> + %0 = "tf.Conv2D"(%arg2, %arg1) <{data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true}> : (tensor<8x32x32x3xf32>, tensor<1x1x3x128xf32>) -> tensor<*xf32> // The result of the conv must be the first input to BiasAdd to be fusable. %1 = "tf.BiasAdd"(%arg0, %0) {data_format = "NHWC"} : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32> %2 = "tf.Elu"(%1) : (tensor<*xf32>) -> tensor<*xf32> @@ -114,7 +114,7 @@ func.func @conv2D_dataFormatMismatch(%arg0: tensor<128xf32>, %arg1: tensor<1x1x3 // CHECK-LABEL: matmulBiasAdd func.func @matmulBiasAdd(%arg0: tensor<64xf32>, %arg1: tensor<8x32xf32>, %arg2: tensor<32x64xf32>) -> (tensor<*xf32>) { - // CHECK: %[[VAL_3:.*]] = "tf._FusedMatMul"(%arg1, %arg2, %arg0) {epsilon = 0.000000e+00 : f32, fused_ops = ["BiasAdd"], transpose_a = false, transpose_b = false} : (tensor<8x32xf32>, tensor<32x64xf32>, tensor<64xf32>) -> tensor<*xf32> + // CHECK: %[[VAL_3:.*]] = "tf._FusedMatMul"(%arg1, %arg2, %arg0) <{epsilon = 0.000000e+00 : f32, fused_ops = ["BiasAdd"], transpose_a = false, transpose_b = false}> : (tensor<8x32xf32>, tensor<32x64xf32>, tensor<64xf32>) -> tensor<*xf32> // CHECK: %[[VAL_4:.*]] = "tf.Identity"(%[[VAL_3]]) : (tensor<*xf32>) -> tensor<*xf32> // CHECK: return %[[VAL_4]] %3 = "tf.MatMul"(%arg1, %arg2) {transpose_a = false, transpose_b = false} : (tensor<8x32xf32>, tensor<32x64xf32>) -> tensor<*xf32> @@ -125,7 +125,7 @@ func.func @matmulBiasAdd(%arg0: tensor<64xf32>, %arg1: tensor<8x32xf32>, %arg2: // CHECK-LABEL: matmulBiasAdd_relu func.func @matmulBiasAdd_relu(%arg0: tensor<64xf32>, %arg1: tensor<8x32xf32>, %arg2: tensor<32x64xf32>) -> (tensor<*xf32>) { - // CHECK: %[[VAL_3:.*]] = "tf._FusedMatMul"(%arg1, %arg2, %arg0) {epsilon = 0.000000e+00 : f32, fused_ops = ["BiasAdd", "Relu"], transpose_a = false, transpose_b = false} : (tensor<8x32xf32>, tensor<32x64xf32>, tensor<64xf32>) -> tensor<*xf32> + // CHECK: %[[VAL_3:.*]] = "tf._FusedMatMul"(%arg1, %arg2, %arg0) <{epsilon = 0.000000e+00 : f32, fused_ops = ["BiasAdd", "Relu"], transpose_a = false, transpose_b = false}> : (tensor<8x32xf32>, tensor<32x64xf32>, tensor<64xf32>) -> tensor<*xf32> // CHECK: %[[VAL_4:.*]] = "tf.Identity"(%[[VAL_3]]) : (tensor<*xf32>) -> tensor<*xf32> // CHECK: return %[[VAL_4]] %3 = "tf.MatMul"(%arg1, %arg2) {transpose_a = false, transpose_b = false} : (tensor<8x32xf32>, tensor<32x64xf32>) -> tensor<*xf32> @@ -137,7 +137,7 @@ func.func @matmulBiasAdd_relu(%arg0: tensor<64xf32>, %arg1: tensor<8x32xf32>, %a // CHECK-LABEL: matmulBiasAdd_relu6 func.func @matmulBiasAdd_relu6(%arg0: tensor<64xf32>, %arg1: tensor<8x32xf32>, %arg2: tensor<32x64xf32>) -> (tensor<*xf32>) { - // CHECK: %[[VAL_3:.*]] = "tf._FusedMatMul"(%arg1, %arg2, %arg0) {epsilon = 0.000000e+00 : f32, fused_ops = ["BiasAdd", "Relu6"], transpose_a = false, transpose_b = false} : (tensor<8x32xf32>, tensor<32x64xf32>, tensor<64xf32>) -> tensor<*xf32> + // CHECK: %[[VAL_3:.*]] = "tf._FusedMatMul"(%arg1, %arg2, %arg0) <{epsilon = 0.000000e+00 : f32, fused_ops = ["BiasAdd", "Relu6"], transpose_a = false, transpose_b = false}> : (tensor<8x32xf32>, tensor<32x64xf32>, tensor<64xf32>) -> tensor<*xf32> // CHECK: %[[VAL_4:.*]] = "tf.Identity"(%[[VAL_3]]) : (tensor<*xf32>) -> tensor<*xf32> // CHECK: return %[[VAL_4]] %3 = "tf.MatMul"(%arg1, %arg2) {transpose_a = false, transpose_b = false} : (tensor<8x32xf32>, tensor<32x64xf32>) -> tensor<*xf32> @@ -149,7 +149,7 @@ func.func @matmulBiasAdd_relu6(%arg0: tensor<64xf32>, %arg1: tensor<8x32xf32>, % // CHECK-LABEL: matmulBiasAdd_elu func.func @matmulBiasAdd_elu(%arg0: tensor<64xf32>, %arg1: tensor<8x32xf32>, %arg2: tensor<32x64xf32>) -> (tensor<*xf32>) { - // CHECK: %[[VAL_3:.*]] = "tf._FusedMatMul"(%arg1, %arg2, %arg0) {epsilon = 0.000000e+00 : f32, fused_ops = ["BiasAdd", "Elu"], transpose_a = false, transpose_b = false} : (tensor<8x32xf32>, tensor<32x64xf32>, tensor<64xf32>) -> tensor<*xf32> + // CHECK: %[[VAL_3:.*]] = "tf._FusedMatMul"(%arg1, %arg2, %arg0) <{epsilon = 0.000000e+00 : f32, fused_ops = ["BiasAdd", "Elu"], transpose_a = false, transpose_b = false}> : (tensor<8x32xf32>, tensor<32x64xf32>, tensor<64xf32>) -> tensor<*xf32> // CHECK: %[[VAL_4:.*]] = "tf.Identity"(%[[VAL_3]]) : (tensor<*xf32>) -> tensor<*xf32> // CHECK: return %[[VAL_4]] %3 = "tf.MatMul"(%arg1, %arg2) {transpose_a = false, transpose_b = false} : (tensor<8x32xf32>, tensor<32x64xf32>) -> tensor<*xf32> diff --git a/tensorflow/compiler/mlir/tensorflow/tests/graphdef2mlir/graph-custom-operation.pbtxt b/tensorflow/compiler/mlir/tensorflow/tests/graphdef2mlir/graph-custom-operation.pbtxt index c2f4d7aab5ca82..e820f366497473 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/graphdef2mlir/graph-custom-operation.pbtxt +++ b/tensorflow/compiler/mlir/tensorflow/tests/graphdef2mlir/graph-custom-operation.pbtxt @@ -54,5 +54,5 @@ versions { # the names are matching between the function definition and the uses / call # site (a numerical suffix may be appended). -# CHECK: "tf.LegacyCall"(%outputs) {_disable_call_shape_inference = false, device = "", f = @foo0} +# CHECK: "tf.LegacyCall"(%outputs) <{_disable_call_shape_inference = false, f = @foo0}> {device = ""} # CHECK: func private @foo0 diff --git a/tensorflow/compiler/mlir/tensorflow/tests/graphdef2mlir/graph-function-call.pbtxt b/tensorflow/compiler/mlir/tensorflow/tests/graphdef2mlir/graph-function-call.pbtxt index f954657765a56e..02dd85f3dc6a34 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/graphdef2mlir/graph-function-call.pbtxt +++ b/tensorflow/compiler/mlir/tensorflow/tests/graphdef2mlir/graph-function-call.pbtxt @@ -68,4 +68,4 @@ library { } # CHECK: func @main -# CHECK: "tf.LegacyCall"(%arg0) {_disable_call_shape_inference = true, _tpu_replicate = "cluster", device = "", f = @test_func_name0} +# CHECK: "tf.LegacyCall"(%arg0) <{_disable_call_shape_inference = true, f = @test_func_name0}> {_tpu_replicate = "cluster", device = ""} diff --git a/tensorflow/compiler/mlir/tensorflow/tests/graphdef2mlir/graph-function-name-bug.pbtxt b/tensorflow/compiler/mlir/tensorflow/tests/graphdef2mlir/graph-function-name-bug.pbtxt index 4b937a17af82b8..7244b60cbd0570 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/graphdef2mlir/graph-function-name-bug.pbtxt +++ b/tensorflow/compiler/mlir/tensorflow/tests/graphdef2mlir/graph-function-name-bug.pbtxt @@ -121,8 +121,8 @@ versions { # Verify that functions from the library are properly imported. # CHECK-LABEL: func @main() { -# CHECK: "tf.LegacyCall"() {_disable_call_shape_inference = false, device = "", f = @foo110} -# CHECK: "tf.LegacyCall"() {_disable_call_shape_inference = false, device = "", f = @foo111} +# CHECK: "tf.LegacyCall"() <{_disable_call_shape_inference = false, f = @foo110}> {device = ""} +# CHECK: "tf.LegacyCall"() <{_disable_call_shape_inference = false, f = @foo111}> {device = ""} # CHECK-LABEL: func private @foo110() # CHECK-LABEL: func private @foo111() diff --git a/tensorflow/compiler/mlir/tensorflow/tests/graphdef2mlir/graph-function-resource-args.pbtxt b/tensorflow/compiler/mlir/tensorflow/tests/graphdef2mlir/graph-function-resource-args.pbtxt index 66847dc63e904f..ea5957a666f2d5 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/graphdef2mlir/graph-function-resource-args.pbtxt +++ b/tensorflow/compiler/mlir/tensorflow/tests/graphdef2mlir/graph-function-resource-args.pbtxt @@ -88,7 +88,7 @@ library { # CHECK: tf_executor.graph # CHECK: "tf.VarHandleOp"() # CHECK: "tf.LegacyCall" -# CHECK-SAME: {_disable_call_shape_inference = true, device = "", f = @test_func_name0} +# CHECK-SAME: <{_disable_call_shape_inference = true, f = @test_func_name0}> {device = ""} # CHECK: tf_executor.fetch # CHECK: return # CHECK: func private @test_func_name0 diff --git a/tensorflow/compiler/mlir/tensorflow/tests/graphdef2mlir/graph-library.pbtxt b/tensorflow/compiler/mlir/tensorflow/tests/graphdef2mlir/graph-library.pbtxt index eb593188888371..f515761b0921ed 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/graphdef2mlir/graph-library.pbtxt +++ b/tensorflow/compiler/mlir/tensorflow/tests/graphdef2mlir/graph-library.pbtxt @@ -54,10 +54,10 @@ versions { # Verify that functions from the library are properly imported. # CHECK-LABEL: func @main() { -# CHECK: "tf.LegacyCall"() {_disable_call_shape_inference = true, device = "", f = @foo0} -# CHECK: "tf.LegacyCall"() {_disable_call_shape_inference = false, device = "", f = @bar0} +# CHECK: "tf.LegacyCall"() <{_disable_call_shape_inference = true, f = @foo0}> {device = ""} +# CHECK: "tf.LegacyCall"() <{_disable_call_shape_inference = false, f = @bar0}> {device = ""} # CHECK-LABEL: func private @foo0() -# CHECK: "tf.LegacyCall"() {_disable_call_shape_inference = false, device = "", f = @bar0} +# CHECK: "tf.LegacyCall"() <{_disable_call_shape_inference = false, f = @bar0}> {device = ""} # CHECK-LABEL: func private @bar0() diff --git a/tensorflow/compiler/mlir/tensorflow/tests/graphdef2mlir/mlir_passthrough_op.pbtxt b/tensorflow/compiler/mlir/tensorflow/tests/graphdef2mlir/mlir_passthrough_op.pbtxt index fd33be7baaada2..868ee809fdd9a0 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/graphdef2mlir/mlir_passthrough_op.pbtxt +++ b/tensorflow/compiler/mlir/tensorflow/tests/graphdef2mlir/mlir_passthrough_op.pbtxt @@ -1,7 +1,7 @@ # RUN: tf-mlir-translate -graphdef-to-mlir -tf-enable-shape-inference-on-import=false %s | FileCheck %s # CHECK:"tf.MlirPassthroughOp" -# CHECK: mlir_module = "\0Afunc @main(%arg0 : tensor<10xf32>, %arg1 : tensor<10xf32>) -> tensor<10x10xf32> {\0A %add = \22tf.Add\22(%arg0, %arg1) : (tensor<10xf32>, tensor<10xf32>) -> tensor<10xf32>\0A %ret = \22magic.op\22(%add, %add) : (tensor<10xf32>, tensor<10xf32>) -> tensor<10x10xf32>\0A return %ret : tensor<10x10xf32>\0A}\0A"} : (tensor<10xf32>, tensor<10xf32>) -> tensor<*xf32> +# CHECK: mlir_module = "\0Afunc @main(%arg0 : tensor<10xf32>, %arg1 : tensor<10xf32>) -> tensor<10x10xf32> {\0A %add = \22tf.Add\22(%arg0, %arg1) : (tensor<10xf32>, tensor<10xf32>) -> tensor<10xf32>\0A %ret = \22magic.op\22(%add, %add) : (tensor<10xf32>, tensor<10xf32>) -> tensor<10x10xf32>\0A return %ret : tensor<10x10xf32>\0A}\0A"}> {device = ""} : (tensor<10xf32>, tensor<10xf32>) -> tensor<*xf32> node { name: "x" diff --git a/tensorflow/compiler/mlir/tensorflow/tests/guarantee-all-funcs-one-use.mlir b/tensorflow/compiler/mlir/tensorflow/tests/guarantee-all-funcs-one-use.mlir index 52f29badb19a2b..0dd181374d5909 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/guarantee-all-funcs-one-use.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/guarantee-all-funcs-one-use.mlir @@ -73,9 +73,9 @@ module { // Test stateful and stateless partitioned calls. // CHECK-LABEL: func @f func.func @f() { - // CHECK: "tf.PartitionedCall"() {config = "", config_proto = "", executor_type = "", f = @g} : () -> () + // CHECK: "tf.PartitionedCall"() <{config = "", config_proto = "", executor_type = "", f = @g}> : () -> () "tf.PartitionedCall"() {config = "", config_proto = "", executor_type = "", f = @g} : () -> () - // CHECK: "tf.StatefulPartitionedCall"() {config = "", config_proto = "", executor_type = "", f = @[[NEWG:.+]]} : () -> () + // CHECK: "tf.StatefulPartitionedCall"() <{config = "", config_proto = "", executor_type = "", f = @[[NEWG:.+]]}> : () -> () "tf.StatefulPartitionedCall"() {config = "", config_proto = "", executor_type = "", f = @g} : () -> () func.return } diff --git a/tensorflow/compiler/mlir/tensorflow/tests/inlining.mlir b/tensorflow/compiler/mlir/tensorflow/tests/inlining.mlir index f47330099fdc26..4b0ba861707096 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/inlining.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/inlining.mlir @@ -71,8 +71,8 @@ func.func @dont_inline_custom_on_duplicated_cases() -> tensor<2xi32> { // CHECK-LABEL: func @inline_shape_cast( // CHECK-SAME: %[[ARG:.*]]: tensor<2xi32> func.func @inline_shape_cast(%arg: tensor<2xi32>) -> tensor<2xi32> { - // CHECK-NEXT: %[[ARG_CAST:.*]] = "tf.Cast"(%[[ARG]]) {Truncate = false} : (tensor<2xi32>) -> tensor<*xi32> - // CHECK-NEXT: %[[RESULT_CAST:.*]] = "tf.Cast"(%[[ARG_CAST]]) {Truncate = false} : (tensor<*xi32>) -> tensor<2xi32> + // CHECK-NEXT: %[[ARG_CAST:.*]] = "tf.Cast"(%[[ARG]]) <{Truncate = false}> : (tensor<2xi32>) -> tensor<*xi32> + // CHECK-NEXT: %[[RESULT_CAST:.*]] = "tf.Cast"(%[[ARG_CAST]]) <{Truncate = false}> : (tensor<*xi32>) -> tensor<2xi32> // CHECK-NEXT: return %[[RESULT_CAST]] %result = "tf.PartitionedCall"(%arg) {config = "", config_proto = "", executor_type = "", f = @inline_shape_cast_callee} : (tensor<2xi32>) -> tensor<2xi32> func.return %result : tensor<2xi32> diff --git a/tensorflow/compiler/mlir/tensorflow/tests/layout_optimization_layout_assignment_to_nchw.mlir b/tensorflow/compiler/mlir/tensorflow/tests/layout_optimization_layout_assignment_to_nchw.mlir index c2dfdc65919583..37461838f4f85e 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/layout_optimization_layout_assignment_to_nchw.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/layout_optimization_layout_assignment_to_nchw.mlir @@ -7,7 +7,7 @@ // CHECK-LABEL: func @transposeConv2D func.func @transposeConv2D(%input: tensor<1x32x32x3xf32>, %filter: tensor<1x1x3x8xf32>) -> tensor<1x7x7x8xf32> { - // CHECK: %[[ARG_PERM:.*]] = "tf.Const"() {value = dense<[0, 3, 1, 2]> : tensor<4xi64>} + // CHECK: %[[ARG_PERM:.*]] = "tf.Const"() <{value = dense<[0, 3, 1, 2]> : tensor<4xi64>}> // CHECK: %[[ARG_TRANSPOSE:[0-9]*]] = "tf.Transpose"(%arg0, %[[ARG_PERM]]) // CHECK: %[[CONV2D:[0-9]*]] = "tf.Conv2D"(%[[ARG_TRANSPOSE]], %arg1) @@ -18,7 +18,7 @@ func.func @transposeConv2D(%input: tensor<1x32x32x3xf32>, %filter: tensor<1x1x3x // CHECK-SAME: strides = [5, 8, 6, 7] // CHECK-SAME: (tensor<1x3x32x32xf32>, tensor<1x1x3x8xf32>) -> tensor<1x8x7x7xf32> - // CHECK: %[[RES_PERM:.*]] = "tf.Const"() {value = dense<[0, 2, 3, 1]> : tensor<4xi64>} + // CHECK: %[[RES_PERM:.*]] = "tf.Const"() <{value = dense<[0, 2, 3, 1]> : tensor<4xi64>}> // CHECK: %[[RES_TRANSPOSE:[0-9]*]] = "tf.Transpose"(%[[CONV2D]], %[[RES_PERM]]) // CHECK: return %[[RES_TRANSPOSE]] @@ -38,7 +38,7 @@ func.func @transposeConv2D(%input: tensor<1x32x32x3xf32>, %filter: tensor<1x1x3x func.func @transposeConv2DWithDefaultAttr(%input: tensor<1x32x32x3xf32>, %filter: tensor<1x1x3x8xf32>) -> tensor { - // CHECK: %[[ARG_PERM:.*]] = "tf.Const"() {value = dense<[0, 3, 1, 2]> : tensor<4xi64>} + // CHECK: %[[ARG_PERM:.*]] = "tf.Const"() <{value = dense<[0, 3, 1, 2]> : tensor<4xi64>}> // CHECK: %[[ARG_TRANSPOSE:[0-9]*]] = "tf.Transpose"(%arg0, %[[ARG_PERM]]) // CHECK: %[[CONV2D:[0-9]*]] = "tf.Conv2D"(%[[ARG_TRANSPOSE]], %arg1) @@ -49,7 +49,7 @@ func.func @transposeConv2DWithDefaultAttr(%input: tensor<1x32x32x3xf32>, %filter // CHECK-SAME: strides = [5, 8, 6, 7] // CHECK-SAME: (tensor<1x3x32x32xf32>, tensor<1x1x3x8xf32>) -> tensor - // CHECK: %[[RES_PERM:.*]] = "tf.Const"() {value = dense<[0, 2, 3, 1]> : tensor<4xi64>} + // CHECK: %[[RES_PERM:.*]] = "tf.Const"() <{value = dense<[0, 2, 3, 1]> : tensor<4xi64>}> // CHECK: %[[RES_TRANSPOSE:[0-9]*]] = "tf.Transpose"(%[[CONV2D]], %[[RES_PERM]]) // CHECK: return %[[RES_TRANSPOSE]] @@ -77,7 +77,7 @@ func.func @transposeConv2DBackpropFilter( // CHECK-SAME: dst_format = "NCHW" // CHECK-SAME: src_format = "NHWC" - // CHECK: %[[ARG_PERM:.*]] = "tf.Const"() {value = dense<[0, 3, 1, 2]> : tensor<4xi64>} + // CHECK: %[[ARG_PERM:.*]] = "tf.Const"() <{value = dense<[0, 3, 1, 2]> : tensor<4xi64>}> // CHECK: %[[IN_TRANSPOSE:[0-9]*]] = "tf.Transpose"(%arg0, %[[ARG_PERM]]) // CHECK: %[[OUT_BP_TRANSPOSE:[0-9]*]] = "tf.Transpose"(%arg2, %[[ARG_PERM]]) @@ -117,7 +117,7 @@ func.func @transposeConv2DBackpropInput( // CHECK-SAME: dst_format = "NCHW" // CHECK-SAME: src_format = "NHWC" - // CHECK: %[[ARG_PERM:.*]] = "tf.Const"() {value = dense<[0, 3, 1, 2]> : tensor<4xi64>} + // CHECK: %[[ARG_PERM:.*]] = "tf.Const"() <{value = dense<[0, 3, 1, 2]> : tensor<4xi64>}> // CHECK: %[[OUT_BP_TRANSPOSE:[0-9]*]] = "tf.Transpose"(%arg2, %[[ARG_PERM]]) // CHECK: %[[CONV2D_BACKPROP:[0-9]*]] = "tf.Conv2DBackpropInput" @@ -130,7 +130,7 @@ func.func @transposeConv2DBackpropInput( // CHECK-SAME: (tensor<4xi32>, tensor<1x1x3x8xf32>, tensor<1x8x32x32xf32>) // CHECK-SAME: -> tensor<1x3x32x32xf32> - // CHECK: %[[RES_PERM:.*]] = "tf.Const"() {value = dense<[0, 2, 3, 1]> : tensor<4xi64>} + // CHECK: %[[RES_PERM:.*]] = "tf.Const"() <{value = dense<[0, 2, 3, 1]> : tensor<4xi64>}> // CHECK: %[[RES_TRANSPOSE:[0-9]*]] = "tf.Transpose"(%[[CONV2D_BACKPROP]], %[[RES_PERM]]) // CHECK: return %[[RES_TRANSPOSE]] @@ -154,7 +154,7 @@ func.func @transposeFusedBatchNormV3( ) -> tensor<1x28x28x64xf32> { // CHECK: %[[ARG_PERM:.*]] = "tf.Const"() - // CHECK-SAME: {value = dense<[0, 3, 1, 2]> : tensor<4xi64>} + // CHECK-SAME: <{value = dense<[0, 3, 1, 2]> : tensor<4xi64>}> // CHECK: %[[ARG_TRANSPOSE:[0-9]*]] = "tf.Transpose"(%arg0, %[[ARG_PERM]]) // CHECK: "tf.FusedBatchNormV3" @@ -164,7 +164,7 @@ func.func @transposeFusedBatchNormV3( // CHECK-SAME: -> (tensor<1x64x28x28xf32>, tensor<64xf32>, // CHECK: %[[RES_PERM:.*]] = "tf.Const"() - // CHECK-SAME: {value = dense<[0, 2, 3, 1]> : tensor<4xi64>} + // CHECK-SAME: <{value = dense<[0, 2, 3, 1]> : tensor<4xi64>}> // CHECK: %[[RES_TRANSPOSE:[0-9]*]] = "tf.Transpose"(%y, %[[RES_PERM]]) // CHECK: return %[[RES_TRANSPOSE]] @@ -192,7 +192,7 @@ func.func @transposeFusedBatchNormGradV3( ) -> tensor<1x28x28x64xf32> { // CHECK: %[[ARG_PERM:.*]] = "tf.Const"() - // CHECK-SAME: {value = dense<[0, 3, 1, 2]> : tensor<4xi64>} + // CHECK-SAME: <{value = dense<[0, 3, 1, 2]> : tensor<4xi64>}> // CHECK: %[[ARG0_TPOSE:[0-9]*]] = "tf.Transpose"(%arg0, %[[ARG_PERM]]) // CHECK: %[[ARG1_TPOSE:[0-9]*]] = "tf.Transpose"(%arg1, %[[ARG_PERM]]) @@ -204,7 +204,7 @@ func.func @transposeFusedBatchNormGradV3( // CHECK-SAME: -> (tensor<1x64x28x28xf32>, // CHECK: %[[RES_PERM:.*]] = "tf.Const"() - // CHECK-SAME: {value = dense<[0, 2, 3, 1]> : tensor<4xi64>} + // CHECK-SAME: <{value = dense<[0, 2, 3, 1]> : tensor<4xi64>}> // CHECK: %[[RES_TPOSE:[0-9]*]] = "tf.Transpose" // CHECK-SAME: (%x_backprop, %[[RES_PERM]]) diff --git a/tensorflow/compiler/mlir/tensorflow/tests/layout_optimization_layout_assignment_to_nhwc.mlir b/tensorflow/compiler/mlir/tensorflow/tests/layout_optimization_layout_assignment_to_nhwc.mlir index 62749d185bfed5..b13da20bb674cd 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/layout_optimization_layout_assignment_to_nhwc.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/layout_optimization_layout_assignment_to_nhwc.mlir @@ -7,7 +7,7 @@ // CHECK-LABEL: func @transposeConv2D func.func @transposeConv2D(%input: tensor<1x3x32x32xf32>, %filter: tensor<1x1x3x8xf32>) -> tensor<1x8x7x6xf32> { - // CHECK: %[[ARG_PERM:.*]] = "tf.Const"() {value = dense<[0, 2, 3, 1]> : tensor<4xi64>} + // CHECK: %[[ARG_PERM:.*]] = "tf.Const"() <{value = dense<[0, 2, 3, 1]> : tensor<4xi64>}> // CHECK: %[[ARG_TRANSPOSE:[0-9]*]] = "tf.Transpose"(%arg0, %[[ARG_PERM]]) // CHECK: %[[CONV2D:[0-9]*]] = "tf.Conv2D"(%[[ARG_TRANSPOSE]], %arg1) @@ -18,7 +18,7 @@ func.func @transposeConv2D(%input: tensor<1x3x32x32xf32>, %filter: tensor<1x1x3x // CHECK-SAME: strides = [5, 7, 8, 6] // CHECK-SAME: (tensor<1x32x32x3xf32>, tensor<1x1x3x8xf32>) -> tensor<1x7x6x8xf32> - // CHECK: %[[RES_PERM:.*]] = "tf.Const"() {value = dense<[0, 3, 1, 2]> : tensor<4xi64>} + // CHECK: %[[RES_PERM:.*]] = "tf.Const"() <{value = dense<[0, 3, 1, 2]> : tensor<4xi64>}> // CHECK: %[[RES_TRANSPOSE:[0-9]*]] = "tf.Transpose"(%[[CONV2D]], %[[RES_PERM]]) // CHECK: return %[[RES_TRANSPOSE]] @@ -41,7 +41,7 @@ func.func @transposeFusedBatchNormV3( ) -> tensor<1x64x28x28xf32> { // CHECK: %[[ARG_PERM:.*]] = "tf.Const"() - // CHECK-SAME: {value = dense<[0, 2, 3, 1]> : tensor<4xi64>} + // CHECK-SAME: <{value = dense<[0, 2, 3, 1]> : tensor<4xi64>}> // CHECK: %[[ARG_TRANSPOSE:[0-9]*]] = "tf.Transpose"(%arg0, %[[ARG_PERM]]) // CHECK: "tf.FusedBatchNormV3" @@ -51,7 +51,7 @@ func.func @transposeFusedBatchNormV3( // CHECK-SAME: -> (tensor<1x28x28x64xf32>, tensor<64xf32>, // CHECK: %[[RES_PERM:.*]] = "tf.Const"() - // CHECK-SAME: {value = dense<[0, 3, 1, 2]> : tensor<4xi64>} + // CHECK-SAME: <{value = dense<[0, 3, 1, 2]> : tensor<4xi64>}> // CHECK: %[[RES_TRANSPOSE:[0-9]*]] = "tf.Transpose"(%y, %[[RES_PERM]]) // CHECK: return %[[RES_TRANSPOSE]] @@ -74,10 +74,10 @@ func.func @transposeFusedBatchNormV3( // CHECK-LABEL: bias_add_nchw func.func @bias_add_nchw(%arg0: tensor<1x256x150x150xf32>, %arg1: tensor<256xf32>) -> tensor<1x256x150x150xf32> { // CHECK: (%[[ARG0:.*]]: tensor<1x256x150x150xf32>, %[[ARG1:.*]]: tensor<256xf32>) - // CHECK: %[[CST:.*]] = "tf.Const"() {value = dense<[0, 2, 3, 1]> : tensor<4xi64>} + // CHECK: %[[CST:.*]] = "tf.Const"() <{value = dense<[0, 2, 3, 1]> : tensor<4xi64>}> // CHECK: %[[R0:.*]] = "tf.Transpose"(%[[ARG0]], %[[CST]]) - // CHECK: %[[R1:.*]] = "tf.BiasAdd"(%[[R0]], %[[ARG1]]) {data_format = "NHWC", device = ""} - // CHECK: %[[CST_0:.*]] = "tf.Const"() {value = dense<[0, 3, 1, 2]> : tensor<4xi64>} + // CHECK: %[[R1:.*]] = "tf.BiasAdd"(%[[R0]], %[[ARG1]]) <{data_format = "NHWC"}> {device = ""} + // CHECK: %[[CST_0:.*]] = "tf.Const"() <{value = dense<[0, 3, 1, 2]> : tensor<4xi64>}> // CHECK: "tf.Transpose"(%[[R1]], %[[CST_0]]) %0 = "tf.BiasAdd"(%arg0, %arg1) {data_format = "NCHW", device = ""} : (tensor<1x256x150x150xf32>, tensor<256xf32>) -> tensor<1x256x150x150xf32> func.return %0 : tensor<1x256x150x150xf32> @@ -85,10 +85,10 @@ func.func @bias_add_nchw(%arg0: tensor<1x256x150x150xf32>, %arg1: tensor<256xf32 // CHECK-LABEL: maxpool_nchw func.func @maxpool_nchw(%arg0: tensor<1x64x112x112xf32>) -> tensor<1x64x56x56xf32> { - // CHECK: %[[CST:.*]] = "tf.Const"() {value = dense<[0, 2, 3, 1]> : tensor<4xi64>} + // CHECK: %[[CST:.*]] = "tf.Const"() <{value = dense<[0, 2, 3, 1]> : tensor<4xi64>}> // CHECK: %[[R0:.*]] = "tf.Transpose"(%arg0, %[[CST]]) - // CHECK: %[[R1:.*]] = "tf.MaxPool"(%[[R0]]) {data_format = "NHWC", explicit_paddings = [], ksize = [1, 3, 3, 1], padding = "SAME", strides = [1, 2, 2, 1]} - // CHECK: %[[CST_0:.*]] = "tf.Const"() {value = dense<[0, 3, 1, 2]> : tensor<4xi64>} + // CHECK: %[[R1:.*]] = "tf.MaxPool"(%[[R0]]) <{data_format = "NHWC", explicit_paddings = [], ksize = [1, 3, 3, 1], padding = "SAME", strides = [1, 2, 2, 1]}> + // CHECK: %[[CST_0:.*]] = "tf.Const"() <{value = dense<[0, 3, 1, 2]> : tensor<4xi64>}> // CHECK: "tf.Transpose"(%[[R1]], %[[CST_0]]) %0 = "tf.MaxPool"(%arg0) { diff --git a/tensorflow/compiler/mlir/tensorflow/tests/layout_optimization_move_transposes_begin.mlir b/tensorflow/compiler/mlir/tensorflow/tests/layout_optimization_move_transposes_begin.mlir index be36e2e13cb7db..be511f962e26a6 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/layout_optimization_move_transposes_begin.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/layout_optimization_move_transposes_begin.mlir @@ -3,7 +3,7 @@ // CHECK-LABEL: func @move_across_single_op func.func @move_across_single_op(%arg0: tensor<1x4x4x8xf32>) -> tensor<1x8x4x4xf32> { - // CHECK: %[[ARG_PERM:.*]] = "tf.Const"() {value = dense<[0, 3, 1, 2]> : tensor<4xi32>} + // CHECK: %[[ARG_PERM:.*]] = "tf.Const"() <{value = dense<[0, 3, 1, 2]> : tensor<4xi32>}> // CHECK: %[[ARG_TRANSPOSE:[0-9]*]] = "tf.Transpose"(%arg0, %[[ARG_PERM]]) // CHECK: %[[TANH:[0-9]*]] = "tf.Tanh"(%[[ARG_TRANSPOSE]]) {{.*}} tensor<1x8x4x4xf32> // CHECK: return %[[TANH]] @@ -18,7 +18,7 @@ func.func @move_across_single_op(%arg0: tensor<1x4x4x8xf32>) -> tensor<1x8x4x4xf // CHECK-LABEL: func @move_across_multiple_ops func.func @move_across_multiple_ops(%arg0: tensor<1x4x4x8xf32>) -> tensor<1x8x4x4xf32> { - // CHECK: %[[ARG_PERM:.*]] = "tf.Const"() {value = dense<[0, 3, 1, 2]> : tensor<4xi32>} + // CHECK: %[[ARG_PERM:.*]] = "tf.Const"() <{value = dense<[0, 3, 1, 2]> : tensor<4xi32>}> // CHECK: %[[ARG_TRANSPOSE:[0-9]*]] = "tf.Transpose"(%arg0, %[[ARG_PERM]]) // CHECK: %[[TANH:[0-9]*]] = "tf.Tanh"(%[[ARG_TRANSPOSE]]) {{.*}} tensor<1x8x4x4xf32> // CHECK: %[[RELU:[0-9]*]] = "tf.Relu"(%[[TANH]]) {{.*}} tensor<1x8x4x4xf32> @@ -36,7 +36,7 @@ func.func @move_across_multiple_ops(%arg0: tensor<1x4x4x8xf32>) -> tensor<1x8x4x // CHECK-LABEL: func @move_across_multi_operand_op func.func @move_across_multi_operand_op(%arg0: tensor<1x4x4x8xf32>, %arg1: tensor<1x4x4x8xf32>) -> tensor<1x8x4x4xf32> { - // CHECK: %[[ARG_PERM:.*]] = "tf.Const"() {value = dense<[0, 3, 1, 2]> : tensor<4xi32>} + // CHECK: %[[ARG_PERM:.*]] = "tf.Const"() <{value = dense<[0, 3, 1, 2]> : tensor<4xi32>}> // CHECK: %[[ARG0_TRANSPOSE:[0-9]*]] = "tf.Transpose"(%arg0, %[[ARG_PERM]]) // CHECK: %[[ARG1_TRANSPOSE:[0-9]*]] = "tf.Transpose"(%arg1, %[[ARG_PERM]]) // CHECK: %[[ADD:[0-9]*]] = "tf.AddV2"(%[[ARG0_TRANSPOSE]], %[[ARG1_TRANSPOSE]]) {{.*}} tensor<1x8x4x4xf32> @@ -52,7 +52,7 @@ func.func @move_across_multi_operand_op(%arg0: tensor<1x4x4x8xf32>, %arg1: tenso // CHECK-LABEL: func @move_with_multiple_uses func.func @move_with_multiple_uses(%arg0: tensor<1x4x4x8xf32>) -> tensor<1x8x4x4xf32> { - // CHECK: %[[ARG_PERM:.*]] = "tf.Const"() {value = dense<[0, 3, 1, 2]> : tensor<4xi32>} + // CHECK: %[[ARG_PERM:.*]] = "tf.Const"() <{value = dense<[0, 3, 1, 2]> : tensor<4xi32>}> // CHECK: %[[ARG_TRANSPOSE:[0-9]*]] = "tf.Transpose"(%arg0, %[[ARG_PERM]]) // CHECK: %[[TANH:[0-9]*]] = "tf.Tanh"(%[[ARG_TRANSPOSE]]) {{.*}} tensor<1x8x4x4xf32> // CHECK: %[[ADD:[0-9]*]] = "tf.AddV2"(%[[TANH]], %[[TANH]]) {{.*}} tensor<1x8x4x4xf32> @@ -78,9 +78,9 @@ func.func @move_transpose_handle_broadcast(%arg0:tensor<8x64xf32>, %arg1:tensor< func.return %3 : tensor<512x64xf32> - // CHECK-DAG: %[[CST_0:.*]] = "tf.Const"() {value = dense<[2, 0, 1]> : tensor<3xi32>} : () -> tensor<3xi32> - // CHECK-DAG: %[[CST_1:.*]] = "tf.Const"() {value = dense<3> : tensor} : () -> tensor - // CHECK-DAG: %[[CST_2:.*]] = "tf.Const"() {value = dense<[512, 64]> : tensor<2xi32>} : () -> tensor<2xi32> + // CHECK-DAG: %[[CST_0:.*]] = "tf.Const"() <{value = dense<[2, 0, 1]> : tensor<3xi32>}> : () -> tensor<3xi32> + // CHECK-DAG: %[[CST_1:.*]] = "tf.Const"() <{value = dense<3> : tensor}> : () -> tensor + // CHECK-DAG: %[[CST_2:.*]] = "tf.Const"() <{value = dense<[512, 64]> : tensor<2xi32>}> : () -> tensor<2xi32> // CHECK: %[[EXPAND_DIMS:.*]] = "tf.ExpandDims"(%arg0, %[[CST_1]]) {device = ""} : (tensor<8x64xf32>, tensor) -> tensor<8x64x1xf32> // CHECK: %[[TRANSPOSE_1:.*]] = "tf.Transpose"(%[[EXPAND_DIMS]], %[[CST_0]]) : (tensor<8x64x1xf32>, tensor<3xi32>) -> tensor<1x8x64xf32> // CHECK: %[[TRANSPOSE_2:.*]] = "tf.Transpose"(%arg1, %[[CST_0]]) : (tensor<8x64x64xf32>, tensor<3xi32>) -> tensor<64x8x64xf32> @@ -97,7 +97,7 @@ func.func @dont_move_transpose_different_ranks(%arg0:tensor<1x1x2x3xf32>, %arg1: func.return %1 : tensor<1x2x1x3xf32> - // CHECK: %[[CST:.*]] = "tf.Const"() {value = dense<[0, 2, 1, 3]> : tensor<4xi32>} : () -> tensor<4xi32> + // CHECK: %[[CST:.*]] = "tf.Const"() <{value = dense<[0, 2, 1, 3]> : tensor<4xi32>}> : () -> tensor<4xi32> // CHECK: %[[ADD:.*]] = "tf.AddV2"(%arg0, %arg1) {device = ""} : (tensor<1x1x2x3xf32>, tensor<2x3xf32>) -> tensor<1x1x2x3xf32> // CHECK: %[[TRANSPOSE:.*]] = "tf.Transpose"(%[[ADD]], %[[CST]]) {device = ""} : (tensor<1x1x2x3xf32>, tensor<4xi32>) -> tensor<1x2x1x3xf32> // CHECK: return %[[TRANSPOSE]] : tensor<1x2x1x3xf32> diff --git a/tensorflow/compiler/mlir/tensorflow/tests/layout_optimization_move_transposes_end.mlir b/tensorflow/compiler/mlir/tensorflow/tests/layout_optimization_move_transposes_end.mlir index 20bf6d65b9ab1d..0bc9a131cfab09 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/layout_optimization_move_transposes_end.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/layout_optimization_move_transposes_end.mlir @@ -4,7 +4,7 @@ // CHECK-LABEL: func @move_across_single_op func.func @move_across_single_op(%arg0: tensor<1x4x4x8xf32>) -> tensor<1x8x4x4xf32> { - // CHECK: %[[RES_PERM:.*]] = "tf.Const"() {value = dense<[0, 3, 1, 2]> : tensor<4xi32>} + // CHECK: %[[RES_PERM:.*]] = "tf.Const"() <{value = dense<[0, 3, 1, 2]> : tensor<4xi32>}> // CHECK: %[[TANH:[0-9]*]] = "tf.Tanh"(%arg0) {{.*}} tensor<1x4x4x8xf32> // CHECK: %[[RES_TRANSPOSE:[0-9]*]] = "tf.Transpose"(%[[TANH]], %[[RES_PERM]]) {{.*}} tensor<1x8x4x4xf32> // CHECK: return %[[RES_TRANSPOSE]] @@ -19,7 +19,7 @@ func.func @move_across_single_op(%arg0: tensor<1x4x4x8xf32>) -> tensor<1x8x4x4xf // CHECK-LABEL: func @move_across_multiple_ops func.func @move_across_multiple_ops(%arg0: tensor<1x4x4x8xf32>) -> tensor<1x8x4x4xf32> { - // CHECK: %[[RES_PERM:.*]] = "tf.Const"() {value = dense<[0, 3, 1, 2]> : tensor<4xi32>} + // CHECK: %[[RES_PERM:.*]] = "tf.Const"() <{value = dense<[0, 3, 1, 2]> : tensor<4xi32>}> // CHECK: %[[TANH:[0-9]*]] = "tf.Tanh"(%arg0) {{.*}} tensor<1x4x4x8xf32> // CHECK: %[[RELU:[0-9]*]] = "tf.Relu"(%[[TANH]]) {{.*}} tensor<1x4x4x8xf32> // CHECK: %[[RES_TRANSPOSE:[0-9]*]] = "tf.Transpose"(%[[RELU]], %[[RES_PERM]]) @@ -36,7 +36,7 @@ func.func @move_across_multiple_ops(%arg0: tensor<1x4x4x8xf32>) -> tensor<1x8x4x // CHECK-LABEL: func @move_across_multi_operand_op func.func @move_across_multi_operand_op(%arg0: tensor<1x4x4x8xf32>, %arg1: tensor<1x4x4x8xf32>) -> tensor<1x8x4x4xf32> { - // CHECK: %[[RES_PERM:.*]] = "tf.Const"() {value = dense<[0, 3, 1, 2]> : tensor<4xi32>} + // CHECK: %[[RES_PERM:.*]] = "tf.Const"() <{value = dense<[0, 3, 1, 2]> : tensor<4xi32>}> // CHECK: %[[ADD:[0-9]*]] = "tf.AddV2"(%arg0, %arg1) {{.*}} tensor<1x4x4x8xf32> // CHECK: %[[RES_TRANSPOSE:[0-9]*]] = "tf.Transpose"(%[[ADD]], %[[RES_PERM]]) // CHECK: return %[[RES_TRANSPOSE]] @@ -52,7 +52,7 @@ func.func @move_across_multi_operand_op(%arg0: tensor<1x4x4x8xf32>, %arg1: tenso // CHECK-LABEL: func @move_across_broadcastable_op func.func @move_across_broadcastable_op(%arg0: tensor<1x4x1x8xf32>, %arg1: tensor<1x4x4x8xf32>) -> tensor<1x8x4x4xf32> { - // CHECK: %[[RES_PERM:.*]] = "tf.Const"() {value = dense<[0, 3, 1, 2]> : tensor<4xi32>} + // CHECK: %[[RES_PERM:.*]] = "tf.Const"() <{value = dense<[0, 3, 1, 2]> : tensor<4xi32>}> // CHECK: %[[ADD:[0-9]*]] = "tf.AddV2"(%arg0, %arg1) : (tensor<1x4x1x8xf32>, tensor<1x4x4x8xf32>) -> tensor<1x4x4x8xf32> // CHECK: %[[RES_TRANSPOSE:[0-9]*]] = "tf.Transpose"(%[[ADD]], %[[RES_PERM]]) // CHECK: return %[[RES_TRANSPOSE]] @@ -68,7 +68,7 @@ func.func @move_across_broadcastable_op(%arg0: tensor<1x4x1x8xf32>, %arg1: tenso // CHECK-LABEL: func @move_across_double_transpose func.func @move_across_double_transpose(%arg0: tensor<1x4x4x8xf32>, %arg1: tensor<1x4x4x8xf32>) -> tensor<1x4x8x4xf32> { - // CHECK: %[[RES_PERM:.*]] = "tf.Const"() {value = dense<[0, 3, 1, 2]> : tensor<4xi32>} + // CHECK: %[[RES_PERM:.*]] = "tf.Const"() <{value = dense<[0, 3, 1, 2]> : tensor<4xi32>}> // CHECK: %[[ADD:[0-9]*]] = "tf.AddV2"(%arg0, %arg1) : (tensor<1x4x4x8xf32>, tensor<1x4x4x8xf32>) -> tensor<1x4x4x8xf32> // CHECK: %[[RES_TRANSPOSE_0:[0-9]*]] = "tf.Transpose"(%[[ADD]], %[[RES_PERM]]) // CHECK: %[[RES_TRANSPOSE_1:[0-9]*]] = "tf.Transpose"(%[[RES_TRANSPOSE_0]], %[[RES_PERM]]) @@ -90,8 +90,8 @@ func.func @fold_into_max_pool(%arg0: tensor<1x64x112x112xf32>) -> tensor<1x56x56 // MaxPool operand transpose must be folded into the op and MaxPool // must use NCHW data format with updated kernel size and strides. - // CHECK: %[[RES_PERM:.*]] = "tf.Const"() {value = dense<[0, 2, 3, 1]> : tensor<4xi32>} - // CHECK: %[[MAX_POOL:[0-9]*]] = "tf.MaxPool"(%arg0) {data_format = "NCHW", ksize = [1, 1, 3, 3], padding = "SAME", strides = [1, 1, 2, 2]} : (tensor<1x64x112x112xf32>) -> tensor<1x64x56x56xf32> + // CHECK: %[[RES_PERM:.*]] = "tf.Const"() <{value = dense<[0, 2, 3, 1]> : tensor<4xi32>}> + // CHECK: %[[MAX_POOL:[0-9]*]] = "tf.MaxPool"(%arg0) <{data_format = "NCHW", ksize = [1, 1, 3, 3], padding = "SAME", strides = [1, 1, 2, 2]}> : (tensor<1x64x112x112xf32>) -> tensor<1x64x56x56xf32> // CHECK: %[[RES_TRANSPOSE:[0-9]*]] = "tf.Transpose"(%[[MAX_POOL]], %[[RES_PERM]]) // CHECK: return %[[RES_TRANSPOSE]] @@ -112,14 +112,14 @@ func.func @fold_into_max_pool(%arg0: tensor<1x64x112x112xf32>) -> tensor<1x56x56 // CHECK-LABEL: func @fold_into_mean func.func @fold_into_mean(%arg0: tensor<1x64x112x112xf32>) -> tensor<1x64xf32> { - // CHECK: %[[RED_IDX:.*]] = "tf.Const"() {value = dense<[2, 3]> : tensor<2xi32>} + // CHECK: %[[RED_IDX:.*]] = "tf.Const"() <{value = dense<[2, 3]> : tensor<2xi32>}> // CHECK: %[[MEAN:[0-9]*]] = "tf.Mean"(%arg0, %[[RED_IDX]]) // CHECK-SAME: (tensor<1x64x112x112xf32>, tensor<2xi32>) -> tensor<1x64xf32> // CHECK: return %[[MEAN]] - // NOFOLD: %[[CST:.*]] = "tf.Const"() {value = dense<[0, 2, 3, 1]> : tensor<4xi32>} + // NOFOLD: %[[CST:.*]] = "tf.Const"() <{value = dense<[0, 2, 3, 1]> : tensor<4xi32>}> // NOFOLD: %[[TRANSPOSE:[0-9]*]] = "tf.Transpose"(%arg0, %[[CST]]) - // NOFOLD: %[[CST_1:.*]] = "tf.Const"() {value = dense<[1, 2]> : tensor<2xi32>} + // NOFOLD: %[[CST_1:.*]] = "tf.Const"() <{value = dense<[1, 2]> : tensor<2xi32>}> // NOFOLD: %[[MEAN:[0-9]*]] = "tf.Mean"(%[[TRANSPOSE]], %[[CST_1]]) // NOFOLD-SAME: (tensor<1x112x112x64xf32>, tensor<2xi32>) -> tensor<1x64xf32> // NOFOLD: return %[[MEAN]] @@ -138,8 +138,8 @@ func.func @fold_into_mean(%arg0: tensor<1x64x112x112xf32>) -> tensor<1x64xf32> { // CHECK-LABEL: func @fold_into_fused_batch_norm func.func @fold_into_fused_batch_norm(%arg0: tensor<1x64x112x112xf32>, %arg1: tensor<64xf32>) -> tensor<1x112x112x64xf32> { - // CHECK: %[[RES_PERM:.*]] = "tf.Const"() {value = dense<[0, 2, 3, 1]> : tensor<4xi32>} - // CHECK: "tf.FusedBatchNormV3"(%arg0, {{.*}} {data_format = "NCHW" + // CHECK: %[[RES_PERM:.*]] = "tf.Const"() <{value = dense<[0, 2, 3, 1]> : tensor<4xi32>}> + // CHECK: "tf.FusedBatchNormV3"(%arg0, {{.*}} <{data_format = "NCHW" // CHECK: %[[RES_TRANSPOSE:[0-9]*]] = "tf.Transpose"(%y, %[[RES_PERM]]) // CHECK: return %[[RES_TRANSPOSE]] @@ -165,9 +165,9 @@ func.func @fold_into_fused_batch_norm(%arg0: tensor<1x64x112x112xf32>, %arg1: te // CHECK-LABEL: func @fold_into_pad_with_extra_uses func.func @fold_into_pad_with_extra_uses(%arg0: tensor<1x2x4x4x3xf32>) -> (tensor<1x2x3x4x4xf32>, tensor<1x2x3x6x6xf32>) { - // CHECK: %[[PERM:.*]] = "tf.Const"() {value = dense<[0, 1, 4, 2, 3]> : tensor<5xi32>} + // CHECK: %[[PERM:.*]] = "tf.Const"() <{value = dense<[0, 1, 4, 2, 3]> : tensor<5xi32>}> // CHECK: %[[TRANSPOSE_OP:[0-9]*]] = "tf.Transpose"(%arg0, %[[PERM]]) - // CHECK: %[[PADDING:.*]] = "tf.Const"() {value = dense<{{\[\[}}0, 0], [0, 0], [1, 1], [1, 1], [0, 0]]> : tensor<5x2xi32>} + // CHECK: %[[PADDING:.*]] = "tf.Const"() <{value = dense<{{\[\[}}0, 0], [0, 0], [1, 1], [1, 1], [0, 0]]> : tensor<5x2xi32>}> // CHECK: %[[PAD_OP:[0-9]*]] = "tf.Pad"(%arg0, %[[PADDING]]) // CHECK: %[[DUP_TRANSPOSE_OP:[0-9]*]] = "tf.Transpose"(%[[PAD_OP]], %[[PERM]]) // CHECK: return %[[TRANSPOSE_OP]], %[[DUP_TRANSPOSE_OP]] diff --git a/tensorflow/compiler/mlir/tensorflow/tests/layout_optimization_to_nhwc.mlir b/tensorflow/compiler/mlir/tensorflow/tests/layout_optimization_to_nhwc.mlir index 5f9256f6424ab1..e82819686948d6 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/layout_optimization_to_nhwc.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/layout_optimization_to_nhwc.mlir @@ -32,7 +32,7 @@ func.func @transpose_resnet_layer(%arg0: tensor, // input // Shuffled paddings. // CHECK: %[[PADDINGS:.*]] = "tf.Const"(){{.*}}[0, 0], [3, 3], [3, 3], [0, 0] // NOFOLD: %[[PADDING:.*]] = "tf.Const"(){{.*}}[0, 0], [0, 0], [3, 3], [3, 3] - // NOFOLD: %[[CST:.*]] = "tf.Const"() {value = dense<[0, 3, 1, 2]> : tensor<4xi32>} : () -> tensor<4xi32> + // NOFOLD: %[[CST:.*]] = "tf.Const"() <{value = dense<[0, 3, 1, 2]> : tensor<4xi32>}> : () -> tensor<4xi32> // NOFOLD: %[[TRANSPOSE:[0-9]*]] = "tf.Transpose"(%arg0, %[[CST]]) : (tensor, tensor<4xi32>) -> tensor // Pad input with new paddings. @@ -151,7 +151,7 @@ func.func @transpose_resnet_layer(%arg0: tensor, // input %16 = "tf.Mean"(%15, %1) : (tensor, tensor<2xi32>) -> tensor // Mean should compute reduction over NHWC spatial dimensions. - // CHECK: %[[MEAN_DIMS:.*]] = "tf.Const"() {value = dense<[1, 2]> : tensor<2xi32>} + // CHECK: %[[MEAN_DIMS:.*]] = "tf.Const"() <{value = dense<[1, 2]> : tensor<2xi32>}> // CHECK: %[[MEAN:[0-9]*]] = "tf.Mean"(%[[RELU]], %[[MEAN_DIMS]]) // CHECK-SAME: (tensor, tensor<2xi32>) -> tensor // CHECK: return %[[MEAN]] : tensor diff --git a/tensorflow/compiler/mlir/tensorflow/tests/legalize_tfg.mlir b/tensorflow/compiler/mlir/tensorflow/tests/legalize_tfg.mlir index c170a0e41c426f..0ff3cd334fc2aa 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/legalize_tfg.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/legalize_tfg.mlir @@ -4,9 +4,9 @@ module { // CHECK: tf_executor.graph tfg.graph #tf_type.version { - // CHECK: tf_executor.island wraps "tf.VarHandleOp"() {_mlir_name = "x", _output_shapes = [#tf_type.shape<>], allowed_devices = [], container = "a", device = "/device:CPU:0", dtype = i64, shape = #tf_type.shape<>, shared_name = "x"} : () -> tensor>> + // CHECK: tf_executor.island wraps "tf.VarHandleOp"() <{container = "a", shared_name = "x"}> {_mlir_name = "x", _output_shapes = [#tf_type.shape<>], allowed_devices = [], device = "/device:CPU:0", dtype = i64, shape = #tf_type.shape<>} : () -> tensor>> %VarHandleOp, %ctl = VarHandleOp device("/CPU:0") name("x") {_output_shapes = [#tf_type.shape<>], allowed_devices = [], container = "a", dtype = i64, shape = #tf_type.shape<>, shared_name = "x"} : () -> (tensor>>) - // CHECK: tf_executor.island wraps "tf.LegacyCall"(%outputs, %outputs) {_disable_call_shape_inference = true, f = @test_func_name0} : (tensor>>, tensor>>) -> tensor<*x!tf_type.resource> + // CHECK: tf_executor.island wraps "tf.LegacyCall"(%outputs, %outputs) <{_disable_call_shape_inference = true, f = @test_func_name0}> : (tensor>>, tensor>>) -> tensor<*x!tf_type.resource> %test_func_name0, %ctl_0 = test_func_name0(%VarHandleOp, %VarHandleOp) name("called") {_disable_call_shape_inference = true, _output_shapes = [#tf_type.shape<*>]} : (tensor>>, tensor>>) -> (tensor<*x!tf_type.resource>) // CHECK: tf_executor.island wraps "tf._Retval"(%outputs_0) {T = !tf_type.resource, _mlir_name = "func_call", index = 0 : i64} : (tensor<*x!tf_type.resource>) -> () %ctl_1 = _Retval(%test_func_name0) name("func_call") {T = !tf_type.resource, index = 0 : i64} : tensor<*x!tf_type.resource> diff --git a/tensorflow/compiler/mlir/tensorflow/tests/lower_quantized.mlir b/tensorflow/compiler/mlir/tensorflow/tests/lower_quantized.mlir index 11c9704ee0074f..eedc235b10f16e 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/lower_quantized.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/lower_quantized.mlir @@ -2,9 +2,9 @@ // CHECK-LABEL: dequantize func.func @dequantize(%arg0: tensor<2x3x!tf_type.qint8>, %min_range: tensor, %max_range: tensor) -> tensor<2x3xf32> { - // CHECK-DAG: %[[HALF_RANGE:.*]] = "tf.Const"() {value = dense<1.280000e+02> : tensor} - // CHECK-DAG: %[[C255:.*]] = "tf.Const"() {value = dense<2.550000e+02> : tensor} - // CHECK-DAG: %[[CAST:.*]] = "tf.Cast"(%arg0) {Truncate = false} + // CHECK-DAG: %[[HALF_RANGE:.*]] = "tf.Const"() <{value = dense<1.280000e+02> : tensor}> + // CHECK-DAG: %[[C255:.*]] = "tf.Const"() <{value = dense<2.550000e+02> : tensor}> + // CHECK-DAG: %[[CAST:.*]] = "tf.Cast"(%arg0) <{Truncate = false}> // CHECK-DAG: %[[SHIFT:.*]] = "tf.AddV2"(%[[CAST]], %[[HALF_RANGE]]) // CHECK-DAG: %[[DRANGE:.*]] = "tf.Sub"(%arg2, %arg1) // CHECK-DAG: %[[SCALE:.*]] = "tf.Div"(%[[DRANGE]], %[[C255:.*]]) @@ -18,8 +18,8 @@ func.func @dequantize(%arg0: tensor<2x3x!tf_type.qint8>, %min_range: tensor // CHECK-LABEL: dequantize_quint8 func.func @dequantize_quint8(%arg0: tensor<2x3x!tf_type.quint8>, %min_range: tensor, %max_range: tensor) -> tensor<2x3xf32> { - // CHECK-NEXT: %[[C255:.*]] = "tf.Const"() {value = dense<2.550000e+02> : tensor} - // CHECK-NEXT: %[[CAST:.*]] = "tf.Cast"(%arg0) {Truncate = false} + // CHECK-NEXT: %[[C255:.*]] = "tf.Const"() <{value = dense<2.550000e+02> : tensor}> + // CHECK-NEXT: %[[CAST:.*]] = "tf.Cast"(%arg0) <{Truncate = false}> // CHECK-NEXT: %[[DRANGE:.*]] = "tf.Sub"(%arg2, %arg1) // CHECK-NEXT: %[[SCALE:.*]] = "tf.Div"(%[[DRANGE]], %[[C255:.*]]) // CHECK-NEXT: %[[SS:.*]] = "tf.Mul"(%[[CAST]], %[[SCALE]]) @@ -32,15 +32,15 @@ func.func @dequantize_quint8(%arg0: tensor<2x3x!tf_type.quint8>, %min_range: ten // CHECK-LABEL: dequantize_to_bf16 func.func @dequantize_to_bf16(%arg0: tensor<2x3x!tf_type.qint8>, %min_range: tensor, %max_range: tensor) -> tensor<2x3xbf16> { - // CHECK-DAG: %[[HALF_RANGE:.*]] = "tf.Const"() {value = dense<1.280000e+02> : tensor} - // CHECK-DAG: %[[C255:.*]] = "tf.Const"() {value = dense<2.550000e+02> : tensor} - // CHECK-DAG: %[[CAST:.*]] = "tf.Cast"(%arg0) {Truncate = false} + // CHECK-DAG: %[[HALF_RANGE:.*]] = "tf.Const"() <{value = dense<1.280000e+02> : tensor}> + // CHECK-DAG: %[[C255:.*]] = "tf.Const"() <{value = dense<2.550000e+02> : tensor}> + // CHECK-DAG: %[[CAST:.*]] = "tf.Cast"(%arg0) <{Truncate = false}> // CHECK-DAG: %[[SHIFT:.*]] = "tf.AddV2"(%[[CAST]], %[[HALF_RANGE]]) // CHECK-DAG: %[[DRANGE:.*]] = "tf.Sub"(%arg2, %arg1) // CHECK-DAG: %[[SCALE:.*]] = "tf.Div"(%[[DRANGE]], %[[C255:.*]]) // CHECK-DAG: %[[SS:.*]] = "tf.Mul"(%[[SHIFT]], %[[SCALE]]) // CHECK-DAG: %[[F32_RESULT:.*]] = "tf.AddV2"(%[[SS]], %arg1) - // CHECK-DAG: %[[RESULT:.*]] = "tf.Cast"(%[[F32_RESULT]]) {Truncate = false} + // CHECK-DAG: %[[RESULT:.*]] = "tf.Cast"(%[[F32_RESULT]]) <{Truncate = false}> %0 = "tf.Dequantize"(%arg0, %min_range, %max_range) : (tensor<2x3x!tf_type.qint8>, tensor, tensor) -> tensor<2x3xbf16> // CHECK-DAG: return %[[RESULT]] diff --git a/tensorflow/compiler/mlir/tensorflow/tests/lower_tf.mlir b/tensorflow/compiler/mlir/tensorflow/tests/lower_tf.mlir index 432195aed3b4d3..83f0b56e398c17 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/lower_tf.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/lower_tf.mlir @@ -2,11 +2,11 @@ // CHECK-LABEL: invert_permutation func.func @invert_permutation(%arg0: tensor<5xi32>) -> tensor<5xi32> { - // CHECK-DAG: %[[UPDATES:.*]] = "tf.Const"() {value = dense<[0, 1, 2, 3, 4]> : tensor<5xi32>} : () -> tensor<5xi32> - // CHECK-DAG: %[[SHAPE:.*]] = "tf.Const"() {value = dense<[5, 1]> : tensor<2xi32>} : () -> tensor<2xi32> - // CHECK-DAG: %[[cst_1:.*]] = "tf.Const"() {value = dense<1> : tensor} : () -> tensor - // CHECK-DAG: %[[cst_2:.*]] = "tf.Const"() {value = dense<1> : tensor<5xi32>} : () -> tensor<5xi32> - // CHECK-DAG: %[[cst_3:.*]] = "tf.Const"() {value = dense<0> : tensor<5xi32>} : () -> tensor<5xi32> + // CHECK-DAG: %[[UPDATES:.*]] = "tf.Const"() <{value = dense<[0, 1, 2, 3, 4]> : tensor<5xi32>}> : () -> tensor<5xi32> + // CHECK-DAG: %[[SHAPE:.*]] = "tf.Const"() <{value = dense<[5, 1]> : tensor<2xi32>}> : () -> tensor<2xi32> + // CHECK-DAG: %[[cst_1:.*]] = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor + // CHECK-DAG: %[[cst_2:.*]] = "tf.Const"() <{value = dense<1> : tensor<5xi32>}> : () -> tensor<5xi32> + // CHECK-DAG: %[[cst_3:.*]] = "tf.Const"() <{value = dense<0> : tensor<5xi32>}> : () -> tensor<5xi32> // CHECK-DAG: %[[INDICES:.*]] = "tf.Reshape"(%arg0, %[[SHAPE]]) : (tensor<5xi32>, tensor<2xi32>) -> tensor<5x1xi32> // CHECK-DAG: %[[INDICES_1:.*]] = "tf.TensorScatterAdd"(%[[cst_3]], %[[INDICES]], %[[cst_2]]) : (tensor<5xi32>, tensor<5x1xi32>, tensor<5xi32>) -> tensor<5xi32> @@ -35,7 +35,7 @@ func.func @invert_permutation_unranked(%arg0: tensor<*xi32>) -> tensor<*xi32> { // CHECK-LABEL: simple_pack // CHECK-SAME: %[[ARG0:.*]]: tensor<3x5xf32>, %[[ARG1:.*]]: tensor<3x5xf32> func.func @simple_pack(%arg0: tensor<3x5xf32>, %arg1: tensor<3x5xf32>) -> tensor<2x3x5xf32> { - // CHECK: %[[AXIS:.*]] = "tf.Const"() {value = dense<0> : tensor} + // CHECK: %[[AXIS:.*]] = "tf.Const"() <{value = dense<0> : tensor}> // CHECK: %[[INP0:.*]] = "tf.ExpandDims"(%[[ARG0]], %[[AXIS]]) : (tensor<3x5xf32>, tensor) -> tensor<1x3x5xf32> // CHECK: %[[INP1:.*]] = "tf.ExpandDims"(%[[ARG1]], %[[AXIS]]) : (tensor<3x5xf32>, tensor) -> tensor<1x3x5xf32> // CHECK: "tf.ConcatV2"(%[[INP0]], %[[INP1]], %[[AXIS]]) : (tensor<1x3x5xf32>, tensor<1x3x5xf32>, tensor) -> tensor<2x3x5xf32> @@ -71,8 +71,8 @@ func.func @squared_difference_complex(%arg0: tensor<3xcomplex>, %arg1: tens // CHECK-LABEL: func @div_no_nan // CHECK-SAME: (%[[X:.*]]: tensor<*xf32>, %[[Y:.*]]: tensor<*xf32>) func.func @div_no_nan(%arg0: tensor<*xf32>, %arg1: tensor<*xf32>) -> tensor<*xf32> { - // CHECK: %[[ZERO:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor} : () -> tensor - // CHECK: %[[IS_ZERO:.*]] = "tf.Equal"(%[[Y]], %[[ZERO]]) {incompatible_shape_error = true} : (tensor<*xf32>, tensor) -> tensor<*xi1> + // CHECK: %[[ZERO:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor}> : () -> tensor + // CHECK: %[[IS_ZERO:.*]] = "tf.Equal"(%[[Y]], %[[ZERO]]) <{incompatible_shape_error = true}> : (tensor<*xf32>, tensor) -> tensor<*xi1> // CHECK: %[[DIV:.*]] = "tf.Div"(%[[X]], %[[Y]]) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32> // CHECK: %[[RESULT:.*]] = "tf.SelectV2"(%[[IS_ZERO]], %[[ZERO]], %[[DIV]]) : (tensor<*xi1>, tensor, tensor<*xf32>) -> tensor<*xf32> %0 = "tf.DivNoNan"(%arg0, %arg1) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32> @@ -96,7 +96,7 @@ func.func @truncate_div_int(%arg0: tensor<*xi32>, %arg1: tensor<*xi32>) // CHECK-SAME: (%[[LHS:.*]]: tensor<*xf32>, %[[RHS:.*]]: tensor<*xf32>) func.func @truncate_div_float(%arg0: tensor<*xf32>, %arg1: tensor<*xf32>) -> tensor<*xf32> { - // CHECK: %[[ZERO:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor} : () -> tensor + // CHECK: %[[ZERO:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor}> : () -> tensor // CHECK: %[[XDIVY:.*]] = "tf.Div"(%[[LHS]], %[[RHS]]) // CHECK: %[[MASK:.*]] = "tf.Less"(%[[XDIVY]], %[[ZERO]]) // CHECK: %[[CEIL:.*]] = "tf.Ceil"(%[[XDIVY]]) @@ -112,8 +112,8 @@ func.func @truncate_div_float(%arg0: tensor<*xf32>, %arg1: tensor<*xf32>) // CHECK-LABEL: func @mul_no_nan // CHECK-SAME: (%[[X:.*]]: tensor<2x3xf32>, %[[Y:.*]]: tensor<3xf32>) func.func @mul_no_nan(%arg0: tensor<2x3xf32>, %arg1: tensor<3xf32>) -> tensor<2x3xf32> { - // CHECK: %[[ZERO:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor} : () -> tensor - // CHECK: %[[IS_ZERO:.*]] = "tf.Equal"(%[[Y]], %[[ZERO]]) {incompatible_shape_error = true} : (tensor<3xf32>, tensor) -> tensor<3xi1> + // CHECK: %[[ZERO:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor}> : () -> tensor + // CHECK: %[[IS_ZERO:.*]] = "tf.Equal"(%[[Y]], %[[ZERO]]) <{incompatible_shape_error = true}> : (tensor<3xf32>, tensor) -> tensor<3xi1> // CHECK: %[[MUL:.*]] = "tf.Mul"(%[[X]], %[[Y]]) : (tensor<2x3xf32>, tensor<3xf32>) -> tensor<2x3xf32> // CHECK: %[[RESULT:.*]] = "tf.SelectV2"(%[[IS_ZERO]], %[[ZERO]], %[[MUL]]) : (tensor<3xi1>, tensor, tensor<2x3xf32>) -> tensor<2x3xf32> %0 = "tf.MulNoNan"(%arg0, %arg1) : (tensor<2x3xf32>, tensor<3xf32>) -> tensor<2x3xf32> @@ -124,9 +124,9 @@ func.func @mul_no_nan(%arg0: tensor<2x3xf32>, %arg1: tensor<3xf32>) -> tensor<2x // CHECK-LABEL: @is_inf func.func @is_inf(%arg0: tensor<3x4xf32>) -> tensor<3x4xi1> { - // CHECK: %[[INF:.*]] = "tf.Const"() {value = dense<0x7F800000> : tensor} : () -> tensor + // CHECK: %[[INF:.*]] = "tf.Const"() <{value = dense<0x7F800000> : tensor}> : () -> tensor // CHECK: %[[ABS:.*]] = "tf.Abs"(%arg0) : (tensor<3x4xf32>) -> tensor<3x4xf32> - // CHECK: %[[RESULT:.*]] = "tf.Equal"(%[[ABS]], %[[INF]]) {incompatible_shape_error = true} : (tensor<3x4xf32>, tensor) -> tensor<3x4xi1> + // CHECK: %[[RESULT:.*]] = "tf.Equal"(%[[ABS]], %[[INF]]) <{incompatible_shape_error = true}> : (tensor<3x4xf32>, tensor) -> tensor<3x4xi1> %0 = "tf.IsInf"(%arg0) : (tensor<3x4xf32>) -> tensor<3x4xi1> // CHECK: return %[[RESULT]] func.return %0 : tensor<3x4xi1> @@ -134,7 +134,7 @@ func.func @is_inf(%arg0: tensor<3x4xf32>) -> tensor<3x4xi1> { // CHECK-LABEL: @is_nan func.func @is_nan(%arg0: tensor<3x4xf32>) -> tensor<3x4xi1> { - // CHECK: %[[RESULT:.*]] = "tf.NotEqual"(%arg0, %arg0) {incompatible_shape_error = true} : (tensor<3x4xf32>, tensor<3x4xf32>) -> tensor<3x4xi1> + // CHECK: %[[RESULT:.*]] = "tf.NotEqual"(%arg0, %arg0) <{incompatible_shape_error = true}> : (tensor<3x4xf32>, tensor<3x4xf32>) -> tensor<3x4xi1> %0 = "tf.IsNan"(%arg0) : (tensor<3x4xf32>) -> tensor<3x4xi1> // CHECK: return %[[RESULT]] func.return %0 : tensor<3x4xi1> @@ -149,7 +149,7 @@ func.func @fill(%arg0: tensor<*xi64>, %arg1: tensor<*xf32>) -> tensor<*xf32> { } func.func @empty(%arg0: tensor) -> tensor<*xf32> { - // CHECK-DAG: [[CST:%.+]] = "tf.Const"() {value = dense<0.000000e+00> : tensor} + // CHECK-DAG: [[CST:%.+]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor}> // CHECK-DAG: [[RES:%.+]] = "tf.BroadcastTo"([[CST]], %arg0) %0 = "tf.Empty"(%arg0) {init = true} : (tensor) -> (tensor<*xf32>) @@ -162,9 +162,9 @@ func.func @empty(%arg0: tensor) -> tensor<*xf32> { func.func @l2_loss(%arg0: tensor) -> tensor { // CHECK-DAG: %[[SQUARE:.*]] = "tf.Mul"(%[[INPUT]], %[[INPUT]]) : (tensor, tensor) -> tensor - // CHECK-DAG: %[[REDUCE_AXES:.*]] = "tf.Const"() {value = dense<[0, 1]> : tensor<2xi64>} - // CHECK-DAG: %[[SUM:.*]] = "tf.Sum"(%[[SQUARE]], %[[REDUCE_AXES]]) {keep_dims = false} : (tensor, tensor<2xi64>) -> tensor - // CHECK-DAG: %[[TWO:.*]] = "tf.Const"() {value = dense<2.000000e+00> : tensor} + // CHECK-DAG: %[[REDUCE_AXES:.*]] = "tf.Const"() <{value = dense<[0, 1]> : tensor<2xi64>}> + // CHECK-DAG: %[[SUM:.*]] = "tf.Sum"(%[[SQUARE]], %[[REDUCE_AXES]]) <{keep_dims = false}> : (tensor, tensor<2xi64>) -> tensor + // CHECK-DAG: %[[TWO:.*]] = "tf.Const"() <{value = dense<2.000000e+00> : tensor}> // CHECK-DAG: %[[LOSS:.*]] = "tf.Div"(%[[SUM]], %[[TWO]]) : (tensor, tensor) -> tensor %0 = "tf.L2Loss"(%arg0) : (tensor) -> tensor @@ -183,7 +183,7 @@ func.func @l2_loss_unranked(%arg0: tensor<*xf32>) -> tensor { // CHECK-LABEL: pack_with_unranked // CHECK-SAME: %[[ARG0:.*]]: tensor, %[[ARG1:.*]]: tensor<*xf32> func.func @pack_with_unranked(%arg0: tensor, %arg1: tensor<*xf32>) -> tensor<*xf32> { - // CHECK: %[[AXIS:.*]] = "tf.Const"() {value = dense<-2> : tensor} + // CHECK: %[[AXIS:.*]] = "tf.Const"() <{value = dense<-2> : tensor}> // CHECK: %[[INP0:.*]] = "tf.ExpandDims"(%[[ARG0]], %[[AXIS]]) : (tensor, tensor) -> tensor // CHECK: %[[INP1:.*]] = "tf.ExpandDims"(%[[ARG1]], %[[AXIS]]) : (tensor<*xf32>, tensor) -> tensor<*xf32> // CHECK: "tf.ConcatV2"(%[[INP0]], %[[INP1]], %[[AXIS]]) : (tensor, tensor<*xf32>, tensor) -> tensor<*xf32> @@ -196,7 +196,7 @@ func.func @pack_with_unranked(%arg0: tensor, %arg1: tensor<*xf32>) -> t func.func @pad(%arg0: tensor<3xf32>) -> tensor<6xf32> { %padding = "tf.Const"() { value = dense<[[1, 2]]> : tensor<1x2xi64> } : () -> tensor<1x2xi64> // CHECK-DAG: [[PAD:%.+]] = "tf.Const"() {{.+}} -> tensor<1x2xi64> - // CHECK-DAG: [[CST:%.+]] = "tf.Const"() {value = dense<0.000000e+00> : tensor} + // CHECK-DAG: [[CST:%.+]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor}> // CHECK: "tf.PadV2"(%arg0, [[PAD]], [[CST]]) %0 = "tf.Pad"(%arg0, %padding) : (tensor<3xf32>, tensor<1x2xi64>) -> tensor<6xf32> func.return %0 : tensor<6xf32> @@ -206,7 +206,7 @@ func.func @pad(%arg0: tensor<3xf32>) -> tensor<6xf32> { func.func @pad_bf16(%arg0: tensor<3xbf16>) -> tensor<6xbf16> { %padding = "tf.Const"() { value = dense<[[1, 2]]> : tensor<1x2xi64> } : () -> tensor<1x2xi64> // CHECK-DAG: [[PAD:%.+]] = "tf.Const"() {{.+}} -> tensor<1x2xi64> - // CHECK-DAG: [[CST:%.+]] = "tf.Const"() {value = dense<0.000000e+00> : tensor} + // CHECK-DAG: [[CST:%.+]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor}> // CHECK: "tf.PadV2"(%arg0, [[PAD]], [[CST]]) %0 = "tf.Pad"(%arg0, %padding) : (tensor<3xbf16>, tensor<1x2xi64>) -> tensor<6xbf16> func.return %0 : tensor<6xbf16> @@ -221,8 +221,8 @@ func.func @add_f32(%arg0: tensor<3xf32>, %arg1: tensor<3xf32>) -> tensor<3xf32> // CHECK-LABEL: func @BiasAddGrad_NHWC func.func @BiasAddGrad_NHWC(%arg0: tensor<2x3x4x5xf32>) -> tensor<5xf32> { - // CHECK: "tf.Const"() {value = dense<[0, 1, 2]> : tensor<3xi64>} - // CHECK: "tf.Sum"({{.*}}) {keep_dims = false} + // CHECK: "tf.Const"() <{value = dense<[0, 1, 2]> : tensor<3xi64>}> + // CHECK: "tf.Sum"({{.*}}) <{keep_dims = false}> %0 = "tf.BiasAddGrad"(%arg0) {data_format = "NHWC"} : (tensor<2x3x4x5xf32>) -> tensor<5xf32> func.return %0 : tensor<5xf32> @@ -230,8 +230,8 @@ func.func @BiasAddGrad_NHWC(%arg0: tensor<2x3x4x5xf32>) -> tensor<5xf32> { // CHECK-LABEL: func @BiasAddGrad_NCHW func.func @BiasAddGrad_NCHW(%arg0: tensor<2x3x4x5xf32>) -> tensor<3xf32> { - // CHECK: "tf.Const"() {value = dense<[0, 2, 3]> : tensor<3xi64>} - // CHECK: "tf.Sum"({{.*}}) {keep_dims = false} + // CHECK: "tf.Const"() <{value = dense<[0, 2, 3]> : tensor<3xi64>}> + // CHECK: "tf.Sum"({{.*}}) <{keep_dims = false}> %0 = "tf.BiasAddGrad"(%arg0) {data_format = "NCHW"} : (tensor<2x3x4x5xf32>) -> tensor<3xf32> func.return %0 : tensor<3xf32> @@ -254,7 +254,7 @@ func.func @BiasAddGrad_unranked(%arg0: tensor<*xf32>) -> tensor { // CHECK-LABEL: func @rsqrt_grad // CHECK-SAME: (%[[ARG0:.*]]: tensor<2xf32>, %[[ARG1:.*]]: tensor<2xf32>) func.func @rsqrt_grad(%arg0: tensor<2xf32>, %arg1: tensor<2xf32>) -> tensor<2xf32> { - // CHECK: %[[CST:.*]] = "tf.Const"() {value = dense<-2.000000e+00> : tensor} + // CHECK: %[[CST:.*]] = "tf.Const"() <{value = dense<-2.000000e+00> : tensor}> // CHECK: %[[LHS2:.*]] = "tf.Mul"(%[[ARG0]], %[[ARG0]]) // CHECK: %[[LHS3:.*]] = "tf.Mul"(%[[LHS2]], %[[ARG0]]) // CHECK: %[[DIV:.*]] = "tf.Div"(%[[ARG1]], %[[CST]]) @@ -279,7 +279,7 @@ func.func @rsqrt_grad_unranked(%arg0: tensor<*xf32>, %arg1: tensor<*xf32>) -> te // CHECK-LABEL: func @sqrt_grad_unranked // CHECK-SAME: (%[[ARG0:.*]]: tensor<*xcomplex>, %[[ARG1:.*]]: tensor<*xcomplex>) func.func @sqrt_grad_unranked(%arg0: tensor<*xcomplex>, %arg1: tensor<*xcomplex>) -> tensor<*xcomplex> { - // CHECK: %[[CST:.*]] = "tf.Const"() {value = dense<(5.000000e-01,0.000000e+00)> : tensor>} : () -> tensor> + // CHECK: %[[CST:.*]] = "tf.Const"() <{value = dense<(5.000000e-01,0.000000e+00)> : tensor>}> : () -> tensor> // CHECK: %[[MUL:.*]] = "tf.Mul"(%arg1, %[[CST]]) : (tensor<*xcomplex>, tensor>) -> tensor<*xcomplex> // CHECK: %[[RET:.*]] = "tf.Div"(%[[MUL]], %arg0) : (tensor<*xcomplex>, tensor<*xcomplex>) -> tensor<*xcomplex> @@ -292,22 +292,22 @@ func.func @sqrt_grad_unranked(%arg0: tensor<*xcomplex>, %arg1: tensor<*xcom // dimension. // CHECK-LABEL: fourdim_space_to_batch_nd func.func @fourdim_space_to_batch_nd(%input: tensor<3x5x7x10xf32>, %block_shape: tensor<2xi64>, %paddings: tensor<2x2xi64>) -> tensor { - // CHECK-DAG: [[PAD00:%.+]] = "tf.Const"() {value = dense<0> : tensor<1x2xi64>} - // CHECK-DAG: [[ZERO_I32:%.+]] = "tf.Const"() {value = dense<0> : tensor} - // CHECK-DAG: [[ZERO_I64:%.+]] = "tf.Const"() {value = dense<0> : tensor} + // CHECK-DAG: [[PAD00:%.+]] = "tf.Const"() <{value = dense<0> : tensor<1x2xi64>}> + // CHECK-DAG: [[ZERO_I32:%.+]] = "tf.Const"() <{value = dense<0> : tensor}> + // CHECK-DAG: [[ZERO_I64:%.+]] = "tf.Const"() <{value = dense<0> : tensor}> // CHECK-DAG: [[FULL_PADDINGS:%.+]] = "tf.ConcatV2"([[PAD00]], %arg2, [[PAD00]], [[ZERO_I64]]) - // CHECK-DAG: [[PAD_DEFAULT:%.+]] = "tf.Const"() {value = dense<0.000000e+00> : tensor} + // CHECK-DAG: [[PAD_DEFAULT:%.+]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor}> // CHECK-DAG: [[PADDED:%.+]] = "tf.PadV2"(%arg0, [[FULL_PADDINGS]], [[PAD_DEFAULT]]) - // CHECK-DAG: [[PADDINGS:%.+]]:2 = "tf.Unpack"([[FULL_PADDINGS]]) {axis = 1 : i64} + // CHECK-DAG: [[PADDINGS:%.+]]:2 = "tf.Unpack"([[FULL_PADDINGS]]) <{axis = 1 : i64}> // CHECK-DAG: [[PADDINGS_SUM:%.+]] = "tf.AddV2"([[PADDINGS]]#0, [[PADDINGS]]#1) - // CHECK-DAG: [[INPUT_SHAPE:%.+]] = "tf.Const"() {value = dense<[3, 5, 7, 10]> : tensor<4xi64>} + // CHECK-DAG: [[INPUT_SHAPE:%.+]] = "tf.Const"() <{value = dense<[3, 5, 7, 10]> : tensor<4xi64>}> // CHECK-DAG: [[PADDED_SHAPE:%.+]] = "tf.AddV2"([[PADDINGS_SUM]], [[INPUT_SHAPE]]) // CHECK-DAG: [[PADDED_SHAPE_SPLITS:%.+]]:4 = "tf.Split"([[ZERO_I32]], [[PADDED_SHAPE]]) // CHECK-DAG: [[BLOCK_SHAPE_SPLITS:%.+]]:2 = "tf.Split"([[ZERO_I32]], %arg1) // CHECK-DAG: [[OUTER_SHAPE_0:%.+]] = "tf.Div"([[PADDED_SHAPE_SPLITS]]#1, [[BLOCK_SHAPE_SPLITS]]#0) // CHECK-DAG: [[OUTER_SHAPE_1:%.+]] = "tf.Div"([[PADDED_SHAPE_SPLITS]]#2, [[BLOCK_SHAPE_SPLITS]]#1) // CHECK-DAG: [[RESHAPED_SHAPE:%.+]] = "tf.ConcatV2"([[PADDED_SHAPE_SPLITS]]#0, [[OUTER_SHAPE_0]], [[BLOCK_SHAPE_SPLITS]]#0, [[OUTER_SHAPE_1]], [[BLOCK_SHAPE_SPLITS]]#1, [[PADDED_SHAPE_SPLITS]]#3, [[ZERO_I64]]) - // CHECK-DAG: [[PERMUTATION:%.+]] = "tf.Const"() {value = dense<[2, 4, 0, 1, 3, 5]> : tensor<6xi64>} + // CHECK-DAG: [[PERMUTATION:%.+]] = "tf.Const"() <{value = dense<[2, 4, 0, 1, 3, 5]> : tensor<6xi64>}> // CHECK-DAG: [[OUTPUT_BATCH_PART:%.+]] = "tf.Mul"([[PADDED_SHAPE_SPLITS]]#0, [[BLOCK_SHAPE_SPLITS]]#0) // CHECK-DAG: [[OUTPUT_BATCH:%.+]] = "tf.Mul"([[OUTPUT_BATCH_PART]], [[BLOCK_SHAPE_SPLITS]]#1) // CHECK-DAG: [[OUTPUT_SHAPE:%.+]] = "tf.ConcatV2"([[OUTPUT_BATCH]], [[OUTER_SHAPE_0]], [[OUTER_SHAPE_1]], [[PADDED_SHAPE_SPLITS]]#3, [[ZERO_I64]]) @@ -336,11 +336,11 @@ func.func @const_paddings_space_to_batch_nd(%arg0: tensor<1x8x2xf32>) -> (tensor %1 = "tf.Const"() {value = dense<[[3, 4]]> : tensor<1x2xi32>} : () -> tensor<1x2xi32> - // CHECK-DAG: [[VAL0:%.+]] = "tf.Const"() {value = dense<[3, 5, 2]> : tensor<3xi64>} - // CHECK-DAG: [[VAL1:%.+]] = "tf.Const"() {value = dense<[1, 5, 3, 2]> : tensor<4xi64>} - // CHECK-DAG: [[VAL2:%.+]] = "tf.Const"() {value = dense<{{\[\[}}0, 0], [3, 4], [0, 0{{\]\]}}> : tensor<3x2xi64>} - // CHECK-DAG: [[VAL3:%.+]] = "tf.Const"() {value = dense<0.000000e+00> : tensor} - // CHECK-DAG: [[VAL4:%.+]] = "tf.Const"() {value = dense<[2, 0, 1, 3]> : tensor<4xi64>} + // CHECK-DAG: [[VAL0:%.+]] = "tf.Const"() <{value = dense<[3, 5, 2]> : tensor<3xi64>}> + // CHECK-DAG: [[VAL1:%.+]] = "tf.Const"() <{value = dense<[1, 5, 3, 2]> : tensor<4xi64>}> + // CHECK-DAG: [[VAL2:%.+]] = "tf.Const"() <{value = dense<{{\[\[}}0, 0], [3, 4], [0, 0{{\]\]}}> : tensor<3x2xi64>}> + // CHECK-DAG: [[VAL3:%.+]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor}> + // CHECK-DAG: [[VAL4:%.+]] = "tf.Const"() <{value = dense<[2, 0, 1, 3]> : tensor<4xi64>}> // CHECK-DAG: [[VAL5:%.+]] = "tf.PadV2"(%arg0, [[VAL2]], [[VAL3]]) // CHECK-SAME: tensor<1x15x2xf32> // CHECK-DAG: [[VAL6:%.+]] = "tf.Reshape"([[VAL5]], [[VAL1]]) @@ -368,14 +368,14 @@ func.func @avoid_lowering_space_to_batch_nd(%arg0: tensor<1x8x2xf32>, %arg1: ten func.func @sixdim_space_to_batch_nd(%input: tensor<3x5x7x9x10x11xf32>, %block_shape: tensor<3xi64>, %paddings: tensor<3x2xi64>) -> tensor { // CHECK-DAG: [[PAD00:%.+]] = "tf.Const"() // CHECK-DAG: [[FULL_PADDINGS:%.+]] = "tf.ConcatV2"([[PAD00]], %arg2, [[PAD00]], [[PAD00]], {{.+}}) - // CHECK-DAG: [[INPUT_SHAPE:%.+]] = "tf.Const"() {value = dense<[3, 5, 7, 9, 10, 11]> : tensor<6xi64>} + // CHECK-DAG: [[INPUT_SHAPE:%.+]] = "tf.Const"() <{value = dense<[3, 5, 7, 9, 10, 11]> : tensor<6xi64>}> // CHECK-DAG: [[PADDED_SHAPE_SPLITS:%.+]]:6 = "tf.Split" // CHECK-DAG: [[BLOCK_SHAPE_SPLITS:%.+]]:3 = "tf.Split" // CHECK-DAG: [[OUTER_SHAPE_0:%.+]] = "tf.Div"([[PADDED_SHAPE_SPLITS]]#1, [[BLOCK_SHAPE_SPLITS]]#0) // CHECK-DAG: [[OUTER_SHAPE_1:%.+]] = "tf.Div"([[PADDED_SHAPE_SPLITS]]#2, [[BLOCK_SHAPE_SPLITS]]#1) // CHECK-DAG: [[OUTER_SHAPE_2:%.+]] = "tf.Div"([[PADDED_SHAPE_SPLITS]]#3, [[BLOCK_SHAPE_SPLITS]]#2) // CHECK-DAG: [[RESHAPED_SHAPE:%.+]] = "tf.ConcatV2"([[PADDED_SHAPE_SPLITS]]#0, [[OUTER_SHAPE_0]], [[BLOCK_SHAPE_SPLITS]]#0, [[OUTER_SHAPE_1]], [[BLOCK_SHAPE_SPLITS]]#1, [[OUTER_SHAPE_2]], [[BLOCK_SHAPE_SPLITS]]#2, [[PADDED_SHAPE_SPLITS]]#4, [[PADDED_SHAPE_SPLITS]]#5, {{.+}}) - // CHECK-DAG: [[PERMUTATION:%.+]] = "tf.Const"() {value = dense<[2, 4, 6, 0, 1, 3, 5, 7, 8]> : tensor<9xi64>} + // CHECK-DAG: [[PERMUTATION:%.+]] = "tf.Const"() <{value = dense<[2, 4, 6, 0, 1, 3, 5, 7, 8]> : tensor<9xi64>}> // CHECK-DAG: [[OUTPUT_BATCH_PART1:%.+]] = "tf.Mul"([[PADDED_SHAPE_SPLITS]]#0, [[BLOCK_SHAPE_SPLITS]]#0) // CHECK-DAG: [[OUTPUT_BATCH_PART2:%.+]] = "tf.Mul"([[OUTPUT_BATCH_PART1]], [[BLOCK_SHAPE_SPLITS]]#1) // CHECK-DAG: [[OUTPUT_BATCH:%.+]] = "tf.Mul"([[OUTPUT_BATCH_PART2]], [[BLOCK_SHAPE_SPLITS]]#2) @@ -386,11 +386,11 @@ func.func @sixdim_space_to_batch_nd(%input: tensor<3x5x7x9x10x11xf32>, %block_sh // CHECK-LABEL: func @batchToSpace func.func @batchToSpace(%arg0: tensor<3x5x2xf32>) -> (tensor<1x8x2xf32>) { - // CHECK-DAG: [[VAL0:%.+]] = "tf.Const"() {value = dense<[3, 1, 5, 2]> : tensor<4xi64>} - // CHECK-DAG: [[VAL1:%.+]] = "tf.Const"() {value = dense<[1, 2, 0, 3]> : tensor<4xi64>} - // CHECK-DAG: [[VAL2:%.+]] = "tf.Const"() {value = dense<[1, 15, 2]> : tensor<3xi64>} - // CHECK-DAG: [[VAL3:%.+]] = "tf.Const"() {value = dense<[0, 3, 0]> : tensor<3xi64>} - // CHECK-DAG: [[VAL4:%.+]] = "tf.Const"() {value = dense<[1, 8, 2]> : tensor<3xi64>} + // CHECK-DAG: [[VAL0:%.+]] = "tf.Const"() <{value = dense<[3, 1, 5, 2]> : tensor<4xi64>}> + // CHECK-DAG: [[VAL1:%.+]] = "tf.Const"() <{value = dense<[1, 2, 0, 3]> : tensor<4xi64>}> + // CHECK-DAG: [[VAL2:%.+]] = "tf.Const"() <{value = dense<[1, 15, 2]> : tensor<3xi64>}> + // CHECK-DAG: [[VAL3:%.+]] = "tf.Const"() <{value = dense<[0, 3, 0]> : tensor<3xi64>}> + // CHECK-DAG: [[VAL4:%.+]] = "tf.Const"() <{value = dense<[1, 8, 2]> : tensor<3xi64>}> // CHECK-DAG: [[VAL5:%.+]] = "tf.Reshape"(%arg0, [[VAL0]]) // CHECK-DAG: [[VAL6:%.+]] = "tf.Transpose"([[VAL5]], [[VAL1]]) // CHECK-DAG: [[VAL7:%.+]] = "tf.Reshape"([[VAL6]], [[VAL2]]) @@ -404,11 +404,11 @@ func.func @batchToSpace(%arg0: tensor<3x5x2xf32>) -> (tensor<1x8x2xf32>) { } func.func @fake_quant_with_min_max_args(%arg0 : tensor) -> tensor { - // CHECK-DAG: [[VAL0:%.+]] = "tf.Const"() {value = dense<1.275000e+02> : tensor} - // CHECK-DAG: [[VAL1:%.+]] = "tf.Const"() {value = dense<1.00392163> : tensor} - // CHECK-DAG: [[VAL2:%.+]] = "tf.Const"() {value = dense<-0.996078491> : tensor} - // CHECK-DAG: [[VAL3:%.+]] = "tf.Const"() {value = dense<0.00784313772> : tensor} - // CHECK-DAG: [[VAL4:%.+]] = "tf.Const"() {value = dense<5.000000e-01> : tensor} + // CHECK-DAG: [[VAL0:%.+]] = "tf.Const"() <{value = dense<1.275000e+02> : tensor}> + // CHECK-DAG: [[VAL1:%.+]] = "tf.Const"() <{value = dense<1.00392163> : tensor}> + // CHECK-DAG: [[VAL2:%.+]] = "tf.Const"() <{value = dense<-0.996078491> : tensor}> + // CHECK-DAG: [[VAL3:%.+]] = "tf.Const"() <{value = dense<0.00784313772> : tensor}> + // CHECK-DAG: [[VAL4:%.+]] = "tf.Const"() <{value = dense<5.000000e-01> : tensor}> // CHECK-DAG: [[VAL5:%.+]] = "tf.ClipByValue"(%arg0, [[VAL2]], [[VAL1]]) // CHECK-DAG: [[VAL6:%.+]] = "tf.Sub"([[VAL5]], [[VAL2]]) // CHECK-DAG: [[VAL7:%.+]] = "tf.Mul"([[VAL6]], [[VAL0]]) @@ -423,11 +423,11 @@ func.func @fake_quant_with_min_max_args(%arg0 : tensor) -> tensor, %arg1 : tensor, %arg2 : tensor) -> tensor { - // CHECK-DAG: %[[ZERO:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor} : () -> tensor - // CHECK-DAG: %[[VAL1:.*]] = "tf.Const"() {value = dense<2.550000e+02> : tensor} : () -> tensor - // CHECK-DAG: %[[VAL2:.*]] = "tf.Const"() {value = dense<2.000000e+00> : tensor} : () -> tensor - // CHECK-DAG: %[[VAL3:.*]] = "tf.Const"() {value = dense<1.000000e+00> : tensor} : () -> tensor - // CHECK-DAG: %[[VAL4:.*]] = "tf.Const"() {value = dense<5.000000e-01> : tensor} : () -> tensor + // CHECK-DAG: %[[ZERO:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor}> : () -> tensor + // CHECK-DAG: %[[VAL1:.*]] = "tf.Const"() <{value = dense<2.550000e+02> : tensor}> : () -> tensor + // CHECK-DAG: %[[VAL2:.*]] = "tf.Const"() <{value = dense<2.000000e+00> : tensor}> : () -> tensor + // CHECK-DAG: %[[VAL3:.*]] = "tf.Const"() <{value = dense<1.000000e+00> : tensor}> : () -> tensor + // CHECK-DAG: %[[VAL4:.*]] = "tf.Const"() <{value = dense<5.000000e-01> : tensor}> : () -> tensor // CHECK-DAG: %[[VAL5:.*]] = "tf.Sub"(%arg2, %arg1) : (tensor, tensor) -> tensor // CHECK-DAG: %[[VAL6:.*]] = "tf.Div"(%[[VAL5]], %[[VAL1]]) : (tensor, tensor) -> tensor // CHECK-DAG: %[[VAL7:.*]] = "tf.Div"(%[[VAL1]], %[[VAL5]]) : (tensor, tensor) -> tensor @@ -436,17 +436,17 @@ func.func @fake_quant_with_min_max_vars(%arg0 : tensor, %arg1 : tensor< // CHECK-DAG: %[[VAL10:.*]] = "tf.Floor"(%[[VAL9]]) : (tensor) -> tensor // CHECK-DAG: %[[VAL11:.*]] = "tf.Sub"(%[[VAL9]], %[[VAL10]]) : (tensor, tensor) -> tensor // CHECK-DAG: %[[VAL12:.*]] = "tf.Greater"(%[[VAL11]], %[[VAL4]]) : (tensor, tensor) -> tensor - // CHECK-DAG: %[[VAL13:.*]] = "tf.Equal"(%[[VAL11]], %[[VAL4]]) {incompatible_shape_error = true} : (tensor, tensor) -> tensor + // CHECK-DAG: %[[VAL13:.*]] = "tf.Equal"(%[[VAL11]], %[[VAL4]]) <{incompatible_shape_error = true}> : (tensor, tensor) -> tensor // CHECK-DAG: %[[VAL14:.*]] = "tf.Mul"(%[[VAL9]], %[[VAL4]]) : (tensor, tensor) -> tensor // CHECK-DAG: %[[VAL15:.*]] = "tf.Floor"(%[[VAL14]]) : (tensor) -> tensor // CHECK-DAG: %[[VAL16:.*]] = "tf.Mul"(%[[VAL15]], %[[VAL2]]) : (tensor, tensor) -> tensor // CHECK-DAG: %[[VAL17:.*]] = "tf.Sub"(%[[VAL10]], %[[VAL16]]) : (tensor, tensor) -> tensor - // CHECK-DAG: %[[VAL18:.*]] = "tf.Equal"(%[[VAL17]], %[[VAL3]]) {incompatible_shape_error = true} : (tensor, tensor) -> tensor + // CHECK-DAG: %[[VAL18:.*]] = "tf.Equal"(%[[VAL17]], %[[VAL3]]) <{incompatible_shape_error = true}> : (tensor, tensor) -> tensor // CHECK-DAG: %[[VAL19:.*]] = "tf.LogicalAnd"(%[[VAL13]], %[[VAL18]]) : (tensor, tensor) -> tensor // CHECK-DAG: %[[VAL20:.*]] = "tf.LogicalOr"(%[[VAL12]], %[[VAL19]]) : (tensor, tensor) -> tensor // CHECK-DAG: %[[VAL21:.*]] = "tf.AddV2"(%[[VAL10]], %[[VAL3]]) : (tensor, tensor) -> tensor // CHECK-DAG: %[[INNER_SELECT:.*]] = "tf.SelectV2"(%[[VAL20]], %[[VAL21]], %[[VAL10]]) : (tensor, tensor, tensor) -> tensor - // CHECK-DAG: %[[IS_ZERO:.*]] = "tf.Equal"(%[[INNER_SELECT]], %[[ZERO]]) {incompatible_shape_error = true} + // CHECK-DAG: %[[IS_ZERO:.*]] = "tf.Equal"(%[[INNER_SELECT]], %[[ZERO]]) <{incompatible_shape_error = true}> // CHECK-DAG: %[[VAL22:.*]] = "tf.SelectV2"(%[[IS_ZERO]], %[[ZERO]], %[[INNER_SELECT]]) // CHECK-DAG: %[[VAL23:.*]] = "tf.ClipByValue"(%[[VAL22]], %[[ZERO]], %[[VAL1]]) : (tensor, tensor, tensor) -> tensor // CHECK-DAG: %[[VAL24:.*]] = "tf.Sub"(%[[ZERO]], %[[VAL23]]) : (tensor, tensor) -> tensor @@ -469,29 +469,29 @@ func.func @fake_quant_with_min_max_vars(%arg0 : tensor, %arg1 : tensor< // CHECK-LABEL: SoftmaxCrossEntropyWithLogits // CHECK-SAME: %[[FEATURES:.*]]: tensor<2x3xf32>, %[[LABELS:.*]]: tensor<2x3xf32> func.func @SoftmaxCrossEntropyWithLogits(%features: tensor<2x3xf32>, %labels: tensor<2x3xf32>) -> (tensor<2xf32>, tensor<2x3xf32>) { - // CHECK-DAG: %[[AXIS:.*]] = "tf.Const"() {value = dense<-1> : tensor<1xi64>} : () -> tensor<1xi64> - // CHECK-DAG: %[[ZERO:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor} : () -> tensor + // CHECK-DAG: %[[AXIS:.*]] = "tf.Const"() <{value = dense<-1> : tensor<1xi64>}> : () -> tensor<1xi64> + // CHECK-DAG: %[[ZERO:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor}> : () -> tensor // CHECK-DAG: %[[NEG_LABELS:.*]] = "tf.Neg"(%[[LABELS]]) : (tensor<2x3xf32>) -> tensor<2x3xf32> // LogSoftmax expansion. - // CHECK-DAG: %[[LOG_SOFTMAX_MAX:.*]] = "tf.Max"(%[[FEATURES]], %[[AXIS]]) {keep_dims = true} : (tensor<2x3xf32>, tensor<1xi64>) -> tensor<2x1xf32> + // CHECK-DAG: %[[LOG_SOFTMAX_MAX:.*]] = "tf.Max"(%[[FEATURES]], %[[AXIS]]) <{keep_dims = true}> : (tensor<2x3xf32>, tensor<1xi64>) -> tensor<2x1xf32> // CHECK-DAG: %[[LOG_SOFTMAX_SHIFTED:.*]] = "tf.Sub"(%[[FEATURES]], %[[LOG_SOFTMAX_MAX]]) : (tensor<2x3xf32>, tensor<2x1xf32>) -> tensor<2x3xf32> // CHECK-DAG: %[[LOG_SOFTMAX_EXP:.*]] = "tf.Exp"(%[[LOG_SOFTMAX_SHIFTED]]) : (tensor<2x3xf32>) -> tensor<2x3xf32> - // CHECK-DAG: %[[LOG_SOFTMAX_SUM:.*]] = "tf.Sum"(%[[LOG_SOFTMAX_EXP]], %[[AXIS]]) {keep_dims = true} : (tensor<2x3xf32>, tensor<1xi64>) -> tensor<2x1xf32> + // CHECK-DAG: %[[LOG_SOFTMAX_SUM:.*]] = "tf.Sum"(%[[LOG_SOFTMAX_EXP]], %[[AXIS]]) <{keep_dims = true}> : (tensor<2x3xf32>, tensor<1xi64>) -> tensor<2x1xf32> // CHECK-DAG: %[[LOG_SOFTMAX_LOG:.*]] = "tf.Log"(%[[LOG_SOFTMAX_SUM]]) : (tensor<2x1xf32>) -> tensor<2x1xf32> // CHECK-DAG: %[[LOG_SOFTMAX:.*]] = "tf.Sub"(%[[LOG_SOFTMAX_SHIFTED]], %[[LOG_SOFTMAX_LOG]]) : (tensor<2x3xf32>, tensor<2x1xf32>) -> tensor<2x3xf32> - // CHECK-DAG: %[[IS_LABEL_ZERO:.*]] = "tf.Equal"(%[[NEG_LABELS]], %[[ZERO]]) {incompatible_shape_error = true} : (tensor<2x3xf32>, tensor) -> tensor<2x3xi1> + // CHECK-DAG: %[[IS_LABEL_ZERO:.*]] = "tf.Equal"(%[[NEG_LABELS]], %[[ZERO]]) <{incompatible_shape_error = true}> : (tensor<2x3xf32>, tensor) -> tensor<2x3xi1> // CHECK-DAG: %[[LOSS_INP:.*]] = "tf.Mul"(%[[LOG_SOFTMAX]], %[[NEG_LABELS]]) : (tensor<2x3xf32>, tensor<2x3xf32>) -> tensor<2x3xf32> // CHECK-DAG: %[[SAFE_LOSS_INP:.*]] = "tf.SelectV2"(%[[IS_LABEL_ZERO]], %[[ZERO]], %[[LOSS_INP]]) : (tensor<2x3xi1>, tensor, tensor<2x3xf32>) -> tensor<2x3xf32> - // CHECK-DAG: %[[LOSS:.*]] = "tf.Sum"(%[[SAFE_LOSS_INP]], %[[AXIS]]) {keep_dims = false} : (tensor<2x3xf32>, tensor<1xi64>) -> tensor<2xf32> + // CHECK-DAG: %[[LOSS:.*]] = "tf.Sum"(%[[SAFE_LOSS_INP]], %[[AXIS]]) <{keep_dims = false}> : (tensor<2x3xf32>, tensor<1xi64>) -> tensor<2xf32> // Softmax expansion. - // CHECK-DAG: %[[SOFTMAX_MAX:.*]] = "tf.Max"(%arg0, %[[AXIS]]) {keep_dims = true} : (tensor<2x3xf32>, tensor<1xi64>) -> tensor<2x1xf32> + // CHECK-DAG: %[[SOFTMAX_MAX:.*]] = "tf.Max"(%arg0, %[[AXIS]]) <{keep_dims = true}> : (tensor<2x3xf32>, tensor<1xi64>) -> tensor<2x1xf32> // CHECK-DAG: %[[SOFTMAX_SHIFTED:.*]] = "tf.Sub"(%[[FEATURES]], %[[SOFTMAX_MAX]]) : (tensor<2x3xf32>, tensor<2x1xf32>) -> tensor<2x3xf32> // CHECK-DAG: %[[SOFTMAX_EXP:.*]] = "tf.Exp"(%[[SOFTMAX_SHIFTED]]) : (tensor<2x3xf32>) -> tensor<2x3xf32> - // CHECK-DAG: %[[SOFTMAX_SUM:.*]] = "tf.Sum"(%[[SOFTMAX_EXP]], %[[AXIS]]) {keep_dims = true} : (tensor<2x3xf32>, tensor<1xi64>) -> tensor<2x1xf32> + // CHECK-DAG: %[[SOFTMAX_SUM:.*]] = "tf.Sum"(%[[SOFTMAX_EXP]], %[[AXIS]]) <{keep_dims = true}> : (tensor<2x3xf32>, tensor<1xi64>) -> tensor<2x1xf32> // CHECK-DAG: %[[SOFTMAX:.*]] = "tf.Div"(%[[SOFTMAX_EXP]], %[[SOFTMAX_SUM]]) : (tensor<2x3xf32>, tensor<2x1xf32>) -> tensor<2x3xf32> // CHECK-DAG: %[[BACKPROP:.*]] = "tf.Sub"(%[[SOFTMAX]], %[[LABELS]]) : (tensor<2x3xf32>, tensor<2x3xf32>) -> tensor<2x3xf32> @@ -534,19 +534,19 @@ func.func @scalar_SoftmaxCrossEntropyWithLogits(%features: tensor, %labels: // CHECK-SAME: %[[FEATURES:.*]]: tensor<2x3xf32>, %[[SPARSE_LABELS:.*]]: tensor<2xi32> func.func @SparseSoftmaxCrossEntropyWithLogits(%features: tensor<2x3xf32>, %labels: tensor<2xi32>) -> (tensor<2xf32>, tensor<2x3xf32>) { // Convert SPARSE_LABELS to dense LABELS. - // CHECK-DAG: %[[DEPTH:.*]] = "tf.Const"() {value = dense<3> : tensor} : () -> tensor - // CHECK-DAG: %[[ONE:.*]] = "tf.Const"() {value = dense<1.000000e+00> : tensor} : () -> tensor - // CHECK-DAG: %[[ZERO:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor} : () -> tensor - // CHECK-DAG: %[[LABELS:.*]] = "tf.OneHot"(%[[SPARSE_LABELS]], %[[DEPTH]], %[[ONE]], %[[ZERO]]) {axis = 1 : i64} : (tensor<2xi32>, tensor, tensor, tensor) -> tensor<2x3xf32> + // CHECK-DAG: %[[DEPTH:.*]] = "tf.Const"() <{value = dense<3> : tensor}> : () -> tensor + // CHECK-DAG: %[[ONE:.*]] = "tf.Const"() <{value = dense<1.000000e+00> : tensor}> : () -> tensor + // CHECK-DAG: %[[ZERO:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor}> : () -> tensor + // CHECK-DAG: %[[LABELS:.*]] = "tf.OneHot"(%[[SPARSE_LABELS]], %[[DEPTH]], %[[ONE]], %[[ZERO]]) <{axis = 1 : i64}> : (tensor<2xi32>, tensor, tensor, tensor) -> tensor<2x3xf32> // Adjust labels to have Nan for out of range labels. - // CHECK-DAG: %[[ZERO_I32:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor + // CHECK-DAG: %[[ZERO_I32:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK-DAG: %[[IS_NEGATIVE:.*]] = "tf.LessEqual"(%[[ZERO_I32]], %arg1) : (tensor, tensor<2xi32>) -> tensor<2xi1> // CHECK-DAG: %[[IS_LESS:.*]] = "tf.Less"(%arg1, %[[DEPTH]]) : (tensor<2xi32>, tensor) -> tensor<2xi1> // CHECK-DAG: %[[IS_WITHIN_RANGE:.*]] = "tf.LogicalAnd"(%[[IS_NEGATIVE]], %[[IS_LESS]]) : (tensor<2xi1>, tensor<2xi1>) -> tensor<2xi1> - // CHECK-DAG: %[[NAN:.*]] = "tf.Const"() {value = dense<0x7FC00000> : tensor} : () -> tensor + // CHECK-DAG: %[[NAN:.*]] = "tf.Const"() <{value = dense<0x7FC00000> : tensor}> : () -> tensor // CHECK-DAG: %[[ZERO_OR_NAN:.*]] = "tf.SelectV2"(%[[IS_WITHIN_RANGE]], %[[ZERO]], %[[NAN]]) : (tensor<2xi1>, tensor, tensor) -> tensor<2xf32> - // CHECK-DAG: %[[NEG_ONE:.*]] = "tf.Const"() {value = dense<-1> : tensor<1xi64>} : () -> tensor<1xi64> + // CHECK-DAG: %[[NEG_ONE:.*]] = "tf.Const"() <{value = dense<-1> : tensor<1xi64>}> : () -> tensor<1xi64> // CHECK-DAG: %[[RESHAPE:.*]] = "tf.ExpandDims"(%[[ZERO_OR_NAN]], %[[NEG_ONE]]) : (tensor<2xf32>, tensor<1xi64>) -> tensor<2x1xf32> // CHECK-DAG: %[[ADJUSTED_LABELS:.*]] = "tf.AddV2"(%[[LABELS]], %[[RESHAPE]]) : (tensor<2x3xf32>, tensor<2x1xf32>) -> tensor<2x3xf32> @@ -589,7 +589,7 @@ func.func @SparseSoftmaxCrossEntropyWithLogits_with_dynamic(%features: tensor<*x // CHECK-LABEL: func @tanhgrad_float // CHECK-SAME: (%[[Y:.*]]: tensor<*xf32>, %[[DY:.*]]: tensor<*xf32>) func.func @tanhgrad_float(%y : tensor<*xf32>, %dy: tensor<*xf32>) -> tensor<*xf32> { - // CHECK: %[[ONE:.*]] = "tf.Const"() {value = dense<1.000000e+00> : tensor} : () -> tensor + // CHECK: %[[ONE:.*]] = "tf.Const"() <{value = dense<1.000000e+00> : tensor}> : () -> tensor // CHECK: %[[Y_SQUARE:.*]] = "tf.Mul"(%[[Y]], %[[Y]]) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32> // CHECK: %[[SUB:.*]] = "tf.Sub"(%[[ONE]], %[[Y_SQUARE]]) : (tensor, tensor<*xf32>) -> tensor<*xf32> // CHECK: %[[RESULT:.*]] = "tf.Mul"(%[[DY]], %[[SUB]]) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32> @@ -610,7 +610,7 @@ func.func @tanhgrad_complex(%y : tensor<*xcomplex>, %dy: tensor<*xcomplex) -> tensor<*xi32> { - // CHECK: %[[ZERO:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor + // CHECK: %[[ZERO:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK: %[[SHAPE:.*]] = "tf.Shape"(%arg0) : (tensor<*xi32>) -> tensor // CHECK: "tf.BroadcastTo"(%[[ZERO]], %[[SHAPE]]) : (tensor, tensor) -> tensor<*xi32> @@ -627,7 +627,7 @@ func.func @ZerosLike_variant(%arg0: tensor>>) -> // CHECK-LABEL: func @OnesLike_unranked func.func @OnesLike_unranked(%arg0: tensor<*xi32>) -> tensor<*xi32> { - // CHECK: %[[ONE:.*]] = "tf.Const"() {value = dense<1> : tensor} : () -> tensor + // CHECK: %[[ONE:.*]] = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor // CHECK: %[[SHAPE:.*]] = "tf.Shape"(%arg0) : (tensor<*xi32>) -> tensor // CHECK: "tf.BroadcastTo"(%[[ONE]], %[[SHAPE]]) : (tensor, tensor) -> tensor<*xi32> @@ -682,8 +682,8 @@ func.func @addN_variant(%arg0: tensor>>, %arg1: t // CHECK-LABEL: func @DynamicStitch_simple func.func @DynamicStitch_simple(%arg0: tensor<2x2xf32>) -> tensor<2x2xf32> { - // CHECK: %[[AXIS:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor - // CHECK: %[[ITEMS:.*]]:2 = "tf.Unpack"(%arg0) {axis = 0 : i64} : (tensor<2x2xf32>) -> (tensor<2xf32>, tensor<2xf32>) + // CHECK: %[[AXIS:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor + // CHECK: %[[ITEMS:.*]]:2 = "tf.Unpack"(%arg0) <{axis = 0 : i64}> : (tensor<2x2xf32>) -> (tensor<2xf32>, tensor<2xf32>) // CHECK-DAG: %[[ITEMS_1:.*]] = "tf.ExpandDims"(%[[ITEMS]]#1, %[[AXIS]]) // CHECK-DAG: %[[ITEMS_0:.*]] = "tf.ExpandDims"(%[[ITEMS]]#0, %[[AXIS]]) // CHECK: %[[RESULT:.*]] = "tf.ConcatV2"(%[[ITEMS_1]], %[[ITEMS_0]], %[[AXIS]]) : (tensor<1x2xf32>, tensor<1x2xf32>, tensor) -> tensor<2x2xf32> @@ -696,12 +696,12 @@ func.func @DynamicStitch_simple(%arg0: tensor<2x2xf32>) -> tensor<2x2xf32> { // CHECK-LABEL: DynamicStitch_scalar_matrix_indices func.func @DynamicStitch_scalar_matrix_indices(%arg0: tensor<2xf32>, %arg1: tensor<2x2x2xf32>) -> (tensor<5x2xf32>) { - // CHECK-DAG: %[[SHAPE:.*]] = "tf.Const"() {value = dense<[-1, 2]> : tensor<2xi64>} : () -> tensor<2xi64> + // CHECK-DAG: %[[SHAPE:.*]] = "tf.Const"() <{value = dense<[-1, 2]> : tensor<2xi64>}> : () -> tensor<2xi64> // CHECK-DAG: %[[INP0:.*]] = "tf.Reshape"(%arg0, %[[SHAPE]]) : (tensor<2xf32>, tensor<2xi64>) -> tensor<1x2xf32> - // CHECK-DAG: %[[ITEMS0:.*]] = "tf.Unpack"(%[[INP0]]) {axis = 0 : i64} : (tensor<1x2xf32>) -> tensor<2xf32> + // CHECK-DAG: %[[ITEMS0:.*]] = "tf.Unpack"(%[[INP0]]) <{axis = 0 : i64}> : (tensor<1x2xf32>) -> tensor<2xf32> // CHECK-DAG: %[[INP1:.*]] = "tf.Reshape"(%arg1, %[[SHAPE]]) : (tensor<2x2x2xf32>, tensor<2xi64>) -> tensor<4x2xf32> - // CHECK-DAG: %[[ITEMS1:.*]]:4 = "tf.Unpack"(%[[INP1]]) {axis = 0 : i64} : (tensor<4x2xf32>) -> (tensor<2xf32>, tensor<2xf32>, tensor<2xf32>, tensor<2xf32>) - // CHECK-DAG: %[[AXIS:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor + // CHECK-DAG: %[[ITEMS1:.*]]:4 = "tf.Unpack"(%[[INP1]]) <{axis = 0 : i64}> : (tensor<4x2xf32>) -> (tensor<2xf32>, tensor<2xf32>, tensor<2xf32>, tensor<2xf32>) + // CHECK-DAG: %[[AXIS:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK-DAG: %[[ITEMS1_3:.*]] = "tf.ExpandDims"(%[[ITEMS1]]#3, %[[AXIS]]) // CHECK-DAG: %[[ITEMS1_2:.*]] = "tf.ExpandDims"(%[[ITEMS1]]#2, %[[AXIS]]) // CHECK-DAG: %[[ITEMS1_1:.*]] = "tf.ExpandDims"(%[[ITEMS1]]#1, %[[AXIS]]) @@ -727,8 +727,8 @@ func.func @DynamicStitch_uint8(%arg0: tensor<2x2xui8>) -> tensor<2x2xui8> { // CHECK-LABEL: func @DynamicStitch_scalar_item func.func @DynamicStitch_scalar_item(%arg0: tensor<2xf32>) -> tensor<2xf32> { - // CHECK-DAG: %[[ITEMS:.*]]:2 = "tf.Unpack"(%arg0) {axis = 0 : i64} : (tensor<2xf32>) -> (tensor, tensor) - // CHECK-DAG: %[[AXIS:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor + // CHECK-DAG: %[[ITEMS:.*]]:2 = "tf.Unpack"(%arg0) <{axis = 0 : i64}> : (tensor<2xf32>) -> (tensor, tensor) + // CHECK-DAG: %[[AXIS:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK-DAG: %[[ITEMS_1:.*]] = "tf.ExpandDims"(%[[ITEMS]]#1, %[[AXIS]]) // CHECK-DAG: %[[ITEMS_0:.*]] = "tf.ExpandDims"(%[[ITEMS]]#0, %[[AXIS]]) // CHECK-DAG: %[[RESULT:.*]] = "tf.ConcatV2"(%[[ITEMS_1]], %[[ITEMS_0]], %[[AXIS]]) : (tensor<1xf32>, tensor<1xf32>, tensor) -> tensor<2xf32> @@ -741,8 +741,8 @@ func.func @DynamicStitch_scalar_item(%arg0: tensor<2xf32>) -> tensor<2xf32> { // CHECK-LABEL: func @DynamicStitch_matrix_item func.func @DynamicStitch_matrix_item(%arg0: tensor<2x2x2xf32>) -> tensor<2x2x2xf32> { - // CHECK-DAG: %[[ITEMS:.*]]:2 = "tf.Unpack"(%arg0) {axis = 0 : i64} : (tensor<2x2x2xf32>) -> (tensor<2x2xf32>, tensor<2x2xf32>) - // CHECK-DAG: %[[AXIS:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor + // CHECK-DAG: %[[ITEMS:.*]]:2 = "tf.Unpack"(%arg0) <{axis = 0 : i64}> : (tensor<2x2x2xf32>) -> (tensor<2x2xf32>, tensor<2x2xf32>) + // CHECK-DAG: %[[AXIS:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK-DAG: %[[ITEMS_1:.*]] = "tf.ExpandDims"(%[[ITEMS]]#1, %[[AXIS]]) // CHECK-DAG: %[[ITEMS_0:.*]] = "tf.ExpandDims"(%[[ITEMS]]#0, %[[AXIS]]) // CHECK-DAG: %[[RESULT:.*]] = "tf.ConcatV2"(%[[ITEMS_1]], %[[ITEMS_0]], %[[AXIS]]) : (tensor<1x2x2xf32>, tensor<1x2x2xf32>, tensor) -> tensor<2x2x2xf32> @@ -762,8 +762,8 @@ func.func @DynamicStitch_dynamic(%arg0: tensor<*xi32>, %arg1: tensor<*xf32>) -> // CHECK-LABEL: func @DynamicStitch_duplicates func.func @DynamicStitch_duplicates(%arg0: tensor<2x2xf32>) -> tensor<1x2xf32> { - // CHECK-DAG: %[[ITEMS:.*]]:2 = "tf.Unpack"(%arg0) {axis = 0 : i64} : (tensor<2x2xf32>) -> (tensor<2xf32>, tensor<2xf32>) - // CHECK-DAG: %[[AXIS:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor + // CHECK-DAG: %[[ITEMS:.*]]:2 = "tf.Unpack"(%arg0) <{axis = 0 : i64}> : (tensor<2x2xf32>) -> (tensor<2xf32>, tensor<2xf32>) + // CHECK-DAG: %[[AXIS:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK-DAG: %[[ITEMS_1:.*]] = "tf.ExpandDims"(%[[ITEMS]]#1, %[[AXIS]]) // CHECK-DAG: %[[RESULT:.*]] = "tf.ConcatV2"(%[[ITEMS_1]], %[[AXIS]]) : (tensor<1x2xf32>, tensor) -> tensor<1x2xf32> // CHECK: return %[[RESULT]] @@ -783,7 +783,7 @@ func.func @ParallelDynamicStitch(%arg0: tensor<2x2xf32>) -> tensor<2x2xf32> { // CHECK-LABEL: @Reciprocal_i32 func.func @Reciprocal_i32(%arg0: tensor<*xi32>) -> tensor<*xi32> { - // CHECK: %[[ONE:.*]] = "tf.Const"() {value = dense<1> : tensor} : () -> tensor + // CHECK: %[[ONE:.*]] = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor // CHECK: "tf.Div"(%[[ONE]], %arg0) : (tensor, tensor<*xi32>) -> tensor<*xi32> %0 = "tf.Reciprocal"(%arg0) : (tensor<*xi32>) -> tensor<*xi32> func.return %0 : tensor<*xi32> @@ -791,7 +791,7 @@ func.func @Reciprocal_i32(%arg0: tensor<*xi32>) -> tensor<*xi32> { // CHECK-LABEL: @Reciprocal_f32 func.func @Reciprocal_f32(%arg0: tensor<*xf32>) -> tensor<*xf32> { - // CHECK: %[[ONE:.*]] = "tf.Const"() {value = dense<1.000000e+00> : tensor} : () -> tensor + // CHECK: %[[ONE:.*]] = "tf.Const"() <{value = dense<1.000000e+00> : tensor}> : () -> tensor // CHECK: "tf.Div"(%[[ONE]], %arg0) : (tensor, tensor<*xf32>) -> tensor<*xf32> %0 = "tf.Reciprocal"(%arg0) : (tensor<*xf32>) -> tensor<*xf32> func.return %0 : tensor<*xf32> @@ -799,7 +799,7 @@ func.func @Reciprocal_f32(%arg0: tensor<*xf32>) -> tensor<*xf32> { // CHECK-LABEL: @Reciprocal_complexf32 func.func @Reciprocal_complexf32(%arg0: tensor<*xcomplex>) -> tensor<*xcomplex> { - // CHECK: %[[ONE:.*]] = "tf.Const"() {value = dense<(1.000000e+00,0.000000e+00)> : tensor>} : () -> tensor> + // CHECK: %[[ONE:.*]] = "tf.Const"() <{value = dense<(1.000000e+00,0.000000e+00)> : tensor>}> : () -> tensor> // CHECK: "tf.Div"(%[[ONE]], %arg0) : (tensor>, tensor<*xcomplex>) -> tensor<*xcomplex> %0 = "tf.Reciprocal"(%arg0) : (tensor<*xcomplex>) -> tensor<*xcomplex> func.return %0 : tensor<*xcomplex> @@ -807,7 +807,7 @@ func.func @Reciprocal_complexf32(%arg0: tensor<*xcomplex>) -> tensor<*xcomp // CHECK-LABEL: @Reciprocal_complexf64 func.func @Reciprocal_complexf64(%arg0: tensor<*xcomplex>) -> tensor<*xcomplex> { - // CHECK: %[[ONE:.*]] = "tf.Const"() {value = dense<(1.000000e+00,0.000000e+00)> : tensor>} : () -> tensor> + // CHECK: %[[ONE:.*]] = "tf.Const"() <{value = dense<(1.000000e+00,0.000000e+00)> : tensor>}> : () -> tensor> // CHECK: "tf.Div"(%[[ONE]], %arg0) : (tensor>, tensor<*xcomplex>) -> tensor<*xcomplex> %0 = "tf.Reciprocal"(%arg0) : (tensor<*xcomplex>) -> tensor<*xcomplex> func.return %0 : tensor<*xcomplex> @@ -816,7 +816,7 @@ func.func @Reciprocal_complexf64(%arg0: tensor<*xcomplex>) -> tensor<*xcomp // Inv is the same as Reciprocal // CHECK-LABEL: @Inv_i32 func.func @Inv_i32(%arg0: tensor<*xi32>) -> tensor<*xi32> { - // CHECK: %[[ONE:.*]] = "tf.Const"() {value = dense<1> : tensor} : () -> tensor + // CHECK: %[[ONE:.*]] = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor // CHECK: "tf.Div"(%[[ONE]], %arg0) : (tensor, tensor<*xi32>) -> tensor<*xi32> %0 = "tf.Inv"(%arg0) : (tensor<*xi32>) -> tensor<*xi32> func.return %0 : tensor<*xi32> @@ -824,7 +824,7 @@ func.func @Inv_i32(%arg0: tensor<*xi32>) -> tensor<*xi32> { // CHECK-LABEL: @ScatterNd func.func @ScatterNd(%arg0: tensor<4x1xi32>, %arg1: tensor<4xf32>) -> tensor<8xf32> { - // CHECK: %[[ZERO:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<8xf32>} : () -> tensor<8xf32> + // CHECK: %[[ZERO:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<8xf32>}> : () -> tensor<8xf32> // CHECK: "tf.TensorScatterAdd"(%[[ZERO]], %arg0, %arg1) : (tensor<8xf32>, tensor<4x1xi32>, tensor<4xf32>) -> tensor<8xf32> %shape = "tf.Const"() {value = dense<[8]> : tensor<1xi32>} : () -> tensor<1xi32> @@ -856,24 +856,24 @@ func.func @round_int(%arg0: tensor<2xi32>) -> tensor<2xi32> { // CHECK-LABEL: @round func.func @round(%arg0: tensor<2xf32>) -> tensor<2xf32> { - // CHECK-DAG: %[[ZERO:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor} : () -> tensor - // CHECK-DAG: %[[HALF:.*]] = "tf.Const"() {value = dense<5.000000e-01> : tensor} : () -> tensor - // CHECK-DAG: %[[TWO:.*]] = "tf.Const"() {value = dense<2.000000e+00> : tensor} : () -> tensor - // CHECK-DAG: %[[ONE:.*]] = "tf.Const"() {value = dense<1.000000e+00> : tensor} : () -> tensor + // CHECK-DAG: %[[ZERO:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor}> : () -> tensor + // CHECK-DAG: %[[HALF:.*]] = "tf.Const"() <{value = dense<5.000000e-01> : tensor}> : () -> tensor + // CHECK-DAG: %[[TWO:.*]] = "tf.Const"() <{value = dense<2.000000e+00> : tensor}> : () -> tensor + // CHECK-DAG: %[[ONE:.*]] = "tf.Const"() <{value = dense<1.000000e+00> : tensor}> : () -> tensor // CHECK: %[[ROUND_VAL:.*]] = "tf.Floor"(%arg0) : (tensor<2xf32>) -> tensor<2xf32> // CHECK: %[[FRACTION:.*]] = "tf.Sub"(%arg0, %[[ROUND_VAL]]) : (tensor<2xf32>, tensor<2xf32>) -> tensor<2xf32> // CHECK: %[[GT:.*]] = "tf.Greater"(%[[FRACTION]], %[[HALF]]) : (tensor<2xf32>, tensor) -> tensor<2xi1> - // CHECK: %[[EQ:.*]] = "tf.Equal"(%[[FRACTION]], %[[HALF]]) {incompatible_shape_error = true} : (tensor<2xf32>, tensor) -> tensor<2xi1> + // CHECK: %[[EQ:.*]] = "tf.Equal"(%[[FRACTION]], %[[HALF]]) <{incompatible_shape_error = true}> : (tensor<2xf32>, tensor) -> tensor<2xi1> // CHECK: %[[MUL1:.*]] = "tf.Mul"(%arg0, %[[HALF]]) : (tensor<2xf32>, tensor) -> tensor<2xf32> // CHECK: %[[FLOOR:.*]] = "tf.Floor"(%[[MUL1]]) : (tensor<2xf32>) -> tensor<2xf32> // CHECK: %[[MUL2:.*]] = "tf.Mul"(%[[FLOOR]], %[[TWO]]) : (tensor<2xf32>, tensor) -> tensor<2xf32> // CHECK: %[[NEAREST_EVEN_INT:.*]] = "tf.Sub"(%[[ROUND_VAL]], %[[MUL2]]) : (tensor<2xf32>, tensor<2xf32>) -> tensor<2xf32> - // CHECK: %[[IS_ODD:.*]] = "tf.Equal"(%[[NEAREST_EVEN_INT]], %[[ONE]]) {incompatible_shape_error = true} : (tensor<2xf32>, tensor) -> tensor<2xi1> + // CHECK: %[[IS_ODD:.*]] = "tf.Equal"(%[[NEAREST_EVEN_INT]], %[[ONE]]) <{incompatible_shape_error = true}> : (tensor<2xf32>, tensor) -> tensor<2xi1> // CHECK: %[[AND:.*]] = "tf.LogicalAnd"(%[[EQ]], %[[IS_ODD]]) : (tensor<2xi1>, tensor<2xi1>) -> tensor<2xi1> // CHECK: %[[OR:.*]] = "tf.LogicalOr"(%[[GT]], %[[AND]]) : (tensor<2xi1>, tensor<2xi1>) -> tensor<2xi1> // CHECK: %[[ADD:.*]] = "tf.AddV2"(%[[ROUND_VAL]], %[[ONE]]) : (tensor<2xf32>, tensor) -> tensor<2xf32> // CHECK: %[[INNER_SELECT:.*]] = "tf.SelectV2"(%[[OR]], %[[ADD]], %[[ROUND_VAL]]) : (tensor<2xi1>, tensor<2xf32>, tensor<2xf32>) -> tensor<2xf32> - // CHECK-DAG: %[[IS_ZERO:.*]] = "tf.Equal"(%[[INNER_SELECT]], %[[ZERO]]) {incompatible_shape_error = true} + // CHECK-DAG: %[[IS_ZERO:.*]] = "tf.Equal"(%[[INNER_SELECT]], %[[ZERO]]) <{incompatible_shape_error = true}> // CHECK-DAG: %[[SELECT:.*]] = "tf.SelectV2"(%[[IS_ZERO]], %[[ZERO]], %[[INNER_SELECT]]) %0 = "tf.Round"(%arg0) : (tensor<2xf32>) -> tensor<2xf32> @@ -890,24 +890,24 @@ func.func @round_dynamic(%arg0: tensor) -> tensor { // CHECK-LABEL: func @rint_dynamic func.func @rint_dynamic(%arg0: tensor) -> tensor { - // CHECK-DAG: %[[ZERO:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor} : () -> tensor - // CHECK-DAG: %[[HALF:.*]] = "tf.Const"() {value = dense<5.000000e-01> : tensor} : () -> tensor - // CHECK-DAG: %[[TWO:.*]] = "tf.Const"() {value = dense<2.000000e+00> : tensor} : () -> tensor - // CHECK-DAG: %[[ONE:.*]] = "tf.Const"() {value = dense<1.000000e+00> : tensor} : () -> tensor + // CHECK-DAG: %[[ZERO:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor}> : () -> tensor + // CHECK-DAG: %[[HALF:.*]] = "tf.Const"() <{value = dense<5.000000e-01> : tensor}> : () -> tensor + // CHECK-DAG: %[[TWO:.*]] = "tf.Const"() <{value = dense<2.000000e+00> : tensor}> : () -> tensor + // CHECK-DAG: %[[ONE:.*]] = "tf.Const"() <{value = dense<1.000000e+00> : tensor}> : () -> tensor // CHECK: %[[ROUND_VAL:.*]] = "tf.Floor"(%arg0) : (tensor) -> tensor // CHECK: %[[FRACTION:.*]] = "tf.Sub"(%arg0, %[[ROUND_VAL]]) : (tensor, tensor) -> tensor // CHECK: %[[GT:.*]] = "tf.Greater"(%[[FRACTION]], %[[HALF]]) : (tensor, tensor) -> tensor - // CHECK: %[[EQ:.*]] = "tf.Equal"(%[[FRACTION]], %[[HALF]]) {incompatible_shape_error = true} : (tensor, tensor) -> tensor + // CHECK: %[[EQ:.*]] = "tf.Equal"(%[[FRACTION]], %[[HALF]]) <{incompatible_shape_error = true}> : (tensor, tensor) -> tensor // CHECK: %[[MUL1:.*]] = "tf.Mul"(%arg0, %[[HALF]]) : (tensor, tensor) -> tensor // CHECK: %[[FLOOR:.*]] = "tf.Floor"(%[[MUL1]]) : (tensor) -> tensor // CHECK: %[[MUL2:.*]] = "tf.Mul"(%[[FLOOR]], %[[TWO]]) : (tensor, tensor) -> tensor // CHECK: %[[NEAREST_EVEN_INT:.*]] = "tf.Sub"(%[[ROUND_VAL]], %[[MUL2]]) : (tensor, tensor) -> tensor - // CHECK: %[[IS_ODD:.*]] = "tf.Equal"(%[[NEAREST_EVEN_INT]], %[[ONE]]) {incompatible_shape_error = true} : (tensor, tensor) -> tensor + // CHECK: %[[IS_ODD:.*]] = "tf.Equal"(%[[NEAREST_EVEN_INT]], %[[ONE]]) <{incompatible_shape_error = true}> : (tensor, tensor) -> tensor // CHECK: %[[AND:.*]] = "tf.LogicalAnd"(%[[EQ]], %[[IS_ODD]]) : (tensor, tensor) -> tensor // CHECK: %[[OR:.*]] = "tf.LogicalOr"(%[[GT]], %[[AND]]) : (tensor, tensor) -> tensor // CHECK: %[[ADD:.*]] = "tf.AddV2"(%[[ROUND_VAL]], %[[ONE]]) : (tensor, tensor) -> tensor // CHECK: %[[INNER_SELECT:.*]] = "tf.SelectV2"(%[[OR]], %[[ADD]], %[[ROUND_VAL]]) : (tensor, tensor, tensor) -> tensor - // CHECK: %[[IS_ZERO:.*]] = "tf.Equal"(%[[INNER_SELECT]], %[[ZERO]]) {incompatible_shape_error = true} + // CHECK: %[[IS_ZERO:.*]] = "tf.Equal"(%[[INNER_SELECT]], %[[ZERO]]) <{incompatible_shape_error = true}> // CHECK: %[[SELECT:.*]] = "tf.SelectV2"(%[[IS_ZERO]], %[[ZERO]], %[[INNER_SELECT]]) %0 = "tf.Rint"(%arg0) : (tensor) -> tensor @@ -937,12 +937,12 @@ func.func @lgamma(%arg0: tensor<4xf32>) -> tensor<4xf32> { func.func @imag_resize_nearest(%arg0: tensor<1x7x7x1xi32>) -> tensor<1x3x3x1xi32> { %shape = "tf.Const"() {device = "", value = dense<3> : tensor<2xi32>} : () -> tensor<2xi32> - // CHECK-DAG: [[VAL0:%.+]] = "tf.Const"() {value = dense<1> : tensor} - // CHECK-DAG: [[VAL1:%.+]] = "tf.Const"() {value = dense<[1, 3, 3, 1]> - // CHECK-DAG: [[VAL2:%.+]] = "tf.Const"() {value = dense<[1, 49, 1]> - // CHECK-DAG: [[VAL3:%.+]] = "tf.Const"() {value = dense<[0, 2, 4, 14, 16, 18, 28, 30, 32]> : tensor<9xi32>} + // CHECK-DAG: [[VAL0:%.+]] = "tf.Const"() <{value = dense<1> : tensor}> + // CHECK-DAG: [[VAL1:%.+]] = "tf.Const"() <{value = dense<[1, 3, 3, 1]> + // CHECK-DAG: [[VAL2:%.+]] = "tf.Const"() <{value = dense<[1, 49, 1]> + // CHECK-DAG: [[VAL3:%.+]] = "tf.Const"() <{value = dense<[0, 2, 4, 14, 16, 18, 28, 30, 32]> : tensor<9xi32>}> // CHECK: [[VAL4:%.+]] = "tf.Reshape"(%arg0, [[VAL2]]) - // CHECK: [[VAL5:%.+]] = "tf.GatherV2"([[VAL4]], [[VAL3]], [[VAL0]]) {batch_dims = 0 : i64} + // CHECK: [[VAL5:%.+]] = "tf.GatherV2"([[VAL4]], [[VAL3]], [[VAL0]]) <{batch_dims = 0 : i64}> // CHECK: [[VAL6:%.+]] = "tf.Reshape"([[VAL5]], [[VAL1]]) // CHECK: return [[VAL6]] %resize = "tf.ResizeNearestNeighbor"(%arg0, %shape) {align_corners = false, device = "", half_pixel_centers = false} : (tensor<1x7x7x1xi32>, tensor<2xi32>) -> tensor<1x3x3x1xi32> @@ -953,17 +953,17 @@ func.func @imag_resize_nearest(%arg0: tensor<1x7x7x1xi32>) -> tensor<1x3x3x1xi32 func.func @imag_resize_nearest_dyn_img(%arg0: tensor<1x?x?x1xi32>) -> tensor<1x3x3x1xi32> { %shape = "tf.Const"() {device = "", value = dense<3> : tensor<2xi32>} : () -> tensor<2xi32> - // CHECK-DAG: [[VAL0:%.+]] = "tf.Const"() {value = dense<1> : tensor} - // CHECK-DAG: [[VAL1:%.+]] = "tf.Const"() {value = dense<[3, 1]> : tensor<2xi32>} - // CHECK-DAG: [[VAL2:%.+]] = "tf.Const"() {value = dense<9> : tensor<1xi32>} - // CHECK-DAG: [[VAL3:%.+]] = "tf.Const"() {value = dense<3> : tensor<1xi32>} - // CHECK-DAG: [[VAL4:%.+]] = "tf.Const"() {value = dense<[1, 3]> : tensor<2xi32>} - // CHECK-DAG: [[VAL5:%.+]] = "tf.Const"() {value = dense<[0.000000e+00, 1.000000e+00, 2.000000e+00]> - // CHECK-DAG: [[VAL6:%.+]] = "tf.Const"() {value = dense<3.000000e+00> : tensor} - // CHECK-DAG: [[VAL7:%.+]] = "tf.Const"() {value = dense<0> : tensor} + // CHECK-DAG: [[VAL0:%.+]] = "tf.Const"() <{value = dense<1> : tensor}> + // CHECK-DAG: [[VAL1:%.+]] = "tf.Const"() <{value = dense<[3, 1]> : tensor<2xi32>}> + // CHECK-DAG: [[VAL2:%.+]] = "tf.Const"() <{value = dense<9> : tensor<1xi32>}> + // CHECK-DAG: [[VAL3:%.+]] = "tf.Const"() <{value = dense<3> : tensor<1xi32>}> + // CHECK-DAG: [[VAL4:%.+]] = "tf.Const"() <{value = dense<[1, 3]> : tensor<2xi32>}> + // CHECK-DAG: [[VAL5:%.+]] = "tf.Const"() <{value = dense<[0.000000e+00, 1.000000e+00, 2.000000e+00]> + // CHECK-DAG: [[VAL6:%.+]] = "tf.Const"() <{value = dense<3.000000e+00> : tensor}> + // CHECK-DAG: [[VAL7:%.+]] = "tf.Const"() <{value = dense<0> : tensor}> // CHECK: [[VAL8:%.+]] = "tf.Shape"(%arg0) // CHECK: [[VAL9:%.+]] = "tf.Cast"([[VAL8]]) - // CHECK: [[VAL10:%.+]]:4 = "tf.Unpack"([[VAL9]]) {axis = 0 : i64} + // CHECK: [[VAL10:%.+]]:4 = "tf.Unpack"([[VAL9]]) <{axis = 0 : i64}> // CHECK: [[VAL11:%.+]] = "tf.Mul"([[VAL10]]#1, [[VAL10]]#2) // CHECK: [[VAL12:%.+]] = "tf.ExpandDims"([[VAL10]]#0, [[VAL7]]) // CHECK: [[VAL13:%.+]] = "tf.ExpandDims"([[VAL10]]#3, [[VAL7]]) @@ -986,7 +986,7 @@ func.func @imag_resize_nearest_dyn_img(%arg0: tensor<1x?x?x1xi32>) -> tensor<1x3 // CHECK: [[VAL30:%.+]] = "tf.ExpandDims"([[VAL10]]#3, [[VAL7]]) // CHECK: [[VAL31:%.+]] = "tf.ConcatV2"([[VAL28]], [[VAL29]], [[VAL30]], [[VAL7]]) // CHECK: [[VAL32:%.+]] = "tf.Reshape"(%arg0, [[VAL31]]) - // CHECK: [[VAL33:%.+]] = "tf.GatherV2"([[VAL32]], [[VAL27]], [[VAL0]]) {batch_dims = 0 : i64} + // CHECK: [[VAL33:%.+]] = "tf.GatherV2"([[VAL32]], [[VAL27]], [[VAL0]]) <{batch_dims = 0 : i64}> // CHECK: [[VAL34:%.+]] = "tf.Reshape"([[VAL33]], [[VAL14]]) // CHECK: return [[VAL34]] %resize = "tf.ResizeNearestNeighbor"(%arg0, %shape) {align_corners = false, device = "", half_pixel_centers = false} : (tensor<1x?x?x1xi32>, tensor<2xi32>) -> tensor<1x3x3x1xi32> @@ -996,17 +996,17 @@ func.func @imag_resize_nearest_dyn_img(%arg0: tensor<1x?x?x1xi32>) -> tensor<1x3 // CHECK-LABEL: func @imag_resize_nearest_full_dyn func.func @imag_resize_nearest_full_dyn(%arg0: tensor<1x?x?x1xi32>, %arg1: tensor<2xi32>) -> tensor<1x?x?x1xi32> { - // CHECK-DAG: [[VAL0:%.+]] = "tf.Const"() {value = dense<1> : tensor} - // CHECK-DAG: [[VAL1:%.+]] = "tf.Const"() {value = dense<0.000000e+00> : tensor} - // CHECK-DAG: [[VAL2:%.+]] = "tf.Const"() {value = dense<1.000000e+00> : tensor} - // CHECK-DAG: [[VAL3:%.+]] = "tf.Const"() {value = dense<1> : tensor<1xi32>} - // CHECK-DAG: [[VAL4:%.+]] = "tf.Const"() {value = dense<1> : tensor<1xi64>} - // CHECK-DAG: [[VAL5:%.+]] = "tf.Const"() {value = dense<0> : tensor} + // CHECK-DAG: [[VAL0:%.+]] = "tf.Const"() <{value = dense<1> : tensor}> + // CHECK-DAG: [[VAL1:%.+]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor}> + // CHECK-DAG: [[VAL2:%.+]] = "tf.Const"() <{value = dense<1.000000e+00> : tensor}> + // CHECK-DAG: [[VAL3:%.+]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}> + // CHECK-DAG: [[VAL4:%.+]] = "tf.Const"() <{value = dense<1> : tensor<1xi64>}> + // CHECK-DAG: [[VAL5:%.+]] = "tf.Const"() <{value = dense<0> : tensor}> // CHECK: [[VAL6:%.+]] = "tf.Shape"(%arg0) // CHECK: [[VAL7:%.+]] = "tf.Cast"([[VAL6]]) - // CHECK: [[VAL8:%.+]]:4 = "tf.Unpack"([[VAL7]]) {axis = 0 : i64} + // CHECK: [[VAL8:%.+]]:4 = "tf.Unpack"([[VAL7]]) <{axis = 0 : i64}> // CHECK: [[VAL9:%.+]] = "tf.Mul"([[VAL8]]#1, [[VAL8]]#2) - // CHECK: [[VAL10:%.+]]:2 = "tf.Unpack"(%arg1) {axis = 0 : i64} + // CHECK: [[VAL10:%.+]]:2 = "tf.Unpack"(%arg1) <{axis = 0 : i64}> // CHECK: [[VAL11:%.+]] = "tf.Mul"([[VAL10]]#0, [[VAL10]]#1) // CHECK: [[VAL12:%.+]] = "tf.ExpandDims"([[VAL8]]#0, [[VAL5]]) // CHECK: [[VAL13:%.+]] = "tf.ExpandDims"([[VAL10]]#0, [[VAL5]]) @@ -1040,7 +1040,7 @@ func.func @imag_resize_nearest_full_dyn(%arg0: tensor<1x?x?x1xi32>, %arg1: tenso // CHECK: [[VAL41:%.+]] = "tf.ExpandDims"([[VAL8]]#3, [[VAL5]]) // CHECK: [[VAL42:%.+]] = "tf.ConcatV2"([[VAL39]], [[VAL40]], [[VAL41]], [[VAL5]]) // CHECK: [[VAL43:%.+]] = "tf.Reshape"(%arg0, [[VAL42]]) - // CHECK: [[VAL44:%.+]] = "tf.GatherV2"([[VAL43]], [[VAL38]], [[VAL0]]) {batch_dims = 0 : i64} + // CHECK: [[VAL44:%.+]] = "tf.GatherV2"([[VAL43]], [[VAL38]], [[VAL0]]) <{batch_dims = 0 : i64}> // CHECK: [[VAL45:%.+]] = "tf.Reshape"([[VAL44]], [[VAL16]]) // CHECK: return [[VAL45]] %resize = "tf.ResizeNearestNeighbor"(%arg0, %arg1) {align_corners = false, device = "", half_pixel_centers = false} : (tensor<1x?x?x1xi32>, tensor<2xi32>) -> tensor<1x?x?x1xi32> @@ -1050,8 +1050,8 @@ func.func @imag_resize_nearest_full_dyn(%arg0: tensor<1x?x?x1xi32>, %arg1: tenso // CHECK-LABEL: func @xdivy // CHECK-SAME: (%[[X:.*]]: tensor<*xf32>, %[[Y:.*]]: tensor<*xf32>) func.func @xdivy(%lhs: tensor<*xf32>, %rhs: tensor<*xf32>) -> tensor<*xf32> { - // CHECK: %[[ZERO:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor} : () -> tensor - // CHECK: %[[IS_ZERO:.*]] = "tf.Equal"(%[[X]], %[[ZERO]]) {incompatible_shape_error = true} : (tensor<*xf32>, tensor) -> tensor<*xi1> + // CHECK: %[[ZERO:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor}> : () -> tensor + // CHECK: %[[IS_ZERO:.*]] = "tf.Equal"(%[[X]], %[[ZERO]]) <{incompatible_shape_error = true}> : (tensor<*xf32>, tensor) -> tensor<*xi1> // CHECK: %[[MUL:.*]] = "tf.Div"(%[[X]], %[[Y]]) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32> // CHECK: %[[RESULT:.*]] = "tf.SelectV2"(%[[IS_ZERO]], %[[X]], %[[MUL]]) : (tensor<*xi1>, tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32> %0 = "tf.Xdivy"(%lhs, %rhs) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32> @@ -1062,8 +1062,8 @@ func.func @xdivy(%lhs: tensor<*xf32>, %rhs: tensor<*xf32>) -> tensor<*xf32> { // CHECK-LABEL: func @xlog1py // CHECK-SAME: (%[[X:.*]]: tensor<*xf32>, %[[Y:.*]]: tensor<*xf32>) func.func @xlog1py(%lhs: tensor<*xf32>, %rhs: tensor<*xf32>) -> tensor<*xf32> { - // CHECK: %[[ZERO:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor} : () -> tensor - // CHECK: %[[IS_ZERO:.*]] = "tf.Equal"(%[[X]], %[[ZERO]]) {incompatible_shape_error = true} : (tensor<*xf32>, tensor) -> tensor<*xi1> + // CHECK: %[[ZERO:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor}> : () -> tensor + // CHECK: %[[IS_ZERO:.*]] = "tf.Equal"(%[[X]], %[[ZERO]]) <{incompatible_shape_error = true}> : (tensor<*xf32>, tensor) -> tensor<*xi1> // CHECK: %[[LOG:.*]] = "tf.Log1p"(%[[Y]]) : (tensor<*xf32>) -> tensor<*xf32> // CHECK: %[[MUL:.*]] = "tf.Mul"(%[[X]], %[[LOG]]) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32> // CHECK: %[[RESULT:.*]] = "tf.SelectV2"(%[[IS_ZERO]], %[[X]], %[[MUL]]) : (tensor<*xi1>, tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32> @@ -1075,8 +1075,8 @@ func.func @xlog1py(%lhs: tensor<*xf32>, %rhs: tensor<*xf32>) -> tensor<*xf32> { // CHECK-LABEL: func @xlogy // CHECK-SAME: (%[[X:.*]]: tensor<*xf32>, %[[Y:.*]]: tensor<*xf32>) func.func @xlogy(%lhs: tensor<*xf32>, %rhs: tensor<*xf32>) -> tensor<*xf32> { - // CHECK: %[[ZERO:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor} : () -> tensor - // CHECK: %[[IS_ZERO:.*]] = "tf.Equal"(%[[X]], %[[ZERO]]) {incompatible_shape_error = true} : (tensor<*xf32>, tensor) -> tensor<*xi1> + // CHECK: %[[ZERO:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor}> : () -> tensor + // CHECK: %[[IS_ZERO:.*]] = "tf.Equal"(%[[X]], %[[ZERO]]) <{incompatible_shape_error = true}> : (tensor<*xf32>, tensor) -> tensor<*xi1> // CHECK: %[[LOG:.*]] = "tf.Log"(%[[Y]]) : (tensor<*xf32>) -> tensor<*xf32> // CHECK: %[[MUL:.*]] = "tf.Mul"(%[[X]], %[[LOG]]) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32> // CHECK: %[[RESULT:.*]] = "tf.SelectV2"(%[[IS_ZERO]], %[[X]], %[[MUL]]) : (tensor<*xi1>, tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32> @@ -1089,9 +1089,9 @@ func.func @xlogy(%lhs: tensor<*xf32>, %rhs: tensor<*xf32>) -> tensor<*xf32> { func.func @size_to_prod_shape_i32(%arg0 : tensor<1x?x2x3xf32>) -> tensor { %0 = "tf.Size"(%arg0) : (tensor<1x?x2x3xf32>) -> tensor func.return %0 : tensor - // CHECK: %[[CONSTANT:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor + // CHECK: %[[CONSTANT:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK: %[[SHAPE:.*]] = "tf.Shape"(%arg0) : (tensor<1x?x2x3xf32>) -> tensor<4xi32> - // CHECK: %[[PROD:.*]] = "tf.Prod"(%[[SHAPE]], %[[CONSTANT]]) {keep_dims = false} : (tensor<4xi32>, tensor) -> tensor + // CHECK: %[[PROD:.*]] = "tf.Prod"(%[[SHAPE]], %[[CONSTANT]]) <{keep_dims = false}> : (tensor<4xi32>, tensor) -> tensor // CHECK: return %[[PROD]] } @@ -1099,9 +1099,9 @@ func.func @size_to_prod_shape_i32(%arg0 : tensor<1x?x2x3xf32>) -> tensor { func.func @size_to_prod_shape_i64(%arg0 : tensor<1x?x2x3xf32>) -> tensor { %0 = "tf.Size"(%arg0) : (tensor<1x?x2x3xf32>) -> tensor func.return %0 : tensor - // CHECK: %[[CONSTANT:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor + // CHECK: %[[CONSTANT:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK: %[[SHAPE:.*]] = "tf.Shape"(%arg0) : (tensor<1x?x2x3xf32>) -> tensor<4xi64> - // CHECK: %[[PROD:.*]] = "tf.Prod"(%[[SHAPE]], %[[CONSTANT]]) {keep_dims = false} : (tensor<4xi64>, tensor) -> tensor + // CHECK: %[[PROD:.*]] = "tf.Prod"(%[[SHAPE]], %[[CONSTANT]]) <{keep_dims = false}> : (tensor<4xi64>, tensor) -> tensor // CHECK: return %[[PROD]] } @@ -1109,9 +1109,9 @@ func.func @size_to_prod_shape_i64(%arg0 : tensor<1x?x2x3xf32>) -> tensor { func.func @is_finite(%arg0: tensor<3x4xf32>) -> tensor<3x4xi1> { %0 = "tf.IsFinite"(%arg0) : (tensor<3x4xf32>) -> tensor<3x4xi1> func.return %0 : tensor<3x4xi1> - // CHECK: %[[ZERO:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor} : () -> tensor + // CHECK: %[[ZERO:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor}> : () -> tensor // CHECK: %[[SUB:.*]] = "tf.Sub"(%arg0, %arg0) : (tensor<3x4xf32>, tensor<3x4xf32>) -> tensor<3x4xf32> - // CHECK: %[[RESULT:.*]] = "tf.Equal"(%[[SUB]], %[[ZERO]]) {incompatible_shape_error = true} : (tensor<3x4xf32>, tensor) -> tensor<3x4xi1> + // CHECK: %[[RESULT:.*]] = "tf.Equal"(%[[SUB]], %[[ZERO]]) <{incompatible_shape_error = true}> : (tensor<3x4xf32>, tensor) -> tensor<3x4xi1> // CHECK: return %[[RESULT]] } @@ -1119,9 +1119,9 @@ func.func @is_finite(%arg0: tensor<3x4xf32>) -> tensor<3x4xi1> { func.func @is_finite_dynamic(%arg0: tensor) -> tensor { %0 = "tf.IsFinite"(%arg0) : (tensor) -> tensor func.return %0 : tensor - // CHECK: %[[ZERO:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor} : () -> tensor + // CHECK: %[[ZERO:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor}> : () -> tensor // CHECK: %[[SUB:.*]] = "tf.Sub"(%arg0, %arg0) : (tensor, tensor) -> tensor - // CHECK: %[[RESULT:.*]] = "tf.Equal"(%[[SUB]], %[[ZERO]]) {incompatible_shape_error = true} : (tensor, tensor) -> tensor + // CHECK: %[[RESULT:.*]] = "tf.Equal"(%[[SUB]], %[[ZERO]]) <{incompatible_shape_error = true}> : (tensor, tensor) -> tensor // CHECK: return %[[RESULT]] } @@ -1131,11 +1131,11 @@ func.func @roll_scalar_axis(%arg0: tensor<3x8x4xi32>) -> tensor<3x8x4xi32> { %0 = "tf.Roll"(%arg0, %shift, %axis) : (tensor<3x8x4xi32>, tensor, tensor) -> tensor<3x8x4xi32> func.return %0 : tensor<3x8x4xi32> // CHECK-LABEL: roll_scalar_axis - // CHECK-DAG: %[[CST:.*]] = "tf.Const"() {value = dense<[0, 6, 0]> : tensor<3xi64>} : () -> tensor<3xi64> - // CHECK-DAG: %[[CST0:.*]] = "tf.Const"() {value = dense<[3, 2, 4]> : tensor<3xi64>} : () -> tensor<3xi64> - // CHECK-DAG: %[[CST1:.*]] = "tf.Const"() {value = dense<0> : tensor<3xi64>} : () -> tensor<3xi64> - // CHECK-DAG: %[[CST2:.*]] = "tf.Const"() {value = dense<[3, 6, 4]> : tensor<3xi64>} : () -> tensor<3xi64> - // CHECK-DAG: %[[CST3:.*]] = "tf.Const"() {value = dense<1> : tensor} : () -> tensor + // CHECK-DAG: %[[CST:.*]] = "tf.Const"() <{value = dense<[0, 6, 0]> : tensor<3xi64>}> : () -> tensor<3xi64> + // CHECK-DAG: %[[CST0:.*]] = "tf.Const"() <{value = dense<[3, 2, 4]> : tensor<3xi64>}> : () -> tensor<3xi64> + // CHECK-DAG: %[[CST1:.*]] = "tf.Const"() <{value = dense<0> : tensor<3xi64>}> : () -> tensor<3xi64> + // CHECK-DAG: %[[CST2:.*]] = "tf.Const"() <{value = dense<[3, 6, 4]> : tensor<3xi64>}> : () -> tensor<3xi64> + // CHECK-DAG: %[[CST3:.*]] = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor // CHECK: %[[SLICE:.*]] = "tf.Slice"(%arg0, %[[CST]], %[[CST0]]) : (tensor<3x8x4xi32>, tensor<3xi64>, tensor<3xi64>) -> tensor<3x2x4xi32> // CHECK: %[[SLICE1:.*]] = "tf.Slice"(%arg0, %[[CST1]], %[[CST2]]) : (tensor<3x8x4xi32>, tensor<3xi64>, tensor<3xi64>) -> tensor<3x6x4xi32> // CHECK: %[[CONCAT:.*]] = "tf.ConcatV2"(%[[SLICE]], %[[SLICE1]], %[[CST3]]) : (tensor<3x2x4xi32>, tensor<3x6x4xi32>, tensor) -> tensor<3x8x4xi32> @@ -1148,11 +1148,11 @@ func.func @roll_1d_axis(%arg0: tensor<3x8x4xi32>) -> tensor<3x8x4xi32> { %0 = "tf.Roll"(%arg0, %shift, %axis) : (tensor<3x8x4xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<3x8x4xi32> func.return %0 : tensor<3x8x4xi32> // CHECK-LABEL: roll_1d_axis - // CHECK-DAG: %[[CST:.*]] = "tf.Const"() {value = dense<[0, 6, 0]> : tensor<3xi64>} : () -> tensor<3xi64> - // CHECK-DAG: %[[CST0:.*]] = "tf.Const"() {value = dense<[3, 2, 4]> : tensor<3xi64>} : () -> tensor<3xi64> - // CHECK-DAG: %[[CST1:.*]] = "tf.Const"() {value = dense<0> : tensor<3xi64>} : () -> tensor<3xi64> - // CHECK-DAG: %[[CST2:.*]] = "tf.Const"() {value = dense<[3, 6, 4]> : tensor<3xi64>} : () -> tensor<3xi64> - // CHECK-DAG: %[[CST3:.*]] = "tf.Const"() {value = dense<1> : tensor} : () -> tensor + // CHECK-DAG: %[[CST:.*]] = "tf.Const"() <{value = dense<[0, 6, 0]> : tensor<3xi64>}> : () -> tensor<3xi64> + // CHECK-DAG: %[[CST0:.*]] = "tf.Const"() <{value = dense<[3, 2, 4]> : tensor<3xi64>}> : () -> tensor<3xi64> + // CHECK-DAG: %[[CST1:.*]] = "tf.Const"() <{value = dense<0> : tensor<3xi64>}> : () -> tensor<3xi64> + // CHECK-DAG: %[[CST2:.*]] = "tf.Const"() <{value = dense<[3, 6, 4]> : tensor<3xi64>}> : () -> tensor<3xi64> + // CHECK-DAG: %[[CST3:.*]] = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor // CHECK: %[[SLICE:.*]] = "tf.Slice"(%arg0, %[[CST]], %[[CST0]]) : (tensor<3x8x4xi32>, tensor<3xi64>, tensor<3xi64>) -> tensor<3x2x4xi32> // CHECK: %[[SLICE1:.*]] = "tf.Slice"(%arg0, %[[CST1]], %[[CST2]]) : (tensor<3x8x4xi32>, tensor<3xi64>, tensor<3xi64>) -> tensor<3x6x4xi32> // CHECK: %[[CONCAT:.*]] = "tf.ConcatV2"(%[[SLICE]], %[[SLICE1]], %[[CST3]]) : (tensor<3x2x4xi32>, tensor<3x6x4xi32>, tensor) -> tensor<3x8x4xi32> @@ -1165,15 +1165,15 @@ func.func @roll_multiple_axis(%arg0: tensor<3x8x4xi32>) -> tensor<3x8x4xi32> { %0 = "tf.Roll"(%arg0, %shift, %axis) : (tensor<3x8x4xi32>, tensor<2xi32>, tensor<2xi32>) -> tensor<3x8x4xi32> func.return %0 : tensor<3x8x4xi32> // CHECK-LABEL: roll_multiple_axis - // CHECK-DAG: %[[CST:.*]] = "tf.Const"() {value = dense<[1, 0, 0]> : tensor<3xi64>} : () -> tensor<3xi64> - // CHECK-DAG: %[[CST0:.*]] = "tf.Const"() {value = dense<[2, 8, 4]> : tensor<3xi64>} : () -> tensor<3xi64> - // CHECK-DAG: %[[CST1:.*]] = "tf.Const"() {value = dense<[1, 8, 4]> : tensor<3xi64>} : () -> tensor<3xi64> - // CHECK-DAG: %[[CST2:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor - // CHECK-DAG: %[[CST3:.*]] = "tf.Const"() {value = dense<[0, 6, 0]> : tensor<3xi64>} : () -> tensor<3xi64> - // CHECK-DAG: %[[CST4:.*]] = "tf.Const"() {value = dense<[3, 2, 4]> : tensor<3xi64>} : () -> tensor<3xi64> - // CHECK-DAG: %[[CST5:.*]] = "tf.Const"() {value = dense<0> : tensor<3xi64>} : () -> tensor<3xi64> - // CHECK-DAG: %[[CST6:.*]] = "tf.Const"() {value = dense<[3, 6, 4]> : tensor<3xi64>} : () -> tensor<3xi64> - // CHECK-DAG: %[[CST7:.*]] = "tf.Const"() {value = dense<1> : tensor} : () -> tensor + // CHECK-DAG: %[[CST:.*]] = "tf.Const"() <{value = dense<[1, 0, 0]> : tensor<3xi64>}> : () -> tensor<3xi64> + // CHECK-DAG: %[[CST0:.*]] = "tf.Const"() <{value = dense<[2, 8, 4]> : tensor<3xi64>}> : () -> tensor<3xi64> + // CHECK-DAG: %[[CST1:.*]] = "tf.Const"() <{value = dense<[1, 8, 4]> : tensor<3xi64>}> : () -> tensor<3xi64> + // CHECK-DAG: %[[CST2:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor + // CHECK-DAG: %[[CST3:.*]] = "tf.Const"() <{value = dense<[0, 6, 0]> : tensor<3xi64>}> : () -> tensor<3xi64> + // CHECK-DAG: %[[CST4:.*]] = "tf.Const"() <{value = dense<[3, 2, 4]> : tensor<3xi64>}> : () -> tensor<3xi64> + // CHECK-DAG: %[[CST5:.*]] = "tf.Const"() <{value = dense<0> : tensor<3xi64>}> : () -> tensor<3xi64> + // CHECK-DAG: %[[CST6:.*]] = "tf.Const"() <{value = dense<[3, 6, 4]> : tensor<3xi64>}> : () -> tensor<3xi64> + // CHECK-DAG: %[[CST7:.*]] = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor // CHECK: %[[SLICE:.*]] = "tf.Slice"(%arg0, %[[CST]], %[[CST0]]) : (tensor<3x8x4xi32>, tensor<3xi64>, tensor<3xi64>) -> tensor<2x8x4xi32> // CHECK: %[[SLICE1:.*]] = "tf.Slice"(%arg0, %[[CST5]], %[[CST1]]) : (tensor<3x8x4xi32>, tensor<3xi64>, tensor<3xi64>) -> tensor<1x8x4xi32> // CHECK: %[[CONCAT:.*]] = "tf.ConcatV2"(%[[SLICE]], %[[SLICE1]], %[[CST2]]) : (tensor<2x8x4xi32>, tensor<1x8x4xi32>, tensor) -> tensor<3x8x4xi32> @@ -1184,8 +1184,8 @@ func.func @roll_multiple_axis(%arg0: tensor<3x8x4xi32>) -> tensor<3x8x4xi32> { } func.func @roll_dynamic_shape(%arg0: tensor) -> tensor { - %axis = "tf.Const"() {value = dense<1> : tensor} : () -> tensor - %shift = "tf.Const"() {value = dense<2> : tensor} : () -> tensor + %axis = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor + %shift = "tf.Const"() <{value = dense<2> : tensor}> : () -> tensor %0 = "tf.Roll"(%arg0, %shift, %axis) : (tensor, tensor, tensor) -> tensor func.return %0 : tensor // CHECK-LABEL: roll_dynamic_shape @@ -1193,7 +1193,7 @@ func.func @roll_dynamic_shape(%arg0: tensor) -> tensor { } func.func @roll_non_constant_axis(%arg0: tensor<3x8x4xi32>, %arg1: tensor) -> tensor<3x8x4xi32> { - %shift = "tf.Const"() {value = dense<2> : tensor} : () -> tensor + %shift = "tf.Const"() <{value = dense<2> : tensor}> : () -> tensor %0 = "tf.Roll"(%arg0, %shift, %arg1) : (tensor<3x8x4xi32>, tensor, tensor) -> tensor<3x8x4xi32> func.return %0 : tensor<3x8x4xi32> // CHECK-LABEL: roll_non_constant_axis @@ -1201,7 +1201,7 @@ func.func @roll_non_constant_axis(%arg0: tensor<3x8x4xi32>, %arg1: tensor) } func.func @roll_non_constant_shift(%arg0: tensor<3x8x4xi32>, %arg1: tensor) -> tensor<3x8x4xi32> { - %axis = "tf.Const"() {value = dense<2> : tensor} : () -> tensor + %axis = "tf.Const"() <{value = dense<2> : tensor}> : () -> tensor %0 = "tf.Roll"(%arg0, %arg1, %axis) : (tensor<3x8x4xi32>, tensor, tensor) -> tensor<3x8x4xi32> func.return %0 : tensor<3x8x4xi32> // CHECK-LABEL: roll_non_constant_shift @@ -1213,9 +1213,9 @@ func.func @scatter_nd_updates(%arg0: tensor<14xf32>, %arg1: tensor<1x1xi32>, %ar func.return %0 : tensor<14xf32> // CHECK-LABEL: scatter_nd_updates - // CHECK-DAG: %[[CST:.*]] = "tf.Const"() {value = dense<1.000000e+00> : tensor} : () -> tensor - // CHECK-DAG: %[[CST0:.*]] = "tf.Const"() {value = dense<1.000000e+00> : tensor<1xf32>} : () -> tensor<1xf32> - // CHECK-DAG: %[[CST1:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<14xf32>} : () -> tensor<14xf32> + // CHECK-DAG: %[[CST:.*]] = "tf.Const"() <{value = dense<1.000000e+00> : tensor}> : () -> tensor + // CHECK-DAG: %[[CST0:.*]] = "tf.Const"() <{value = dense<1.000000e+00> : tensor<1xf32>}> : () -> tensor<1xf32> + // CHECK-DAG: %[[CST1:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<14xf32>}> : () -> tensor<14xf32> // CHECK: %[[SCATTER:.*]] = "tf.TensorScatterAdd"(%cst_1, %arg1, %[[CST0]]) : (tensor<14xf32>, tensor<1x1xi32>, tensor<1xf32>) -> tensor<14xf32> // CHECK: %[[SUB:.*]] = "tf.Sub"(%[[CST]], %[[SCATTER]]) : (tensor, tensor<14xf32>) -> tensor<14xf32> // CHECK: %[[MUL:.*]] = "tf.Mul"(%[[SUB]], %arg0) : (tensor<14xf32>, tensor<14xf32>) -> tensor<14xf32> @@ -1229,17 +1229,17 @@ func.func @scatter_nd_updates_bool(%arg0: tensor<1x24xi1>, %arg1: tensor<1x2x2xi func.return %0 : tensor<1x24xi1> // CHECK-LABEL: scatter_nd_updates_bool( -// CHECK-DAG: %[[CST:.*]] = "tf.Const"() {value = dense<1> : tensor} : () -> tensor -// CHECK-DAG: %[[CST0:.*]] = "tf.Const"() {value = dense<1> : tensor<1x2xi32>} : () -> tensor<1x2xi32> -// CHECK-DAG: %[[CST1:.*]] = "tf.Const"() {value = dense<0> : tensor<1x24xi32>} : () -> tensor<1x24xi32> -// CHECK: %[[CAST0:.*]] = "tf.Cast"(%arg0) {Truncate = false} : (tensor<1x24xi1>) -> tensor<1x24xi32> -// CHECK: %[[CAST1:.*]] = "tf.Cast"(%arg2) {Truncate = false} : (tensor<1x2xi1>) -> tensor<1x2xi32> +// CHECK-DAG: %[[CST:.*]] = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor +// CHECK-DAG: %[[CST0:.*]] = "tf.Const"() <{value = dense<1> : tensor<1x2xi32>}> : () -> tensor<1x2xi32> +// CHECK-DAG: %[[CST1:.*]] = "tf.Const"() <{value = dense<0> : tensor<1x24xi32>}> : () -> tensor<1x24xi32> +// CHECK: %[[CAST0:.*]] = "tf.Cast"(%arg0) <{Truncate = false}> : (tensor<1x24xi1>) -> tensor<1x24xi32> +// CHECK: %[[CAST1:.*]] = "tf.Cast"(%arg2) <{Truncate = false}> : (tensor<1x2xi1>) -> tensor<1x2xi32> // CHECK: %[[SCATTER:.*]] = "tf.TensorScatterAdd"(%[[CST1]], %arg1, %[[CST0]]) : (tensor<1x24xi32>, tensor<1x2x2xi32>, tensor<1x2xi32>) -> tensor<1x24xi32> // CHECK: %[[SUB:.*]] = "tf.Sub"(%[[CST]], %[[SCATTER]]) : (tensor, tensor<1x24xi32>) -> tensor<1x24xi32> // CHECK: %[[MUL:.*]] = "tf.Mul"(%[[SUB]], %[[CAST0]]) : (tensor<1x24xi32>, tensor<1x24xi32>) -> tensor<1x24xi32> // CHECK: %[[SCATTER1:.*]] = "tf.TensorScatterAdd"(%[[CST1]], %arg1, %[[CAST1]]) : (tensor<1x24xi32>, tensor<1x2x2xi32>, tensor<1x2xi32>) -> tensor<1x24xi32> // CHECK: %[[ADD:.*]] = "tf.AddV2"(%[[MUL]], %[[SCATTER1]]) : (tensor<1x24xi32>, tensor<1x24xi32>) -> tensor<1x24xi32> -// CHECK: %[[CAST2:.*]] = "tf.Cast"(%[[ADD]]) {Truncate = false} : (tensor<1x24xi32>) -> tensor<1x24xi1> +// CHECK: %[[CAST2:.*]] = "tf.Cast"(%[[ADD]]) <{Truncate = false}> : (tensor<1x24xi32>) -> tensor<1x24xi1> // CHECK: return %[[CAST2]] : tensor<1x24xi1> } @@ -1250,11 +1250,11 @@ func.func @scatter_nd_updates_bool(%arg0: tensor<1x24xi1>, %arg1: tensor<1x2x2xi // CHECK-LABEL: func @simple_softmax // CHECK-SAME: (%[[ARG0:.*]]: tensor<2x3xf32>) func.func @simple_softmax(%arg0: tensor<2x3xf32>) -> tensor<2x3xf32> { - // CHECK-DAG: %[[AXIS:.*]] = "tf.Const"() {value = dense<-1> : tensor<1xi64>} : () -> tensor<1xi64> - // CHECK-DAG: %[[MAX:.*]] = "tf.Max"(%[[ARG0]], %[[AXIS]]) {keep_dims = true} : (tensor<2x3xf32>, tensor<1xi64>) -> tensor<2x1xf32> + // CHECK-DAG: %[[AXIS:.*]] = "tf.Const"() <{value = dense<-1> : tensor<1xi64>}> : () -> tensor<1xi64> + // CHECK-DAG: %[[MAX:.*]] = "tf.Max"(%[[ARG0]], %[[AXIS]]) <{keep_dims = true}> : (tensor<2x3xf32>, tensor<1xi64>) -> tensor<2x1xf32> // CHECK-DAG: %[[SHIFTED:.*]] = "tf.Sub"(%[[ARG0]], %[[MAX]]) : (tensor<2x3xf32>, tensor<2x1xf32>) -> tensor<2x3xf32> // CHECK-DAG: %[[EXP:.*]] = "tf.Exp"(%[[SHIFTED]]) : (tensor<2x3xf32>) -> tensor<2x3xf32> - // CHECK-DAG: %[[SUM:.*]] = "tf.Sum"(%[[EXP]], %[[AXIS]]) {keep_dims = true} : (tensor<2x3xf32>, tensor<1xi64>) -> tensor<2x1xf32> + // CHECK-DAG: %[[SUM:.*]] = "tf.Sum"(%[[EXP]], %[[AXIS]]) <{keep_dims = true}> : (tensor<2x3xf32>, tensor<1xi64>) -> tensor<2x1xf32> // CHECK-DAG: %[[RESULT:.*]] = "tf.Div"(%[[EXP]], %[[SUM]]) : (tensor<2x3xf32>, tensor<2x1xf32>) -> tensor<2x3xf32> // CHECK: return %[[RESULT]] %0 = "tf.Softmax"(%arg0) : (tensor<2x3xf32>) -> tensor<2x3xf32> @@ -1277,11 +1277,11 @@ func.func @unranked_softmax(%arg0: tensor<*xf32>) -> tensor<*xf32> { // CHECK-LABEL: func @simple_logsoftmax // CHECK-SAME: (%[[ARG0:.*]]: tensor<2x3xf32>) func.func @simple_logsoftmax(%arg0: tensor<2x3xf32>) -> tensor<2x3xf32> { - // CHECK-DAG: %[[AXIS:.*]] = "tf.Const"() {value = dense<-1> : tensor<1xi64>} : () -> tensor<1xi64> - // CHECK-DAG: %[[MAX:.*]] = "tf.Max"(%[[ARG0]], %[[AXIS]]) {keep_dims = true} : (tensor<2x3xf32>, tensor<1xi64>) -> tensor<2x1xf32> + // CHECK-DAG: %[[AXIS:.*]] = "tf.Const"() <{value = dense<-1> : tensor<1xi64>}> : () -> tensor<1xi64> + // CHECK-DAG: %[[MAX:.*]] = "tf.Max"(%[[ARG0]], %[[AXIS]]) <{keep_dims = true}> : (tensor<2x3xf32>, tensor<1xi64>) -> tensor<2x1xf32> // CHECK-DAG: %[[SHIFTED:.*]] = "tf.Sub"(%[[ARG0]], %[[MAX]]) : (tensor<2x3xf32>, tensor<2x1xf32>) -> tensor<2x3xf32> // CHECK-DAG: %[[EXP:.*]] = "tf.Exp"(%[[SHIFTED]]) : (tensor<2x3xf32>) -> tensor<2x3xf32> - // CHECK-DAG: %[[SUM:.*]] = "tf.Sum"(%[[EXP]], %[[AXIS]]) {keep_dims = true} : (tensor<2x3xf32>, tensor<1xi64>) -> tensor<2x1xf32> + // CHECK-DAG: %[[SUM:.*]] = "tf.Sum"(%[[EXP]], %[[AXIS]]) <{keep_dims = true}> : (tensor<2x3xf32>, tensor<1xi64>) -> tensor<2x1xf32> // CHECK-DAG: %[[LOG:.*]] = "tf.Log"(%[[SUM]]) : (tensor<2x1xf32>) -> tensor<2x1xf32> // CHECK-DAG: %[[RESULT:.*]] = "tf.Sub"(%[[SHIFTED]], %[[LOG]]) : (tensor<2x3xf32>, tensor<2x1xf32>) -> tensor<2x3xf32> // CHECK: return %[[RESULT]] @@ -1299,10 +1299,10 @@ func.func @unranked_logsoftmax(%arg0: tensor<*xf32>) -> tensor<*xf32> { // CHECK-LABEL: func @selu // CHECK-SAME: (%[[FEATURES:.*]]: tensor<1x4x4x3xf32>) -> tensor<1x4x4x3xf32> { func.func @selu(%arg0: tensor<1x4x4x3xf32>) -> tensor<1x4x4x3xf32> { - // CHECK-DAG: %[[ZERO:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor} : () -> tensor - // CHECK-DAG: %[[SCALE:.*]] = "tf.Const"() {value = dense<1.05070102> : tensor} : () -> tensor - // CHECK-DAG: %[[SCALED_ALPHA:.*]] = "tf.Const"() {value = dense<1.75809932> : tensor} : () -> tensor - // CHECK-NEXT: %[[ONE:.*]] = "tf.Const"() {value = dense<1.000000e+00> : tensor} : () -> tensor + // CHECK-DAG: %[[ZERO:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor}> : () -> tensor + // CHECK-DAG: %[[SCALE:.*]] = "tf.Const"() <{value = dense<1.05070102> : tensor}> : () -> tensor + // CHECK-DAG: %[[SCALED_ALPHA:.*]] = "tf.Const"() <{value = dense<1.75809932> : tensor}> : () -> tensor + // CHECK-NEXT: %[[ONE:.*]] = "tf.Const"() <{value = dense<1.000000e+00> : tensor}> : () -> tensor // CHECK-DAG: %[[PRED:.*]] = "tf.Greater"(%[[FEATURES]], %[[ZERO]]) : (tensor<1x4x4x3xf32>, tensor) -> tensor<1x4x4x3xi1> // CHECK-NEXT: %[[SCALED_FEATURES:.*]] = "tf.Mul"(%[[FEATURES]], %[[SCALE]]) : (tensor<1x4x4x3xf32>, tensor) -> tensor<1x4x4x3xf32> // CHECK-NEXT: %[[EXP:.*]] = "tf.Exp"(%[[FEATURES]]) : (tensor<1x4x4x3xf32>) -> tensor<1x4x4x3xf32> @@ -1317,9 +1317,9 @@ func.func @selu(%arg0: tensor<1x4x4x3xf32>) -> tensor<1x4x4x3xf32> { // CHECK-LABEL: func @selu_grad // CHECK-SAME: (%[[GRADIENTS:.*]]: tensor<4x8xf32>, %[[FEATURES:.*]]: tensor<4x8xf32>) -> tensor<4x8xf32> { func.func @selu_grad(%gradients: tensor<4x8xf32>, %features: tensor<4x8xf32>) -> tensor<4x8xf32> { - // CHECK-DAG: %[[ZERO:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor} : () -> tensor - // CHECK-DAG: %[[SCALE:.*]] = "tf.Const"() {value = dense<1.05070102> : tensor} : () -> tensor - // CHECK-DAG: %[[SCALED_ALPHA:.*]] = "tf.Const"() {value = dense<1.75809932> : tensor} : () -> tensor + // CHECK-DAG: %[[ZERO:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor}> : () -> tensor + // CHECK-DAG: %[[SCALE:.*]] = "tf.Const"() <{value = dense<1.05070102> : tensor}> : () -> tensor + // CHECK-DAG: %[[SCALED_ALPHA:.*]] = "tf.Const"() <{value = dense<1.75809932> : tensor}> : () -> tensor // CHECK-DAG: %[[PRED:.*]] = "tf.Greater"(%[[FEATURES]], %[[ZERO]]) : (tensor<4x8xf32>, tensor) -> tensor<4x8xi1> // CHECK-NEXT: %[[SCALED_GRADIENTS:.*]] = "tf.Mul"(%[[GRADIENTS]], %[[SCALE]]) : (tensor<4x8xf32>, tensor) -> tensor<4x8xf32> // CHECK-NEXT: %[[FEATURES_PLUS_SCALED_ALPHA:.*]] = "tf.AddV2"(%[[FEATURES]], %[[SCALED_ALPHA]]) : (tensor<4x8xf32>, tensor) -> tensor<4x8xf32> @@ -1335,7 +1335,7 @@ func.func @selu_grad(%gradients: tensor<4x8xf32>, %features: tensor<4x8xf32>) -> func.func @expm1(%arg0: tensor<3x4xf32>) -> tensor<3x4xf32> { %0 = "tf.Expm1"(%arg0) : (tensor<3x4xf32>) -> tensor<3x4xf32> func.return %0 : tensor<3x4xf32> - // CHECK: %[[ONE:.*]] = "tf.Const"() {value = dense<1.000000e+00> : tensor} : () -> tensor + // CHECK: %[[ONE:.*]] = "tf.Const"() <{value = dense<1.000000e+00> : tensor}> : () -> tensor // CHECK: %[[EXP:.*]] = "tf.Exp"(%[[ARG0]]) : (tensor<3x4xf32>) -> tensor<3x4xf32> // CHECK: %[[RESULT:.*]] = "tf.Sub"(%[[EXP]], %[[ONE]]) : (tensor<3x4xf32>, tensor) -> tensor<3x4xf32> // CHECK: return %[[RESULT]] @@ -1344,11 +1344,11 @@ func.func @expm1(%arg0: tensor<3x4xf32>) -> tensor<3x4xf32> { // CHECK-LABEL: func @matrix_band_part // CHECK-SAME: (%[[INPUT:.*]]: tensor<4x5xf32>, %[[NUM_LOWER:.*]]: tensor, %[[NUM_UPPER:.*]]: tensor) -> tensor<4x5xf32> { func.func @matrix_band_part(%input: tensor<4x5xf32>, %num_lower: tensor, %num_upper: tensor) -> tensor<4x5xf32> { - // CHECK-DAG: %[[ZERO:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor - // CHECK-DAG: %[[OFFSET:.*]] = "tf.Const"() {{.+}} : () -> tensor<4x5xi64> - // CHECK-DAG: %[[M:.*]] = "tf.Const"() {value = dense<4> : tensor} : () -> tensor - // CHECK-DAG: %[[N:.*]] = "tf.Const"() {value = dense<5> : tensor} : () -> tensor - // CHECK-DAG: %[[ZEROS_LIKE:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor<4x5xf32>} : () -> tensor<4x5xf32> + // CHECK-DAG: %[[ZERO:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor + // CHECK-DAG: %[[OFFSET:.*]] = "tf.Const"() <{{.+}}> : () -> tensor<4x5xi64> + // CHECK-DAG: %[[M:.*]] = "tf.Const"() <{value = dense<4> : tensor}> : () -> tensor + // CHECK-DAG: %[[N:.*]] = "tf.Const"() <{value = dense<5> : tensor}> : () -> tensor + // CHECK-DAG: %[[ZEROS_LIKE:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<4x5xf32>}> : () -> tensor<4x5xf32> // CHECK-DAG: %[[LE:.*]] = "tf.Less"(%[[NUM_LOWER]], %[[ZERO]]) : (tensor, tensor) -> tensor // CHECK-DAG: %[[NUM_LOWER_OR_M:.*]] = "tf.SelectV2"(%[[LE]], %[[M]], %[[NUM_LOWER]]) : (tensor, tensor, tensor) -> tensor // CHECK-DAG: %[[LE1:.*]] = "tf.Less"(%[[NUM_UPPER]], %[[ZERO]]) : (tensor, tensor) -> tensor @@ -1373,17 +1373,17 @@ func.func @rank3_matrix_band_part(%input: tensor, %num_lower: tensor< // CHECK-LABEL: func @dynamic_shape_matrix_band_part // CHECK-SAME: (%[[INPUT:.*]]: tensor, %[[NUM_LOWER:.*]]: tensor, %[[NUM_UPPER:.*]]: tensor) -> tensor { func.func @dynamic_shape_matrix_band_part(%input: tensor, %num_lower: tensor, %num_upper: tensor) -> tensor { - // CHECK-DAG: %[[ZERO:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor - // CHECK-DAG: %[[ONE:.*]] = "tf.Const"() {value = dense<1> : tensor} : () -> tensor - // CHECK-DAG: %[[NEG_ONE:.*]] = "tf.Const"() {value = dense<-1> : tensor} : () -> tensor - // CHECK-DAG: %[[ZERO_1D:.*]] = "tf.Const"() {value = dense<0> : tensor<1xi32>} : () -> tensor<1xi32> - // CHECK-DAG: %[[ONE_1D:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32> - // CHECK-DAG: %[[TWO_1D:.*]] = "tf.Const"() {value = dense<2> : tensor<1xi32>} : () -> tensor<1xi32> - // CHECK-DAG: %[[ZERO_F32:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor} : () -> tensor + // CHECK-DAG: %[[ZERO:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor + // CHECK-DAG: %[[ONE:.*]] = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor + // CHECK-DAG: %[[NEG_ONE:.*]] = "tf.Const"() <{value = dense<-1> : tensor}> : () -> tensor + // CHECK-DAG: %[[ZERO_1D:.*]] = "tf.Const"() <{value = dense<0> : tensor<1xi32>}> : () -> tensor<1xi32> + // CHECK-DAG: %[[ONE_1D:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}> : () -> tensor<1xi32> + // CHECK-DAG: %[[TWO_1D:.*]] = "tf.Const"() <{value = dense<2> : tensor<1xi32>}> : () -> tensor<1xi32> + // CHECK-DAG: %[[ZERO_F32:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor}> : () -> tensor // CHECK-DAG: %[[SHAPE:.*]] = "tf.Shape"(%[[INPUT]]) : (tensor) -> tensor<2xi32> - // CHECK-DAG: %[[M:.*]] = "tf.StridedSlice"(%[[SHAPE]], %[[ZERO_1D]], %[[ONE_1D]], %[[ONE_1D]]) {begin_mask = 0 : i64, ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 1 : i64} : (tensor<2xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor + // CHECK-DAG: %[[M:.*]] = "tf.StridedSlice"(%[[SHAPE]], %[[ZERO_1D]], %[[ONE_1D]], %[[ONE_1D]]) <{begin_mask = 0 : i64, ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 1 : i64}> : (tensor<2xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor // CHECK-DAG: %[[SHAPE1:.*]] = "tf.Shape"(%[[INPUT]]) : (tensor) -> tensor<2xi32> - // CHECK-DAG: %[[N:.*]] = "tf.StridedSlice"(%[[SHAPE1]], %[[ONE_1D]], %[[TWO_1D]], %[[ONE_1D]]) {begin_mask = 0 : i64, ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 1 : i64} : (tensor<2xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor + // CHECK-DAG: %[[N:.*]] = "tf.StridedSlice"(%[[SHAPE1]], %[[ONE_1D]], %[[TWO_1D]], %[[ONE_1D]]) <{begin_mask = 0 : i64, ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 1 : i64}> : (tensor<2xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor // CHECK-DAG: %[[LE:.*]] = "tf.Less"(%[[NUM_LOWER]], %[[ZERO]]) : (tensor, tensor) -> tensor // CHECK-DAG: %[[NUM_LOWER_OR_M:.*]] = "tf.SelectV2"(%[[LE]], %[[M]], %[[NUM_LOWER]]) : (tensor, tensor, tensor) -> tensor // CHECK-DAG: %[[LE1:.*]] = "tf.Less"(%[[NUM_UPPER]], %[[ZERO]]) : (tensor, tensor) -> tensor diff --git a/tensorflow/compiler/mlir/tensorflow/tests/mark_ops_for_outside_compilation.mlir b/tensorflow/compiler/mlir/tensorflow/tests/mark_ops_for_outside_compilation.mlir index 837c37b9a71be8..ca1e4c99549d94 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/mark_ops_for_outside_compilation.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/mark_ops_for_outside_compilation.mlir @@ -131,7 +131,7 @@ func.func @ignore_const_foldable_ops(%arg0: tensor) -> () { // CHECK-LABEL: func @op_string_result func.func @op_string_result() -> tensor { %0 = "tf_device.cluster"() ({ - // CHECK: "tf.Const"() {value = dense<1> : tensor} + // CHECK: "tf.Const"() <{value = dense<1> : tensor}> // CHECK-NOT: _xla_outside_compilation // CHECK: "tf.Const" // CHECK-SAME: _xla_outside_compilation @@ -148,7 +148,7 @@ func.func @op_string_result() -> tensor { // CHECK-LABEL: func @op_string_operand func.func @op_string_operand(%arg0: tensor) -> tensor { %0 = "tf_device.cluster"() ({ - // CHECK: "tf.Const"() {value = dense<1> : tensor} + // CHECK: "tf.Const"() <{value = dense<1> : tensor}> // CHECK-NOT: _xla_outside_compilation // CHECK: "tf.StringToNumber" // CHECK-SAME: _xla_outside_compilation @@ -166,7 +166,7 @@ func.func @op_string_operand(%arg0: tensor) -> tensor { // CHECK-LABEL: func @op_string_operand_string_result func.func @op_string_operand_string_result(%arg0: tensor) -> tensor { %0 = "tf_device.cluster"() ({ - // CHECK: "tf.Const"() {value = dense<1> : tensor} + // CHECK: "tf.Const"() <{value = dense<1> : tensor}> // CHECK-NOT: _xla_outside_compilation // CHECK: "tf.Identity" // CHECK-SAME: _xla_outside_compilation @@ -187,7 +187,7 @@ func.func @op_string_operand_string_result(%arg0: tensor) -> te // CHECK-LABEL: func @ops_inside_tf_if_outside_compiled func.func @ops_inside_tf_if_outside_compiled(%arg0: tensor, %arg1: tensor) -> tensor { %0 = "tf_device.cluster"() ({ - // CHECK: "tf.Const"() {value = dense<1> : tensor} + // CHECK: "tf.Const"() <{value = dense<1> : tensor}> // CHECK-NOT: _xla_outside_compilation // CHECK: "tf.IfRegion" // CHECK: "tf.StringToNumber" @@ -212,22 +212,22 @@ func.func @ops_inside_tf_if_outside_compiled(%arg0: tensor, %arg1: tensor, %arg1: tensor) -> tensor { %0 = "tf_device.cluster"() ({ - // CHECK: "tf.Const"() {value = dense<1> : tensor} + // CHECK: "tf.Const"() <{value = dense<1> : tensor}> // CHECK-NOT: _xla_outside_compilation // CHECK: "tf.IfRegion" + // CHECK: <{is_stateless // CHECK-NOT: _xla_outside_compilation %1 = "tf.Const"() {value = dense<1> : tensor} : () -> tensor %2 = "tf.IfRegion"(%arg0) ({ %3 = "tf.Const"() {value = dense<1.0> : tensor} : () -> tensor "tf.Yield"(%3) : (tensor) -> () }, { - // CHECK: "tf.Const"() {_xla_outside_compilation = "auto0", value = dense<"1.0"> : tensor} + // CHECK: "tf.Const"() <{value = dense<"1.0"> : tensor}> {_xla_outside_compilation = "auto0"} // CHECK-NEXT: "tf.StringToNumber" // CHECK-SAME: _xla_outside_compilation %4 = "tf.Const"() {value = dense<"1.0"> : tensor} : () -> tensor %5 = "tf.StringToNumber"(%4) {out_type = f32} : (tensor) -> tensor "tf.Yield"(%5) : (tensor) -> () - // CHECK: {is_stateless }) {is_stateless = true} : (tensor) -> (tensor) %6 = "tf.Identity"(%2) : (tensor) -> tensor tf_device.return %6: tensor @@ -241,34 +241,34 @@ func.func @if_region_string_op(%arg0: tensor, %arg1: tensor) -> tenso // CHECK-LABEL: func @nested_if_region_string_op func.func @nested_if_region_string_op(%arg0: tensor, %arg1: tensor) -> tensor { %0 = "tf_device.cluster"() ({ - // CHECK: "tf.Const"() {value = dense<1> : tensor} + // CHECK: "tf.Const"() <{value = dense<1> : tensor}> // CHECK-NOT: _xla_outside_compilation // CHECK: "tf.IfRegion" + // CHECK: <{is_stateless // CHECK-NOT: _xla_outside_compilation %1 = "tf.Const"() {value = dense<1> : tensor} : () -> tensor %2 = "tf.IfRegion"(%arg0) ({ %3 = "tf.Const"() {value = dense<1.0> : tensor} : () -> tensor "tf.Yield"(%3) : (tensor) -> () }, { - // CHECK: "tf.Const"() {value = dense : tensor} + // CHECK: "tf.Const"() <{value = dense : tensor}> + // CHECK: <{is_stateless // CHECK-NOT: _xla_outside_compilation %4 = "tf.Const"() {value = dense : tensor} : () -> tensor %5 = "tf.IfRegion"(%4)({ - // CHECK: "tf.Const"() {_xla_outside_compilation = "auto0", value = dense<"1.0"> : tensor} + // CHECK: "tf.Const"() <{value = dense<"1.0"> : tensor}> {_xla_outside_compilation = "auto0"} // CHECK-NEXT: "tf.StringToNumber" // CHECK-SAME: _xla_outside_compilation %6 = "tf.Const"() {value = dense<"1.0"> : tensor} : () -> tensor %7 = "tf.StringToNumber"(%6) {out_type = f32} : (tensor) -> tensor "tf.Yield"(%7) : (tensor) -> () }, { - // CHECK: "tf.Const"() {value = dense<1.000000e+00> : tensor} + // CHECK: "tf.Const"() <{value = dense<1.000000e+00> : tensor}> // CHECK-NOT: _xla_outside_compilation %8 = "tf.Const"() {value = dense<1.000000e+00> : tensor} : () -> tensor "tf.Yield"(%8) : (tensor) -> () - // CHECK: {is_stateless }){is_stateless = true} : (tensor) -> (tensor) "tf.Yield"(%5) : (tensor) -> () - // CHECK: {is_stateless }) {is_stateless = true} : (tensor) -> (tensor) %9 = "tf.Identity"(%2) : (tensor) -> tensor tf_device.return %9: tensor @@ -282,7 +282,7 @@ func.func @nested_if_region_string_op(%arg0: tensor, %arg1: tensor) - // CHECK-LABEL: func @ops_inside_while_outside_compiled func.func @ops_inside_while_outside_compiled(%arg0: tensor, %arg1: tensor) -> tensor { %0 = "tf_device.cluster"() ({ - // CHECK: "tf.Const"() {value = dense<1.000000e+00> : tensor} + // CHECK: "tf.Const"() <{value = dense<1.000000e+00> : tensor}> // CHECK-NOT: _xla_outside_compilation // CHECK: "tf.WhileRegion" // CHECK: "tf.StringToNumber" @@ -313,9 +313,10 @@ func.func @ops_inside_while_outside_compiled(%arg0: tensor, %arg1: tensor, %arg1: tensor) -> tensor { %0 = "tf_device.cluster"() ({ - // CHECK: "tf.Const"() {value = dense<1.000000e+00> : tensor} + // CHECK: "tf.Const"() <{value = dense<1.000000e+00> : tensor}> // CHECK-NOT: _xla_outside_compilation // CHECK: "tf.WhileRegion" + // CHECK: <{is_stateless = true %1 = "tf.Const"() {value = dense<1.0> : tensor} : () -> tensor %2:2 = "tf.WhileRegion"(%1, %arg0) ({ ^bb0(%carg0: tensor, %carg1: tensor): @@ -329,10 +330,9 @@ func.func @while_region_unsupported_op(%arg0: tensor, %arg1: tensor : tensor} : () -> tensor - // CHECK: "tf.Const"() {value = dense<1.000000e+00> : tensor} + // CHECK: "tf.Const"() <{value = dense<1.000000e+00> : tensor}> %4 = "tf.Const"() {value = dense<1.0> : tensor} : () -> tensor "tf.Yield"(%4, %sub) : (tensor, tensor) -> () - // CHECK: {is_stateless = true }) {is_stateless = true} : (tensor, tensor) -> (tensor, tensor) // CHECK: "tf.Identity" // CHECK-NOT: _xla_outside_compilation diff --git a/tensorflow/compiler/mlir/tensorflow/tests/merge_control_flow.mlir b/tensorflow/compiler/mlir/tensorflow/tests/merge_control_flow.mlir index bd5f805b810f69..a9e4fc902401b6 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/merge_control_flow.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/merge_control_flow.mlir @@ -807,19 +807,19 @@ func.func @nested_IfRegions_with_same_predicate_same_block_level_merged() { func.func @two_overlapped_if_groups_with_no_dependency_merged() { // CHECK: tf_device.cluster // CHECK: "tf.IfRegion" - // CHECK: "tf.Const"() {value = dense<1.000000e+00> : tensor} : () -> tensor - // CHECK "tf.Const"() {value = dense<5.000000e+00> : tensor} : () -> tensor - // CHECK "tf.Const"() {value = dense<9.000000e+00> : tensor} : () -> tensor - // CHECK: "tf.Const"() {value = dense<2.000000e+00> : tensor} : () -> tensor - // CHECK "tf.Const"() {value = dense<6.000000e+00> : tensor} : () -> tensor - // CHECK "tf.Const"() {value = dense<1.000000e+01> : tensor} : () -> tensor + // CHECK: "tf.Const"() <{value = dense<1.000000e+00> : tensor}> : () -> tensor + // CHECK "tf.Const"() <{value = dense<5.000000e+00> : tensor}> : () -> tensor + // CHECK "tf.Const"() <{value = dense<9.000000e+00> : tensor}> : () -> tensor + // CHECK: "tf.Const"() <{value = dense<2.000000e+00> : tensor}> : () -> tensor + // CHECK "tf.Const"() <{value = dense<6.000000e+00> : tensor}> : () -> tensor + // CHECK "tf.Const"() <{value = dense<1.000000e+01> : tensor}> : () -> tensor // CHECK: "tf.IfRegion" - // CHECK: "tf.Const"() {value = dense<3.000000e+00> : tensor} : () -> tensor - // CHECK "tf.Const"() {value = dense<7.000000e+00> : tensor} : () -> tensor - // CHECK "tf.Const"() {value = dense<1.100000e+01> : tensor} : () -> tensor - // CHECK: "tf.Const"() {value = dense<4.000000e+00> : tensor} : () -> tensor - // CHECK "tf.Const"() {value = dense<8.000000e+00> : tensor} : () -> tensor - // CHECK "tf.Const"() {value = dense<1.200000e+01> : tensor} : () -> tensor + // CHECK: "tf.Const"() <{value = dense<3.000000e+00> : tensor}> : () -> tensor + // CHECK "tf.Const"() <{value = dense<7.000000e+00> : tensor}> : () -> tensor + // CHECK "tf.Const"() <{value = dense<1.100000e+01> : tensor}> : () -> tensor + // CHECK: "tf.Const"() <{value = dense<4.000000e+00> : tensor}> : () -> tensor + // CHECK "tf.Const"() <{value = dense<8.000000e+00> : tensor}> : () -> tensor + // CHECK "tf.Const"() <{value = dense<1.200000e+01> : tensor}> : () -> tensor // CHECK-NOT: "tf.IfRegion" "tf_device.cluster"() ({ %0 = "tf.Const"() {value = dense : tensor} : () -> tensor @@ -888,12 +888,12 @@ func.func @two_overlapped_if_groups_with_no_dependency_merged() { // CHECK: "tf.E" // CHECK: "tf.F" // CHECK: "tf.IfRegion" - // CHECK: "tf.Const"() {value = dense<1.000000e+00> : tensor} : () -> tensor - // CHECK: "tf.Const"() {value = dense<3.000000e+00> : tensor} : () -> tensor - // CHECK: "tf.Const"() {value = dense<5.000000e+00> : tensor} : () -> tensor - // CHECK: "tf.Const"() {value = dense<2.000000e+00> : tensor} : () -> tensor - // CHECK: "tf.Const"() {value = dense<4.000000e+00> : tensor} : () -> tensor - // CHECK; "tf.Const"() {value = dense<6.000000e+00> : tensor} : () -> tensor + // CHECK: "tf.Const"() <{value = dense<1.000000e+00> : tensor}> : () -> tensor + // CHECK: "tf.Const"() <{value = dense<3.000000e+00> : tensor}> : () -> tensor + // CHECK: "tf.Const"() <{value = dense<5.000000e+00> : tensor}> : () -> tensor + // CHECK: "tf.Const"() <{value = dense<2.000000e+00> : tensor}> : () -> tensor + // CHECK: "tf.Const"() <{value = dense<4.000000e+00> : tensor}> : () -> tensor + // CHECK; "tf.Const"() <{value = dense<6.000000e+00> : tensor}> : () -> tensor // CHECK-NOT: "tf.IfRegion" func.func @two_overlapped_if_groups_with_dependency_not_merged_for_first_if_region_group() { "tf_device.cluster"() ({ diff --git a/tensorflow/compiler/mlir/tensorflow/tests/mlprogram.mlir b/tensorflow/compiler/mlir/tensorflow/tests/mlprogram.mlir index 5b0958aa37696b..6ded59b51ad8d4 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/mlprogram.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/mlprogram.mlir @@ -123,7 +123,7 @@ module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, p // CHECK-LABEL @lowers_string_ops // CHECK-DAG: ml_program.global public @vars.Variable_1([]) : tensor func.func @lowers_string_ops(%arg0: tensor<128xi32>, %arg1: tensor<128xi32>, %arg2: tensor<128x1xi32>, %arg3: tensor<128x90xi32>, %arg4: tensor<128x90xi32>, %arg5: tensor<128x90xi32>, %arg6: tensor<128x90x64xf32>, %arg7: tensor<128x90x64xf32>) -> tensor { - // CHECK: %0 = ml_program.global_load @vars.Variable_1 : tensor + // CHECK: %[[v0:.*]] = ml_program.global_load @vars.Variable_1 : tensor %0 = tf_executor.graph { %outputs_4, %control_5 = tf_executor.island wraps "tf.VarHandleOp"() {container = "", shared_name = "Variable"} : () -> tensor>> %outputs_10, %control_11 = tf_executor.island wraps "tf.VarHandleOp"() {container = "", shared_name = "Variable_1"} : () -> tensor>> diff --git a/tensorflow/compiler/mlir/tensorflow/tests/prepare_tpu_computation_for_tf_export.mlir b/tensorflow/compiler/mlir/tensorflow/tests/prepare_tpu_computation_for_tf_export.mlir index 45ee57ad75d72a..021cad3b78be8f 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/prepare_tpu_computation_for_tf_export.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/prepare_tpu_computation_for_tf_export.mlir @@ -3,8 +3,8 @@ // CHECK-LABEL: @ShardingAttr func.func @ShardingAttr(%arg0: tensor<128x10xf32> {mhlo.sharding = "\08\03\1A\02\01\02\22\02\00\01"}, %arg1: tensor<10x1024xf32> {mhlo.sharding = "\08\01\1A\01\01\22\01\00"}, %arg2: tensor<128x1024xf32> {mhlo.sharding = ""}) -> (tensor<128x10xf32>, tensor<10x1024xf32>, tensor<128x1024xf32>) { - // CHECK: %[[SHARDED_ARG0:.*]] = "tf.XlaSharding"(%arg0) {_XlaSharding = "\08\03\1A\02\01\02\22\02\00\01", sharding = "\08\03\1A\02\01\02\22\02\00\01"} - // CHECK: %[[SHARDED_ARG1:.*]] = "tf.XlaSharding"(%arg1) {_XlaSharding = "\08\01\1A\01\01\22\01\00", sharding = "\08\01\1A\01\01\22\01\00"} + // CHECK: %[[SHARDED_ARG0:.*]] = "tf.XlaSharding"(%arg0) <{_XlaSharding = "\08\03\1A\02\01\02\22\02\00\01", sharding = "\08\03\1A\02\01\02\22\02\00\01"}> + // CHECK: %[[SHARDED_ARG1:.*]] = "tf.XlaSharding"(%arg1) <{_XlaSharding = "\08\01\1A\01\01\22\01\00", sharding = "\08\01\1A\01\01\22\01\00"}> // CHECK: "tf.Identity"(%[[SHARDED_ARG1]]) %0 = "tf.Identity"(%arg1) : (tensor<10x1024xf32>) -> tensor<10x1024xf32> diff --git a/tensorflow/compiler/mlir/tensorflow/tests/promote_resources_to_args.mlir b/tensorflow/compiler/mlir/tensorflow/tests/promote_resources_to_args.mlir index 14f20b6cf5a18a..faf2a960aed7d8 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/promote_resources_to_args.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/promote_resources_to_args.mlir @@ -62,7 +62,7 @@ func.func @main(%arg0: tensor) { func.func @main(%arg0: tensor) -> tensor<2xf32> { // CHECK-NOT: "tf.VarHandleOp" // CHECK-NOT: "tf.ReadVariableOp" - // CHECK: %[[CONST:.*]] = "tf.Const"() {value = dense<4.200000e+01> : tensor} + // CHECK: %[[CONST:.*]] = "tf.Const"() <{value = dense<4.200000e+01> : tensor}> // CHECK: %[[ADD1:[0-9]*]] = "tf.AddV2"(%arg1, %[[CONST]]) // CHECK: %[[ADD2:[0-9]*]] = "tf.AddV2"(%[[ADD1]], %arg1) // CHECK: %[[PACK:[0-9]*]] = "tf.Pack"(%[[CONST]], %[[ADD2]]) @@ -133,7 +133,7 @@ func.func @main(%arg0: tensor) -> tensor<2xf32> { // CHECK-LABEL: func @main(%arg0: tensor) -> (tensor<2xf32>, tensor {tf.resource_name = "x"}) func.func @main(%arg0: tensor) -> tensor<2xf32> { // CHECK-NOT: "tf.AssignVariableOp" - // CHECK: %[[CONST:.*]] = "tf.Const"() {value = dense<4.200000e+01> : tensor} + // CHECK: %[[CONST:.*]] = "tf.Const"() <{value = dense<4.200000e+01> : tensor}> // CHECK: %[[ADD1:[0-9]*]] = "tf.AddV2"(%[[CONST]], %[[CONST]]) // CHECK: %[[ADD2:[0-9]*]] = "tf.AddV2"(%[[ADD1]], %[[ADD1]]) // CHECK: %[[PACK:[0-9]*]] = "tf.Pack"(%[[CONST]], %[[ADD2]]) @@ -222,7 +222,7 @@ func.func @main(%arg0: tensor>>, %arg1: tensor func.func @main(%arg0: tensor>>, %arg1: tensor) { %0 = "tf.Const"() {value = dense<4.200000e+01> : tensor} : () -> tensor "tf.AssignVariableOp"(%arg0, %0) : (tensor>>, tensor) -> () - // CHECK: %[[CONST:.*]] = "tf.Const"() {value = dense<1.050000e+03> : tensor} + // CHECK: %[[CONST:.*]] = "tf.Const"() <{value = dense<1.050000e+03> : tensor}> %1 = "tf.Const"() {value = dense<1.050000e+03> : tensor} : () -> tensor "tf.AssignVariableOp"(%arg0, %1) : (tensor>>, tensor) -> () // CHECK-NEXT: return %[[CONST]] : tensor @@ -241,7 +241,7 @@ func.func @main(%arg0: tensor>>, %arg1: tensor func.func @main(%arg0: tensor>>, %arg1: tensor) -> tensor { %0 = "tf.Const"() {value = dense<4.200000e+01> : tensor} : () -> tensor "tf.AssignVariableOp"(%arg0, %0) : (tensor>>, tensor) -> () - // CHECK: %[[CONST:.*]] = "tf.Const"() {value = dense<1.050000e+03> : tensor} + // CHECK: %[[CONST:.*]] = "tf.Const"() <{value = dense<1.050000e+03> : tensor}> %1 = "tf.Const"() {value = dense<1.050000e+03> : tensor} : () -> tensor "tf.AssignVariableOp"(%arg0, %1) : (tensor>>, tensor) -> () // CHECK-NEXT: return %[[CONST]], %[[CONST]] : tensor, tensor @@ -257,13 +257,13 @@ func.func @main(%arg0: tensor>>, %arg1: tensor // CHECK-SAME: %arg1: tensor // CHECK-SAME: -> (tensor, tensor) func.func @main(%arg0: tensor>>, %arg1: tensor) -> tensor { - // CHECK-NEXT: %[[CONST_0:.*]] = "tf.Const"() {value = dense<4.200000e+01> : tensor} + // CHECK-NEXT: %[[CONST_0:.*]] = "tf.Const"() <{value = dense<4.200000e+01> : tensor}> %0 = "tf.Const"() {value = dense<4.200000e+01> : tensor} : () -> tensor "tf.AssignVariableOp"(%arg0, %0) : (tensor>>, tensor) -> () %1 = "tf.ReadVariableOp"(%arg0) : (tensor>>) -> tensor // CHECK-NEXT: %[[ADD:[a-z0-9]+]] = "tf.AddV2"(%[[CONST_0]], %[[CONST_0]]) %2 = "tf.AddV2"(%1, %1) : (tensor, tensor) -> tensor - // CHECK-NEXT: %[[CONST_1:.*]] = "tf.Const"() {value = dense<1.050000e+03> : tensor} + // CHECK-NEXT: %[[CONST_1:.*]] = "tf.Const"() <{value = dense<1.050000e+03> : tensor}> %3 = "tf.Const"() {value = dense<1.050000e+03> : tensor} : () -> tensor "tf.AssignVariableOp"(%arg0, %3) : (tensor>>, tensor) -> () // CHECK-NEXT: return %[[ADD]], %[[CONST_1]] : tensor, tensor diff --git a/tensorflow/compiler/mlir/tensorflow/tests/region-control-flow-to-functional.mlir b/tensorflow/compiler/mlir/tensorflow/tests/region-control-flow-to-functional.mlir index e26a299a85c66d..eff3e38ab5ace2 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/region-control-flow-to-functional.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/region-control-flow-to-functional.mlir @@ -7,11 +7,11 @@ // CHECK-NEXT: "tf.Abs" func.func @testSimple(%arg0: tensor, %arg1: tensor<*xf32>) -> tensor<*xf32> { // CHECK: "tf.If" - // CHECK-SAME: _attr0 = false - // CHECK-SAME: _xla_propagate_compile_time_consts = true // CHECK-NOT: attr1 // CHECK-SAME: else_branch = @test_else_name // CHECK-SAME: then_branch = @test_then_name + // CHECK-SAME: _attr0 = false + // CHECK-SAME: _xla_propagate_compile_time_consts = true %0 = "tf.IfRegion"(%arg0) ({ %1 = "tf.Abs"(%arg1) : (tensor<*xf32>) -> tensor<*xf32> "tf.Yield"(%1) : (tensor<*xf32>) -> () @@ -31,10 +31,10 @@ func.func @testSimple(%arg0: tensor, %arg1: tensor<*xf32>) -> tensor<*xf32> // CHECK-NEXT: "tf.Abs" func.func @testSimpleEmptyBranchNames(%arg0: tensor, %arg1: tensor<*xf32>) -> tensor<*xf32> { // CHECK: "tf.If" - // CHECK-SAME: _attr0 = false // CHECK-NOT: attr1 // CHECK-SAME: else_branch = @tf.IfRegion_else // CHECK-SAME: then_branch = @tf.IfRegion_then + // CHECK-SAME: _attr0 = false %0 = "tf.IfRegion"(%arg0) ({ %1 = "tf.Abs"(%arg1) : (tensor<*xf32>) -> tensor<*xf32> "tf.Yield"(%1) : (tensor<*xf32>) -> () @@ -78,7 +78,7 @@ func.func @testIfCondition(%arg0: tensor, %arg1: tensor<2xf32>) -> tensor<2x // CHECK-NEXT: constant dense<0.0 func.func @testIfConstant(%arg0: tensor) -> tensor<2xf32> { %cst_zero = arith.constant dense<0.0> : tensor<2xf32> - // CHECK: "tf.If"(%arg0) {{.*}} else_branch = @tf.IfRegion_else{{.+}}then_branch = @tf.IfRegion_then + // CHECK: "tf.If"(%arg0) <{else_branch = @tf.IfRegion_else{{.+}}then_branch = @tf.IfRegion_then %0 = "tf.IfRegion"(%arg0) ({ "tf.Yield"(%cst_zero) : (tensor<2xf32>) -> () }, { @@ -98,7 +98,7 @@ func.func @testIfConstant(%arg0: tensor) -> tensor<2xf32> { // CHECK: func private @tf.IfRegion1_then // CHECK-NEXT: "tf.LogicalNot" // CHECK-NEXT: "tf.Asin" -// CHECK-NEXT: "tf.If"({{.+}}) {{.*}} else_branch = @tf.IfRegion_else, {{.+}} then_branch = @tf.IfRegion_then} +// CHECK-NEXT: "tf.If"({{.+}}) <{else_branch = @tf.IfRegion_else, {{.+}} then_branch = @tf.IfRegion_then} // CHECK: func private @tf.IfRegion_else // CHECK-NEXT: "tf.Neg" @@ -106,7 +106,7 @@ func.func @testIfConstant(%arg0: tensor) -> tensor<2xf32> { // CHECK-NEXT: "tf.Abs" func.func @testNested(%arg0: tensor, %arg1: tensor<*xf32>) -> tensor<*xf32> { - // CHECK: "tf.If"({{.+}}) {{.*}} else_branch = @tf.IfRegion1_else, {{.+}} then_branch = @tf.IfRegion1_then} + // CHECK: "tf.If"({{.+}}) <{else_branch = @tf.IfRegion1_else, {{.+}} then_branch = @tf.IfRegion1_then} %0 = "tf.IfRegion"(%arg0) ({ // Outer Then %cond = "tf.LogicalNot"(%arg0) : (tensor) -> tensor @@ -137,7 +137,7 @@ func.func @testNested(%arg0: tensor, %arg1: tensor<*xf32>) -> tensor<*xf32> func.func private @testIf1Then(tensor<*xf32>) -> tensor<*xf32> func.func private @testIf1Else(tensor<*xf32>) -> tensor<*xf32> func.func @testIf1Result(%arg0: tensor, %arg1: tensor<*xf32>) -> tensor<*xf32> { - // CHECK: "tf.If"({{.+}}) {{.*}} else_branch = @testIf1Else, {{.+}} then_branch = @testIf1Then} + // CHECK: "tf.If"({{.+}}) <{else_branch = @testIf1Else, {{.+}} then_branch = @testIf1Then} %0 = "tf.IfRegion"(%arg0) ({ %1 = func.call @testIf1Then(%arg1) : (tensor<*xf32>) -> tensor<*xf32> "tf.Yield"(%1) : (tensor<*xf32>) -> () @@ -155,7 +155,7 @@ func.func @testIf1Result(%arg0: tensor, %arg1: tensor<*xf32>) -> tensor<*xf3 func.func private @testIf1Then(tensor<*xf32>) -> tensor<*xf32> func.func private @testIf1Else(tensor<*xf32>) -> tensor<*xf32> func.func @testIf2Result(%arg0: tensor, %arg1: tensor<2xf32>) -> tensor<2xf32> { - // CHECK: "tf.If"({{.+}}) {{.*}} else_branch = @testIf1Else, {{.+}} then_branch = @testIf1Then} + // CHECK: "tf.If"({{.+}}) <{else_branch = @testIf1Else, {{.+}} then_branch = @testIf1Then} %0 = "tf.IfRegion"(%arg0) ({ %1 = "tf.Cast"(%arg1) {Truncate = false} : (tensor<2xf32>) -> tensor<*xf32> %2 = func.call @testIf1Then(%1) : (tensor<*xf32>) -> tensor<*xf32> @@ -175,7 +175,7 @@ func.func @testIf2Result(%arg0: tensor, %arg1: tensor<2xf32>) -> tensor<2xf3 func.func private @testIf1Then(tensor<*xf32>) -> tensor<*xf32> func.func private @testIf1Else(tensor<*xf32>) -> tensor<*xf32> func.func @testIf2Result(%arg0: tensor, %arg1: tensor<2xf32>) -> tensor<2xf32> { - // CHECK: "tf.If"({{.+}}) {{.*}} else_branch = @testIf1Else, {{.+}} then_branch = @testIf1Then} + // CHECK: "tf.If"({{.+}}) <{else_branch = @testIf1Else, {{.+}} then_branch = @testIf1Then} %0 = "tf.IfRegion"(%arg0) ({ %1 = "tf.Cast"(%arg1) {Truncate = false} : (tensor<2xf32>) -> tensor %2 = "tf.Cast"(%1) {Truncate = false} : (tensor) -> tensor<*xf32> @@ -197,8 +197,8 @@ func.func @testIf2Result(%arg0: tensor, %arg1: tensor<2xf32>) -> tensor<2xf3 func.func private @testIf1Then(tensor<*xf32>) -> tensor<*xf32> func.func private @testIf1Else(tensor<*xf32>) -> tensor<*xf32> func.func @testIfExternIncompatibleCastTrivialTransform(%arg0: tensor, %arg1: tensor<2xi64>) -> tensor<2xf32> { - // CHECK: %[[CAST:.*]] = "tf.Cast"(%arg1) {Truncate = false} : (tensor<2xi64>) -> tensor<*xf32> - // CHECK: "tf.If"(%arg0, %[[CAST]]) {{.*}} else_branch = @testIf1Else, {{.+}} then_branch = @testIf1Then} + // CHECK: %[[CAST:.*]] = "tf.Cast"(%arg1) <{Truncate = false}> : (tensor<2xi64>) -> tensor<*xf32> + // CHECK: "tf.If"(%arg0, %[[CAST]]) <{else_branch = @testIf1Else, {{.+}} then_branch = @testIf1Then} %1 = "tf.Cast"(%arg1) {Truncate = false} : (tensor<2xi64>) -> tensor<*xf32> %0 = "tf.IfRegion"(%arg0) ({ %2 = func.call @testIf1Then(%1) : (tensor<*xf32>) -> tensor<*xf32> @@ -221,7 +221,7 @@ func.func @testIfExternIncompatibleCastTrivialTransform(%arg0: tensor, %arg1 func.func private @testIf1Then(tensor<*xf32>) -> tensor<*xf32> func.func private @testIf1Else(tensor<*xf32>) -> tensor<*xf32> func.func @testIfIncompatibleCastTrivialTransform(%arg0: tensor, %arg1: tensor<2xi64>) -> tensor<2xf32> { - // CHECK: "tf.If"(%arg0, %arg1) {{.*}} else_branch = @tf.IfRegion_else{{.+}}then_branch = @tf.IfRegion_then} + // CHECK: "tf.If"(%arg0, %arg1) <{else_branch = @tf.IfRegion_else{{.+}}then_branch = @tf.IfRegion_then} %0 = "tf.IfRegion"(%arg0) ({ %1 = "tf.Cast"(%arg1) {Truncate = false} : (tensor<2xi64>) -> tensor<*xf32> %2 = func.call @testIf1Then(%1) : (tensor<*xf32>) -> tensor<*xf32> @@ -341,11 +341,11 @@ func.func @testCase(%arg0: tensor, %arg1: tensor, %arg1 : tensor) -> tensor<*xf32> { // CHECK: [[Result:%.*]]:2 = "tf.While"(%arg0, %arg1) - // CHECK-SAME: _attr0 = false - // CHECK-SAME: _xla_propagate_compile_time_consts = true // CHECK-NOT: attr1 // CHECK-SAME: body = @tf.WhileRegion_body // CHECK-SAME: cond = @tf.WhileRegion_cond + // CHECK-SAME: _attr0 = false + // CHECK-SAME: _xla_propagate_compile_time_consts = true %0:2 = "tf.WhileRegion"(%arg0, %arg1) ( { // condition, check if count has reached 0 @@ -379,7 +379,7 @@ func.func @testValidWhileRegion(%arg0 : tensor<*xf32>, %arg1 : tensor) -> t // CHECK: "tf.NotEqual" // CHECK-LABEL: testWhileRegionTypeMismatch func.func @testWhileRegionTypeMismatch(%arg0 : tensor<*xf32>, %arg1 : tensor) -> tensor<*xf32> { - // CHECK: [[Result:%.*]]:2 = "tf.While"(%arg0, %arg1) {{.*}} body = @tf.WhileRegion_body, cond = @tf.WhileRegion_cond + // CHECK: [[Result:%.*]]:2 = "tf.While"(%arg0, %arg1) <{body = @tf.WhileRegion_body, cond = @tf.WhileRegion_cond %0:2 = "tf.WhileRegion"(%arg0, %arg1) ( { // condition, check if count has reached 0 @@ -415,7 +415,7 @@ func.func @testWhileRegionTypeMismatch(%arg0 : tensor<*xf32>, %arg1 : tensor, %arg1 : tensor) -> tensor<*xf32> { %zero = arith.constant dense<0> : tensor %one = arith.constant dense<1> : tensor - // CHECK: [[Result:%.*]]:2 = "tf.While"(%arg0, %arg1) {{.*}} body = @tf.WhileRegion_body, cond = @tf.WhileRegion_cond + // CHECK: [[Result:%.*]]:2 = "tf.While"(%arg0, %arg1) <{body = @tf.WhileRegion_body, cond = @tf.WhileRegion_cond %0:2 = "tf.WhileRegion"(%arg0, %arg1) ( { ^bb0(%carg0: tensor<4xf32>, %carg1: tensor): @@ -448,7 +448,7 @@ func.func @testWhileRegionConstantSink(%arg0 : tensor<*xf32>, %arg1 : tensor, %arg1 : tensor, %arg2 : tensor) -> tensor<*xf32> { %cst = arith.constant dense<4> : tensor %limit = "tf.Add"(%arg2, %cst) : (tensor, tensor) -> tensor - // CHECK: [[Result:%.*]]:3 = "tf.While"(%arg0, %arg1, %{{.+}} body = @tf.WhileRegion_body, cond = @tf.WhileRegion_cond + // CHECK: [[Result:%.*]]:3 = "tf.While"(%arg0, %arg1, %{{.+}} <{body = @tf.WhileRegion_body, cond = @tf.WhileRegion_cond %0:2 = "tf.WhileRegion"(%arg0, %arg1) ( { ^bb0(%carg0: tensor<*xf32>, %carg1: tensor): @@ -485,7 +485,7 @@ func.func @testWhileRegionExternInBody(%arg0 : tensor<*xf32>, %arg1 : tensor : tensor %cst = arith.constant dense<4> : tensor %stride = "tf.Add"(%arg2, %cst) : (tensor, tensor) -> tensor - // CHECK: [[Result:%.*]]:3 = "tf.While"(%arg0, %arg1, %{{.+}} body = @tf.WhileRegion_body, cond = @tf.WhileRegion_cond + // CHECK: [[Result:%.*]]:3 = "tf.While"(%arg0, %arg1, %{{.+}} <{body = @tf.WhileRegion_body, cond = @tf.WhileRegion_cond %0:2 = "tf.WhileRegion"(%arg0, %arg1) ( { ^bb0(%carg0: tensor<*xf32>, %carg1: tensor): @@ -516,7 +516,7 @@ func.func @testWhileRegionExternInBodyAndCond(%arg0 : tensor<*xf32>, %arg1 : ten %stride = "tf.Add"(%arg2, %cst) : (tensor, tensor) -> tensor %cst1 = arith.constant dense<44> : tensor %limit = "tf.Add"(%arg2, %cst1) : (tensor, tensor) -> tensor - // CHECK: [[Result:%.*]]:4 = "tf.While"(%arg0, %arg1, %{{.+}}, %{{.+}} body = @tf.WhileRegion_body, cond = @tf.WhileRegion_cond + // CHECK: [[Result:%.*]]:4 = "tf.While"(%arg0, %arg1, %{{.+}}, %{{.+}} <{body = @tf.WhileRegion_body, cond = @tf.WhileRegion_cond %0:2 = "tf.WhileRegion"(%arg0, %arg1) ( { ^bb0(%carg0: tensor<*xf32>, %carg1: tensor): @@ -545,7 +545,7 @@ func.func @testWhileRegionExternInBodyAndCond(%arg0 : tensor<*xf32>, %arg1 : ten func.func @testWhileRegionSameExternInBodyAndCond(%arg0 : tensor<*xf32>, %arg1 : tensor, %arg2 : tensor) -> tensor<*xf32> { %cst = arith.constant dense<4> : tensor %stride = "tf.Add"(%arg2, %cst) : (tensor, tensor) -> tensor - // CHECK: [[Result:%.*]]:3 = "tf.While"(%arg0, %arg1, %{{.+}} body = @tf.WhileRegion_body, cond = @tf.WhileRegion_cond + // CHECK: [[Result:%.*]]:3 = "tf.While"(%arg0, %arg1, %{{.+}} <{body = @tf.WhileRegion_body, cond = @tf.WhileRegion_cond %0:2 = "tf.WhileRegion"(%arg0, %arg1) ( { ^bb0(%carg0: tensor<*xf32>, %carg1: tensor): @@ -573,7 +573,7 @@ func.func @testWhileRegionSameExternInBodyAndCond(%arg0 : tensor<*xf32>, %arg1 : func.func private @while_cond(%arg0 : tensor<*xf32>, %arg1 : tensor) -> tensor func.func private @while_body(%arg0 : tensor<*xf32>, %arg1 : tensor) -> (tensor<*xf32>, tensor) func.func @testWhileRegionTrivial(%arg0 : tensor<*xf32>, %arg1 : tensor) -> tensor<*xf32> { - // CHECK: [[Result:%.*]]:2 = "tf.While"(%arg0, %arg1) {{.*}} body = @while_body, cond = @while_cond + // CHECK: [[Result:%.*]]:2 = "tf.While"(%arg0, %arg1) <{body = @while_body, cond = @while_cond %0:2 = "tf.WhileRegion"(%arg0, %arg1) ( { ^bb0(%carg0: tensor<*xf32>, %carg1: tensor): @@ -600,7 +600,7 @@ func.func @testWhileRegionTrivial(%arg0 : tensor<*xf32>, %arg1 : tensor) -> func.func private @while_cond(%arg0 : tensor<4xf32>, %arg1 : tensor) -> tensor func.func private @while_body(%arg0 : tensor<4xf32>, %arg1 : tensor) -> (tensor<4xf32>, tensor) func.func @testWhileRegionTrivialCasts(%arg0 : tensor<*xf32>, %arg1 : tensor) -> tensor<*xf32> { - // CHECK: [[Result:%.*]]:2 = "tf.While"(%arg0, %arg1) {{.*}} body = @while_body, cond = @while_cond + // CHECK: [[Result:%.*]]:2 = "tf.While"(%arg0, %arg1) <{body = @while_body, cond = @while_cond %0:2 = "tf.WhileRegion"(%arg0, %arg1) ( { ^bb0(%carg0: tensor<*xf32>, %carg1: tensor): @@ -629,7 +629,7 @@ func.func @testWhileRegionTrivialCasts(%arg0 : tensor<*xf32>, %arg1 : tensor, %arg1 : tensor) -> tensor func.func private @while_body(%arg0 : tensor<4xf32>, %arg1 : tensor) -> (tensor<4xf32>, tensor) func.func @testWhileRegionTrivialMultipleCasts(%arg0 : tensor<*xf32>, %arg1 : tensor) -> tensor<*xf32> { - // CHECK: [[Result:%.*]]:2 = "tf.While"(%arg0, %arg1) {{.*}} body = @while_body, cond = @while_cond + // CHECK: [[Result:%.*]]:2 = "tf.While"(%arg0, %arg1) <{body = @while_body, cond = @while_cond %0:2 = "tf.WhileRegion"(%arg0, %arg1) ( { ^bb0(%carg0: tensor<*xf32>, %carg1: tensor): @@ -662,7 +662,7 @@ func.func @testWhileRegionTrivialMultipleCasts(%arg0 : tensor<*xf32>, %arg1 : te func.func private @while_cond(%arg0 : tensor<4xf32>, %arg1 : tensor) -> tensor func.func private @while_body(%arg0 : tensor<4xf32>, %arg1 : tensor) -> (tensor<4xi64>, tensor) func.func @testWhileRegionIncompatibleCast(%arg0 : tensor<*xi64>, %arg1 : tensor) -> tensor<*xi64> { - // CHECK: [[Result:%.*]]:2 = "tf.While"(%arg0, %arg1) {{.*}} body = @tf.WhileRegion_body, cond = @tf.WhileRegion_cond + // CHECK: [[Result:%.*]]:2 = "tf.While"(%arg0, %arg1) <{body = @tf.WhileRegion_body, cond = @tf.WhileRegion_cond %0:2 = "tf.WhileRegion"(%arg0, %arg1) ( { ^bb0(%carg0: tensor<*xi64>, %carg1: tensor): @@ -694,7 +694,7 @@ func.func private @while_cond(%arg0 : tensor<*xf32>, %arg1 : tensor) -> ten func.func private @while_body(%arg0 : tensor<*xf32>, %arg1 : tensor, %arg2 : tensor<*xf32>) -> (tensor<*xf32>, tensor) func.func @testWhileRegionExtern(%arg0 : tensor<*xf32>, %arg1 : tensor) -> tensor<*xf32> { %ext = "tf.Neg"(%arg0) : (tensor<*xf32>) -> tensor<*xf32> - // CHECK: [[Result:%.*]]:3 = "tf.While"(%arg0, %arg1, %{{.+}} body = @tf.WhileRegion_body, cond = @tf.WhileRegion_cond + // CHECK: [[Result:%.*]]:3 = "tf.While"(%arg0, %arg1, %{{.+}} <{body = @tf.WhileRegion_body, cond = @tf.WhileRegion_cond %0:2 = "tf.WhileRegion"(%arg0, %arg1) ( { ^bb0(%carg0: tensor<*xf32>, %carg1: tensor): @@ -723,7 +723,7 @@ func.func @testWhileRegionExtern(%arg0 : tensor<*xf32>, %arg1 : tensor) -> func.func private @while_cond(%arg0 : tensor, %arg1 : tensor<*xf32>) -> tensor func.func private @while_body(%arg0 : tensor<*xf32>, %arg1 : tensor) -> (tensor<*xf32>, tensor) func.func @testWhileRegionBlockArgMismatch(%arg0 : tensor<*xf32>, %arg1 : tensor) -> tensor<*xf32> { - // CHECK: [[Result:%.*]]:2 = "tf.While"(%arg0, %arg1) {{.*}} body = @tf.WhileRegion_body, cond = @tf.WhileRegion_cond + // CHECK: [[Result:%.*]]:2 = "tf.While"(%arg0, %arg1) <{body = @tf.WhileRegion_body, cond = @tf.WhileRegion_cond %0:2 = "tf.WhileRegion"(%arg0, %arg1) ( { ^bb0(%carg0: tensor<*xf32>, %carg1: tensor): @@ -750,7 +750,7 @@ func.func @testWhileRegionBlockArgMismatch(%arg0 : tensor<*xf32>, %arg1 : tensor func.func private @while_cond(%arg0 : tensor<*xf32>, %arg1 : tensor) -> tensor func.func private @while_body(%arg0 : tensor<*xf32>, %arg1 : tensor) -> (tensor<*xf32>, tensor) func.func @testWhileRegionTrivial(%arg0 : tensor<*xf32>, %arg1 : tensor) -> tensor<*xf32> { - // CHECK: [[Result:%.*]]:2 = "tf.While"(%arg0, %arg1) {{.*}} body = @while_body, cond = @while_cond + // CHECK: [[Result:%.*]]:2 = "tf.While"(%arg0, %arg1) <{body = @while_body, cond = @while_cond %0:2 = "tf.WhileRegion"(%arg0, %arg1) ( { ^bb0(%carg0: tensor<*xf32>, %carg1: tensor): @@ -830,11 +830,11 @@ func.func @testOverrideIfRegionXlaPropageCompileTimeConsts(%arg0: tensor, %a // CHECK-LABEL: testValidWhileRegion func.func @testValidWhileRegion(%arg0 : tensor<*xf32>, %arg1 : tensor) -> tensor<*xf32> { // CHECK: [[Result:%.*]]:2 = "tf.While"(%arg0, %arg1) - // CHECK-SAME: _attr0 = false - // CHECK-SAME: _xla_propagate_compile_time_consts = true // CHECK-NOT: attr1 // CHECK-SAME: body = @tf.WhileRegion_body // CHECK-SAME: cond = @tf.WhileRegion_cond + // CHECK-SAME: _attr0 = false + // CHECK-SAME: _xla_propagate_compile_time_consts = true %0:2 = "tf.WhileRegion"(%arg0, %arg1) ( { // condition, check if count has reached 0 @@ -881,3 +881,82 @@ func.func @testPassThroughCond(%arg0 : tensor<*xf32>, %arg1 : tensor) -> te ) { is_stateless = false, _attr0 = false, attr1 = "hello"} : (tensor<*xf32>, tensor) -> (tensor<*xf32>, tensor) func.return %0#0 : tensor<*xf32> } + +// ----- + +func.func @init(%arg0: tensor<4xf32>) -> tensor<7xf32> { + %0 = builtin.unrealized_conversion_cast to tensor<7xf32> + return %0 : tensor<7xf32> +} +func.func @next(%arg0: tensor<7xf32>, %arg1: tensor<3xf32>) -> tensor<6xf32> { + %0 = builtin.unrealized_conversion_cast to tensor<6xf32> + return %0 : tensor<6xf32> +} +func.func @finalize(%arg0: tensor<6xf32>, %arg1: tensor<2xf32>) -> tensor<5xf32> { + %0 = builtin.unrealized_conversion_cast to tensor<5xf32> + return %0 : tensor<5xf32> +} + +// CHECK-LABEL: testGeneratorDatasetRegion +func.func @testGeneratorDatasetRegion(%arg0: tensor<4xf32>, %arg1: tensor<3xf32>, %arg2: tensor, %arg3: tensor<2xf32>) { + // CHECK: "tf.GeneratorDataset" + // CHECK-DAG: @init + // CHECK-DAG: @next + // CHECK-DAG: @finalize + // CHECK: return + %0 = "tf.GeneratorDatasetRegion"(%arg0, %arg1, %arg2, %arg3) ({ + ^bb0(%arg4: tensor<4xf32>): + %1 = func.call @init(%arg4) : (tensor<4xf32>) -> tensor<7xf32> + "tf.Yield"(%1) : (tensor<7xf32>) -> () + }, { + ^bb0(%arg4: tensor<7xf32>, %arg5: tensor<3xf32>): + %1 = func.call @next(%arg4, %arg5) : (tensor<7xf32>, tensor<3xf32>) -> tensor<6xf32> + "tf.Yield"(%1) : (tensor<6xf32>) -> () + }, { + ^bb0(%arg4: tensor<6xf32>, %arg5: tensor<2xf32>): + %1 = func.call @finalize(%arg4, %arg5) : (tensor<6xf32>, tensor<2xf32>) -> tensor<5xf32> + "tf.Yield"(%1) : (tensor<5xf32>) -> () + }) {device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0", metadata = "", operandSegmentSizes = array, output_shapes = [#tf_type.shape<>], output_types = [!tf_type.string]} : (tensor<4xf32>, tensor<3xf32>, tensor, tensor<2xf32>) -> tensor + return +} + +// ----- + +func.func @init(%arg0: tensor<4xf32>) -> tensor<7xf32> { + %0 = builtin.unrealized_conversion_cast to tensor<7xf32> + return %0 : tensor<7xf32> +} +func.func @next(%arg0: tensor<3xf32>, %arg1: tensor<7xf32>) -> tensor<6xf32> { + %0 = builtin.unrealized_conversion_cast to tensor<6xf32> + return %0 : tensor<6xf32> +} +func.func @finalize(%arg0: tensor<6xf32>, %arg1: tensor<2xf32>) -> tensor<5xf32> { + %0 = builtin.unrealized_conversion_cast to tensor<5xf32> + return %0 : tensor<5xf32> +} + +// CHECK-LABEL: testGeneratorDatasetRegionWithComplexBlocks +func.func @testGeneratorDatasetRegionWithComplexBlocks(%arg0: tensor<4xf32>, %arg1: tensor<3xf32>, %arg2: tensor, %arg3: tensor<2xf32>) { + // CHECK: "tf.GeneratorDataset" + // CHECK-NOT: @init + // CHECK-NOT: @next + // CHECK-NOT: @finalize + // CHECK: -> tensor + // CHECK: return + %0 = "tf.GeneratorDatasetRegion"(%arg0, %arg1, %arg2, %arg3) ({ + ^bb0(%arg4: tensor<4xf32>): + %sum = "tf.Add"(%arg4, %arg4) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32> + %1 = func.call @init(%sum) : (tensor<4xf32>) -> tensor<7xf32> + "tf.Yield"(%1) : (tensor<7xf32>) -> () + }, { + ^bb0(%arg4: tensor<7xf32>, %arg5: tensor<3xf32>): + %1 = func.call @next(%arg5, %arg4) : (tensor<3xf32>, tensor<7xf32>) -> tensor<6xf32> + "tf.Yield"(%1) : (tensor<6xf32>) -> () + }, { + ^bb0(%arg4: tensor<6xf32>, %arg5: tensor<2xf32>): + %1 = func.call @finalize(%arg4, %arg5) : (tensor<6xf32>, tensor<2xf32>) -> tensor<5xf32> + %sum = "tf.Add"(%1, %1) : (tensor<5xf32>, tensor<5xf32>) -> tensor<5xf32> + "tf.Yield"(%sum) : (tensor<5xf32>) -> () + }) {device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0", metadata = "", operandSegmentSizes = array, output_shapes = [#tf_type.shape<>], output_types = [!tf_type.string]} : (tensor<4xf32>, tensor<3xf32>, tensor, tensor<2xf32>) -> tensor + return +} diff --git a/tensorflow/compiler/mlir/tensorflow/tests/replicate_to_island.mlir b/tensorflow/compiler/mlir/tensorflow/tests/replicate_to_island.mlir index 8e0e4558b851f0..8fec2a5bb55223 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/replicate_to_island.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/replicate_to_island.mlir @@ -268,10 +268,10 @@ func.func @device_ordinals() { func.return } -// CHECK: tf_executor.island wraps "tf.Const"() {_parallel_execution_ids = "r0:0", value = dense<1> : tensor} -// CHECK: tf_executor.island wraps "tf.Const"() {_parallel_execution_ids = "r0:0", value = dense<3> : tensor} -// CHECK: tf_executor.island wraps "tf.Const"() {_parallel_execution_ids = "r0:1", value = dense<2> : tensor} -// CHECK: tf_executor.island wraps "tf.Const"() {_parallel_execution_ids = "r0:1", value = dense<4> : tensor} +// CHECK: tf_executor.island wraps "tf.Const"() <{value = dense<1> : tensor}> {_parallel_execution_ids = "r0:0"} +// CHECK: tf_executor.island wraps "tf.Const"() <{value = dense<3> : tensor}> {_parallel_execution_ids = "r0:0"} +// CHECK: tf_executor.island wraps "tf.Const"() <{value = dense<2> : tensor}> {_parallel_execution_ids = "r0:1"} +// CHECK: tf_executor.island wraps "tf.Const"() <{value = dense<4> : tensor}> {_parallel_execution_ids = "r0:1"} // ----- // Tests parallel_execute nested inside replicate diff --git a/tensorflow/compiler/mlir/tensorflow/tests/resource_op_lifting.mlir b/tensorflow/compiler/mlir/tensorflow/tests/resource_op_lifting.mlir index 02c144467d92c5..9c00c8ee5f849c 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/resource_op_lifting.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/resource_op_lifting.mlir @@ -221,7 +221,7 @@ func.func @named_internal_resource() -> tensor<*xi32> { // CHECK-LABEL: func @cluster_with_loop func.func @cluster_with_loop() -> () { - // CHECK: %[[COUNT:.*]] = "tf.Const"() {value = dense<10> : tensor} + // CHECK: %[[COUNT:.*]] = "tf.Const"() <{value = dense<10> : tensor}> %0 = "tf.Const"() {value = dense<10> : tensor} : () -> tensor // CHECK: %[[VH:.*]] = "tf.VarHandleOp"() %1 = "tf.VarHandleOp"() {container = "c", shared_name = "v"} : () -> tensor<*x!tf_type.resource>> @@ -253,7 +253,7 @@ func.func @while_body(%arg0: tensor, %arg1: tensor<*x!tf_type.resource, tensor) -> tensor "tf.AssignVariableOp"(%arg1, %add1) : (tensor<*x!tf_type.resource>>, tensor) -> () - // CHECK-NEXT: %[[DELTA:.*]] = "tf.Const"() {value = dense<-1> : tensor} + // CHECK-NEXT: %[[DELTA:.*]] = "tf.Const"() <{value = dense<-1> : tensor}> %constant = "tf.Const"() {value = dense<-1> : tensor} : () -> tensor // CHECK-NEXT: %[[ADD2:.*]] = "tf.AddV2"(%[[BARG0]], %[[DELTA]]) %add2 = "tf.AddV2"(%arg0, %constant) : (tensor, tensor) -> tensor @@ -299,7 +299,7 @@ func.func @while_body(%arg0: tensor<*x!tf_type.resource>>) -> (tenso // CHECK-NEXT: return %[[CONST]] func.return %arg0 : tensor<*x!tf_type.resource>> } -// CHECK: func @while_cond(%arg0: tensor) +// CHECK: func @while_cond(%[[CARG0:.*]]: tensor) func.func @while_cond(%arg0: tensor<*x!tf_type.resource>>) -> tensor { %id = "tf.Identity"(%arg0) : (tensor<*x!tf_type.resource>>) -> tensor<*x!tf_type.resource>> %read = "tf.ReadVariableOp"(%id) : (tensor<*x!tf_type.resource>>) -> tensor @@ -935,7 +935,7 @@ func.func @cluster_with_caseregion(%arg0: tensor) -> tensor<4xf32> { // CHECK-LABEL: func @cluster_with_whileregion func.func @cluster_with_whileregion() -> () { - // CHECK: %[[COUNT:.*]] = "tf.Const"() {value = dense<10> : tensor} + // CHECK: %[[COUNT:.*]] = "tf.Const"() <{value = dense<10> : tensor}> // CHECK: %[[VH:.*]] = "tf.VarHandleOp"() // CHECK: %[[READ:.*]] = "tf.ReadVariableOp"(%[[VH]]) // CHECK: %[[CLUSTER:.*]] = "tf_device.cluster"() @@ -959,7 +959,7 @@ func.func @cluster_with_whileregion() -> () { // CHECK: (%[[BARG0:.+]]: tensor, %[[BARG1:.+]]: tensor): // CHECK: %[[ADD0:.*]] = "tf.AddV2"(%[[BARG1]], %[[BARG1]]) // CHECK-NEXT: %[[ADD1:.*]] = "tf.AddV2"(%[[ADD0]], %[[ADD0]]) - // CHECK-NEXT: %[[DELTA:.*]] = "tf.Const"() {value = dense<-1> : tensor} + // CHECK-NEXT: %[[DELTA:.*]] = "tf.Const"() <{value = dense<-1> : tensor}> // CHECK-NEXT: %[[ADD2:.*]] = "tf.AddV2"(%[[BARG0]], %[[DELTA]]) // CHECK-NEXT: "tf.Yield"(%[[ADD2]], %[[ADD1]]) ^bb1(%barg0: tensor, %barg1: !tf_ref, %barg2: !tf_ref, %barg3: !tf_ref): @@ -1046,7 +1046,7 @@ func.func @if_region_with_store_in_then(%arg0: tensor) { %0 = "tf.VarHandleOp"() {container = "c", shared_name = "v"} : () -> tensor<*x!tf_type.resource>> "tf_device.cluster"() ({ "tf.IfRegion"(%arg0) ({ - // CHECK: %[[CONST:.*]] = "tf.Const"() {value = dense<0.000000e+00> + // CHECK: %[[CONST:.*]] = "tf.Const"() <{value = dense<0.000000e+00> // CHECK: "tf.Yield"(%[[CONST]]) %constant = "tf.Const"() {value = dense<0.0> : tensor<4xf32>} : () -> tensor<4xf32> "tf.AssignVariableOp"(%0, %constant) : (tensor<*x!tf_type.resource>>, tensor<4xf32>) -> () @@ -1074,13 +1074,13 @@ func.func @if_region_with_store_in_both(%arg0: tensor) { %0 = "tf.VarHandleOp"() {container = "c", shared_name = "v"} : () -> tensor<*x!tf_type.resource>> "tf_device.cluster"() ({ "tf.IfRegion"(%arg0) ({ - // CHECK: %[[CONST:.*]] = "tf.Const"() {value = dense<0.000000e+00> + // CHECK: %[[CONST:.*]] = "tf.Const"() <{value = dense<0.000000e+00> // CHECK: "tf.Yield"(%[[CONST]]) %constant = "tf.Const"() {value = dense<0.0> : tensor<4xf32>} : () -> tensor<4xf32> "tf.AssignVariableOp"(%0, %constant) : (tensor<*x!tf_type.resource>>, tensor<4xf32>) -> () "tf.Yield"() : () -> () }, { - // CHECK: %[[CONST:.*]] = "tf.Const"() {value = dense<1.000000e+00> + // CHECK: %[[CONST:.*]] = "tf.Const"() <{value = dense<1.000000e+00> // CHECK: "tf.Yield"(%[[CONST]]) %constant = "tf.Const"() {value = dense<1.0> : tensor<4xf32>} : () -> tensor<4xf32> "tf.AssignVariableOp"(%0, %constant) : (tensor<*x!tf_type.resource>>, tensor<4xf32>) -> () @@ -1123,8 +1123,8 @@ func.func @test_unsupported_resource_op() -> tensor<*xi32> { // to not be lifted and arg1 to be lifted. // CHECK-LABEL: func @test_unsupported_resource_op_in_if func.func @test_unsupported_resource_op_in_if(%arg0: tensor) -> tensor<*xi32> { - // CHECK: [[VH0:%.*]] = "tf.VarHandleOp"() {container = "c", shared_name = "v"} - // CHECK: [[VH1:%.*]] = "tf.VarHandleOp"() {container = "d", shared_name = "w"} + // CHECK: [[VH0:%.*]] = "tf.VarHandleOp"() <{container = "c", shared_name = "v"}> + // CHECK: [[VH1:%.*]] = "tf.VarHandleOp"() <{container = "d", shared_name = "w"}> // CHECK-NOT: "tf.ReadVariableOp"([[VH0]]) // CHECK: [[READ1:%.*]] = "tf.ReadVariableOp"([[VH1]]) // CHECK-NOT: "tf.ReadVariableOp"([[VH0]]) diff --git a/tensorflow/compiler/mlir/tensorflow/tests/rewrite_tpu_embedding_ops.mlir b/tensorflow/compiler/mlir/tensorflow/tests/rewrite_tpu_embedding_ops.mlir index 0f1344ce73a8d0..0099129136f79a 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/rewrite_tpu_embedding_ops.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/rewrite_tpu_embedding_ops.mlir @@ -2,8 +2,8 @@ // CHECK-LABEL: func @recv_tpu_embedding_activations func.func @recv_tpu_embedding_activations() -> (tensor<512x256xf32>) { - // CHECK: %[[DATA:.*]] = "tf.XlaRecvTPUEmbeddingDeduplicationData"() {config = {{.*}}} : () -> tensor - // CHECK: %[[RESULT:.*]] = "tf.XlaRecvTPUEmbeddingActivations"(%[[DATA]]) {config = {{.*}}} : (tensor) -> tensor<512x256xf32> + // CHECK: %[[DATA:.*]] = "tf.XlaRecvTPUEmbeddingDeduplicationData"() <{config = {{.*}}}> : () -> tensor + // CHECK: %[[RESULT:.*]] = "tf.XlaRecvTPUEmbeddingActivations"(%[[DATA]]) <{config = {{.*}}}> : (tensor) -> tensor<512x256xf32> // CHECK: return %[[RESULT]] // CHECK-NOT: tf.RecvTPUEmbeddingActivations // CHECK-NOT: tf.SendTPUEmbeddingGradients @@ -14,8 +14,8 @@ func.func @recv_tpu_embedding_activations() -> (tensor<512x256xf32>) { // CHECK-LABEL: func @send_tpu_embedding_gradients func.func @send_tpu_embedding_gradients(%arg0: tensor<512x256xf32>) -> () { - // CHECK: %[[DATA:.*]] = "tf.XlaRecvTPUEmbeddingDeduplicationData"() {config = {{.*}}} : () -> tensor - // CHECK: "tf.XlaSendTPUEmbeddingGradients"(%arg0, %[[DATA]]) {config = {{.*}}, operandSegmentSizes = array} : (tensor<512x256xf32>, tensor) -> () + // CHECK: %[[DATA:.*]] = "tf.XlaRecvTPUEmbeddingDeduplicationData"() <{config = {{.*}}}> : () -> tensor + // CHECK: "tf.XlaSendTPUEmbeddingGradients"(%arg0, %[[DATA]]) <{config = {{.*}}, operandSegmentSizes = array}> : (tensor<512x256xf32>, tensor) -> () // CHECK-NOT: tf.SendTPUEmbeddingGradients // CHECK-NOT: tf.RecvTPUEmbeddingActivations diff --git a/tensorflow/compiler/mlir/tensorflow/tests/shape_inference.mlir b/tensorflow/compiler/mlir/tensorflow/tests/shape_inference.mlir index 45d5ba99fc5dea..8e9218354bad4d 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/shape_inference.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/shape_inference.mlir @@ -131,7 +131,8 @@ module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, pr // CHECK-SAME: -> tensor<1x2x3xf32> func.func @shape_from_if_to_region_bodies_to_output(%arg0: tensor, %arg1: tensor<1x2x3xf32>) -> tensor<*xf32> { %unshaped = "tf.Cast"(%arg1) : (tensor<1x2x3xf32>) -> tensor<*xf32> - %0 = "tf.IfRegion"(%arg0) ({ + // CHECK: <{is_stateless = true}> + %0 = "tf.IfRegion"(%arg0) <{is_stateless = true}> ({ // CHECK: "tf.Add"{{.+}}(tensor<1x2x3xf32>, tensor<1x2x3xf32>) -> tensor<1x2x3xf32> // CHECK: "tf.Yield"{{.+}}(tensor<1x2x3xf32>) -> () %1 = "tf.Add"(%unshaped, %unshaped) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32> @@ -141,8 +142,8 @@ module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, pr // CHECK: "tf.Yield"{{.+}}(tensor<1x2x3xf32>) -> () %2 = "tf.Sub"(%unshaped, %unshaped) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32> "tf.Yield"(%2) : (tensor<*xf32>) -> () - // CHECK: {is_stateless = true} : (tensor) -> tensor<1x2x3xf32> - }) {is_stateless = true} : (tensor) -> tensor<*xf32> + // CHECK: (tensor) -> tensor<1x2x3xf32> + }) : (tensor) -> tensor<*xf32> // CHECK: return {{.*}} : tensor<1x2x3xf32> func.return %0 : tensor<*xf32> } @@ -176,7 +177,8 @@ module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, pr // CHECK-SAME: -> tensor<1x2x3xf32> func.func @shape_from_case_to_region_bodies_to_output(%arg0: tensor, %arg1: tensor<1x2x3xf32>) -> tensor<*xf32> { %unshaped = "tf.Cast"(%arg1) : (tensor<1x2x3xf32>) -> tensor<*xf32> - %0 = "tf.CaseRegion"(%arg0) ({ + // CHECK: <{is_stateless = true}> + %0 = "tf.CaseRegion"(%arg0) <{is_stateless = true}> ({ // CHECK: "tf.Add"{{.+}}(tensor<1x2x3xf32>, tensor<1x2x3xf32>) -> tensor<1x2x3xf32> // CHECK: "tf.Yield"{{.+}}(tensor<1x2x3xf32>) -> () %1 = "tf.Add"(%unshaped, %unshaped) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32> @@ -186,8 +188,8 @@ module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, pr // CHECK: "tf.Yield"{{.+}}(tensor<1x2x3xf32>) -> () %2 = "tf.Sub"(%unshaped, %unshaped) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32> "tf.Yield"(%2) : (tensor<*xf32>) -> () - // CHECK: {is_stateless = true} : (tensor) -> tensor<1x2x3xf32> - }) {is_stateless = true} : (tensor) -> tensor<*xf32> + // CHECK: (tensor) -> tensor<1x2x3xf32> + }) : (tensor) -> tensor<*xf32> // CHECK: return {{.*}} : tensor<1x2x3xf32> func.return %0 : tensor<*xf32> } @@ -243,7 +245,8 @@ module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, pr func.func @shape_from_while_operands_to_cond_body_to_while_results(%arg0: tensor, %arg1: tensor<1x2x3xf32>) -> tensor<*xf32> { %unshaped = "tf.Cast"(%arg1) : (tensor<1x2x3xf32>) -> tensor<*xf32> // CHECK: "tf.WhileRegion" - %0:2 = "tf.WhileRegion"(%arg0, %unshaped) ({ + // CHECK: <{is_stateless = true}> + %0:2 = "tf.WhileRegion"(%arg0, %unshaped) <{is_stateless = true}> ({ // CHECK: {{.*}}({{.+}}: tensor, {{.+}}: tensor<1x2x3xf32>): ^bb0(%carg0: tensor, %carg1: tensor<*xf32>): %limit = arith.constant dense<5> : tensor @@ -258,8 +261,8 @@ module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, pr %neg = "tf.Neg"(%barg1) : (tensor<*xf32>) -> tensor<*xf32> // CHECK: "tf.Yield"{{.+}}, {{.+}}) : (tensor, tensor<1x2x3xf32>) -> () "tf.Yield"(%sub, %neg) : (tensor, tensor<*xf32>) -> () - // CHECK: {is_stateless = true} : (tensor, tensor<1x2x3xf32>) -> (tensor, tensor<1x2x3xf32>) - }) {is_stateless = true} : (tensor, tensor<*xf32>) -> (tensor, tensor<*xf32>) + // CHECK: (tensor, tensor<1x2x3xf32>) -> (tensor, tensor<1x2x3xf32>) + }) : (tensor, tensor<*xf32>) -> (tensor, tensor<*xf32>) // CHECK: return {{.+}}#1 : tensor<1x2x3xf32> func.return %0#1 : tensor<*xf32> } @@ -752,7 +755,7 @@ module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, pr // CHECK-LABEL: replace_tensor_list_element_shape func.func @replace_tensor_list_element_shape() { - // CHECK: %[[ELEMENT_SHAPE:.*]] = "tf.Const"() {value = dense<[-1, 1]> : tensor<2xi32>} + // CHECK: %[[ELEMENT_SHAPE:.*]] = "tf.Const"() <{value = dense<[-1, 1]> : tensor<2xi32>}> %elem_shape = "tf.Const"() {value = dense<[-1, 1]> : tensor<2xi32>} : () -> tensor<2xi32> %size = "tf.Const"() {value = dense<10> : tensor} : () -> tensor %elem = "tf._SomeOp"() : () -> tensor<16x1xf32> @@ -767,7 +770,7 @@ module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, pr "tf._SomeOtherOp"(%shape_32, %shape_64) : (tensor, tensor) -> () func.return } - + // CHECK-LABEL: refine_pop_back_results_from_operands func.func @refine_pop_back_results_from_operands(%arg0: tensor>>, %arg1: tensor<1xi32>) -> (tensor, tensor<*xi32>) { %0, %1 = "tf.TensorListPopBack"(%arg0, %arg1) : (tensor>>, tensor<1xi32>) -> (tensor, tensor<*xi32>) @@ -876,7 +879,7 @@ module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, pr tf_device.return %2 : tensor<1x8x2xf32> // CHECK: () -> tensor<1x8x2xf32> }) {device = "/device:CPU:0"} : () -> tensor<*xf32> - // CHECK: "tf.Cast"(%{{.*}}) {Truncate = false} : (tensor<1x8x2xf32>) -> tensor<*xf32> + // CHECK: "tf.Cast"(%{{.*}}) <{Truncate = false}> : (tensor<1x8x2xf32>) -> tensor<*xf32> // CHECK: (tensor, tensor<1x8x2xf32>) -> (tensor<1x8x1xf32>, tensor<1x8x1xf32>) %3:2 = "tf.Split"(%0, %1) {device = ""} : (tensor, tensor<*xf32>) -> (tensor<*xf32>, tensor<*xf32>) %4 = tensor.cast %1 : tensor<*xf32> to tensor @@ -891,7 +894,7 @@ module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, pr tf_device.return %2 : tensor<1x8x2xf32> // CHECK: () -> tensor<1x8x2xf32> }) : () -> tensor<*xf32> - // CHECK: "tf.Cast"(%{{.*}}) {Truncate = false} : (tensor<1x8x2xf32>) -> tensor<*xf32> + // CHECK: "tf.Cast"(%{{.*}}) <{Truncate = false}> : (tensor<1x8x2xf32>) -> tensor<*xf32> // CHECK: (tensor, tensor<1x8x2xf32>) -> (tensor<1x8x1xf32>, tensor<1x8x1xf32>) %3:2 = "tf.Split"(%0, %1) {device = ""} : (tensor, tensor<*xf32>) -> (tensor<*xf32>, tensor<*xf32>) %4 = tensor.cast %1 : tensor<*xf32> to tensor @@ -980,7 +983,8 @@ module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, pr %0:4 = "tf.While"(%arg0, %arg1, %arg2, %arg3) {cond = @while_shape_invariant_cond_func_propagate, body = @while_shape_invariant_body_func_propagate, is_stateless = false, shape_invariant} : (tensor<4xf32>, tensor>>, tensor>>, tensor<1xi32>) -> (tensor<*xf32>, tensor<*x!tf_type.resource>, tensor, tensor) // CHECK: "tf.WhileRegion" - %1:4 = "tf.WhileRegion"(%arg0, %arg1, %arg2, %arg3) ({ + // CHECK-SAME: shape_invariant + %1:4 = "tf.WhileRegion"(%arg0, %arg1, %arg2, %arg3) <{is_stateless = false, shape_invariant}> ({ // CHECK-NEXT: ^{{.+}}({{%.+}}: tensor<*xf32>, {{%.+}}: tensor<*x!tf_type.resource>>, {{%.+}}: tensor>>, {{%.+}}: tensor): ^cond(%carg0: tensor<*xf32>, %carg1: tensor<*x!tf_type.resource>, %carg2: tensor, %carg3: tensor): %2 = "tf.Const"() {value = dense : tensor} : () -> tensor @@ -992,10 +996,9 @@ module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, pr // CHECK: "tf.Yield" // CHECK-SAME: (tensor<*xf32>, tensor<*x!tf_type.resource>>, tensor>>, tensor) -> () "tf.Yield"(%barg0, %barg1, %barg2, %2) : (tensor<*xf32>, tensor<*x!tf_type.resource>, tensor, tensor) -> () - // CHECK-NEXT: shape_invariant - // CHECK-SAME: (tensor<4xf32>, tensor>>, tensor>>, tensor<1xi32>) + // CHECK-NEXT: (tensor<4xf32>, tensor>>, tensor>>, tensor<1xi32>) // CHECK-SAME: -> (tensor<*xf32>, tensor<*x!tf_type.resource>>, tensor>>, tensor) - }) {is_stateless = false, shape_invariant} : (tensor<4xf32>, tensor>>, tensor>>, tensor<1xi32>) -> (tensor<*xf32>, tensor<*x!tf_type.resource>, tensor, tensor) + }) : (tensor<4xf32>, tensor>>, tensor>>, tensor<1xi32>) -> (tensor<*xf32>, tensor<*x!tf_type.resource>, tensor, tensor) func.return %0#0, %0#1, %0#2, %0#3, %1#0, %1#1, %1#2, %1#3 : tensor<*xf32>, tensor<*x!tf_type.resource>, tensor, tensor, tensor<*xf32>, tensor<*x!tf_type.resource>, tensor, tensor } @@ -1028,7 +1031,8 @@ module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, pr %0 = "tf.While"(%arg0) {cond = @while_shape_invariant_cond_func_different_dims, body = @while_shape_invariant_body_func_different_dims, is_stateless = false, shape_invariant} : (tensor<1x2x3xf32>) -> tensor<1x8x3xf32> // CHECK: "tf.WhileRegion" - %1 = "tf.WhileRegion"(%arg0) ({ + // CHECK-SAME: shape_invariant + %1 = "tf.WhileRegion"(%arg0) <{is_stateless = false, shape_invariant}> ({ // CHECK-NEXT: ^{{.+}}({{%.+}}: tensor<1x?x3xf32>): ^cond(%carg0: tensor<*xf32>): %2 = "tf.Const"() {value = dense : tensor} : () -> tensor @@ -1040,10 +1044,9 @@ module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, pr // CHECK: "tf.Yield" // CHECK-SAME: (tensor<1x?x3xf32>) -> () "tf.Yield"(%2) : (tensor<*xf32>) -> () - // CHECK-NEXT: shape_invariant - // CHECK-SAME: (tensor<1x2x3xf32>) + // CHECK-NEXT: (tensor<1x2x3xf32>) // CHECK-SAME: -> tensor<1x8x3xf32> - }) {is_stateless = false, shape_invariant} : (tensor<1x2x3xf32>) -> tensor<1x8x3xf32> + }) : (tensor<1x2x3xf32>) -> tensor<1x8x3xf32> func.return %0, %1 : tensor<1x8x3xf32>, tensor<1x8x3xf32> } @@ -1076,7 +1079,8 @@ module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, pr %0 = "tf.While"(%arg0) {cond = @while_shape_invariant_cond_func_body_result_propagate, body = @while_shape_invariant_body_func_body_result_propagate, is_stateless = false, shape_invariant} : (tensor<*x!tf_type.resource>>) -> tensor<*x!tf_type.resource> // CHECK: "tf.WhileRegion" - %1 = "tf.WhileRegion"(%arg0) ({ + // CHECK-SAME: shape_invariant + %1 = "tf.WhileRegion"(%arg0) <{is_stateless = false, shape_invariant}> ({ // CHECK-NEXT: ^{{.+}}({{%.+}}: tensor<*x!tf_type.resource>>): ^cond(%carg0: tensor<*x!tf_type.resource>>): %2 = "tf.Const"() {value = dense : tensor} : () -> tensor @@ -1088,10 +1092,9 @@ module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, pr // CHECK: "tf.Yield" // CHECK-SAME: (tensor<*x!tf_type.resource>>) -> () "tf.Yield"(%2) : (tensor<*x!tf_type.resource>>) -> () - // CHECK-NEXT: shape_invariant - // CHECK-SAME: (tensor<*x!tf_type.resource>>) + // CHECK-NEXT: (tensor<*x!tf_type.resource>>) // CHECK-SAME: -> tensor<*x!tf_type.resource>> - }) {is_stateless = false, shape_invariant} : (tensor<*x!tf_type.resource>>) -> tensor<*x!tf_type.resource> + }) : (tensor<*x!tf_type.resource>>) -> tensor<*x!tf_type.resource> func.return %0, %1 : tensor<*x!tf_type.resource>, tensor<*x!tf_type.resource> } @@ -1333,7 +1336,7 @@ module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, pr %cst = "tf.Const"() {value = dense<0> : tensor<4x2xi32>} : () -> tensor<4x2xi32> %cst_0 = "tf.Const"() {value = dense<[2, 2, 1, 1]> : tensor<4xi32>} : () -> tensor<4xi32> %cst_1 = "tf.Const"() {value = dense<[2, 3, 1, 1]> : tensor<4xi32>} : () -> tensor<4xi32> - // CHECK: %0 = "tf.XlaSelectAndScatter"(%arg0, %cst_1, %cst_0, %cst, %arg1, %arg2) {scatter = @add_scatter, select = @ge_select} : (tensor<4x5x1x1xbf16>, tensor<4xi32>, tensor<4xi32>, tensor<4x2xi32>, tensor<2x2x1x1xbf16>, tensor) -> tensor<4x5x1x1xbf16> + // CHECK: %0 = "tf.XlaSelectAndScatter"(%arg0, %cst_1, %cst_0, %cst, %arg1, %arg2) <{scatter = @add_scatter, select = @ge_select}> : (tensor<4x5x1x1xbf16>, tensor<4xi32>, tensor<4xi32>, tensor<4x2xi32>, tensor<2x2x1x1xbf16>, tensor) -> tensor<4x5x1x1xbf16> %0 = "tf.XlaSelectAndScatter"(%arg0, %cst_1, %cst_0, %cst, %arg1, %arg2) {scatter = @add_scatter, select = @ge_select} : (tensor<4x5x1x1xbf16>, tensor<4xi32>, tensor<4xi32>, tensor<4x2xi32>, tensor<2x2x1x1xbf16>, tensor) -> tensor func.return %0 : tensor } @@ -1373,7 +1376,7 @@ module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, pr %cst_1 = "tf.Const"() {value = dense<2> : tensor<1xi32>} : () -> tensor<1xi32> %cst_2 = "tf.Const"() {value = dense<3> : tensor<1xi32>} : () -> tensor<1xi32> %cst_3 = "tf.Const"() {value = dense<4> : tensor<1xi32>} : () -> tensor<1xi32> - // CHECK: 0 = "tf.XlaReduceWindow"(%arg0, %arg1, %cst_0, %cst_1, %cst_2, %cst_3, %cst) {computation = @sum_reducer3} : (tensor<7xf32>, tensor, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1x2xi32>) -> tensor<10xf32> + // CHECK: 0 = "tf.XlaReduceWindow"(%arg0, %arg1, %cst_0, %cst_1, %cst_2, %cst_3, %cst) <{computation = @sum_reducer3}> : (tensor<7xf32>, tensor, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1x2xi32>) -> tensor<10xf32> %0 = "tf.XlaReduceWindow"(%arg0, %arg1, %cst_0, %cst_1, %cst_2, %cst_3, %cst) {computation = @sum_reducer3} : (tensor<7xf32>, tensor, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1x2xi32>) -> tensor func.return %0 : tensor } @@ -1813,7 +1816,7 @@ module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, pr func.func @infer_var_handle_op_from_assigns() -> tensor<1xi8> { %cst = arith.constant dense<1> : tensor<1xi8> %0 = "tf.VarHandleOp"() {container = "", shared_name = "bar"} : () -> tensor>> - // CHECK: "tf.VarHandleOp"() {container = "", shared_name = "bar"} : () -> tensor>> + // CHECK: "tf.VarHandleOp"() <{container = "", shared_name = "bar"}> : () -> tensor>> "tf.AssignVariableOp"(%0, %cst) : (tensor>>, tensor<1xi8>) -> () func.return %cst : tensor<1xi8> } @@ -1822,7 +1825,7 @@ module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, pr func.func @infer_var_handle_op_from_read() -> tensor<1xi8> { %cst = arith.constant dense<1> : tensor<1xi8> %0 = "tf.VarHandleOp"() {container = "", shared_name = "bar"} : () -> tensor>> - // CHECK: "tf.VarHandleOp"() {container = "", shared_name = "bar"} : () -> tensor>> + // CHECK: "tf.VarHandleOp"() <{container = "", shared_name = "bar"}> : () -> tensor>> %read = "tf.ReadVariableOp"(%0) : (tensor>>) -> tensor<1xi8> func.return %read : tensor<1xi8> } @@ -1831,7 +1834,7 @@ module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, pr func.func @do_not_infer_var_handle_op_when_custom_op_uses_it() -> tensor<1xi8> { %cst = arith.constant dense<1> : tensor<1xi8> %0 = "tf.VarHandleOp"() {container = "", shared_name = "bar"} : () -> tensor>> - // CHECK: "tf.VarHandleOp"() {container = "", shared_name = "bar"} : () -> tensor>> + // CHECK: "tf.VarHandleOp"() <{container = "", shared_name = "bar"}> : () -> tensor>> %read = "tf.ReadVariableOp"(%0) : (tensor>>) -> tensor<1xi8> %1 = "tf.MyCustomOp"(%0) : (tensor>>) -> tensor<4xi8> func.return %read : tensor<1xi8> @@ -1916,7 +1919,7 @@ module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, pr %lhs_dilation = "tf.Const"() {value = dense<[4, 1, 1]> : tensor<3xi32>} : () -> tensor<3xi32> %padding = "tf.Const"() {value = dense<0> : tensor<3x2xi32>} : () -> tensor<3x2xi32> %strides = "tf.Const"() {value = dense<[3, 1, 1]> : tensor<3xi32>} : () -> tensor<3xi32> - // CHECK: %0 = "tf.XlaConvV2"(%arg0, %arg1, %cst_3, %cst_2, %cst_1, %cst_0, %cst) {dimension_numbers = "\18\03 \042\03\00\01\02@\04P\04Z\03\01\02\03b\03\01\02\03", precision_config = ""} : (tensor<8x?x?x?x16xf32>, tensor<4x3x3x16x16xf32>, tensor<3xi32>, tensor<3x2xi32>, tensor<3xi32>, tensor<3xi32>, tensor) -> tensor<8x4x14x14x16xf32> + // CHECK: %0 = "tf.XlaConvV2"(%arg0, %arg1, %cst_3, %cst_2, %cst_1, %cst_0, %cst) <{dimension_numbers = "\18\03 \042\03\00\01\02@\04P\04Z\03\01\02\03b\03\01\02\03", precision_config = ""}> : (tensor<8x?x?x?x16xf32>, tensor<4x3x3x16x16xf32>, tensor<3xi32>, tensor<3x2xi32>, tensor<3xi32>, tensor<3xi32>, tensor) -> tensor<8x4x14x14x16xf32> %0 = "tf.XlaConvV2"(%lhs, %rhs, %strides, %padding, %lhs_dilation, %rhs_dilation, %feature_group_count) {dimension_numbers = "\18\03 \042\03\00\01\02@\04P\04Z\03\01\02\03b\03\01\02\03", precision_config = ""} : (tensor<8x?x?x?x16xf32>, tensor<4x3x3x16x16xf32>, tensor<3xi32>, tensor<3x2xi32>, tensor<3xi32>, tensor<3xi32>, tensor) -> tensor<8x4x14x14x16xf32> func.return %0 : tensor<8x4x14x14x16xf32> } @@ -1928,7 +1931,7 @@ module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, pr %lhs_dilation = "tf.Const"() {value = dense<[4, 1, 1]> : tensor<3xi32>} : () -> tensor<3xi32> %padding = "tf.Const"() {value = dense<0> : tensor<3x2xi32>} : () -> tensor<3x2xi32> %strides = "tf.Const"() {value = dense<[3, 1, 1]> : tensor<3xi32>} : () -> tensor<3xi32> - // CHECK: %0 = "tf.XlaConvV2"(%arg0, %arg1, %cst_3, %cst_2, %cst_1, %cst_0, %cst) {dimension_numbers = "\18\03 \042\03\00\01\02@\04P\04Z\03\01\02\03b\03\01\02\03", precision_config = ""} : (tensor<8x4x16x16x16xf32>, tensor, tensor<3xi32>, tensor<3x2xi32>, tensor<3xi32>, tensor<3xi32>, tensor) -> tensor<8x4x14x14x16xf32> + // CHECK: %0 = "tf.XlaConvV2"(%arg0, %arg1, %cst_3, %cst_2, %cst_1, %cst_0, %cst) <{dimension_numbers = "\18\03 \042\03\00\01\02@\04P\04Z\03\01\02\03b\03\01\02\03", precision_config = ""}> : (tensor<8x4x16x16x16xf32>, tensor, tensor<3xi32>, tensor<3x2xi32>, tensor<3xi32>, tensor<3xi32>, tensor) -> tensor<8x4x14x14x16xf32> %0 = "tf.XlaConvV2"(%lhs, %rhs, %strides, %padding, %lhs_dilation, %rhs_dilation, %feature_group_count) {dimension_numbers = "\18\03 \042\03\00\01\02@\04P\04Z\03\01\02\03b\03\01\02\03", precision_config = ""} : (tensor<8x4x16x16x16xf32>, tensor, tensor<3xi32>, tensor<3x2xi32>, tensor<3xi32>, tensor<3xi32>, tensor) -> tensor<8x4x14x14x16xf32> func.return %0 : tensor<8x4x14x14x16xf32> } @@ -1941,7 +1944,7 @@ module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, pr %cst_2 = "tf.Const"() {value = dense<2> : tensor<1xi32>} : () -> tensor<1xi32> %cst_3 = "tf.Const"() {value = dense<1> : tensor} : () -> tensor %0 = tf_executor.graph { - // CHECK: "tf.XlaConvV2"(%arg0, %arg1, %cst, %cst_0, %cst_1, %cst_2, %cst_3) {_XlaHasReferenceVars = false, device = "/job:localhost/replica:0/task:0/device:XLA_CPU:0", dimension_numbers = "\18\012\01\02@\01P\01Z\01\02b\01\02", precision_config = "\0A\02\01\01"} : (tensor<*xf32>, tensor<*xf32>, tensor<1xi32>, tensor<1x2xi32>, tensor<1xi32>, tensor<1xi32>, tensor) -> tensor<*xf32> + // CHECK: "tf.XlaConvV2"(%arg0, %arg1, %cst, %cst_0, %cst_1, %cst_2, %cst_3) <{dimension_numbers = "\18\012\01\02@\01P\01Z\01\02b\01\02", precision_config = "\0A\02\01\01"}> {_XlaHasReferenceVars = false, device = "/job:localhost/replica:0/task:0/device:XLA_CPU:0"} : (tensor<*xf32>, tensor<*xf32>, tensor<1xi32>, tensor<1x2xi32>, tensor<1xi32>, tensor<1xi32>, tensor) -> tensor<*xf32> %outputs, %control = tf_executor.island wraps "tf.XlaConvV2"(%arg0, %arg1, %cst, %cst_0, %cst_1, %cst_2, %cst_3) {_XlaHasReferenceVars = false, device = "/job:localhost/replica:0/task:0/device:XLA_CPU:0", dimension_numbers = "\18\012\01\02@\01P\01Z\01\02b\01\02", precision_config = "\0A\02\01\01"} : (tensor<*xf32>, tensor<*xf32>, tensor<1xi32>, tensor<1x2xi32>, tensor<1xi32>, tensor<1xi32>, tensor) -> tensor<*xf32> tf_executor.fetch %outputs : tensor<*xf32> } @@ -1955,7 +1958,7 @@ module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, pr %lhs_dilation = "tf.Const"() {value = dense<[4, 1, 1]> : tensor<3xi32>} : () -> tensor<3xi32> %padding = "tf.Const"() {value = dense<0> : tensor<3x2xi32>} : () -> tensor<3x2xi32> %strides = "tf.Const"() {value = dense<[3, 1, 1]> : tensor<3xi32>} : () -> tensor<3xi32> - // CHECK: %0 = "tf.XlaConvV2"(%arg0, %arg1, %cst_3, %cst_2, %cst_1, %cst_0, %cst) {dimension_numbers = "\18\03 \042\03\00\01\02@\04P\04Z\03\01\02\03b\03\01\02\03", precision_config = ""} : (tensor<8x4x16x16x16xf32>, tensor<4x3x3x16x16xf32>, tensor<3xi32>, tensor<3x2xi32>, tensor<3xi32>, tensor<3xi32>, tensor) -> tensor<8x4x14x14x16xf32> + // CHECK: %0 = "tf.XlaConvV2"(%arg0, %arg1, %cst_3, %cst_2, %cst_1, %cst_0, %cst) <{dimension_numbers = "\18\03 \042\03\00\01\02@\04P\04Z\03\01\02\03b\03\01\02\03", precision_config = ""}> : (tensor<8x4x16x16x16xf32>, tensor<4x3x3x16x16xf32>, tensor<3xi32>, tensor<3x2xi32>, tensor<3xi32>, tensor<3xi32>, tensor) -> tensor<8x4x14x14x16xf32> %0 = "tf.XlaConvV2"(%lhs, %rhs, %strides, %padding, %lhs_dilation, %rhs_dilation, %feature_group_count) {dimension_numbers = "\18\03 \042\03\00\01\02@\04P\04Z\03\01\02\03b\03\01\02\03", precision_config = ""} : (tensor<8x4x16x16x16xf32>, tensor<4x3x3x16x16xf32>, tensor<3xi32>, tensor<3x2xi32>, tensor<3xi32>, tensor<3xi32>, tensor) -> tensor<8x4x14x14x16xf32> func.return %0 : tensor<8x4x14x14x16xf32> } @@ -1967,7 +1970,7 @@ module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, pr %lhs_dilation = "tf.Const"() {value = dense<[4, 1, 1]> : tensor<3xi32>} : () -> tensor<3xi32> %padding = "tf.Const"() {value = dense<0> : tensor<3x2xi32>} : () -> tensor<3x2xi32> %strides = "tf.Const"() {value = dense<[3, 1, 1]> : tensor<3xi32>} : () -> tensor<3xi32> - // CHECK: %0 = "tf.XlaConvV2"(%arg0, %arg1, %cst_3, %cst_2, %cst_1, %cst_0, %cst) {dimension_numbers = "\18\03 \042\03\00\01\02@\04P\04Z\03\01\02\03b\03\01\02\03", precision_config = ""} : (tensor<8x4x16x16x16xf32>, tensor<4x3x3x16x16xf32>, tensor<3xi32>, tensor<3x2xi32>, tensor<3xi32>, tensor<3xi32>, tensor) -> tensor<8x4x14x14x16xf32> + // CHECK: %0 = "tf.XlaConvV2"(%arg0, %arg1, %cst_3, %cst_2, %cst_1, %cst_0, %cst) <{dimension_numbers = "\18\03 \042\03\00\01\02@\04P\04Z\03\01\02\03b\03\01\02\03", precision_config = ""}> : (tensor<8x4x16x16x16xf32>, tensor<4x3x3x16x16xf32>, tensor<3xi32>, tensor<3x2xi32>, tensor<3xi32>, tensor<3xi32>, tensor) -> tensor<8x4x14x14x16xf32> %0 = "tf.XlaConvV2"(%lhs, %rhs, %strides, %padding, %lhs_dilation, %rhs_dilation, %feature_group_count) {dimension_numbers = "\18\03 \042\03\00\01\02@\04P\04Z\03\01\02\03b\03\01\02\03", precision_config = ""} : (tensor<8x4x16x16x16xf32>, tensor<4x3x3x16x16xf32>, tensor<3xi32>, tensor<3x2xi32>, tensor<3xi32>, tensor<3xi32>, tensor) -> tensor func.return %0 : tensor } @@ -1979,7 +1982,7 @@ module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, pr %lhs_dilation = "tf.Const"() {value = dense<[4, 1, 1]> : tensor<3xi32>} : () -> tensor<3xi32> %padding = "tf.Const"() {value = dense<0> : tensor<3x2xi32>} : () -> tensor<3x2xi32> %strides = "tf.Const"() {value = dense<[3, 1, 1]> : tensor<3xi32>} : () -> tensor<3xi32> - // CHECK: %0 = "tf.XlaConvV2"(%arg0, %arg1, %cst_3, %cst_2, %cst_1, %cst_0, %cst) {dimension_numbers = "\18\03 \042\03\00\01\02@\04P\04Z\03\01\02\03b\03\01\02\03", precision_config = ""} : (tensor<8x4x16x16x16xf16>, tensor<4x3x3x16x16xf16>, tensor<3xi32>, tensor<3x2xi32>, tensor<3xi32>, tensor<3xi32>, tensor) -> tensor<8x4x14x14x16xf32> + // CHECK: %0 = "tf.XlaConvV2"(%arg0, %arg1, %cst_3, %cst_2, %cst_1, %cst_0, %cst) <{dimension_numbers = "\18\03 \042\03\00\01\02@\04P\04Z\03\01\02\03b\03\01\02\03", precision_config = ""}> : (tensor<8x4x16x16x16xf16>, tensor<4x3x3x16x16xf16>, tensor<3xi32>, tensor<3x2xi32>, tensor<3xi32>, tensor<3xi32>, tensor) -> tensor<8x4x14x14x16xf32> %0 = "tf.XlaConvV2"(%lhs, %rhs, %strides, %padding, %lhs_dilation, %rhs_dilation, %feature_group_count) {dimension_numbers = "\18\03 \042\03\00\01\02@\04P\04Z\03\01\02\03b\03\01\02\03", precision_config = ""} : (tensor<8x4x16x16x16xf16>, tensor<4x3x3x16x16xf16>, tensor<3xi32>, tensor<3x2xi32>, tensor<3xi32>, tensor<3xi32>, tensor) -> tensor func.return %0 : tensor } @@ -1991,7 +1994,7 @@ module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, pr %lhs_dilation = "tf.Const"() {value = dense<[4, 1, 1]> : tensor<3xi32>} : () -> tensor<3xi32> %padding = "tf.Const"() {value = dense<0> : tensor<3x2xi32>} : () -> tensor<3x2xi32> %strides = "tf.Const"() {value = dense<[3, 1, 1]> : tensor<3xi64>} : () -> tensor<3xi64> - // CHECK: %0 = "tf.XlaConvV2"(%arg0, %arg1, %cst_3, %cst_2, %cst_1, %cst_0, %cst) {dimension_numbers = "\18\03 \042\03\00\01\02@\04P\04Z\03\01\02\03b\03\01\02\03", precision_config = ""} : (tensor<8x4x16x16x16xf32>, tensor<4x3x3x16x16xf32>, tensor<3xi64>, tensor<3x2xi32>, tensor<3xi32>, tensor<3xi32>, tensor) -> tensor<8x4x14x14x16xf32> + // CHECK: %0 = "tf.XlaConvV2"(%arg0, %arg1, %cst_3, %cst_2, %cst_1, %cst_0, %cst) <{dimension_numbers = "\18\03 \042\03\00\01\02@\04P\04Z\03\01\02\03b\03\01\02\03", precision_config = ""}> : (tensor<8x4x16x16x16xf32>, tensor<4x3x3x16x16xf32>, tensor<3xi64>, tensor<3x2xi32>, tensor<3xi32>, tensor<3xi32>, tensor) -> tensor<8x4x14x14x16xf32> %0 = "tf.XlaConvV2"(%lhs, %rhs, %strides, %padding, %lhs_dilation, %rhs_dilation, %feature_group_count) {dimension_numbers = "\18\03 \042\03\00\01\02@\04P\04Z\03\01\02\03b\03\01\02\03", precision_config = ""} : (tensor<8x4x16x16x16xf32>, tensor<4x3x3x16x16xf32>, tensor<3xi64>, tensor<3x2xi32>, tensor<3xi32>, tensor<3xi32>, tensor) -> tensor func.return %0 : tensor } diff --git a/tensorflow/compiler/mlir/tensorflow/tests/split_into_island_per_op.mlir b/tensorflow/compiler/mlir/tensorflow/tests/split_into_island_per_op.mlir index 7307bff6e69914..4428811b7df371 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/split_into_island_per_op.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/split_into_island_per_op.mlir @@ -116,7 +116,7 @@ func.func @dangling_print(%arg0: tensor<*xi32>, %arg1: tensor) -> (tensor<* // CHECK: %[[GRAPH:.*]]:2 = tf_executor.graph { // CHECK: %[[ADD1:.*]], %[[ADD1_control:.*]] = tf_executor.island wraps "tf.Add"(%arg0, %arg1) // CHECK: %[[ADD2:.*]], %[[ADD2_control:.*]] = tf_executor.island wraps "tf.Add"(%[[ADD1]], %arg1) -// CHECK: %[[PRINT:.*]], %[[PRINT_control:.*]] = tf_executor.island wraps "tf.Print"(%[[ADD2]]) {message = "add result"} +// CHECK: %[[PRINT:.*]], %[[PRINT_control:.*]] = tf_executor.island wraps "tf.Print"(%[[ADD2]]) <{message = "add result"}> // CHECK: tf_executor.fetch %[[ADD1]], %[[ADD2]] : // CHECK: } // CHECK: return %[[GRAPH]]#0, %[[GRAPH]]#1 @@ -186,7 +186,7 @@ func.func @non_aliasing_reads_writes( // CHECK: %[[READ0:.*]], %[[READ0_CONTROL:.*]] = tf_executor.island wraps "tf.ReadVariableOp"(%arg0) // CHECK: %[[ASSIGN0_CONTROL:.*]] = tf_executor.island wraps "tf.AssignVariableOp"(%arg0, %arg2) // CHECK: %[[READ1:.*]], %[[READ1_CONTROL:.*]] = tf_executor.island wraps "tf.ReadVariableOp"(%arg1) -// CHECK: %[[VH0:.*]], %[[VH0_CONTROL:.*]] = tf_executor.island wraps "tf.VarHandleOp"() {container = "c", shared_name = "v0"} +// CHECK: %[[VH0:.*]], %[[VH0_CONTROL:.*]] = tf_executor.island wraps "tf.VarHandleOp"() <{container = "c", shared_name = "v0"}> // CHECK: %[[READ2:.*]], %[[READ2_CONTROL:.*]] = tf_executor.island wraps "tf.ReadVariableOp"(%[[VH0]]) // CHECK: %[[ASSIGN1_CONTROL:.*]] = tf_executor.island wraps "tf.AssignVariableOp"(%arg1, %[[READ0]]) // CHECK: %[[ASSIGN2_CONTROL:.*]] = tf_executor.island wraps "tf.AssignVariableOp"(%arg0, %[[READ2]]) @@ -214,8 +214,8 @@ func.func @unknown_side_effecting_op(%arg0: tensor<32xf32>) -> () { // CHECK-LABEL: func @unknown_side_effecting_op // CHECK: tf_executor.graph { -// CHECK: %[[VH0:.*]], %[[VH0_CONTROL:.*]] = tf_executor.island wraps "tf.VarHandleOp"() {container = "c", shared_name = "v0"} -// CHECK: %[[VH1:.*]], %[[VH1_CONTROL:.*]] = tf_executor.island wraps "tf.VarHandleOp"() {container = "c", shared_name = "v1"} +// CHECK: %[[VH0:.*]], %[[VH0_CONTROL:.*]] = tf_executor.island wraps "tf.VarHandleOp"() <{container = "c", shared_name = "v0"}> +// CHECK: %[[VH1:.*]], %[[VH1_CONTROL:.*]] = tf_executor.island wraps "tf.VarHandleOp"() <{container = "c", shared_name = "v1"}> // CHECK: %[[READ0:.*]], %[[READ0_CONTROL:.*]] = tf_executor.island wraps "tf.ReadVariableOp"(%[[VH0]]) // CHECK: %[[ASSIGN0_CONTROL:.*]] = tf_executor.island wraps "tf.AssignVariableOp"(%[[VH1]], %arg0) // CHECK: %[[UNKNOWN_CONTROL:.*]] = tf_executor.island wraps "tf._UnknownSideEffectingOp_"() @@ -432,4 +432,4 @@ func.func @else_function() { tf_executor.fetch } func.return -} \ No newline at end of file +} diff --git a/tensorflow/compiler/mlir/tensorflow/tests/stack_ops_decomposition.mlir b/tensorflow/compiler/mlir/tensorflow/tests/stack_ops_decomposition.mlir index ee663e98ae9674..907d512b2571b6 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/stack_ops_decomposition.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/stack_ops_decomposition.mlir @@ -4,15 +4,15 @@ // CHECK-LABEL: func @main func.func @main() -> tensor { - // CHECK-NEXT: "tf.Const"() {value = dense<10> : tensor} + // CHECK-NEXT: "tf.Const"() <{value = dense<10> : tensor}> %max_size = "tf.Const"() {value = dense<10> : tensor} : () -> tensor - // CHECK-NEXT: %[[ZERO_SCALAR:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor + // CHECK-NEXT: %[[ZERO_SCALAR:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK-NEXT: %[[CAST_ZERO:.*]] = "tf.Cast"(%[[ZERO_SCALAR]]) : (tensor) -> tensor - // CHECK-NEXT: %[[CONST10:.*]] = "tf.Const"() {value = dense<10> : tensor<1xi32>} : () -> tensor<1xi32> + // CHECK-NEXT: %[[CONST10:.*]] = "tf.Const"() <{value = dense<10> : tensor<1xi32>}> : () -> tensor<1xi32> // CHECK-NEXT: %[[BROADCAST:.*]] = "tf.BroadcastTo"(%[[CAST_ZERO]], %[[CONST10]]) : (tensor, tensor<1xi32>) -> tensor<10xf32> // CHECK-NEXT: %[[BUFFER:.*]] = "tf.MlirLocalVarOp"() : () -> tensor>> // CHECK-NEXT: %[[SIZE:.*]] = "tf.MlirLocalVarOp"() : () -> tensor>> - // CHECK-NEXT: %[[ZERO:.*]] = "tf.Const"() {value = dense<0> : tensor<1xi32>} : () -> tensor<1xi32> + // CHECK-NEXT: %[[ZERO:.*]] = "tf.Const"() <{value = dense<0> : tensor<1xi32>}> : () -> tensor<1xi32> // CHECK-NEXT: "tf.AssignVariableOp"(%[[SIZE]], %[[ZERO]]) // CHECK-NEXT: "tf.AssignVariableOp"(%[[BUFFER]], %[[BROADCAST]]) %stack = "tf.StackV2"(%max_size) {elem_type = f32, stack_name = "s"} : (tensor) -> tensor @@ -21,22 +21,22 @@ func.func @main() -> tensor { %elem = "tf._SomeOp"() : () -> tensor // CHECK-NEXT: %[[READ_VAL:.*]] = "tf.ReadVariableOp"(%[[BUFFER]]) // CHECK-NEXT: %[[READ_SIZE:.*]] = "tf.ReadVariableOp"(%[[SIZE]]) - // CHECK-NEXT: %[[UPDATE_SHAPE:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32> + // CHECK-NEXT: %[[UPDATE_SHAPE:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}> : () -> tensor<1xi32> // CHECK-NEXT: %[[UPDATE_SLICE:.*]] = "tf.Reshape"(%[[PUSHVAL]], %[[UPDATE_SHAPE]]) : (tensor, tensor<1xi32>) -> tensor<1xf32> // CHECK-NEXT: %[[UPDATE:.*]] = "tf.XlaDynamicUpdateSlice"(%[[READ_VAL]], %[[UPDATE_SLICE]], %[[READ_SIZE]]) : (tensor<10xf32>, tensor<1xf32>, tensor<1xi32>) -> tensor<10xf32> // CHECK-NEXT: "tf.AssignVariableOp"(%[[BUFFER]], %[[UPDATE]]) : (tensor>>, tensor<10xf32>) -> () - // CHECK-NEXT: %[[CONST1:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32> + // CHECK-NEXT: %[[CONST1:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}> : () -> tensor<1xi32> // CHECK-NEXT: %[[NEW_SIZE:.*]] = "tf.AddV2"(%[[READ_SIZE]], %[[CONST1]]) : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32> // CHECK-NEXT: "tf.AssignVariableOp"(%[[SIZE]], %[[NEW_SIZE]]) : (tensor>>, tensor<1xi32>) -> () %push = "tf.StackPushV2"(%id, %elem) {swap_memory = false} : (tensor, tensor) -> tensor %pop = "tf.StackPopV2"(%stack) : (tensor) -> tensor // CHECK-NEXT: %[[READ_VAL1:.*]] = "tf.ReadVariableOp"(%[[BUFFER]]) // CHECK-NEXT: %[[READ_SIZE1:.*]] = "tf.ReadVariableOp"(%[[SIZE]]) - // CHECK-NEXT: %[[CONST1_1:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32> + // CHECK-NEXT: %[[CONST1_1:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}> : () -> tensor<1xi32> // CHECK-NEXT: %[[SUB:.*]] = "tf.Sub"(%[[READ_SIZE1]], %[[CONST1_1]]) - // CHECK-NEXT: %[[SLICE_SIZE:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32> + // CHECK-NEXT: %[[SLICE_SIZE:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}> : () -> tensor<1xi32> // CHECK-NEXT: %[[SLICE:.*]] = "tf.Slice"(%[[READ_VAL1]], %[[SUB]], %[[SLICE_SIZE]]) : (tensor<10xf32>, tensor<1xi32>, tensor<1xi32>) -> tensor<1xf32> - // CHECK-NEXT: %[[ELEM_SHAPE:.*]] = "tf.Const"() {value = dense<> : tensor<0xi32>} : () -> tensor<0xi32> + // CHECK-NEXT: %[[ELEM_SHAPE:.*]] = "tf.Const"() <{value = dense<> : tensor<0xi32>}> : () -> tensor<0xi32> // CHECK-NEXT: %[[ELEM:.*]] = "tf.Reshape"(%[[SLICE]], %[[ELEM_SHAPE]]) : (tensor<1xf32>, tensor<0xi32>) -> tensor // CHECK-NEXT: "tf.AssignVariableOp"(%[[SIZE]], %[[SUB]]) : (tensor>>, tensor<1xi32>) -> () "tf.StackCloseV2"(%stack) : (tensor) -> () @@ -50,14 +50,14 @@ func.func @main() -> tensor { // CHECK-LABEL: func @main func.func @main() -> tensor<2xi32> { - // CHECK-NEXT: "tf.Const"() {value = dense<10> : tensor} : () -> tensor + // CHECK-NEXT: "tf.Const"() <{value = dense<10> : tensor}> : () -> tensor %size = "tf.Const"() {value = dense<10> : tensor} : () -> tensor - // CHECK-NEXT: %[[ZERO_CONST:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor - // CHECK-NEXT: %[[STACK_SHAPE:.*]] = "tf.Const"() {value = dense<[10, 2]> : tensor<2xi32>} : () -> tensor<2xi32> + // CHECK-NEXT: %[[ZERO_CONST:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor + // CHECK-NEXT: %[[STACK_SHAPE:.*]] = "tf.Const"() <{value = dense<[10, 2]> : tensor<2xi32>}> : () -> tensor<2xi32> // CHECK-NEXT: %[[BROADCAST:.*]] = "tf.BroadcastTo"(%[[ZERO_CONST]], %[[STACK_SHAPE]]) : (tensor, tensor<2xi32>) -> tensor<10x2xi32> // CHECK-NEXT: %[[BUFFER:.*]] = "tf.MlirLocalVarOp"() : () -> tensor>> // CHECK-NEXT: %[[SIZE:.*]] = "tf.MlirLocalVarOp"() : () -> tensor>> - // CHECK-NEXT: %[[ZERO_SIZE:.*]] = "tf.Const"() {value = dense<0> : tensor<1xi32>} : () -> tensor<1xi32> + // CHECK-NEXT: %[[ZERO_SIZE:.*]] = "tf.Const"() <{value = dense<0> : tensor<1xi32>}> : () -> tensor<1xi32> // CHECK-NEXT: "tf.AssignVariableOp"(%[[SIZE]], %[[ZERO_SIZE]]) : (tensor>>, tensor<1xi32>) -> () // CHECK-NEXT: "tf.AssignVariableOp"(%[[BUFFER]], %[[BROADCAST]]) : (tensor>>, tensor<10x2xi32>) -> () %stack = "tf.StackV2"(%size) {elem_type = i32, stack_name = "s"} : (tensor) -> tensor @@ -65,14 +65,14 @@ func.func @main() -> tensor<2xi32> { %elem = "tf._SomeOp"() : () -> tensor<2xi32> // CHECK-NEXT: %[[STACK_VAL:.*]] = "tf.ReadVariableOp"(%[[BUFFER]]) : (tensor>>) -> tensor<10x2xi32> // CHECK-NEXT: %[[STACK_SIZE:.*]] = "tf.ReadVariableOp"(%[[SIZE]]) : (tensor>>) -> tensor<1xi32> - // CHECK-NEXT: %[[UPDATE_SHAPE:.*]] = "tf.Const"() {value = dense<[1, 2]> : tensor<2xi32>} : () -> tensor<2xi32> + // CHECK-NEXT: %[[UPDATE_SHAPE:.*]] = "tf.Const"() <{value = dense<[1, 2]> : tensor<2xi32>}> : () -> tensor<2xi32> // CHECK-NEXT: %[[UPDATE_SLICE:.*]] = "tf.Reshape"(%[[PUSH_VAL]], %[[UPDATE_SHAPE]]) : (tensor<2xi32>, tensor<2xi32>) -> tensor<1x2xi32> - // CHECK-NEXT: %[[ZERO_INDS:.*]] = "tf.Const"() {value = dense<0> : tensor<1xi32>} : () -> tensor<1xi32> - // CHECK-NEXT: %[[CONCAT_DIM:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor + // CHECK-NEXT: %[[ZERO_INDS:.*]] = "tf.Const"() <{value = dense<0> : tensor<1xi32>}> : () -> tensor<1xi32> + // CHECK-NEXT: %[[CONCAT_DIM:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK-NEXT: %[[CONCAT_OFFETS:.*]] = "tf.ConcatV2"(%[[STACK_SIZE]], %[[ZERO_INDS]], %[[CONCAT_DIM]]) : (tensor<1xi32>, tensor<1xi32>, tensor) -> tensor<2xi32> // CHECK-NEXT: %[[UPDATE:.*]] = "tf.XlaDynamicUpdateSlice"(%[[STACK_VAL]], %[[UPDATE_SLICE]], %[[CONCAT_OFFETS]]) : (tensor<10x2xi32>, tensor<1x2xi32>, tensor<2xi32>) -> tensor<10x2xi32> // CHECK-NEXT: "tf.AssignVariableOp"(%[[BUFFER]], %[[UPDATE]]) : (tensor>>, tensor<10x2xi32>) -> () - // CHECK-NEXT: %[[CONST1:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32> + // CHECK-NEXT: %[[CONST1:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}> : () -> tensor<1xi32> // CHECK-NEXT: %[[NEW_SIZE:.*]] = "tf.AddV2"(%[[STACK_SIZE]], %[[CONST1]]) : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32> // CHECK-NEXT: "tf.AssignVariableOp"(%[[SIZE]], %[[NEW_SIZE]]) : (tensor>>, tensor<1xi32>) -> () %push = "tf.StackPushV2"(%stack, %elem) {swap_memory = false} : (tensor, tensor<2xi32>) -> tensor<2xi32> @@ -102,7 +102,7 @@ func.func @main() -> () { } // CHECK: func @while_body(%[[BARG0:.*]]: tensor>>, %[[BARG1:.*]]: tensor, %[[BARG2:.*]]: tensor>>) func.func @while_body(%arg0: tensor, %arg1: tensor) -> (tensor, tensor) { - // CHECK: %[[CONST1:.*]] = "tf.Const"() {value = dense<1> : tensor} : () -> tensor + // CHECK: %[[CONST1:.*]] = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor %const1 = "tf.Const"() {value = dense<1> : tensor} : () -> tensor // CHECK: %[[SUB:.*]] = "tf.Sub"(%[[BARG1]], %[[CONST1]]) %sub = "tf.Sub"(%arg1, %const1) : (tensor, tensor) -> tensor @@ -143,7 +143,7 @@ func.func @main() -> () { }, { // CHECK: ^bb0(%[[BARG0:.*]]: tensor ^bb0(%barg0: tensor): - // CHECK: %[[CONST1:.*]] = "tf.Const"() {value = dense<1> : tensor} : () -> tensor + // CHECK: %[[CONST1:.*]] = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor %const1 = "tf.Const"() {value = dense<1> : tensor} : () -> tensor // CHECK: %[[SUB:.*]] = "tf.Sub"(%[[BARG0]], %[[CONST1]]) %sub = "tf.Sub"(%barg0, %const1) : (tensor, tensor) -> tensor @@ -185,7 +185,7 @@ func.func @main(%arg0: tensor) -> () { // CHECK: tf.AssignVariableOp // CHECK: tf.AssignVariableOp %stack = "tf.StackV2"(%max_size) {elem_type = f32, stack_name = "s"} : (tensor) -> tensor - // CHECK: %[[CASE_OUTPUT:.*]] = "tf.CaseRegion"(%[[BRANCH_INDEX]]) ({ + // CHECK: %[[CASE_OUTPUT:.*]] = "tf.CaseRegion"(%[[BRANCH_INDEX]]) {{.*}} ({ %case_op = "tf.CaseRegion"(%arg0) ({ %elem = "tf._SomeOp"() : () -> tensor // CHECK-NOT: tf.StackPushV2 diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tensor_array_ops_decomposition.mlir b/tensorflow/compiler/mlir/tensorflow/tests/tensor_array_ops_decomposition.mlir index 06ecfd401328ed..6adc432958e06a 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/tensor_array_ops_decomposition.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/tensor_array_ops_decomposition.mlir @@ -10,9 +10,9 @@ func.func @main() -> tensor<3xf32> { // CHECK: %[[VAR:.*]] = "tf.MlirLocalVarOp"() : () -> tensor>> // CHECK: "tf.AssignVariableOp"(%[[VAR]], %[[BUFFER]]) %ta:2 = "tf.TensorArrayV3"(%size) {dtype = f32, element_shape = #tf_type.shape<3>, dynamic_size = false, clear_after_read = true, identical_element_shapes = true, tensor_array_name = "ta"} : (tensor) -> (tensor, tensor) - // CHECK: %[[IND:.*]] = "tf.Const"() {value = dense<1> : tensor} : () -> tensor + // CHECK: %[[IND:.*]] = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor %index = "tf.Const"() {value = dense<1> : tensor} : () -> tensor - // CHECK: %[[VAL:.*]] = "tf.Const"() {value = dense<[1.000000e+00, 2.000000e+00, 3.000000e+00]> : tensor<3xf32>} : () -> tensor<3xf32> + // CHECK: %[[VAL:.*]] = "tf.Const"() <{value = dense<[1.000000e+00, 2.000000e+00, 3.000000e+00]> : tensor<3xf32>}> : () -> tensor<3xf32> %value = "tf.Const"() {value = dense<[1.0, 2.0, 3.0]> : tensor<3xf32>} : () -> tensor<3xf32> // CHECK: %[[READ_VAR:.*]] = "tf.ReadVariableOp"(%[[VAR]]) // CHECK: %[[UPDATE_SLICE:.*]] = "tf.Reshape"(%[[VAL]] @@ -46,7 +46,7 @@ func.func @main() -> tensor { %index = "tf.Const"() {value = dense<1> : tensor} : () -> tensor %value = "tf.Const"() {value = dense<[1.0, 2.0, 3.0]> : tensor<3xf32>} : () -> tensor<3xf32> %write = "tf.TensorArrayWriteV3"(%ta#0, %index, %value, %ta#1) : (tensor, tensor, tensor<3xf32>, tensor) -> tensor - // CHECK: %[[SIZE:.*]] = "tf.Const"() {value = dense<5> : tensor} : () -> tensor + // CHECK: %[[SIZE:.*]] = "tf.Const"() <{value = dense<5> : tensor}> : () -> tensor %size_out = "tf.TensorArraySizeV3"(%ta#0, %write) : (tensor, tensor) -> tensor // CHECK: return %[[SIZE]] : tensor func.return %size_out : tensor @@ -110,7 +110,7 @@ func.func @main() -> () { // CHECK: %[[READ:.*]] = "tf.ReadVariableOp"(%[[VAR]]) : (tensor>>) -> tensor<5x3xf32> // CHECK: %[[CONCAT_RESHAPE:.*]] = "tf.Reshape"(%[[READ]], // CHECK-SAME: -> tensor<15xf32> - // CHECK: %[[LENS:.*]] = "tf.Const"() {value = dense<3> : tensor<5xi64>} : () -> tensor<5xi64> + // CHECK: %[[LENS:.*]] = "tf.Const"() <{value = dense<3> : tensor<5xi64>}> : () -> tensor<5xi64> %concat:2 = "tf.TensorArrayConcatV3"(%ta#0, %ta#1) {element_shape_except0 = #tf_type.shape<*>} : (tensor, tensor) -> (tensor<*xf32>, tensor<*xi64>) // CHECK: %[[SPLIT_RESHAPE:.*]] = "tf.Reshape"(%[[CONCAT_RESHAPE]], // CHECK-SAME: -> tensor<5x3xf32> @@ -153,33 +153,33 @@ func.func @main() -> () { // CHECK: %[[VAR:.*]] = "tf.MlirLocalVarOp"() : () -> tensor>> // CHECK: "tf.AssignVariableOp" %ta:2 = "tf.TensorArrayV3"(%size) {dtype = f32, element_shape = #tf_type.shape<3>, dynamic_size = false, clear_after_read = true, identical_element_shapes = true, tensor_array_name = "ta"} : (tensor) -> (tensor, tensor) - // CHECK: %[[INDS:.*]] = "tf.Const"() {value = dense<[2, 1]> : tensor<2xi32>} : () -> tensor<2xi32> + // CHECK: %[[INDS:.*]] = "tf.Const"() <{value = dense<[2, 1]> : tensor<2xi32>}> : () -> tensor<2xi32> %indices = "tf.Const"() {value = dense<[2, 1]> : tensor<2xi32>} : () -> tensor<2xi32> // CHECK: %[[READ:.*]] = "tf.ReadVariableOp"(%[[VAR]]) : (tensor>>) -> tensor<5x3xf32> - // CHECK: %[[AXIS:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor + // CHECK: %[[AXIS:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK: %[[GATHER:.*]] = "tf.GatherV2"(%[[READ]], %[[INDS]], %[[AXIS]]) : (tensor<5x3xf32>, tensor<2xi32>, tensor) -> tensor<2x3xf32> %gather = "tf.TensorArrayGatherV3"(%ta#0, %indices, %ta#1) {element_shape = #tf_type.shape<*>} : (tensor, tensor<2xi32>, tensor) -> tensor<*xf32> // CHECK: %[[READ2:.*]] = "tf.ReadVariableOp"(%[[VAR]]) : (tensor>>) -> tensor<5x3xf32> - // CHECK-DAG: %[[SLICE_SIZE:.*]] = "tf.Const"() {value = dense<[1, 3]> : tensor<2xi32>} : () -> tensor<2xi32> - // CHECK-DAG: %[[IND_SLICE0_START:.*]] = "tf.Const"() {value = dense<0> : tensor<1xi32>} : () -> tensor<1xi32> - // CHECK-DAG: %[[IND_SLICE0_SIZE:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32> + // CHECK-DAG: %[[SLICE_SIZE:.*]] = "tf.Const"() <{value = dense<[1, 3]> : tensor<2xi32>}> : () -> tensor<2xi32> + // CHECK-DAG: %[[IND_SLICE0_START:.*]] = "tf.Const"() <{value = dense<0> : tensor<1xi32>}> : () -> tensor<1xi32> + // CHECK-DAG: %[[IND_SLICE0_SIZE:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}> : () -> tensor<1xi32> // CHECK: %[[IND_SLICE0:.*]] = "tf.Slice"(%[[INDS]], %[[IND_SLICE0_START]], %[[IND_SLICE0_SIZE]]) : (tensor<2xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32> // CHECK: %[[SLICE0_START:.*]] = "tf.ConcatV2"(%[[IND_SLICE0]], // CHECK: %[[OLD_SLICE0:.*]] = "tf.Slice"(%[[READ2]], %[[SLICE0_START]], // CHECK-SAME: (tensor<5x3xf32>, tensor<2xi32>, tensor<2xi32>) -> tensor<1x3xf32> - // CHECK: %[[UPDATE_SLICE0_START:.*]] = "tf.Const"() {value = dense<0> : tensor<2xi32>} : () -> tensor<2xi32> + // CHECK: %[[UPDATE_SLICE0_START:.*]] = "tf.Const"() <{value = dense<0> : tensor<2xi32>}> : () -> tensor<2xi32> // CHECK: %[[UPDATE_SLICE0:.*]] = "tf.Slice"(%[[GATHER]], %[[UPDATE_SLICE0_START]], %[[SLICE_SIZE]]) : (tensor<2x3xf32>, tensor<2xi32>, tensor<2xi32>) -> tensor<1x3xf32> // CHECK: %[[ADD0:.*]] = "tf.AddV2"(%[[OLD_SLICE0]], %[[UPDATE_SLICE0]]) // CHECK: %[[UPDATE0:.*]] = "tf.XlaDynamicUpdateSlice"(%[[READ2]], %[[ADD0]] // CHECK-SAME: (tensor<5x3xf32>, tensor<1x3xf32>, tensor<2xi32>) -> tensor<5x3xf32> - // CHECK-DAG: %[[IND_SLICE1_START:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32> - // CHECK-DAG: %[[IND_SLICE1_SIZE:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32> + // CHECK-DAG: %[[IND_SLICE1_START:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}> : () -> tensor<1xi32> + // CHECK-DAG: %[[IND_SLICE1_SIZE:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}> : () -> tensor<1xi32> // CHECK: %[[IND_SLICE1:.*]] = "tf.Slice"(%[[INDS]], %[[IND_SLICE1_START]], %[[IND_SLICE1_SIZE]]) : (tensor<2xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32> // CHECK: %[[SLICE1_START:.*]] = "tf.ConcatV2"(%[[IND_SLICE1]], // CHECK: %[[OLD_SLICE1:.*]] = "tf.Slice"(%[[UPDATE0]], %[[SLICE1_START]], // CHECK-SAME: (tensor<5x3xf32>, tensor<2xi32>, tensor<2xi32>) -> tensor<1x3xf32> - // CHECK: %[[UPDATE_SLICE1_START:.*]] = "tf.Const"() {value = dense<[1, 0]> : tensor<2xi32>} : () -> tensor<2xi32> + // CHECK: %[[UPDATE_SLICE1_START:.*]] = "tf.Const"() <{value = dense<[1, 0]> : tensor<2xi32>}> : () -> tensor<2xi32> // CHECK: %[[UPDATE_SLICE1:.*]] = "tf.Slice"(%[[GATHER]], %[[UPDATE_SLICE1_START]], %[[SLICE_SIZE]]) : (tensor<2x3xf32>, tensor<2xi32>, tensor<2xi32>) -> tensor<1x3xf32> // CHECK: %[[ADD1:.*]] = "tf.AddV2"(%[[OLD_SLICE1]], %[[UPDATE_SLICE1]]) // CHECK: %[[UPDATE1:.*]] = "tf.XlaDynamicUpdateSlice"(%[[UPDATE0]], %[[ADD1]] @@ -200,7 +200,7 @@ func.func @main() { // CHECK: "tf.AssignVariableOp"(%[[VAR]], %ta:2 = "tf.TensorArrayV3"(%size) {dtype = f32, element_shape = #tf_type.shape<3>, dynamic_size = false, clear_after_read = true, identical_element_shapes = true, tensor_array_name = "ta"} : (tensor) -> (tensor, tensor) %index = "tf.Const"() {value = dense<1> : tensor} : () -> tensor - // CHECK: %[[VALUE:.*]] = "tf.Const"() {value = dense<[1.000000e+00, 2.000000e+00, 3.000000e+00]> : tensor<3xf32>} : () -> tensor<3xf32> + // CHECK: %[[VALUE:.*]] = "tf.Const"() <{value = dense<[1.000000e+00, 2.000000e+00, 3.000000e+00]> : tensor<3xf32>}> : () -> tensor<3xf32> %value = "tf.Const"() {value = dense<[1.0, 2.0, 3.0]> : tensor<3xf32>} : () -> tensor<3xf32> // CHECK: %[[GVAR1:.*]] = "tf.MlirLocalVarOp"() : () -> tensor>> // CHECK: "tf.AssignVariableOp"(%[[GVAR1]], @@ -240,7 +240,7 @@ func.func @main() { // CHECK-LABEL: func @main func.func @main() -> () { - // CHECK: %[[SIZE:.*]] = "tf.Const"() {value = dense<5> : tensor} : () -> tensor + // CHECK: %[[SIZE:.*]] = "tf.Const"() <{value = dense<5> : tensor}> : () -> tensor %size = "tf.Const"() {value = dense<5> : tensor} : () -> tensor %index = "tf.Const"() {value = dense<1> : tensor} : () -> tensor // CHECK: %[[VAR:.*]] = "tf.MlirLocalVarOp"() : () -> tensor>> @@ -257,7 +257,7 @@ func.func @main() -> () { } // CHECK: func @while_body(%[[BARG0:.*]]: tensor>>, %[[BARG1:.*]]: tensor, %[[BARG2:.*]]: tensor>>) func.func @while_body(%arg0: tensor, %arg1: tensor) -> (tensor, tensor) { - // CHECK: %[[CONST1:.*]] = "tf.Const"() {value = dense<1> : tensor} : () -> tensor + // CHECK: %[[CONST1:.*]] = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor %const1 = "tf.Const"() {value = dense<1> : tensor} : () -> tensor // CHECK: %[[SUB:.*]] = "tf.Sub"(%[[BARG1]], %[[CONST1]]) %sub = "tf.Sub"(%arg1, %const1) : (tensor, tensor) -> tensor @@ -288,7 +288,7 @@ func.func @while_cond(%arg0: tensor, %arg1: tensor) -> t // CHECK-LABEL: func @main func.func @main() -> () { - // CHECK: %[[SIZE:.*]] = "tf.Const"() {value = dense<5> : tensor} : () -> tensor + // CHECK: %[[SIZE:.*]] = "tf.Const"() <{value = dense<5> : tensor}> : () -> tensor %size = "tf.Const"() {value = dense<5> : tensor} : () -> tensor %index = "tf.Const"() {value = dense<1> : tensor} : () -> tensor // CHECK: %[[VAR:.*]] = "tf.MlirLocalVarOp"() : () -> tensor>> @@ -328,7 +328,7 @@ func.func @then_branch(%arg0: tensor) -> tensor>>, %[[EARG1:.*]]: tensor>>, %[[EARG2:.*]]: tensor>>) func.func @else_branch(%arg0: tensor) -> tensor { - // CHECK: %[[CONST1:.*]] = "tf.Const"() {value = dense<1> : tensor} : () -> tensor + // CHECK: %[[CONST1:.*]] = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor %const1 = "tf.Const"() {value = dense<1> : tensor} : () -> tensor %elem = "tf._SomeOp"() : () -> tensor<3xf32> %flow = "tf.Const"() {value = dense<1.0> : tensor} : () -> tensor @@ -348,14 +348,14 @@ func.func @else_branch(%arg0: tensor) -> tensor () { - // CHECK: %[[SIZE:.*]] = "tf.Const"() {value = dense<5> : tensor} + // CHECK: %[[SIZE:.*]] = "tf.Const"() <{value = dense<5> : tensor}> %size = "tf.Const"() {value = dense<5> : tensor} : () -> tensor %index = "tf.Const"() {value = dense<1> : tensor} : () -> tensor // CHECK: %[[VAR:.*]] = "tf.MlirLocalVarOp"() : () -> tensor>> // CHECK-NOT: tf.TensorArrayV3 %ta:2 = "tf.TensorArrayV3"(%size) {dtype = f32, element_shape = #tf_type.shape<3>, dynamic_size = false, clear_after_read = true, identical_element_shapes = true, tensor_array_name = "ta"} : (tensor) -> (tensor, tensor) - // CHECK: %[[FLOW_INIT:.*]] = "tf.Const"() {value = dense<0.000000e+00> : tensor} - // CHECK: %[[WHILE:.*]]:2 = "tf.WhileRegion"(%[[FLOW_INIT]], %[[SIZE]]) ({ + // CHECK: %[[FLOW_INIT:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor}> + // CHECK: %[[WHILE:.*]]:2 = "tf.WhileRegion"(%[[FLOW_INIT]], %[[SIZE]]) {{.*}} ({ %while:2 = "tf.WhileRegion"(%ta#1, %size) ({ // CHECK: ^bb0(%[[BARG0:.*]]: tensor, %[[BARG1:.*]]: tensor): ^bb0(%barg0: tensor, %barg1: tensor): @@ -402,8 +402,8 @@ func.func @main(%arg0: tensor) -> () { // CHECK: "tf.AssignVariableOp"(%[[TA_BUFFER]] // CHECK-NOT: tf.TensorArrayV3 %ta:2 = "tf.TensorArrayV3"(%size) {dtype = f32, element_shape = #tf_type.shape<3>, dynamic_size = false, clear_after_read = true, identical_element_shapes = true, tensor_array_name = "ta"} : (tensor) -> (tensor, tensor) - // CHECK: "tf.IfRegion"(%[[PRED]]) ({ - %case_op = "tf.IfRegion"(%arg0) ({ + // CHECK: "tf.IfRegion"(%[[PRED]]) <{is_stateless = false}> ({ + %case_op = "tf.IfRegion"(%arg0) <{is_stateless = false}> ({ // CHECK: %[[TA_VAL:.*]] = "tf.ReadVariableOp"(%[[TA_BUFFER]]) // CHECK: "tf.Slice"(%[[TA_VAL]] // CHECK-NOT: tf.TensorArrayReadV3 @@ -420,8 +420,8 @@ func.func @main(%arg0: tensor) -> () { %elem = "tf._SomeOp"() : () -> tensor<3xf32> %write = "tf.TensorArrayWriteV3"(%ta#0, %idx, %elem, %ta#1) : (tensor, tensor, tensor<3xf32>, tensor) -> tensor "tf.Yield"(%write) : (tensor) -> () - // CHECK: }) {is_stateless = false} : (tensor) -> tensor - }) {is_stateless = false} : (tensor) -> tensor + // CHECK: }) : (tensor) -> tensor + }) : (tensor) -> tensor %idx = "tf.Const"() {value = dense<6> : tensor} : () -> tensor // CHECK-NOT: tf.TensorArrayReadV3 %read_val = "tf.TensorArrayReadV3"(%ta#0, %idx, %case_op) : (tensor, tensor, tensor) -> tensor<3xf32> @@ -436,7 +436,7 @@ func.func @main(%arg0: tensor) -> () { // CHECK-LABEL: func @main func.func @main() -> () { - // CHECK: %[[SIZE:.*]] = "tf.Const"() {value = dense<5> : tensor} : () -> tensor + // CHECK: %[[SIZE:.*]] = "tf.Const"() <{value = dense<5> : tensor}> : () -> tensor %size = "tf.Const"() {value = dense<5> : tensor} : () -> tensor %index = "tf.Const"() {value = dense<1> : tensor} : () -> tensor // CHECK: %[[VAR:.*]] = "tf.MlirLocalVarOp"() : () -> tensor>> @@ -486,7 +486,7 @@ func.func @callee(%arg0: tensor) -> tensor // CHECK-LABEL: func @main func.func @main() -> () { - // CHECK: %[[SIZE:.*]] = "tf.Const"() {value = dense<5> : tensor} : () -> tensor + // CHECK: %[[SIZE:.*]] = "tf.Const"() <{value = dense<5> : tensor}> : () -> tensor %size = "tf.Const"() {value = dense<5> : tensor} : () -> tensor %index = "tf.Const"() {value = dense<1> : tensor} : () -> tensor // CHECK: %[[VAR:.*]] = "tf.MlirLocalVarOp"() : () -> tensor>> @@ -543,7 +543,7 @@ func.func @callee() -> tensor { // CHECK: "tf.MlirLocalVarOp"() : () -> tensor>> // CHECK: "tf.AssignVariableOp" %ta:2 = "tf.TensorArrayV3"(%size) {dtype = f32, element_shape = #tf_type.shape<>, dynamic_size = false, clear_after_read = true, identical_element_shapes = true, tensor_array_name = "ta"} : (tensor) -> (tensor, tensor) - // CHECK: %[[SIZE:.*]] = "tf.Const"() {value = dense<5> : tensor} : () -> tensor + // CHECK: %[[SIZE:.*]] = "tf.Const"() <{value = dense<5> : tensor}> : () -> tensor %size_out = "tf.TensorArraySizeV3"(%ta#0, %ta#1) : (tensor, tensor) -> tensor // CHECK: return %[[SIZE]] : tensor func.return %size_out : tensor @@ -553,7 +553,7 @@ func.func @callee() -> tensor { // CHECK-LABEL: func @main func.func @main() -> () { - // CHECK: "tf.PartitionedCall"() {config = "", config_proto = "", executor_type = "", f = @callee} : () -> tensor<*xf32> + // CHECK: "tf.PartitionedCall"() <{config = "", config_proto = "", executor_type = "", f = @callee}> : () -> tensor<*xf32> %call = "tf.PartitionedCall"() {config = "", config_proto = "", executor_type = "", f = @callee} : () -> (tensor<*xf32>) func.return } @@ -567,7 +567,7 @@ func.func private @callee() -> (tensor<*xf32>) { // CHECK: "tf.AssignVariableOp"(%[[LOCAL_VAR]], %[[UPDATE]]) : (tensor>>, tensor<5x3xf32>) -> () %flow = "tf.TensorArrayWriteV3"(%ta#0, %index, %value, %ta#1) : (tensor>>, tensor, tensor<3xf32>, tensor) -> tensor // CHECK: %[[SLICE:.*]] = "tf.Slice" - // CHECK: %[[ELEM_SHAPE:.*]] = "tf.Const"() {value = dense<3> : tensor<1xi32>} + // CHECK: %[[ELEM_SHAPE:.*]] = "tf.Const"() <{value = dense<3> : tensor<1xi32>}> // CHECK: %[[ELEM:.*]] = "tf.Reshape"(%[[SLICE]], %[[ELEM_SHAPE]]) %val = "tf.TensorArrayReadV3"(%ta#0, %index, %ta#1) : (tensor>>, tensor, tensor) -> tensor<*xf32> // CHECK: %[[CAST:.*]] = tensor.cast %[[ELEM]] : tensor<3xf32> to tensor<*xf32> @@ -604,7 +604,7 @@ func.func private @callee(%arg0: tensor) -> tensor : tensor} : () -> tensor // CHECK: %[[BR_INDEX:.*]] = "tf.SomeOp"() : () -> tensor %branch_index = "tf.SomeOp"() : () -> tensor - // CHECK: "tf.CaseRegion"(%[[BR_INDEX]]) ({ + // CHECK: "tf.CaseRegion"(%[[BR_INDEX]]) {{.*}} ({ "tf.CaseRegion"(%branch_index) ({ // CHECK: %[[READ_GVAR:.*]] = "tf.ReadVariableOp"(%[[GVAR]]) // CHECK: %[[UPDATE:.*]] = "tf.XlaDynamicUpdateSlice"(%[[READ_GVAR]], diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tensor_list_ops_decomposition.mlir b/tensorflow/compiler/mlir/tensorflow/tests/tensor_list_ops_decomposition.mlir index 1177f9f6f9688c..6fb95988b66dd2 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/tensor_list_ops_decomposition.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/tensor_list_ops_decomposition.mlir @@ -4,34 +4,34 @@ // CHECK-LABEL: func @main func.func @main() -> (tensor, tensor) { - // CHECK-NEXT: "tf.Const"() {value = dense<> : tensor<0xi32>} + // CHECK-NEXT: "tf.Const"() <{value = dense<> : tensor<0xi32>}> %elem_shape = "tf.Const"() {value = dense<> : tensor<0xi32>} : () -> tensor<0xi32> - // CHECK-NEXT: "tf.Const"() {value = dense<10> : tensor} + // CHECK-NEXT: "tf.Const"() <{value = dense<10> : tensor}> %max_size = "tf.Const"() {value = dense<10> : tensor} : () -> tensor - // CHECK-NEXT: %[[ZERO_SCALAR:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor + // CHECK-NEXT: %[[ZERO_SCALAR:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK-NEXT: %[[CAST_ZERO:.*]] = "tf.Cast"(%[[ZERO_SCALAR]]) : (tensor) -> tensor - // CHECK-NEXT: %[[CONST10:.*]] = "tf.Const"() {value = dense<10> : tensor<1xi32>} : () -> tensor<1xi32> + // CHECK-NEXT: %[[CONST10:.*]] = "tf.Const"() <{value = dense<10> : tensor<1xi32>}> : () -> tensor<1xi32> // CHECK-NEXT: %[[BROADCAST:.*]] = "tf.BroadcastTo"(%[[CAST_ZERO]], %[[CONST10]]) : (tensor, tensor<1xi32>) -> tensor<10xf32> - // CHECK-NEXT: %[[ZERO:.*]] = "tf.Const"() {value = dense<0> : tensor<1xi32>} : () -> tensor<1xi32> + // CHECK-NEXT: %[[ZERO:.*]] = "tf.Const"() <{value = dense<0> : tensor<1xi32>}> : () -> tensor<1xi32> %tl = "tf.EmptyTensorList"(%elem_shape, %max_size) : (tensor<0xi32>, tensor) -> tensor>> %id = "tf.Identity"(%tl) : (tensor>>) -> tensor>> // CHECK-NEXT: %[[PUSHVAL:.*]] = "tf._SomeOp"() %elem = "tf._SomeOp"() : () -> tensor - // CHECK-NEXT: %[[UPDATE_SHAPE:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32> + // CHECK-NEXT: %[[UPDATE_SHAPE:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}> : () -> tensor<1xi32> // CHECK-NEXT: %[[UPDATE_SLICE:.*]] = "tf.Reshape"(%[[PUSHVAL]], %[[UPDATE_SHAPE]]) : (tensor, tensor<1xi32>) -> tensor<1xf32> // CHECK-NEXT: %[[UPDATE:.*]] = "tf.XlaDynamicUpdateSlice"(%[[BROADCAST]], %[[UPDATE_SLICE]], %[[ZERO]]) : (tensor<10xf32>, tensor<1xf32>, tensor<1xi32>) -> tensor<10xf32> - // CHECK-NEXT: %[[CONST1:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32> + // CHECK-NEXT: %[[CONST1:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}> : () -> tensor<1xi32> // CHECK-NEXT: %[[NEW_SIZE:.*]] = "tf.AddV2"(%[[ZERO]], %[[CONST1]]) : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32> %push = "tf.TensorListPushBack"(%id, %elem) : (tensor>>, tensor) -> tensor>> // CHECK-NEXT: %[[COPY:.*]] = "tf.Identity"(%[[UPDATE]]) - // CHECK-NEXT: %[[CONST1_1:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32> + // CHECK-NEXT: %[[CONST1_1:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}> : () -> tensor<1xi32> // CHECK-NEXT: %[[SUB:.*]] = "tf.Sub"(%[[NEW_SIZE]], %[[CONST1_1]]) - // CHECK-NEXT: %[[SLICE_SIZE:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32> + // CHECK-NEXT: %[[SLICE_SIZE:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}> : () -> tensor<1xi32> // CHECK-NEXT: %[[SLICE:.*]] = "tf.Slice"(%[[COPY]], %[[SUB]], %[[SLICE_SIZE]]) : (tensor<10xf32>, tensor<1xi32>, tensor<1xi32>) -> tensor<1xf32> - // CHECK-NEXT: %[[ELEM_SHAPE:.*]] = "tf.Const"() {value = dense<> : tensor<0xi32>} : () -> tensor<0xi32> + // CHECK-NEXT: %[[ELEM_SHAPE:.*]] = "tf.Const"() <{value = dense<> : tensor<0xi32>}> : () -> tensor<0xi32> // CHECK-NEXT: %[[ELEM:.*]] = "tf.Reshape"(%[[SLICE]], %[[ELEM_SHAPE]]) : (tensor<1xf32>, tensor<0xi32>) -> tensor %pop:2 = "tf.TensorListPopBack"(%push, %elem_shape) : (tensor>>, tensor<0xi32>) -> (tensor>>, tensor) - // CHECK-NEXT: %[[SCALAR_SHAPE:.*]] = "tf.Const"() {value = dense<> : tensor<0xi32>} + // CHECK-NEXT: %[[SCALAR_SHAPE:.*]] = "tf.Const"() <{value = dense<> : tensor<0xi32>}> // CHECK-NEXT: %[[LENGTH:.*]] = "tf.Reshape"(%[[NEW_SIZE]], %[[SCALAR_SHAPE]]) %length = "tf.TensorListLength"(%push) : (tensor>>) -> tensor // CHECK-NEXT: return %[[ELEM]], %[[LENGTH]] : tensor, tensor @@ -46,30 +46,30 @@ func.func @main() -> (tensor, tensor) { // CHECK-LABEL: func @main // CHECK-SAME: (%[[ARG0:.*]]: tensor) -> (tensor, tensor<10xf32>, tensor) func.func @main(%arg0: tensor) -> (tensor, tensor<10xf32>, tensor) { - // CHECK-NEXT: "tf.Const"() {value = dense<> : tensor<0xi32>} + // CHECK-NEXT: "tf.Const"() <{value = dense<> : tensor<0xi32>}> %elem_shape = "tf.Const"() {value = dense<> : tensor<0xi32>} : () -> tensor<0xi32> - // CHECK-NEXT: %[[NUM:.*]] = "tf.Const"() {value = dense<10> : tensor} + // CHECK-NEXT: %[[NUM:.*]] = "tf.Const"() <{value = dense<10> : tensor}> %num = "tf.Const"() {value = dense<10> : tensor} : () -> tensor - // CHECK-NEXT: %[[ZERO_SCALAR:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor + // CHECK-NEXT: %[[ZERO_SCALAR:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK-NEXT: %[[CAST_ZERO:.*]] = "tf.Cast"(%[[ZERO_SCALAR]]) : (tensor) -> tensor - // CHECK-NEXT: %[[CONST10:.*]] = "tf.Const"() {value = dense<10> : tensor<1xi32>} : () -> tensor<1xi32> + // CHECK-NEXT: %[[CONST10:.*]] = "tf.Const"() <{value = dense<10> : tensor<1xi32>}> : () -> tensor<1xi32> // CHECK-NEXT: %[[BROADCAST:.*]] = "tf.BroadcastTo"(%[[CAST_ZERO]], %[[CONST10]]) : (tensor, tensor<1xi32>) -> tensor<10xf32> - // CHECK-NEXT: %[[SIZE_SHAPE:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi32>} + // CHECK-NEXT: %[[SIZE_SHAPE:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}> // CHECK-NEXT: %[[SIZE:.*]] = "tf.Reshape"(%[[NUM]], %[[SIZE_SHAPE]]) %tl = "tf.TensorListReserve"(%elem_shape, %num) : (tensor<0xi32>, tensor) -> tensor>> // CHECK-NEXT: %[[SETVAL:.*]] = "tf._SomeOp"() %elem = "tf._SomeOp"() : () -> tensor - // CHECK-NEXT: %[[SIZE_SHAPE1:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi32>} + // CHECK-NEXT: %[[SIZE_SHAPE1:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}> // CHECK-NEXT: %[[SET_INDEX:.*]] = "tf.Reshape"(%[[ARG0]], %[[SIZE_SHAPE1]]) : (tensor, tensor<1xi32>) -> tensor<1xi32> - // CHECK-NEXT: %[[UPDATE_SHAPE:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32> + // CHECK-NEXT: %[[UPDATE_SHAPE:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}> : () -> tensor<1xi32> // CHECK-NEXT: %[[UPDATE_SLICE:.*]] = "tf.Reshape"(%[[SETVAL]], %[[UPDATE_SHAPE]]) : (tensor, tensor<1xi32>) -> tensor<1xf32> // CHECK-NEXT: %[[UPDATE:.*]] = "tf.XlaDynamicUpdateSlice"(%[[BROADCAST]], %[[UPDATE_SLICE]], %[[SET_INDEX]]) : (tensor<10xf32>, tensor<1xf32>, tensor<1xi32>) -> tensor<10xf32> %set = "tf.TensorListSetItem"(%tl, %arg0, %elem) : (tensor>>, tensor, tensor) -> tensor>> - // CHECK-NEXT: %[[SIZE_SHAPE2:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi32>} + // CHECK-NEXT: %[[SIZE_SHAPE2:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}> // CHECK-NEXT: %[[GET_INDEX:.*]] = "tf.Reshape"(%[[ARG0]], %[[SIZE_SHAPE2]]) : (tensor, tensor<1xi32>) -> tensor<1xi32> - // CHECK-NEXT: %[[SLICE_SIZE:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32> + // CHECK-NEXT: %[[SLICE_SIZE:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}> : () -> tensor<1xi32> // CHECK-NEXT: %[[SLICE:.*]] = "tf.Slice"(%[[UPDATE]], %[[GET_INDEX]], %[[SLICE_SIZE]]) : (tensor<10xf32>, tensor<1xi32>, tensor<1xi32>) -> tensor<1xf32> - // CHECK-NEXT: %[[ELEM_SHAPE:.*]] = "tf.Const"() {value = dense<> : tensor<0xi32>} : () -> tensor<0xi32> + // CHECK-NEXT: %[[ELEM_SHAPE:.*]] = "tf.Const"() <{value = dense<> : tensor<0xi32>}> : () -> tensor<0xi32> // CHECK-NEXT: %[[ELEM:.*]] = "tf.Reshape"(%[[SLICE]], %[[ELEM_SHAPE]]) : (tensor<1xf32>, tensor<0xi32>) -> tensor %get = "tf.TensorListGetItem"(%set, %arg0, %elem_shape) : (tensor>>, tensor, tensor<0xi32>) -> tensor // CHECK-NEXT: %[[ADDN:.*]] = "tf.AddN"(%[[UPDATE]], %[[BROADCAST]]) : (tensor<10xf32>, tensor<10xf32>) -> tensor<10xf32> @@ -79,7 +79,7 @@ func.func @main(%arg0: tensor) -> (tensor, tensor<10xf32>, tensor // CHECK-NEXT: %[[ADDN2:.*]] = "tf.AddN"(%[[ADDN]], %[[ZEROS_LIKE]]) : (tensor<10xf32>, tensor<10xf32>) -> tensor<10xf32> %addn2 = "tf.AddN"(%addn, %zeros-like) : (tensor>>, tensor>>) -> tensor>> %stack = "tf.TensorListStack"(%addn2, %elem_shape) : (tensor>>, tensor<0xi32>) -> tensor<10xf32> - // CHECK-NEXT: %[[LEN:.*]] = "tf.Const"() {value = dense<10> : tensor} : () -> tensor + // CHECK-NEXT: %[[LEN:.*]] = "tf.Const"() <{value = dense<10> : tensor}> : () -> tensor %length = "tf.TensorListLength"(%addn2) : (tensor>>) -> tensor // CHECK-NEXT: return %[[ELEM]], %[[ADDN2]], %[[LEN]] : tensor, tensor<10xf32>, tensor func.return %get, %stack, %length : tensor, tensor<10xf32>, tensor @@ -92,16 +92,16 @@ func.func @main(%arg0: tensor) -> (tensor, tensor<10xf32>, tensor // CHECK-LABEL: func @main // CHECK-SAME: (%[[ARG0:.*]]: tensor, %[[ARG1:.*]]: tensor<10xf32>) -> tensor func.func @main(%arg0: tensor, %arg1: tensor<10xf32>) -> tensor { - // CHECK-NEXT: "tf.Const"() {value = dense<> : tensor<0xi32>} + // CHECK-NEXT: "tf.Const"() <{value = dense<> : tensor<0xi32>}> %elem_shape = "tf.Const"() {value = dense<> : tensor<0xi32>} : () -> tensor<0xi32> // CHECK-NEXT: %[[BUFFER:.*]] = "tf.Identity"(%[[ARG1]]) : (tensor<10xf32>) -> tensor<10xf32> - // CHECK-NEXT: %[[SIZE:.*]] = "tf.Const"() {value = dense<10> : tensor<1xi32>} : () -> tensor<1xi32> + // CHECK-NEXT: %[[SIZE:.*]] = "tf.Const"() <{value = dense<10> : tensor<1xi32>}> : () -> tensor<1xi32> %tl = "tf.TensorListFromTensor"(%arg1, %elem_shape) : (tensor<10xf32>, tensor<0xi32>) -> tensor>> - // CHECK-NEXT: %[[SIZE_SHAPE:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi32>} + // CHECK-NEXT: %[[SIZE_SHAPE:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}> // CHECK-NEXT: %[[GET_INDEX:.*]] = "tf.Reshape"(%[[ARG0]], %[[SIZE_SHAPE]]) : (tensor, tensor<1xi32>) -> tensor<1xi32> - // CHECK-NEXT: %[[SLICE_SIZE:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32> + // CHECK-NEXT: %[[SLICE_SIZE:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}> : () -> tensor<1xi32> // CHECK-NEXT: %[[SLICE:.*]] = "tf.Slice"(%[[BUFFER]], %[[GET_INDEX]], %[[SLICE_SIZE]]) : (tensor<10xf32>, tensor<1xi32>, tensor<1xi32>) -> tensor<1xf32> - // CHECK-NEXT: %[[ELEM_SHAPE:.*]] = "tf.Const"() {value = dense<> : tensor<0xi32>} : () -> tensor<0xi32> + // CHECK-NEXT: %[[ELEM_SHAPE:.*]] = "tf.Const"() <{value = dense<> : tensor<0xi32>}> : () -> tensor<0xi32> // CHECK-NEXT: %[[ELEM:.*]] = "tf.Reshape"(%[[SLICE]], %[[ELEM_SHAPE]]) : (tensor<1xf32>, tensor<0xi32>) -> tensor %get = "tf.TensorListGetItem"(%tl, %arg0, %elem_shape) : (tensor>>, tensor, tensor<0xi32>) -> tensor // CHECK-NEXT: return %[[ELEM]] : tensor @@ -116,7 +116,7 @@ func.func @main(%arg0: tensor, %arg1: tensor<10xf32>) -> tensor { func.func @main(%arg0: tensor<10x8x9xf32>) -> tensor<2xi64> { %elem_shape = "tf.Const"() {value = dense<[8, 9]> : tensor<2xi32>} : () -> tensor<2xi32> %tl = "tf.TensorListFromTensor"(%arg0, %elem_shape) : (tensor<10x8x9xf32>, tensor<2xi32>) -> tensor>> - // CHECK: %[[SHAPE:.*]] = "tf.Const"() {value = dense<[8, 9]> : tensor<2xi64>} : () -> tensor<2xi64> + // CHECK: %[[SHAPE:.*]] = "tf.Const"() <{value = dense<[8, 9]> : tensor<2xi64>}> : () -> tensor<2xi64> %shape = "tf.TensorListElementShape"(%tl) : (tensor>>) -> tensor<2xi64> // CHECK-NEXT: return %[[SHAPE]] : tensor<2xi64> func.return %shape: tensor<2xi64> @@ -132,7 +132,7 @@ func.func @main(%arg0: tensor<10x8x9xf32>, %arg1: tensor<3xi32>) -> tensor<3x8x9 %elem_shape = "tf.Const"() {value = dense<[8, 9]> : tensor<2xi32>} : () -> tensor<2xi32> // CHECK: %[[BUFFER:.*]] = "tf.Identity"(%[[ARG0]]) : (tensor<10x8x9xf32>) -> tensor<10x8x9xf32> %tl = "tf.TensorListFromTensor"(%arg0, %elem_shape) : (tensor<10x8x9xf32>, tensor<2xi32>) -> tensor>> - // CHECK: %[[AXIS:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor + // CHECK: %[[AXIS:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK: %[[GATHER:.*]] = "tf.GatherV2"(%[[BUFFER]], %[[ARG1]], %[[AXIS]]) : (tensor<10x8x9xf32>, tensor<3xi32>, tensor) -> tensor<3x8x9xf32> %gather = "tf.TensorListGather"(%tl, %arg1, %elem_shape) : (tensor>>, tensor<3xi32>, tensor<2xi32>) -> tensor<3x8x9xf32> // CHECK-NEXT: return %[[GATHER]] : tensor<3x8x9xf32> @@ -149,7 +149,7 @@ func.func @main(%arg0: tensor<10x8x9xf32>, %arg1: tensor<5xi32>, %arg2: tensor<5 %elem_shape = "tf.Const"() {value = dense<[8, 9]> : tensor<2xi32>} : () -> tensor<2xi32> // CHECK: %[[BUFFER:.*]] = "tf.Identity"(%[[ARG0]]) : (tensor<10x8x9xf32>) -> tensor<10x8x9xf32> %tl = "tf.TensorListFromTensor"(%arg0, %elem_shape) : (tensor<10x8x9xf32>, tensor<2xi32>) -> tensor>> - // CHECK: %[[IND_SHAPE:.*]] = "tf.Const"() {value = dense<[5, 1]> : tensor<2xi32>} : () -> tensor<2xi32> + // CHECK: %[[IND_SHAPE:.*]] = "tf.Const"() <{value = dense<[5, 1]> : tensor<2xi32>}> : () -> tensor<2xi32> // CHECK: %[[IND_RESHPE:.*]] = "tf.Reshape"(%[[ARG1]], %[[IND_SHAPE]]) : (tensor<5xi32>, tensor<2xi32>) -> tensor<5x1xi32> // CHECK: %[[SC:.*]] = "tf.TensorScatterUpdate"(%[[BUFFER]], %[[IND_RESHPE]], %[[ARG2]]) : (tensor<10x8x9xf32>, tensor<5x1xi32>, tensor<5x8x9xf32>) -> tensor<10x8x9xf32> %scatter = "tf.TensorListScatterIntoExistingList"(%tl, %arg2, %arg1) : (tensor>>, tensor<5x8x9xf32>, tensor<5xi32>) -> tensor>> @@ -179,14 +179,14 @@ func.func @main() -> () { } // CHECK: func @while_body(%[[BARG0:.*]]: tensor<10xf32>, %[[BARG1:.*]]: tensor, %[[BARG2:.*]]: tensor<1xi32>) func.func @while_body(%arg0: tensor>>, %arg1: tensor) -> (tensor>>, tensor) { - // CHECK: %[[CONST1:.*]] = "tf.Const"() {value = dense<1> : tensor} : () -> tensor + // CHECK: %[[CONST1:.*]] = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor %const1 = "tf.Const"() {value = dense<1> : tensor} : () -> tensor // CHECK: %[[SUB:.*]] = "tf.Sub"(%[[BARG1]], %[[CONST1]]) %sub = "tf.Sub"(%arg1, %const1) : (tensor, tensor) -> tensor %elem = "tf._SomeOp"() : () -> tensor // CHECK-NOT: "tf.TensorListPushBack" // CHECK: %[[UPDATE:.*]] = "tf.XlaDynamicUpdateSlice" - // CHECK: %[[CONST1:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32> + // CHECK: %[[CONST1:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}> : () -> tensor<1xi32> // CHECK: %[[ADD:.*]] = "tf.AddV2"(%[[BARG2]], %[[CONST1]]) // CHECK-NOT: "tf.TensorListPushBack" %push = "tf.TensorListPushBack"(%arg0, %elem) : (tensor>>, tensor) -> tensor>> @@ -222,7 +222,7 @@ func.func @if_then(%arg0: tensor>>) -> tensor tensor // CHECK-NOT: "tf.TensorListPushBack" // CHECK: %[[UPDATE:.*]] = "tf.XlaDynamicUpdateSlice" - // CHECK: %[[CONST1:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32> + // CHECK: %[[CONST1:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}> : () -> tensor<1xi32> // CHECK: %[[ADD:.*]] = "tf.AddV2"(%[[TARG1]], %[[CONST1]]) // CHECK-NOT: "tf.TensorListPushBack" %push = "tf.TensorListPushBack"(%arg0, %elem) : (tensor>>, tensor) -> tensor>> @@ -234,11 +234,11 @@ func.func @if_else(%arg0: tensor>>) -> tensor : tensor<0xi32>} : () -> tensor<0xi32> // CHECK-NOT: "tf.TensorListPopBack" // CHECK: %[[COPY:.*]] = "tf.Identity"(%[[EARG0]]) - // CHECK: %[[CONST1_1:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32> + // CHECK: %[[CONST1_1:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}> : () -> tensor<1xi32> // CHECK: %[[SUB:.*]] = "tf.Sub"(%[[EARG1]], %[[CONST1_1]]) - // CHECK: %[[SLICE_SIZE:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32> + // CHECK: %[[SLICE_SIZE:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}> : () -> tensor<1xi32> // CHECK: %[[SLICE:.*]] = "tf.Slice"(%[[COPY]], %[[SUB]], %[[SLICE_SIZE]]) : (tensor<10xf32>, tensor<1xi32>, tensor<1xi32>) -> tensor<1xf32> - // CHECK: %[[ELEM_SHAPE:.*]] = "tf.Const"() {value = dense<> : tensor<0xi32>} : () -> tensor<0xi32> + // CHECK: %[[ELEM_SHAPE:.*]] = "tf.Const"() <{value = dense<> : tensor<0xi32>}> : () -> tensor<0xi32> // CHECK: %[[ELEM:.*]] = "tf.Reshape"(%[[SLICE]], %[[ELEM_SHAPE]]) : (tensor<1xf32>, tensor<0xi32>) -> tensor // CHECK-NOT: "tf.TensorListPopBack" %pop:2 = "tf.TensorListPopBack"(%arg0, %elem_shape) : (tensor>>, tensor<0xi32>) -> (tensor>>, tensor) @@ -269,7 +269,7 @@ func.func @branch_0(%arg0: tensor>>) -> tensor tensor // CHECK-NOT: "tf.TensorListPushBack" // CHECK: %[[UPDATE:.*]] = "tf.XlaDynamicUpdateSlice" - // CHECK: %[[CONST1:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32> + // CHECK: %[[CONST1:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}> : () -> tensor<1xi32> // CHECK: %[[ADD:.*]] = "tf.AddV2"(%[[TARG1]], %[[CONST1]]) // CHECK-NOT: "tf.TensorListPushBack" %push = "tf.TensorListPushBack"(%arg0, %elem) : (tensor>>, tensor) -> tensor>> @@ -281,11 +281,11 @@ func.func @branch_1(%arg0: tensor>>) -> tensor : tensor<0xi32>} : () -> tensor<0xi32> // CHECK-NOT: "tf.TensorListPopBack" // CHECK: %[[COPY:.*]] = "tf.Identity"(%[[EARG0]]) - // CHECK: %[[CONST1_1:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32> + // CHECK: %[[CONST1_1:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}> : () -> tensor<1xi32> // CHECK: %[[SUB:.*]] = "tf.Sub"(%[[EARG1]], %[[CONST1_1]]) - // CHECK: %[[SLICE_SIZE:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32> + // CHECK: %[[SLICE_SIZE:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}> : () -> tensor<1xi32> // CHECK: %[[SLICE:.*]] = "tf.Slice"(%[[COPY]], %[[SUB]], %[[SLICE_SIZE]]) : (tensor<10xf32>, tensor<1xi32>, tensor<1xi32>) -> tensor<1xf32> - // CHECK: %[[ELEM_SHAPE:.*]] = "tf.Const"() {value = dense<> : tensor<0xi32>} : () -> tensor<0xi32> + // CHECK: %[[ELEM_SHAPE:.*]] = "tf.Const"() <{value = dense<> : tensor<0xi32>}> : () -> tensor<0xi32> // CHECK: %[[ELEM:.*]] = "tf.Reshape"(%[[SLICE]], %[[ELEM_SHAPE]]) : (tensor<1xf32>, tensor<0xi32>) -> tensor // CHECK-NOT: "tf.TensorListPopBack" %pop:2 = "tf.TensorListPopBack"(%arg0, %elem_shape) : (tensor>>, tensor<0xi32>) -> (tensor>>, tensor) @@ -297,11 +297,11 @@ func.func @branch_2(%arg0: tensor>>) -> tensor : tensor<0xi32>} : () -> tensor<0xi32> // CHECK-NOT: "tf.TensorListPopBack" // CHECK: %[[COPY:.*]] = "tf.Identity"(%[[EARG0]]) - // CHECK: %[[CONST1_1:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32> + // CHECK: %[[CONST1_1:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}> : () -> tensor<1xi32> // CHECK: %[[SUB:.*]] = "tf.Sub"(%[[EARG1]], %[[CONST1_1]]) - // CHECK: %[[SLICE_SIZE:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32> + // CHECK: %[[SLICE_SIZE:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}> : () -> tensor<1xi32> // CHECK: %[[SLICE:.*]] = "tf.Slice"(%[[COPY]], %[[SUB]], %[[SLICE_SIZE]]) : (tensor<10xf32>, tensor<1xi32>, tensor<1xi32>) -> tensor<1xf32> - // CHECK: %[[ELEM_SHAPE:.*]] = "tf.Const"() {value = dense<> : tensor<0xi32>} : () -> tensor<0xi32> + // CHECK: %[[ELEM_SHAPE:.*]] = "tf.Const"() <{value = dense<> : tensor<0xi32>}> : () -> tensor<0xi32> // CHECK: %[[ELEM:.*]] = "tf.Reshape"(%[[SLICE]], %[[ELEM_SHAPE]]) : (tensor<1xf32>, tensor<0xi32>) -> tensor // CHECK-NOT: "tf.TensorListPopBack" %pop:2 = "tf.TensorListPopBack"(%arg0, %elem_shape) : (tensor>>, tensor<0xi32>) -> (tensor>>, tensor) @@ -317,7 +317,7 @@ func.func @main() -> tensor { %size = "tf.Const"() {value = dense<10> : tensor} : () -> tensor // CHECK-NOT: tf.EmptyTensorList %tl = "tf.EmptyTensorList"(%elem_shape, %size) : (tensor<0xi32>, tensor) -> tensor>> - %while_op:2 = "tf.WhileRegion"(%tl, %size) ({ + %while_op:2 = "tf.WhileRegion"(%tl, %size) <{is_stateless = false}> ({ // CHECK: ^bb0(%[[CARG0:.*]]: tensor<10xf32>, %[[CARG1:.*]]: tensor, %[[CARG2:.*]]: tensor<1xi32>): ^bb0(%arg0: tensor>>, %arg1: tensor): // CHECK: %[[PRED:.*]] = "tf._SomeOp"() @@ -327,7 +327,7 @@ func.func @main() -> tensor { }, { // CHECK: ^bb0(%[[CARG0:.*]]: tensor<10xf32>, %[[CARG1:.*]]: tensor, %[[CARG2:.*]]: tensor<1xi32>): ^bb0(%arg0: tensor>>, %arg1: tensor): - // CHECK: %[[CST:.*]] = "tf.Const"() {value = dense<1> : tensor} : () -> tensor + // CHECK: %[[CST:.*]] = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor // CHECK: %[[SUB:.*]] = "tf.Sub"(%[[CARG1]], %[[CST]]) // CHECK: %[[ELEM:.*]] = "tf._SomeOp"() : () -> tensor %cst = "tf.Const"() {value = dense<1> : tensor} : () -> tensor @@ -335,14 +335,14 @@ func.func @main() -> tensor { %elem = "tf._SomeOp"() : () -> tensor // CHECK-NOT: "tf.TensorListPushBack" // CHECK: %[[UPDATE:.*]] = "tf.XlaDynamicUpdateSlice"(%[[CARG0]] - // CHECK: %[[ONE:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi32>} + // CHECK: %[[ONE:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}> // CHECK: %[[ADD:.*]] = "tf.AddV2"(%[[CARG2]], %[[ONE]]) // CHECK-NOT: "tf.TensorListPushBack" // CHECK: "tf.Yield"(%[[UPDATE]], %[[SUB]], %[[ADD]]) - // CHECK: }) {is_stateless = false} + // CHECK: }) %push = "tf.TensorListPushBack"(%arg0, %elem) : (tensor>>, tensor) -> tensor>> "tf.Yield"(%push, %sub) : (tensor>>, tensor) -> () - }) {is_stateless = false} : (tensor>>, tensor) -> (tensor>>, tensor) + }) : (tensor>>, tensor) -> (tensor>>, tensor) // CHECK: "tf.Slice" // CHECK-NOT: tf.TensorListPopBack %pop:2 = "tf.TensorListPopBack"(%while_op#0, %elem_shape) : (tensor>>, tensor<0xi32>) -> (tensor>>, tensor) @@ -356,27 +356,27 @@ func.func @main(%arg0: tensor) -> () { %elem_shape = "tf.Const"() {value = dense<> : tensor<0xi32>} : () -> tensor<0xi32> %max_size = "tf.Const"() {value = dense<10> : tensor} : () -> tensor %tl = "tf.EmptyTensorList"(%elem_shape, %max_size) : (tensor<0xi32>, tensor) -> tensor>> - // CHECK: %[[ZERO:.*]] = "tf.Const"() {value = dense<0> : tensor} + // CHECK: %[[ZERO:.*]] = "tf.Const"() <{value = dense<0> : tensor}> // CHECK: %[[ZERO_F32:.*]] = "tf.Cast"(%[[ZERO]]) - // CHECK: %[[MAX_SIZE:.*]] = "tf.Const"() {value = dense<10> : tensor<1xi32>} + // CHECK: %[[MAX_SIZE:.*]] = "tf.Const"() <{value = dense<10> : tensor<1xi32>}> // CHECK: %[[BUFFER:.*]] = "tf.BroadcastTo"(%[[ZERO_F32]], %[[MAX_SIZE]]) - // CHECK: %[[BUFFER_SIZE:.*]] = "tf.Const"() {value = dense<0> : tensor<1xi32>} + // CHECK: %[[BUFFER_SIZE:.*]] = "tf.Const"() <{value = dense<0> : tensor<1xi32>}> // CHECK-NOT: tf.EmptyTensorList %if_op = "tf.IfRegion"(%arg0) ({ %elem = "tf._SomeOp"() : () -> tensor // CHECK: %[[UPDATE:.*]] = "tf.XlaDynamicUpdateSlice" - // CHECK: %[[ONE:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32> + // CHECK: %[[ONE:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}> : () -> tensor<1xi32> // CHECK: %[[ADD:.*]] = "tf.AddV2"(%[[BUFFER_SIZE]], %[[ONE]]) // CHECK-NOT: "tf.TensorListPushBack" %push = "tf.TensorListPushBack"(%tl, %elem) : (tensor>>, tensor) -> tensor>> "tf.Yield" (%push) : (tensor>>) -> () }, { // CHECK: %[[COPY:.*]] = "tf.Identity"(%[[BUFFER]]) - // CHECK: %[[ONE:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi32>} + // CHECK: %[[ONE:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}> // CHECK: %[[SUB:.*]] = "tf.Sub"(%[[BUFFER_SIZE]], %[[ONE]]) - // CHECK: %[[SLICE_SIZE:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi32>} + // CHECK: %[[SLICE_SIZE:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}> // CHECK: %[[SLICE:.*]] = "tf.Slice"(%[[COPY]], %[[SUB]], %[[SLICE_SIZE]]) - // CHECK: %[[ELEM_SHAPE:.*]] = "tf.Const"() {value = dense<> : tensor<0xi32>} + // CHECK: %[[ELEM_SHAPE:.*]] = "tf.Const"() <{value = dense<> : tensor<0xi32>}> // CHECK: %[[ELEM:.*]] = "tf.Reshape"(%[[SLICE]], %[[ELEM_SHAPE]]) // CHECK-NOT: "tf.TensorListPopBack" %pop:2 = "tf.TensorListPopBack"(%tl, %elem_shape) : (tensor>>, tensor<0xi32>) -> (tensor>>, tensor) @@ -397,28 +397,28 @@ func.func @main(%arg0: tensor) -> () { func.func @main(%arg0: tensor) -> () { %elem_shape = "tf.Const"() {value = dense<> : tensor<0xi32>} : () -> tensor<0xi32> %max_size = "tf.Const"() {value = dense<10> : tensor} : () -> tensor - // CHECK: %[[ZERO:.*]] = "tf.Const"() {value = dense<0> : tensor} + // CHECK: %[[ZERO:.*]] = "tf.Const"() <{value = dense<0> : tensor}> // CHECK: %[[ZERO_F32:.*]] = "tf.Cast"(%[[ZERO]]) - // CHECK: %[[MAX_SIZE:.*]] = "tf.Const"() {value = dense<10> : tensor<1xi32>} + // CHECK: %[[MAX_SIZE:.*]] = "tf.Const"() <{value = dense<10> : tensor<1xi32>}> // CHECK: %[[BUFFER:.*]] = "tf.BroadcastTo"(%[[ZERO_F32]], %[[MAX_SIZE]]) - // CHECK: %[[BUFFER_SIZE:.*]] = "tf.Const"() {value = dense<0> : tensor<1xi32>} + // CHECK: %[[BUFFER_SIZE:.*]] = "tf.Const"() <{value = dense<0> : tensor<1xi32>}> // CHECK-NOT: tf.EmptyTensorList %tl = "tf.EmptyTensorList"(%elem_shape, %max_size) : (tensor<0xi32>, tensor) -> tensor>> %case_op = "tf.CaseRegion"(%arg0) ({ %elem = "tf._SomeOp"() : () -> tensor // CHECK: %[[UPDATE:.*]] = "tf.XlaDynamicUpdateSlice" - // CHECK: %[[ONE:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32> + // CHECK: %[[ONE:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}> : () -> tensor<1xi32> // CHECK: %[[ADD:.*]] = "tf.AddV2"(%[[BUFFER_SIZE]], %[[ONE]]) // CHECK-NOT: "tf.TensorListPushBack" %push = "tf.TensorListPushBack"(%tl, %elem) : (tensor>>, tensor) -> tensor>> "tf.Yield" (%push) : (tensor>>) -> () }, { // CHECK: %[[COPY:.*]] = "tf.Identity"(%[[BUFFER]]) - // CHECK: %[[ONE:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi32>} + // CHECK: %[[ONE:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}> // CHECK: %[[SUB:.*]] = "tf.Sub"(%[[BUFFER_SIZE]], %[[ONE]]) - // CHECK: %[[SLICE_SIZE:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi32>} + // CHECK: %[[SLICE_SIZE:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}> // CHECK: %[[SLICE:.*]] = "tf.Slice"(%[[COPY]], %[[SUB]], %[[SLICE_SIZE]]) - // CHECK: %[[ELEM_SHAPE:.*]] = "tf.Const"() {value = dense<> : tensor<0xi32>} + // CHECK: %[[ELEM_SHAPE:.*]] = "tf.Const"() <{value = dense<> : tensor<0xi32>}> // CHECK: %[[ELEM:.*]] = "tf.Reshape"(%[[SLICE]], %[[ELEM_SHAPE]]) // CHECK-NOT: "tf.TensorListPopBack" %pop:2 = "tf.TensorListPopBack"(%tl, %elem_shape) : (tensor>>, tensor<0xi32>) -> (tensor>>, tensor) @@ -426,11 +426,11 @@ func.func @main(%arg0: tensor) -> () { "tf.Yield" (%pop#0) : (tensor>>) -> () }, { // CHECK: %[[COPY:.*]] = "tf.Identity"(%[[BUFFER]]) - // CHECK: %[[ONE:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi32>} + // CHECK: %[[ONE:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}> // CHECK: %[[SUB:.*]] = "tf.Sub"(%[[BUFFER_SIZE]], %[[ONE]]) - // CHECK: %[[SLICE_SIZE:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi32>} + // CHECK: %[[SLICE_SIZE:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}> // CHECK: %[[SLICE:.*]] = "tf.Slice"(%[[COPY]], %[[SUB]], %[[SLICE_SIZE]]) - // CHECK: %[[ELEM_SHAPE:.*]] = "tf.Const"() {value = dense<> : tensor<0xi32>} + // CHECK: %[[ELEM_SHAPE:.*]] = "tf.Const"() <{value = dense<> : tensor<0xi32>}> // CHECK: %[[ELEM:.*]] = "tf.Reshape"(%[[SLICE]], %[[ELEM_SHAPE]]) // CHECK-NOT: "tf.TensorListPopBack" %pop:2 = "tf.TensorListPopBack"(%tl, %elem_shape) : (tensor>>, tensor<0xi32>) -> (tensor>>, tensor) @@ -482,7 +482,7 @@ func.func @callee(%arg0: tensor>>, %arg1: tensor, %[[ARG1:.*]]: tensor, %[[ARG2:.*]]: tensor<1xi32>) -> (tensor<10xf32>, tensor<1xi32>) // CHECK-NOT: "tf.TensorListPushBack" // CHECK: %[[UPDATE:.*]] = "tf.XlaDynamicUpdateSlice" -// CHECK: %[[CONST1:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32> +// CHECK: %[[CONST1:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}> : () -> tensor<1xi32> // CHECK: %[[ADD:.*]] = "tf.AddV2"(%[[ARG2]], %[[CONST1]]) // CHECK-NOT: "tf.TensorListPushBack" // CHECK: return %[[UPDATE]], %[[ADD]] @@ -520,7 +520,7 @@ func.func private @callee(%arg0: tensor>>, %arg1: t // CHECK-NOT: "tf.TensorListPushBack" // CHECK: %[[UPDATE:.*]] = "tf.XlaDynamicUpdateSlice" - // CHECK: %[[CONST1:.*]] = "tf.Const"() {value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32> + // CHECK: %[[CONST1:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}> : () -> tensor<1xi32> // CHECK: %[[ADD:.*]] = "tf.AddV2"(%[[ARG2]], %[[CONST1]]) // CHECK-NOT: "tf.TensorListPushBack" %push = "tf.TensorListPushBack"(%arg0, %elem) : (tensor>>, tensor) -> tensor>> diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tf-ops.mlir b/tensorflow/compiler/mlir/tensorflow/tests/tf-ops.mlir index 22de9ee736426a..7605f0360625fe 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/tf-ops.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/tf-ops.mlir @@ -1185,7 +1185,7 @@ func.func @testInvalidIfOp(tensor, tensor<*xf32>) -> tensor<2xf32> { // Test invalid tf.Yield operation (parent should be IfRegion) func.func @testInvalidYieldOp(%arg0: f32) -> () { - // expected-error @+1 {{'tf.Yield' op expects parent op to be one of 'tf.CaseRegion, tf.IfRegion, tf.WhileRegion'}} + // expected-error @+1 {{'tf.Yield' op expects parent op to be one of 'tf.CaseRegion, tf.IfRegion, tf.WhileRegion, tf.GeneratorDatasetRegion'}} "tf.Yield"(%arg0) : (f32) -> () } @@ -5180,3 +5180,41 @@ func.func @test_xla_call_module_with_invalid_symbol() { "tf.XlaCallModule"() {Sout = [], device = "", dim_args_spec = [], function_list = [@undefined_function], module = "", platforms = [], version = 4 : i64} : () -> () func.return } + +// ----- + +func.func @init(%arg0: tensor<4xf32>) -> tensor<7xf32> { + %0 = builtin.unrealized_conversion_cast to tensor<7xf32> + return %0 : tensor<7xf32> +} + +func.func @next(%arg0: tensor<7xf32>, %arg1: tensor<3xf32>) -> tensor<6xf32> { + %0 = builtin.unrealized_conversion_cast to tensor<6xf32> + return %0 : tensor<6xf32> +} + +func.func @finalize(%arg0: tensor<6xf32>, %arg1: tensor<2xf32>) -> tensor<5xf32> { + %0 = builtin.unrealized_conversion_cast to tensor<5xf32> + return %0 : tensor<5xf32> +} + +// CHECK-LABEL: func @testGeneratorDataset +func.func @testGeneratorDataset(%arg0: tensor<4xf32>, + %arg1: tensor<3xf32>, + %arg2: tensor, + %arg3: tensor<2xf32>) -> tensor { + %0 = "tf.GeneratorDataset"(%arg0, %arg1, %arg2, %arg3) { + device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0", + finalize_func = @finalize, + init_func = @init, + next_func = @next, + operandSegmentSizes = array, + output_shapes = [#tf_type.shape<>], + output_types = [!tf_type.string], + metadata = ""} : ( + tensor<4xf32>, + tensor<3xf32>, + tensor, + tensor<2xf32>) -> tensor + return %0 : tensor +} diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tf_data_fuse_map_and_batch.mlir b/tensorflow/compiler/mlir/tensorflow/tests/tf_data_fuse_map_and_batch.mlir index af605e92aed4e8..c6f63db25acc68 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/tf_data_fuse_map_and_batch.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/tf_data_fuse_map_and_batch.mlir @@ -6,7 +6,7 @@ func.func @fuse_map_and_batch() -> tensor attributes {tf.entry %0 = "tf.Const"() {value = dense<5> : tensor} : () -> tensor %1 = "tf.Const"() {value = dense : tensor} : () -> tensor %2 = "tf.Const"() {value = dense<[0, 1, 2]> : tensor<3xi32>} : () -> tensor<3xi32> - // CHECK: %[[NPC:.*]] = "tf.Const"() {value = dense<1> : tensor} + // CHECK: %[[NPC:.*]] = "tf.Const"() <{value = dense<1> : tensor}> // CHECK: %[[TSLICE:.*]] = "tf.TensorSliceDataset" %3 = "tf.TensorSliceDataset"(%2) {device = "", output_shapes = [#tf_type.shape<>], metadata = ""} : (tensor<3xi32>) -> tensor<*x!tf_type.variant> // CHECK: "tf.MapAndBatchDataset"(%[[TSLICE]], %[[BSIZE:.*]], %[[NPC]] diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tf_optimize.mlir b/tensorflow/compiler/mlir/tensorflow/tests/tf_optimize.mlir index 3dacd11a31b0bc..d0a5a74edc5b14 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/tf_optimize.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/tf_optimize.mlir @@ -12,7 +12,7 @@ func.func @fuseMulIntoConv2d(%arg0: tensor<1x112x112x3xf32>) -> tensor<1x28x23x2 // CHECK-SAME: [1.000000e+00, 4.000000e+00], [3.000000e+00, 8.000000e+00], [5.000000e+00, 1.200000e+01] // CHECK-SAME: [7.000000e+00, 1.600000e+01], [9.000000e+00, 2.000000e+01], [1.100000e+01, 2.400000e+01] // CHECK-SAME: [1.300000e+01, 2.800000e+01], [1.500000e+01, 3.200000e+01], [1.700000e+01, 3.600000e+01] - // CHECK: %[[CONV:.*]] = "tf.Conv2D"(%arg0, %[[CST]]) {data_format = "NHWC", dilations = [1, 2, 3, 1], explicit_paddings = [], padding = "SAME", strides = [1, 4, 5, 1], use_cudnn_on_gpu = true} + // CHECK: %[[CONV:.*]] = "tf.Conv2D"(%arg0, %[[CST]]) <{data_format = "NHWC", dilations = [1, 2, 3, 1], explicit_paddings = [], padding = "SAME", strides = [1, 4, 5, 1], use_cudnn_on_gpu = true}> // CHECK: return %[[CONV]] : tensor<1x28x23x2xf32> } @@ -26,7 +26,7 @@ func.func @notfuseMulIntoConv2d(%arg0: tensor<1x112x112x3xf32>) -> tensor<1x28x2 func.return %1 : tensor<1x28x23x2xf32> // CHECK: %cst_0 = arith.constant dense<3.000000e+00> : tensor<23x2xf32> - // CHECK: %0 = "tf.Conv2D"(%arg0, %cst) {T = "tfdtype$DT_FLOAT", data_format = "NHWC", dilations = [1, 2, 3, 1], padding = "SAME", strides = [1, 4, 5, 1]} + // CHECK: %0 = "tf.Conv2D"(%arg0, %cst) <{data_format = "NHWC", dilations = [1, 2, 3, 1], padding = "SAME", strides = [1, 4, 5, 1]}> {T = "tfdtype$DT_FLOAT"} // CHECK: %1 = "tf.Mul"(%0, %cst_0) : (tensor<1x28x23x2xf32>, tensor<23x2xf32>) -> tensor<1x28x23x2xf32> // CHECK: return %1 : tensor<1x28x23x2xf32> } @@ -40,8 +40,8 @@ func.func @simplifyBroadcastReshape(%arg0: tensor<1x8x1x1x1x1x1x18xbf16>) -> ten %98 = "tf.Reshape"(%97, %cst_2) : (tensor<1x8x6x1x6x1x1x18xbf16>, tensor<4xi64>) -> tensor<8x6x6x18xbf16> func.return %98 : tensor<8x6x6x18xbf16> - // CHECK-DAG: %[[CST:.*]] = "tf.Const"() {value = dense<[8, 1, 1, 18]> : tensor<4xi64>} : () -> tensor<4xi64> - // CHECK-DAG: %[[CST1:.*]] = "tf.Const"() {value = dense<[8, 6, 6, 18]> : tensor<4xi64>} : () -> tensor<4xi64> + // CHECK-DAG: %[[CST:.*]] = "tf.Const"() <{value = dense<[8, 1, 1, 18]> : tensor<4xi64>}> : () -> tensor<4xi64> + // CHECK-DAG: %[[CST1:.*]] = "tf.Const"() <{value = dense<[8, 6, 6, 18]> : tensor<4xi64>}> : () -> tensor<4xi64> // CHECK: %[[RESHAPE:.*]] = "tf.Reshape"(%arg0, %[[CST]]) : (tensor<1x8x1x1x1x1x1x18xbf16>, tensor<4xi64>) -> tensor<8x1x1x18xbf16> // CHECK: %[[BROADCAST:.*]] = "tf.BroadcastTo"(%[[RESHAPE]], %[[CST1]]) : (tensor<8x1x1x18xbf16>, tensor<4xi64>) -> tensor<8x6x6x18xbf16> // CHECK: return %[[BROADCAST]] : tensor<8x6x6x18xbf16> @@ -55,8 +55,8 @@ func.func @simplifyBroadcastReshapeExtraDims(%arg0: tensor<1x8x1x1x1x1x1x18xbf16 %98 = "tf.Reshape"(%97, %cst_2) : (tensor<7x1x8x6x1x6x1x1x18xbf16>, tensor<5xi64>) -> tensor<7x8x6x6x18xbf16> func.return %98 : tensor<7x8x6x6x18xbf16> - // CHECK-DAG: %[[CST:.*]] = "tf.Const"() {value = dense<[1, 8, 1, 1, 18]> : tensor<5xi64>} : () -> tensor<5xi64> - // CHECK-DAG: %[[CST1:.*]] = "tf.Const"() {value = dense<[7, 8, 6, 6, 18]> : tensor<5xi64>} : () -> tensor<5xi64> + // CHECK-DAG: %[[CST:.*]] = "tf.Const"() <{value = dense<[1, 8, 1, 1, 18]> : tensor<5xi64>}> : () -> tensor<5xi64> + // CHECK-DAG: %[[CST1:.*]] = "tf.Const"() <{value = dense<[7, 8, 6, 6, 18]> : tensor<5xi64>}> : () -> tensor<5xi64> // CHECK: %[[RESHAPE:.*]] = "tf.Reshape"(%arg0, %[[CST]]) : (tensor<1x8x1x1x1x1x1x18xbf16>, tensor<5xi64>) -> tensor<1x8x1x1x18xbf16> // CHECK: %[[BROADCAST:.*]] = "tf.BroadcastTo"(%[[RESHAPE]], %[[CST1]]) : (tensor<1x8x1x1x18xbf16>, tensor<5xi64>) -> tensor<7x8x6x6x18xbf16> // CHECK: return %[[BROADCAST]] : tensor<7x8x6x6x18xbf16> @@ -70,8 +70,8 @@ func.func @simplifyBroadcastReshapeOnes(%arg0: tensor<1x1x1x1x1x1x1x18xbf16>) -> %98 = "tf.Reshape"(%97, %cst_2) : (tensor<1x1x6x1x6x1x1x18xbf16>, tensor<5xi64>) -> tensor<1x6x1x6x18xbf16> func.return %98 : tensor<1x6x1x6x18xbf16> - // CHECK-DAG: %[[CST:.*]] = "tf.Const"() {value = dense<[1, 1, 1, 1, 18]> : tensor<5xi64>} : () -> tensor<5xi64> - // CHECK-DAG: %[[CST1:.*]] = "tf.Const"() {value = dense<[1, 6, 1, 6, 18]> : tensor<5xi64>} : () -> tensor<5xi64> + // CHECK-DAG: %[[CST:.*]] = "tf.Const"() <{value = dense<[1, 1, 1, 1, 18]> : tensor<5xi64>}> : () -> tensor<5xi64> + // CHECK-DAG: %[[CST1:.*]] = "tf.Const"() <{value = dense<[1, 6, 1, 6, 18]> : tensor<5xi64>}> : () -> tensor<5xi64> // CHECK: %[[RESHAPE:.*]] = "tf.Reshape"(%arg0, %[[CST]]) : (tensor<1x1x1x1x1x1x1x18xbf16>, tensor<5xi64>) -> tensor<1x1x1x1x18xbf16> // CHECK: %[[BROADCAST:.*]] = "tf.BroadcastTo"(%[[RESHAPE]], %[[CST1]]) : (tensor<1x1x1x1x18xbf16>, tensor<5xi64>) -> tensor<1x6x1x6x18xbf16> // CHECK: return %[[BROADCAST]] : tensor<1x6x1x6x18xbf16> diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/basic_v1.py b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/basic_v1.py index 9b71d46edf5fda..2d620c8df759ce 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/basic_v1.py +++ b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/basic_v1.py @@ -38,7 +38,7 @@ # CHECK-SAME: attributes {{.*}} tf_saved_model.exported_names = ["key"] # CHECK-NEXT: [[R0:%.*]] = "tf.ReadVariableOp"([[ARG1]]) {{{.*}}} : (tensor>>) -> tensor<1x3xf32> -# CHECK-NEXT: [[R1:%.*]] = "tf.MatMul"([[ARG0]], [[R0]]) {{{.*}}} : (tensor<3x1xf32>, tensor<1x3xf32>) -> tensor<3x3xf32> +# CHECK-NEXT: [[R1:%.*]] = "tf.MatMul"([[ARG0]], [[R0]]) <{{{.*}}}> {device = ""} : (tensor<3x1xf32>, tensor<1x3xf32>) -> tensor<3x3xf32> # CHECK-NEXT: return [[R1]] : tensor<3x3xf32> diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/hash_table_v1.py b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/hash_table_v1.py index b3cd46c2d60517..744f302dac610b 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/hash_table_v1.py +++ b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/hash_table_v1.py @@ -35,7 +35,7 @@ # CHECK: func @[[init]] # CHECK-SAME: tf_saved_model.initializer_type = "init_op" # CHECK-NEXT: [[R6:%.*]] = "tf.Const"() -# CHECK-NEXT: [[R5:%.*]] = "tf.Const"() {device = "", value = dense<[1, 2, +# CHECK-NEXT: [[R5:%.*]] = "tf.Const"() <{value = dense<[1, 2, # CHECK-NEXT: [[R7:%.*]] = "tf.HashTableV2"() # CHECK-SAME: shared_name = "[[hash_table:.*]]" # CHECK-NEXT: "tf.LookupTableImportV2"([[R7]], [[R5]], [[R6]]) diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/include_variables_in_init_v1.py b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/include_variables_in_init_v1.py index 2f99fae8d8d286..bc8457aab0f63c 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/include_variables_in_init_v1.py +++ b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/include_variables_in_init_v1.py @@ -38,7 +38,7 @@ # CHECK-SAME: tf_saved_model.exported_names = ["__tf_saved_model_session_initializer_init"] # CHECK-SAME: tf_saved_model.initializer_type = "init_op" # CHECK-SAME: } -# CHECK-DAG: %[[CST_0:.*]] = "tf.Const"() {{{.*dense<.*> : tensor<2xi32>.*}}} : () -> tensor<2xi32> +# CHECK-DAG: %[[CST_0:.*]] = "tf.Const"() <{{{.*dense<.*> : tensor<2xi32>.*}}}> {{{.*}}} : () -> tensor<2xi32> # CHECK: %[[RAND_STD_NORMAL:.*]] = "tf.RandomStandardNormal"(%[[CST_0]]) # CHECK: "tf.AssignVariableOp"(%[[ARG_0]], %[[RAND_STD_NORMAL]]){{.*}}: (tensor>>, tensor<1x3xf32>) -> () # CHECK: return @@ -50,7 +50,7 @@ # CHECK-SAME: -> (tensor<3x3xf32> {tf_saved_model.index_path = ["r"]}) # CHECK-SAME: attributes {{.*}} tf_saved_model.exported_names = ["key"] # CHECK-NEXT: %[[READ_VAR_0:.*]] = "tf.ReadVariableOp"(%[[ARG_2]]) {{{.*}}} : (tensor>>) -> tensor<1x3xf32> -# CHECK-NEXT: %[[MATMUL_0:.*]] = "tf.MatMul"(%[[ARG_1]], %[[READ_VAR_0]]) {{{.*}}} : (tensor<3x1xf32>, tensor<1x3xf32>) -> tensor<3x3xf32> +# CHECK-NEXT: %[[MATMUL_0:.*]] = "tf.MatMul"(%[[ARG_1]], %[[READ_VAR_0]]) <{{{.*}}}> {{{.*}}} : (tensor<3x1xf32>, tensor<1x3xf32>) -> tensor<3x3xf32> # CHECK-NEXT: return %[[MATMUL_0]] : tensor<3x3xf32> diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/multi_variables_v1.py b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/multi_variables_v1.py index 4ce47a3aec61ef..6eb8cbe20b0fbf 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/multi_variables_v1.py +++ b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/multi_variables_v1.py @@ -29,7 +29,7 @@ # CHECK-NEXT: [[R0:%.*]] = "tf.ReadVariableOp"([[ARG0]]) {{{.*}}} : (tensor>>) -> tensor<5x3xf32> # CHECK-NEXT: [[R1:%.*]] = "tf.ReadVariableOp"([[ARG1]]) {{{.*}}} : (tensor>>) -> tensor<3x5xf32> -# CHECK-NEXT: [[R2:%.*]] = "tf.MatMul"([[R0]], [[R1]]) {{{.*}}} : (tensor<5x3xf32>, tensor<3x5xf32>) -> tensor<5x5xf32> +# CHECK-NEXT: [[R2:%.*]] = "tf.MatMul"([[R0]], [[R1]]) <{{{.*}}}> {{{.*}}} : (tensor<5x3xf32>, tensor<3x5xf32>) -> tensor<5x5xf32> def Test(): diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/remove_init_variable_v1.py b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/remove_init_variable_v1.py index 123a11a1883521..261ff5577eb209 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/remove_init_variable_v1.py +++ b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/remove_init_variable_v1.py @@ -39,7 +39,7 @@ # CHECK-SAME: attributes {{.*}} tf_saved_model.exported_names = ["key"] # CHECK-NEXT: [[R0:%.*]] = "tf.ReadVariableOp"([[ARG1]]) {{{.*}}} : (tensor>>) -> tensor<1x3xf32> -# CHECK-NEXT: [[R1:%.*]] = "tf.MatMul"([[ARG0]], [[R0]]) {{{.*}}} : (tensor<3x1xf32>, tensor<1x3xf32>) -> tensor<3x3xf32> +# CHECK-NEXT: [[R1:%.*]] = "tf.MatMul"([[ARG0]], [[R0]]) <{{{.*}}}> {{{.*}}} : (tensor<3x1xf32>, tensor<1x3xf32>) -> tensor<3x3xf32> # CHECK-NEXT: return [[R1]] : tensor<3x3xf32> diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_asset_sinking.mlir b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_asset_sinking.mlir index 2638aab86b890b..94e0da360795af 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_asset_sinking.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_asset_sinking.mlir @@ -10,8 +10,8 @@ module @asset attributes {tf_saved_model.semantics} { // CHECK: func @init() func.func @init(%arg0: tensor {tf_saved_model.bound_input = @asset0}, %arg1: tensor {tf_saved_model.bound_input = @asset1}) attributes {tf_saved_model.exported_names = ["init"]} { - // CHECK-DAG: %[[ASSET0:.*]] = "tf.Const"() {value = dense<"foo/bar/assets/test0.txt"> : tensor} - // CHECK-DAG: %[[ASSET1:.*]] = "tf.Const"() {value = dense<"foo/bar/assets/test1.txt"> : tensor} + // CHECK-DAG: %[[ASSET0:.*]] = "tf.Const"() <{value = dense<"foo/bar/assets/test0.txt"> : tensor}> + // CHECK-DAG: %[[ASSET1:.*]] = "tf.Const"() <{value = dense<"foo/bar/assets/test1.txt"> : tensor}> // CHECK: %[[VAR0:.*]] = "tf.VarHandleOp"() %0 = "tf.VarHandleOp"() {container = "", shared_name = "var0"} : () -> tensor>> diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_freeze_assets.mlir b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_freeze_assets.mlir index eb4aed85564dab..982ace4360ae98 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_freeze_assets.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_freeze_assets.mlir @@ -11,7 +11,7 @@ module attributes {tf_saved_model.semantics} { attributes {tf_saved_model.exported_names = ["f"]} { %0 = "tf.HashTableV2"() {container = "", device = "", key_dtype = !tf_type.string, shared_name = "", use_node_name_sharing = false, value_dtype = i64} : () -> tensor "tf.InitializeTableFromTextFileV2"(%0, %arg0) {delimiter = "\09", device = "", key_index = -2 : i64, offset = 0 : i64, value_index = -1 : i64, vocab_size = 437 : i64} : (tensor, tensor) -> () - // CHECK: [[CST:%.+]] = "tf.Const"() {value = dense<"assets/table.txt"> : tensor<1x!tf_type.string>} : () -> tensor<1x!tf_type.string> + // CHECK: [[CST:%.+]] = "tf.Const"() <{value = dense<"assets/table.txt"> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string> // CHECK: [[HASHTABLE:%.+]] = "tf.HashTableV2"() // CHECK: "tf.InitializeTableFromTextFileV2"([[HASHTABLE]], [[CST]]) func.return @@ -69,8 +69,8 @@ module attributes {tf_saved_model.semantics} { "tf.InitializeTableFromTextFileV2"(%0, %arg0) {delimiter = "\09", device = "", key_index = -2 : i64, offset = 0 : i64, value_index = -1 : i64, vocab_size = 437 : i64} : (tensor, tensor) -> () %1 = "tf.HashTableV2"() {container = "", device = "", key_dtype = !tf_type.string, shared_name = "", use_node_name_sharing = false, value_dtype = i64} : () -> tensor "tf.InitializeTableFromTextFileV2"(%1, %arg1) {delimiter = "\09", device = "", key_index = -2 : i64, offset = 0 : i64, value_index = -1 : i64, vocab_size = 437 : i64} : (tensor, tensor) -> () - // CHECK-DAG: [[CST_1:%.+]] = "tf.Const"() {value = dense<"assets/table2.txt"> : tensor<1x!tf_type.string>} : () -> tensor<1x!tf_type.string> - // CHECK-DAG: [[CST:%.+]] = "tf.Const"() {value = dense<"assets/table.txt"> : tensor<1x!tf_type.string>} : () -> tensor<1x!tf_type.string> + // CHECK-DAG: [[CST_1:%.+]] = "tf.Const"() <{value = dense<"assets/table2.txt"> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string> + // CHECK-DAG: [[CST:%.+]] = "tf.Const"() <{value = dense<"assets/table.txt"> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string> // CHECK: [[HASHTABLE:%.+]] = "tf.HashTableV2"() // CHECK: "tf.InitializeTableFromTextFileV2"([[HASHTABLE]], [[CST]]) // CHECK: [[HASHTABLE_1:%.+]] = "tf.HashTableV2"() diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_freeze_global_tensors.mlir b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_freeze_global_tensors.mlir index 67a5439524f8e3..9c0f9b2eddb29b 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_freeze_global_tensors.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_freeze_global_tensors.mlir @@ -12,7 +12,7 @@ module attributes {tf_saved_model.semantics} { func.func @f(%arg0: tensor>> {tf_saved_model.bound_input = @v}) attributes {tf_saved_model.exported_names = ["f"]} { %val = "tf.ReadVariableOp"(%arg0) : (tensor>>) -> tensor - // CHECK: "tf.Const"() {value = dense<1.000000e+00> : tensor} + // CHECK: "tf.Const"() <{value = dense<1.000000e+00> : tensor}> func.return } } @@ -67,7 +67,7 @@ module attributes {tf_saved_model.semantics} { } func.func private @f_callee(%arg0: tensor>>) { - // CHECK: "tf.Const"() {value = dense<2.100000e+01> : tensor} + // CHECK: "tf.Const"() <{value = dense<2.100000e+01> : tensor}> func.return } } @@ -90,7 +90,7 @@ module attributes {tf_saved_model.semantics} { func.func private @g_callee(%arg0: tensor>>) { %val = "tf.ReadVariableOp"(%arg0) : (tensor>>) -> tensor - // CHECK: "tf.Const"() {value = dense<3.200000e+01> : tensor} + // CHECK: "tf.Const"() <{value = dense<3.200000e+01> : tensor}> func.return } } @@ -146,10 +146,10 @@ module attributes {tf_saved_model.semantics} { func.func @f(%arg1: tensor>> {tf_saved_model.bound_input = @"v"}, %arg2: tensor>> {tf_saved_model.bound_input = @"v2"}) attributes {tf_saved_model.exported_names = ["f"]} { - // CHECK-DAG: "tf.Const"() {value = dense<1.000000e+00> : tensor} + // CHECK-DAG: "tf.Const"() <{value = dense<1.000000e+00> : tensor}> %0 = "tf.ReadVariableOp"(%arg1) {device = ""} : (tensor>>) -> tensor - // CHECK-DAG: "tf.Const"() {value = dense<2.000000e+00> : tensor} + // CHECK-DAG: "tf.Const"() <{value = dense<2.000000e+00> : tensor}> %1 = "tf.ReadVariableOp"(%arg2) {device = ""} : (tensor>>) -> tensor func.return } diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_mark_initialized_variables.mlir b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_mark_initialized_variables.mlir index 51368f29fa76bd..4f6967990eb162 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_mark_initialized_variables.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model_mark_initialized_variables.mlir @@ -15,14 +15,14 @@ module attributes {tf_saved_model.semantics, tf_saved_model.under_construction} func.return %4 : tensor<100x50xf32> } // CHECK: "tf.VarHandleOp" - // CHECK-SAME: _is_initialized = true // CHECK-SAME: shared_name = "var1" - // CHECK: "tf.VarHandleOp" // CHECK-SAME: _is_initialized = true + // CHECK: "tf.VarHandleOp" // CHECK-SAME: shared_name = "var2" + // CHECK-SAME: _is_initialized = true // CHECK: "tf.VarHandleOp" - // CHECK-SAME: _is_initialized = false // CHECK-SAME: shared_name = "var3" + // CHECK-SAME: _is_initialized = false // INVALID-NOT: _is_initialized } diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tpu-dynamic-layout-pass.mlir b/tensorflow/compiler/mlir/tensorflow/tests/tpu-dynamic-layout-pass.mlir index 6d655ab4ffc95a..553bfa0955106b 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/tpu-dynamic-layout-pass.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/tpu-dynamic-layout-pass.mlir @@ -14,8 +14,8 @@ func.func @non_replicated(%arg0: tensor<*x!tf_type.resource> {tf.device = "/devi mlir_module = "..."} : () -> (tensor, tensor<2x!tf_type.string>) tf_device.return %1#0, %1#1 : tensor, tensor<2x!tf_type.string> }) {device = "/device:CPU:0"} : () -> (tensor, tensor<2x!tf_type.string>) - // CHECK-DAG: %[[LAYOUT0:.*]] = "tf.TPUGetLayoutOp"(%[[COMPILE]]#1) {index = 0 : i64, is_output = false} - // CHECK-DAG: %[[LAYOUT1:.*]] = "tf.TPUGetLayoutOp"(%[[COMPILE]]#1) {index = 1 : i64, is_output = false} + // CHECK-DAG: %[[LAYOUT0:.*]] = "tf.TPUGetLayoutOp"(%[[COMPILE]]#1) <{index = 0 : i64, is_output = false}> + // CHECK-DAG: %[[LAYOUT1:.*]] = "tf.TPUGetLayoutOp"(%[[COMPILE]]#1) <{index = 1 : i64, is_output = false}> // CHECK: %[[ITER:.*]]:2 = "tf.IteratorGetNext" %2:2 = "tf.IteratorGetNext"(%arg0) {device = "/device:CPU:0"} : (tensor<*x!tf_type.resource>) -> (tensor<3x3x1x32xf32>, tensor<3x3x1x32xf32>) @@ -249,8 +249,8 @@ func.func @replicated(%arg0: tensor<*x!tf_type.resource> {tf.device = "/device:C mlir_module = "..."} : () -> (tensor, tensor<2x!tf_type.string>) tf_device.return %1#0, %1#1 : tensor, tensor<2x!tf_type.string> }) {device = "/device:CPU:0"} : () -> (tensor, tensor<2x!tf_type.string>) - // CHECK-DAG: %[[LAYOUT0:.*]] = "tf.TPUGetLayoutOp"(%[[COMPILE]]#1) {index = 0 : i64, is_output = false} - // CHECK-DAG: %[[LAYOUT1:.*]] = "tf.TPUGetLayoutOp"(%[[COMPILE]]#1) {index = 1 : i64, is_output = false} + // CHECK-DAG: %[[LAYOUT0:.*]] = "tf.TPUGetLayoutOp"(%[[COMPILE]]#1) <{index = 0 : i64, is_output = false}> + // CHECK-DAG: %[[LAYOUT1:.*]] = "tf.TPUGetLayoutOp"(%[[COMPILE]]#1) <{index = 1 : i64, is_output = false}> // CHECK: %[[ITER1:.*]]:2 = "tf.IteratorGetNext" %3:2 = "tf.IteratorGetNext"(%arg0) {device = "/device:CPU:0"} : (tensor<*x!tf_type.resource>) -> (tensor<3x3x1x32xf32>, tensor<3x3x1x32xf32>) @@ -294,8 +294,8 @@ func.func @replicated_packed(%arg0: tensor<*x!tf_type.resource> {tf.device = "/d mlir_module = "..."} : () -> (tensor, tensor<2x!tf_type.string>) tf_device.return %1#0, %1#1 : tensor, tensor<2x!tf_type.string> }) {device = "/device:CPU:0"} : () -> (tensor, tensor<2x!tf_type.string>) - // CHECK-DAG: %[[LAYOUT0:.*]] = "tf.TPUGetLayoutOp"(%[[COMPILE]]#1) {index = 0 : i64, is_output = false} - // CHECK-DAG: %[[LAYOUT1:.*]] = "tf.TPUGetLayoutOp"(%[[COMPILE]]#1) {index = 1 : i64, is_output = false} + // CHECK-DAG: %[[LAYOUT0:.*]] = "tf.TPUGetLayoutOp"(%[[COMPILE]]#1) <{index = 0 : i64, is_output = false}> + // CHECK-DAG: %[[LAYOUT1:.*]] = "tf.TPUGetLayoutOp"(%[[COMPILE]]#1) <{index = 1 : i64, is_output = false}> // CHECK-DAG: %[[COPY0:.*]] = "tf.TPUCopyWithLayout"(%[[ITER0]]#0, %[[LAYOUT0]]) {device = "/device:TPU:0"} // CHECK-DAG: %[[COPY1:.*]] = "tf.TPUCopyWithLayout"(%[[ITER0]]#1, %[[LAYOUT1]]) {device = "/device:TPU:0"} @@ -332,8 +332,8 @@ func.func @replicated(%arg0: tensor<*x!tf_type.resource> {tf.device = "/device:C mlir_module = "..."} : () -> (tensor, tensor<2x!tf_type.string>) tf_device.return %1#0, %1#1 : tensor, tensor<2x!tf_type.string> }) {device = "/device:CPU:0"} : () -> (tensor, tensor<2x!tf_type.string>) - // CHECK-DAG: %[[LAYOUT0:.*]] = "tf.TPUGetLayoutOp"(%[[COMPILE]]#1) {index = 0 : i64, is_output = false} - // CHECK-DAG: %[[LAYOUT1:.*]] = "tf.TPUGetLayoutOp"(%[[COMPILE]]#1) {index = 1 : i64, is_output = false} + // CHECK-DAG: %[[LAYOUT0:.*]] = "tf.TPUGetLayoutOp"(%[[COMPILE]]#1) <{index = 0 : i64, is_output = false}> + // CHECK-DAG: %[[LAYOUT1:.*]] = "tf.TPUGetLayoutOp"(%[[COMPILE]]#1) <{index = 1 : i64, is_output = false}> // CHECK: %[[ITER1:.*]] = "tf.IteratorGetNext" %3 = "tf.IteratorGetNext"(%arg1) {device = "/device:CPU:0"} : (tensor<*x!tf_type.resource>) -> tensor<3x3x1x32xf32> @@ -415,8 +415,8 @@ func.func @parallel_execute(%arg0: tensor<*x!tf_type.resource> {tf.device = "/de %1:3 = "tf._TPUCompileMlir"() {NumDynamicShapes = 0 : i64, metadata = "\0A\09\08\01\12\05\12\03\08\80\01\18\01 \02", mlir_module = "..."} : () -> (tensor, tensor<2x!tf_type.string>, tensor<2x!tf_type.string>) tf_device.return %1#0, %1#1, %1#2 : tensor, tensor<2x!tf_type.string>, tensor<2x!tf_type.string> }) {device = "/device:CPU:0"} : () -> (tensor, tensor<2x!tf_type.string>, tensor<2x!tf_type.string>) - // CHECK-DAG: %[[LAYOUT0:.*]] = "tf.TPUGetLayoutOp"(%[[COMPILE]]#1) {index = 0 : i64, is_output = false} - // CHECK-DAG: %[[LAYOUT1:.*]] = "tf.TPUGetLayoutOp"(%[[COMPILE]]#2) {index = 0 : i64, is_output = false} + // CHECK-DAG: %[[LAYOUT0:.*]] = "tf.TPUGetLayoutOp"(%[[COMPILE]]#1) <{index = 0 : i64, is_output = false}> + // CHECK-DAG: %[[LAYOUT1:.*]] = "tf.TPUGetLayoutOp"(%[[COMPILE]]#2) <{index = 0 : i64, is_output = false}> // CHECK: %[[ITER:.*]]:2 = "tf.IteratorGetNext" %2:2 = "tf.IteratorGetNext"(%arg0) {device = "/device:CPU:0"} : (tensor<*x!tf_type.resource>) -> (tensor<128xf32>, tensor<128xf32>) // CHECK: "tf.TPUCompileSucceededAssert"(%[[COMPILE]]#0) @@ -481,8 +481,8 @@ func.func @replicated_parallel_execute(%arg0: tensor<*x!tf_type.resource> {tf.de %1:3 = "tf._TPUCompileMlir"() {NumDynamicShapes = 0 : i64, metadata = "\0A\09\08\01\12\05\12\03\08\80\01\18\02 \02", mlir_module = "..."} : () -> (tensor, tensor<2x!tf_type.string>, tensor<2x!tf_type.string>) tf_device.return %1#0, %1#1, %1#2 : tensor, tensor<2x!tf_type.string>, tensor<2x!tf_type.string> }) {device = "/device:CPU:0"} : () -> (tensor, tensor<2x!tf_type.string>, tensor<2x!tf_type.string>) - // CHECK-DAG: %[[LAYOUT0:.*]] = "tf.TPUGetLayoutOp"(%[[COMPILE]]#1) {index = 0 : i64, is_output = false} - // CHECK-DAG: %[[LAYOUT1:.*]] = "tf.TPUGetLayoutOp"(%[[COMPILE]]#2) {index = 0 : i64, is_output = false} + // CHECK-DAG: %[[LAYOUT0:.*]] = "tf.TPUGetLayoutOp"(%[[COMPILE]]#1) <{index = 0 : i64, is_output = false}> + // CHECK-DAG: %[[LAYOUT1:.*]] = "tf.TPUGetLayoutOp"(%[[COMPILE]]#2) <{index = 0 : i64, is_output = false}> // CHECK-DAG: %[[ITER0:.*]]:2 = "tf.IteratorGetNext"(%[[ARG0]]) // CHECK-DAG: %[[ITER1:.*]]:2 = "tf.IteratorGetNext"(%[[ARG1]]) %2:2 = "tf.IteratorGetNext"(%arg0) {device = "/device:CPU:0"} : (tensor<*x!tf_type.resource>) -> (tensor<128xf32>, tensor<128xf32>) diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tpu-multiple-while-body-func.mlir b/tensorflow/compiler/mlir/tensorflow/tests/tpu-multiple-while-body-func.mlir index 1762e9f7889fd4..9746ca6919a8fe 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/tpu-multiple-while-body-func.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/tpu-multiple-while-body-func.mlir @@ -1,5 +1,5 @@ // RUN: tf-opt %s -tf-tpu-bridge 2>&1 | FileCheck %s -// RUN: tf-opt %s -tf-cluster-tpu-bridge-v1 2>&1 | FileCheck %s +// RUN: tf-opt %s -tf-cluster-tpu-bridge-v1 -tfrt-lower-cluster-to-runtime-ops-tpu -tf-dialect-to-executor-v1 2>&1 | FileCheck %s // This test verifies there is no warning about shape inference failure in TPU // bridge in handling multiple usage of the same function. diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tpu_cluster_formation.mlir b/tensorflow/compiler/mlir/tensorflow/tests/tpu_cluster_formation.mlir index ed4375091167a4..1d3c1b6f3cf518 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/tpu_cluster_formation.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/tpu_cluster_formation.mlir @@ -942,7 +942,7 @@ func.func @missing_metadata() { // CHECK-LABEL: func @const_with_attrs func.func @const_with_attrs(%arg0: tensor<*xi32>, %arg1: tensor) -> (tensor, tensor) { - // CHECK: %{{[a-z0-9_]*}} = "tf.Const"() {value = dense<-1> : tensor<1xi32>} : () -> tensor<1xi32> + // CHECK: %{{[a-z0-9_]*}} = "tf.Const"() <{value = dense<-1> : tensor<1xi32>}> : () -> tensor<1xi32> // CHECK-NEXT: %{{[a-z0-9_]*}} = "tf.Reshape"(%arg0 // CHECK-NEXT: %{{.*}} = "tf_device.cluster"() ({ %minus_one = "tf.Const"() {_replication_info = "cluster", diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tpu_space_to_depth_pass.mlir b/tensorflow/compiler/mlir/tensorflow/tests/tpu_space_to_depth_pass.mlir index 2125f2877fb8cf..4d290b71fca1c0 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/tpu_space_to_depth_pass.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/tpu_space_to_depth_pass.mlir @@ -15,7 +15,7 @@ module attributes {tf.devices = {"/job:localhost/replica:0/task:0/device:CPU:0" %0 = "tf.Const"() {value = dense<1> : tensor} : () -> tensor // CHECK: %[[INPUT:.*]] = "tf.IteratorGetNext" %1 = "tf.IteratorGetNext"(%arg5) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor) -> tensor<2x224x224x3xf32> - // CHECK-DAG: %[[SPACETODEPTH0:.*]] = "tf.SpaceToDepth"([[INPUT:.*]]) {block_size = 2 : i64, data_format = "NHWC"} : (tensor<2x224x224x3xf32>) -> tensor<2x112x112x12xf32> + // CHECK-DAG: %[[SPACETODEPTH0:.*]] = "tf.SpaceToDepth"([[INPUT:.*]]) <{block_size = 2 : i64, data_format = "NHWC"}> : (tensor<2x224x224x3xf32>) -> tensor<2x112x112x12xf32> %2 = "tf.AddV2"(%arg2, %arg3) {device = ""} : (tensor, tensor) -> tensor %3 = "tf.ReadVariableOp"(%arg6) : (tensor>>) -> tensor<7x7x3x64xf32> %4 = "tf.ReadVariableOp"(%arg8) : (tensor>>) -> tensor @@ -61,18 +61,18 @@ module attributes {tf.devices = {"/job:localhost/replica:0/task:0/device:CPU:0" // CHECK-SAME: strides = [1, 1, 1, 1] // CHECK-SAME: (tensor<2x115x115x12xf32>, tensor<4xi32>, tensor<2x112x112x64xf32>) -> tensor<4x4x12x64xf32> %7 = "tf.Conv2DBackpropFilter"(%5, %2, %6) {data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "VALID", strides = [1, 2, 2, 1], use_cudnn_on_gpu = true} : (tensor<2x230x230x3xf32>, tensor<4xi32>, tensor<2x112x112x64xf32>) -> tensor<7x7x3x64xf32> - // CHECK: %[[CONST0:.*]] = "tf.Const"() {value = dense< + // CHECK: %[[CONST0:.*]] = "tf.Const"() <{value = dense< // CHECK-SAME: [4, 4, 2, 2, 3, 64] // CHECK: %[[RESHAPE0:.*]] = "tf.Reshape"(%[[BACKPROP:.*]], %[[CONST0:.*]]) : (tensor<4x4x12x64xf32>, tensor<6xi64>) -> tensor<4x4x2x2x3x64xf32> - // CHECK: %[[CONST1:.*]] = "tf.Const"() {value = dense< + // CHECK: %[[CONST1:.*]] = "tf.Const"() <{value = dense< // CHECK-SAME: [0, 2, 1, 3, 4, 5] // CHECK: %[[TRANSPOSE0:.*]] = "tf.Transpose"(%[[RESHAPE0:.*]], %[[CONST1:.*]]) : (tensor<4x4x2x2x3x64xf32>, tensor<6xi32>) -> tensor<4x2x4x2x3x64xf32> - // CHECK: %[[CONST2:.*]] = "tf.Const"() {value = dense< + // CHECK: %[[CONST2:.*]] = "tf.Const"() <{value = dense< // CHECK-SAME: [8, 8, 3, 64] // CHECK: %[[RESHAPE1:.*]] = "tf.Reshape"(%[[TRANSPOSE1:.*]], %[[CONST2:.*]]) : (tensor<4x2x4x2x3x64xf32>, tensor<4xi64>) -> tensor<8x8x3x64xf32> - // CHECK: %[[CONST3:.*]] = "tf.Const"() {value = dense< + // CHECK: %[[CONST3:.*]] = "tf.Const"() <{value = dense< // CHECK-SAME: [7, 7, 3, 64] - // CHECK: %[[CONST4:.*]] = "tf.Const"() {value = dense< + // CHECK: %[[CONST4:.*]] = "tf.Const"() <{value = dense< // CHECK-SAME: 0 // CHECK: %[[SLICE0:.*]] = "tf.Slice"(%[[RESHAPE1:.*]], %[[CONST4:.*]], %[[CONST3:.*]]) : (tensor<8x8x3x64xf32>, tensor<4xi64>, tensor<4xi32>) -> tensor<7x7x3x64xf32> %8 = "tf.CrossReplicaSum"(%7, %1) : (tensor<7x7x3x64xf32>, tensor<1x1xi32>) -> tensor<7x7x3x64xf32> @@ -90,10 +90,10 @@ module attributes {tf.devices = {"/job:localhost/replica:0/task:0/device:CPU:0" module attributes {tf.devices = {"/job:localhost/replica:0/task:0/device:COMPOSITE:0" = {}, "/job:localhost/replica:0/task:0/device:CPU:0" = {}, "/job:localhost/replica:0/task:0/device:TPU:0" = {}, "/job:localhost/replica:0/task:0/device:TPU:1" = {}, "/job:localhost/replica:0/task:0/device:TPU_SYSTEM:0" = {}}, tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 458 : i32}} { func.func @main(%arg0: tensor<*x!tf_type.resource> {tf._user_specified_name = "iterator", tf.device = "/job:localhost/replica:0/task:0/device:CPU:0"}, %arg1: tensor {tf._user_specified_name = "iterator", tf.device = "/job:localhost/replica:0/task:0/device:CPU:0"}, %arg2: tensor<*x!tf_type.resource> {tf._user_specified_name = "iterator", tf.device = "/job:localhost/replica:0/task:0/device:CPU:0"}, %arg3: tensor {tf._user_specified_name = "iterator", tf.device = "/job:localhost/replica:0/task:0/device:CPU:0"}, %arg4: tensor<*x!tf_type.resource> {tf._user_specified_name = "iterator", tf.device = "/job:localhost/replica:0/task:0/device:CPU:0"}, %arg5: tensor {tf._user_specified_name = "iterator", tf.device = "/job:localhost/replica:0/task:0/device:CPU:0"}, %arg6: tensor<*x!tf_type.resource>> {tf._composite_device = "/job:localhost/replica:0/task:0/device:COMPOSITE:0", tf.device = "/job:localhost/replica:0/task:0/device:COMPOSITE:0"}, %arg7: tensor<*x!tf_type.resource>> {tf._composite_device = "/job:localhost/replica:0/task:0/device:COMPOSITE:0", tf.device = "/job:localhost/replica:0/task:0/device:COMPOSITE:0"}, %arg8: tensor<*x!tf_type.resource>> {tf._composite_device = "/job:localhost/replica:0/task:0/device:COMPOSITE:0", tf.device = "/job:localhost/replica:0/task:0/device:COMPOSITE:0"}, %arg9: tensor<*x!tf_type.resource>> {tf._composite_device = "/job:localhost/replica:0/task:0/device:COMPOSITE:0", tf.device = "/job:localhost/replica:0/task:0/device:COMPOSITE:0"}, %arg10: tensor<*x!tf_type.resource>> {tf._composite_device = "/job:localhost/replica:0/task:0/device:COMPOSITE:0", tf.device = "/job:localhost/replica:0/task:0/device:COMPOSITE:0"}, %arg11: tensor<*x!tf_type.resource>> {tf._composite_device = "/job:localhost/replica:0/task:0/device:COMPOSITE:0", tf.device = "/job:localhost/replica:0/task:0/device:COMPOSITE:0"}, %arg12: tensor<*x!tf_type.resource>> {tf._composite_device = "/job:localhost/replica:0/task:0/device:COMPOSITE:0", tf.device = "/job:localhost/replica:0/task:0/device:COMPOSITE:0"}) attributes {tf.entry_function = {control_outputs = "IteratorGetNext,IteratorGetNext_1,CrossReplicaSum,AssignAddVariableOp,CrossReplicaSum_1,AssignAddVariableOp_1,CrossReplicaSum_2,AssignAddVariableOp_2,CrossReplicaSum_3,AssignAddVariableOp_3", inputs = "iterator,iterator_1,iterator_2,iterator_3,iterator_4,iterator_5,resnet50_conv1_conv2d_conv1_kernel_140365606309224_handle_inputs_0,resnet50_fc1000_matmul_fc1000_kernel_140365944145960_handle_inputs_0,resnet50_fc1000_biasadd_fc1000_bias_140365944146240_handle_inputs_0,total_140366323758976_handle_inputs_0,count_140366323759312_handle_inputs_0,total_140366323760264_handle_inputs_0,count_140366323760600_handle_inputs_0", outputs = ""}} { // CHECK: %[[INPUT00:.*]] = "tf.IteratorGetNext" - // CHECK-DAG: %[[SPACETODEPTH00:.*]] = "tf.SpaceToDepth"([[INPUT00:.*]]#0) {block_size = 2 : i64, data_format = "NHWC"} : (tensor<2x224x224x3xf32>) -> tensor<2x112x112x12xf32> + // CHECK-DAG: %[[SPACETODEPTH00:.*]] = "tf.SpaceToDepth"([[INPUT00:.*]]#0) <{block_size = 2 : i64, data_format = "NHWC"}> : (tensor<2x224x224x3xf32>) -> tensor<2x112x112x12xf32> %0:2 = "tf.IteratorGetNext"(%arg2) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<*x!tf_type.resource>) -> (tensor<2x224x224x3xf32>, tensor<2x1xf32>) // CHECK: %[[INPUT01:.*]] = "tf.IteratorGetNext" - // CHECK-DAG: %[[SPACETODEPTH01:.*]] = "tf.SpaceToDepth"([[INPUT01:.*]]#0) {block_size = 2 : i64, data_format = "NHWC"} : (tensor<2x224x224x3xf32>) -> tensor<2x112x112x12xf32> + // CHECK-DAG: %[[SPACETODEPTH01:.*]] = "tf.SpaceToDepth"([[INPUT01:.*]]#0) <{block_size = 2 : i64, data_format = "NHWC"}> : (tensor<2x224x224x3xf32>) -> tensor<2x112x112x12xf32> %1:2 = "tf.IteratorGetNext"(%arg4) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<*x!tf_type.resource>) -> (tensor<2x224x224x3xf32>, tensor<2x1xf32>) tf_device.replicate([%0#0, %1#0] as %arg13: tensor<2x224x224x3xf32>, [%0#1, %1#1] as %arg14: tensor<2x1xf32>, %arg6 as %arg15: tensor<*x!tf_type.resource>>, %arg8 as %arg16: tensor<*x!tf_type.resource>>, %arg7 as %arg17: tensor<*x!tf_type.resource>>, %arg9 as %arg18: tensor<*x!tf_type.resource>>, %arg10 as %arg19: tensor<*x!tf_type.resource>>, %arg11 as %arg20: tensor<*x!tf_type.resource>>, %arg12 as %arg21: tensor<*x!tf_type.resource>>) {_mirrored_variable_indices = [2, 3, 4, 5, 6, 7, 8], devices = {TPU_REPLICATED_CORE_0 = ["/device:TPU:0", "/device:TPU:1"]}, n = 2 : i32} { %2 = "tf.ReadVariableOp"(%arg15) : (tensor<*x!tf_type.resource>>) -> tensor<7x7x3x64xf32> @@ -167,7 +167,7 @@ module attributes {tf.devices = {"/job:localhost/replica:0/task:0/device:COMPOSI module attributes {tf.devices = {"/job:localhost/replica:0/task:0/device:COMPOSITE:0" = {}, "/job:localhost/replica:0/task:0/device:CPU:0" = {}, "/job:localhost/replica:0/task:0/device:TPU:0" = {}, "/job:localhost/replica:0/task:0/device:TPU:1" = {}, "/job:localhost/replica:0/task:0/device:TPU_SYSTEM:0" = {}}, tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 458 : i32}} { func.func @main(%arg0: tensor<*x!tf_type.resource> {tf._user_specified_name = "iterator", tf.device = "/job:localhost/replica:0/task:0/device:CPU:0"}, %arg1: tensor {tf._user_specified_name = "iterator", tf.device = "/job:localhost/replica:0/task:0/device:CPU:0"}, %arg2: tensor<*x!tf_type.resource> {tf._user_specified_name = "iterator", tf.device = "/job:localhost/replica:0/task:0/device:CPU:0"}, %arg3: tensor {tf._user_specified_name = "iterator", tf.device = "/job:localhost/replica:0/task:0/device:CPU:0"}, %arg4: tensor<*x!tf_type.resource> {tf._user_specified_name = "iterator", tf.device = "/job:localhost/replica:0/task:0/device:CPU:0"}, %arg5: tensor {tf._user_specified_name = "iterator", tf.device = "/job:localhost/replica:0/task:0/device:CPU:0"}, %arg6: tensor<*x!tf_type.resource>> {tf._composite_device = "/job:localhost/replica:0/task:0/device:COMPOSITE:0", tf.device = "/job:localhost/replica:0/task:0/device:COMPOSITE:0"}, %arg7: tensor<*x!tf_type.resource>> {tf._composite_device = "/job:localhost/replica:0/task:0/device:COMPOSITE:0", tf.device = "/job:localhost/replica:0/task:0/device:COMPOSITE:0"}, %arg8: tensor<*x!tf_type.resource>> {tf._composite_device = "/job:localhost/replica:0/task:0/device:COMPOSITE:0", tf.device = "/job:localhost/replica:0/task:0/device:COMPOSITE:0"}, %arg9: tensor<*x!tf_type.resource>> {tf._composite_device = "/job:localhost/replica:0/task:0/device:COMPOSITE:0", tf.device = "/job:localhost/replica:0/task:0/device:COMPOSITE:0"}, %arg10: tensor<*x!tf_type.resource>> {tf._composite_device = "/job:localhost/replica:0/task:0/device:COMPOSITE:0", tf.device = "/job:localhost/replica:0/task:0/device:COMPOSITE:0"}, %arg11: tensor<*x!tf_type.resource>> {tf._composite_device = "/job:localhost/replica:0/task:0/device:COMPOSITE:0", tf.device = "/job:localhost/replica:0/task:0/device:COMPOSITE:0"}, %arg12: tensor<*x!tf_type.resource>> {tf._composite_device = "/job:localhost/replica:0/task:0/device:COMPOSITE:0", tf.device = "/job:localhost/replica:0/task:0/device:COMPOSITE:0"}) attributes {tf.entry_function = {control_outputs = "IteratorGetNext,IteratorGetNext_1,CrossReplicaSum,AssignAddVariableOp,CrossReplicaSum_1,AssignAddVariableOp_1,CrossReplicaSum_2,AssignAddVariableOp_2,CrossReplicaSum_3,AssignAddVariableOp_3", inputs = "iterator,iterator_1,iterator_2,iterator_3,iterator_4,iterator_5,resnet50_conv1_conv2d_conv1_kernel_140365606309224_handle_inputs_0,resnet50_fc1000_matmul_fc1000_kernel_140365944145960_handle_inputs_0,resnet50_fc1000_biasadd_fc1000_bias_140365944146240_handle_inputs_0,total_140366323758976_handle_inputs_0,count_140366323759312_handle_inputs_0,total_140366323760264_handle_inputs_0,count_140366323760600_handle_inputs_0", outputs = ""}} { // CHECK: %[[INPUT00:.*]] = "tf.IteratorGetNext" - // CHECK-DAG: %[[SPACETODEPTH00:.*]] = "tf.SpaceToDepth"([[INPUT00:.*]]#0) {block_size = 2 : i64, data_format = "NHWC"} : (tensor<2x224x224x3xf32>) -> tensor<2x112x112x12xf32> + // CHECK-DAG: %[[SPACETODEPTH00:.*]] = "tf.SpaceToDepth"([[INPUT00:.*]]#0) <{block_size = 2 : i64, data_format = "NHWC"}> : (tensor<2x224x224x3xf32>) -> tensor<2x112x112x12xf32> %0:2 = "tf.IteratorGetNext"(%arg2) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<*x!tf_type.resource>) -> (tensor<2x224x224x3xf32>, tensor<2x1xf32>) tf_device.replicate(%0#0 as %arg13: tensor<2x224x224x3xf32>, %0#1 as %arg14: tensor<2x1xf32>, %arg6 as %arg15: tensor<*x!tf_type.resource>>, %arg8 as %arg16: tensor<*x!tf_type.resource>>, %arg7 as %arg17: tensor<*x!tf_type.resource>>, %arg9 as %arg18: tensor<*x!tf_type.resource>>, %arg10 as %arg19: tensor<*x!tf_type.resource>>, %arg11 as %arg20: tensor<*x!tf_type.resource>>, %arg12 as %arg21: tensor<*x!tf_type.resource>>) {_mirrored_variable_indices = [2, 3, 4, 5, 6, 7, 8], devices = {TPU_REPLICATED_CORE_0 = ["/device:TPU:0", "/device:TPU:1"]}, n = 2 : i32} { %2 = "tf.ReadVariableOp"(%arg15) : (tensor<*x!tf_type.resource>>) -> tensor<7x7x3x64xf32> diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tpu_update_embedding_enqueue_op_inputs.mlir b/tensorflow/compiler/mlir/tensorflow/tests/tpu_update_embedding_enqueue_op_inputs.mlir index d3276c28394408..d1437a6e5a20ad 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/tpu_update_embedding_enqueue_op_inputs.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/tpu_update_embedding_enqueue_op_inputs.mlir @@ -15,7 +15,7 @@ func.func @check_enqueue_ops_update_for_eval(%arg0: tensor, %arg1: tens // CHECK: %[[CONST_0:.*]] = "tf.Const"() %0 = "tf.Const"() {value = dense<[]> : tensor<0xf32>} : () -> tensor<0xf32> - // CHECK: %[[CONST_MODE:.*]] = "tf.Const"() {_xla_outside_compilation = "0", value = dense<"inference"> : tensor} : () -> tensor + // CHECK: %[[CONST_MODE:.*]] = "tf.Const"() <{value = dense<"inference"> : tensor}> {_xla_outside_compilation = "0"} : () -> tensor // CHECK: "tf.EnqueueTPUEmbeddingSparseTensorBatch"(%[[ARG_0]], %[[ARG_1]], %[[ARG_2]], %[[ARG_3]], %[[ARG_4]], %[[ARG_5]], %[[CONST_0]], %[[CONST_0]], %[[CONST_0]], %[[CONST_MODE]]) "tf.EnqueueTPUEmbeddingSparseTensorBatch"(%arg0, %arg1, %arg2, %arg3, %arg4, %arg5, %0, %0, %0, %arg7) {_tpu_embedding_layer = "call1", _xla_outside_compilation = "0", combiners = ["mean", "sum"], device_ordinal = -1 : i64, max_sequence_lengths = [0, 0, 0], table_ids = [1, 1, 0]} : (tensor, tensor, tensor, tensor, tensor, tensor, tensor<0xf32>, tensor<0xf32>, tensor<0xf32>, tensor) -> () %2:2 = "tf.RecvTPUEmbeddingActivations"() {_tpu_embedding_layer = "call1", config = "\0A\0B\0C\0D"} : () -> (tensor<2x2xf32>, tensor<4x4xf32>) @@ -43,7 +43,7 @@ func.func @check_enqueue_ops_update_for_training(%arg0: tensor, %arg1: %3 = "tf.Const"() {value = dense<0.0> : tensor<4x4xf32>} : () -> tensor<4x4xf32> "tf.SendTPUEmbeddingGradients"(%2, %3) {_tpu_embedding_layer = "call1", config = "\0A\0B\0C\0D", operandSegmentSizes = array} : (tensor<2x2xf32>, tensor<4x4xf32>) -> () - // CHECK: %[[CONST_MODE:.*]] = "tf.Const"() {_xla_outside_compilation = "0", value = dense<"train"> : tensor} : () -> tensor + // CHECK: %[[CONST_MODE:.*]] = "tf.Const"() <{value = dense<"train"> : tensor}> {_xla_outside_compilation = "0"} : () -> tensor // CHECK: "tf.EnqueueTPUEmbeddingSparseTensorBatch"(%[[ARG_0]], %[[ARG_1]], %[[ARG_2]], %[[ARG_3]], %[[ARG_4]], %[[ARG_5]], %[[CONST_0]], %[[CONST_0]], %[[CONST_0]], %[[CONST_MODE]]) "tf.EnqueueTPUEmbeddingSparseTensorBatch"(%arg0, %arg1, %arg2, %arg3, %arg4, %arg5, %0, %0, %0, %arg7) {_tpu_embedding_layer = "call1", _xla_outside_compilation = "0", combiners = ["mean", "sum"], device_ordinal = -1 : i64, max_sequence_lengths = [0, 0, 0], table_ids = [1, 1, 0]} : (tensor, tensor, tensor, tensor, tensor, tensor, tensor<0xf32>, tensor<0xf32>, tensor<0xf32>, tensor) -> () %4:2 = "tf.RecvTPUEmbeddingActivations"() {_tpu_embedding_layer = "call1", config = "\0A\0B\0C\0D"} : () -> (tensor<2x2xf32>, tensor<4x4xf32>) diff --git a/tensorflow/compiler/mlir/tensorflow/tests/unroll-batch-matmul.mlir b/tensorflow/compiler/mlir/tensorflow/tests/unroll-batch-matmul.mlir index ec2d36d126775b..4333e79e0ee430 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/unroll-batch-matmul.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/unroll-batch-matmul.mlir @@ -7,12 +7,12 @@ func.func @batchMatMulTwoDim(%arg0: tensor<2x3x4x5xf32>, %arg1: tensor<2x3x5x6xf func.return %0 : tensor<2x3x4x6xf32> // CHECK-LABEL: batchMatMulTwoDim - // CHECK-DAG: %[[LHS_RESHAPED_SHAPE:.*]] = "tf.Const"() {value = dense<[6, 4, 5]> : tensor<3xi64>} - // CHECK-DAG: %[[RHS_RESHAPED_SHAPE:.*]] = "tf.Const"() {value = dense<[6, 5, 6]> : tensor<3xi64>} - // CHECK-DAG: %[[SPLITTING_AXIS:.*]] = "tf.Const"() {value = dense<0> : tensor} - // CHECK-DAG: %[[MATMUL_LHS_SHAPE:.*]] = "tf.Const"() {value = dense<[4, 5]> : tensor<2xi64>} - // CHECK-DAG: %[[MATMUL_RHS_SHAPE:.*]] = "tf.Const"() {value = dense<[5, 6]> : tensor<2xi64>} - // CHECK-DAG: %[[RESULT_SHAPE:.*]] = "tf.Const"() {value = dense<[2, 3, 4, 6]> : tensor<4xi64>} + // CHECK-DAG: %[[LHS_RESHAPED_SHAPE:.*]] = "tf.Const"() <{value = dense<[6, 4, 5]> : tensor<3xi64>}> + // CHECK-DAG: %[[RHS_RESHAPED_SHAPE:.*]] = "tf.Const"() <{value = dense<[6, 5, 6]> : tensor<3xi64>}> + // CHECK-DAG: %[[SPLITTING_AXIS:.*]] = "tf.Const"() <{value = dense<0> : tensor}> + // CHECK-DAG: %[[MATMUL_LHS_SHAPE:.*]] = "tf.Const"() <{value = dense<[4, 5]> : tensor<2xi64>}> + // CHECK-DAG: %[[MATMUL_RHS_SHAPE:.*]] = "tf.Const"() <{value = dense<[5, 6]> : tensor<2xi64>}> + // CHECK-DAG: %[[RESULT_SHAPE:.*]] = "tf.Const"() <{value = dense<[2, 3, 4, 6]> : tensor<4xi64>}> // CHECK: %[[LHS_RESHAPED:.*]] = "tf.Reshape"(%arg0, %[[LHS_RESHAPED_SHAPE]]) : (tensor<2x3x4x5xf32>, tensor<3xi64>) -> tensor<6x4x5xf32> // CHECK: %[[LHS_SPLIT:.*]]:6 = "tf.Split"(%[[SPLITTING_AXIS]], %[[LHS_RESHAPED]]) : (tensor, tensor<6x4x5xf32>) -> (tensor<1x4x5xf32>, tensor<1x4x5xf32>, tensor<1x4x5xf32>, tensor<1x4x5xf32>, tensor<1x4x5xf32>, tensor<1x4x5xf32>) @@ -32,14 +32,14 @@ func.func @batchMatMulTwoDim(%arg0: tensor<2x3x4x5xf32>, %arg1: tensor<2x3x5x6xf // CHECK: %[[RHS_5:.*]] = "tf.Reshape"(%[[RHS_SPLIT]]#4, %[[MATMUL_RHS_SHAPE]]) : (tensor<1x5x6xf32>, tensor<2xi64>) -> tensor<5x6xf32> // CHECK: %[[RHS_6:.*]] = "tf.Reshape"(%[[RHS_SPLIT]]#5, %[[MATMUL_RHS_SHAPE]]) : (tensor<1x5x6xf32>, tensor<2xi64>) -> tensor<5x6xf32> - // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%[[LHS_1]], %[[RHS_1]]) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_2:.*]] = "tf.MatMul"(%[[LHS_2]], %[[RHS_2]]) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_3:.*]] = "tf.MatMul"(%[[LHS_3]], %[[RHS_3]]) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_4:.*]] = "tf.MatMul"(%[[LHS_4]], %[[RHS_4]]) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_5:.*]] = "tf.MatMul"(%[[LHS_5]], %[[RHS_5]]) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_6:.*]] = "tf.MatMul"(%[[LHS_6]], %[[RHS_6]]) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%[[LHS_1]], %[[RHS_1]]) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_2:.*]] = "tf.MatMul"(%[[LHS_2]], %[[RHS_2]]) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_3:.*]] = "tf.MatMul"(%[[LHS_3]], %[[RHS_3]]) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_4:.*]] = "tf.MatMul"(%[[LHS_4]], %[[RHS_4]]) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_5:.*]] = "tf.MatMul"(%[[LHS_5]], %[[RHS_5]]) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_6:.*]] = "tf.MatMul"(%[[LHS_6]], %[[RHS_6]]) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_PACKED:.*]] = "tf.Pack"(%[[MATMUL_1]], %[[MATMUL_2]], %[[MATMUL_3]], %[[MATMUL_4]], %[[MATMUL_5]], %[[MATMUL_6]]) {axis = 0 : i64} : (tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>) -> tensor<6x4x6xf32> + // CHECK: %[[MATMUL_PACKED:.*]] = "tf.Pack"(%[[MATMUL_1]], %[[MATMUL_2]], %[[MATMUL_3]], %[[MATMUL_4]], %[[MATMUL_5]], %[[MATMUL_6]]) <{axis = 0 : i64}> : (tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>) -> tensor<6x4x6xf32> // CHECK: %[[RESULT:.*]] = "tf.Reshape"(%[[MATMUL_PACKED]], %[[RESULT_SHAPE]]) : (tensor<6x4x6xf32>, tensor<4xi64>) -> tensor<2x3x4x6xf32> // CHECK: return %[[RESULT]] : tensor<2x3x4x6xf32> } @@ -51,12 +51,12 @@ func.func @batchMatMulTwoDimAdjXY(%arg0: tensor<2x3x5x4xf32>, %arg1: tensor<2x3x func.return %0 : tensor<2x3x4x6xf32> // CHECK-LABEL: batchMatMulTwoDimAdjXY - // CHECK-DAG: %[[LHS_RESHAPED_SHAPE:.*]] = "tf.Const"() {value = dense<[6, 5, 4]> : tensor<3xi64>} - // CHECK-DAG: %[[RHS_RESHAPED_SHAPE:.*]] = "tf.Const"() {value = dense<[6, 6, 5]> : tensor<3xi64>} - // CHECK-DAG: %[[SPLITTING_AXIS:.*]] = "tf.Const"() {value = dense<0> : tensor} - // CHECK-DAG: %[[MATMUL_LHS_SHAPE:.*]] = "tf.Const"() {value = dense<[5, 4]> : tensor<2xi64>} - // CHECK-DAG: %[[MATMUL_RHS_SHAPE:.*]] = "tf.Const"() {value = dense<[6, 5]> : tensor<2xi64>} - // CHECK-DAG: %[[RESULT_SHAPE:.*]] = "tf.Const"() {value = dense<[2, 3, 4, 6]> : tensor<4xi64>} + // CHECK-DAG: %[[LHS_RESHAPED_SHAPE:.*]] = "tf.Const"() <{value = dense<[6, 5, 4]> : tensor<3xi64>}> + // CHECK-DAG: %[[RHS_RESHAPED_SHAPE:.*]] = "tf.Const"() <{value = dense<[6, 6, 5]> : tensor<3xi64>}> + // CHECK-DAG: %[[SPLITTING_AXIS:.*]] = "tf.Const"() <{value = dense<0> : tensor}> + // CHECK-DAG: %[[MATMUL_LHS_SHAPE:.*]] = "tf.Const"() <{value = dense<[5, 4]> : tensor<2xi64>}> + // CHECK-DAG: %[[MATMUL_RHS_SHAPE:.*]] = "tf.Const"() <{value = dense<[6, 5]> : tensor<2xi64>}> + // CHECK-DAG: %[[RESULT_SHAPE:.*]] = "tf.Const"() <{value = dense<[2, 3, 4, 6]> : tensor<4xi64>}> // CHECK: %[[LHS_RESHAPED:.*]] = "tf.Reshape"(%arg0, %[[LHS_RESHAPED_SHAPE]]) : (tensor<2x3x5x4xf32>, tensor<3xi64>) -> tensor<6x5x4xf32> // CHECK: %[[LHS_SPLIT:.*]]:6 = "tf.Split"(%[[SPLITTING_AXIS]], %[[LHS_RESHAPED]]) : (tensor, tensor<6x5x4xf32>) -> (tensor<1x5x4xf32>, tensor<1x5x4xf32>, tensor<1x5x4xf32>, tensor<1x5x4xf32>, tensor<1x5x4xf32>, tensor<1x5x4xf32>) @@ -76,14 +76,14 @@ func.func @batchMatMulTwoDimAdjXY(%arg0: tensor<2x3x5x4xf32>, %arg1: tensor<2x3x // CHECK: %[[RHS_5:.*]] = "tf.Reshape"(%[[RHS_SPLIT]]#4, %[[MATMUL_RHS_SHAPE]]) : (tensor<1x6x5xf32>, tensor<2xi64>) -> tensor<6x5xf32> // CHECK: %[[RHS_6:.*]] = "tf.Reshape"(%[[RHS_SPLIT]]#5, %[[MATMUL_RHS_SHAPE]]) : (tensor<1x6x5xf32>, tensor<2xi64>) -> tensor<6x5xf32> - // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%[[LHS_1]], %[[RHS_1]]) {transpose_a = true, transpose_b = true} : (tensor<5x4xf32>, tensor<6x5xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_2:.*]] = "tf.MatMul"(%[[LHS_2]], %[[RHS_2]]) {transpose_a = true, transpose_b = true} : (tensor<5x4xf32>, tensor<6x5xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_3:.*]] = "tf.MatMul"(%[[LHS_3]], %[[RHS_3]]) {transpose_a = true, transpose_b = true} : (tensor<5x4xf32>, tensor<6x5xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_4:.*]] = "tf.MatMul"(%[[LHS_4]], %[[RHS_4]]) {transpose_a = true, transpose_b = true} : (tensor<5x4xf32>, tensor<6x5xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_5:.*]] = "tf.MatMul"(%[[LHS_5]], %[[RHS_5]]) {transpose_a = true, transpose_b = true} : (tensor<5x4xf32>, tensor<6x5xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_6:.*]] = "tf.MatMul"(%[[LHS_6]], %[[RHS_6]]) {transpose_a = true, transpose_b = true} : (tensor<5x4xf32>, tensor<6x5xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%[[LHS_1]], %[[RHS_1]]) <{transpose_a = true, transpose_b = true}> : (tensor<5x4xf32>, tensor<6x5xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_2:.*]] = "tf.MatMul"(%[[LHS_2]], %[[RHS_2]]) <{transpose_a = true, transpose_b = true}> : (tensor<5x4xf32>, tensor<6x5xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_3:.*]] = "tf.MatMul"(%[[LHS_3]], %[[RHS_3]]) <{transpose_a = true, transpose_b = true}> : (tensor<5x4xf32>, tensor<6x5xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_4:.*]] = "tf.MatMul"(%[[LHS_4]], %[[RHS_4]]) <{transpose_a = true, transpose_b = true}> : (tensor<5x4xf32>, tensor<6x5xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_5:.*]] = "tf.MatMul"(%[[LHS_5]], %[[RHS_5]]) <{transpose_a = true, transpose_b = true}> : (tensor<5x4xf32>, tensor<6x5xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_6:.*]] = "tf.MatMul"(%[[LHS_6]], %[[RHS_6]]) <{transpose_a = true, transpose_b = true}> : (tensor<5x4xf32>, tensor<6x5xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_PACKED:.*]] = "tf.Pack"(%[[MATMUL_1]], %[[MATMUL_2]], %[[MATMUL_3]], %[[MATMUL_4]], %[[MATMUL_5]], %[[MATMUL_6]]) {axis = 0 : i64} : (tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>) -> tensor<6x4x6xf32> + // CHECK: %[[MATMUL_PACKED:.*]] = "tf.Pack"(%[[MATMUL_1]], %[[MATMUL_2]], %[[MATMUL_3]], %[[MATMUL_4]], %[[MATMUL_5]], %[[MATMUL_6]]) <{axis = 0 : i64}> : (tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>) -> tensor<6x4x6xf32> // CHECK: %[[RESULT:.*]] = "tf.Reshape"(%[[MATMUL_PACKED]], %[[RESULT_SHAPE]]) : (tensor<6x4x6xf32>, tensor<4xi64>) -> tensor<2x3x4x6xf32> // CHECK: return %[[RESULT]] : tensor<2x3x4x6xf32> } @@ -95,9 +95,9 @@ func.func @batchMatMulOneDim(%arg0: tensor<3x4x5xf32>, %arg1: tensor<3x5x6xf32>) func.return %0 : tensor<3x4x6xf32> // CHECK-LABEL: batchMatMulOneDim - // CHECK-DAG: %[[SPLITTING_AXIS:.*]] = "tf.Const"() {value = dense<0> : tensor} - // CHECK-DAG: %[[MATMUL_LHS_SHAPE:.*]] = "tf.Const"() {value = dense<[4, 5]> : tensor<2xi64>} - // CHECK-DAG: %[[MATMUL_RHS_SHAPE:.*]] = "tf.Const"() {value = dense<[5, 6]> : tensor<2xi64>} + // CHECK-DAG: %[[SPLITTING_AXIS:.*]] = "tf.Const"() <{value = dense<0> : tensor}> + // CHECK-DAG: %[[MATMUL_LHS_SHAPE:.*]] = "tf.Const"() <{value = dense<[4, 5]> : tensor<2xi64>}> + // CHECK-DAG: %[[MATMUL_RHS_SHAPE:.*]] = "tf.Const"() <{value = dense<[5, 6]> : tensor<2xi64>}> // CHECK: %[[LHS_RESHAPED:.*]]:3 = "tf.Split"(%[[SPLITTING_AXIS]], %arg0) : (tensor, tensor<3x4x5xf32>) -> (tensor<1x4x5xf32>, tensor<1x4x5xf32>, tensor<1x4x5xf32>) // CHECK: %[[LHS_1:.*]] = "tf.Reshape"(%[[LHS_RESHAPED]]#0, %[[MATMUL_LHS_SHAPE]]) : (tensor<1x4x5xf32>, tensor<2xi64>) -> tensor<4x5xf32> @@ -109,11 +109,11 @@ func.func @batchMatMulOneDim(%arg0: tensor<3x4x5xf32>, %arg1: tensor<3x5x6xf32>) // CHECK: %[[RHS_2:.*]] = "tf.Reshape"(%[[RHS_RESHAPED]]#1, %[[MATMUL_RHS_SHAPE]]) : (tensor<1x5x6xf32>, tensor<2xi64>) -> tensor<5x6xf32> // CHECK: %[[RHS_3:.*]] = "tf.Reshape"(%[[RHS_RESHAPED]]#2, %[[MATMUL_RHS_SHAPE]]) : (tensor<1x5x6xf32>, tensor<2xi64>) -> tensor<5x6xf32> - // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%[[LHS_1]], %[[RHS_1]]) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_2:.*]] = "tf.MatMul"(%[[LHS_2]], %[[RHS_2]]) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_3:.*]] = "tf.MatMul"(%[[LHS_3]], %[[RHS_3]]) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%[[LHS_1]], %[[RHS_1]]) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_2:.*]] = "tf.MatMul"(%[[LHS_2]], %[[RHS_2]]) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_3:.*]] = "tf.MatMul"(%[[LHS_3]], %[[RHS_3]]) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_PACKED:.*]] = "tf.Pack"(%[[MATMUL_1]], %[[MATMUL_2]], %[[MATMUL_3]]) {axis = 0 : i64} : (tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>) -> tensor<3x4x6xf32> + // CHECK: %[[MATMUL_PACKED:.*]] = "tf.Pack"(%[[MATMUL_1]], %[[MATMUL_2]], %[[MATMUL_3]]) <{axis = 0 : i64}> : (tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>) -> tensor<3x4x6xf32> // CHECK: return %[[MATMUL_PACKED]] : tensor<3x4x6xf32> } @@ -124,16 +124,16 @@ func.func @batchMatMulSingleBatch(%arg0: tensor<1x4x5xf32>, %arg1: tensor<1x5x6x func.return %0 : tensor<1x4x6xf32> // CHECK-LABEL: batchMatMulSingleBatch - // CHECK-DAG: %[[MATMUL_LHS_SHAPE:.*]] = "tf.Const"() {value = dense<[4, 5]> : tensor<2xi64>} : () -> tensor<2xi64> - // CHECK-DAG: %[[MATMUL_RHS_SHAPE:.*]] = "tf.Const"() {value = dense<[5, 6]> : tensor<2xi64>} : () -> tensor<2xi64> + // CHECK-DAG: %[[MATMUL_LHS_SHAPE:.*]] = "tf.Const"() <{value = dense<[4, 5]> : tensor<2xi64>}> : () -> tensor<2xi64> + // CHECK-DAG: %[[MATMUL_RHS_SHAPE:.*]] = "tf.Const"() <{value = dense<[5, 6]> : tensor<2xi64>}> : () -> tensor<2xi64> // CHECK: %[[LHS_1:.*]] = "tf.Reshape"(%arg0, %[[MATMUL_LHS_SHAPE]]) : (tensor<1x4x5xf32>, tensor<2xi64>) -> tensor<4x5xf32> // CHECK: %[[RHS_2:.*]] = "tf.Reshape"(%arg1, %[[MATMUL_RHS_SHAPE]]) : (tensor<1x5x6xf32>, tensor<2xi64>) -> tensor<5x6xf32> - // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%[[LHS_1]], %[[RHS_2]]) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%[[LHS_1]], %[[RHS_2]]) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_PACKED:.*]] = "tf.Pack"(%[[MATMUL_1]]) {axis = 0 : i64} : (tensor<4x6xf32>) -> tensor<1x4x6xf32> + // CHECK: %[[MATMUL_PACKED:.*]] = "tf.Pack"(%[[MATMUL_1]]) <{axis = 0 : i64}> : (tensor<4x6xf32>) -> tensor<1x4x6xf32> // CHECK: return %[[MATMUL_PACKED]] : tensor<1x4x6xf32> } @@ -144,19 +144,19 @@ func.func @batchMatMulUnbatchedLeft(%arg0: tensor<4x5xf32>, %arg1: tensor<3x5x6x func.return %0 : tensor<3x4x6xf32> // CHECK-LABEL: batchMatMulUnbatchedLeft - // CHECK-DAG: %[[SPLITTING_AXIS:.*]] = "tf.Const"() {value = dense<0> : tensor} - // CHECK-DAG: %[[MATMUL_RHS_SHAPE:.*]] = "tf.Const"() {value = dense<[5, 6]> : tensor<2xi64>} + // CHECK-DAG: %[[SPLITTING_AXIS:.*]] = "tf.Const"() <{value = dense<0> : tensor}> + // CHECK-DAG: %[[MATMUL_RHS_SHAPE:.*]] = "tf.Const"() <{value = dense<[5, 6]> : tensor<2xi64>}> // CHECK: %[[RHS_RESHAPED:.*]]:3 = "tf.Split"(%[[SPLITTING_AXIS]], %arg1) : (tensor, tensor<3x5x6xf32>) -> (tensor<1x5x6xf32>, tensor<1x5x6xf32>, tensor<1x5x6xf32>) // CHECK: %[[RHS_1:.*]] = "tf.Reshape"(%[[RHS_RESHAPED]]#0, %[[MATMUL_RHS_SHAPE]]) : (tensor<1x5x6xf32>, tensor<2xi64>) -> tensor<5x6xf32> // CHECK: %[[RHS_2:.*]] = "tf.Reshape"(%[[RHS_RESHAPED]]#1, %[[MATMUL_RHS_SHAPE]]) : (tensor<1x5x6xf32>, tensor<2xi64>) -> tensor<5x6xf32> // CHECK: %[[RHS_3:.*]] = "tf.Reshape"(%[[RHS_RESHAPED]]#2, %[[MATMUL_RHS_SHAPE]]) : (tensor<1x5x6xf32>, tensor<2xi64>) -> tensor<5x6xf32> - // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%arg0, %[[RHS_1]]) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_2:.*]] = "tf.MatMul"(%arg0, %[[RHS_2]]) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_3:.*]] = "tf.MatMul"(%arg0, %[[RHS_3]]) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%arg0, %[[RHS_1]]) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_2:.*]] = "tf.MatMul"(%arg0, %[[RHS_2]]) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_3:.*]] = "tf.MatMul"(%arg0, %[[RHS_3]]) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_PACKED:.*]] = "tf.Pack"(%[[MATMUL_1]], %[[MATMUL_2]], %[[MATMUL_3]]) {axis = 0 : i64} : (tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>) -> tensor<3x4x6xf32> + // CHECK: %[[MATMUL_PACKED:.*]] = "tf.Pack"(%[[MATMUL_1]], %[[MATMUL_2]], %[[MATMUL_3]]) <{axis = 0 : i64}> : (tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>) -> tensor<3x4x6xf32> // CHECK: return %[[MATMUL_PACKED]] : tensor<3x4x6xf32> } @@ -167,19 +167,19 @@ func.func @batchMatMulUnbatchedRight(%arg0: tensor<3x4x5xf32>, %arg1: tensor<5x6 func.return %0 : tensor<3x4x6xf32> // CHECK-LABEL: batchMatMulUnbatchedRight - // CHECK-DAG: %[[SPLITTING_AXIS:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor - // CHECK-DAG: %[[MATMUL_LHS_SHAPE:.*]] = "tf.Const"() {value = dense<[4, 5]> : tensor<2xi64>} : () -> tensor<2xi64> + // CHECK-DAG: %[[SPLITTING_AXIS:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor + // CHECK-DAG: %[[MATMUL_LHS_SHAPE:.*]] = "tf.Const"() <{value = dense<[4, 5]> : tensor<2xi64>}> : () -> tensor<2xi64> // CHECK: %[[LHS_SPLIT:.*]]:3 = "tf.Split"(%[[SPLITTING_AXIS]], %arg0) : (tensor, tensor<3x4x5xf32>) -> (tensor<1x4x5xf32>, tensor<1x4x5xf32>, tensor<1x4x5xf32>) // CHECK: %[[LHS_1:.*]] = "tf.Reshape"(%[[LHS_SPLIT]]#0, %[[MATMUL_LHS_SHAPE]]) : (tensor<1x4x5xf32>, tensor<2xi64>) -> tensor<4x5xf32> // CHECK: %[[LHS_2:.*]] = "tf.Reshape"(%[[LHS_SPLIT]]#1, %[[MATMUL_LHS_SHAPE]]) : (tensor<1x4x5xf32>, tensor<2xi64>) -> tensor<4x5xf32> // CHECK: %[[LHS_3:.*]] = "tf.Reshape"(%[[LHS_SPLIT]]#2, %[[MATMUL_LHS_SHAPE]]) : (tensor<1x4x5xf32>, tensor<2xi64>) -> tensor<4x5xf32> - // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%[[LHS_1]], %arg1) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_2:.*]] = "tf.MatMul"(%[[LHS_2]], %arg1) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_3:.*]] = "tf.MatMul"(%[[LHS_3]], %arg1) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%[[LHS_1]], %arg1) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_2:.*]] = "tf.MatMul"(%[[LHS_2]], %arg1) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_3:.*]] = "tf.MatMul"(%[[LHS_3]], %arg1) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_PACKED:.*]] = "tf.Pack"(%[[MATMUL_1]], %[[MATMUL_2]], %[[MATMUL_3]]) {axis = 0 : i64} : (tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>) -> tensor<3x4x6xf32> + // CHECK: %[[MATMUL_PACKED:.*]] = "tf.Pack"(%[[MATMUL_1]], %[[MATMUL_2]], %[[MATMUL_3]]) <{axis = 0 : i64}> : (tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>) -> tensor<3x4x6xf32> // CHECK: return %[[MATMUL_PACKED]] : tensor<3x4x6xf32> } @@ -190,7 +190,7 @@ func.func @batchMatMulMatrix(%arg0: tensor<4x5xf32>, %arg1: tensor<5x6xf32>) -> func.return %0 : tensor<4x6xf32> // CHECK-LABEL: batchMatMulMatrix - // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%arg0, %arg1) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%arg0, %arg1) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> // CHECK: return %[[MATMUL_1]] : tensor<4x6xf32> } @@ -201,7 +201,7 @@ func.func @batchMatMulMatrixAdjXY(%arg0: tensor<5x4xf32>, %arg1: tensor<6x5xf32> func.return %0 : tensor<4x6xf32> // CHECK-LABEL: batchMatMulMatrixAdjXY - // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%arg0, %arg1) {transpose_a = true, transpose_b = true} : (tensor<5x4xf32>, tensor<6x5xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%arg0, %arg1) <{transpose_a = true, transpose_b = true}> : (tensor<5x4xf32>, tensor<6x5xf32>) -> tensor<4x6xf32> // CHECK: return %[[MATMUL_1]] : tensor<4x6xf32> } @@ -213,12 +213,12 @@ func.func @batchMatMulV2TwoDim(%arg0: tensor<2x3x4x5xf32>, %arg1: tensor<2x3x5x6 func.return %0 : tensor<2x3x4x6xf32> // CHECK-LABEL: batchMatMulV2TwoDim - // CHECK-DAG: %[[LHS_RESHAPED_SHAPE:.*]] = "tf.Const"() {value = dense<[6, 4, 5]> : tensor<3xi64>} - // CHECK-DAG: %[[RHS_RESHAPED_SHAPE:.*]] = "tf.Const"() {value = dense<[6, 5, 6]> : tensor<3xi64>} - // CHECK-DAG: %[[SPLITTING_AXIS:.*]] = "tf.Const"() {value = dense<0> : tensor} - // CHECK-DAG: %[[MATMUL_LHS_SHAPE:.*]] = "tf.Const"() {value = dense<[4, 5]> : tensor<2xi64>} - // CHECK-DAG: %[[MATMUL_RHS_SHAPE:.*]] = "tf.Const"() {value = dense<[5, 6]> : tensor<2xi64>} - // CHECK-DAG: %[[RESULT_SHAPE:.*]] = "tf.Const"() {value = dense<[2, 3, 4, 6]> : tensor<4xi64>} + // CHECK-DAG: %[[LHS_RESHAPED_SHAPE:.*]] = "tf.Const"() <{value = dense<[6, 4, 5]> : tensor<3xi64>}> + // CHECK-DAG: %[[RHS_RESHAPED_SHAPE:.*]] = "tf.Const"() <{value = dense<[6, 5, 6]> : tensor<3xi64>}> + // CHECK-DAG: %[[SPLITTING_AXIS:.*]] = "tf.Const"() <{value = dense<0> : tensor}> + // CHECK-DAG: %[[MATMUL_LHS_SHAPE:.*]] = "tf.Const"() <{value = dense<[4, 5]> : tensor<2xi64>}> + // CHECK-DAG: %[[MATMUL_RHS_SHAPE:.*]] = "tf.Const"() <{value = dense<[5, 6]> : tensor<2xi64>}> + // CHECK-DAG: %[[RESULT_SHAPE:.*]] = "tf.Const"() <{value = dense<[2, 3, 4, 6]> : tensor<4xi64>}> // CHECK: %[[LHS_RESHAPED:.*]] = "tf.Reshape"(%arg0, %[[LHS_RESHAPED_SHAPE]]) : (tensor<2x3x4x5xf32>, tensor<3xi64>) -> tensor<6x4x5xf32> // CHECK: %[[LHS_SPLIT:.*]]:6 = "tf.Split"(%[[SPLITTING_AXIS]], %[[LHS_RESHAPED]]) : (tensor, tensor<6x4x5xf32>) -> (tensor<1x4x5xf32>, tensor<1x4x5xf32>, tensor<1x4x5xf32>, tensor<1x4x5xf32>, tensor<1x4x5xf32>, tensor<1x4x5xf32>) @@ -238,14 +238,14 @@ func.func @batchMatMulV2TwoDim(%arg0: tensor<2x3x4x5xf32>, %arg1: tensor<2x3x5x6 // CHECK: %[[RHS_5:.*]] = "tf.Reshape"(%[[RHS_SPLIT]]#4, %[[MATMUL_RHS_SHAPE]]) : (tensor<1x5x6xf32>, tensor<2xi64>) -> tensor<5x6xf32> // CHECK: %[[RHS_6:.*]] = "tf.Reshape"(%[[RHS_SPLIT]]#5, %[[MATMUL_RHS_SHAPE]]) : (tensor<1x5x6xf32>, tensor<2xi64>) -> tensor<5x6xf32> - // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%[[LHS_1]], %[[RHS_1]]) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_2:.*]] = "tf.MatMul"(%[[LHS_2]], %[[RHS_2]]) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_3:.*]] = "tf.MatMul"(%[[LHS_3]], %[[RHS_3]]) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_4:.*]] = "tf.MatMul"(%[[LHS_4]], %[[RHS_4]]) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_5:.*]] = "tf.MatMul"(%[[LHS_5]], %[[RHS_5]]) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_6:.*]] = "tf.MatMul"(%[[LHS_6]], %[[RHS_6]]) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%[[LHS_1]], %[[RHS_1]]) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_2:.*]] = "tf.MatMul"(%[[LHS_2]], %[[RHS_2]]) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_3:.*]] = "tf.MatMul"(%[[LHS_3]], %[[RHS_3]]) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_4:.*]] = "tf.MatMul"(%[[LHS_4]], %[[RHS_4]]) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_5:.*]] = "tf.MatMul"(%[[LHS_5]], %[[RHS_5]]) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_6:.*]] = "tf.MatMul"(%[[LHS_6]], %[[RHS_6]]) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_PACKED:.*]] = "tf.Pack"(%[[MATMUL_1]], %[[MATMUL_2]], %[[MATMUL_3]], %[[MATMUL_4]], %[[MATMUL_5]], %[[MATMUL_6]]) {axis = 0 : i64} : (tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>) -> tensor<6x4x6xf32> + // CHECK: %[[MATMUL_PACKED:.*]] = "tf.Pack"(%[[MATMUL_1]], %[[MATMUL_2]], %[[MATMUL_3]], %[[MATMUL_4]], %[[MATMUL_5]], %[[MATMUL_6]]) <{axis = 0 : i64}> : (tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>) -> tensor<6x4x6xf32> // CHECK: %[[RESULT:.*]] = "tf.Reshape"(%[[MATMUL_PACKED]], %[[RESULT_SHAPE]]) : (tensor<6x4x6xf32>, tensor<4xi64>) -> tensor<2x3x4x6xf32> // CHECK: return %[[RESULT]] : tensor<2x3x4x6xf32> } @@ -257,12 +257,12 @@ func.func @batchMatMulV2TwoDimAdjXY(%arg0: tensor<2x3x5x4xf32>, %arg1: tensor<2x func.return %0 : tensor<2x3x4x6xf32> // CHECK-LABEL: batchMatMulV2TwoDimAdjXY - // CHECK-DAG: %[[LHS_RESHAPED_SHAPE:.*]] = "tf.Const"() {value = dense<[6, 5, 4]> : tensor<3xi64>} - // CHECK-DAG: %[[RHS_RESHAPED_SHAPE:.*]] = "tf.Const"() {value = dense<[6, 6, 5]> : tensor<3xi64>} - // CHECK-DAG: %[[SPLITTING_AXIS:.*]] = "tf.Const"() {value = dense<0> : tensor} - // CHECK-DAG: %[[MATMUL_LHS_SHAPE:.*]] = "tf.Const"() {value = dense<[5, 4]> : tensor<2xi64>} - // CHECK-DAG: %[[MATMUL_RHS_SHAPE:.*]] = "tf.Const"() {value = dense<[6, 5]> : tensor<2xi64>} - // CHECK-DAG: %[[RESULT_SHAPE:.*]] = "tf.Const"() {value = dense<[2, 3, 4, 6]> : tensor<4xi64>} + // CHECK-DAG: %[[LHS_RESHAPED_SHAPE:.*]] = "tf.Const"() <{value = dense<[6, 5, 4]> : tensor<3xi64>}> + // CHECK-DAG: %[[RHS_RESHAPED_SHAPE:.*]] = "tf.Const"() <{value = dense<[6, 6, 5]> : tensor<3xi64>}> + // CHECK-DAG: %[[SPLITTING_AXIS:.*]] = "tf.Const"() <{value = dense<0> : tensor}> + // CHECK-DAG: %[[MATMUL_LHS_SHAPE:.*]] = "tf.Const"() <{value = dense<[5, 4]> : tensor<2xi64>}> + // CHECK-DAG: %[[MATMUL_RHS_SHAPE:.*]] = "tf.Const"() <{value = dense<[6, 5]> : tensor<2xi64>}> + // CHECK-DAG: %[[RESULT_SHAPE:.*]] = "tf.Const"() <{value = dense<[2, 3, 4, 6]> : tensor<4xi64>}> // CHECK: %[[LHS_RESHAPED:.*]] = "tf.Reshape"(%arg0, %[[LHS_RESHAPED_SHAPE]]) : (tensor<2x3x5x4xf32>, tensor<3xi64>) -> tensor<6x5x4xf32> // CHECK: %[[LHS_SPLIT:.*]]:6 = "tf.Split"(%[[SPLITTING_AXIS]], %[[LHS_RESHAPED]]) : (tensor, tensor<6x5x4xf32>) -> (tensor<1x5x4xf32>, tensor<1x5x4xf32>, tensor<1x5x4xf32>, tensor<1x5x4xf32>, tensor<1x5x4xf32>, tensor<1x5x4xf32>) @@ -282,14 +282,14 @@ func.func @batchMatMulV2TwoDimAdjXY(%arg0: tensor<2x3x5x4xf32>, %arg1: tensor<2x // CHECK: %[[RHS_5:.*]] = "tf.Reshape"(%[[RHS_SPLIT]]#4, %[[MATMUL_RHS_SHAPE]]) : (tensor<1x6x5xf32>, tensor<2xi64>) -> tensor<6x5xf32> // CHECK: %[[RHS_6:.*]] = "tf.Reshape"(%[[RHS_SPLIT]]#5, %[[MATMUL_RHS_SHAPE]]) : (tensor<1x6x5xf32>, tensor<2xi64>) -> tensor<6x5xf32> - // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%[[LHS_1]], %[[RHS_1]]) {transpose_a = true, transpose_b = true} : (tensor<5x4xf32>, tensor<6x5xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_2:.*]] = "tf.MatMul"(%[[LHS_2]], %[[RHS_2]]) {transpose_a = true, transpose_b = true} : (tensor<5x4xf32>, tensor<6x5xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_3:.*]] = "tf.MatMul"(%[[LHS_3]], %[[RHS_3]]) {transpose_a = true, transpose_b = true} : (tensor<5x4xf32>, tensor<6x5xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_4:.*]] = "tf.MatMul"(%[[LHS_4]], %[[RHS_4]]) {transpose_a = true, transpose_b = true} : (tensor<5x4xf32>, tensor<6x5xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_5:.*]] = "tf.MatMul"(%[[LHS_5]], %[[RHS_5]]) {transpose_a = true, transpose_b = true} : (tensor<5x4xf32>, tensor<6x5xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_6:.*]] = "tf.MatMul"(%[[LHS_6]], %[[RHS_6]]) {transpose_a = true, transpose_b = true} : (tensor<5x4xf32>, tensor<6x5xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%[[LHS_1]], %[[RHS_1]]) <{transpose_a = true, transpose_b = true}> : (tensor<5x4xf32>, tensor<6x5xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_2:.*]] = "tf.MatMul"(%[[LHS_2]], %[[RHS_2]]) <{transpose_a = true, transpose_b = true}> : (tensor<5x4xf32>, tensor<6x5xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_3:.*]] = "tf.MatMul"(%[[LHS_3]], %[[RHS_3]]) <{transpose_a = true, transpose_b = true}> : (tensor<5x4xf32>, tensor<6x5xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_4:.*]] = "tf.MatMul"(%[[LHS_4]], %[[RHS_4]]) <{transpose_a = true, transpose_b = true}> : (tensor<5x4xf32>, tensor<6x5xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_5:.*]] = "tf.MatMul"(%[[LHS_5]], %[[RHS_5]]) <{transpose_a = true, transpose_b = true}> : (tensor<5x4xf32>, tensor<6x5xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_6:.*]] = "tf.MatMul"(%[[LHS_6]], %[[RHS_6]]) <{transpose_a = true, transpose_b = true}> : (tensor<5x4xf32>, tensor<6x5xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_PACKED:.*]] = "tf.Pack"(%[[MATMUL_1]], %[[MATMUL_2]], %[[MATMUL_3]], %[[MATMUL_4]], %[[MATMUL_5]], %[[MATMUL_6]]) {axis = 0 : i64} : (tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>) -> tensor<6x4x6xf32> + // CHECK: %[[MATMUL_PACKED:.*]] = "tf.Pack"(%[[MATMUL_1]], %[[MATMUL_2]], %[[MATMUL_3]], %[[MATMUL_4]], %[[MATMUL_5]], %[[MATMUL_6]]) <{axis = 0 : i64}> : (tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>) -> tensor<6x4x6xf32> // CHECK: %[[RESULT:.*]] = "tf.Reshape"(%[[MATMUL_PACKED]], %[[RESULT_SHAPE]]) : (tensor<6x4x6xf32>, tensor<4xi64>) -> tensor<2x3x4x6xf32> // CHECK: return %[[RESULT]] : tensor<2x3x4x6xf32> } @@ -301,12 +301,12 @@ func.func @batchMatMulV2Broadcast(%arg0: tensor<2x1x4x5xf32>, %arg1: tensor<1x3x func.return %0 : tensor<2x3x4x6xf32> // CHECK-LABEL: batchMatMulV2Broadcast - // CHECK-DAG: %[[LHS_RESHAPED_SHAPE:.*]] = "tf.Const"() {value = dense<[2, 4, 5]> : tensor<3xi64>} - // CHECK-DAG: %[[RHS_RESHAPED_SHAPE:.*]] = "tf.Const"() {value = dense<[3, 5, 6]> : tensor<3xi64>} - // CHECK-DAG: %[[SPLITTING_AXIS:.*]] = "tf.Const"() {value = dense<0> : tensor} - // CHECK-DAG: %[[MATMUL_LHS_SHAPE:.*]] = "tf.Const"() {value = dense<[4, 5]> : tensor<2xi64>} - // CHECK-DAG: %[[MATMUL_RHS_SHAPE:.*]] = "tf.Const"() {value = dense<[5, 6]> : tensor<2xi64>} - // CHECK-DAG: %[[RESULT_SHAPE:.*]] = "tf.Const"() {value = dense<[2, 3, 4, 6]> : tensor<4xi64>} + // CHECK-DAG: %[[LHS_RESHAPED_SHAPE:.*]] = "tf.Const"() <{value = dense<[2, 4, 5]> : tensor<3xi64>}> + // CHECK-DAG: %[[RHS_RESHAPED_SHAPE:.*]] = "tf.Const"() <{value = dense<[3, 5, 6]> : tensor<3xi64>}> + // CHECK-DAG: %[[SPLITTING_AXIS:.*]] = "tf.Const"() <{value = dense<0> : tensor}> + // CHECK-DAG: %[[MATMUL_LHS_SHAPE:.*]] = "tf.Const"() <{value = dense<[4, 5]> : tensor<2xi64>}> + // CHECK-DAG: %[[MATMUL_RHS_SHAPE:.*]] = "tf.Const"() <{value = dense<[5, 6]> : tensor<2xi64>}> + // CHECK-DAG: %[[RESULT_SHAPE:.*]] = "tf.Const"() <{value = dense<[2, 3, 4, 6]> : tensor<4xi64>}> // CHECK: %[[LHS_RESHAPED:.*]] = "tf.Reshape"(%arg0, %[[LHS_RESHAPED_SHAPE]]) : (tensor<2x1x4x5xf32>, tensor<3xi64>) -> tensor<2x4x5xf32> // CHECK: %[[LHS_SPLIT:.*]]:2 = "tf.Split"(%[[SPLITTING_AXIS]], %[[LHS_RESHAPED]]) : (tensor, tensor<2x4x5xf32>) -> (tensor<1x4x5xf32>, tensor<1x4x5xf32>) @@ -319,14 +319,14 @@ func.func @batchMatMulV2Broadcast(%arg0: tensor<2x1x4x5xf32>, %arg1: tensor<1x3x // CHECK: %[[RHS_2:.*]] = "tf.Reshape"(%[[RHS_SPLIT]]#1, %[[MATMUL_RHS_SHAPE]]) : (tensor<1x5x6xf32>, tensor<2xi64>) -> tensor<5x6xf32> // CHECK: %[[RHS_3:.*]] = "tf.Reshape"(%[[RHS_SPLIT]]#2, %[[MATMUL_RHS_SHAPE]]) : (tensor<1x5x6xf32>, tensor<2xi64>) -> tensor<5x6xf32> - // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%[[LHS_1]], %[[RHS_1]]) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_2:.*]] = "tf.MatMul"(%[[LHS_1]], %[[RHS_2]]) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_3:.*]] = "tf.MatMul"(%[[LHS_1]], %[[RHS_3]]) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_4:.*]] = "tf.MatMul"(%[[LHS_2]], %[[RHS_1]]) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_5:.*]] = "tf.MatMul"(%[[LHS_2]], %[[RHS_2]]) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_6:.*]] = "tf.MatMul"(%[[LHS_2]], %[[RHS_3]]) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%[[LHS_1]], %[[RHS_1]]) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_2:.*]] = "tf.MatMul"(%[[LHS_1]], %[[RHS_2]]) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_3:.*]] = "tf.MatMul"(%[[LHS_1]], %[[RHS_3]]) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_4:.*]] = "tf.MatMul"(%[[LHS_2]], %[[RHS_1]]) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_5:.*]] = "tf.MatMul"(%[[LHS_2]], %[[RHS_2]]) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_6:.*]] = "tf.MatMul"(%[[LHS_2]], %[[RHS_3]]) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_PACKED:.*]] = "tf.Pack"(%[[MATMUL_1]], %[[MATMUL_2]], %[[MATMUL_3]], %[[MATMUL_4]], %[[MATMUL_5]], %[[MATMUL_6]]) {axis = 0 : i64} : (tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>) -> tensor<6x4x6xf32> + // CHECK: %[[MATMUL_PACKED:.*]] = "tf.Pack"(%[[MATMUL_1]], %[[MATMUL_2]], %[[MATMUL_3]], %[[MATMUL_4]], %[[MATMUL_5]], %[[MATMUL_6]]) <{axis = 0 : i64}> : (tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>) -> tensor<6x4x6xf32> // CHECK: %[[RESULT:.*]] = "tf.Reshape"(%[[MATMUL_PACKED]], %[[RESULT_SHAPE]]) : (tensor<6x4x6xf32>, tensor<4xi64>) -> tensor<2x3x4x6xf32> // CHECK: return %[[RESULT]] : tensor<2x3x4x6xf32> } @@ -338,9 +338,9 @@ func.func @batchMatMulV2OneDim(%arg0: tensor<3x4x5xf32>, %arg1: tensor<3x5x6xf32 func.return %0 : tensor<3x4x6xf32> // CHECK-LABEL: batchMatMulV2OneDim - // CHECK-DAG: %[[SPLITTING_AXIS:.*]] = "tf.Const"() {value = dense<0> : tensor} - // CHECK-DAG: %[[MATMUL_LHS_SHAPE:.*]] = "tf.Const"() {value = dense<[4, 5]> : tensor<2xi64>} - // CHECK-DAG: %[[MATMUL_RHS_SHAPE:.*]] = "tf.Const"() {value = dense<[5, 6]> : tensor<2xi64>} + // CHECK-DAG: %[[SPLITTING_AXIS:.*]] = "tf.Const"() <{value = dense<0> : tensor}> + // CHECK-DAG: %[[MATMUL_LHS_SHAPE:.*]] = "tf.Const"() <{value = dense<[4, 5]> : tensor<2xi64>}> + // CHECK-DAG: %[[MATMUL_RHS_SHAPE:.*]] = "tf.Const"() <{value = dense<[5, 6]> : tensor<2xi64>}> // CHECK: %[[LHS_RESHAPED:.*]]:3 = "tf.Split"(%[[SPLITTING_AXIS]], %arg0) : (tensor, tensor<3x4x5xf32>) -> (tensor<1x4x5xf32>, tensor<1x4x5xf32>, tensor<1x4x5xf32>) // CHECK: %[[LHS_1:.*]] = "tf.Reshape"(%[[LHS_RESHAPED]]#0, %[[MATMUL_LHS_SHAPE]]) : (tensor<1x4x5xf32>, tensor<2xi64>) -> tensor<4x5xf32> @@ -352,11 +352,11 @@ func.func @batchMatMulV2OneDim(%arg0: tensor<3x4x5xf32>, %arg1: tensor<3x5x6xf32 // CHECK: %[[RHS_2:.*]] = "tf.Reshape"(%[[RHS_RESHAPED]]#1, %[[MATMUL_RHS_SHAPE]]) : (tensor<1x5x6xf32>, tensor<2xi64>) -> tensor<5x6xf32> // CHECK: %[[RHS_3:.*]] = "tf.Reshape"(%[[RHS_RESHAPED]]#2, %[[MATMUL_RHS_SHAPE]]) : (tensor<1x5x6xf32>, tensor<2xi64>) -> tensor<5x6xf32> - // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%[[LHS_1]], %[[RHS_1]]) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_2:.*]] = "tf.MatMul"(%[[LHS_2]], %[[RHS_2]]) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_3:.*]] = "tf.MatMul"(%[[LHS_3]], %[[RHS_3]]) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%[[LHS_1]], %[[RHS_1]]) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_2:.*]] = "tf.MatMul"(%[[LHS_2]], %[[RHS_2]]) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_3:.*]] = "tf.MatMul"(%[[LHS_3]], %[[RHS_3]]) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_PACKED:.*]] = "tf.Pack"(%[[MATMUL_1]], %[[MATMUL_2]], %[[MATMUL_3]]) {axis = 0 : i64} : (tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>) -> tensor<3x4x6xf32> + // CHECK: %[[MATMUL_PACKED:.*]] = "tf.Pack"(%[[MATMUL_1]], %[[MATMUL_2]], %[[MATMUL_3]]) <{axis = 0 : i64}> : (tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>) -> tensor<3x4x6xf32> // CHECK: return %[[MATMUL_PACKED]] : tensor<3x4x6xf32> } @@ -367,16 +367,16 @@ func.func @batchMatMulV2SingleBatch(%arg0: tensor<1x4x5xf32>, %arg1: tensor<1x5x func.return %0 : tensor<1x4x6xf32> // CHECK-LABEL: batchMatMulV2SingleBatch - // CHECK-DAG: %[[MATMUL_LHS_SHAPE:.*]] = "tf.Const"() {value = dense<[4, 5]> : tensor<2xi64>} : () -> tensor<2xi64> - // CHECK-DAG: %[[MATMUL_RHS_SHAPE:.*]] = "tf.Const"() {value = dense<[5, 6]> : tensor<2xi64>} : () -> tensor<2xi64> + // CHECK-DAG: %[[MATMUL_LHS_SHAPE:.*]] = "tf.Const"() <{value = dense<[4, 5]> : tensor<2xi64>}> : () -> tensor<2xi64> + // CHECK-DAG: %[[MATMUL_RHS_SHAPE:.*]] = "tf.Const"() <{value = dense<[5, 6]> : tensor<2xi64>}> : () -> tensor<2xi64> // CHECK: %[[LHS_1:.*]] = "tf.Reshape"(%arg0, %[[MATMUL_LHS_SHAPE]]) : (tensor<1x4x5xf32>, tensor<2xi64>) -> tensor<4x5xf32> // CHECK: %[[RHS_2:.*]] = "tf.Reshape"(%arg1, %[[MATMUL_RHS_SHAPE]]) : (tensor<1x5x6xf32>, tensor<2xi64>) -> tensor<5x6xf32> - // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%[[LHS_1]], %[[RHS_2]]) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%[[LHS_1]], %[[RHS_2]]) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_PACKED:.*]] = "tf.Pack"(%[[MATMUL_1]]) {axis = 0 : i64} : (tensor<4x6xf32>) -> tensor<1x4x6xf32> + // CHECK: %[[MATMUL_PACKED:.*]] = "tf.Pack"(%[[MATMUL_1]]) <{axis = 0 : i64}> : (tensor<4x6xf32>) -> tensor<1x4x6xf32> // CHECK: return %[[MATMUL_PACKED]] : tensor<1x4x6xf32> } @@ -387,19 +387,19 @@ func.func @batchMatMulV2UnbatchedLeft(%arg0: tensor<4x5xf32>, %arg1: tensor<3x5x func.return %0 : tensor<3x4x6xf32> // CHECK-LABEL: batchMatMulV2UnbatchedLeft - // CHECK-DAG: %[[SPLITTING_AXIS:.*]] = "tf.Const"() {value = dense<0> : tensor} - // CHECK-DAG: %[[MATMUL_RHS_SHAPE:.*]] = "tf.Const"() {value = dense<[5, 6]> : tensor<2xi64>} + // CHECK-DAG: %[[SPLITTING_AXIS:.*]] = "tf.Const"() <{value = dense<0> : tensor}> + // CHECK-DAG: %[[MATMUL_RHS_SHAPE:.*]] = "tf.Const"() <{value = dense<[5, 6]> : tensor<2xi64>}> // CHECK: %[[RHS_RESHAPED:.*]]:3 = "tf.Split"(%[[SPLITTING_AXIS]], %arg1) : (tensor, tensor<3x5x6xf32>) -> (tensor<1x5x6xf32>, tensor<1x5x6xf32>, tensor<1x5x6xf32>) // CHECK: %[[RHS_1:.*]] = "tf.Reshape"(%[[RHS_RESHAPED]]#0, %[[MATMUL_RHS_SHAPE]]) : (tensor<1x5x6xf32>, tensor<2xi64>) -> tensor<5x6xf32> // CHECK: %[[RHS_2:.*]] = "tf.Reshape"(%[[RHS_RESHAPED]]#1, %[[MATMUL_RHS_SHAPE]]) : (tensor<1x5x6xf32>, tensor<2xi64>) -> tensor<5x6xf32> // CHECK: %[[RHS_3:.*]] = "tf.Reshape"(%[[RHS_RESHAPED]]#2, %[[MATMUL_RHS_SHAPE]]) : (tensor<1x5x6xf32>, tensor<2xi64>) -> tensor<5x6xf32> - // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%arg0, %[[RHS_1]]) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_2:.*]] = "tf.MatMul"(%arg0, %[[RHS_2]]) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_3:.*]] = "tf.MatMul"(%arg0, %[[RHS_3]]) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%arg0, %[[RHS_1]]) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_2:.*]] = "tf.MatMul"(%arg0, %[[RHS_2]]) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_3:.*]] = "tf.MatMul"(%arg0, %[[RHS_3]]) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_PACKED:.*]] = "tf.Pack"(%[[MATMUL_1]], %[[MATMUL_2]], %[[MATMUL_3]]) {axis = 0 : i64} : (tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>) -> tensor<3x4x6xf32> + // CHECK: %[[MATMUL_PACKED:.*]] = "tf.Pack"(%[[MATMUL_1]], %[[MATMUL_2]], %[[MATMUL_3]]) <{axis = 0 : i64}> : (tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>) -> tensor<3x4x6xf32> // CHECK: return %[[MATMUL_PACKED]] : tensor<3x4x6xf32> } @@ -410,19 +410,19 @@ func.func @batchMatMulV2UnbatchedRight(%arg0: tensor<3x4x5xf32>, %arg1: tensor<5 func.return %0 : tensor<3x4x6xf32> // CHECK-LABEL: batchMatMulV2UnbatchedRight - // CHECK-DAG: %[[SPLITTING_AXIS:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor - // CHECK-DAG: %[[MATMUL_LHS_SHAPE:.*]] = "tf.Const"() {value = dense<[4, 5]> : tensor<2xi64>} : () -> tensor<2xi64> + // CHECK-DAG: %[[SPLITTING_AXIS:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor + // CHECK-DAG: %[[MATMUL_LHS_SHAPE:.*]] = "tf.Const"() <{value = dense<[4, 5]> : tensor<2xi64>}> : () -> tensor<2xi64> // CHECK: %[[LHS_SPLIT:.*]]:3 = "tf.Split"(%[[SPLITTING_AXIS]], %arg0) : (tensor, tensor<3x4x5xf32>) -> (tensor<1x4x5xf32>, tensor<1x4x5xf32>, tensor<1x4x5xf32>) // CHECK: %[[LHS_1:.*]] = "tf.Reshape"(%[[LHS_SPLIT]]#0, %[[MATMUL_LHS_SHAPE]]) : (tensor<1x4x5xf32>, tensor<2xi64>) -> tensor<4x5xf32> // CHECK: %[[LHS_2:.*]] = "tf.Reshape"(%[[LHS_SPLIT]]#1, %[[MATMUL_LHS_SHAPE]]) : (tensor<1x4x5xf32>, tensor<2xi64>) -> tensor<4x5xf32> // CHECK: %[[LHS_3:.*]] = "tf.Reshape"(%[[LHS_SPLIT]]#2, %[[MATMUL_LHS_SHAPE]]) : (tensor<1x4x5xf32>, tensor<2xi64>) -> tensor<4x5xf32> - // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%[[LHS_1]], %arg1) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_2:.*]] = "tf.MatMul"(%[[LHS_2]], %arg1) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_3:.*]] = "tf.MatMul"(%[[LHS_3]], %arg1) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%[[LHS_1]], %arg1) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_2:.*]] = "tf.MatMul"(%[[LHS_2]], %arg1) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_3:.*]] = "tf.MatMul"(%[[LHS_3]], %arg1) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_PACKED:.*]] = "tf.Pack"(%[[MATMUL_1]], %[[MATMUL_2]], %[[MATMUL_3]]) {axis = 0 : i64} : (tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>) -> tensor<3x4x6xf32> + // CHECK: %[[MATMUL_PACKED:.*]] = "tf.Pack"(%[[MATMUL_1]], %[[MATMUL_2]], %[[MATMUL_3]]) <{axis = 0 : i64}> : (tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>) -> tensor<3x4x6xf32> // CHECK: return %[[MATMUL_PACKED]] : tensor<3x4x6xf32> } @@ -433,7 +433,7 @@ func.func @batchMatMulV2Matrix(%arg0: tensor<4x5xf32>, %arg1: tensor<5x6xf32>) - func.return %0 : tensor<4x6xf32> // CHECK-LABEL: batchMatMulV2Matrix - // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%arg0, %arg1) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%arg0, %arg1) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> // CHECK: return %[[MATMUL_1]] : tensor<4x6xf32> } @@ -444,7 +444,7 @@ func.func @batchMatMulV2MatrixAdjXY(%arg0: tensor<5x4xf32>, %arg1: tensor<6x5xf3 func.return %0 : tensor<4x6xf32> // CHECK-LABEL: batchMatMulV2MatrixAdjXY - // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%arg0, %arg1) {transpose_a = true, transpose_b = true} : (tensor<5x4xf32>, tensor<6x5xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%arg0, %arg1) <{transpose_a = true, transpose_b = true}> : (tensor<5x4xf32>, tensor<6x5xf32>) -> tensor<4x6xf32> // CHECK: return %[[MATMUL_1]] : tensor<4x6xf32> } @@ -455,7 +455,7 @@ func.func @batchMatMulV2DynamicSize(%arg0: tensor, %arg1: tensor // CHECK-LABEL: batchMatMulV2DynamicSize - // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%arg0, %arg1) {transpose_a = false, transpose_b = false} : (tensor, tensor) -> tensor + // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%arg0, %arg1) <{transpose_a = false, transpose_b = false}> : (tensor, tensor) -> tensor // CHECK: return %[[MATMUL_1]] : tensor } @@ -467,12 +467,12 @@ func.func @batchMatMulV3TwoDim(%arg0: tensor<2x3x4x5xf32>, %arg1: tensor<2x3x5x6 func.return %0 : tensor<2x3x4x6xf32> // CHECK-LABEL: batchMatMulV3TwoDim - // CHECK-DAG: %[[LHS_RESHAPED_SHAPE:.*]] = "tf.Const"() {value = dense<[6, 4, 5]> : tensor<3xi64>} - // CHECK-DAG: %[[RHS_RESHAPED_SHAPE:.*]] = "tf.Const"() {value = dense<[6, 5, 6]> : tensor<3xi64>} - // CHECK-DAG: %[[SPLITTING_AXIS:.*]] = "tf.Const"() {value = dense<0> : tensor} - // CHECK-DAG: %[[MATMUL_LHS_SHAPE:.*]] = "tf.Const"() {value = dense<[4, 5]> : tensor<2xi64>} - // CHECK-DAG: %[[MATMUL_RHS_SHAPE:.*]] = "tf.Const"() {value = dense<[5, 6]> : tensor<2xi64>} - // CHECK-DAG: %[[RESULT_SHAPE:.*]] = "tf.Const"() {value = dense<[2, 3, 4, 6]> : tensor<4xi64>} + // CHECK-DAG: %[[LHS_RESHAPED_SHAPE:.*]] = "tf.Const"() <{value = dense<[6, 4, 5]> : tensor<3xi64>}> + // CHECK-DAG: %[[RHS_RESHAPED_SHAPE:.*]] = "tf.Const"() <{value = dense<[6, 5, 6]> : tensor<3xi64>}> + // CHECK-DAG: %[[SPLITTING_AXIS:.*]] = "tf.Const"() <{value = dense<0> : tensor}> + // CHECK-DAG: %[[MATMUL_LHS_SHAPE:.*]] = "tf.Const"() <{value = dense<[4, 5]> : tensor<2xi64>}> + // CHECK-DAG: %[[MATMUL_RHS_SHAPE:.*]] = "tf.Const"() <{value = dense<[5, 6]> : tensor<2xi64>}> + // CHECK-DAG: %[[RESULT_SHAPE:.*]] = "tf.Const"() <{value = dense<[2, 3, 4, 6]> : tensor<4xi64>}> // CHECK: %[[LHS_RESHAPED:.*]] = "tf.Reshape"(%arg0, %[[LHS_RESHAPED_SHAPE]]) : (tensor<2x3x4x5xf32>, tensor<3xi64>) -> tensor<6x4x5xf32> // CHECK: %[[LHS_SPLIT:.*]]:6 = "tf.Split"(%[[SPLITTING_AXIS]], %[[LHS_RESHAPED]]) : (tensor, tensor<6x4x5xf32>) -> (tensor<1x4x5xf32>, tensor<1x4x5xf32>, tensor<1x4x5xf32>, tensor<1x4x5xf32>, tensor<1x4x5xf32>, tensor<1x4x5xf32>) @@ -492,14 +492,14 @@ func.func @batchMatMulV3TwoDim(%arg0: tensor<2x3x4x5xf32>, %arg1: tensor<2x3x5x6 // CHECK: %[[RHS_5:.*]] = "tf.Reshape"(%[[RHS_SPLIT]]#4, %[[MATMUL_RHS_SHAPE]]) : (tensor<1x5x6xf32>, tensor<2xi64>) -> tensor<5x6xf32> // CHECK: %[[RHS_6:.*]] = "tf.Reshape"(%[[RHS_SPLIT]]#5, %[[MATMUL_RHS_SHAPE]]) : (tensor<1x5x6xf32>, tensor<2xi64>) -> tensor<5x6xf32> - // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%[[LHS_1]], %[[RHS_1]]) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_2:.*]] = "tf.MatMul"(%[[LHS_2]], %[[RHS_2]]) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_3:.*]] = "tf.MatMul"(%[[LHS_3]], %[[RHS_3]]) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_4:.*]] = "tf.MatMul"(%[[LHS_4]], %[[RHS_4]]) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_5:.*]] = "tf.MatMul"(%[[LHS_5]], %[[RHS_5]]) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_6:.*]] = "tf.MatMul"(%[[LHS_6]], %[[RHS_6]]) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%[[LHS_1]], %[[RHS_1]]) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_2:.*]] = "tf.MatMul"(%[[LHS_2]], %[[RHS_2]]) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_3:.*]] = "tf.MatMul"(%[[LHS_3]], %[[RHS_3]]) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_4:.*]] = "tf.MatMul"(%[[LHS_4]], %[[RHS_4]]) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_5:.*]] = "tf.MatMul"(%[[LHS_5]], %[[RHS_5]]) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_6:.*]] = "tf.MatMul"(%[[LHS_6]], %[[RHS_6]]) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_PACKED:.*]] = "tf.Pack"(%[[MATMUL_1]], %[[MATMUL_2]], %[[MATMUL_3]], %[[MATMUL_4]], %[[MATMUL_5]], %[[MATMUL_6]]) {axis = 0 : i64} : (tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>) -> tensor<6x4x6xf32> + // CHECK: %[[MATMUL_PACKED:.*]] = "tf.Pack"(%[[MATMUL_1]], %[[MATMUL_2]], %[[MATMUL_3]], %[[MATMUL_4]], %[[MATMUL_5]], %[[MATMUL_6]]) <{axis = 0 : i64}> : (tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>) -> tensor<6x4x6xf32> // CHECK: %[[RESULT:.*]] = "tf.Reshape"(%[[MATMUL_PACKED]], %[[RESULT_SHAPE]]) : (tensor<6x4x6xf32>, tensor<4xi64>) -> tensor<2x3x4x6xf32> // CHECK: return %[[RESULT]] : tensor<2x3x4x6xf32> } @@ -511,12 +511,12 @@ func.func @batchMatMulV3TwoDimAdjXY(%arg0: tensor<2x3x5x4xf32>, %arg1: tensor<2x func.return %0 : tensor<2x3x4x6xf32> // CHECK-LABEL: batchMatMulV3TwoDimAdjXY - // CHECK-DAG: %[[LHS_RESHAPED_SHAPE:.*]] = "tf.Const"() {value = dense<[6, 5, 4]> : tensor<3xi64>} - // CHECK-DAG: %[[RHS_RESHAPED_SHAPE:.*]] = "tf.Const"() {value = dense<[6, 6, 5]> : tensor<3xi64>} - // CHECK-DAG: %[[SPLITTING_AXIS:.*]] = "tf.Const"() {value = dense<0> : tensor} - // CHECK-DAG: %[[MATMUL_LHS_SHAPE:.*]] = "tf.Const"() {value = dense<[5, 4]> : tensor<2xi64>} - // CHECK-DAG: %[[MATMUL_RHS_SHAPE:.*]] = "tf.Const"() {value = dense<[6, 5]> : tensor<2xi64>} - // CHECK-DAG: %[[RESULT_SHAPE:.*]] = "tf.Const"() {value = dense<[2, 3, 4, 6]> : tensor<4xi64>} + // CHECK-DAG: %[[LHS_RESHAPED_SHAPE:.*]] = "tf.Const"() <{value = dense<[6, 5, 4]> : tensor<3xi64>}> + // CHECK-DAG: %[[RHS_RESHAPED_SHAPE:.*]] = "tf.Const"() <{value = dense<[6, 6, 5]> : tensor<3xi64>}> + // CHECK-DAG: %[[SPLITTING_AXIS:.*]] = "tf.Const"() <{value = dense<0> : tensor}> + // CHECK-DAG: %[[MATMUL_LHS_SHAPE:.*]] = "tf.Const"() <{value = dense<[5, 4]> : tensor<2xi64>}> + // CHECK-DAG: %[[MATMUL_RHS_SHAPE:.*]] = "tf.Const"() <{value = dense<[6, 5]> : tensor<2xi64>}> + // CHECK-DAG: %[[RESULT_SHAPE:.*]] = "tf.Const"() <{value = dense<[2, 3, 4, 6]> : tensor<4xi64>}> // CHECK: %[[LHS_RESHAPED:.*]] = "tf.Reshape"(%arg0, %[[LHS_RESHAPED_SHAPE]]) : (tensor<2x3x5x4xf32>, tensor<3xi64>) -> tensor<6x5x4xf32> // CHECK: %[[LHS_SPLIT:.*]]:6 = "tf.Split"(%[[SPLITTING_AXIS]], %[[LHS_RESHAPED]]) : (tensor, tensor<6x5x4xf32>) -> (tensor<1x5x4xf32>, tensor<1x5x4xf32>, tensor<1x5x4xf32>, tensor<1x5x4xf32>, tensor<1x5x4xf32>, tensor<1x5x4xf32>) @@ -536,14 +536,14 @@ func.func @batchMatMulV3TwoDimAdjXY(%arg0: tensor<2x3x5x4xf32>, %arg1: tensor<2x // CHECK: %[[RHS_5:.*]] = "tf.Reshape"(%[[RHS_SPLIT]]#4, %[[MATMUL_RHS_SHAPE]]) : (tensor<1x6x5xf32>, tensor<2xi64>) -> tensor<6x5xf32> // CHECK: %[[RHS_6:.*]] = "tf.Reshape"(%[[RHS_SPLIT]]#5, %[[MATMUL_RHS_SHAPE]]) : (tensor<1x6x5xf32>, tensor<2xi64>) -> tensor<6x5xf32> - // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%[[LHS_1]], %[[RHS_1]]) {transpose_a = true, transpose_b = true} : (tensor<5x4xf32>, tensor<6x5xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_2:.*]] = "tf.MatMul"(%[[LHS_2]], %[[RHS_2]]) {transpose_a = true, transpose_b = true} : (tensor<5x4xf32>, tensor<6x5xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_3:.*]] = "tf.MatMul"(%[[LHS_3]], %[[RHS_3]]) {transpose_a = true, transpose_b = true} : (tensor<5x4xf32>, tensor<6x5xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_4:.*]] = "tf.MatMul"(%[[LHS_4]], %[[RHS_4]]) {transpose_a = true, transpose_b = true} : (tensor<5x4xf32>, tensor<6x5xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_5:.*]] = "tf.MatMul"(%[[LHS_5]], %[[RHS_5]]) {transpose_a = true, transpose_b = true} : (tensor<5x4xf32>, tensor<6x5xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_6:.*]] = "tf.MatMul"(%[[LHS_6]], %[[RHS_6]]) {transpose_a = true, transpose_b = true} : (tensor<5x4xf32>, tensor<6x5xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%[[LHS_1]], %[[RHS_1]]) <{transpose_a = true, transpose_b = true}> : (tensor<5x4xf32>, tensor<6x5xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_2:.*]] = "tf.MatMul"(%[[LHS_2]], %[[RHS_2]]) <{transpose_a = true, transpose_b = true}> : (tensor<5x4xf32>, tensor<6x5xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_3:.*]] = "tf.MatMul"(%[[LHS_3]], %[[RHS_3]]) <{transpose_a = true, transpose_b = true}> : (tensor<5x4xf32>, tensor<6x5xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_4:.*]] = "tf.MatMul"(%[[LHS_4]], %[[RHS_4]]) <{transpose_a = true, transpose_b = true}> : (tensor<5x4xf32>, tensor<6x5xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_5:.*]] = "tf.MatMul"(%[[LHS_5]], %[[RHS_5]]) <{transpose_a = true, transpose_b = true}> : (tensor<5x4xf32>, tensor<6x5xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_6:.*]] = "tf.MatMul"(%[[LHS_6]], %[[RHS_6]]) <{transpose_a = true, transpose_b = true}> : (tensor<5x4xf32>, tensor<6x5xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_PACKED:.*]] = "tf.Pack"(%[[MATMUL_1]], %[[MATMUL_2]], %[[MATMUL_3]], %[[MATMUL_4]], %[[MATMUL_5]], %[[MATMUL_6]]) {axis = 0 : i64} : (tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>) -> tensor<6x4x6xf32> + // CHECK: %[[MATMUL_PACKED:.*]] = "tf.Pack"(%[[MATMUL_1]], %[[MATMUL_2]], %[[MATMUL_3]], %[[MATMUL_4]], %[[MATMUL_5]], %[[MATMUL_6]]) <{axis = 0 : i64}> : (tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>) -> tensor<6x4x6xf32> // CHECK: %[[RESULT:.*]] = "tf.Reshape"(%[[MATMUL_PACKED]], %[[RESULT_SHAPE]]) : (tensor<6x4x6xf32>, tensor<4xi64>) -> tensor<2x3x4x6xf32> // CHECK: return %[[RESULT]] : tensor<2x3x4x6xf32> } @@ -555,12 +555,12 @@ func.func @batchMatMulV3Broadcast(%arg0: tensor<2x1x4x5xf32>, %arg1: tensor<1x3x func.return %0 : tensor<2x3x4x6xf32> // CHECK-LABEL: batchMatMulV3Broadcast - // CHECK-DAG: %[[LHS_RESHAPED_SHAPE:.*]] = "tf.Const"() {value = dense<[2, 4, 5]> : tensor<3xi64>} - // CHECK-DAG: %[[RHS_RESHAPED_SHAPE:.*]] = "tf.Const"() {value = dense<[3, 5, 6]> : tensor<3xi64>} - // CHECK-DAG: %[[SPLITTING_AXIS:.*]] = "tf.Const"() {value = dense<0> : tensor} - // CHECK-DAG: %[[MATMUL_LHS_SHAPE:.*]] = "tf.Const"() {value = dense<[4, 5]> : tensor<2xi64>} - // CHECK-DAG: %[[MATMUL_RHS_SHAPE:.*]] = "tf.Const"() {value = dense<[5, 6]> : tensor<2xi64>} - // CHECK-DAG: %[[RESULT_SHAPE:.*]] = "tf.Const"() {value = dense<[2, 3, 4, 6]> : tensor<4xi64>} + // CHECK-DAG: %[[LHS_RESHAPED_SHAPE:.*]] = "tf.Const"() <{value = dense<[2, 4, 5]> : tensor<3xi64>}> + // CHECK-DAG: %[[RHS_RESHAPED_SHAPE:.*]] = "tf.Const"() <{value = dense<[3, 5, 6]> : tensor<3xi64>}> + // CHECK-DAG: %[[SPLITTING_AXIS:.*]] = "tf.Const"() <{value = dense<0> : tensor}> + // CHECK-DAG: %[[MATMUL_LHS_SHAPE:.*]] = "tf.Const"() <{value = dense<[4, 5]> : tensor<2xi64>}> + // CHECK-DAG: %[[MATMUL_RHS_SHAPE:.*]] = "tf.Const"() <{value = dense<[5, 6]> : tensor<2xi64>}> + // CHECK-DAG: %[[RESULT_SHAPE:.*]] = "tf.Const"() <{value = dense<[2, 3, 4, 6]> : tensor<4xi64>}> // CHECK: %[[LHS_RESHAPED:.*]] = "tf.Reshape"(%arg0, %[[LHS_RESHAPED_SHAPE]]) : (tensor<2x1x4x5xf32>, tensor<3xi64>) -> tensor<2x4x5xf32> // CHECK: %[[LHS_SPLIT:.*]]:2 = "tf.Split"(%[[SPLITTING_AXIS]], %[[LHS_RESHAPED]]) : (tensor, tensor<2x4x5xf32>) -> (tensor<1x4x5xf32>, tensor<1x4x5xf32>) @@ -573,14 +573,14 @@ func.func @batchMatMulV3Broadcast(%arg0: tensor<2x1x4x5xf32>, %arg1: tensor<1x3x // CHECK: %[[RHS_2:.*]] = "tf.Reshape"(%[[RHS_SPLIT]]#1, %[[MATMUL_RHS_SHAPE]]) : (tensor<1x5x6xf32>, tensor<2xi64>) -> tensor<5x6xf32> // CHECK: %[[RHS_3:.*]] = "tf.Reshape"(%[[RHS_SPLIT]]#2, %[[MATMUL_RHS_SHAPE]]) : (tensor<1x5x6xf32>, tensor<2xi64>) -> tensor<5x6xf32> - // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%[[LHS_1]], %[[RHS_1]]) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_2:.*]] = "tf.MatMul"(%[[LHS_1]], %[[RHS_2]]) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_3:.*]] = "tf.MatMul"(%[[LHS_1]], %[[RHS_3]]) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_4:.*]] = "tf.MatMul"(%[[LHS_2]], %[[RHS_1]]) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_5:.*]] = "tf.MatMul"(%[[LHS_2]], %[[RHS_2]]) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_6:.*]] = "tf.MatMul"(%[[LHS_2]], %[[RHS_3]]) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%[[LHS_1]], %[[RHS_1]]) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_2:.*]] = "tf.MatMul"(%[[LHS_1]], %[[RHS_2]]) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_3:.*]] = "tf.MatMul"(%[[LHS_1]], %[[RHS_3]]) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_4:.*]] = "tf.MatMul"(%[[LHS_2]], %[[RHS_1]]) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_5:.*]] = "tf.MatMul"(%[[LHS_2]], %[[RHS_2]]) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_6:.*]] = "tf.MatMul"(%[[LHS_2]], %[[RHS_3]]) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_PACKED:.*]] = "tf.Pack"(%[[MATMUL_1]], %[[MATMUL_2]], %[[MATMUL_3]], %[[MATMUL_4]], %[[MATMUL_5]], %[[MATMUL_6]]) {axis = 0 : i64} : (tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>) -> tensor<6x4x6xf32> + // CHECK: %[[MATMUL_PACKED:.*]] = "tf.Pack"(%[[MATMUL_1]], %[[MATMUL_2]], %[[MATMUL_3]], %[[MATMUL_4]], %[[MATMUL_5]], %[[MATMUL_6]]) <{axis = 0 : i64}> : (tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>) -> tensor<6x4x6xf32> // CHECK: %[[RESULT:.*]] = "tf.Reshape"(%[[MATMUL_PACKED]], %[[RESULT_SHAPE]]) : (tensor<6x4x6xf32>, tensor<4xi64>) -> tensor<2x3x4x6xf32> // CHECK: return %[[RESULT]] : tensor<2x3x4x6xf32> } @@ -592,9 +592,9 @@ func.func @batchMatMulV3OneDim(%arg0: tensor<3x4x5xf32>, %arg1: tensor<3x5x6xf32 func.return %0 : tensor<3x4x6xf32> // CHECK-LABEL: batchMatMulV3OneDim - // CHECK-DAG: %[[SPLITTING_AXIS:.*]] = "tf.Const"() {value = dense<0> : tensor} - // CHECK-DAG: %[[MATMUL_LHS_SHAPE:.*]] = "tf.Const"() {value = dense<[4, 5]> : tensor<2xi64>} - // CHECK-DAG: %[[MATMUL_RHS_SHAPE:.*]] = "tf.Const"() {value = dense<[5, 6]> : tensor<2xi64>} + // CHECK-DAG: %[[SPLITTING_AXIS:.*]] = "tf.Const"() <{value = dense<0> : tensor}> + // CHECK-DAG: %[[MATMUL_LHS_SHAPE:.*]] = "tf.Const"() <{value = dense<[4, 5]> : tensor<2xi64>}> + // CHECK-DAG: %[[MATMUL_RHS_SHAPE:.*]] = "tf.Const"() <{value = dense<[5, 6]> : tensor<2xi64>}> // CHECK: %[[LHS_RESHAPED:.*]]:3 = "tf.Split"(%[[SPLITTING_AXIS]], %arg0) : (tensor, tensor<3x4x5xf32>) -> (tensor<1x4x5xf32>, tensor<1x4x5xf32>, tensor<1x4x5xf32>) // CHECK: %[[LHS_1:.*]] = "tf.Reshape"(%[[LHS_RESHAPED]]#0, %[[MATMUL_LHS_SHAPE]]) : (tensor<1x4x5xf32>, tensor<2xi64>) -> tensor<4x5xf32> @@ -606,11 +606,11 @@ func.func @batchMatMulV3OneDim(%arg0: tensor<3x4x5xf32>, %arg1: tensor<3x5x6xf32 // CHECK: %[[RHS_2:.*]] = "tf.Reshape"(%[[RHS_RESHAPED]]#1, %[[MATMUL_RHS_SHAPE]]) : (tensor<1x5x6xf32>, tensor<2xi64>) -> tensor<5x6xf32> // CHECK: %[[RHS_3:.*]] = "tf.Reshape"(%[[RHS_RESHAPED]]#2, %[[MATMUL_RHS_SHAPE]]) : (tensor<1x5x6xf32>, tensor<2xi64>) -> tensor<5x6xf32> - // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%[[LHS_1]], %[[RHS_1]]) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_2:.*]] = "tf.MatMul"(%[[LHS_2]], %[[RHS_2]]) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_3:.*]] = "tf.MatMul"(%[[LHS_3]], %[[RHS_3]]) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%[[LHS_1]], %[[RHS_1]]) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_2:.*]] = "tf.MatMul"(%[[LHS_2]], %[[RHS_2]]) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_3:.*]] = "tf.MatMul"(%[[LHS_3]], %[[RHS_3]]) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_PACKED:.*]] = "tf.Pack"(%[[MATMUL_1]], %[[MATMUL_2]], %[[MATMUL_3]]) {axis = 0 : i64} : (tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>) -> tensor<3x4x6xf32> + // CHECK: %[[MATMUL_PACKED:.*]] = "tf.Pack"(%[[MATMUL_1]], %[[MATMUL_2]], %[[MATMUL_3]]) <{axis = 0 : i64}> : (tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>) -> tensor<3x4x6xf32> // CHECK: return %[[MATMUL_PACKED]] : tensor<3x4x6xf32> } @@ -621,16 +621,16 @@ func.func @batchMatMulV3SingleBatch(%arg0: tensor<1x4x5xf32>, %arg1: tensor<1x5x func.return %0 : tensor<1x4x6xf32> // CHECK-LABEL: batchMatMulV3SingleBatch - // CHECK-DAG: %[[MATMUL_LHS_SHAPE:.*]] = "tf.Const"() {value = dense<[4, 5]> : tensor<2xi64>} : () -> tensor<2xi64> - // CHECK-DAG: %[[MATMUL_RHS_SHAPE:.*]] = "tf.Const"() {value = dense<[5, 6]> : tensor<2xi64>} : () -> tensor<2xi64> + // CHECK-DAG: %[[MATMUL_LHS_SHAPE:.*]] = "tf.Const"() <{value = dense<[4, 5]> : tensor<2xi64>}> : () -> tensor<2xi64> + // CHECK-DAG: %[[MATMUL_RHS_SHAPE:.*]] = "tf.Const"() <{value = dense<[5, 6]> : tensor<2xi64>}> : () -> tensor<2xi64> // CHECK: %[[LHS_1:.*]] = "tf.Reshape"(%arg0, %[[MATMUL_LHS_SHAPE]]) : (tensor<1x4x5xf32>, tensor<2xi64>) -> tensor<4x5xf32> // CHECK: %[[RHS_2:.*]] = "tf.Reshape"(%arg1, %[[MATMUL_RHS_SHAPE]]) : (tensor<1x5x6xf32>, tensor<2xi64>) -> tensor<5x6xf32> - // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%[[LHS_1]], %[[RHS_2]]) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%[[LHS_1]], %[[RHS_2]]) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_PACKED:.*]] = "tf.Pack"(%[[MATMUL_1]]) {axis = 0 : i64} : (tensor<4x6xf32>) -> tensor<1x4x6xf32> + // CHECK: %[[MATMUL_PACKED:.*]] = "tf.Pack"(%[[MATMUL_1]]) <{axis = 0 : i64}> : (tensor<4x6xf32>) -> tensor<1x4x6xf32> // CHECK: return %[[MATMUL_PACKED]] : tensor<1x4x6xf32> } @@ -641,19 +641,19 @@ func.func @batchMatMulV3UnbatchedLeft(%arg0: tensor<4x5xf32>, %arg1: tensor<3x5x func.return %0 : tensor<3x4x6xf32> // CHECK-LABEL: batchMatMulV3UnbatchedLeft - // CHECK-DAG: %[[SPLITTING_AXIS:.*]] = "tf.Const"() {value = dense<0> : tensor} - // CHECK-DAG: %[[MATMUL_RHS_SHAPE:.*]] = "tf.Const"() {value = dense<[5, 6]> : tensor<2xi64>} + // CHECK-DAG: %[[SPLITTING_AXIS:.*]] = "tf.Const"() <{value = dense<0> : tensor}> + // CHECK-DAG: %[[MATMUL_RHS_SHAPE:.*]] = "tf.Const"() <{value = dense<[5, 6]> : tensor<2xi64>}> // CHECK: %[[RHS_RESHAPED:.*]]:3 = "tf.Split"(%[[SPLITTING_AXIS]], %arg1) : (tensor, tensor<3x5x6xf32>) -> (tensor<1x5x6xf32>, tensor<1x5x6xf32>, tensor<1x5x6xf32>) // CHECK: %[[RHS_1:.*]] = "tf.Reshape"(%[[RHS_RESHAPED]]#0, %[[MATMUL_RHS_SHAPE]]) : (tensor<1x5x6xf32>, tensor<2xi64>) -> tensor<5x6xf32> // CHECK: %[[RHS_2:.*]] = "tf.Reshape"(%[[RHS_RESHAPED]]#1, %[[MATMUL_RHS_SHAPE]]) : (tensor<1x5x6xf32>, tensor<2xi64>) -> tensor<5x6xf32> // CHECK: %[[RHS_3:.*]] = "tf.Reshape"(%[[RHS_RESHAPED]]#2, %[[MATMUL_RHS_SHAPE]]) : (tensor<1x5x6xf32>, tensor<2xi64>) -> tensor<5x6xf32> - // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%arg0, %[[RHS_1]]) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_2:.*]] = "tf.MatMul"(%arg0, %[[RHS_2]]) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_3:.*]] = "tf.MatMul"(%arg0, %[[RHS_3]]) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%arg0, %[[RHS_1]]) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_2:.*]] = "tf.MatMul"(%arg0, %[[RHS_2]]) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_3:.*]] = "tf.MatMul"(%arg0, %[[RHS_3]]) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_PACKED:.*]] = "tf.Pack"(%[[MATMUL_1]], %[[MATMUL_2]], %[[MATMUL_3]]) {axis = 0 : i64} : (tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>) -> tensor<3x4x6xf32> + // CHECK: %[[MATMUL_PACKED:.*]] = "tf.Pack"(%[[MATMUL_1]], %[[MATMUL_2]], %[[MATMUL_3]]) <{axis = 0 : i64}> : (tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>) -> tensor<3x4x6xf32> // CHECK: return %[[MATMUL_PACKED]] : tensor<3x4x6xf32> } @@ -664,19 +664,19 @@ func.func @batchMatMulV3UnbatchedRight(%arg0: tensor<3x4x5xf32>, %arg1: tensor<5 func.return %0 : tensor<3x4x6xf32> // CHECK-LABEL: batchMatMulV3UnbatchedRight - // CHECK-DAG: %[[SPLITTING_AXIS:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor - // CHECK-DAG: %[[MATMUL_LHS_SHAPE:.*]] = "tf.Const"() {value = dense<[4, 5]> : tensor<2xi64>} : () -> tensor<2xi64> + // CHECK-DAG: %[[SPLITTING_AXIS:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor + // CHECK-DAG: %[[MATMUL_LHS_SHAPE:.*]] = "tf.Const"() <{value = dense<[4, 5]> : tensor<2xi64>}> : () -> tensor<2xi64> // CHECK: %[[LHS_SPLIT:.*]]:3 = "tf.Split"(%[[SPLITTING_AXIS]], %arg0) : (tensor, tensor<3x4x5xf32>) -> (tensor<1x4x5xf32>, tensor<1x4x5xf32>, tensor<1x4x5xf32>) // CHECK: %[[LHS_1:.*]] = "tf.Reshape"(%[[LHS_SPLIT]]#0, %[[MATMUL_LHS_SHAPE]]) : (tensor<1x4x5xf32>, tensor<2xi64>) -> tensor<4x5xf32> // CHECK: %[[LHS_2:.*]] = "tf.Reshape"(%[[LHS_SPLIT]]#1, %[[MATMUL_LHS_SHAPE]]) : (tensor<1x4x5xf32>, tensor<2xi64>) -> tensor<4x5xf32> // CHECK: %[[LHS_3:.*]] = "tf.Reshape"(%[[LHS_SPLIT]]#2, %[[MATMUL_LHS_SHAPE]]) : (tensor<1x4x5xf32>, tensor<2xi64>) -> tensor<4x5xf32> - // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%[[LHS_1]], %arg1) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_2:.*]] = "tf.MatMul"(%[[LHS_2]], %arg1) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_3:.*]] = "tf.MatMul"(%[[LHS_3]], %arg1) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%[[LHS_1]], %arg1) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_2:.*]] = "tf.MatMul"(%[[LHS_2]], %arg1) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_3:.*]] = "tf.MatMul"(%[[LHS_3]], %arg1) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> - // CHECK: %[[MATMUL_PACKED:.*]] = "tf.Pack"(%[[MATMUL_1]], %[[MATMUL_2]], %[[MATMUL_3]]) {axis = 0 : i64} : (tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>) -> tensor<3x4x6xf32> + // CHECK: %[[MATMUL_PACKED:.*]] = "tf.Pack"(%[[MATMUL_1]], %[[MATMUL_2]], %[[MATMUL_3]]) <{axis = 0 : i64}> : (tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>) -> tensor<3x4x6xf32> // CHECK: return %[[MATMUL_PACKED]] : tensor<3x4x6xf32> } @@ -687,7 +687,7 @@ func.func @batchMatMulV3Matrix(%arg0: tensor<4x5xf32>, %arg1: tensor<5x6xf32>) - func.return %0 : tensor<4x6xf32> // CHECK-LABEL: batchMatMulV3Matrix - // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%arg0, %arg1) {transpose_a = false, transpose_b = false} : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%arg0, %arg1) <{transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32> // CHECK: return %[[MATMUL_1]] : tensor<4x6xf32> } @@ -698,7 +698,7 @@ func.func @batchMatMulV3MatrixAdjXY(%arg0: tensor<5x4xf32>, %arg1: tensor<6x5xf3 func.return %0 : tensor<4x6xf32> // CHECK-LABEL: batchMatMulV3MatrixAdjXY - // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%arg0, %arg1) {transpose_a = true, transpose_b = true} : (tensor<5x4xf32>, tensor<6x5xf32>) -> tensor<4x6xf32> + // CHECK: %[[MATMUL_1:.*]] = "tf.MatMul"(%arg0, %arg1) <{transpose_a = true, transpose_b = true}> : (tensor<5x4xf32>, tensor<6x5xf32>) -> tensor<4x6xf32> // CHECK: return %[[MATMUL_1]] : tensor<4x6xf32> } diff --git a/tensorflow/compiler/mlir/tensorflow/tests/update_control_dependencies.mlir b/tensorflow/compiler/mlir/tensorflow/tests/update_control_dependencies.mlir index 263a6762238823..09931d9fc2adb6 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/update_control_dependencies.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/update_control_dependencies.mlir @@ -80,7 +80,7 @@ func.func @incorrect_control_deps_replaced_with_correct_controls(%arg0: tensor<* // CHECK: %[[GRAPH:.*]]:2 = tf_executor.graph { // CHECK: %[[ADD1:.*]], %[[ADD1_control:.*]] = tf_executor.island wraps "tf.Add"(%arg0, %arg1) // CHECK: %[[ADD2:.*]], %[[ADD2_control:.*]] = tf_executor.island wraps "tf.Add"(%[[ADD1]], %arg1) -// CHECK: %[[PRINT:.*]], %[[PRINT_control:.*]] = tf_executor.island wraps "tf.Print"(%[[ADD2]]) {message = "add2 result"} +// CHECK: %[[PRINT:.*]], %[[PRINT_control:.*]] = tf_executor.island wraps "tf.Print"(%[[ADD2]]) <{message = "add2 result"}> // CHECK: tf_executor.fetch %[[ADD1]], %[[ADD2]], %[[PRINT_control]] : // CHECK: } // CHECK: return %[[GRAPH]]#0, %[[GRAPH]]#1 @@ -99,7 +99,7 @@ func.func @trailing_print(%arg0: tensor<*xi32>, %arg1: tensor) -> (tensor<* // CHECK: %[[GRAPH:.*]]:2 = tf_executor.graph { // CHECK: %[[ADD1:.*]], %[[ADD1_control:.*]] = tf_executor.island wraps "tf.Add"(%arg0, %arg1) // CHECK: %[[ADD2:.*]], %[[ADD2_control:.*]] = tf_executor.island wraps "tf.Add"(%[[ADD1]], %arg1) -// CHECK: %[[PRINT:.*]], %[[PRINT_control:.*]] = tf_executor.island wraps "tf.Print"(%[[ADD2]]) {message = "add2 result"} +// CHECK: %[[PRINT:.*]], %[[PRINT_control:.*]] = tf_executor.island wraps "tf.Print"(%[[ADD2]]) <{message = "add2 result"}> // CHECK: tf_executor.fetch %[[ADD1]], %[[ADD2]], %[[PRINT_control]] : // CHECK: } // CHECK: return %[[GRAPH]]#0, %[[GRAPH]]#1 @@ -127,7 +127,7 @@ func.func @non_aliasing_reads_writes( // CHECK-DAG: %[[READ0:.*]], %[[READ0_CONTROL:.*]] = tf_executor.island wraps "tf.ReadVariableOp"(%arg0) // CHECK-DAG: %[[ASSIGN0_CONTROL:.*]] = tf_executor.island(%[[READ0_CONTROL]]) wraps "tf.AssignVariableOp"(%arg0, %arg2) // CHECK-DAG: %[[READ1:.*]], %[[READ1_CONTROL:.*]] = tf_executor.island wraps "tf.ReadVariableOp"(%arg1) -// CHECK: %[[VH0:.*]], %[[VH0_CONTROL:.*]] = tf_executor.island wraps "tf.VarHandleOp"() {container = "c", shared_name = "v0"} +// CHECK: %[[VH0:.*]], %[[VH0_CONTROL:.*]] = tf_executor.island wraps "tf.VarHandleOp"() <{container = "c", shared_name = "v0"}> // CHECK: %[[READ2:.*]], %[[READ2_CONTROL:.*]] = tf_executor.island wraps "tf.ReadVariableOp"(%[[VH0]]) // CHECK: %[[ASSIGN1_CONTROL:.*]] = tf_executor.island(%[[READ1_CONTROL]]) wraps "tf.AssignVariableOp"(%arg1, %[[READ0]]) // CHECK: %[[ASSIGN2_CONTROL:.*]] = tf_executor.island(%[[ASSIGN0_CONTROL]]) wraps "tf.AssignVariableOp"(%arg0, %[[READ2]]) @@ -151,8 +151,8 @@ func.func @unknown_side_effecting_op(%arg0: tensor<32xf32>) { } // CHECK-LABEL: func @unknown_side_effecting_op // CHECK: tf_executor.graph { -// CHECK: %[[VH0:.*]], %[[VH0_CONTROL:.*]] = tf_executor.island wraps "tf.VarHandleOp"() {container = "c", shared_name = "v0"} -// CHECK: %[[VH1:.*]], %[[VH1_CONTROL:.*]] = tf_executor.island wraps "tf.VarHandleOp"() {container = "c", shared_name = "v1"} +// CHECK: %[[VH0:.*]], %[[VH0_CONTROL:.*]] = tf_executor.island wraps "tf.VarHandleOp"() <{container = "c", shared_name = "v0"}> +// CHECK: %[[VH1:.*]], %[[VH1_CONTROL:.*]] = tf_executor.island wraps "tf.VarHandleOp"() <{container = "c", shared_name = "v1"}> // CHECK: %[[READ0:.*]], %[[READ0_CONTROL:.*]] = tf_executor.island wraps "tf.ReadVariableOp"(%[[VH0]]) // CHECK: %[[ASSIGN0_CONTROL:.*]] = tf_executor.island wraps "tf.AssignVariableOp"(%[[VH1]], %arg0) // CHECK: %[[UNKNOWN_CONTROL:.*]] = tf_executor.island(%[[READ0_CONTROL]], %[[ASSIGN0_CONTROL]]) wraps "tf._UnknownSideEffectingOp_"() @@ -170,7 +170,7 @@ func.func @single_op_island_forward_block_arg(%arg0: tensor) -> (t } // CHECK-LABEL: func @single_op_island_forward_block_arg // CHECK: tf_executor.graph { -// CHECK: %[[outputs:.*]], %[[control:.*]] = tf_executor.island wraps "tf.Const"() {value = dense<0.000000e+00> : tensor<2048xf32>} : () -> tensor<2048xf32> +// CHECK: %[[outputs:.*]], %[[control:.*]] = tf_executor.island wraps "tf.Const"() <{value = dense<0.000000e+00> : tensor<2048xf32>}> : () -> tensor<2048xf32> // CHECK: tf_executor.fetch %[[outputs]], %arg0 : tensor<2048xf32>, tensor func.func @tpu_load_embedding_ops_sink_controls(%arg0: tensor<*x!tf_type.resource>>, %arg1: tensor<*x!tf_type.resource>>, %arg2: tensor<*x!tf_type.resource>>, %arg3: tensor<*x!tf_type.resource>>) { @@ -194,13 +194,13 @@ func.func @tpu_load_embedding_ops_sink_controls(%arg0: tensor<*x!tf_type.resourc // CHECK: %[[outputs:.*]], %[[control:.*]] = tf_executor.island wraps "tf.ReadVariableOp"(%arg0) {device = ""} : (tensor<*x!tf_type.resource>>) -> tensor<8xf32> // CHECK: %[[outputs_0:.*]], %[[control_1:.*]] = tf_executor.island wraps "tf.ReadVariableOp"(%arg1) {device = ""} : (tensor<*x!tf_type.resource>>) -> tensor<8xf32> // CHECK: %[[outputs_2:.*]], %[[control_3:.*]] = tf_executor.island wraps "tf.ReadVariableOp"(%arg2) {device = ""} : (tensor<*x!tf_type.resource>>) -> tensor<8xf32> -// CHECK: %[[control_4:.*]] = tf_executor.island wraps "tf.LoadTPUEmbeddingAdagradParameters"(%[[outputs]], %[[outputs_0]]) {config = "", num_shards = 1 : i64, shard_id = 0 : i64, table_id = -1 : i64, table_name = "table1"} : (tensor<8xf32>, tensor<8xf32>) -> () +// CHECK: %[[control_4:.*]] = tf_executor.island wraps "tf.LoadTPUEmbeddingAdagradParameters"(%[[outputs]], %[[outputs_0]]) <{config = "", num_shards = 1 : i64, shard_id = 0 : i64, table_id = -1 : i64, table_name = "table1"}> : (tensor<8xf32>, tensor<8xf32>) -> () // CHECK: %[[outputs_5:.*]], %[[control_6:.*]] = tf_executor.island wraps "tf.ReadVariableOp"(%arg3) {device = ""} : (tensor<*x!tf_type.resource>>) -> tensor<8xf32> -// CHECK: %[[control_7:.*]] = tf_executor.island wraps "tf.LoadTPUEmbeddingAdagradParameters"(%[[outputs_2]], %[[outputs_5]]) {config = "", num_shards = 1 : i64, shard_id = 0 : i64, table_id = -1 : i64, table_name = "table2"} : (tensor<8xf32>, tensor<8xf32>) -> () +// CHECK: %[[control_7:.*]] = tf_executor.island wraps "tf.LoadTPUEmbeddingAdagradParameters"(%[[outputs_2]], %[[outputs_5]]) <{config = "", num_shards = 1 : i64, shard_id = 0 : i64, table_id = -1 : i64, table_name = "table2"}> : (tensor<8xf32>, tensor<8xf32>) -> () // CHECK: %[[control_8:.*]] = tf_executor.island(%[[control]], %[[control_1]], %[[control_3]], %[[control_4]], %[[control_6]], %[[control_7]]) wraps "tf.UnknownOp"() : () -> () // CHECK: %[[control_9:.*]] = tf_executor.island(%[[control_8]]) wraps "tf.UnknownOp"() : () -> () -// CHECK: %[[control_10:.*]] = tf_executor.island(%[[control_9]]) wraps "tf.LoadTPUEmbeddingAdagradParameters"(%[[outputs]], %[[outputs_0]]) {config = "", num_shards = 1 : i64, shard_id = 0 : i64, table_id = -1 : i64, table_name = "table3"} : (tensor<8xf32>, tensor<8xf32>) -> () -// CHECK: %[[control_11:.*]] = tf_executor.island(%[[control_9]]) wraps "tf.LoadTPUEmbeddingAdagradParameters"(%[[outputs_2]], %[[outputs_5]]) {config = "", num_shards = 1 : i64, shard_id = 0 : i64, table_id = -1 : i64, table_name = "table4"} : (tensor<8xf32>, tensor<8xf32>) -> () +// CHECK: %[[control_10:.*]] = tf_executor.island(%[[control_9]]) wraps "tf.LoadTPUEmbeddingAdagradParameters"(%[[outputs]], %[[outputs_0]]) <{config = "", num_shards = 1 : i64, shard_id = 0 : i64, table_id = -1 : i64, table_name = "table3"}> : (tensor<8xf32>, tensor<8xf32>) -> () +// CHECK: %[[control_11:.*]] = tf_executor.island(%[[control_9]]) wraps "tf.LoadTPUEmbeddingAdagradParameters"(%[[outputs_2]], %[[outputs_5]]) <{config = "", num_shards = 1 : i64, shard_id = 0 : i64, table_id = -1 : i64, table_name = "table4"}> : (tensor<8xf32>, tensor<8xf32>) -> () // CHECK: tf_executor.fetch %[[control_10]], %[[control_11]] : !tf_executor.control, !tf_executor.control // ----- diff --git a/tensorflow/compiler/mlir/tensorflow/tests/xla_call_module_deserialization.mlir b/tensorflow/compiler/mlir/tensorflow/tests/xla_call_module_deserialization.mlir index 982b44fd86ba9f..82f999c46261a1 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/xla_call_module_deserialization.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/xla_call_module_deserialization.mlir @@ -18,9 +18,9 @@ module { // CHECK-SAME: %[[ARG0:.*]]: tensor<10xi32>, %[[ARG1:.*]]: tensor<10xi32> func.func @main(%arg0: tensor<10xi32>, %arg1: tensor<10xi32>) -> tensor<10xi32> { // CHECK: %[[RESULT:.*]] = "tf.XlaCallModule"(%[[ARG0]], %[[ARG1]]) - // CHECK-SAME: _entry_function = @main0, // CHECK-NOT: function_list // CHECK-SAME: module = "" + // CHECK-SAME: _entry_function = @main0, // `module` is stablehlo bytecode for: // func.func @main(%arg0: tensor {jax.arg_info = "x", mhlo.sharding = "{replicated}"}, %arg1: tensor<*xi32>) -> (tensor {jax.result_info = ""}) { @@ -36,9 +36,9 @@ module { // CHECK-SAME: %[[ARG0:.*]]: tensor<10xi32>, %[[ARG1:.*]]: tensor<10xi32> func.func @foo(%arg0: tensor<10xi32>, %arg1: tensor<10xi32>) -> tensor<10xi32> { // CHECK: %[[RESULT:.*]] = "tf.XlaCallModule"(%[[ARG0]], %[[ARG1]]) - // CHECK-SAME: _entry_function = @main1, // CHECK-NOT: function_list // CHECK-SAME: module = "" + // CHECK-SAME: _entry_function = @main1, // `module` is stablehlo bytecode for: // func.func @main(%arg0: tensor {jax.arg_info = "x", mhlo.sharding = "{replicated}"}, %arg1: tensor<*xi32>) -> (tensor {jax.result_info = ""}) { diff --git a/tensorflow/compiler/mlir/tensorflow/tests/xla_call_module_round_trip.mlir b/tensorflow/compiler/mlir/tensorflow/tests/xla_call_module_round_trip.mlir index 83be1fdbe42285..f19d295c0b33ae 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/xla_call_module_round_trip.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/xla_call_module_round_trip.mlir @@ -13,12 +13,12 @@ module { func.func @main(%arg0: tensor<10xi32>, %arg1: tensor<10xi32>) -> tensor<10xi32> { // CHECK: %[[RESULT:.*]] = "tf.XlaCallModule"(%[[ARG0]], %[[ARG1]]) // CHECK-SAME: Sout = [#tf_type.shape] - // CHECK-SAME: _entry_function = @main0 - // CHECK-SAME: _stablehlo_module_attrs = {} // CHECK-NOT: function_list // CHECK-SAME: module = "" // CHECK-SAME: platforms = [] // CHECK-SAME: version = 5 + // CHECK-SAME: _entry_function = @main0 + // CHECK-SAME: _stablehlo_module_attrs = {} %0 = "tf.XlaCallModule"(%arg0, %arg1) {Sout = [#tf_type.shape], dim_args_spec = [], _entry_function = @main0, module = "", platforms = [], version = 5 : i64} : (tensor<10xi32>, tensor<10xi32>) -> tensor<10xi32> // CHECK: return %[[RESULT]] diff --git a/tensorflow/compiler/mlir/tensorflow/tests/xla_cluster_formation.mlir b/tensorflow/compiler/mlir/tensorflow/tests/xla_cluster_formation.mlir index 282588a4bd2952..6b0f700c65eb36 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/xla_cluster_formation.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/xla_cluster_formation.mlir @@ -91,8 +91,8 @@ func.func @func(%arg0: tensor) -> tensor { // Check that we encapsulate the function body of entry functions with compilation markers, and not the included partitioned calls with the markers. // CHECK-LABEL: func.func @entry_function_with_compilation_markers(%arg0: tensor) -> tensor attributes {_xla_compile_device_type = "GPU", allow_soft_placement = true, device = "/device:GPU:0", tf.entry_function = {}} { // CHECK: %0 = "tf_device.cluster"() ({ -// CHECK: %1 = "tf.StatefulPartitionedCall"(%arg0) {_xla_compile_device_type = "GPU", config = "", config_proto = "", executor_type = "", f = @stateful_pcall_func} : (tensor) -> tensor -// CHECK: %cst = "tf.Const"() {value = dense<5> : tensor} : () -> tensor +// CHECK: %1 = "tf.StatefulPartitionedCall"(%arg0) <{config = "", config_proto = "", executor_type = "", f = @stateful_pcall_func}> {_xla_compile_device_type = "GPU"} : (tensor) -> tensor +// CHECK: %cst = "tf.Const"() <{value = dense<5> : tensor}> : () -> tensor // CHECK: %2 = "tf.Add"(%1, %cst) : (tensor, tensor) -> tensor // CHECK: tf_device.return %2 : tensor // CHECK: }) {_cluster_outlined_function_name = "entry_function_with_compilation_markers_cluster_func", _xla_compile_device_type = "GPU", allow_soft_placement = true, device = "/device:GPU:0"} : () -> tensor diff --git a/tensorflow/compiler/mlir/tensorflow/tests/xla_rewrite.mlir b/tensorflow/compiler/mlir/tensorflow/tests/xla_rewrite.mlir index 914c753c1ad447..4e0284816a7644 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/xla_rewrite.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/xla_rewrite.mlir @@ -4,7 +4,7 @@ module attributes {tf.devices = ["/job:worker/replica:0/task:0/device:CPU:0", "/job:worker/replica:0/task:0/device:GPU:0"]} { // CHECK-LABEL: func.func @convert_cluster_func func.func @convert_cluster_func(%arg0: tensor) -> tensor { - // CHECK: "tf.XlaLaunch"(%arg0) {function = @func, operandSegmentSizes = array} : (tensor) -> tensor + // CHECK: "tf.XlaLaunch"(%arg0) <{function = @func, operandSegmentSizes = array}> : (tensor) -> tensor %0 = "tf_device.cluster_func"(%arg0) {func = @func} : (tensor) -> tensor func.return %0 : tensor } @@ -19,7 +19,7 @@ module attributes {tf.devices = ["/job:worker/replica:0/task:0/device:CPU:0", "/ module attributes {tf.devices = ["/job:worker/replica:0/task:0/device:CPU:0", "/job:worker/replica:0/task:0/device:GPU:0"]} { // CHECK-LABEL: func.func @convert_cluster_func_with_resources_in_order func.func @convert_cluster_func_with_resources_in_order(%arg0: tensor, %arg1: tensor) -> tensor { - // CHECK: "tf.XlaLaunch"(%arg1, %arg0) {function = @func_with_resources_in_order, operandSegmentSizes = array} : (tensor, tensor) -> tensor + // CHECK: "tf.XlaLaunch"(%arg1, %arg0) <{function = @func_with_resources_in_order, operandSegmentSizes = array}> : (tensor, tensor) -> tensor %0 = "tf_device.cluster_func"(%arg1, %arg0) {func = @func_with_resources_in_order} : (tensor, tensor) -> (tensor) func.return %0 : tensor } @@ -34,9 +34,9 @@ module attributes {tf.devices = ["/job:worker/replica:0/task:0/device:CPU:0", "/ module attributes {tf.devices = ["/job:worker/replica:0/task:0/device:CPU:0", "/job:worker/replica:0/task:0/device:GPU:0"]} { // CHECK-LABEL: func.func @convert_cluster_func_with_resources func.func @convert_cluster_func_with_resources(%arg0: tensor, %arg1: tensor) -> tensor { - // CHECK: "tf.XlaLaunch"(%arg1, %arg0) {function = @func_with_resources, operandSegmentSizes = array} : (tensor, tensor) -> tensor + // CHECK: "tf.XlaLaunch"(%arg1, %arg0) <{function = @func_with_resources, operandSegmentSizes = array}> : (tensor, tensor) -> tensor %0 = "tf_device.cluster_func"(%arg0, %arg1) {func = @func_with_resources} : (tensor, tensor) -> tensor - // CHECK: "tf.XlaLaunch"(%arg1, %arg0) {function = @func_with_resources, operandSegmentSizes = array} : (tensor, tensor) -> tensor + // CHECK: "tf.XlaLaunch"(%arg1, %arg0) <{function = @func_with_resources, operandSegmentSizes = array}> : (tensor, tensor) -> tensor %1 = "tf_device.cluster_func"(%arg0, %arg1) {func = @func_with_resources} : (tensor, tensor) -> tensor return %0 : tensor } diff --git a/tensorflow/compiler/mlir/tensorflow/tests/xla_rewrite_v2.mlir b/tensorflow/compiler/mlir/tensorflow/tests/xla_rewrite_v2.mlir index 7d225cec72502a..e79eb9f0b1fbab 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/xla_rewrite_v2.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/xla_rewrite_v2.mlir @@ -5,7 +5,7 @@ module attributes {tf.devices = ["/job:worker/replica:0/task:0/device:CPU:0", "/ // CHECK-LABEL: func.func @convert_cluster_func func.func @convert_cluster_func(%arg0: tensor) -> tensor { // CHECK: "tf_device.launch"() - // CHECK: "tf._XlaCompile"(%arg0) {function = @func, must_compile = true, operandSegmentSizes = array} : (tensor) -> (tensor<3x!tf_type.string>, tensor) + // CHECK: "tf._XlaCompile"(%arg0) <{function = @func, must_compile = true, operandSegmentSizes = array}> : (tensor) -> (tensor<3x!tf_type.string>, tensor) // CHECK: {device = "/job:localhost/replica:0/task:0/device:GPU:0"} // CHECK: "tf_device.launch"() // CHECK: "tf._XlaRun"(%arg0, %0#0) : (tensor, tensor<3x!tf_type.string>) -> tensor @@ -25,7 +25,7 @@ module attributes {tf.devices = ["/job:worker/replica:0/task:0/device:CPU:0", "/ // CHECK-LABEL: func.func @convert_cluster_func_with_resources_in_order func.func @convert_cluster_func_with_resources_in_order(%arg0: tensor, %arg1: tensor) -> tensor { // CHECK: "tf_device.launch"() - // CHECK: "tf._XlaCompile"(%arg1, %arg0) {function = @func_with_resources_in_order, must_compile = true, operandSegmentSizes = array} : (tensor, tensor) + // CHECK: "tf._XlaCompile"(%arg1, %arg0) <{function = @func_with_resources_in_order, must_compile = true, operandSegmentSizes = array}> : (tensor, tensor) // CHECK: {device = "/job:localhost/replica:0/task:0/device:GPU:0"} // CHECK: "tf_device.launch"() // CHECK: "tf._XlaRun"(%arg1, %arg0, %0#0) : (tensor, tensor, tensor<3x!tf_type.string>) -> tensor @@ -45,14 +45,14 @@ module attributes {tf.devices = ["/job:worker/replica:0/task:0/device:CPU:0", "/ // CHECK-LABEL: func.func @convert_cluster_func_with_resources func.func @convert_cluster_func_with_resources(%arg0: tensor, %arg1: tensor) -> tensor { // CHECK: "tf_device.launch"() - // CHECK: "tf._XlaCompile"(%arg1, %arg0) {function = @func_with_resources_1, must_compile = true, operandSegmentSizes = array} : (tensor, tensor) -> (tensor<3x!tf_type.string>, tensor) + // CHECK: "tf._XlaCompile"(%arg1, %arg0) <{function = @func_with_resources_1, must_compile = true, operandSegmentSizes = array}> : (tensor, tensor) -> (tensor<3x!tf_type.string>, tensor) // CHECK: {device = "/job:localhost/replica:0/task:0/device:GPU:0"} // CHECK: "tf_device.launch"() // CHECK: "tf._XlaRun"(%arg1, %arg0, %0#0) : (tensor, tensor, tensor<3x!tf_type.string>) -> tensor // CHECK: {device = "/job:localhost/replica:0/task:0/device:GPU:0"} : () -> tensor %0 = "tf_device.cluster_func"(%arg0, %arg1) {func = @func_with_resources_1, device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor, tensor) -> tensor // CHECK: "tf_device.launch"() - // CHECK: "tf._XlaCompile"(%arg1, %arg0) {function = @func_with_resources_2, must_compile = true, operandSegmentSizes = array} : (tensor, tensor) -> (tensor<3x!tf_type.string>, tensor) + // CHECK: "tf._XlaCompile"(%arg1, %arg0) <{function = @func_with_resources_2, must_compile = true, operandSegmentSizes = array}> : (tensor, tensor) -> (tensor<3x!tf_type.string>, tensor) // CHECK: {device = "/job:localhost/replica:0/task:0/device:GPU:0"} // CHECK: "tf_device.launch"() // CHECK: "tf._XlaRun"(%arg1, %arg0, %2#0) : (tensor, tensor, tensor<3x!tf_type.string>) -> tensor @@ -77,7 +77,7 @@ module attributes {tf.devices = ["/job:worker/replica:0/task:0/device:CPU:0", "/ module attributes {tf.devices = ["/job:localhost/replica:0/task:0/device:CPU:0"], tf.versions = {producer = 888 : i32}} { func.func @outside_compilation_in_generic_pipeline(%arg0: tensor<2xi32>) -> tensor<2xi32> { // CHECK: tf_device.launch - // CHECK: "tf._XlaCompile"() {function = @func, must_compile = true, operandSegmentSizes = array} + // CHECK: "tf._XlaCompile"() <{function = @func, must_compile = true, operandSegmentSizes = array}> // CHECK: {device = "/job:localhost/replica:0/task:0/device:GPU:0"} // CHECK: tf_device.parallel_execute // CHECK: tf_device.launch diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/BUILD b/tensorflow/compiler/mlir/tensorflow/transforms/BUILD index d69de7494e4b5b..00ab20970f113d 100644 --- a/tensorflow/compiler/mlir/tensorflow/transforms/BUILD +++ b/tensorflow/compiler/mlir/tensorflow/transforms/BUILD @@ -76,7 +76,7 @@ cc_library( ":bridge", ":tensorflow_passes", ":tf_saved_model_passes", - "//tensorflow/compiler/mlir/tf2xla:compile_mlir_util", + "//tensorflow/compiler/mlir/tf2xla/internal:clustering_bridge_passes", "//tensorflow/compiler/mlir/tf2xla/transforms:legalize_tf", "@llvm-project//llvm:Support", "@llvm-project//mlir:IR", diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/bridge.cc b/tensorflow/compiler/mlir/tensorflow/transforms/bridge.cc index d657fc940bc3b5..3e0c19fb87acc3 100644 --- a/tensorflow/compiler/mlir/tensorflow/transforms/bridge.cc +++ b/tensorflow/compiler/mlir/tensorflow/transforms/bridge.cc @@ -31,10 +31,6 @@ limitations under the License. #include "tensorflow/compiler/mlir/tensorflow/utils/data_dumper_logger_config.h" #include "tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.h" #include "tensorflow/compiler/mlir/tensorflow/utils/error_util.h" -#include "tensorflow/compiler/mlir/tf2xla/api/v1/tf_dialect_to_executor.h" -#include "tensorflow/compiler/mlir/tf2xla/api/v2/cluster_tf.h" -#include "tensorflow/compiler/mlir/tf2xla/api/v2/device_type.pb.h" -#include "tensorflow/compiler/mlir/tf2xla/api/v2/tf_dialect_to_executor.h" #include "tensorflow/compiler/mlir/tf2xla/internal/clustering_bridge_passes.h" #include "tensorflow/compiler/mlir/tf2xla/internal/inference/inference_passes.h" #include "tensorflow/compiler/mlir/tf2xla/internal/logging_hooks.h" @@ -48,68 +44,9 @@ limitations under the License. namespace mlir { namespace TFTPU { -namespace { constexpr char kBridgeComponent[] = "TFXLABridge"; -// Run the TF XLA Bridge based on the input pipeline, which can be either TPU -// bridge pipeline or non TPU bridge pipeline. -tensorflow::Status RunTFXLABridge( - ModuleOp module, - llvm::function_ref pipeline_builder, - llvm::StringRef module_name = llvm::StringRef()) { - // Explicitly check that the TensorFlow dialect can constant fold ops. - // Constant folding is essential for the bridge. Without this check, the - // bridge may fail with an error that is difficult to understand and not - // actionable. - if (!TF::TensorFlowDialect::HasConstantFoldHook()) { - return tensorflow::errors::Internal( - "TensorFlow dialect missing constant fold hook in TFXLA bridge phase " - "1; this could happen if the binary doesn't link the constant fold " - "hook registration library."); - } - - PassManager bridge(module.getContext()); - ::tensorflow::applyTensorflowAndCLOptions(bridge); - - // Populate a passmanager with the list of passes that implement the bridge. - pipeline_builder(bridge); - - mlir::StatusScopedDiagnosticHandler diag_handler( - module.getContext(), /*propagate=*/false, - /*filter_stack=*/!VLOG_IS_ON(1)); - - if (VLOG_IS_ON(1) || - DEBUG_DATA_DUMPER()->ShouldDump(module_name.str(), kDebugGroupMain)) { - ::tensorflow::DumpMlirOpToFile( - DEBUG_DATA_DUMPER()->GetDumpFilename(module_name.str(), kDebugGroupMain, - "tf_xla_bridge_before"), - module, llvm::StringRef(), &bridge); - } - - if (VLOG_IS_ON(2) || - DEBUG_DATA_DUMPER()->ShouldDump(module_name.str(), - kDebugGroupBridgePhase1Clustering)) { - ::tensorflow::tf2xla::internal::EnablePassIRPrinting( - bridge, kDebugGroupBridgePhase1Clustering, module_name); - } - - LogicalResult result = bridge.run(module); - (void)result; - - if (VLOG_IS_ON(1) || - DEBUG_DATA_DUMPER()->ShouldDump(module_name.str(), kDebugGroupMain)) { - ::tensorflow::DumpMlirOpToFile( - DEBUG_DATA_DUMPER()->GetDumpFilename(module_name.str(), kDebugGroupMain, - "tf_xla_bridge_after"), - module, llvm::StringRef(), &bridge); - } - - return diag_handler.ConsumeStatus(); -} - -} // namespace - void CreateTPUBridgePipeline(OpPassManager &pm, llvm::StringRef module_name) { pm.addPass(CreateTPUValidateInputsPass()); pm.addNestedPass( @@ -151,17 +88,5 @@ tensorflow::Status RunBridgeWithStandardPipeline(ModuleOp module, return diag_handler.ConsumeStatus(); } -void CreateTFXLABridgePipeline(OpPassManager &pm) { - tensorflow::tf2xla::internal::AddNonTPUBridgeClusteringPipelinePasses(pm); -} - -tensorflow::Status RunTFXLABridge(ModuleOp module, - llvm::StringRef module_name) { - // CPU == GPU here, so both are equivalent. - return tensorflow::tf2xla::v2::RunFunctionTf2xlaClusteringBridge( - module, tensorflow::tf2xla::v2::XLA_GPU_JIT, - /*is_in_fallback_enabled_mode=*/false, module_name); -} - } // namespace TF } // namespace mlir diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/bridge.h b/tensorflow/compiler/mlir/tensorflow/transforms/bridge.h index 72875f2143b75e..671235095b4d7c 100644 --- a/tensorflow/compiler/mlir/tensorflow/transforms/bridge.h +++ b/tensorflow/compiler/mlir/tensorflow/transforms/bridge.h @@ -40,11 +40,6 @@ ABSL_DEPRECATED( tensorflow::Status RunBridgeWithStandardPipeline(ModuleOp module, bool enable_logging, bool enable_inliner); - -// Runs all passes for non TPU (GPU and CPU) graph. -ABSL_DEPRECATED("Use tf2xla::v2::RunFunctionTf2xlaClusteringBridge instead.") -tensorflow::Status RunTFXLABridge( - ModuleOp module, llvm::StringRef module_name = llvm::StringRef()); } // namespace TF } // namespace mlir diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/host_runtime/BUILD b/tensorflow/compiler/mlir/tensorflow/transforms/host_runtime/BUILD index fcf2b1bca58552..359dd5c4624712 100644 --- a/tensorflow/compiler/mlir/tensorflow/transforms/host_runtime/BUILD +++ b/tensorflow/compiler/mlir/tensorflow/transforms/host_runtime/BUILD @@ -29,14 +29,20 @@ cc_library( "//tensorflow/compiler/mlir/tensorflow/transforms:verify_no_outside_compilation_markers_pass", "//tensorflow/core:framework", "//tensorflow/core:lib_proto_parsing", + "//tensorflow/core/platform:error_payloads", + "//tensorflow/core/platform:status", "//tensorflow/core/tpu:tpu_defs", + "@com_google_absl//absl/base:core_headers", "@com_google_absl//absl/log", + "@com_google_absl//absl/status", "@llvm-project//llvm:Support", "@llvm-project//mlir:FuncDialect", "@llvm-project//mlir:IR", "@llvm-project//mlir:Pass", "@llvm-project//mlir:Support", "@llvm-project//mlir:Transforms", + "@local_tsl//tsl/platform:error_logging", + "@local_tsl//tsl/platform:errors", ], ) @@ -56,6 +62,7 @@ tf_cc_test( "//tensorflow/core:framework", "//tensorflow/core:lib", "//tensorflow/core:test", + "//tensorflow/core/lib/monitoring:cell_reader", "//tensorflow/core/platform:resource_loader", "//tensorflow/core/tpu:tpu_defs", "@com_google_absl//absl/status", diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/host_runtime/lower_cluster_to_runtime_ops.cc b/tensorflow/compiler/mlir/tensorflow/transforms/host_runtime/lower_cluster_to_runtime_ops.cc index 4dfe7c1fc19179..cba2b05cc2e78a 100644 --- a/tensorflow/compiler/mlir/tensorflow/transforms/host_runtime/lower_cluster_to_runtime_ops.cc +++ b/tensorflow/compiler/mlir/tensorflow/transforms/host_runtime/lower_cluster_to_runtime_ops.cc @@ -17,6 +17,7 @@ limitations under the License. #include #include "absl/log/log.h" +#include "absl/status/status.h" #include "llvm/ADT/StringRef.h" #include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project #include "mlir/IR/BuiltinOps.h" // from @llvm-project @@ -29,10 +30,15 @@ limitations under the License. #include "tensorflow/compiler/mlir/tensorflow/utils/data_dumper_logger_config.h" #include "tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.h" #include "tensorflow/compiler/mlir/tensorflow/utils/error_util.h" +#include "tensorflow/core/framework/metrics.h" #include "tensorflow/core/framework/types.h" +#include "tensorflow/core/platform/error_payloads.h" +#include "tensorflow/core/platform/status.h" #include "tensorflow/core/tpu/tpu_defs.h" #include "tensorflow/core/util/debug_data_dumper.h" #include "tsl/framework/device_type.h" +#include "tsl/platform/error_logging.h" +#include "tsl/platform/errors.h" namespace tensorflow { namespace tfrt_compiler { @@ -111,6 +117,39 @@ void CreateNonTPULowerClusterToRuntimeOpsPassPipeline( AddNonTPULowerClusterToRuntimeOpsPassPipeline(pm, /*module_name=*/""); } +// TODO(b/306728216): Move this out of the Bridge component and into a Host +// runtime component. +tensorflow::Status RecordIfErrorStatus(const std::string error_prefix, + tsl::DeviceType device_type, + absl::Status status) { + if (status.ok()) { + return status; + } + + VLOG(2) << error_prefix << " " << status; + tensorflow::metrics::UpdateTfMlirBridgeFirstPhaseCounter( + device_type.type_string(), /*bridge_version=*/"v2", + /*fallback_enabled=*/false, + /*result=*/"failure"); + + constexpr char kBridgeComponent[] = "TFXLABridge"; + std::string bridge_subcomponent = "TFXLA_PHASE_ONE_MLIR_TPU_BRIDGE"; + + tsl::OkOrSetErrorCounterPayload( + tensorflow::core::platform::ErrorSourceProto::MLIR_BRIDGE_PHASE_1, + status); + + if (device_type != DeviceType(DEVICE_TPU_XLA_JIT)) { + bridge_subcomponent = "TFXLA_PHASE_ONE_MLIR_CPU/GPU_BRIDGE"; + } + + tsl::error_logging::Log(kBridgeComponent, bridge_subcomponent, + status.ToString()) + .IgnoreError(); + + return status; +} + absl::Status RunLowerClusterToRuntimeOpsPassPipeline( mlir::ModuleOp module, tsl::DeviceType xla_device_type, llvm::StringRef module_name) { @@ -154,7 +193,12 @@ absl::Status RunLowerClusterToRuntimeOpsPassPipeline( module, llvm::StringRef(), &runtime_lowering); } - return diag_handler.ConsumeStatus(); + auto result_status = diag_handler.ConsumeStatus(); + TF_RETURN_IF_ERROR( + RecordIfErrorStatus(/*error_prefix=*/"lower_cluster_to_runtime", + xla_device_type, result_status)); + + return absl::OkStatus(); } // TODO(b/305211853): Unify the CPU/TPU/GPU Execution Ops and thus these two diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/host_runtime/lower_cluster_to_runtime_ops_test.cc b/tensorflow/compiler/mlir/tensorflow/transforms/host_runtime/lower_cluster_to_runtime_ops_test.cc index b58a13d679adb6..e0e376e200e992 100644 --- a/tensorflow/compiler/mlir/tensorflow/transforms/host_runtime/lower_cluster_to_runtime_ops_test.cc +++ b/tensorflow/compiler/mlir/tensorflow/transforms/host_runtime/lower_cluster_to_runtime_ops_test.cc @@ -15,6 +15,7 @@ limitations under the License. #include "tensorflow/compiler/mlir/tensorflow/transforms/host_runtime/lower_cluster_to_runtime_ops.h" +#include #include #include @@ -33,6 +34,7 @@ limitations under the License. #include "tensorflow/compiler/mlir/register_common_dialects.h" #include "tensorflow/compiler/mlir/tensorflow/ir/tf_device.h" #include "tensorflow/compiler/tf2xla/xla_op_registry.h" +#include "tensorflow/core/lib/monitoring/cell_reader.h" #include "tensorflow/core/platform/env.h" #include "tensorflow/core/platform/resource_loader.h" #include "tensorflow/core/platform/test.h" @@ -51,6 +53,7 @@ using mlir::ModuleOp; using mlir::OpPassManager; using mlir::OwningOpRef; using mlir::func::FuncOp; +using ::tensorflow::monitoring::testing::CellReader; using tsl::DeviceType; std::string TestDataPath() { @@ -58,6 +61,9 @@ std::string TestDataPath() { "tensorflow/compiler/mlir/tensorflow/transforms/host_runtime/testdata/"); } +static constexpr char kCompilationStreamz[] = + "/tensorflow/core/tf_mlir_bridge_first_phase_count"; + class LowerClusterToRuntimeOpsTest : public ::testing::Test { public: LowerClusterToRuntimeOpsTest() { @@ -154,11 +160,17 @@ TEST_F(LowerClusterToRuntimeOpsTest, LowersClusterOpsGPU) { } TEST_F(LowerClusterToRuntimeOpsTest, ErrorsWithBadCluster) { + CellReader compilation_status(kCompilationStreamz); + TF_ASSERT_OK(CreateMlirModule("malformed_cluster.mlir")); EXPECT_FALSE(RunLowerClusterToRuntimeOpsPassPipeline( *mlir_module_, DeviceType(DEVICE_TPU_XLA_JIT)) .ok()); + + EXPECT_EQ(compilation_status.Delta("XLA_TPU_JIT", "v2", "fallback_disabled", + "failure"), + 1); } TEST_F(LowerClusterToRuntimeOpsTest, DumpsPipelinePasses) { diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/mlprogram.cc b/tensorflow/compiler/mlir/tensorflow/transforms/mlprogram.cc index aa1a74a1a5a1fb..523cdf290e6772 100644 --- a/tensorflow/compiler/mlir/tensorflow/transforms/mlprogram.cc +++ b/tensorflow/compiler/mlir/tensorflow/transforms/mlprogram.cc @@ -24,13 +24,14 @@ limitations under the License. #include "tensorflow/compiler/mlir/tensorflow/transforms/bridge.h" #include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h" #include "tensorflow/compiler/mlir/tensorflow/transforms/tf_saved_model_passes.h" +#include "tensorflow/compiler/mlir/tf2xla/internal/clustering_bridge_passes.h" #include "tensorflow/compiler/mlir/tf2xla/transforms/passes.h" #include "xla/mlir_hlo/mhlo/transforms/passes.h" namespace tensorflow { void PopulateLowerToMlProgramAndHloPipeline(mlir::OpPassManager& pm) { - mlir::TF::CreateTFXLABridgePipeline(pm); + tensorflow::tf2xla::internal::AddNonTPUBridgeClusteringPipelinePasses(pm); // Remove unused global tensors, or make then immutable if possible. pm.addPass(mlir::tf_saved_model::CreateOptimizeGlobalTensorsPass()); diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/passes.h b/tensorflow/compiler/mlir/tensorflow/transforms/passes.h index 247ea10dd68b50..caa76907a3ee45 100644 --- a/tensorflow/compiler/mlir/tensorflow/transforms/passes.h +++ b/tensorflow/compiler/mlir/tensorflow/transforms/passes.h @@ -318,10 +318,6 @@ std::unique_ptr> CreatePrintPass( // Moves TPUCompileMlir ops as far to the front as possible. std::unique_ptr> CreateMoveTpuCompileToFrontPass(); -// Populates the supplied passmanager with the passes required to run the -// CPU/GPU bridge. -void CreateTFXLABridgePipeline(OpPassManager& pm); - //===----------------------------------------------------------------------===// // XlaCallModule //===----------------------------------------------------------------------===// diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/region_control_flow_to_functional.cc b/tensorflow/compiler/mlir/tensorflow/transforms/region_control_flow_to_functional.cc index d81ee1e8e33b1a..78f2c3e0423124 100644 --- a/tensorflow/compiler/mlir/tensorflow/transforms/region_control_flow_to_functional.cc +++ b/tensorflow/compiler/mlir/tensorflow/transforms/region_control_flow_to_functional.cc @@ -17,9 +17,13 @@ limitations under the License. // the TensorFlow dialect to their functional counterparts, i.e., // tf.IfRegion -> tf.If and tf.WhileRegion -> tf.While +#include +#include #include +#include #include "llvm/ADT/STLExtras.h" +#include "llvm/ADT/SetVector.h" #include "llvm/ADT/SmallVector.h" #include "llvm/Support/Casting.h" #include "llvm/Support/FormatVariadic.h" @@ -27,18 +31,21 @@ limitations under the License. #include "mlir/IR/Attributes.h" // from @llvm-project #include "mlir/IR/Builders.h" // from @llvm-project #include "mlir/IR/BuiltinOps.h" // from @llvm-project +#include "mlir/IR/MLIRContext.h" // from @llvm-project +#include "mlir/IR/Matchers.h" // from @llvm-project #include "mlir/IR/Operation.h" // from @llvm-project +#include "mlir/IR/Region.h" // from @llvm-project #include "mlir/IR/SymbolTable.h" // from @llvm-project #include "mlir/IR/TypeUtilities.h" // from @llvm-project #include "mlir/IR/Value.h" // from @llvm-project #include "mlir/IR/Verifier.h" // from @llvm-project #include "mlir/IR/Visitors.h" // from @llvm-project #include "mlir/Pass/Pass.h" // from @llvm-project -#include "mlir/Pass/PassRegistry.h" // from @llvm-project +#include "mlir/Support/LLVM.h" // from @llvm-project +#include "mlir/Support/LogicalResult.h" // from @llvm-project #include "mlir/Transforms/RegionUtils.h" // from @llvm-project #include "tensorflow/compiler/mlir/op_or_arg_name_mapper.h" #include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h" -#include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h" #include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h" #include "tensorflow/compiler/mlir/tensorflow/utils/attribute_utils.h" @@ -69,6 +76,8 @@ struct RegionControlFlowToFunctional CaseRegionOp case_region); LogicalResult ConvertWhileOp(SymbolTableCollection& symbol_table, WhileRegionOp while_region); + LogicalResult ConvertGeneratorDatasetOp(SymbolTableCollection& symbol_table, + GeneratorDatasetRegionOp regional); // Get unique name by using the loc to name mapping. std::string GetName(Operation* op, StringRef suffix); @@ -124,6 +133,37 @@ void CopyAndOverrideAttributes(Operation* src, Operation* dst, dst->setAttr(kXlaPropagateCompileTimeConsts, builder->getBoolAttr(true)); } +// If the region only does a single function call whose operands / returns match +// exactly the block args and results, return the name of the called function. +std::optional UnwrapSingleFunctionCall(Region& region) { + // The pattern we're matching is + // ^block(arg0, arg1, ..., argN): + // r0, r1, ..., rN = func.call @foo(arg0, arg1, ..., argN) + // "tf.yield"(r0, r1, ..., rN) + if (!region.hasOneBlock()) return std::nullopt; + Block& block = region.front(); + if (std::distance(block.begin(), block.end()) != 2) return std::nullopt; + TF::YieldOp yield = + llvm::dyn_cast_or_null(block.getTerminator()); + if (!yield) return std::nullopt; + func::CallOp call = llvm::dyn_cast_or_null(*block.begin()); + if (!call) return std::nullopt; + if (block.getNumArguments() != call.getNumOperands() || + call.getNumResults() != yield.getNumOperands()) + return std::nullopt; + for (auto [arg, operand] : + llvm::zip(block.getArguments(), call.getOperands())) { + if (arg != operand) return std::nullopt; + } + for (auto [ret, operand] : + llvm::zip(call.getResults(), yield.getOperands())) { + if (ret != operand) return std::nullopt; + } + SymbolRefAttr symbol = call.getCallableForCallee().get(); + if (!symbol) return std::nullopt; + return symbol.getLeafReference(); +} + // Extracts the contents of a region with a single block into a new function. // `extern_values` is the set of external values that the region refers to. // Returns the name of the newly created function. @@ -135,7 +175,13 @@ StringRef ExtractSingleBlockRegion( SymbolTableCollection& symbol_table, Region& region, StringRef name, llvm::SmallVectorImpl& extern_values, llvm::SmallVectorImpl& worklist, - bool extern_values_passthrough, bool only_one_return_value) { + bool extern_values_passthrough, bool only_one_return_value, + bool allow_return_of_existing = false) { + if (allow_return_of_existing && extern_values.empty()) { + auto existing = UnwrapSingleFunctionCall(region); + if (existing) return *existing; + } + ModuleOp module = region.getParentOfType(); auto builder = OpBuilder::atBlockBegin(module.getBody()); auto loc = region.getParentOp()->getLoc(); @@ -524,6 +570,52 @@ LogicalResult RegionControlFlowToFunctional::ConvertWhileOp( return success(); } +// Transform GeneratorDatasetRegion to GeneratorDatasetOp. +LogicalResult RegionControlFlowToFunctional::ConvertGeneratorDatasetOp( + SymbolTableCollection& symbol_table, GeneratorDatasetRegionOp regional) { + mlir::MLIRContext* ctx = regional.getContext(); + std::string init_name, next_name, finalize_name; + + llvm::SmallVector extern_values = + CollectExternValues(regional.getRegions()); + + if (!extern_values.empty()) return failure(); + + init_name = GetName(regional, "_init"); + init_name = ExtractSingleBlockRegion(symbol_table, regional.getInit(), + init_name, extern_values, worklist, + /*extern_values_passthrough=*/false, + /*only_one_return_value=*/false, + /*allow_return_of_existing=*/true); + + next_name = GetName(regional, "_next"); + next_name = ExtractSingleBlockRegion(symbol_table, regional.getNext(), + next_name, extern_values, worklist, + /*extern_values_passthrough=*/false, + /*only_one_return_value=*/false, + /*allow_return_of_existing=*/true); + + finalize_name = GetName(regional, "_finalize"); + finalize_name = + ExtractSingleBlockRegion(symbol_table, regional.getFinalize(), + finalize_name, extern_values, worklist, + /*extern_values_passthrough=*/false, + /*only_one_return_value=*/false, + /*allow_return_of_existing=*/true); + + auto new_op = OpBuilder(regional).create( + regional.getLoc(), regional->getResultTypes(), + regional.getInitFuncOtherArgs(), regional.getNextFuncOtherArgs(), + regional.getFinalizeFuncOtherArgs(), SymbolRefAttr::get(ctx, init_name), + SymbolRefAttr::get(ctx, next_name), + SymbolRefAttr::get(ctx, finalize_name), regional.getOutputTypes(), + regional.getOutputShapes(), regional.getMetadata()); + + regional->replaceAllUsesWith(new_op->getResults()); + regional->erase(); + return success(); +} + void RegionControlFlowToFunctional::runOnOperation() { ModuleOp module = getOperation(); SymbolTableCollection symbol_table; @@ -549,6 +641,11 @@ void RegionControlFlowToFunctional::runOnOperation() { op->emitOpError() << "failed to convert to functional form"; return WalkResult::interrupt(); } + } else if (auto gen = llvm::dyn_cast(op)) { + if (failed(ConvertGeneratorDatasetOp(symbol_table, gen))) { + op->emitOpError() << "failed to convert to functional form"; + return WalkResult::interrupt(); + } } return WalkResult::advance(); }); diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/shape_inference.cc b/tensorflow/compiler/mlir/tensorflow/transforms/shape_inference.cc index 581e04a3b8f643..ea51b928562a46 100644 --- a/tensorflow/compiler/mlir/tensorflow/transforms/shape_inference.cc +++ b/tensorflow/compiler/mlir/tensorflow/transforms/shape_inference.cc @@ -1257,7 +1257,7 @@ bool ShapeInference::InferShapeForXlaCallModule(XlaCallModuleOp op) { bool changed = false; for (auto [result, type] : - llvm::zip(op.getResults(), loader->output_types())) { + llvm::zip(op.getResults(), loader->OutputTypes())) { auto ranked = type.dyn_cast(); if (ranked == nullptr) { LLVM_DEBUG(llvm::dbgs() diff --git a/tensorflow/compiler/mlir/tf2xla/api/v1/BUILD b/tensorflow/compiler/mlir/tf2xla/api/v1/BUILD index bd581f689586ec..d0653b9677d5c7 100644 --- a/tensorflow/compiler/mlir/tf2xla/api/v1/BUILD +++ b/tensorflow/compiler/mlir/tf2xla/api/v1/BUILD @@ -233,14 +233,19 @@ cc_library( "//tensorflow/compiler/mlir/tensorflow/transforms:verify_no_outside_compilation_markers_pass", "//tensorflow/compiler/mlir/tf2xla/internal:logging_hooks", "//tensorflow/core:framework", + "//tensorflow/core/platform:error_payloads", "//tensorflow/core/platform:status", "@com_google_absl//absl/base:core_headers", + "@com_google_absl//absl/log", + "@com_google_absl//absl/status", "@llvm-project//llvm:Support", "@llvm-project//mlir:FuncDialect", "@llvm-project//mlir:IR", "@llvm-project//mlir:Pass", "@llvm-project//mlir:Support", "@llvm-project//mlir:Transforms", + "@local_tsl//tsl/lib/monitoring:counter", + "@local_tsl//tsl/platform:error_logging", "@local_tsl//tsl/platform:status", ], ) @@ -255,6 +260,7 @@ tf_cc_test( deps = [ ":tf_dialect_to_executor", "//tensorflow/compiler/mlir:register_common_dialects", + "//tensorflow/core/lib/monitoring:cell_reader", "//tensorflow/core/platform:resource_loader", "@com_google_absl//absl/status", "@com_google_absl//absl/strings", diff --git a/tensorflow/compiler/mlir/tf2xla/api/v1/cluster_tf.cc b/tensorflow/compiler/mlir/tf2xla/api/v1/cluster_tf.cc index 105b9016bb3c24..2f8469ee3f6f69 100644 --- a/tensorflow/compiler/mlir/tf2xla/api/v1/cluster_tf.cc +++ b/tensorflow/compiler/mlir/tf2xla/api/v1/cluster_tf.cc @@ -202,41 +202,6 @@ absl::Status RunClusteringPipelineOnSubmodule( return absl::OkStatus(); } -absl::Status RunLowerToRuntimeOpsOnSubmodule(ModuleOp parent_module, - bool is_in_fallback_enabled_mode) { - int num_submodules = 0; - absl::Status runtime_lowering_status; - parent_module.walk([&](ModuleOp submodule) { - if (submodule == parent_module) return mlir::WalkResult::advance(); - num_submodules++; - runtime_lowering_status = - tensorflow::tfrt_compiler::RunLowerClusterToRuntimeOpsPassPipeline( - submodule, tsl::DeviceType(DEVICE_TPU_XLA_JIT)); - if (num_submodules > 1) { - return mlir::WalkResult::interrupt(); - } - - return mlir::WalkResult::advance(); - }); - - if (num_submodules > 1) { - auto num_submodules_error = absl::InternalError( - "Lower to runtime has more than one submodule. Erroring out."); - TF_RETURN_IF_ERROR(RecordStatusIfError( - /*error_prefix=*/"V1 Lowering to runtime has more than one submodule:", - is_in_fallback_enabled_mode, num_submodules_error)); - } - - if (!runtime_lowering_status.ok()) { - TF_RETURN_IF_ERROR(RecordStatusIfError( - /*error_prefix=*/ - "Errored running lowering cluster ops to runtime ops pipeline:", - is_in_fallback_enabled_mode, runtime_lowering_status)); - } - - return absl::OkStatus(); -} - tensorflow::Status RunSessionTf2xlaClusteringBridge( ModuleOp module, bool is_in_fallback_enabled_mode) { VLOG(2) << "TPU Sessions Bridge called stack trace is " @@ -254,35 +219,12 @@ tensorflow::Status RunSessionTf2xlaClusteringBridge( TF_RETURN_IF_ERROR( RunClusteringPipelineOnSubmodule(module, is_in_fallback_enabled_mode)); - TF_RETURN_IF_ERROR( - RunLowerToRuntimeOpsOnSubmodule(module, is_in_fallback_enabled_mode)); - - Status export_preparation_status = RunTFXLABridge( - module, - [](OpPassManager &pm) { - pm.addPass( - mlir::tf_executor::CreateTFExecutorTPUV1IslandInliningPass()); - // There are cases where we don't consume all compilation and - // replication attributes like we do for the V2 pipeline, so we need to - // convert them from unified to legacy attributes before they get - // exposed to outside of the bridge. - pm.addNestedPass( - mlir::TFTPU:: - CreateConvertToLegacyCompileAndReplicateAttributesPass()); - }, - /*module_name=*/"", - /*dump_prefix=*/"tf_xla_bridge_v1_export_preparation"); - - TF_RETURN_IF_ERROR(RecordStatusIfError( - /*error_prefix=*/"Bridge Export Preparation Failed:", - is_in_fallback_enabled_mode, export_preparation_status)); - tensorflow::metrics::UpdateTfMlirBridgeFirstPhaseCounter( /*device_type=*/"tpu", /*bridge_version=*/"v1", /*n_fallback_enabled*/ is_in_fallback_enabled_mode, /*result=*/"success"); - return tensorflow::tf2xla::v1::ExportFromTensorflowDialectToExecutor(module); + return absl::OkStatus(); } // Registers a pipeline builder function for TF TPU V1 bridge. diff --git a/tensorflow/compiler/mlir/tf2xla/api/v1/cluster_tf.h b/tensorflow/compiler/mlir/tf2xla/api/v1/cluster_tf.h index 441e509ed187db..e27ec14770b2be 100644 --- a/tensorflow/compiler/mlir/tf2xla/api/v1/cluster_tf.h +++ b/tensorflow/compiler/mlir/tf2xla/api/v1/cluster_tf.h @@ -29,11 +29,11 @@ namespace v1 { // These transformations take as input a Tensorflow Graph as an MLIR Module // and transforms the module in place to cluster the given ops for compilation // that is compatible with the given device_type. The MLIR should be in the TF -// Executor Dialect for graph nodes and edges. Individual Op inside a node -// should be the Tensorflow Dialect. The output MLIR is in the TF Executor -// Dialect. The input MLIR should not have infeed and outfeed ops, which are -// unsupported via this API. -// Returns OkStatus if passed, otherwise an error. +// Executor Dialect for graph nodes and edges or TF Functional. It will convert +// to TF Functional internally. Individual Op inside a node should be the +// Tensorflow Dialect. The output MLIR is in the TF Functional Dialect. The +// input MLIR should not have infeed and outfeed ops, which are unsupported via +// this API. Returns OkStatus if passed, otherwise an error. tensorflow::Status RunSessionTf2xlaClusteringBridge( mlir::ModuleOp module, bool is_in_fallback_enabled_mode); diff --git a/tensorflow/compiler/mlir/tf2xla/api/v1/cluster_tf_test.cc b/tensorflow/compiler/mlir/tf2xla/api/v1/cluster_tf_test.cc index 3b195bd3977fe1..44eafb25f579c8 100644 --- a/tensorflow/compiler/mlir/tf2xla/api/v1/cluster_tf_test.cc +++ b/tensorflow/compiler/mlir/tf2xla/api/v1/cluster_tf_test.cc @@ -88,18 +88,6 @@ TEST_F(SessionClusterTensorflowDialectTest, ClustersTf) { compilation_status.Delta("tpu", "v1", "fallback_disabled", "success"), 1); } -// Required for now due to the Bridge API, but this should be separated out -// later. -TEST_F(SessionClusterTensorflowDialectTest, - RunsTensorflowDialectToTensorflowExecutor) { - TF_ASSERT_OK(CreateMlirModule("invalid_executor.mlir")); - - EXPECT_FALSE( - RunSessionTf2xlaClusteringBridge(*mlir_module_, - /*is_in_fallback_enabled_mode=*/false) - .ok()); -} - TEST_F(SessionClusterTensorflowDialectTest, FailsWithMultipleSubmodules) { CellReader compilation_status(kCompilationStreamz); diff --git a/tensorflow/compiler/mlir/tf2xla/api/v1/tf_dialect_to_executor.cc b/tensorflow/compiler/mlir/tf2xla/api/v1/tf_dialect_to_executor.cc index 941dd362a6c8d4..236282f625e20a 100644 --- a/tensorflow/compiler/mlir/tf2xla/api/v1/tf_dialect_to_executor.cc +++ b/tensorflow/compiler/mlir/tf2xla/api/v1/tf_dialect_to_executor.cc @@ -19,6 +19,8 @@ limitations under the License. #include #include +#include "absl/log/log.h" +#include "absl/status/status.h" #include "llvm/ADT/StringRef.h" #include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project #include "mlir/IR/BuiltinOps.h" // from @llvm-project @@ -32,8 +34,11 @@ limitations under the License. #include "tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.h" #include "tensorflow/compiler/mlir/tensorflow/utils/error_util.h" #include "tensorflow/compiler/mlir/tf2xla/internal/logging_hooks.h" +#include "tensorflow/core/platform/error_payloads.h" #include "tensorflow/core/platform/status.h" #include "tensorflow/core/util/debug_data_dumper.h" +#include "tsl/lib/monitoring/counter.h" +#include "tsl/platform/error_logging.h" #include "tsl/platform/status.h" namespace tensorflow { @@ -47,6 +52,15 @@ using mlir::Pass; using mlir::PassManager; using mlir::func::FuncOp; +auto *tf_dialect_to_executor_dialect_status = tsl::monitoring::Counter<1>::New( + "/tensorflow/core/tf2xla/api/v1/tf_dialect_to_executor_dialect_status", + "Counts how often a successful export from TF Dialect to Executor Dialect " + "is", + "status"); + +constexpr char kExportSuccess[] = "success"; +constexpr char kExportFailed[] = "failed"; + namespace { void AddTfDialectToExecutorPasses(OpPassManager &pm) { @@ -55,6 +69,14 @@ void AddTfDialectToExecutorPasses(OpPassManager &pm) { pm.addPass(mlir::CreateBreakUpIslandsPass()); }; + pm.addPass(mlir::tf_executor::CreateTFExecutorTPUV1IslandInliningPass()); + // There are cases where we don't consume all compilation and + // replication attributes like we do for the V2 pipeline, so we need to + // convert them from unified to legacy attributes before they get + // exposed to outside of the bridge. + pm.addNestedPass( + mlir::TFTPU::CreateConvertToLegacyCompileAndReplicateAttributesPass()); + pm.addPass(mlir::TF::CreateTFRegionControlFlowToFunctional()); add_pass(mlir::CreateFunctionalToExecutorDialectConversionPass()); add_pass(mlir::TFDevice::CreateReplicateToIslandPass( @@ -75,6 +97,30 @@ void AddTfDialectToExecutorPasses(OpPassManager &pm) { pm.addPass(mlir::TF::CreateVerifySuitableForExportPass()); } +tensorflow::Status RecordStatusIfError(absl::Status status) { + if (status.ok()) { + return absl::OkStatus(); + } + + VLOG(1) << "Failed to export from TF Dialect to TF Executor Dialect. " + << status; + tf_dialect_to_executor_dialect_status->GetCell(kExportFailed)->IncrementBy(1); + + constexpr char bridge_subcomponent[] = + "TFXLA_TF_FUNCTIONAL_TO_EXECUTOR_EXPORT_v1"; + constexpr char kBridgeComponent[] = "TFXLABridge"; + + tsl::OkOrSetErrorCounterPayload( + tensorflow::core::platform::ErrorSourceProto::MLIR_BRIDGE_PHASE_1, + status); + + tsl::error_logging::Log(kBridgeComponent, bridge_subcomponent, + status.ToString()) + .IgnoreError(); + + return status; +} + } // namespace tensorflow::Status ExportFromTensorflowDialectToExecutor( @@ -116,6 +162,13 @@ tensorflow::Status ExportFromTensorflowDialectToExecutor( module, llvm::StringRef(), &tf_to_executor); } + if (result.failed()) { + return RecordStatusIfError(diag_handler.ConsumeStatus()); + } + + tf_dialect_to_executor_dialect_status->GetCell(kExportSuccess) + ->IncrementBy(1); + return diag_handler.ConsumeStatus(); } diff --git a/tensorflow/compiler/mlir/tf2xla/api/v1/tf_dialect_to_executor_test.cc b/tensorflow/compiler/mlir/tf2xla/api/v1/tf_dialect_to_executor_test.cc index 80a770169a88b2..38393d3753146e 100644 --- a/tensorflow/compiler/mlir/tf2xla/api/v1/tf_dialect_to_executor_test.cc +++ b/tensorflow/compiler/mlir/tf2xla/api/v1/tf_dialect_to_executor_test.cc @@ -15,6 +15,7 @@ limitations under the License. #include "tensorflow/compiler/mlir/tf2xla/api/v1/tf_dialect_to_executor.h" +#include #include #include @@ -26,6 +27,7 @@ limitations under the License. #include "mlir/IR/OwningOpRef.h" // from @llvm-project #include "mlir/Parser/Parser.h" // from @llvm-project #include "tensorflow/compiler/mlir/register_common_dialects.h" +#include "tensorflow/core/lib/monitoring/cell_reader.h" #include "tensorflow/core/platform/resource_loader.h" #include "tsl/lib/core/status_test_util.h" #include "tsl/platform/status.h" @@ -41,6 +43,10 @@ using mlir::DialectRegistry; using mlir::MLIRContext; using mlir::ModuleOp; using mlir::OwningOpRef; +using tensorflow::monitoring::testing::CellReader; + +static constexpr char kCompilationStreamz[] = + "/tensorflow/core/tf2xla/api/v1/tf_dialect_to_executor_dialect_status"; std::string TestDataPath() { return tensorflow::GetDataDependencyFilepath( @@ -73,15 +79,23 @@ class TensorflowDialectToExecutorTest : public ::testing::Test { }; TEST_F(TensorflowDialectToExecutorTest, ConvertsToExecutor) { + CellReader compilation_status(kCompilationStreamz); + TF_ASSERT_OK(CreateMlirModule("empty_func.mlir")); TF_EXPECT_OK(ExportFromTensorflowDialectToExecutor(*mlir_module_)); + + EXPECT_EQ(compilation_status.Delta("success"), 1); } TEST_F(TensorflowDialectToExecutorTest, ErrorsWhenCannotConvert) { + CellReader compilation_status(kCompilationStreamz); + TF_ASSERT_OK(CreateMlirModule("invalid_executor.mlir")); EXPECT_FALSE(ExportFromTensorflowDialectToExecutor(*mlir_module_).ok()); + + EXPECT_EQ(compilation_status.Delta("failed"), 1); } } // namespace diff --git a/tensorflow/compiler/mlir/tf2xla/api/v2/BUILD b/tensorflow/compiler/mlir/tf2xla/api/v2/BUILD index b673432a9f28b0..03445b8a9c341d 100644 --- a/tensorflow/compiler/mlir/tf2xla/api/v2/BUILD +++ b/tensorflow/compiler/mlir/tf2xla/api/v2/BUILD @@ -171,7 +171,9 @@ cc_library( "//tensorflow/compiler/mlir/tensorflow/transforms:verify_no_outside_compilation_markers_pass", "//tensorflow/compiler/mlir/tf2xla/internal:logging_hooks", "//tensorflow/core:framework", + "//tensorflow/core/platform:error_payloads", "//tensorflow/core/platform:status", + "@com_google_absl//absl/log", "@com_google_absl//absl/status", "@llvm-project//llvm:Support", "@llvm-project//mlir:FuncDialect", @@ -180,6 +182,8 @@ cc_library( "@llvm-project//mlir:Support", "@llvm-project//mlir:Transforms", "@local_tsl//tsl/lib/monitoring:counter", + "@local_tsl//tsl/platform:error_logging", + "@local_tsl//tsl/platform:errors", "@local_tsl//tsl/platform:status", ], ) diff --git a/tensorflow/compiler/mlir/tf2xla/api/v2/cluster_tf.cc b/tensorflow/compiler/mlir/tf2xla/api/v2/cluster_tf.cc index b07d0340705e01..4d507cf22625ed 100644 --- a/tensorflow/compiler/mlir/tf2xla/api/v2/cluster_tf.cc +++ b/tensorflow/compiler/mlir/tf2xla/api/v2/cluster_tf.cc @@ -200,19 +200,6 @@ tensorflow::Status RunNonTPUBridge(ModuleOp module, is_in_fallback_enabled_mode, device_type, clustering_status)); - Status runtime_lowering_status = - tensorflow::tfrt_compiler::RunLowerClusterToRuntimeOpsPassPipeline( - module, tsl::DeviceType("XLA_GPU_JIT"), module_name); - TF_RETURN_IF_ERROR(RecordIfErrorStatus(/*error_prefix=*/"runtime_lowering_v2", - is_in_fallback_enabled_mode, - device_type, runtime_lowering_status)); - - Status export_status = - tensorflow::tf2xla::v2::ExportFromTensorflowDialectToExecutor( - module, module_name); - TF_RETURN_IF_ERROR(RecordIfErrorStatus(/*error_prefix=*/"export_to_executor", - is_in_fallback_enabled_mode, - device_type, export_status)); tensorflow::metrics::UpdateTfMlirBridgeFirstPhaseCounter( device_type, /*bridge_version=*/"v2", is_in_fallback_enabled_mode, /*result=*/"success"); diff --git a/tensorflow/compiler/mlir/tf2xla/api/v2/cluster_tf.h b/tensorflow/compiler/mlir/tf2xla/api/v2/cluster_tf.h index d597ca89e8bb23..e1298ac53560d3 100644 --- a/tensorflow/compiler/mlir/tf2xla/api/v2/cluster_tf.h +++ b/tensorflow/compiler/mlir/tf2xla/api/v2/cluster_tf.h @@ -31,13 +31,14 @@ namespace v2 { // API. These transformations take as input a Tensorflow Graph as an MLIR Module // and transforms the module in place to cluster the given ops for compilation // that is compatible with the given device_type. The MLIR should be in the TF -// Executor Dialect for graph nodes and edges. Individual Op inside a node -// should be the Tensorflow Dialect. The output MLIR is in the TF Executor -// Dialect. Returns OkStatus if passed, otherwise an error. +// Executor Dialect for graph nodes and edges or be in TF Functional already. +// Individual Op inside a node should be the Tensorflow Functional Dialect. The +// output MLIR is in the TF Functional Dialect. Returns OkStatus if passed, +// otherwise an error. // // Inputs: // module - The MLIR Module that will be clustered. Expected to be in TF -// Executor Dialect +// Executor Dialect or TF Functional Dialect. Will convert to TF Functional. // . device_type - The device type to cluster for. // is_in_fallback_enabled_mode - Whether this was called with fallback to the // non-MLIR Bridge. This is just for logging purposes and doesn't affect diff --git a/tensorflow/compiler/mlir/tf2xla/api/v2/cluster_tf_test.cc b/tensorflow/compiler/mlir/tf2xla/api/v2/cluster_tf_test.cc index e9a6de55693f8f..d00d8b43d9e790 100644 --- a/tensorflow/compiler/mlir/tf2xla/api/v2/cluster_tf_test.cc +++ b/tensorflow/compiler/mlir/tf2xla/api/v2/cluster_tf_test.cc @@ -15,6 +15,7 @@ limitations under the License. #include "tensorflow/compiler/mlir/tf2xla/api/v2/cluster_tf.h" +#include #include #include @@ -132,13 +133,6 @@ TEST_F(FunctionClusterTensorflowDialectTest, ClustersTFCPU) { FuncOp main = mlir_module_->lookupSymbol("main"); ASSERT_TRUE(main); - bool has_graph_op = false; - main.walk([&](mlir::tf_executor::GraphOp graph) { - has_graph_op = true; - return WalkResult::advance(); - }); - - EXPECT_TRUE(has_graph_op); EXPECT_EQ( compilation_status.Delta("cpu/gpu", "v2", "fallback_disabled", "success"), 1); @@ -156,14 +150,6 @@ TEST_F(FunctionClusterTensorflowDialectTest, ClustersTFGPU) { FuncOp main = mlir_module_->lookupSymbol("main"); ASSERT_TRUE(main); - bool has_graph_op = false; - main.walk([&](mlir::tf_executor::GraphOp graph) { - has_graph_op = true; - return WalkResult::advance(); - }); - - EXPECT_TRUE(has_graph_op); - EXPECT_EQ( compilation_status.Delta("cpu/gpu", "v2", "fallback_disabled", "success"), 1); diff --git a/tensorflow/compiler/mlir/tf2xla/api/v2/tf_dialect_to_executor.cc b/tensorflow/compiler/mlir/tf2xla/api/v2/tf_dialect_to_executor.cc index da42af23d62772..69f1c0e20a5e1b 100644 --- a/tensorflow/compiler/mlir/tf2xla/api/v2/tf_dialect_to_executor.cc +++ b/tensorflow/compiler/mlir/tf2xla/api/v2/tf_dialect_to_executor.cc @@ -18,6 +18,7 @@ limitations under the License. #include #include +#include "absl/log/log.h" #include "absl/status/status.h" #include "llvm/ADT/StringRef.h" #include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project @@ -31,9 +32,11 @@ limitations under the License. #include "tensorflow/compiler/mlir/tensorflow/utils/data_dumper_logger_config.h" #include "tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.h" #include "tensorflow/compiler/mlir/tf2xla/internal/logging_hooks.h" +#include "tensorflow/core/platform/error_payloads.h" #include "tensorflow/core/platform/status.h" #include "tensorflow/core/util/debug_data_dumper.h" #include "tsl/lib/monitoring/counter.h" +#include "tsl/platform/error_logging.h" #include "tsl/platform/status.h" namespace tensorflow { @@ -93,6 +96,30 @@ void AddTfDialectToExecutorPasses(OpPassManager &pm) { pm.addPass(mlir::TF::CreateVerifySuitableForExportPass()); } +tensorflow::Status RecordStatusIfError(absl::Status status) { + if (status.ok()) { + return absl::OkStatus(); + } + + tf_dialect_to_executor_dialect_status->GetCell(kExportFailed)->IncrementBy(1); + VLOG(1) << "Failed to export from TF Dialect to TF Executor Dialect. " + << status; + + constexpr char bridge_subcomponent[] = + "TFXLA_TF_FUNCTIONAL_TO_EXECUTOR_EXPORT_v2"; + constexpr char kBridgeComponent[] = "TFXLABridge"; + + tsl::OkOrSetErrorCounterPayload( + tensorflow::core::platform::ErrorSourceProto::MLIR_BRIDGE_PHASE_1, + status); + + tsl::error_logging::Log(kBridgeComponent, bridge_subcomponent, + status.ToString()) + .IgnoreError(); + + return status; +} + } // namespace tensorflow::Status ExportFromTensorflowDialectToExecutor( @@ -128,12 +155,10 @@ tensorflow::Status ExportFromTensorflowDialectToExecutor( module, llvm::StringRef(), &tf_to_executor); } - if (!result.succeeded()) { - tf_dialect_to_executor_dialect_status->GetCell(kExportFailed) - ->IncrementBy(1); - - return absl::InternalError( - "Failed to export from TF Dialect to TF Executor Dialect."); + if (result.failed()) { + return RecordStatusIfError( + absl::InternalError("Failed to export from TF Dialect to TF Executor " + "Dialect. Read LLVM Pipeline Error")); } tf_dialect_to_executor_dialect_status->GetCell(kExportSuccess) diff --git a/tensorflow/compiler/mlir/tf2xla/internal/passes/BUILD b/tensorflow/compiler/mlir/tf2xla/internal/passes/BUILD new file mode 100644 index 00000000000000..9abed32b2fac8b --- /dev/null +++ b/tensorflow/compiler/mlir/tf2xla/internal/passes/BUILD @@ -0,0 +1,95 @@ +load("//tensorflow/compiler/mlir:glob_lit_test.bzl", "glob_lit_tests") +load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library") +load("//tensorflow:tensorflow.bzl", "tf_cc_test") +load("//tensorflow:tensorflow.default.bzl", "filegroup", "get_compatible_with_portable") +load("//tensorflow/core/platform:rules_cc.bzl", "cc_library") + +package( + # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"], + default_visibility = [ + "//tensorflow/compiler/mlir:__pkg__", + "//tensorflow/compiler/mlir/tf2xla/internal:__subpackages__", + ], + licenses = ["notice"], +) + +cc_library( + name = "clustering_passes", + srcs = [ + "verify_clustering_pass.cc", + ], + hdrs = [ + "clustering_passes.h", + ], + textual_hdrs = [ + "clustering_passes.h.inc", + ], + deps = [ + ":clustering_passes_inc_gen", + "//tensorflow/compiler/mlir/tensorflow", + "//tensorflow/core:framework", + "//tensorflow/core/transforms/toposort:Pass", + "@com_google_absl//absl/log", + "@com_google_absl//absl/strings", + "@llvm-project//llvm:Support", + "@llvm-project//mlir:FuncDialect", + "@llvm-project//mlir:IR", + "@llvm-project//mlir:Pass", + "@llvm-project//mlir:Support", + "@llvm-project//mlir:Transforms", + ], +) + +gentbl_cc_library( + name = "clustering_passes_inc_gen", + compatible_with = get_compatible_with_portable(), + tbl_outs = [ + ( + [ + "-gen-pass-decls", + "-name=TFXLABridge", + ], + "clustering_passes.h.inc", + ), + ], + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "clustering_passes.td", + deps = [ + "@llvm-project//mlir:PassBaseTdFiles", + ], +) + +tf_cc_test( + name = "verify_clustering_pass_test", + srcs = ["verify_clustering_pass_test.cc"], + deps = [ + ":clustering_passes", + "//tensorflow/compiler/mlir/tf2xla/transforms:test_utils", + "@com_google_absl//absl/strings", + "@com_google_googletest//:gtest_main", + "@llvm-project//mlir:FuncDialect", + "@llvm-project//mlir:IR", + "@llvm-project//mlir:Pass", + "@llvm-project//mlir:Support", + "@local_tsl//tsl/platform:statusor", + ], +) + +glob_lit_tests( + name = "all_tests", + data = [":test_utilities"], + driver = "@llvm-project//mlir:run_lit.sh", + test_file_exts = [ + "mlir", + ], +) + +# Bundle together all of the test utilities that are used by tests. +filegroup( + name = "test_utilities", + testonly = True, + data = [ + "//tensorflow/compiler/mlir:tf-opt", + "@llvm-project//llvm:FileCheck", + ], +) diff --git a/tensorflow/compiler/mlir/tf2xla/internal/passes/clustering_passes.h b/tensorflow/compiler/mlir/tf2xla/internal/passes/clustering_passes.h new file mode 100644 index 00000000000000..0e477c86b5405b --- /dev/null +++ b/tensorflow/compiler/mlir/tf2xla/internal/passes/clustering_passes.h @@ -0,0 +1,37 @@ +/* Copyright 2023 The TensorFlow Authors. All Rights Reserved. +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + http://www.apache.org/licenses/LICENSE-2.0 +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#ifndef TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_PASSES_CLUSTERING_PASSES_H_ +#define TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_PASSES_CLUSTERING_PASSES_H_ + +#include + +#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project +#include "mlir/Pass/Pass.h" // from @llvm-project + +namespace tensorflow { +namespace tf2xla { +namespace internal { + +// Verifies that all MLIR Ops have the expected attributes. +std::unique_ptr> +CreateVerifyClusteringPass(); + +#define GEN_PASS_REGISTRATION +#define GEN_PASS_DECL_VERIFYCLUSTERINGPASS +#include "tensorflow/compiler/mlir/tf2xla/internal/passes/clustering_passes.h.inc" + +} // namespace internal +} // namespace tf2xla +} // namespace tensorflow + +#endif // TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_PASSES_CLUSTERING_PASSES_H_ diff --git a/tensorflow/compiler/mlir/tf2xla/internal/passes/clustering_passes.td b/tensorflow/compiler/mlir/tf2xla/internal/passes/clustering_passes.td new file mode 100644 index 00000000000000..c1431369c6e0f3 --- /dev/null +++ b/tensorflow/compiler/mlir/tf2xla/internal/passes/clustering_passes.td @@ -0,0 +1,26 @@ +/* Copyright 2023 The TensorFlow Authors. All Rights Reserved. +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + http://www.apache.org/licenses/LICENSE-2.0 +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ +include "mlir/Pass/PassBase.td" + +def VerifyClusteringPass : Pass<"verify-clustering-pass", "mlir::func::FuncOp"> { + + let summary = "Verify that the Bridge output is correct and errors if verification fails."; + + let description = [{ + Verifies whether clustering has resulted in the expected invariants. These + include verifying that clusters have been created and have been outside + compiled, the result is device agnostic and in TF functional dialect & + that the device attribute exists. + }]; + + let constructor = "tensorflow::tf2xla::internal::CreateVerifyClusteringPass()"; +} diff --git a/tensorflow/compiler/mlir/tf2xla/internal/passes/verify_clustering_pass.cc b/tensorflow/compiler/mlir/tf2xla/internal/passes/verify_clustering_pass.cc new file mode 100644 index 00000000000000..2c053dbbb245e5 --- /dev/null +++ b/tensorflow/compiler/mlir/tf2xla/internal/passes/verify_clustering_pass.cc @@ -0,0 +1,68 @@ +/* Copyright 2023 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include +#include +#include + +#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project +#include "mlir/IR/Operation.h" // from @llvm-project +#include "mlir/IR/Visitors.h" // from @llvm-project +#include "mlir/Pass/Pass.h" // from @llvm-project + +namespace tensorflow { +namespace tf2xla { +namespace internal { + +namespace { + +#define GEN_PASS_DEF_VERIFYCLUSTERINGPASS +#include "tensorflow/compiler/mlir/tf2xla/internal/passes/clustering_passes.h.inc" + +class VerifyClusteringPass + : public impl::VerifyClusteringPassBase { + public: + void runOnOperation() override; +}; + +void VerifyClusteringPass::runOnOperation() { + std::set valid_namespaces = {"tf", "func", "return"}; + mlir::Operation* func_op = getOperation(); + + auto walk_result = func_op->walk([&](mlir::Operation* op) { + if (valid_namespaces.find(op->getDialect()->getNamespace().str()) == + valid_namespaces.end()) { + std::string error = "op is in dialect " + + op->getDialect()->getNamespace().str() + + " not in tf functional dialect"; + op->emitError() << error; + return mlir::WalkResult::interrupt(); + } + return mlir::WalkResult::advance(); + }); + + if (walk_result.wasInterrupted()) { + signalPassFailure(); + } +} +} // namespace + +std::unique_ptr> +CreateVerifyClusteringPass() { + return std::make_unique(); +} +} // namespace internal +} // namespace tf2xla +} // namespace tensorflow diff --git a/tensorflow/compiler/mlir/tf2xla/internal/passes/verify_clustering_pass_test.cc b/tensorflow/compiler/mlir/tf2xla/internal/passes/verify_clustering_pass_test.cc new file mode 100644 index 00000000000000..6767a00c2fb4a6 --- /dev/null +++ b/tensorflow/compiler/mlir/tf2xla/internal/passes/verify_clustering_pass_test.cc @@ -0,0 +1,87 @@ +/* Copyright 2023 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include + +#include +#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project +#include "mlir/IR/BuiltinOps.h" // from @llvm-project +#include "mlir/IR/MLIRContext.h" // from @llvm-project +#include "mlir/IR/OwningOpRef.h" // from @llvm-project +#include "mlir/Pass/PassManager.h" // from @llvm-project +#include "mlir/Support/LogicalResult.h" // from @llvm-project +#include "tensorflow/compiler/mlir/tf2xla/internal/passes/clustering_passes.h" +#include "tensorflow/compiler/mlir/tf2xla/transforms/test_utils.h" +#include "tsl/platform/statusor.h" + +namespace tensorflow { +namespace tf2xla { +namespace internal { + +namespace { + +using mlir::mhlo::test::GetMlirModuleFromString; + +class VerifyClusteringPassTest : public testing::Test { + protected: + void CreateModule(const char* module_string) { + TF_ASSERT_OK_AND_ASSIGN(module_, + GetMlirModuleFromString(module_string, &context_)); + pm_ = std::make_unique(&context_); + pm_->addNestedPass(CreateVerifyClusteringPass()); + } + + mlir::LogicalResult Run() { return pm_->run(module_.get()); } + + private: + mlir::MLIRContext context_; + mlir::OwningOpRef module_; + std::unique_ptr pm_; +}; + +TEST_F(VerifyClusteringPassTest, OnlyTfFunctionalPasses) { + static constexpr char kMlirModuleStr[] = R"( + module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} { + func.func @main() -> tensor<1xi32> { + %0 = "tf.Const"() {value = dense<1000> : tensor<1xi32>} : () -> tensor<1xi32> + return %0 : tensor<1xi32> + } + })"; + CreateModule(kMlirModuleStr); + + auto result = Run(); + + EXPECT_TRUE(result.succeeded()); +} + +TEST_F(VerifyClusteringPassTest, NotTfFunctionalFails) { + static constexpr char kMlirModuleStr[] = R"( + module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} { + func.func @main() -> tensor<3x32x32x3xf32> { + %0 = mhlo.constant dense<2.550000e+02> : tensor<3x32x32x3xf32> + return %0 : tensor<3x32x32x3xf32> + } + })"; + CreateModule(kMlirModuleStr); + + auto result = Run(); + + EXPECT_TRUE(result.failed()); +} + +} // namespace +} // namespace internal +} // namespace tf2xla +} // namespace tensorflow diff --git a/tensorflow/compiler/mlir/tf2xla/internal/passes/verify_clustering_pass_test.mlir b/tensorflow/compiler/mlir/tf2xla/internal/passes/verify_clustering_pass_test.mlir new file mode 100644 index 00000000000000..23e60242621f37 --- /dev/null +++ b/tensorflow/compiler/mlir/tf2xla/internal/passes/verify_clustering_pass_test.mlir @@ -0,0 +1,16 @@ +// RUN: tf-opt -verify-clustering-pass -split-input-file -verify-diagnostics %s | FileCheck %s +// Tests the VerifyClusteringPass Pass, ensures that an error is thrown when validation fails. + +func.func @testNotTfDialect(%arg0: tensor<1x32x10x32xi32>, %arg1: tensor<32xi32>) -> tensor<1x32x10x32xi32> { + // expected-error@below {{op is in dialect chlo not in tf functional dialect}} + %0 = "chlo.broadcast_add"(%arg0, %arg1) {broadcast_dimensions = dense<3> : tensor<1xi64>} : (tensor<1x32x10x32xi32>, tensor<32xi32>) -> tensor<1x32x10x32xi32> + func.return %0 : tensor<1x32x10x32xi32> +} + +// ----- + +// CHECK-LABEL: func @testTFDialect +func.func @testTFDialect(%arg0: tensor<4x?x!tf_type.stringref>) -> tensor<4x2x!tf_type.string> { + %0 = "tf.Identity"(%arg0) : (tensor<4x?x!tf_type.stringref>) -> tensor<4x2x!tf_type.string> + func.return %0 : tensor<4x2x!tf_type.string> +} \ No newline at end of file diff --git a/tensorflow/compiler/mlir/tf2xla/tests/legalize-tf-binary-elementwise.mlir b/tensorflow/compiler/mlir/tf2xla/tests/legalize-tf-binary-elementwise.mlir index 30dc8ecbbf647b..01dc4701923675 100644 --- a/tensorflow/compiler/mlir/tf2xla/tests/legalize-tf-binary-elementwise.mlir +++ b/tensorflow/compiler/mlir/tf2xla/tests/legalize-tf-binary-elementwise.mlir @@ -266,28 +266,28 @@ func.func @equal_broadcast(%arg0: tensor<1xi32>, %arg1: tensor<1x2xi32>) -> tens // CHECK-LABEL: func @equal_broadcast_no_incompatible_shapes_error func.func @equal_broadcast_no_incompatible_shapes_error(%arg0: tensor<2xi32>, %arg1: tensor<1x2xi32>) -> tensor<1x2xi1> { - // CHECK-NEXT: "tf.Equal"(%arg0, %arg1) {incompatible_shape_error = true} + // CHECK-NEXT: "tf.Equal"(%arg0, %arg1) <{incompatible_shape_error = true}> %0 = "tf.Equal"(%arg0, %arg1) { incompatible_shape_error = false } : (tensor<2xi32>, tensor<1x2xi32>) -> tensor<1x2xi1> func.return %0: tensor<1x2xi1> } // CHECK-LABEL: func @equal_incompatible_shape_broadcastable func.func @equal_incompatible_shape_broadcastable(%arg0: tensor, %arg1: tensor<1xi32>) -> tensor { - // CHECK-NEXT: "tf.Equal"(%arg0, %arg1) {incompatible_shape_error = true} + // CHECK-NEXT: "tf.Equal"(%arg0, %arg1) <{incompatible_shape_error = true}> %0 = "tf.Equal"(%arg0, %arg1) { incompatible_shape_error = false } : (tensor, tensor<1xi32>) -> tensor func.return %0: tensor } // CHECK-LABEL: func @equal_incompatible_shape_dynamic func.func @equal_incompatible_shape_dynamic(%arg0: tensor<2xi32>, %arg1: tensor) -> tensor<*xi1> { - // CHECK-NEXT: "tf.Equal"(%arg0, %arg1) {incompatible_shape_error = false} + // CHECK-NEXT: "tf.Equal"(%arg0, %arg1) <{incompatible_shape_error = false}> %0 = "tf.Equal"(%arg0, %arg1) { incompatible_shape_error = false } : (tensor<2xi32>, tensor) -> tensor<*xi1> func.return %0: tensor<*xi1> } // CHECK-LABEL: func @equal_incompatible_shape_both_dynamic func.func @equal_incompatible_shape_both_dynamic(%arg0: tensor, %arg1: tensor) -> tensor<*xi1> { - // CHECK-NEXT: "tf.Equal"(%arg0, %arg1) {incompatible_shape_error = false} + // CHECK-NEXT: "tf.Equal"(%arg0, %arg1) <{incompatible_shape_error = false}> %0 = "tf.Equal"(%arg0, %arg1) { incompatible_shape_error = false } : (tensor, tensor) -> tensor<*xi1> func.return %0: tensor<*xi1> } diff --git a/tensorflow/compiler/mlir/tf2xla/tests/legalize-tf-with-tf2xla-hlo-importer.mlir b/tensorflow/compiler/mlir/tf2xla/tests/legalize-tf-with-tf2xla-hlo-importer.mlir index 32ef8cba3b832c..a732c6d61281ca 100644 --- a/tensorflow/compiler/mlir/tf2xla/tests/legalize-tf-with-tf2xla-hlo-importer.mlir +++ b/tensorflow/compiler/mlir/tf2xla/tests/legalize-tf-with-tf2xla-hlo-importer.mlir @@ -701,9 +701,8 @@ module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, pr // Verifies that the following functions are added from xla_call_module. Note this must be at the end of the file. // CHECK: func.func private @main.2(%arg0: tensor {mhlo.sharding = "{replicated}"}) -> tensor { - // CHECK: %0 = mhlo.bitcast_convert %arg0 : (tensor) -> tensor - // CHECK: %1 = mhlo.sine %0 : tensor - // CHECK: return %1 : tensor + // CHECK: %0 = mhlo.sine %arg0 : tensor + // CHECK: return %0 : tensor // CHECK: } } diff --git a/tensorflow/compiler/mlir/tf2xla/tests/tfxla_device_specific_transformations_cpu.mlir b/tensorflow/compiler/mlir/tf2xla/tests/tfxla_device_specific_transformations_cpu.mlir index c06ba74641ae09..67aa69d5806aaf 100644 --- a/tensorflow/compiler/mlir/tf2xla/tests/tfxla_device_specific_transformations_cpu.mlir +++ b/tensorflow/compiler/mlir/tf2xla/tests/tfxla_device_specific_transformations_cpu.mlir @@ -4,9 +4,9 @@ module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, pr // CHECK-LABEL: stateless_op func.func @stateless_op() -> tensor { - // CHECK: %cst = "tf.Const"() {value = dense<1> : tensor} : () -> tensor + // CHECK: %cst = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor %0 = "tf.StatelessRandomGetAlg"() {device = ""} : () -> tensor return %0 : tensor } -} \ No newline at end of file +} diff --git a/tensorflow/compiler/mlir/tf2xla/tests/tfxla_device_specific_transformations_gpu.mlir b/tensorflow/compiler/mlir/tf2xla/tests/tfxla_device_specific_transformations_gpu.mlir index f051960fb5a2c0..4d5da14a519b61 100644 --- a/tensorflow/compiler/mlir/tf2xla/tests/tfxla_device_specific_transformations_gpu.mlir +++ b/tensorflow/compiler/mlir/tf2xla/tests/tfxla_device_specific_transformations_gpu.mlir @@ -4,9 +4,9 @@ module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, pr // CHECK-LABEL: stateless_op func.func @stateless_op() -> tensor { - // CHECK: %cst = "tf.Const"() {value = dense<1> : tensor} : () -> tensor + // CHECK: %cst = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor %0 = "tf.StatelessRandomGetAlg"() {device = ""} : () -> tensor return %0 : tensor } -} \ No newline at end of file +} diff --git a/tensorflow/compiler/mlir/tf2xla/transforms/BUILD b/tensorflow/compiler/mlir/tf2xla/transforms/BUILD index b4f5d00b7a9e37..ed0429ad242c94 100644 --- a/tensorflow/compiler/mlir/tf2xla/transforms/BUILD +++ b/tensorflow/compiler/mlir/tf2xla/transforms/BUILD @@ -1,10 +1,10 @@ # Description: # TF2XLA Bridge transforms -load("//tensorflow/core/platform:rules_cc.bzl", "cc_library") -load("//tensorflow:tensorflow.bzl", "tf_cc_test") load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library") +load("//tensorflow:tensorflow.bzl", "tf_cc_test") load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable") +load("//tensorflow/core/platform:rules_cc.bzl", "cc_library") load("@local_tsl//tsl/platform:build_config_root.bzl", "if_static") package( @@ -487,13 +487,12 @@ tf_cc_test( "//tensorflow/compiler/mlir/tensorflow", "//tensorflow/compiler/tf2xla:xla_op_registry", "//tensorflow/compiler/tf2xla/kernels:xla_ops", - "//tensorflow/core/tpu:tpu_defs", "@com_google_absl//absl/status", "@com_google_googletest//:gtest_main", "@llvm-project//mlir:FuncDialect", "@llvm-project//mlir:IR", + "@llvm-project//mlir:Support", "@local_tsl//tsl/lib/core:status_test_util", - "@local_tsl//tsl/platform:errors", "@local_tsl//tsl/platform:status", "@local_tsl//tsl/platform:statusor", ], diff --git a/tensorflow/compiler/mlir/tf2xla/transforms/legalization_op_config_test.cc b/tensorflow/compiler/mlir/tf2xla/transforms/legalization_op_config_test.cc index 24f7711539a5b4..ade2b5faa73c8a 100644 --- a/tensorflow/compiler/mlir/tf2xla/transforms/legalization_op_config_test.cc +++ b/tensorflow/compiler/mlir/tf2xla/transforms/legalization_op_config_test.cc @@ -20,21 +20,22 @@ limitations under the License. #include #include -#include #include #include "absl/status/status.h" #include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project #include "mlir/IR/BuiltinOps.h" // from @llvm-project +#include "mlir/IR/DialectRegistry.h" // from @llvm-project +#include "mlir/IR/OperationSupport.h" // from @llvm-project +#include "mlir/IR/OwningOpRef.h" // from @llvm-project +#include "mlir/IR/PatternMatch.h" // from @llvm-project +#include "mlir/Support/TypeID.h" // from @llvm-project #include "tensorflow/compiler/mlir/register_common_dialects.h" -#include "tensorflow/compiler/mlir/tensorflow/dialect_registration.h" #include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h" #include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h" #include "tensorflow/compiler/mlir/tf2xla/transforms/passes.h" #include "tensorflow/compiler/mlir/tf2xla/transforms/test_utils.h" #include "tensorflow/compiler/tf2xla/xla_op_registry.h" -#include "tensorflow/core/tpu/tpu_defs.h" #include "tsl/lib/core/status_test_util.h" -#include "tsl/platform/errors.h" #include "tsl/platform/status.h" #include "tsl/platform/statusor.h" @@ -131,7 +132,7 @@ TEST_F(LegalizationOpConfigTest, CountLoweringsSet) { // a new op, we should expect these to change too. EXPECT_EQ(mlir_lowering_count, 67); EXPECT_EQ(tf2xla_fallback_count, 315); - EXPECT_EQ(non_categorized_count, 420); + EXPECT_EQ(non_categorized_count, 422); } // Just a counter test to see which ops have duplicate lowerings. This isn't a diff --git a/tensorflow/compiler/mlir/tf_mlir_opt_main.cc b/tensorflow/compiler/mlir/tf_mlir_opt_main.cc index 4a454c326550ae..f803230ea4f504 100644 --- a/tensorflow/compiler/mlir/tf_mlir_opt_main.cc +++ b/tensorflow/compiler/mlir/tf_mlir_opt_main.cc @@ -27,6 +27,7 @@ limitations under the License. #include "tensorflow/compiler/mlir/tensorflow/transforms/tf_graph_optimization_pass.h" #include "tensorflow/compiler/mlir/tensorflow/utils/mlprogram_util.h" #include "tensorflow/compiler/mlir/tf2xla/api/v1/compile_mlir_util.h" +#include "tensorflow/compiler/mlir/tf2xla/internal/passes/clustering_passes.h" #include "tensorflow/compiler/mlir/tf2xla/transforms/passes.h" #include "tensorflow/compiler/mlir/tosa/tf_passes.h" #include "tensorflow/compiler/mlir/tosa/tf_tfl_passes.h" @@ -55,6 +56,7 @@ int main(int argc, char **argv) { mlir::mhlo::registerLegalizeTfPasses(); mlir::mhlo::registerTfXlaPasses(); mlir::quant::stablehlo::registerBridgePasses(); + tensorflow::tf2xla::internal::registerTFXLABridgePasses(); mlir::tosa::registerLegalizeTosaPasses(); mlir::tosa::registerTFtoTOSALegalizationPipeline(); mlir::tosa::registerTFLtoTOSALegalizationPipeline(); diff --git a/tensorflow/compiler/mlir/tfr/examples/mnist/BUILD b/tensorflow/compiler/mlir/tfr/examples/mnist/BUILD index ee79b37a73d1cc..135bc20b970bc5 100644 --- a/tensorflow/compiler/mlir/tfr/examples/mnist/BUILD +++ b/tensorflow/compiler/mlir/tfr/examples/mnist/BUILD @@ -81,6 +81,7 @@ distribute_py_strict_test( "notap", # The test is too long to run as part of llvm presubmits (b/173661843). "notpu", # Takes too long (b/192305423) "notsan", # Not needed, and there were issues with timeouts. + "requires-net:external", ], # TODO(b/175056184): Re-enable xla_enable_strict_auto_jit once the issues diff --git a/tensorflow/compiler/mlir/tfr/tests/canonicalize.mlir b/tensorflow/compiler/mlir/tfr/tests/canonicalize.mlir index 77508b60046151..912385f36b6f85 100644 --- a/tensorflow/compiler/mlir/tfr/tests/canonicalize.mlir +++ b/tensorflow/compiler/mlir/tfr/tests/canonicalize.mlir @@ -41,7 +41,7 @@ func.func @constant_tensor_array() -> !tfr.tensor { %1 = "tfr.constant_tensor"(%0) : (!tfr.attr) -> !tfr.tensor func.return %1 : !tfr.tensor -// CHECK-NEXT: %[[RES:.*]] = "tf.Const"() {value = dense<[1, -1, 3]> : tensor<3xi64>} : () -> tensor<3xi64> +// CHECK-NEXT: %[[RES:.*]] = "tf.Const"() <{value = dense<[1, -1, 3]> : tensor<3xi64>}> : () -> tensor<3xi64> // CHECK-NEXT: "tfr.cast"(%[[RES]]) : (tensor<3xi64>) -> !tfr.tensor // CHECK-NEXT: return } @@ -54,7 +54,7 @@ func.func @constant_tensor_scalar() -> !tfr.tensor { %1 = "tfr.constant_tensor"(%0) : (i32) -> !tfr.tensor func.return %1 : !tfr.tensor -// CHECK-NEXT: %[[RES:.*]] = "tf.Const"() {value = dense<42> : tensor} : () -> tensor +// CHECK-NEXT: %[[RES:.*]] = "tf.Const"() <{value = dense<42> : tensor}> : () -> tensor // CHECK-NEXT: "tfr.cast"(%[[RES]]) : (tensor) -> !tfr.tensor // CHECK-NEXT: return } @@ -83,7 +83,7 @@ func.func @quant_raw_data(%arg0: tensor<1x10x!quant.uniform>) -> // CHECK-LABEL: quant_raw_data_with_list func.func @quant_raw_data_with_list(%arg0: !tfr.tensor, %arg1: !tfr.tensor) -> !tfr.tensor { - %cst_1 = "tf.Const"() {value = dense<1> : tensor} : () -> tensor + %cst_1 = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor %1 = "tfr.cast"(%arg0) : (!tfr.tensor) -> tensor<1x4x4x3x!quant.uniform> %2 = "tfr.cast"(%arg1) : (!tfr.tensor) -> tensor<1x3x4x3x!quant.uniform> %3 = "tfr.cast"(%2) : (tensor<1x3x4x3x!quant.uniform>) -> !tfr.tensor @@ -94,7 +94,7 @@ func.func @quant_raw_data_with_list(%arg0: !tfr.tensor, %arg1: !tfr.tensor) -> ! %8 = tfr.call @tf__concat(%7, %6) : (!tfr.tensor, !tfr.tensor_list) -> !tfr.tensor func.return %8 : !tfr.tensor -// CHECK: %[[CONST_0:.*]] = "tf.Const"() {value = dense<1> : tensor} : () -> tensor +// CHECK: %[[CONST_0:.*]] = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor // CHECK: %[[BUILD_LIST_0:.*]] = "tfr.build_list"(%arg1, %arg0) : (!tfr.tensor, !tfr.tensor) -> !tfr.tensor_list // CHECK: %[[CAST_0:.*]] = "tfr.cast"(%[[CONST_0]]) : (tensor) -> !tfr.tensor // CHECK: %[[CONCAT_O:.*]] = tfr.call @tf__concat(%[[CAST_0]], %[[BUILD_LIST_0]]) : (!tfr.tensor, !tfr.tensor_list) -> !tfr.tensor @@ -131,8 +131,8 @@ func.func @quant_qparam(%arg0: tensor<1x10x!quant.uniform>) -> ( %2 = "tfr.cast"(%zp) : (!tfr.tensor) -> tensor func.return %1, %2 : tensor, tensor -// CHECK-DAG: %[[scale:.*]] = "tf.Const"() {value = dense<1.000000e-01> : tensor} -// CHECK-DAG: %[[zp:.*]] = "tf.Const"() {value = dense<42> : tensor} : () -> tensor +// CHECK-DAG: %[[scale:.*]] = "tf.Const"() <{value = dense<1.000000e-01> : tensor}> +// CHECK-DAG: %[[zp:.*]] = "tf.Const"() <{value = dense<42> : tensor}> : () -> tensor // CHECK: return %[[scale]], %[[zp]] } @@ -144,8 +144,8 @@ func.func @quant_qparam_per_channel(%arg0: tensor<1x3x!quant.uniform tensor<3xi32> func.return %1, %2 : tensor<3xf32>, tensor<3xi32> -// CHECK-DAG: %[[scale:.*]] = "tf.Const"() {value = dense<[1.000000e-01, 2.000000e-01, 3.000000e-01]> : tensor<3xf32>} -// CHECK-DAG: %[[zp:.*]] = "tf.Const"() {value = dense<[1, 2, 3]> : tensor<3xi32>} : () -> tensor<3xi32> +// CHECK-DAG: %[[scale:.*]] = "tf.Const"() <{value = dense<[1.000000e-01, 2.000000e-01, 3.000000e-01]> : tensor<3xf32>}> +// CHECK-DAG: %[[zp:.*]] = "tf.Const"() <{value = dense<[1, 2, 3]> : tensor<3xi32>}> : () -> tensor<3xi32> // CHECK: return %[[scale]], %[[zp]] } @@ -168,9 +168,9 @@ func.func @redundant_cast_with_different_element_type(%arg0: tensor<*xf32>) -> ( %2 = "tfr.cast"(%0) : (!tfr.tensor) -> tensor<2xi32> func.return %1, %2 : tensor<*xi32>, tensor<2xi32> -// CHECK: %[[tf_cast_unranked:.*]] = "tf.Cast"(%arg0) {Truncate = false} : (tensor<*xf32>) -> tensor<*xi32> -// CHECK: %[[ensure_shape:.*]] = "tf.EnsureShape"(%arg0) {shape = #tf_type.shape<2>} : (tensor<*xf32>) -> tensor<2xf32> -// CHECK: %[[tf_cast_ranked:.*]] = "tf.Cast"(%[[ensure_shape]]) {Truncate = false} : (tensor<2xf32>) -> tensor<2xi32> +// CHECK: %[[tf_cast_unranked:.*]] = "tf.Cast"(%arg0) <{Truncate = false}> : (tensor<*xf32>) -> tensor<*xi32> +// CHECK: %[[ensure_shape:.*]] = "tf.EnsureShape"(%arg0) <{shape = #tf_type.shape<2>}> : (tensor<*xf32>) -> tensor<2xf32> +// CHECK: %[[tf_cast_ranked:.*]] = "tf.Cast"(%[[ensure_shape]]) <{Truncate = false}> : (tensor<2xf32>) -> tensor<2xi32> // CHECK: return %[[tf_cast_unranked]], %[[tf_cast_ranked]] : tensor<*xi32>, tensor<2xi32> } @@ -185,7 +185,7 @@ func.func @redundant_cast_with_quant_type(%arg0: tensor<10x!quant.uniform // CHECK: %[[CAST_0:.*]] = "tfr.cast"(%arg0) : (tensor<10x!quant.uniform>) -> !tfr.tensor // CHECK: %[[CAST_1:.*]] = "tfr.cast"(%[[CAST_0]]) : (!tfr.tensor) -> tensor<10xi8> -// CHECK: %[[CAST_2:.*]] = "tf.Cast"(%[[CAST_1]]) {Truncate = false} : (tensor<10xi8>) -> tensor<10xi32> +// CHECK: %[[CAST_2:.*]] = "tf.Cast"(%[[CAST_1]]) <{Truncate = false}> : (tensor<10xi8>) -> tensor<10xi32> // CHECK: return %[[CAST_2]] : tensor<10xi32> } diff --git a/tensorflow/compiler/mlir/tfr/tests/decompose.mlir b/tensorflow/compiler/mlir/tfr/tests/decompose.mlir index 0dcd363164f760..eb35c3a30e79e6 100644 --- a/tensorflow/compiler/mlir/tfr/tests/decompose.mlir +++ b/tensorflow/compiler/mlir/tfr/tests/decompose.mlir @@ -227,8 +227,8 @@ func.func @decompose_quant_scale_factor() -> (!tfr.tensor, !tfr.tensor) { %list2 = "tfr.build_list"(%input_scale_tensor, %perchannel_scale_tensor) : (!tfr.tensor, !tfr.tensor) -> !tfr.tensor_list %perchannel = "tfr.quant_scale_factor"(%output_scale, %list2) : (f32, !tfr.tensor_list) -> !tfr.tensor func.return %out, %perchannel : !tfr.tensor, !tfr.tensor -// CHECK-DAG: %[[scale_factors:.*]] = "tf.Const"() {value = dense<[1.000000e+00, 1.000000e+01]> : tensor<2xf32>} : () -> tensor<2xf32> -// CHECK-DAG: %[[scale_factor:.*]] = "tf.Const"() {value = dense<1.000000e+00> : tensor} : () -> tensor +// CHECK-DAG: %[[scale_factors:.*]] = "tf.Const"() <{value = dense<[1.000000e+00, 1.000000e+01]> : tensor<2xf32>}> : () -> tensor<2xf32> +// CHECK-DAG: %[[scale_factor:.*]] = "tf.Const"() <{value = dense<1.000000e+00> : tensor}> : () -> tensor // CHECK: %[[cast:.*]] = "tfr.cast"(%[[scale_factor]]) : (tensor) -> !tfr.tensor // CHECK: %[[cast_perchannel:.*]] = "tfr.cast"(%[[scale_factors]]) : (tensor<2xf32>) -> !tfr.tensor // CHECK: return %[[cast]], %[[cast_perchannel]] : !tfr.tensor, !tfr.tensor @@ -245,8 +245,8 @@ func.func @decompose_quant_scale_factor_invalid() -> !tfr.tensor { %out = "tfr.quant_scale_factor"(%output_scale, %list) : (f32, !tfr.tensor_list) -> !tfr.tensor func.return %out : !tfr.tensor // CHECK-DAG: %[[cst_0:.*]] = arith.constant 1.000000e-01 : f32 -// CHECK-DAG: %[[cst_1:.*]] = "tf.Const"() {value = dense<2.500000e-01> : tensor} : () -> tensor -// CHECK-DAG: %[[cst_2:.*]] = "tf.Const"() {value = dense<4.000000e-01> : tensor} : () -> tensor +// CHECK-DAG: %[[cst_1:.*]] = "tf.Const"() <{value = dense<2.500000e-01> : tensor}> : () -> tensor +// CHECK-DAG: %[[cst_2:.*]] = "tf.Const"() <{value = dense<4.000000e-01> : tensor}> : () -> tensor // CHECK: %[[tfrcast0:.*]] = "tfr.cast"(%[[cst_1]]) : (tensor) -> !tfr.tensor // CHECK: %[[tfrcast1:.*]] = "tfr.cast"(%[[cst_2]]) : (tensor) -> !tfr.tensor // CHECK: %[[list:.*]] = "tfr.build_list"(%[[tfrcast0]], %[[tfrcast1]], %[[tfrcast0]]) : (!tfr.tensor, !tfr.tensor, !tfr.tensor) -> !tfr.tensor_list @@ -265,9 +265,9 @@ func.func @decompose_quant_rescale(%arg0: tensor<2xi32>) -> !tfr.tensor { // CHECK-DAG: %[[f32:.*]] = tfr.constant f32 -> !tfr.attr // CHECK-DAG: %[[i32:.*]] = tfr.constant i32 -> !tfr.attr -// CHECK-DAG: %[[scale_cst:.*]] = "tf.Const"() {value = dense<1.000000e+00> : tensor} : () -> tensor +// CHECK-DAG: %[[scale_cst:.*]] = "tf.Const"() <{value = dense<1.000000e+00> : tensor}> : () -> tensor // CHECK-DAG: %false = arith.constant false -// CHECK-DAG: %[[zp_cst:.*]] = "tf.Const"() {value = dense<67> : tensor} : () -> tensor +// CHECK-DAG: %[[zp_cst:.*]] = "tf.Const"() <{value = dense<67> : tensor}> : () -> tensor // CHECK: %[[zp:.*]] = "tfr.cast"(%[[zp_cst]]) : (tensor) -> !tfr.tensor // CHECK: %[[scale:.*]] = "tfr.cast"(%[[scale_cst]]) : (tensor) -> !tfr.tensor // CHECK: %[[input:.*]] = "tfr.cast"(%arg0) : (tensor<2xi32>) -> !tfr.tensor diff --git a/tensorflow/compiler/mlir/tfr/tests/end2end.mlir b/tensorflow/compiler/mlir/tfr/tests/end2end.mlir index 6a49a0dbbb1d13..0654b216e0f165 100644 --- a/tensorflow/compiler/mlir/tfr/tests/end2end.mlir +++ b/tensorflow/compiler/mlir/tfr/tests/end2end.mlir @@ -17,7 +17,7 @@ func.func @my_rsqrt(%arg0: tensor<2x3xf32>) -> tensor<3x2x3xf32> { // CHECK-NEXT: %[[RE:.*]] = "tf.RiscReciprocal"(%arg0) : (tensor<2x3xf32>) -> tensor<*xf32> // CHECK-NEXT: %[[SQRT:.*]] = "tf.RiscSqrt"(%[[RE]]) : (tensor<*xf32>) -> tensor<*xf32> -// CHECK-NEXT: %[[ES:.*]] = "tf.EnsureShape"(%[[SQRT]]) {shape = #tf_type.shape<3x2x3>} : (tensor<*xf32>) -> tensor<3x2x3xf32> +// CHECK-NEXT: %[[ES:.*]] = "tf.EnsureShape"(%[[SQRT]]) <{shape = #tf_type.shape<3x2x3>}> : (tensor<*xf32>) -> tensor<3x2x3xf32> // CHECK-NEXT: return %[[ES]] : tensor<3x2x3xf32> } @@ -26,11 +26,11 @@ func.func @my_leaky_relu(%arg0: tensor<2x3xf32>) -> tensor<3x2x3xf32> { %0 = "tf.MyLeakyRelu"(%arg0) {alpha=3.0 : f32} : (tensor<2x3xf32>) -> tensor<3x2x3xf32> func.return %0 : tensor<3x2x3xf32> -// CHECK-NEXT: %[[ALPHA:.*]] = "tf.Const"() {value = dense<3.000000e+00> : tensor} : () -> tensor +// CHECK-NEXT: %[[ALPHA:.*]] = "tf.Const"() <{value = dense<3.000000e+00> : tensor}> : () -> tensor // CHECK-NEXT: %[[SHAPE:.*]] = "tf.RiscShape"(%arg0) {T = i32} : (tensor<2x3xf32>) -> tensor<*xi32> // CHECK-NEXT: %[[ALPHA1:.*]] = "tf.RiscBroadcast"(%[[ALPHA]], %[[SHAPE]]) : (tensor, tensor<*xi32>) -> tensor<*xf32> // CHECK-NEXT: %[[MAX:.*]] = "tf.RiscMaximum"(%arg0, %[[ALPHA1]]) : (tensor<2x3xf32>, tensor<*xf32>) -> tensor<*xf32> -// CHECK-NEXT: %[[ES:.*]] = "tf.EnsureShape"(%[[MAX]]) {shape = #tf_type.shape<3x2x3>} : (tensor<*xf32>) -> tensor<3x2x3xf32> +// CHECK-NEXT: %[[ES:.*]] = "tf.EnsureShape"(%[[MAX]]) <{shape = #tf_type.shape<3x2x3>}> : (tensor<*xf32>) -> tensor<3x2x3xf32> // CHECK-NEXT: return %[[ES]] : tensor<3x2x3xf32> } @@ -39,11 +39,11 @@ func.func @my_leaky_relu_with_default(%arg0: tensor<2x3xf32>) -> tensor<3x2x3xf3 %0 = "tf.MyLeakyRelu"(%arg0) : (tensor<2x3xf32>) -> tensor<3x2x3xf32> func.return %0 : tensor<3x2x3xf32> -// CHECK-NEXT: %[[ALPHA:.*]] = "tf.Const"() {value = dense<2.000000e-01> : tensor} : () -> tensor +// CHECK-NEXT: %[[ALPHA:.*]] = "tf.Const"() <{value = dense<2.000000e-01> : tensor}> : () -> tensor // CHECK-NEXT: %[[SHAPE:.*]] = "tf.RiscShape"(%arg0) {T = i32} : (tensor<2x3xf32>) -> tensor<*xi32> // CHECK-NEXT: %[[ALPHA1:.*]] = "tf.RiscBroadcast"(%[[ALPHA]], %[[SHAPE]]) : (tensor, tensor<*xi32>) -> tensor<*xf32> // CHECK-NEXT: %[[MAX:.*]] = "tf.RiscMaximum"(%arg0, %[[ALPHA1]]) : (tensor<2x3xf32>, tensor<*xf32>) -> tensor<*xf32> -// CHECK-NEXT: %[[ES:.*]] = "tf.EnsureShape"(%[[MAX]]) {shape = #tf_type.shape<3x2x3>} : (tensor<*xf32>) -> tensor<3x2x3xf32> +// CHECK-NEXT: %[[ES:.*]] = "tf.EnsureShape"(%[[MAX]]) <{shape = #tf_type.shape<3x2x3>}> : (tensor<*xf32>) -> tensor<3x2x3xf32> // CHECK-NEXT: return %[[ES]] : tensor<3x2x3xf32> } @@ -53,7 +53,7 @@ func.func @my_cast(%arg0: tensor<2x3xf32>) -> tensor<2x3xi32> { func.return %0 : tensor<2x3xi32> // CHECK-NEXT: %[[CAST:.*]] = "tf.RiscCast"(%arg0) {Tout = i32} : (tensor<2x3xf32>) -> tensor<*xi32> -// CHECK-NEXT: %[[ES:.*]] = "tf.EnsureShape"(%[[CAST]]) {shape = #tf_type.shape<2x3>} : (tensor<*xi32>) -> tensor<2x3xi32> +// CHECK-NEXT: %[[ES:.*]] = "tf.EnsureShape"(%[[CAST]]) <{shape = #tf_type.shape<2x3>}> : (tensor<*xi32>) -> tensor<2x3xi32> // CHECK-NEXT: return %[[ES]] : tensor<2x3xi32> } @@ -62,9 +62,9 @@ func.func @my_pack_single_input(%arg0: tensor<2x3xf32>) -> tensor<3x2x3xf32> { %0 = "tf.MyPack"(%arg0) {N=1:i32, axis=0:i32} : (tensor<2x3xf32>) -> tensor<3x2x3xf32> func.return %0 : tensor<3x2x3xf32> -// CHECK-NEXT: %[[AXIS:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor +// CHECK-NEXT: %[[AXIS:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK-NEXT: %[[ED:.*]] = "tf.ExpandDims"(%arg0, %[[AXIS]]) : (tensor<2x3xf32>, tensor) -> tensor<*xf32> -// CHECK-NEXT: %[[ES:.*]] = "tf.EnsureShape"(%[[ED]]) {shape = #tf_type.shape<3x2x3>} : (tensor<*xf32>) -> tensor<3x2x3xf32> +// CHECK-NEXT: %[[ES:.*]] = "tf.EnsureShape"(%[[ED]]) <{shape = #tf_type.shape<3x2x3>}> : (tensor<*xf32>) -> tensor<3x2x3xf32> // CHECK-NEXT: return %[[ES]] : tensor<3x2x3xf32> } @@ -73,13 +73,13 @@ func.func @my_pack_multiple_inputs(%arg0: tensor<2x3xf32>, %arg1: tensor<2x3xf32 %0 = "tf.MyPack"(%arg0, %arg1, %arg2) {N=3:i32, axis=0:i32} : (tensor<2x3xf32>, tensor<2x3xf32>, tensor<2x3xf32>) -> tensor<3x2x3xf32> func.return %0 : tensor<3x2x3xf32> -// CHECK-NEXT: %[[AXIS:.*]] = "tf.Const"() {value = dense<0> : tensor} : () -> tensor +// CHECK-NEXT: %[[AXIS:.*]] = "tf.Const"() <{value = dense<0> : tensor}> : () -> tensor // CHECK-NEXT: %[[ED0:.*]] = "tf.ExpandDims"(%arg0, %[[AXIS]]) : (tensor<2x3xf32>, tensor) -> tensor<*xf32> // CHECK-NEXT: %[[ED1:.*]] = "tf.ExpandDims"(%arg1, %[[AXIS]]) : (tensor<2x3xf32>, tensor) -> tensor<*xf32> // CHECK-NEXT: %[[CC0:.*]] = "tf.RiscConcat"(%[[ED0]], %[[ED1]]) {axis = 0 : i32} : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32> // CHECK-NEXT: %[[ED2:.*]] = "tf.ExpandDims"(%arg2, %[[AXIS]]) : (tensor<2x3xf32>, tensor) -> tensor<*xf32> // CHECK-NEXT: %[[CC1:.*]] = "tf.RiscConcat"(%[[CC0]], %[[ED2]]) {axis = 0 : i32} : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32> -// CHECK-NEXT: %[[ES:.*]] = "tf.EnsureShape"(%[[CC1]]) {shape = #tf_type.shape<3x2x3>} : (tensor<*xf32>) -> tensor<3x2x3xf32> +// CHECK-NEXT: %[[ES:.*]] = "tf.EnsureShape"(%[[CC1]]) <{shape = #tf_type.shape<3x2x3>}> : (tensor<*xf32>) -> tensor<3x2x3xf32> // CHECK-NEXT: return %[[ES]] : tensor<3x2x3xf32> } @@ -98,7 +98,7 @@ func.func @my_add_n_multiple_inputs(%arg0: tensor<2x3xf32>, %arg1: tensor<2x3xf3 // CHECK-NEXT: %[[ADD0:.*]] = "tf.RiscAdd"(%arg0, %arg1) : (tensor<2x3xf32>, tensor<2x3xf32>) -> tensor<*xf32> // CHECK-NEXT: %[[ADD1:.*]] = "tf.RiscAdd"(%[[ADD0]], %arg2) : (tensor<*xf32>, tensor<2x3xf32>) -> tensor<*xf32> -// CHECK-NEXT: %[[ES:.*]] = "tf.EnsureShape"(%[[ADD1]]) {shape = #tf_type.shape<2x3>} : (tensor<*xf32>) -> tensor<2x3xf32> +// CHECK-NEXT: %[[ES:.*]] = "tf.EnsureShape"(%[[ADD1]]) <{shape = #tf_type.shape<2x3>}> : (tensor<*xf32>) -> tensor<2x3xf32> // CHECK-NEXT: return %[[ES]] : tensor<2x3xf32> } @@ -112,10 +112,10 @@ func.func @my_map_and_batch_dataset(%input: tensor<*x!tf_type.variant>, : (tensor<*x!tf_type.variant>, tensor<*xf32>, tensor<*xi32>) -> tensor<*x!tf_type.variant> func.return %0 : tensor<*x!tf_type.variant> -// CHECK-DAG: %[[BATCH:.*]] = "tf.Const"() {value = dense<1000> : tensor} : () -> tensor -// CHECK-DAG: %[[PARAL:.*]] = "tf.Const"() {value = dense<8> : tensor} : () -> tensor -// CHECK-DAG: %[[KEEP:.*]] = "tf.Const"() {value = dense : tensor} : () -> tensor -// CHECK: %[[CAST:.*]] = "tf.Cast"(%arg2) {Truncate = false} : (tensor<*xi32>) -> tensor<*xf32> +// CHECK-DAG: %[[BATCH:.*]] = "tf.Const"() <{value = dense<1000> : tensor}> : () -> tensor +// CHECK-DAG: %[[PARAL:.*]] = "tf.Const"() <{value = dense<8> : tensor}> : () -> tensor +// CHECK-DAG: %[[KEEP:.*]] = "tf.Const"() <{value = dense : tensor}> : () -> tensor +// CHECK: %[[CAST:.*]] = "tf.Cast"(%arg2) <{Truncate = false}> : (tensor<*xi32>) -> tensor<*xf32> // CHECK: %[[RET:.*]] = "tf.MapAndBatchDatasetV0"(%arg0, %[[BATCH]], %[[PARAL]], %[[KEEP]], %arg1, %[[CAST]]) // CHECK-SAME: {f = @__some_func, output_shapes = [#tf_type.shape<>], output_types = [f32], preserve_cardinality = true} : (tensor<*x!tf_type.variant>, tensor, tensor, tensor, tensor<*xf32>, tensor<*xf32>) -> tensor<*x!tf_type.variant> // CHECK: return %[[RET]] : tensor<*x!tf_type.variant> diff --git a/tensorflow/compiler/mlir/tfr/tests/raise_to_tf.mlir b/tensorflow/compiler/mlir/tfr/tests/raise_to_tf.mlir index 4be59b531365f2..14d21277dce2f1 100644 --- a/tensorflow/compiler/mlir/tfr/tests/raise_to_tf.mlir +++ b/tensorflow/compiler/mlir/tfr/tests/raise_to_tf.mlir @@ -17,7 +17,7 @@ func.func @decompose_tf_same(%arg0: tensor<1x2x3x4x!tf_type.string>) -> tensor<1 func.return %2 : tensor<1x2x3x4x!tf_type.string> // CHECK: %[[id:.*]] = "tf.RiscSame"(%arg0) : (tensor<1x2x3x4x!tf_type.string>) -> tensor<*x!tf_type.string> -// CHECK: %[[es:.*]] = "tf.EnsureShape"(%[[id]]) {shape = #tf_type.shape<1x2x3x4>} : (tensor<*x!tf_type.string>) -> tensor<1x2x3x4x!tf_type.string> +// CHECK: %[[es:.*]] = "tf.EnsureShape"(%[[id]]) <{shape = #tf_type.shape<1x2x3x4>}> : (tensor<*x!tf_type.string>) -> tensor<1x2x3x4x!tf_type.string> // CHECK: return %[[es]] : tensor<1x2x3x4x!tf_type.string> } @@ -32,7 +32,7 @@ func.func @decompose_tf_consecutive(%arg0: tensor<1x2x3x4x!tf_type.string>, %arg // CHECK: %[[id0:.*]] = "tf.RiscSame"(%arg0) : (tensor<1x2x3x4x!tf_type.string>) -> tensor<*x!tf_type.string> // CHECK: %[[id2:.*]] = "tf.RiscSame"(%arg2) : (tensor) -> tensor<*xf32> -// CHECK: %[[es:.*]] = "tf.EnsureShape"(%[[id2]]) {shape = #tf_type.shape<>} : (tensor<*xf32>) -> tensor +// CHECK: %[[es:.*]] = "tf.EnsureShape"(%[[id2]]) <{shape = #tf_type.shape<>}> : (tensor<*xf32>) -> tensor // CHECK: return %[[es]] : tensor } @@ -47,7 +47,7 @@ func.func @decompose_tf_concat_n(%arg0: tensor, %arg1: tensor, %arg2: func.return %4 : tensor<3xf32> // CHECK: %[[concat:.*]] = "tf.RiscConcat"(%arg0, %arg1, %arg2) : (tensor, tensor, tensor) -> tensor<*xf32> -// CHECK: %[[es:.*]] = "tf.EnsureShape"(%[[concat]]) {shape = #tf_type.shape<3>} : (tensor<*xf32>) -> tensor<3xf32> +// CHECK: %[[es:.*]] = "tf.EnsureShape"(%[[concat]]) <{shape = #tf_type.shape<3>}> : (tensor<*xf32>) -> tensor<3xf32> // CHECK: return %[[es]] : tensor<3xf32> } @@ -62,7 +62,7 @@ func.func @decompose_tf_split(%arg0: tensor<3xf32>) -> (tensor) { func.return %4 : tensor // CHECK: %[[split:.*]]:3 = "tf.RiscSplit"(%arg0) {N = 3 : i32} : (tensor<3xf32>) -> (tensor<*xf32>, tensor<*xf32>, tensor<*xf32>) -// CHECK: %[[es:.*]] = "tf.EnsureShape"(%[[split]]#0) {shape = #tf_type.shape<>} : (tensor<*xf32>) -> tensor +// CHECK: %[[es:.*]] = "tf.EnsureShape"(%[[split]]#0) <{shape = #tf_type.shape<>}> : (tensor<*xf32>) -> tensor // CHECK: return %[[es]] : tensor } @@ -75,7 +75,7 @@ func.func @decompose_tf_cast(%arg0: tensor) -> tensor { func.return %4 : tensor // CHECK: %[[tfcast:.*]] = "tf.RiscCast"(%arg0) {K = i32} : (tensor) -> tensor<*xi32> -// CHECK: %[[es:.*]] = "tf.EnsureShape"(%[[tfcast]]) {shape = #tf_type.shape<>} : (tensor<*xi32>) -> tensor +// CHECK: %[[es:.*]] = "tf.EnsureShape"(%[[tfcast]]) <{shape = #tf_type.shape<>}> : (tensor<*xi32>) -> tensor // CHECK: return %[[es]] : tensor } @@ -87,7 +87,7 @@ func.func @convert_to_scalar_tensor() -> tensor { %4 = "tfr.cast"(%cst) : (!tfr.tensor) -> tensor func.return %4 : tensor -// CHECK: %[[cst:.*]] = "tf.Const"() {value = dense<3.000000e+00> : tensor} : () -> tensor +// CHECK: %[[cst:.*]] = "tf.Const"() <{value = dense<3.000000e+00> : tensor}> : () -> tensor // CHECK: return %[[cst]] : tensor } @@ -100,7 +100,7 @@ func.func @attribute_propagate(%arg0: tensor) -> tensor { func.return %4 : tensor // CHECK: %[[tfcast:.*]] = "tf.RiscCast"(%arg0) {K = i32, _tpu_replicate, device = "hello"} : (tensor) -> tensor<*xi32> -// CHECK: %[[es:.*]] = "tf.EnsureShape"(%[[tfcast]]) {shape = #tf_type.shape<>} : (tensor<*xi32>) -> tensor +// CHECK: %[[es:.*]] = "tf.EnsureShape"(%[[tfcast]]) <{shape = #tf_type.shape<>}> : (tensor<*xi32>) -> tensor // CHECK: return %[[es]] : tensor } @@ -111,7 +111,7 @@ func.func @fixed_element_attribute(%arg0: tensor<2xf32>) -> tensor<2xi1> { %2 = "tfr.cast"(%1) : (!tfr.tensor) -> tensor<2xi1> func.return %2 : tensor<2xi1> // CHECK: %[[positive:.*]] = "tf.Positive"(%arg0) : (tensor<2xf32>) -> tensor<*xi1> -// CHECK: %[[res:.*]] = "tf.EnsureShape"(%[[positive]]) {shape = #tf_type.shape<2>} : (tensor<*xi1>) -> tensor<2xi1> +// CHECK: %[[res:.*]] = "tf.EnsureShape"(%[[positive]]) <{shape = #tf_type.shape<2>}> : (tensor<*xi1>) -> tensor<2xi1> // CHECK: return %[[res]] : tensor<2xi1> } diff --git a/tensorflow/compiler/mlir/tfr/tests/rewrite_quantized_io.mlir b/tensorflow/compiler/mlir/tfr/tests/rewrite_quantized_io.mlir index af2f843144ce5f..05823bfad27005 100644 --- a/tensorflow/compiler/mlir/tfr/tests/rewrite_quantized_io.mlir +++ b/tensorflow/compiler/mlir/tfr/tests/rewrite_quantized_io.mlir @@ -22,8 +22,8 @@ func.func @remove_quantized_io( %1 = "tf.Intermediate"(%arg1) : (tensor<1x5xf32>) -> tensor<1x5xf32> func.return %0, %1 : tensor<1x10x!quant.uniform>, tensor<1x5xf32> -// CHECK-DAG: %[[scale:.*]] = "tf.Const"() {value = dense<1.000000e-01> : tensor} : () -> tensor -// CHECK-DAG: %[[zp:.*]] = "tf.Const"() {value = dense<-128> : tensor} : () -> tensor +// CHECK-DAG: %[[scale:.*]] = "tf.Const"() <{value = dense<1.000000e-01> : tensor}> : () -> tensor +// CHECK-DAG: %[[zp:.*]] = "tf.Const"() <{value = dense<-128> : tensor}> : () -> tensor // CHECK: %[[quant:.*]] = "tfr.cast"(%arg0) : (tensor<1x10xi8>) -> !tfr.tensor // CHECK: %[[scale_cast:.*]] = "tfr.cast"(%[[scale]]) // CHECK: %[[zp_cast:.*]] = "tfr.cast"(%[[zp]]) diff --git a/tensorflow/compiler/mlir/tfrt/BUILD b/tensorflow/compiler/mlir/tfrt/BUILD index 5a3626061d6050..50b19af44ecf52 100644 --- a/tensorflow/compiler/mlir/tfrt/BUILD +++ b/tensorflow/compiler/mlir/tfrt/BUILD @@ -368,6 +368,7 @@ cc_library( "//tensorflow/compiler/mlir/tensorflow/transforms/host_runtime:lower_cluster_to_runtime_ops", "//tensorflow/compiler/mlir/tf2xla/api/v2:cluster_tf", "//tensorflow/compiler/mlir/tf2xla/api/v2:tf_dialect_to_executor", + "//tensorflow/compiler/tf2xla:xla_op_registry", "//tensorflow/core:framework", "//tensorflow/core:lib", "//tensorflow/core/common_runtime:function_body", diff --git a/tensorflow/compiler/mlir/tfrt/tests/fuse_tpu_compile_and_execute_ops.mlir b/tensorflow/compiler/mlir/tfrt/tests/fuse_tpu_compile_and_execute_ops.mlir index 823cbbc6f24cf7..75eff829e60952 100644 --- a/tensorflow/compiler/mlir/tfrt/tests/fuse_tpu_compile_and_execute_ops.mlir +++ b/tensorflow/compiler/mlir/tfrt/tests/fuse_tpu_compile_and_execute_ops.mlir @@ -11,7 +11,7 @@ func.func private @test_fuse_tpu_ops(%arg0: tensor<*xi32>, %arg1: tensor<*x!tf_t // CHECK-NOT: tf.TPUExecuteOp // CHECK-NEXT: %0 = "tf.ReadVariableOp"(%arg1) - // CHECK: [[key:%.*]], [[exec_result:%.*]] = "tf.TPUCompileMlirAndExecute"(%arg0, %0) {metadata = "metadata", mlir_module = "mlir_module", operandSegmentSizes = array, operands_with_static_shape = [], producer_name = "default"} : (tensor<*xi32>, tensor<*xi32>) -> (tensor<3x!tf_type.string>, tensor<*xi32>) + // CHECK: [[key:%.*]], [[exec_result:%.*]] = "tf.TPUCompileMlirAndExecute"(%arg0, %0) <{metadata = "metadata", mlir_module = "mlir_module", operandSegmentSizes = array, operands_with_static_shape = [], producer_name = "default"}> : (tensor<*xi32>, tensor<*xi32>) -> (tensor<3x!tf_type.string>, tensor<*xi32>) // CHECK-NEXT: return [[exec_result]] : tensor<*xi32> %0 = "tf.ReadVariableOp"(%arg1) {device = "/CPU:0"} : (tensor<*x!tf_type.resource>) -> tensor<*xi32> @@ -38,8 +38,8 @@ func.func private @test_outside_compilation(%arg0: tensor<*xi32>, %arg1: tensor< // CHECK-NOT: tf.TPUExecuteOp // CHECK-NEXT: %0 = "tf.ReadVariableOp"(%arg1) - // CHECK: [[key:%.*]], [[exec_result:%.*]] = "tf.TPUCompileMlirAndExecute"(%arg0, %0) {metadata = "metadata", mlir_module = "mlir_module", operandSegmentSizes = array, operands_with_static_shape = [], producer_name = "default"} : (tensor<*xi32>, tensor<*xi32>) -> (tensor<3x!tf_type.string>, tensor<*xi32>) - // CHECK-NEXT: "tf._XlaSendFromHost"(%arg0, %0, [[key]]) {_xla_has_host_transfer = true, device = "/job:localhost/replica:0/task:0/device:CPU:0", device_ordinal = 0 : i64, key = "host_compute_channel_0_retvals"} : (tensor<*xi32>, tensor<*xi32>, tensor<3x!tf_type.string>) -> () + // CHECK: [[key:%.*]], [[exec_result:%.*]] = "tf.TPUCompileMlirAndExecute"(%arg0, %0) <{metadata = "metadata", mlir_module = "mlir_module", operandSegmentSizes = array, operands_with_static_shape = [], producer_name = "default"}> : (tensor<*xi32>, tensor<*xi32>) -> (tensor<3x!tf_type.string>, tensor<*xi32>) + // CHECK-NEXT: "tf._XlaSendFromHost"(%arg0, %0, [[key]]) <{device_ordinal = 0 : i64, key = "host_compute_channel_0_retvals"}> {_xla_has_host_transfer = true, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<*xi32>, tensor<*xi32>, tensor<3x!tf_type.string>) -> () // CHECK-NEXT: return [[exec_result]] : tensor<*xi32> %0 = "tf.ReadVariableOp"(%arg1) {device = "/CPU:0"} : (tensor<*x!tf_type.resource>) -> tensor<*xi32> %1 = "tf.Shape"(%arg0) {device = "/CPU:0"} : (tensor<*xi32>) -> tensor @@ -69,8 +69,8 @@ func.func private @test_fuse_dynamic_dimension_ops(%arg0: tensor, %arg1 // CHECK: [[read_result:%.*]] = "tf.ReadVariableOp"(%arg1) // CHECK: [[shape_result_1:%.*]] = "tf.Shape"(%arg0) {device = "/CPU:0"} : (tensor) -> tensor // CHECK: [[shape_result_2:%.*]] = "tf.Shape"([[read_result]]) {device = "/CPU:0"} : (tensor<*xi32>) -> tensor - // CHECK: [[key:%.*]], [[exec_result:%.*]] = "tf.TPUCompileMlirAndExecute"(%arg0, [[shape_result_2]], %0, %0, %arg2, %arg4, %arg3) {metadata = "metadata", mlir_module = "mlir_module", operandSegmentSizes = array, operands_with_static_shape = [0 : i32, 1 : i32, 3 : i32], producer_name = "default"} : (tensor, tensor, tensor<*xi32>, tensor<*xi32>, tensor<2xi64>, tensor, tensor) -> (tensor<3x!tf_type.string>, tensor<*xi32>) - // CHECK: [[key_1:%.*]], [[exec_result_1:%.*]] = "tf.TPUCompileMlirAndExecute"(%arg0, %2, %0, %1) {metadata = "metadata", mlir_module = "mlir_module", operandSegmentSizes = array, operands_with_static_shape = [], producer_name = "default"} : (tensor, tensor, tensor<*xi32>, tensor) -> (tensor<3x!tf_type.string>, tensor<*xi32>) + // CHECK: [[key:%.*]], [[exec_result:%.*]] = "tf.TPUCompileMlirAndExecute"(%arg0, [[shape_result_2]], %0, %0, %arg2, %arg4, %arg3) <{metadata = "metadata", mlir_module = "mlir_module", operandSegmentSizes = array, operands_with_static_shape = [0 : i32, 1 : i32, 3 : i32], producer_name = "default"}> : (tensor, tensor, tensor<*xi32>, tensor<*xi32>, tensor<2xi64>, tensor, tensor) -> (tensor<3x!tf_type.string>, tensor<*xi32>) + // CHECK: [[key_1:%.*]], [[exec_result_1:%.*]] = "tf.TPUCompileMlirAndExecute"(%arg0, %2, %0, %1) <{metadata = "metadata", mlir_module = "mlir_module", operandSegmentSizes = array, operands_with_static_shape = [], producer_name = "default"}> : (tensor, tensor, tensor<*xi32>, tensor) -> (tensor<3x!tf_type.string>, tensor<*xi32>) // CHECK-NEXT: return [[exec_result]] : tensor<*xi32> %0 = "tf.ReadVariableOp"(%arg1) {device = "/CPU:0"} : (tensor<*x!tf_type.resource>) -> tensor<*xi32> %dyn_arg0 = "tf.SetStaticDimensionBounds" (%arg0, %arg2) :(tensor, tensor<2xi64>) -> tensor diff --git a/tensorflow/compiler/mlir/tfrt/tests/hoist_invariant_ops.mlir b/tensorflow/compiler/mlir/tfrt/tests/hoist_invariant_ops.mlir index 28c08a1a369f0e..e6d5aec8285a0b 100644 --- a/tensorflow/compiler/mlir/tfrt/tests/hoist_invariant_ops.mlir +++ b/tensorflow/compiler/mlir/tfrt/tests/hoist_invariant_ops.mlir @@ -5,16 +5,16 @@ module attributes {tf_saved_model.semantics} { // Test hoisting varhandle op. // CHECK-LABEL: func @_tfrt_resource_init -// CHECK: [[handle:%.*]] = "tf.VarHandleOp"() {container = "", shared_name = "x"} : () -> tensor>> +// CHECK: [[handle:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "x"}> : () -> tensor>> // CHECK: [[x:%.*]] = "tf.ReadVariableOp"([[handle]]) {device = "/CPU:0", dtype = i32} : (tensor>>) -> tensor -// CHECK: "tf._TfrtSetResource"([[x]]) {device = "/CPU:0", index = 0 : i64} : (tensor) -> () +// CHECK: "tf._TfrtSetResource"([[x]]) <{index = 0 : i64}> {device = "/CPU:0"} : (tensor) -> () // CHECK-LABEL: func @test_hoist_varhandleop func.func @hoist_varhandleop(%arg: tensor {tf_saved_model.index_path = ["input"]}) -> (tensor {tf_saved_model.index_path = ["r"]}) attributes {tf_saved_model.exported_names = ["test_hoist_varhandleop"]} { // CHECK-NOT: tf.VarHandleOp // CHECK-NOT: tf.ReadVariableOp - // CHECK: [[v:%.*]] = "tf._TfrtGetResource"() {container = [""], device = "/CPU:0", indices = [0], shared_name = [""]} : () -> tensor + // CHECK: [[v:%.*]] = "tf._TfrtGetResource"() <{container = [""], indices = [0], shared_name = [""]}> {device = "/CPU:0"} : () -> tensor // CHECK: [[r:%.*]] = "tf.AddV2"({{.*}}, [[v]]) {device = "/CPU:0"} : (tensor, tensor) -> tensor // CHECK: return [[r]] %handle = "tf.VarHandleOp"() {container = "", shared_name = "x"} : () -> tensor>> @@ -34,16 +34,16 @@ module attributes {tf_saved_model.semantics} { // CHECK-LABEL: func @_tfrt_resource_init // CHECK: [[handle:%.*]] = "tf.HashTableV2"() // CHECK-SAME: shared_name = "x" -// CHECK: "tf._TfrtSetResource"([[handle]]) {device = "/job:localhost/replica:0/task:0/device:CPU:0", index = [[handle_idx:.*]] : i64} +// CHECK: "tf._TfrtSetResource"([[handle]]) <{index = [[handle_idx:.*]] : i64}> {device = "/job:localhost/replica:0/task:0/device:CPU:0"} // CHECK: [[x:%.*]] = "tf.LookupTableSizeV2"([[handle]]) -// CHECK: "tf._TfrtSetResource"([[x]]) {device = "/job:localhost/replica:0/task:0/device:CPU:0", index = [[size_idx:.*]] : i64} : (tensor) -> () +// CHECK: "tf._TfrtSetResource"([[x]]) <{index = [[size_idx:.*]] : i64}> {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor) -> () // CHECK: func @test_hoist_hash_table func.func @hoist_hash_table(%arg: tensor {tf_saved_model.index_path = ["input"]}, %default: tensor {tf_saved_model.index_path = ["default"]}) -> (tensor {tf_saved_model.index_path = ["r"]}, tensor<*xi64> {tf_saved_model.index_path = ["r1"]}) attributes {tf_saved_model.exported_names = ["test_hoist_hash_table"]} { // CHECK-NOT: tf.HashTableV2 // CHECK-NOT: tf.LookupTableSizeV2 - // CHECK: [[v:%.*]]:2 = "tf._TfrtGetResource"() {container = ["", ""], device = "/job:localhost/replica:0/task:0/device:CPU:0", indices = [0, 1], shared_name = [{{.*}}, {{.*}}]} + // CHECK: [[v:%.*]]:2 = "tf._TfrtGetResource"() <{container = ["", ""], indices = [0, 1], shared_name = [{{.*}}, {{.*}}]}> {device = "/job:localhost/replica:0/task:0/device:CPU:0"} // CHECK: [[r:%.*]] = "tf.LookupTableFindV2"([[v]]#[[handle_idx]] // CHECK: return [[v]]#[[size_idx]], [[r]] %0 = "tf.HashTableV2"() {container = "", device = "", key_dtype = !tf_type.string, shared_name = "x", use_node_name_sharing = false, value_dtype = i64} : () -> tensor @@ -61,17 +61,17 @@ module attributes {tf_saved_model.semantics} { // Test hoisting const op. // CHECK-LABEL: func @_tfrt_resource_init -// CHECK: [[const:%.*]] = "tf.Const"() {device = "/CPU:0", value = dense<0> : tensor} : () -> tensor +// CHECK: [[const:%.*]] = "tf.Const"() <{value = dense<0> : tensor}> {device = "/CPU:0"} : () -> tensor // CHECK: [[x:%.*]] = "tf.AddV2"([[const]], [[const]]) {device = "/CPU:0"} : (tensor, tensor) -> tensor -// CHECK: "tf._TfrtSetResource"([[x]]) {device = "/CPU:0", index = 0 : i64} : (tensor) -> () -// CHECK: [[const_1:%.*]] = "tf.Const"() {device = "/CPU:0", value = dense<1> : tensor} : () -> tensor -// CHECK: "tf._TfrtSetResource"([[const_1]]) {device = "/CPU:0", index = 1 : i64} : (tensor) -> () +// CHECK: "tf._TfrtSetResource"([[x]]) <{index = 0 : i64}> {device = "/CPU:0"} : (tensor) -> () +// CHECK: [[const_1:%.*]] = "tf.Const"() <{value = dense<1> : tensor}> {device = "/CPU:0"} : () -> tensor +// CHECK: "tf._TfrtSetResource"([[const_1]]) <{index = 1 : i64}> {device = "/CPU:0"} : (tensor) -> () // CHECK-LABEL: func @test_hoist_const func.func @hoist_const(%arg: tensor {tf_saved_model.index_path = ["input"]}) -> (tensor {tf_saved_model.index_path = ["r"]}) attributes {tf_saved_model.exported_names = ["test_hoist_const"]} { // CHECK-NOT: tf.Const - // CHECK: [[v:%.*]] = "tf._TfrtGetResource"() {container = [""], device = "/CPU:0", indices = [0], shared_name = [""]} : () -> tensor + // CHECK: [[v:%.*]] = "tf._TfrtGetResource"() <{container = [""], indices = [0], shared_name = [""]}> {device = "/CPU:0"} : () -> tensor // CHECK-NEXT: "tf.AddV2"({{.*}}, [[v]]) {device = "/CPU:0"} : (tensor, tensor) -> tensor // CHECK-NEXT: return %const = "tf.Const"() {device = "/CPU:0", value = dense<0> : tensor} : () -> tensor @@ -84,7 +84,7 @@ func.func @hoist_const(%arg: tensor {tf_saved_model.index_path = ["input"]} func.func @hoist_const_return(%arg: tensor {tf_saved_model.index_path = ["input"]}) -> (tensor {tf_saved_model.index_path = ["r"]}) attributes {tf_saved_model.exported_names = ["test_hoist_const_return"]} { // CHECK-NOT: tf.Const - // CHECK: [[v:%.*]] = "tf._TfrtGetResource"() {container = [""], device = "/CPU:0", indices = [1], shared_name = [""]} : () -> tensor + // CHECK: [[v:%.*]] = "tf._TfrtGetResource"() <{container = [""], indices = [1], shared_name = [""]}> {device = "/CPU:0"} : () -> tensor // CHECK-NEXT: return [[v]] %const = "tf.Const"() {device = "/CPU:0", value = dense<1> : tensor} : () -> tensor func.return %const : tensor @@ -99,17 +99,17 @@ module attributes {tf_saved_model.semantics} { // Test hoisting write side-effect ops. // CHECK-LABEL: func @_tfrt_resource_init -// CHECK: [[const:%.*]] = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<0> : tensor} : () -> tensor -// CHECK: "tf._TfrtSetResource"([[const]]) {device = "/job:localhost/replica:0/task:0/device:CPU:0", index = [[const_idx:.*]] : i64} : (tensor) -> () -// CHECK: [[handle:%.*]] = "tf.VarHandleOp"() {container = "", shared_name = "x"} : () -> tensor>> -// CHECK: "tf._TfrtSetResource"([[handle]]) {device = "/job:localhost/replica:0/task:0/device:CPU:0", index = [[handle_idx:.*]] : i64} : (tensor>>) -> () +// CHECK: [[const:%.*]] = "tf.Const"() <{value = dense<0> : tensor}> {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : () -> tensor +// CHECK: "tf._TfrtSetResource"([[const]]) <{index = [[const_idx:.*]] : i64}> {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor) -> () +// CHECK: [[handle:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "x"}> : () -> tensor>> +// CHECK: "tf._TfrtSetResource"([[handle]]) <{index = [[handle_idx:.*]] : i64}> {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor>>) -> () // CHECK: func @test_hoist_var_read_write func.func @hoist_var_read_write() -> (tensor {tf_saved_model.index_path = ["x"]}, tensor {tf_saved_model.index_path = ["r"]}) attributes {tf_saved_model.exported_names = ["test_hoist_var_read_write"]} { // CHECK-NOT: tf.Const // CHECK-NOT: tf.VarHandleOp - // CHECK: [[v:%.*]]:2 = "tf._TfrtGetResource"() {container = ["", ""], device = "/job:localhost/replica:0/task:0/device:CPU:0", indices = [0, 1], shared_name = [{{.*}}, {{.*}}]} : () -> ({{.*}}) + // CHECK: [[v:%.*]]:2 = "tf._TfrtGetResource"() <{container = ["", ""], indices = [0, 1], shared_name = [{{.*}}, {{.*}}]}> {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : () -> ({{.*}}) // CHECK: [[x:%.*]] = "tf.ReadVariableOp"([[v]]#[[handle_idx]]) {device = "/CPU:0", dtype = i32} : (tensor>>) -> tensor // CHECK-NEXT: "tf.AssignVariable"([[v]]#[[handle_idx]], [[v]]#[[const_idx]]) {device = "/CPU:0"} : (tensor>>, tensor) -> () // CHECK-NEXT: [[r:%.*]] = "tf.ReadVariableOp"([[v]]#[[handle_idx]]) {device = "/CPU:0", dtype = i32} : (tensor>>) -> tensor @@ -131,13 +131,13 @@ module attributes {tf_saved_model.semantics} { // Test not hoisting read variable op that used by control flow ops if var handle op and read variable op are separated, but still hoists const ops and var handle ops. // CHECK-LABEL: func @_tfrt_resource_init -// CHECK: [[handle:%.*]] = "tf.VarHandleOp"() {container = "", shared_name = "x"} : () -> tensor>> +// CHECK: [[handle:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "x"}> : () -> tensor>> // CHECK: "tf._TfrtSetResource"([[handle]]) // CHECK-SAME: index = [[handle_index:.*]] -// CHECK: [[handle1:%.*]] = "tf.VarHandleOp"() {container = "", shared_name = "x"} : () -> tensor>> +// CHECK: [[handle1:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "x"}> : () -> tensor>> // CHECK: "tf._TfrtSetResource"([[handle1]]) // CHECK-SAME: index = [[handle1_index:.*]] -// CHECK: [[const:%.*]] = "tf.Const"() {device = "/CPU:0", value = dense : tensor} : () -> tensor +// CHECK: [[const:%.*]] = "tf.Const"() <{value = dense : tensor}> {device = "/CPU:0"} : () -> tensor // CHECK: "tf._TfrtSetResource"([[const]]) // CHECK-SAME: index = [[const_index:.*]] func.func private @some_func( @@ -164,7 +164,7 @@ func.func @not_hoist_if(%arg: tensor {tf_saved_model.index_path = ["input"] attributes {tf_saved_model.exported_names = ["test_not_hoist_if"]} { %handle = "tf.VarHandleOp"() {container = "", shared_name = "x"} : () -> tensor>> // CHECK-NOT: tf.Const - // CHECK: "tf._TfrtGetResource"() + // CHECK: "tf._TfrtGetResource"() %cond = "tf.Const"() {device = "/CPU:0", value = dense : tensor} : () -> tensor // CHECK: tf.If %x = "tf.If"(%cond, %handle) {then_branch = @some_func, else_branch = @some_func, is_stateless = false} : (tensor, tensor>>) -> tensor @@ -185,7 +185,7 @@ func.func private @batched_function(%arg0: tensor<1x3xf32>) -> tensor<1x3xf32> attributes {tf._input_shapes = [#tf_type.shape<1x3>, #tf_type.shape<*>], tf.signature.is_stateful} { // CHECK-NOT: tf.VarHandleOp // CHECK-NOT: tf.ReadVariableOp - // CHECK: "tf._TfrtGetResource"() + // CHECK: "tf._TfrtGetResource"() %0 = "tf.VarHandleOp"() {device = "/device:CPU:0", container = "", shared_name = "variable"} : () -> tensor>> %1 = "tf.ReadVariableOp"(%0) {device = "/device:CPU:0"} : (tensor>>) -> tensor<1x3xf32> %2 = "tf.AddV2"(%arg0, %1) {device = "/device:CPU:0"} : (tensor<1x3xf32>, tensor<1x3xf32>) -> tensor<1x3xf32> @@ -197,7 +197,7 @@ func.func private @batched_function(%arg0: tensor<1x3xf32>) -> tensor<1x3xf32> func.func @main(%arg0: tensor<1x3xf32> {tf_saved_model.index_path = ["input"]}) -> (tensor<*xf32> {tf_saved_model.index_path = ["r"]}) attributes {tf_saved_model.exported_names = ["main"]} { // CHECK-NOT: tf.VarHandleOp - // CHECK: "tf._TfrtGetResource"() + // CHECK: "tf._TfrtGetResource"() %0 = "tf.VarHandleOp"() {device = "/device:CPU:0", container = "", shared_name = "variable"} : () -> tensor>> // CHECK: "tf.BatchFunction"(%arg0, %0) // CHECK: operandSegmentSizes = array @@ -288,4 +288,4 @@ func.func private @func2(%arg: tensor) -> tensor { func.return %r : tensor } -} \ No newline at end of file +} diff --git a/tensorflow/compiler/mlir/tfrt/tests/hoist_invariant_ops_mlrt.mlir b/tensorflow/compiler/mlir/tfrt/tests/hoist_invariant_ops_mlrt.mlir index 7b797b357a1539..7f82726b0740fa 100644 --- a/tensorflow/compiler/mlir/tfrt/tests/hoist_invariant_ops_mlrt.mlir +++ b/tensorflow/compiler/mlir/tfrt/tests/hoist_invariant_ops_mlrt.mlir @@ -7,17 +7,17 @@ module attributes {tf_saved_model.semantics} { // CHECK-LABEL: func @_tfrt_resource_init // CHECK: [[handle:%.*]] = "tf.HashTableV2"() // CHECK-SAME: shared_name = "x" -// CHECK: "tf._TfrtSetResource"([[handle]]) {device = "/job:localhost/replica:0/task:0/device:CPU:0", index = [[handle_id:.*]] : i64} +// CHECK: "tf._TfrtSetResource"([[handle]]) <{index = [[handle_id:.*]] : i64}> {device = "/job:localhost/replica:0/task:0/device:CPU:0"} // CHECK: [[x:%.*]] = "tf.LookupTableSizeV2"([[handle]]) -// CHECK: "tf._TfrtSetResource"([[x]]) {device = "/job:localhost/replica:0/task:0/device:CPU:0", index = [[size_id:.*]] : i64} : (tensor) -> () +// CHECK: "tf._TfrtSetResource"([[x]]) <{index = [[size_id:.*]] : i64}> {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor) -> () // CHECK: func @test_hoist_hash_table func.func @hoist_hash_table(%arg: tensor {tf_saved_model.index_path = ["input"]}, %default: tensor {tf_saved_model.index_path = ["default"]}) -> (tensor {tf_saved_model.index_path = ["r"]}, tensor<*xi64> {tf_saved_model.index_path = ["r1"]}) attributes {tf_saved_model.exported_names = ["test_hoist_hash_table"]} { // CHECK-NOT: tf.HashTableV2 // CHECK-NOT: tf.LookupTableSizeV2 - // CHECK-DAG: [[v0:%.*]] = "tf._TfrtGetResource"() {container = [""], device = "/job:localhost/replica:0/task:0/device:CPU:0", indices = [[[handle_id]]], shared_name = [{{.*}}]} - // CHECK-DAG: [[v1:%.*]] = "tf._TfrtGetResource"() {container = [""], device = "/job:localhost/replica:0/task:0/device:CPU:0", indices = [[[size_id]]], shared_name = [{{.*}}]} + // CHECK-DAG: [[v0:%.*]] = "tf._TfrtGetResource"() <{container = [""], indices = [[[handle_id]]], shared_name = [{{.*}}]}> {device = "/job:localhost/replica:0/task:0/device:CPU:0"} + // CHECK-DAG: [[v1:%.*]] = "tf._TfrtGetResource"() <{container = [""], indices = [[[size_id]]], shared_name = [{{.*}}]}> {device = "/job:localhost/replica:0/task:0/device:CPU:0"} // CHECK-DAG: [[r:%.*]] = "tf.LookupTableFindV2"([[v0]] // CHECK-DAG: return [[v1]], [[r]] %0 = "tf.HashTableV2"() {container = "", device = "", key_dtype = !tf_type.string, shared_name = "x", use_node_name_sharing = false, value_dtype = i64} : () -> tensor diff --git a/tensorflow/compiler/mlir/tfrt/tests/mlrt/async_while.mlir b/tensorflow/compiler/mlir/tfrt/tests/mlrt/async_while.mlir index b13096c9ea0fb2..d7fee9dee62358 100644 --- a/tensorflow/compiler/mlir/tfrt/tests/mlrt/async_while.mlir +++ b/tensorflow/compiler/mlir/tfrt/tests/mlrt/async_while.mlir @@ -22,11 +22,11 @@ func.func private @"map/while_body"(%loop_count: tensor, %max_iterations: t } // CHECK-LABEL: func.func private @"map/while_body/TfMlrtAsyncWhileBody"(%arg0: !mlrt.promise, %arg1: !mlrt.future, %arg2: !mlrt.promise, %arg3: !mlrt.future, %arg4: !mlrt.promise, %arg5: tensor, %arg6: tensor, %arg7: tensor<*xf32>) { -// CHECK-NEXT: %cst = "tf.Const"() {value = dense<1> : tensor} : () -> tensor +// CHECK-NEXT: %cst = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor // CHECK-NEXT: %0 = "tf_mlrt.tf_await"(%arg1) : (!mlrt.future) -> tensor // CHECK-NEXT: %1 = "tf.AddV2"(%0, %cst) : (tensor, tensor) -> tensor // CHECK-NEXT: "tf_mlrt.tf_promise"(%arg2, %1) : (!mlrt.promise, tensor) -> () -// CHECK-NEXT: %2 = "tf.PartitionedCall"(%1, %arg5) {config = "", config_proto = "", executor_type = "", f = @"map/while_cond/TfMlrtAsyncWhilePredicate"} : (tensor, tensor) -> tensor +// CHECK-NEXT: %2 = "tf.PartitionedCall"(%1, %arg5) <{config = "", config_proto = "", executor_type = "", f = @"map/while_cond/TfMlrtAsyncWhilePredicate"}> : (tensor, tensor) -> tensor // CHECK-NEXT: "tf_mlrt.tf_promise"(%arg0, %2) : (!mlrt.promise, tensor) -> () // CHECK-NEXT: %3 = "tf.TensorArrayReadV3"(%arg6, %0, %arg7) : (tensor, tensor, tensor<*xf32>) -> tensor<3x3xf32> // CHECK-NEXT: %4 = "tf_mlrt.tf_await"(%arg3) : (!mlrt.future) -> tensor<3x3xf32> @@ -37,7 +37,7 @@ func.func private @"map/while_body"(%loop_count: tensor, %max_iterations: t //CHECK-LABEL: func.func @serving_default func.func @serving_default(%max_iterations: tensor, %array_handle: tensor, %array_flow: tensor<*xf32>, %matrix: tensor<3x3xf32>) -> (tensor<3x3xf32>, tensor<*xf32>) { %cst_0 = "tf.Const"() {value = dense<0> : tensor} : () -> tensor - // CHECK: %0 = "tf.PartitionedCall"(%cst, %arg0) {config = "", config_proto = "", executor_type = "", f = @"map/while_cond/TfMlrtAsyncWhilePredicate"} : (tensor, tensor) -> tensor + // CHECK: %0 = "tf.PartitionedCall"(%cst, %arg0) <{config = "", config_proto = "", executor_type = "", f = @"map/while_cond/TfMlrtAsyncWhilePredicate"}> : (tensor, tensor) -> tensor // CHECK-NEXT: %1:6 = tf_mlrt.tf_async_while @"map/while_body/TfMlrtAsyncWhileBody"(%0, %cst, %arg3, %arg0, %arg1, %arg2) {invariant_size = 3 : i32} : (tensor, tensor, tensor<3x3xf32>, tensor, tensor, tensor<*xf32>) -> (!mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future) %1:5 = "tf.While"(%cst_0, %max_iterations, %array_handle, %array_flow, %matrix) {body= @"map/while_body", cond = @"map/while_cond", is_stateless = false, parallel_iterations = 10 : i64, shape_invariant} : (tensor, tensor, tensor, tensor<*xf32>, tensor<3x3xf32>) -> (tensor, tensor, tensor, tensor<*xf32>, tensor<3x3xf32>) // CHECK-NEXT: %2 = "tf_mlrt.tf_await"(%1#5) : (!mlrt.future) -> tensor<*xf32> @@ -50,10 +50,10 @@ func.func @serving_default(%max_iterations: tensor, %array_handle: tensor, %array_handle: tensor, %array_flow: tensor<*xf32>, %matrix: tensor<3x3xf32>, %array_handle_2: tensor, %array_flow_2: tensor<*xf32>, %matrix_2: tensor<3x3xf32>) -> (tensor<3x3xf32>, tensor<3x3xf32>) { %cst_0 = "tf.Const"() {value = dense<0> : tensor} : () -> tensor - // CHECK: %0 = "tf.PartitionedCall"(%cst, %arg0) {config = "", config_proto = "", executor_type = "", f = @"map/while_cond/TfMlrtAsyncWhilePredicate"} : (tensor, tensor) -> tensor + // CHECK: %0 = "tf.PartitionedCall"(%cst, %arg0) <{config = "", config_proto = "", executor_type = "", f = @"map/while_cond/TfMlrtAsyncWhilePredicate"}> : (tensor, tensor) -> tensor // CHECK-NEXT: %1:6 = tf_mlrt.tf_async_while @"map/while_body/TfMlrtAsyncWhileBody"(%0, %cst, %arg3, %arg0, %arg1, %arg2) {invariant_size = 3 : i32} : (tensor, tensor, tensor<3x3xf32>, tensor, tensor, tensor<*xf32>) -> (!mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future) %1:5 = "tf.While"(%cst_0, %max_iterations, %array_handle, %array_flow, %matrix) {body= @"map/while_body", cond = @"map/while_cond", is_stateless = false, parallel_iterations = 10 : i64, shape_invariant} : (tensor, tensor, tensor, tensor<*xf32>, tensor<3x3xf32>) -> (tensor, tensor, tensor, tensor<*xf32>, tensor<3x3xf32>) - // CHECK: %2 = "tf.PartitionedCall"(%cst, %arg0) {config = "", config_proto = "", executor_type = "", f = @"map/while_cond/TfMlrtAsyncWhilePredicate"} : (tensor, tensor) -> tensor + // CHECK: %2 = "tf.PartitionedCall"(%cst, %arg0) <{config = "", config_proto = "", executor_type = "", f = @"map/while_cond/TfMlrtAsyncWhilePredicate"}> : (tensor, tensor) -> tensor // CHECK-NEXT: %3:6 = tf_mlrt.tf_async_while @"map/while_body/TfMlrtAsyncWhileBody"(%2, %cst, %arg6, %arg0, %arg4, %arg5) {invariant_size = 3 : i32} : (tensor, tensor, tensor<3x3xf32>, tensor, tensor, tensor<*xf32>) -> (!mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future) %2:5 = "tf.While"(%cst_0, %max_iterations, %array_handle_2, %array_flow_2, %matrix_2) {body= @"map/while_body", cond = @"map/while_cond", is_stateless = false, parallel_iterations = 10 : i64, shape_invariant} : (tensor, tensor, tensor, tensor<*xf32>, tensor<3x3xf32>) -> (tensor, tensor, tensor, tensor<*xf32>, tensor<3x3xf32>) // CHECK-NEXT: %4 = "tf_mlrt.tf_await"(%1#2) : (!mlrt.future) -> tensor<3x3xf32> @@ -128,11 +128,11 @@ func.func private @"random/while_body"(%loop_count: tensor, %max_iterations } // CHECK-LABEL: func.func private @"random/while_body/TfMlrtAsyncWhileBody_1"(%arg0: !mlrt.promise, %arg1: !mlrt.future, %arg2: !mlrt.promise, %arg3: !mlrt.future, %arg4: !mlrt.promise, %arg5: tensor, %arg6: tensor, %arg7: tensor<*xf32>) { -// CHECK-NEXT: %cst = "tf.Const"() {value = dense<1> : tensor} : () -> tensor +// CHECK-NEXT: %cst = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor // CHECK-NEXT: %0 = "tf_mlrt.tf_await"(%arg1) : (!mlrt.future) -> tensor // CHECK-NEXT: %1 = "tf.AddV2"(%0, %cst) : (tensor, tensor) -> tensor // CHECK-NEXT: "tf_mlrt.tf_promise"(%arg2, %1) : (!mlrt.promise, tensor) -> () -// CHECK-NEXT: %2 = "tf.PartitionedCall"(%1, %arg5) {config = "", config_proto = "", executor_type = "", f = @"random/while_cond/TfMlrtAsyncWhilePredicate_0"} : (tensor, tensor) -> tensor +// CHECK-NEXT: %2 = "tf.PartitionedCall"(%1, %arg5) <{config = "", config_proto = "", executor_type = "", f = @"random/while_cond/TfMlrtAsyncWhilePredicate_0"}> : (tensor, tensor) -> tensor // CHECK-NEXT: "tf_mlrt.tf_promise"(%arg0, %2) : (!mlrt.promise, tensor) -> () // CHECK-NEXT: %3 = "tf.TensorArrayReadV3"(%arg6, %0, %arg7) : (tensor, tensor, tensor<*xf32>) -> tensor<3x3xf32> // CHECK-NEXT: %4 = "tf_mlrt.tf_await"(%arg3) : (!mlrt.future) -> tensor<3x3xf32> @@ -143,7 +143,7 @@ func.func private @"random/while_body"(%loop_count: tensor, %max_iterations //CHECK-LABEL: func.func @random_serving_default func.func @random_serving_default(%max_iterations: tensor, %array_handle: tensor, %array_flow: tensor<*xf32>, %matrix: tensor<3x3xf32>) -> (tensor<3x3xf32>, tensor<*xf32>) { %cst_0 = "tf.Const"() {value = dense<0> : tensor} : () -> tensor - // CHECK: %0 = "tf.PartitionedCall"(%cst, %arg0) {config = "", config_proto = "", executor_type = "", f = @"random/while_cond/TfMlrtAsyncWhilePredicate_0"} : (tensor, tensor) -> tensor + // CHECK: %0 = "tf.PartitionedCall"(%cst, %arg0) <{config = "", config_proto = "", executor_type = "", f = @"random/while_cond/TfMlrtAsyncWhilePredicate_0"}> : (tensor, tensor) -> tensor // CHECK-NEXT: %1:6 = tf_mlrt.tf_async_while @"random/while_body/TfMlrtAsyncWhileBody_1"(%0, %cst, %arg3, %arg0, %arg1, %arg2) {invariant_size = 3 : i32} : (tensor, tensor, tensor<3x3xf32>, tensor, tensor, tensor<*xf32>) -> (!mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future) %1:5 = "tf.While"(%cst_0, %max_iterations, %array_handle, %array_flow, %matrix) {body= @"random/while_body", cond = @"random/while_cond", is_stateless = false, parallel_iterations = 10 : i64, shape_invariant} : (tensor, tensor, tensor, tensor<*xf32>, tensor<3x3xf32>) -> (tensor, tensor, tensor, tensor<*xf32>, tensor<3x3xf32>) // CHECK-NEXT: %2 = "tf_mlrt.tf_await"(%1#5) : (!mlrt.future) -> tensor<*xf32> @@ -166,7 +166,7 @@ func.func private @"sort_map/while_cond"(%loop_count: tensor, %max_iteratio // CHECK-NEXT: return %0 : tensor // CHECK-LABEL: func.func private @"sort_map/while_body" -// CHECK-NEXT: %cst = "tf.Const"() {value = dense<1> : tensor} : () -> tensor +// CHECK-NEXT: %cst = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor // CHECK-NEXT: %0 = "tf.AddV2"(%arg0, %cst) : (tensor, tensor) -> tensor // CHECK-NEXT: %1 = "tf.TensorArrayReadV3"(%arg2, %arg0, %arg3) : (tensor, tensor, tensor<*xf32>) -> tensor<3x3xf32> // CHECK-NEXT: %2 = "tf.TensorArrayReadV3"(%arg5, %arg0, %arg6) : (tensor, tensor, tensor<*xf32>) -> tensor<3x3xf32> @@ -177,7 +177,7 @@ func.func private @"sort_map/while_cond"(%loop_count: tensor, %max_iteratio // CHECK-NEXT: %7 = "tf.Identity"(%4) : (tensor) -> tensor // CHECK-NEXT: %8 = "tf.MatMul"(%5, %arg7) : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32> // CHECK-NEXT: %9 = "tf.Select"(%7, %8, %arg7) : (tensor, tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32> -// CHECK-NEXT: return %0, %arg1, %arg2, %arg3, %6, %arg5, %arg6, %9, %arg8 +// CHECK-NEXT: return %0, %arg1, %arg2, %arg3, %6, %arg5, %arg6, %9, %arg8 func.func private @"sort_map/while_body"(%loop_count: tensor, %max_iterations: tensor, %handle: tensor, %flow_in: tensor<*xf32>, %matrix: tensor<3x3xf32>, %handle_2: tensor, %flow_in_2: tensor<*xf32>, %matrix_2: tensor<3x3xf32>, %bound: tensor) -> (tensor, tensor, tensor, tensor<*xf32>, tensor<3x3xf32>, tensor, tensor<*xf32>, tensor<3x3xf32>, tensor) { %cst_1 = "tf.Const"() {value = dense<1> : tensor} : () -> tensor %updated_loop_count = "tf.AddV2"(%loop_count, %cst_1) : (tensor, tensor) -> tensor @@ -194,11 +194,11 @@ func.func private @"sort_map/while_body"(%loop_count: tensor, %max_iteratio } // CHECK-LABEL: func.func private @"sort_map/while_body/TfMlrtAsyncWhileBody" -// CHECK-NEXT: %cst = "tf.Const"() {value = dense<1> : tensor} : () -> tensor +// CHECK-NEXT: %cst = "tf.Const"() <{value = dense<1> : tensor}> : () -> tensor // CHECK-NEXT: %0 = "tf_mlrt.tf_await"(%arg1) : (!mlrt.future) -> tensor // CHECK-NEXT: %1 = "tf.AddV2"(%0, %cst) : (tensor, tensor) -> tensor // CHECK-NEXT: "tf_mlrt.tf_promise"(%arg2, %1) : (!mlrt.promise, tensor) -> () -// CHECK-NEXT: %2 = "tf.PartitionedCall"(%1, %arg7) {config = "", config_proto = "", executor_type = "", f = @"sort_map/while_cond/TfMlrtAsyncWhilePredicate"} : (tensor, tensor) -> tensor +// CHECK-NEXT: %2 = "tf.PartitionedCall"(%1, %arg7) <{config = "", config_proto = "", executor_type = "", f = @"sort_map/while_cond/TfMlrtAsyncWhilePredicate"}> : (tensor, tensor) -> tensor // CHECK-NEXT: "tf_mlrt.tf_promise"(%arg0, %2) : (!mlrt.promise, tensor) -> () // CHECK-NEXT: %3 = "tf.TensorArrayReadV3"(%arg8, %0, %arg9) : (tensor, tensor, tensor<*xf32>) -> tensor<3x3xf32> // CHECK-NEXT: %4 = "tf.TensorArrayReadV3"(%arg10, %0, %arg11) : (tensor, tensor, tensor<*xf32>) -> tensor<3x3xf32> @@ -218,7 +218,7 @@ func.func private @"sort_map/while_body"(%loop_count: tensor, %max_iteratio //CHECK-LABEL: func.func @sort_serving_default func.func @sort_serving_default(%max_iterations: tensor, %array_handle: tensor, %array_flow: tensor<*xf32>, %matrix: tensor<3x3xf32>, %bound: tensor) -> (tensor<3x3xf32>, tensor<3x3xf32>, tensor<*xf32>) { %cst_0 = "tf.Const"() {value = dense<0> : tensor} : () -> tensor - // CHECK: %0 = "tf.PartitionedCall"(%cst, %arg0) {config = "", config_proto = "", executor_type = "", f = @"sort_map/while_cond/TfMlrtAsyncWhilePredicate"} : (tensor, tensor) -> tensor + // CHECK: %0 = "tf.PartitionedCall"(%cst, %arg0) <{config = "", config_proto = "", executor_type = "", f = @"sort_map/while_cond/TfMlrtAsyncWhilePredicate"}> : (tensor, tensor) -> tensor // CHECK-NEXT: %1:10 = tf_mlrt.tf_async_while @"sort_map/while_body/TfMlrtAsyncWhileBody"(%0, %cst, %arg3, %arg3, %arg0, %arg1, %arg2, %arg1, %arg2, %arg4) {invariant_size = 6 : i32} : (tensor, tensor, tensor<3x3xf32>, tensor<3x3xf32>, tensor, tensor, tensor<*xf32>, tensor, tensor<*xf32>, tensor) -> (!mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future) %1:9 = "tf.While"(%cst_0, %max_iterations, %array_handle, %array_flow, %matrix , %array_handle, %array_flow, %matrix, %bound) {body= @"sort_map/while_body", cond = @"sort_map/while_cond", is_stateless = false, parallel_iterations = 10 : i64, shape_invariant} : (tensor, tensor, tensor, tensor<*xf32>, tensor<3x3xf32>, tensor, tensor<*xf32>, tensor<3x3xf32>,tensor) -> (tensor, tensor, tensor, tensor<*xf32>, tensor<3x3xf32>, tensor, tensor<*xf32>, tensor<3x3xf32>, tensor) // CHECK-NEXT: %2 = "tf_mlrt.tf_await"(%1#6) : (!mlrt.future) -> tensor<*xf32> diff --git a/tensorflow/compiler/mlir/tfrt/tests/mlrt/while_to_map_fn.mlir b/tensorflow/compiler/mlir/tfrt/tests/mlrt/while_to_map_fn.mlir index 27c92289a5bddf..2b1f5fc9b17a4e 100644 --- a/tensorflow/compiler/mlir/tfrt/tests/mlrt/while_to_map_fn.mlir +++ b/tensorflow/compiler/mlir/tfrt/tests/mlrt/while_to_map_fn.mlir @@ -37,7 +37,7 @@ func.func private @"map/while_body"(%arg0: tensor, %arg1: tensor, %arg // CHECK-SAME: (%arg0: !mlrt.future, %arg1: !mlrt.promise, %arg2: tensor, %arg3: tensor, %arg4: tensor) // CHECK: [[det:%.*]] = "tf.MatrixDeterminant" // CHECK-NEXT: [[ta_0:%.*]] = "tf_mlrt.tf_await"(%arg0) : (!mlrt.future) -> tensor>> -// CHECK-NEXT: [[ta_1:%.*]] = "tf.TensorListSetItem"([[ta_0]], %arg3, [[det]]) { +// CHECK-NEXT: [[ta_1:%.*]] = "tf.TensorListSetItem"([[ta_0]], %arg3, [[det]]) <{ // CHECK-NEXT: "tf_mlrt.tf_promise"(%arg1, [[ta_1]]) : (!mlrt.promise, tensor>>) -> () // CHECK-NEXT: return @@ -53,7 +53,7 @@ func.func @serving_default(%arg0: tensor {tf.device = "/job:localhost/rep // CHECK-SAME: {body_fn = @"map/while_body/MapFnBody", num_tensor_list_or_flow_in = 1 : i32} // CHECK-NOT: tf.While %1:4 = "tf.While"(%cst, %cst, %0, %arg0) {_lower_using_switch_merge = true, _num_original_outputs = 6 : i64, _read_only_resource_inputs = [], _xla_propagate_compile_time_consts = true, body = @"map/while_body", cond = @"map/while_cond", device = "/job:localhost/replica:0/task:0/device:CPU:0", is_stateless = true, parallel_iterations = 4 : i64, shape_invariant} : (tensor, tensor, tensor>>, tensor) -> (tensor, tensor, tensor>>, tensor) - // CHECK-NEXT: "tf.TensorListStack"([[map_fn_result]], %cst_0) { + // CHECK-NEXT: "tf.TensorListStack"([[map_fn_result]], %cst_0) <{ %2 = "tf.TensorListStack"(%1#2, %cst_0) {device = "/job:localhost/replica:0/task:0/device:CPU:0", num_elements = 3 : i64} : (tensor>>, tensor<0xi32>) -> tensor<3xf32> return %2 : tensor<3xf32> } @@ -458,7 +458,7 @@ func.func private @tf.WhileRegion2_body(%arg0: tensor<*xi32>) -> (tensor>>, tensor, tensor) -> tensor // CHECK: TensorArrayGatherV3 %6 = "tf.TensorArrayGatherV3"(%handle_14, %2, %4#3) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<2x!tf_type.resource>>, tensor, tensor) -> tensor - return %5, %6 : tensor, tensor + return %5, %6 : tensor, tensor } // ----- diff --git a/tensorflow/compiler/mlir/tfrt/tests/sink_in_invariant_ops.mlir b/tensorflow/compiler/mlir/tfrt/tests/sink_in_invariant_ops.mlir index 2b0344cd7202a2..42e2e7ccb5086a 100644 --- a/tensorflow/compiler/mlir/tfrt/tests/sink_in_invariant_ops.mlir +++ b/tensorflow/compiler/mlir/tfrt/tests/sink_in_invariant_ops.mlir @@ -93,8 +93,8 @@ module attributes {tf_saved_model.semantics} { // CHECK-LABEL: func private @batched_function func.func private @batched_function(%arg0: tensor>>, %arg1: tensor>>) -> tensor<1x3xf32> attributes {tf._input_shapes = [#tf_type.shape<1x3>, #tf_type.shape<*>], tf.signature.is_stateful} { - // CHECK-DAG: [[handle1:%.*]] = "tf.VarHandleOp"() {{{.*}}, shared_name = "variable1"} - // CHECK-DAG: [[handle2:%.*]] = "tf.VarHandleOp"() {{{.*}}, shared_name = "variable2"} + // CHECK-DAG: [[handle1:%.*]] = "tf.VarHandleOp"() <{{{.*}}, shared_name = "variable1"}> + // CHECK-DAG: [[handle2:%.*]] = "tf.VarHandleOp"() <{{{.*}}, shared_name = "variable2"}> // CHECK: "tf.ReadVariableOp"([[handle1]]) // CHECK: "tf.ReadVariableOp"([[handle2]]) %0 = "tf.ReadVariableOp"(%arg0) {device = "/device:CPU:0"} : (tensor>>) -> tensor<1x3xf32> @@ -298,7 +298,7 @@ func.func private @nested_batched_function(%arg0: tensor<1x3xf32>, %arg1: tensor } // CHECK-LABEL: func @main -func.func @main(%arg0: tensor<1x3xf32> {tf_saved_model.index_path = ["input"]}) -> (tensor<*xf32> {tf_saved_model.index_path = ["r"]}) +func.func @main(%arg0: tensor<1x3xf32> {tf_saved_model.index_path = ["input"]}) -> (tensor<*xf32> {tf_saved_model.index_path = ["r"]}) attributes {tf_saved_model.exported_names = ["main"]} { // CHECK: [[handle:%.*]] = "tf.VarHandleOp"() %0 = "tf.VarHandleOp"() {device = "/device:CPU:0", container = "", shared_name = "variable"} : () -> tensor>> diff --git a/tensorflow/compiler/mlir/tfrt/tests/tf_to_corert/merge_tf_if_ops.mlir b/tensorflow/compiler/mlir/tfrt/tests/tf_to_corert/merge_tf_if_ops.mlir index f79c31963b3693..cb1990762f0b20 100644 --- a/tensorflow/compiler/mlir/tfrt/tests/tf_to_corert/merge_tf_if_ops.mlir +++ b/tensorflow/compiler/mlir/tfrt/tests/tf_to_corert/merge_tf_if_ops.mlir @@ -51,8 +51,8 @@ func.func @nested_if_op_else_1(%cond: tensor, %x: tensor, %y: tensor, [[y:%.*]]: tensor) -// CHECK-DAG: [[cst:%.*]] = "tf.Const"() {value = dense<1> : tensor} -// CHECK-DAG: [[cst_0:%.*]] = "tf.Const"() {value = dense<2> : tensor} +// CHECK-DAG: [[cst:%.*]] = "tf.Const"() <{value = dense<1> : tensor}> +// CHECK-DAG: [[cst_0:%.*]] = "tf.Const"() <{value = dense<2> : tensor}> // CHECK: [[r0:%.*]] = "tf.AddV2"([[x]], [[cst]]) // CHECK: [[r1:%.*]] = "tf.AddV2"([[y]], [[r0]]) // CHECK: [[r2:%.*]] = "tf.AddV2"([[x]], [[cst_0]]) @@ -63,7 +63,7 @@ func.func @nested_if_op_else_1(%cond: tensor, %x: tensor, %y: tensor, [[y:%.*]]: tensor, [[cond:%.*]]: tensor) func.func @merge_stateless(%x: tensor, %y: tensor, %cond: tensor) -> (tensor, tensor, tensor) { // CHECK-NEXT: [[res:%.*]]:3 = "tf.If"([[cond]], [[x]], [[y]]) - // CHECK-SAME: {else_branch = @merge_stateless_merged_if_0_0_else, is_stateless = true, then_branch = @merge_stateless_merged_if_0_0_then} + // CHECK-SAME: <{else_branch = @merge_stateless_merged_if_0_0_else, is_stateless = true, then_branch = @merge_stateless_merged_if_0_0_then}> // CHECK-SAME: (tensor, tensor, tensor) -> (tensor, tensor, tensor) // CHECK-NEXT: return [[res]]#0, [[res]]#1, [[res]]#2 %0, %1 = "tf.If"(%cond, %x, %y) {else_branch = @no_side_effect_else_0, then_branch = @no_side_effect_then_0, is_stateless = true} : (tensor, tensor, tensor) -> (tensor, tensor) @@ -83,8 +83,8 @@ func.func @merge_stateless(%x: tensor, %y: tensor, %cond: tensor) // CHECK-LABEL: func private @merge_nested_if_op_merged_if_0_0_else_merged_if_1_0_else // CHECK-SAME: ([[x:%.*]]: tensor, [[y:%.*]]: tensor) -// CHECK-NEXT: [[cst:%.*]] = "tf.Const"() {value = dense<2> : tensor} -// CHECK-NEXT: [[cst_0:%.*]] = "tf.Const"() {value = dense<1> : tensor} +// CHECK-NEXT: [[cst:%.*]] = "tf.Const"() <{value = dense<2> : tensor}> +// CHECK-NEXT: [[cst_0:%.*]] = "tf.Const"() <{value = dense<1> : tensor}> // CHECK-NEXT: [[r0:%.*]] = "tf.AddV2"([[x]], [[cst_0]]) // CHECK-NEXT: [[r1:%.*]] = "tf.AddV2"([[y]], [[r0]]) // CHECK-NEXT: [[r2:%.*]] = "tf.AddV2"([[x]], [[cst]]) @@ -93,14 +93,14 @@ func.func @merge_stateless(%x: tensor, %y: tensor, %cond: tensor) // CHECK-LABEL: func private @merge_nested_if_op_merged_if_0_0_else // CHECK-SAME: ([[cond:%.*]]: tensor, [[x:%.*]]: tensor, [[y:%.*]]: tensor) -// CHECK-NEXT: [[r0:%.*]]:3 = "tf.If"(%arg0, %arg1, %arg2) {else_branch = @merge_nested_if_op_merged_if_0_0_else_merged_if_1_0_else, is_stateless = true, then_branch = @merge_nested_if_op_merged_if_0_0_else_merged_if_1_0_then} : (tensor, tensor, tensor) -> (tensor, tensor, tensor) +// CHECK-NEXT: [[r0:%.*]]:3 = "tf.If"(%arg0, %arg1, %arg2) <{else_branch = @merge_nested_if_op_merged_if_0_0_else_merged_if_1_0_else, is_stateless = true, then_branch = @merge_nested_if_op_merged_if_0_0_else_merged_if_1_0_then}> : (tensor, tensor, tensor) -> (tensor, tensor, tensor) // CHECK-NEXT: return [[r0]]#0, [[r0]]#1, [[r0]]#2 // CHECK-LABEL: func @merge_nested_if_op // CHECK-SAME: ([[x:%.*]]: tensor, [[y:%.*]]: tensor, [[cond:%.*]]: tensor, [[nested_cond:%.*]]: tensor) func.func @merge_nested_if_op(%x: tensor, %y: tensor, %cond: tensor, %nested_cond: tensor) -> (tensor, tensor, tensor) { // CHECK-NEXT: [[res:%.*]]:3 = "tf.If"([[cond]], [[nested_cond]], [[x]], [[y]]) - // CHECK-SAME: {else_branch = @merge_nested_if_op_merged_if_0_0_else, is_stateless = true, then_branch = @merge_nested_if_op_merged_if_0_0_then} + // CHECK-SAME: <{else_branch = @merge_nested_if_op_merged_if_0_0_else, is_stateless = true, then_branch = @merge_nested_if_op_merged_if_0_0_then}> // CHECK-SAME: (tensor, tensor, tensor, tensor) -> (tensor, tensor, tensor) // CHECK-NEXT: return [[res]]#0, [[res]]#1, [[res]]#2 %0, %1 = "tf.If"(%cond, %nested_cond, %x, %y) {else_branch = @nested_if_op_else_0, then_branch = @nested_if_op_then_0, is_stateless = true} : (tensor, tensor, tensor, tensor) -> (tensor, tensor) @@ -122,7 +122,7 @@ func.func @merge_side_effect(%x: tensor, %y: tensor, %cond: tensor // CHECK-LABEL: func @multiple_uses func.func @multiple_uses(%x: tensor, %y: tensor, %cond: tensor) -> (tensor, tensor, tensor) { // CHECK-NEXT: tf.If - // CHECK-SAME: {else_branch = @multiple_uses_merged_if_0_0_else, is_stateless = true, then_branch = @multiple_uses_merged_if_0_0_then} + // CHECK-SAME: <{else_branch = @multiple_uses_merged_if_0_0_else, is_stateless = true, then_branch = @multiple_uses_merged_if_0_0_then}> %0, %1 = "tf.If"(%cond, %x, %y) {else_branch = @no_side_effect_else_0, then_branch = @no_side_effect_then_0, is_stateless = true} : (tensor, tensor, tensor) -> (tensor, tensor) %2 = "tf.If"(%cond, %x, %y) {else_branch = @no_side_effect_else_1, then_branch = @no_side_effect_then_1, is_stateless = true} : (tensor, tensor, tensor) -> tensor func.return %0, %1, %2 : tensor, tensor, tensor diff --git a/tensorflow/compiler/mlir/tfrt/transforms/mlrt/import_model.cc b/tensorflow/compiler/mlir/tfrt/transforms/mlrt/import_model.cc index 28072642a22746..6e04fe1c1e23a1 100644 --- a/tensorflow/compiler/mlir/tfrt/transforms/mlrt/import_model.cc +++ b/tensorflow/compiler/mlir/tfrt/transforms/mlrt/import_model.cc @@ -14,7 +14,9 @@ limitations under the License. ==============================================================================*/ #include "tensorflow/compiler/mlir/tfrt/transforms/mlrt/import_model.h" +#include #include +#include #include "absl/log/log.h" #include "absl/status/status.h" @@ -49,10 +51,10 @@ namespace tensorflow { namespace mlrt_compiler { StatusOr ConvertTfMlirToBytecode( - const TfrtCompileOptions& options, - const tfrt_stub::FallbackState& fallback_state, mlir::ModuleOp module, - tfrt_stub::ModelRuntimeContext& model_context, - mlir::OwningOpRef* module_with_op_keys) { + const TfrtCompileOptions& options, tfrt_stub::FallbackState& fallback_state, + mlir::ModuleOp module, tfrt_stub::ModelRuntimeContext& model_context, + mlir::OwningOpRef* module_with_op_keys, + std::vector* added_xla_function_names) { mlrt::bc::Buffer bytecode_buffer; TF_RETURN_IF_ERROR(ConvertTfMlirToRuntimeExecutable( options, module, @@ -127,7 +129,7 @@ StatusOr ConvertTfMlirToBytecode( bytecode_buffer = std::move(*statusor); return OkStatus(); }, - model_context)); + model_context, &fallback_state, added_xla_function_names)); return bytecode_buffer; } diff --git a/tensorflow/compiler/mlir/tfrt/transforms/mlrt/import_model.h b/tensorflow/compiler/mlir/tfrt/transforms/mlrt/import_model.h index 87dc685694235a..ef9caeb47337a2 100644 --- a/tensorflow/compiler/mlir/tfrt/transforms/mlrt/import_model.h +++ b/tensorflow/compiler/mlir/tfrt/transforms/mlrt/import_model.h @@ -33,10 +33,10 @@ namespace mlrt_compiler { // // This is for initial conversion. StatusOr ConvertTfMlirToBytecode( - const TfrtCompileOptions& options, - const tfrt_stub::FallbackState& fallback_state, mlir::ModuleOp module, - tfrt_stub::ModelRuntimeContext& model_context, - mlir::OwningOpRef* module_with_op_keys = nullptr); + const TfrtCompileOptions& options, tfrt_stub::FallbackState& fallback_state, + mlir::ModuleOp module, tfrt_stub::ModelRuntimeContext& model_context, + mlir::OwningOpRef* module_with_op_keys = nullptr, + std::vector* added_xla_function_names = nullptr); // Converts an MLIR `module_with_op_keys` in TF dialect to MLRT's bytecode // format, with op costs from `cost_recorder`. diff --git a/tensorflow/compiler/mlir/tfrt/translate/import_model.cc b/tensorflow/compiler/mlir/tfrt/translate/import_model.cc index 1b933933de6673..9a8d41a4db7d03 100644 --- a/tensorflow/compiler/mlir/tfrt/translate/import_model.cc +++ b/tensorflow/compiler/mlir/tfrt/translate/import_model.cc @@ -44,6 +44,7 @@ limitations under the License. #include "tensorflow/compiler/mlir/tfrt/transforms/tfrt_pipeline_options.h" #include "tensorflow/compiler/mlir/tfrt/transforms/tpu_passes.h" #include "tensorflow/compiler/mlir/tfrt/translate/tfrt_compile_options.h" +#include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "tensorflow/core/common_runtime/function_body.h" #include "tensorflow/core/common_runtime/function_def_utils.h" #include "tensorflow/core/platform/env.h" @@ -233,7 +234,17 @@ Status ConvertTfMlirToRuntimeExecutable( "Failed to process TPUPartitionedCallOp for fallback execution")); } } else if (options.device_target == TfrtDeviceInfraTarget::kGpu) { - TF_RETURN_IF_ERROR(mlir::TF::RunTFXLABridge(module)); + TF_RETURN_IF_ERROR( + tensorflow::tf2xla::v2::RunFunctionTf2xlaClusteringBridge( + module, tf2xla::v2::DeviceType::XLA_GPU_JIT, + /*is_in_fallback_enabled_mode=*/false)); + + TF_RETURN_IF_ERROR( + tensorflow::tfrt_compiler::RunLowerClusterToRuntimeOpsPassPipeline( + module, tsl::DeviceType(DEVICE_GPU_XLA_JIT))); + + TF_RETURN_IF_ERROR( + tensorflow::tf2xla::v2::ExportFromTensorflowDialectToExecutor(module)); if (options.serialize_mlir_module_to_aot_packages) { const std::string mlir_string = SerializeMlirModule(module); diff --git a/tensorflow/compiler/tests/xla_call_module_test.py b/tensorflow/compiler/tests/xla_call_module_test.py index d46ab642bc81b1..ea31b4c072e607 100644 --- a/tensorflow/compiler/tests/xla_call_module_test.py +++ b/tensorflow/compiler/tests/xla_call_module_test.py @@ -92,12 +92,12 @@ def f(x): self._assertOpOutputMatchesExpected(f, (x,), (np.sin(np.cos(x)),)) - def test_basic_with_token(self): + def test_basic_with_token_v8(self): x = np.array([1.0, 2.0, 3.0], dtype=np.float32) def f(x): # sin(cos(x)) - module, version = serialize(""" + module, _ = serialize(""" module @jit_f.0 { func.func public @main(%arg0: !stablehlo.token, %arg1: tensor<3xf32>) -> (!stablehlo.token, tensor<3xf32>) { %0 = stablehlo.cosine %arg1 : tensor<3xf32> @@ -105,6 +105,32 @@ def f(x): return %arg0, %1 : !stablehlo.token, tensor<3xf32> } } +""") + return xla.call_module( + [x], + version=8, # Version 8 uses only one prefix token + module=module, + Tout=[x.dtype], + Sout=[x.shape], + has_token_input_output=True, # Version 8 cares about this + platforms=[self.testing_platform()], + ) + + self._assertOpOutputMatchesExpected(f, (x,), (np.sin(np.cos(x)),)) + + def test_basic_with_multiple_tokens(self): + x = np.array([1.0, 2.0, 3.0], dtype=np.float32) + + def f(x): + # sin(cos(x)) + module, version = serialize(""" +module @jit_f.0 { + func.func public @main(%arg0: !stablehlo.token {jax.token = true}, %arg1: !stablehlo.token {jax.token = true}, %arg2: tensor<3xf32>) -> (!stablehlo.token, !stablehlo.token, tensor<3xf32>) { + %0 = stablehlo.cosine %arg2 : tensor<3xf32> + %1 = stablehlo.sine %0 : tensor<3xf32> + return %arg0, %arg1, %1 : !stablehlo.token, !stablehlo.token, tensor<3xf32> + } +} """) return xla.call_module( [x], @@ -112,7 +138,31 @@ def f(x): module=module, Tout=[x.dtype], Sout=[x.shape], - has_token_input_output=True, + platforms=[self.testing_platform()], + ) + + self._assertOpOutputMatchesExpected(f, (x,), (np.sin(np.cos(x)),)) + + def test_basic_with_tokens_preceeded_by_other_args(self): + x = np.array([1.0, 2.0, 3.0], dtype=np.float32) + + def f(x): + # sin(cos(x)) + module, version = serialize(""" +module @jit_f.0 { + func.func public @main(%arg0: tensor, %arg1: !stablehlo.token {jax.token = true}, %arg2: !stablehlo.token {jax.token = true}, %arg3: tensor<3xf32>) -> (!stablehlo.token, !stablehlo.token, tensor<3xf32>) { + %0 = stablehlo.cosine %arg3 : tensor<3xf32> + %1 = stablehlo.sine %0 : tensor<3xf32> + return %arg1, %arg2, %1 : !stablehlo.token, !stablehlo.token, tensor<3xf32> + } +} +""") + return xla.call_module( + [np.int32(0), x], + version=version, + module=module, + Tout=[x.dtype], + Sout=[x.shape], platforms=[self.testing_platform()], ) @@ -183,7 +233,7 @@ def f(x): # x: f32[2, b] %0, %1 = call @dyn_main(%arg0_new, %arg1) : (tensor<{dim_var_type}>, tensor<2x?xf32>) -> (tensor<2x?xf32>, tensor<{dim_var_type}>) return %0, %1 : tensor<2x?xf32>, tensor<{dim_var_type}> }} - func.func private @dyn_main(%arg0: tensor<{dim_var_type}>, %arg1: tensor<2x?xf32>) -> (tensor<2x?xf32>, tensor<{dim_var_type}>) {{ + func.func private @dyn_main(%arg0: tensor<{dim_var_type}> {{jax.global_constant = "b"}}, %arg1: tensor<2x?xf32>) -> (tensor<2x?xf32>, tensor<{dim_var_type}>) {{ %0 = stablehlo.sine %arg1 : tensor<2x?xf32> return %0, %arg0 : tensor<2x?xf32>, tensor<{dim_var_type}> }} @@ -278,7 +328,7 @@ def test_platforms_basic(self, *, platform_idx_type: str): # returns x + 2. on CPU, x + 3. on GPU (CUDA or ROCM) and x + 4. on TPU module, version = serialize(f""" module @jit_f.0 {{ - func.func public @main(%arg_platform_idx: tensor<{platform_idx_type}>, %arg0: tensor) -> tensor {{ + func.func public @main(%arg_platform_idx: tensor<{platform_idx_type}> {{jax.global_constant = "_platform_index"}}, %arg0: tensor) -> tensor {{ %0 = stablehlo.convert %arg_platform_idx : (tensor<{platform_idx_type}>) -> tensor %to_add = "stablehlo.case"(%0) ({{ %cpu_val = stablehlo.constant dense<2.> : tensor @@ -319,7 +369,7 @@ def test_platforms_unknown_custom_call(self): # returns x + 2. on CPU, x + 3. on GPU, and x + 4. on TPU module, version = serialize(""" module @jit_f.0 { - func.func public @main(%arg_platform_idx: tensor, %arg0: tensor) -> tensor { + func.func public @main(%arg_platform_idx: tensor {jax.global_constant = "_platform_index"}, %arg0: tensor) -> tensor { %to_add = "stablehlo.case"(%arg_platform_idx) ({ %cpu_val = stablehlo.constant dense<2.> : tensor stablehlo.return %cpu_val : tensor @@ -358,13 +408,13 @@ def test_platforms_and_poly(self): module, version = serialize(""" module @jit_f_jax attributes {jax.uses_shape_polymorphism = true} { - func.func public @main(%arg_platform_idx: tensor, %arg0: tensor) -> (tensor) { + func.func public @main(%arg_platform_idx: tensor {jax.global_constant = "_platform_index"}, %arg0: tensor) -> (tensor) { %0 = stablehlo.get_dimension_size %arg0, dim = 0 : (tensor) -> tensor %5 = call @_wrapped_jax_export_main(%arg_platform_idx, %0, %arg0) : (tensor, tensor, tensor) -> tensor return %5 : tensor } - func.func private @_wrapped_jax_export_main(%arg_platform_idx: tensor, %arg0: tensor, %arg1: tensor) -> (tensor) { + func.func private @_wrapped_jax_export_main(%arg_platform_idx: tensor {jax.global_constant = "_platform_index"}, %arg0: tensor {jax.global_constant = "b"}, %arg1: tensor) -> (tensor) { %to_add = "stablehlo.case"(%arg_platform_idx) ({ %cpu_val = stablehlo.constant dense<2.> : tensor stablehlo.return %cpu_val : tensor @@ -395,6 +445,52 @@ def f(x): ) self._assertOpOutputMatchesExpected(f, (x,), (expected_value,)) + + def test_platforms_and_poly_and_tokens(self): + if test.is_built_with_rocm(): + self.skipTest('Currently failing on ROCm due to mismatch') + x = np.arange(6, dtype=np.float32) + # returns x + 2. on CPU, x + 3. on GPU (CUDA or ROCM) and x + 4. on TPU + + module, version = serialize(""" +module @jit_f_jax attributes {jax.uses_shape_polymorphism = true} { + func.func public @main(%arg_platform_idx: tensor {jax.global_constant = "_platform_index"}, %arg_tok: !stablehlo.token {jax.token = true}, %arg0: tensor) -> (!stablehlo.token, tensor) { + %0 = stablehlo.get_dimension_size %arg0, dim = 0 : (tensor) -> tensor + %5:2 = call @_wrapped_jax_export_main(%arg_platform_idx, %0, %arg_tok, %arg0) : (tensor, tensor, !stablehlo.token, tensor) -> (!stablehlo.token, tensor) + return %5#0, %5#1 : !stablehlo.token, tensor + } + + func.func private @_wrapped_jax_export_main(%arg_platform_idx: tensor {jax.global_constant = "_platform_index"}, %arg0: tensor {jax.global_constant = "b"}, %arg_tok: !stablehlo.token {jax.token = true}, %arg1: tensor) -> (!stablehlo.token, tensor) { + %to_add = "stablehlo.case"(%arg_platform_idx) ({ + %cpu_val = stablehlo.constant dense<2.> : tensor + stablehlo.return %cpu_val : tensor + }, { + %gpu_val = stablehlo.constant dense<3.> : tensor + stablehlo.return %gpu_val : tensor + }, { + %tpu_val = stablehlo.constant dense<4.> : tensor + stablehlo.return %tpu_val : tensor + }) : (tensor) -> tensor + %1 = stablehlo.reshape %arg0 : (tensor) -> tensor<1xi32> + %3 = stablehlo.dynamic_broadcast_in_dim %to_add, %1, dims = [] : (tensor, tensor<1xi32>) -> tensor + %4 = stablehlo.add %3, %arg1 : tensor + return %arg_tok, %4 : !stablehlo.token, tensor + } +} +""") + platforms = ['CPU', 'CUDA', 'ROCM', 'TPU'] + def f(x): + return xla.call_module([x], version=version, + module=module, + Tout=[np.float32], + Sout=[()], + platforms=platforms) + + expected_value = ( + x + dict(CPU=2.0, CUDA=3.0, ROCM=3.0, TPU=4.0)[self.testing_platform()] + ) + self._assertOpOutputMatchesExpected(f, (x,), (expected_value,)) + # A module used for testing errors related to use of "platforms". platforms_errors_module_str = """ module @jit_f.0 { @@ -403,7 +499,7 @@ def f(x): } } """ - + def platforms_errors_helper( self, *, @@ -742,7 +838,7 @@ def f(x): # x: f32[b, 5] %0 = call @dyn_main(%arg0_new, %arg1) : (tensor, tensor) -> tensor return %0 : tensor } - func.func private @dyn_main(%arg0: tensor, %arg1: tensor) -> tensor { + func.func private @dyn_main(%arg0: tensor {jax.global_constant = "b"}, %arg1: tensor) -> tensor { %0 = stablehlo.reshape %arg0 : (tensor) -> tensor<1xi32> %1 = "stablehlo.dynamic_iota"(%0) {iota_dimension = 0 : i64} : (tensor<1xi32>) -> tensor return %1 : tensor @@ -790,7 +886,7 @@ def f(x): # x: f32[b, 3] %0 = call @dyn_main(%arg0_new, %arg1) : (tensor, tensor) -> tensor return %0 : tensor } - func.func private @dyn_main(%arg0: tensor, %arg1: tensor) -> tensor { + func.func private @dyn_main(%arg0: tensor {jax.global_constant = "b"}, %arg1: tensor) -> tensor { %0 = stablehlo.constant dense<3> : tensor %1 = stablehlo.multiply %arg0, %0 : tensor %2 = stablehlo.reshape %1 : (tensor) -> tensor<1xi32> @@ -819,7 +915,7 @@ def f(x): # x: f32[b, 4] %0 = call @dyn_main(%arg0_new, %arg1) : (tensor, tensor) -> tensor return %0 : tensor } - func.func private @dyn_main(%arg0: tensor, %arg1: tensor) -> tensor { + func.func private @dyn_main(%arg0: tensor {jax.global_constant = "b"}, %arg1: tensor) -> tensor { %0 = stablehlo.constant dense<0> : tensor %1 = stablehlo.constant dense<0> : tensor<1xi64> %2 = stablehlo.reshape %arg0 : (tensor) -> tensor<1xi32> @@ -850,7 +946,7 @@ def f(x): # x: f32[b, 4] %0 = call @dyn_main(%arg0_new, %arg1) : (tensor, tensor) -> tensor<4xf32> return %0 : tensor<4xf32> } - func.func private @dyn_main(%arg0: tensor, %arg1: tensor) -> tensor<4xf32> { + func.func private @dyn_main(%arg0: tensor {jax.global_constant = "b"}, %arg1: tensor) -> tensor<4xf32> { %0 = stablehlo.constant dense<-1> : tensor %1 = stablehlo.add %arg0, %0 : tensor %2 = stablehlo.reshape %1 : (tensor) -> tensor<1xi32> @@ -887,7 +983,7 @@ def f(x, idx): # x: f32[b, 4] idx: i32 %0 = call @dyn_main(%arg0_new, %arg1, %arg2) : (tensor, tensor, tensor) -> tensor return %0 : tensor } - func.func private @dyn_main(%arg0: tensor, %arg1: tensor, %arg2: tensor) -> tensor { + func.func private @dyn_main(%arg0: tensor {jax.global_constant = "b"}, %arg1: tensor, %arg2: tensor) -> tensor { %0 = stablehlo.constant dense<0> : tensor %1 = stablehlo.compare LT, %arg2, %0, SIGNED : (tensor, tensor) -> tensor %2 = stablehlo.add %arg2, %arg0 : tensor @@ -920,7 +1016,7 @@ def f(x, y): # x: f32[b, 4] y: f32[2, b, 4] %0, %1 = call @dyn_main(%arg0_new, %arg1, %arg2) : (tensor, tensor, tensor<2x?x4xf32>) -> (tensor<2x?x4xf32>, tensor<2x?x4xf32>) return %0, %1 : tensor<2x?x4xf32>, tensor<2x?x4xf32> } - func.func private @dyn_main(%arg0: tensor, %arg1: tensor, %arg2: tensor<2x?x4xf32>) -> (tensor<2x?x4xf32>, tensor<2x?x4xf32>) { + func.func private @dyn_main(%arg0: tensor {jax.global_constant = "b"}, %arg1: tensor, %arg2: tensor<2x?x4xf32>) -> (tensor<2x?x4xf32>, tensor<2x?x4xf32>) { %0 = stablehlo.constant dense<2> : tensor<1xi32> %2 = stablehlo.reshape %arg0 : (tensor) -> tensor<1xi32> %3 = stablehlo.constant dense<4> : tensor<1xi32> @@ -952,7 +1048,7 @@ def f(x): # x: i32[b] %0 = call @dyn_main(%arg0_new, %arg1) : (tensor, tensor) -> tensor return %0 : tensor } - func.func private @dyn_main(%arg0: tensor, %arg1: tensor) -> tensor { + func.func private @dyn_main(%arg0: tensor {jax.global_constant = "b"}, %arg1: tensor) -> tensor { %0 = stablehlo.constant dense<0> : tensor %1 = stablehlo.reduce(%arg1 init: %0) across dimensions = [0] : (tensor, tensor) -> tensor reducer(%arg2: tensor, %arg3: tensor) { @@ -984,7 +1080,7 @@ def f(x): # x: f32[b, 5] %0 = call @dyn_main(%arg0_new, %arg1) : (tensor, tensor) -> tensor return %0 : tensor } - func.func private @dyn_main(%arg0: tensor, %arg1: tensor) -> tensor { + func.func private @dyn_main(%arg0: tensor {jax.global_constant = "b"}, %arg1: tensor) -> tensor { %0 = stablehlo.constant dense<0.000000e+00> : tensor %1 = stablehlo.reduce(%arg1 init: %0) across dimensions = [1] : (tensor, tensor) -> tensor reducer(%arg2: tensor, %arg3: tensor) { @@ -1020,11 +1116,11 @@ def f(x): # x: f32[b] %0 = call @dyn_main(%arg0_new, %arg1) : (tensor, tensor) -> tensor return %0 : tensor } - func.func private @dyn_main(%arg0: tensor, %arg1: tensor) -> tensor { + func.func private @dyn_main(%arg0: tensor {jax.global_constant = "b"}, %arg1: tensor) -> tensor { %0 = call @f(%arg0, %arg1) : (tensor, tensor) -> tensor return %0 : tensor } - func.func private @f(%arg0: tensor, %arg1: tensor) -> tensor { + func.func private @f(%arg0: tensor {jax.global_constant = "b"}, %arg1: tensor) -> tensor { %0 = stablehlo.reshape %arg0 : (tensor) -> tensor<1xi32> %1 = "stablehlo.dynamic_iota"(%0) {iota_dimension = 0 : i64} : (tensor<1xi32>) -> tensor return %1 : tensor @@ -1051,7 +1147,7 @@ def f(x): # x: f32[b] %0 = call @dyn_main(%arg0_new, %arg1) : (tensor, tensor) -> tensor return %0 : tensor } - func.func private @dyn_main(%arg0: tensor, %arg1: tensor) -> tensor { + func.func private @dyn_main(%arg0: tensor {jax.global_constant = "b"}, %arg1: tensor) -> tensor { return %arg1 : tensor } } @@ -1081,7 +1177,7 @@ def f(x): # x: f32[b] %0, %1 = call @dyn_main(%arg0_new, %arg1) : (tensor, tensor) -> (tensor, tensor) return %0, %1 : tensor, tensor } - func.func private @dyn_main(%arg0: tensor, %arg1: tensor) -> (tensor, tensor) { + func.func private @dyn_main(%arg0: tensor {jax.global_constant = "b"}, %arg1: tensor) -> (tensor, tensor) { %0 = stablehlo.constant dense<0> : tensor %1:2 = "stablehlo.while"(%arg1, %0) ({ ^bb0(%arg2: tensor, %arg3: tensor): @@ -1123,7 +1219,7 @@ def f(x): # x: f32[b] %0 = call @dyn_main(%arg0_new, %arg1) : (tensor, tensor) -> tensor return %0 : tensor }} - func.func private @dyn_main(%arg0: tensor, %arg1: tensor) -> tensor {{ + func.func private @dyn_main(%arg0: tensor {{jax.global_constant = "b"}}, %arg1: tensor) -> tensor {{ return %arg1 : tensor }} }} @@ -1311,7 +1407,6 @@ def f(x, y): Sout=[res.shape], platforms=[self.testing_platform()], function_list=(foo,), - has_token_input_output=True, ) self._assertOpOutputMatchesExpected(f, (x, y), (res,)) diff --git a/tensorflow/compiler/tf2xla/BUILD b/tensorflow/compiler/tf2xla/BUILD index 9c039e03135ab3..bd8b7a7a68f5ce 100644 --- a/tensorflow/compiler/tf2xla/BUILD +++ b/tensorflow/compiler/tf2xla/BUILD @@ -1151,6 +1151,7 @@ cc_library( visibility = [":internal"], deps = [ ":tf2xla_defs", + ":xla_op_registry", "//tensorflow/compiler/jit:flags", "//tensorflow/compiler/mlir:mlir_graph_optimization_pass", "//tensorflow/compiler/mlir/tensorflow", @@ -1159,6 +1160,7 @@ cc_library( "//tensorflow/compiler/mlir/tensorflow/transforms/host_runtime:lower_cluster_to_runtime_ops", "//tensorflow/compiler/mlir/tf2xla:mlir_bridge_rollout_policy", "//tensorflow/compiler/mlir/tf2xla/api/v1:cluster_tf", + "//tensorflow/compiler/mlir/tf2xla/api/v1:tf_dialect_to_executor", "//tensorflow/compiler/mlir/tf2xla/api/v2:cluster_tf", "//tensorflow/compiler/mlir/tf2xla/api/v2:tf_dialect_to_executor", "//tensorflow/core:core_cpu", @@ -1171,6 +1173,7 @@ cc_library( "@com_google_absl//absl/status", "@llvm-project//llvm:Support", "@llvm-project//mlir:FuncDialect", + "@llvm-project//mlir:IR", ], alwayslink = 1, ) diff --git a/tensorflow/compiler/tf2xla/kernels/BUILD b/tensorflow/compiler/tf2xla/kernels/BUILD index e8d7cd052b5fba..e1605b5c3d3bdf 100644 --- a/tensorflow/compiler/tf2xla/kernels/BUILD +++ b/tensorflow/compiler/tf2xla/kernels/BUILD @@ -369,6 +369,7 @@ tf_kernel_library( "//tensorflow/core:framework", "//tensorflow/core:protos_all_cc", "//tensorflow/core/tpu:tpu_defs", + "@com_google_absl//absl/log", "@com_google_absl//absl/status", "@com_google_absl//absl/strings", "@llvm-project//llvm:Support", diff --git a/tensorflow/compiler/tf2xla/kernels/xla_call_module_loader.cc b/tensorflow/compiler/tf2xla/kernels/xla_call_module_loader.cc index d8e5fcefdda67b..f602e66fd9fa55 100644 --- a/tensorflow/compiler/tf2xla/kernels/xla_call_module_loader.cc +++ b/tensorflow/compiler/tf2xla/kernels/xla_call_module_loader.cc @@ -80,11 +80,11 @@ constexpr int kVersionStartSupportCallTFGraph = 5; constexpr int kVersionStartSupportDisabledChecks = 6; constexpr int kVersionStartSupportShapeAssertions = 7; constexpr int kVersionStartSupportUsesShapePolymorphismAttr = 8; +constexpr int kVersionStartSupportEffects = 9; constexpr int kVersionMinimumSupported = kVersionStartStableHloCompatibility; // This should match xla.py:call_module_maximum_supported_version -constexpr int kVersionMaximumSupported = - kVersionStartSupportUsesShapePolymorphismAttr; +constexpr int kVersionMaximumSupported = kVersionStartSupportEffects; constexpr llvm::StringRef kDisabledCheckPlatform = "platform"; @@ -141,6 +141,11 @@ tsl::Status SetPlatformIndex(mlir::func::FuncOp main, int platform_index) { } // namespace +bool IsTokenType(mlir::Type type) { + return type.isa() || + type.isa(); +} + tsl::StatusOr> XlaCallModuleLoader::Create( mlir::MLIRContext *context, int version, std::string module_str, std::vector disabled_checks, @@ -236,7 +241,8 @@ tsl::Status XlaCallModuleLoader::RefineDynamicShapes( // Refine 'main' argument types to use static input types instead. The main // arguments may occur as return values, or as inputs to called functions, // and changing their types may invalidate the module. To prevent this - // we insert dummy conversion ops as the sole uses of the main arguments. + // we insert dummy conversion ops as the sole uses of the main arguments, for + // the arguments that are not tokens and have dynamic shape. // If we use stablehlo.convert, we end up with "convert 3xf32 -> *xf32" // after we set the static shapes for the main arguments. The "convert" // op does not support unranked result for ranked inputs. So, we use @@ -246,9 +252,16 @@ tsl::Status XlaCallModuleLoader::RefineDynamicShapes( op_builder.setInsertionPointToStart(&main_body); for (auto i = 0; i < main_body.getNumArguments(); ++i) { mlir::BlockArgument arg = main_body.getArgument(i); - auto convert_op = op_builder.create( - arg.getLoc(), arg.getType(), arg); - arg.replaceAllUsesExcept(convert_op, convert_op); + mlir::Type arg_type = arg.getType(); + if (IsTokenType(arg_type)) { + continue; + } + auto ranked_arg_type = arg_type.dyn_cast(); + if (!ranked_arg_type || !ranked_arg_type.hasStaticShape()) { + auto convert_op = op_builder.create( + arg.getLoc(), arg_type, arg); + arg.replaceAllUsesExcept(convert_op, convert_op); + } } auto static_array_output_types = llvm::to_vector(main_.getResultTypes()); @@ -376,7 +389,19 @@ tsl::Status XlaCallModuleLoader::LoadAndPreprocessModule( } mlir::Block &main_body = main_.front(); - int nr_token_arguments = main_has_token_input_output ? 1 : 0; + + int nr_token_arguments = llvm::count_if(InputTypes(), IsTokenType); + if (version < kVersionStartSupportEffects) { + bool has_token_at_start = (nr_token_arguments == 1 && + IsTokenType(main_.getArgument(0).getType())); + if (main_has_token_input_output != has_token_at_start) { + return absl::InvalidArgumentError(absl::StrCat( + "Expected a token at start iff main_has_token_input_output. ", + "Found main function type ", + mlir::debugString(main_.getFunctionType()), + " and main_has_token_input_output = ", main_has_token_input_output)); + } + } int nr_platform_args = (platform_index_ >= 0 ? 1 : 0); if (num_invocation_args != main_body.getNumArguments() - nr_token_arguments) { return absl::InvalidArgumentError(absl::StrCat( diff --git a/tensorflow/compiler/tf2xla/kernels/xla_call_module_loader.h b/tensorflow/compiler/tf2xla/kernels/xla_call_module_loader.h index 33de164d0e453f..bb2e73e331a5a9 100644 --- a/tensorflow/compiler/tf2xla/kernels/xla_call_module_loader.h +++ b/tensorflow/compiler/tf2xla/kernels/xla_call_module_loader.h @@ -26,11 +26,15 @@ limitations under the License. #include "mlir/IR/MLIRContext.h" // from @llvm-project #include "mlir/IR/OwningOpRef.h" // from @llvm-project #include "mlir/IR/TypeRange.h" // from @llvm-project +#include "stablehlo/dialect/StablehloOps.h" // from @stablehlo #include "xla/client/xla_computation.h" +#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h" #include "tsl/platform/statusor.h" namespace tensorflow { +bool IsTokenType(mlir::Type type); + class XlaCallModuleLoader { public: static tsl::StatusOr> Create( @@ -39,8 +43,11 @@ class XlaCallModuleLoader { std::vector platforms, std::string loading_platform, int num_invocation_args, bool main_has_token_input_output); - int nr_outputs() { return main_.getNumResults(); } - mlir::TypeRange output_types() { return main_.getResultTypes(); } + int NrInputs() { return main_.getNumArguments(); } + mlir::TypeRange InputTypes() { return main_.getArgumentTypes(); } + + int NrOutputs() { return main_.getNumResults(); } + mlir::TypeRange OutputTypes() { return main_.getResultTypes(); } // Refines the dynamic module arguments based on the static argument shapes. // This assumes that the module has a "main" function without dimension args, diff --git a/tensorflow/compiler/tf2xla/kernels/xla_call_module_op.cc b/tensorflow/compiler/tf2xla/kernels/xla_call_module_op.cc index 200b24951ed68b..049142034f90b6 100644 --- a/tensorflow/compiler/tf2xla/kernels/xla_call_module_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/xla_call_module_op.cc @@ -19,10 +19,12 @@ limitations under the License. #include #include +#include "absl/log/log.h" #include "absl/status/status.h" #include "absl/strings/str_cat.h" #include "absl/strings/str_join.h" #include "llvm/ADT/ArrayRef.h" +#include "llvm/ADT/STLExtras.h" #include "llvm/ADT/SmallVector.h" #include "mlir/Dialect/Func/Extensions/AllExtensions.h" // from @llvm-project #include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project @@ -150,7 +152,36 @@ class XlaCallModuleOp : public XlaOpKernel { OP_REQUIRES_OK(ctx, ctx->GetAttr("disabled_checks", &disabled_checks)); std::vector platforms; OP_REQUIRES_OK(ctx, ctx->GetAttr("platforms", &platforms)); + // TODO(necula): change this to OP_REQUIRES_OK when 6 months have passed + // since we added the function_list and has_token_input_output + // attributes (May 25, 2023). + bool main_has_token_input_output = false; + if (!ctx->GetAttr("has_token_input_output", &main_has_token_input_output) + .ok()) { + // Whether the StableHLO module's main function has token input/output as + // the first argument and the first result. + // This is used only prior to version 9; afterwards, we just look for + // tokens among the types of the arguments and results, and we support + // multiple tokens, not necessarily at the start. + main_has_token_input_output = false; + } + if (!ctx->GetAttr("function_list", &function_list_).ok()) { + function_list_.clear(); + } + if (VLOG_IS_ON(3)) { + VLOG(3) << "Initializing XlaCallModuleOp (version = " << version + << ", platforms = [" << absl::StrJoin(platforms, ", ") + << "], has_token_input_output = " << main_has_token_input_output + << ", disabled_checks = [" << absl::StrJoin(disabled_checks, ", ") + << "], " + << "function_list = [" + << absl::StrJoin(function_list_, ",", + [](std::string *out, NameAttrList x) { + absl::StrAppend(out, x.name()); + }) + << "])"; + } string loading_device_type = ctx->device_type().type_string(); string loading_platform = ""; if (loading_device_type == DEVICE_CPU_XLA_JIT) { @@ -171,31 +202,21 @@ class XlaCallModuleOp : public XlaOpKernel { absl::UnimplementedError(absl::StrCat( "Unexpected device type ", loading_device_type))); } - VLOG(3) << "Initialized XlaCallModuleOp on " << loading_platform; - if (!ctx->GetAttr("has_token_input_output", &module_has_token_input_output_) - .ok()) { - module_has_token_input_output_ = false; - } + VLOG(3) << "Initializing XlaCallModuleOp on " << loading_platform; { auto loader = XlaCallModuleLoader::Create( &context_, version, std::move(module_str), std::move(disabled_checks), std::move(platforms), loading_platform, /*num_invocation_args=*/ctx->num_inputs(), - module_has_token_input_output_); + main_has_token_input_output); OP_REQUIRES_OK(ctx, loader.status()); loader_ = *std::move(loader); } OP_REQUIRES_OK(ctx, loader_->ValidateDialect()); - if (!ctx->GetAttr("function_list", &function_list_).ok()) { - function_list_.clear(); - } - - if (!ctx->GetAttr(kXlaTokenInputNodesAttrName, &token_input_nodes_).ok()) { - token_input_nodes_.clear(); - op_has_token_input_output_ = false; - } else { - op_has_token_input_output_ = !token_input_nodes_.empty(); + if (!ctx->GetAttr(kXlaTokenInputNodesAttrName, &op_token_input_nodes_) + .ok()) { + op_token_input_nodes_.clear(); } if (!ctx->GetAttr(kXlaOriginalOutsideCompilationNodeName, &original_node_name_) @@ -213,13 +234,15 @@ class XlaCallModuleOp : public XlaOpKernel { xla::XlaBuilder *const b = ctx->builder(); std::vector input_shapes; - if (module_has_token_input_output_) { - input_shapes.push_back(xla::ShapeUtil::MakeTokenShape()); - } - for (int i = 0; i < ctx->num_inputs(); ++i) { - auto shape = ctx->InputXlaShape(i); - OP_REQUIRES_OK(ctx, shape.status()); - input_shapes.push_back(*std::move(shape)); + int next_actual_input = 0; + for (mlir::Type inputType : loader_->InputTypes()) { + if (IsTokenType(inputType)) { + input_shapes.push_back(xla::ShapeUtil::MakeTokenShape()); + } else { + auto shape = ctx->InputXlaShape(next_actual_input++); + OP_REQUIRES_OK(ctx, shape.status()); + input_shapes.push_back(*std::move(shape)); + } } OP_REQUIRES_OK(ctx, loader_->RefineDynamicShapes(input_shapes)); OP_REQUIRES_OK(ctx, loader_->ValidateStaticShapes()); @@ -228,27 +251,30 @@ class XlaCallModuleOp : public XlaOpKernel { OP_REQUIRES_OK(ctx, LowerTfFunctionCalls(ctx)); } + xla::XlaOp token_input; + if (!op_token_input_nodes_.empty()) { + std::vector token_inputs; + for (const string &node_name : op_token_input_nodes_) { + auto token = compiler->GetNodeToken(node_name); + OP_REQUIRES_OK(ctx, token.status()); + token_inputs.push_back(token.value()); + } + token_input = xla::AfterAll(b, token_inputs); + } + std::vector inputs; - if (module_has_token_input_output_) { - // The main function expects a token input at the start. - if (!token_input_nodes_.empty()) { - std::vector token_inputs; - for (const string &node_name : token_input_nodes_) { - auto token = compiler->GetNodeToken(node_name); - OP_REQUIRES_OK(ctx, token.status()); - token_inputs.push_back(token.value()); + next_actual_input = 0; + for (mlir::Type inputType : loader_->InputTypes()) { + if (IsTokenType(inputType)) { + if (token_input.IsUninitialized()) { + // Generate a dummy token if the XlaCallModule doesn't take one. + token_input = xla::CreateToken(b); } - inputs.push_back(xla::AfterAll(b, token_inputs)); + inputs.push_back(token_input); } else { - // Generate a dummy token if the main function expects a token but the - // XlaCallModule doesn't take one. - inputs.push_back(xla::CreateToken(b)); + inputs.push_back(ctx->Input(next_actual_input++)); } } - for (int i = 0, end = ctx->num_inputs(); i < end; ++i) { - inputs.push_back(ctx->Input(i)); - } - auto xla_computation = loader_->ToXlaComputation(); OP_REQUIRES_OK(ctx, xla_computation.status()); @@ -266,47 +292,58 @@ class XlaCallModuleOp : public XlaOpKernel { hlo_module->ToString(options))); } - xla::XlaOp output = xla::Call(b, *xla_computation, inputs); + xla::XlaOp computation_output = xla::Call(b, *xla_computation, inputs); // Check that the resulting computation returns the expected shape - OP_REQUIRES_VALUE(xla::Shape found_output_shape, ctx, b->GetShape(output)); + OP_REQUIRES_VALUE(xla::Shape found_output_shape, ctx, + b->GetShape(computation_output)); VLOG(3) << "XlaCallModule compiled output shape : " << xla::ShapeUtil::HumanString(found_output_shape); - - std::vector outputs; - if (loader_->nr_outputs() == 1) { - outputs.push_back(output); + std::vector computation_outputs; + if (loader_->NrOutputs() == 1) { + computation_outputs.push_back(computation_output); } else { - for (int i = 0; i < loader_->nr_outputs(); ++i) { - outputs.push_back(xla::GetTupleElement(output, i)); + for (int i = 0; i < loader_->NrOutputs(); ++i) { + computation_outputs.push_back( + xla::GetTupleElement(computation_output, i)); } } - xla::XlaOp token_output; - if (module_has_token_input_output_) { - // The main function returns a token as the first output. - token_output = outputs.front(); - outputs.erase(outputs.begin()); - auto shape = b->GetShape(token_output); + // Collect the token outputs and set the non-token outputs + std::vector token_outputs; + int next_actual_output = 0; + for (auto it : llvm::enumerate(loader_->OutputTypes())) { + int i = it.index(); + mlir::Type output_type = it.value(); + auto shape = b->GetShape(computation_outputs[i]); OP_REQUIRES_OK(ctx, shape.status()); - OP_REQUIRES(ctx, shape->IsToken(), - absl::FailedPreconditionError( - absl::StrCat("Token output is not token type: ", - xla::ShapeUtil::HumanString(*shape)))); + if (IsTokenType(output_type)) { + OP_REQUIRES(ctx, shape->IsToken(), + absl::FailedPreconditionError(absl::StrCat( + "Token output at index ", i, " is not token type: ", + xla::ShapeUtil::HumanString(*shape)))); + token_outputs.push_back(computation_outputs[i]); + } else { + OP_REQUIRES(ctx, !shape->IsToken(), + absl::FailedPreconditionError(absl::StrCat( + "Non-token output at index ", i, " is a token type: ", + xla::ShapeUtil::HumanString(*shape)))); + ctx->SetOutput(next_actual_output++, computation_outputs[i]); + } } - if (op_has_token_input_output_) { - if (token_output.IsUninitialized()) { - // The main function does not return any token, but the XlaCallModule is - // expected to return one. Create a dummy token. - token_output = xla::CreateToken(b); + + if (!op_token_input_nodes_.empty()) { + xla::XlaOp token_output = token_input; + if (!token_outputs.empty()) { + token_output = xla::AfterAll(b, token_outputs); + } else { + if (token_output.IsUninitialized()) { + token_output = xla::CreateToken(b); + } } OP_REQUIRES_OK(ctx, compiler->SetNodeToken(original_node_name_, token_output)); } - - for (int i = 0; i < outputs.size(); ++i) { - ctx->SetOutput(i, outputs[i]); - } } private: @@ -404,7 +441,7 @@ class XlaCallModuleOp : public XlaOpKernel { options.always_return_tuple = true; options.is_entry_computation = false; // Propagate tokens from XlaCallModule to inner computation. - options.add_token_input_output = op_has_token_input_output_; + options.add_token_input_output = !op_token_input_nodes_.empty(); XlaCompiler::CompilationResult result; TF_RETURN_IF_ERROR( @@ -518,11 +555,8 @@ class XlaCallModuleOp : public XlaOpKernel { std::unique_ptr loader_; std::vector function_list_; - // Whether the StableHLO module's main function has token input/output. - bool module_has_token_input_output_; // Whether the XlaCallModule op has token input/output. - bool op_has_token_input_output_; - std::vector token_input_nodes_; + std::vector op_token_input_nodes_; std::string original_node_name_; }; diff --git a/tensorflow/compiler/tf2xla/mlir_bridge_pass.cc b/tensorflow/compiler/tf2xla/mlir_bridge_pass.cc index 6df7dded2383e5..b69aa2be5e2e0c 100644 --- a/tensorflow/compiler/tf2xla/mlir_bridge_pass.cc +++ b/tensorflow/compiler/tf2xla/mlir_bridge_pass.cc @@ -22,18 +22,23 @@ limitations under the License. #include "absl/log/log.h" #include "absl/status/status.h" #include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project +#include "mlir/IR/BuiltinOps.h" // from @llvm-project +#include "mlir/IR/Visitors.h" // from @llvm-project +#include "tensorflow/compiler/mlir/mlir_graph_optimization_pass.h" #include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h" #include "tensorflow/compiler/mlir/tensorflow/ir/tf_structs.h" -#include "tensorflow/compiler/mlir/tensorflow/transforms/bridge.h" #include "tensorflow/compiler/mlir/tensorflow/transforms/host_runtime/lower_cluster_to_runtime_ops.h" #include "tensorflow/compiler/mlir/tensorflow/utils/device_util.h" #include "tensorflow/compiler/mlir/tf2xla/api/v1/cluster_tf.h" +#include "tensorflow/compiler/mlir/tf2xla/api/v1/tf_dialect_to_executor.h" #include "tensorflow/compiler/mlir/tf2xla/api/v2/cluster_tf.h" #include "tensorflow/compiler/mlir/tf2xla/api/v2/tf_dialect_to_executor.h" #include "tensorflow/compiler/tf2xla/tf2xla_defs.h" +#include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "tensorflow/core/common_runtime/device_set.h" #include "tensorflow/core/framework/metrics.h" #include "tensorflow/core/lib/monitoring/gauge.h" +#include "tensorflow/core/platform/status.h" #include "tensorflow/core/public/session_options.h" #include "tensorflow/core/tpu/tpu_defs.h" #include "tensorflow/core/util/device_name_utils.h" @@ -51,6 +56,8 @@ auto* mlir_bridge_gauge_v2 = monitoring::Gauge::New( namespace { +using ::mlir::ModuleOp; + bool HasTPUDevice(mlir::ModuleOp module) { mlir::TF::RuntimeDevices devices; if (failed(GetDevicesFromOp(module.getOperation(), &devices))) return false; @@ -132,6 +139,33 @@ bool HasTPUPartitionedCallOpInModule(mlir::ModuleOp module) { return has_tpu_partitioned_call; } +// V1 Compat Bridge extracts out a program into a submodule and runs clustering +// only on the submodule. +absl::Status RunLowerToRuntimeOpsOnSubmodule(ModuleOp parent_module, + bool is_in_fallback_enabled_mode) { + int num_submodules = 0; + absl::Status runtime_lowering_status; + parent_module.walk([&](ModuleOp submodule) { + if (submodule == parent_module) return mlir::WalkResult::advance(); + num_submodules++; + runtime_lowering_status = + tensorflow::tfrt_compiler::RunLowerClusterToRuntimeOpsPassPipeline( + submodule, tsl::DeviceType(DEVICE_TPU_XLA_JIT)); + if (num_submodules > 1) { + return mlir::WalkResult::interrupt(); + } + + return mlir::WalkResult::advance(); + }); + + if (num_submodules > 1) { + return absl::InternalError( + "Lower to runtime has more than one submodule. Erroring out."); + } + + return runtime_lowering_status; +} + } // namespace // Analyzes the user requested policy as well as the contents of the graph and @@ -270,8 +304,8 @@ Status MlirBridgePass::Run(const std::string& function_name, return OkStatus(); } + bool fallback_enabled = false; if (run_tpu_bridge) { - bool fallback_enabled = false; if (pass_state == MlirOptimizationPassState::FallbackEnabled) { // We set `uses_uninitialized_resource_args` to false here because the // first phase of the bridge is not affected by uninitialized resource @@ -293,14 +327,18 @@ Status MlirBridgePass::Run(const std::string& function_name, TF_RETURN_IF_ERROR( tensorflow::tfrt_compiler::RunLowerClusterToRuntimeOpsPassPipeline( module, tsl::DeviceType(DEVICE_TPU_XLA_JIT))); - + } else { + VLOG(1) << "Running GPU/CPU Bridge"; TF_RETURN_IF_ERROR( - tensorflow::tf2xla::v2::ExportFromTensorflowDialectToExecutor(module)); + tensorflow::tf2xla::v2::RunFunctionTf2xlaClusteringBridge( + module, tf2xla::v2::DeviceType::XLA_GPU_JIT, fallback_enabled)); - return absl::OkStatus(); + TF_RETURN_IF_ERROR( + tensorflow::tfrt_compiler::RunLowerClusterToRuntimeOpsPassPipeline( + module, tsl::DeviceType(DEVICE_GPU_XLA_JIT))); } - VLOG(1) << "Running MLIR CPU/GPU Bridge"; - return mlir::TF::RunTFXLABridge(module, function_name); + + return tensorflow::tf2xla::v2::ExportFromTensorflowDialectToExecutor(module); } MlirOptimizationPassState MlirBridgeV1CompatPass::GetPassState( @@ -402,10 +440,19 @@ Status MlirBridgeV1CompatPass::Run(const GraphOptimizationPassOptions& options, } VLOG(1) << "Running MLIR TPU Bridge V1 Compat"; - mlir_bridge_gauge_v1->GetCell()->Set(true); - return tensorflow::tf2xla::v1::RunSessionTf2xlaClusteringBridge( - module, fallback_enabled); + TF_RETURN_IF_ERROR(tensorflow::tf2xla::v1::RunSessionTf2xlaClusteringBridge( + module, fallback_enabled)); + + auto lower_cluster_to_runtime_ops_pass_pipeline = + RunLowerToRuntimeOpsOnSubmodule(module, fallback_enabled); + if (!lower_cluster_to_runtime_ops_pass_pipeline.ok()) { + VLOG(1) << "Error while lowering cluster to runtime ops: " + << lower_cluster_to_runtime_ops_pass_pipeline; + return lower_cluster_to_runtime_ops_pass_pipeline; + } + + return tensorflow::tf2xla::v1::ExportFromTensorflowDialectToExecutor(module); } } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/ops/xla_ops.cc b/tensorflow/compiler/tf2xla/ops/xla_ops.cc index 4aa2811585c8c8..480dc474410359 100644 --- a/tensorflow/compiler/tf2xla/ops/xla_ops.cc +++ b/tensorflow/compiler/tf2xla/ops/xla_ops.cc @@ -1392,7 +1392,10 @@ function_list: This list contains the TensorFlow FunctionDefs that are used by has_token_input_output: If true, the embedded StableHLO module's main function must take a `!stablehlo.token` as its first argument and returns a token as its first result. This can be used in conjunction with the TF2XLA's side - effect mechanism in order to model side effects. + effect mechanism in order to model side effects. This is used only in versions + prior to version 9. After that, the number and position of tokens among + the arguments and results are obtained from the main function type. This + allows us to support more than one token and not necessarily at the start. disabled_checks: A list of strings describing the safety checks that were disabled at serialization time. This attribute was added in version 6. For more details see diff --git a/tensorflow/compiler/tf2xla/python/xla.py b/tensorflow/compiler/tf2xla/python/xla.py index 80eecfb0d2f7e2..27940b7fb92c17 100644 --- a/tensorflow/compiler/tf2xla/python/xla.py +++ b/tensorflow/compiler/tf2xla/python/xla.py @@ -669,7 +669,7 @@ def call_module_maximum_supported_version(): See versioning details documentation for the XlaCallModule op at: https://github.com/search?q=repo%3Atensorflow%2Ftensorflow+path%3Axla_call_module+%22int+VERSION_MAXIMUM_SUPPORTED%22&type=code """ - return 8 + return 9 # pylint: enable=g-doc-args # pylint: enable=g-doc-return-or-yield diff --git a/tensorflow/core/api_def/BUILD b/tensorflow/core/api_def/BUILD index 4884972676b864..b3efda8f4791b7 100644 --- a/tensorflow/core/api_def/BUILD +++ b/tensorflow/core/api_def/BUILD @@ -6,12 +6,12 @@ # :python_api_def # :java_api_def -load("//tensorflow/core/platform:rules_cc.bzl", "cc_library") load( "//tensorflow:tensorflow.bzl", "tf_cc_binary", "tf_cc_test", ) +load("//tensorflow/core/platform:rules_cc.bzl", "cc_library") load( "//third_party/mkl:build_defs.bzl", "if_mkl", @@ -116,5 +116,8 @@ tf_cc_test( "//tensorflow/core:test", "//tensorflow/core:test_main", "//tensorflow/core/platform:resource_loader", + "//tensorflow/core/tpu/ops:sparse_core_ops", + "//tensorflow/core/tpu/ops:sparse_core_preprocess_ops", + "//tensorflow/core/tpu/ops:tpu_copy_with_dynamic_shape_op", ], ) diff --git a/tensorflow/core/api_def/base_api/api_def_ConvertToCooTensor.pbtxt b/tensorflow/core/api_def/base_api/api_def_ConvertToCooTensor.pbtxt new file mode 100644 index 00000000000000..c9d629514f49f7 --- /dev/null +++ b/tensorflow/core/api_def/base_api/api_def_ConvertToCooTensor.pbtxt @@ -0,0 +1,4 @@ +op { + graph_op_name: "ConvertToCooTensor" + visibility: HIDDEN +} diff --git a/tensorflow/core/api_def/base_api/api_def_GetMinibatchSplitsWithPhysicalReplica.pbtxt b/tensorflow/core/api_def/base_api/api_def_GetMinibatchSplitsWithPhysicalReplica.pbtxt new file mode 100644 index 00000000000000..e402d2bf92de67 --- /dev/null +++ b/tensorflow/core/api_def/base_api/api_def_GetMinibatchSplitsWithPhysicalReplica.pbtxt @@ -0,0 +1,4 @@ +op { + graph_op_name: "GetMinibatchSplitsWithPhysicalReplica" + visibility: HIDDEN +} diff --git a/tensorflow/core/api_def/base_api/api_def_GetMinibatchesInCsrWithPhysicalReplica.pbtxt b/tensorflow/core/api_def/base_api/api_def_GetMinibatchesInCsrWithPhysicalReplica.pbtxt new file mode 100644 index 00000000000000..49493ee2254354 --- /dev/null +++ b/tensorflow/core/api_def/base_api/api_def_GetMinibatchesInCsrWithPhysicalReplica.pbtxt @@ -0,0 +1,4 @@ +op { + graph_op_name: "GetMinibatchesInCsrWithPhysicalReplica" + visibility: HIDDEN +} diff --git a/tensorflow/core/api_def/base_api/api_def_StoreMinibatchStatisticsInFdo.pbtxt b/tensorflow/core/api_def/base_api/api_def_StoreMinibatchStatisticsInFdo.pbtxt new file mode 100644 index 00000000000000..9545ffdf9c4b7e --- /dev/null +++ b/tensorflow/core/api_def/base_api/api_def_StoreMinibatchStatisticsInFdo.pbtxt @@ -0,0 +1,4 @@ +op { + graph_op_name: "StoreMinibatchStatisticsInFdo" + visibility: HIDDEN +} diff --git a/tensorflow/core/api_def/base_api/api_def_TPUAnnotateTensorsWithDynamicShape.pbtxt b/tensorflow/core/api_def/base_api/api_def_TPUAnnotateTensorsWithDynamicShape.pbtxt new file mode 100644 index 00000000000000..84ac58c2e7e5ea --- /dev/null +++ b/tensorflow/core/api_def/base_api/api_def_TPUAnnotateTensorsWithDynamicShape.pbtxt @@ -0,0 +1,4 @@ +op { + graph_op_name: "TPUAnnotateTensorsWithDynamicShape" + visibility: HIDDEN +} diff --git a/tensorflow/core/api_def/base_api/api_def_TPUCopyWithDynamicShape.pbtxt b/tensorflow/core/api_def/base_api/api_def_TPUCopyWithDynamicShape.pbtxt new file mode 100644 index 00000000000000..423e1a22244a4b --- /dev/null +++ b/tensorflow/core/api_def/base_api/api_def_TPUCopyWithDynamicShape.pbtxt @@ -0,0 +1,8 @@ +op { + graph_op_name: "TPUCopyWithDynamicShape" + visibility: HIDDEN + summary: <