diff --git a/build_tools/bazel_to_cmake/bazel_to_cmake_targets.py b/build_tools/bazel_to_cmake/bazel_to_cmake_targets.py index a5333f9eadbb..e192a3d74183 100644 --- a/build_tools/bazel_to_cmake/bazel_to_cmake_targets.py +++ b/build_tools/bazel_to_cmake/bazel_to_cmake_targets.py @@ -42,6 +42,7 @@ "@llvm-project//mlir:ShapeTransforms": ["MLIRShapeOpsTransforms"], "@llvm-project//mlir:SideEffects": ["MLIRSideEffectInterfaces"], "@llvm-project//mlir:SPIRVDialect": ["MLIRSPIRV"], + "@llvm-project//mlir:TosaDialect": ["MLIRTosa"], "@llvm-project//mlir:mlir_c_runner_utils": ["MLIRExecutionEngine"], "@llvm-project//mlir:mlir-translate": ["mlir-translate"], "@llvm-project//mlir:MlirTableGenMain": ["MLIRTableGen"], diff --git a/iree/compiler/Dialect/Flow/Transforms/BUILD b/iree/compiler/Dialect/Flow/Transforms/BUILD index 65e408d53e00..c32aba1e2f64 100644 --- a/iree/compiler/Dialect/Flow/Transforms/BUILD +++ b/iree/compiler/Dialect/Flow/Transforms/BUILD @@ -79,6 +79,7 @@ cc_library( "@llvm-project//mlir:StandardOps", "@llvm-project//mlir:Support", "@llvm-project//mlir:TensorDialect", + "@llvm-project//mlir:TosaToLinalg", "@llvm-project//mlir:TransformUtils", "@llvm-project//mlir:Transforms", "@mlir-hlo//:chlo_legalize_to_hlo", diff --git a/iree/compiler/Dialect/Flow/Transforms/CMakeLists.txt b/iree/compiler/Dialect/Flow/Transforms/CMakeLists.txt index 55a482546f7d..4c2f4acad222 100644 --- a/iree/compiler/Dialect/Flow/Transforms/CMakeLists.txt +++ b/iree/compiler/Dialect/Flow/Transforms/CMakeLists.txt @@ -58,6 +58,7 @@ iree_cc_library( MLIRStandard MLIRSupport MLIRTensor + MLIRTosaToLinalg MLIRTransformUtils MLIRTransforms iree::base::signature_mangle diff --git a/iree/compiler/Dialect/Flow/Transforms/Passes.cpp b/iree/compiler/Dialect/Flow/Transforms/Passes.cpp index 231e1f76cf4e..5669b70eeea8 100644 --- a/iree/compiler/Dialect/Flow/Transforms/Passes.cpp +++ b/iree/compiler/Dialect/Flow/Transforms/Passes.cpp @@ -20,6 +20,7 @@ #include "iree/compiler/Dialect/Shape/Conversion/Passes.h" #include "iree/compiler/Dialect/Shape/Transforms/Passes.h" #include "mlir-hlo/Dialect/mhlo/transforms/passes.h" +#include "mlir/Conversion/TosaToLinalg/TosaToLinalg.h" #include "mlir/Dialect/Shape/Transforms/Passes.h" #include "mlir/Pass/PassOptions.h" #include "mlir/Pass/PassRegistry.h" @@ -68,6 +69,9 @@ void buildFlowTransformPassPipeline(OpPassManager &passManager) { passManager.addNestedPass(mhlo::createLegalizeControlFlowPass()); passManager.addNestedPass(createHLOPreprocessingPass()); + // Convert TOSA ops to Linalg-on-tensor ops. + passManager.addNestedPass(tosa::createTosaToLinalgOnTensors()); + // Run passes to remove shape constraints. HLO lowering inserts them, but they // are not desired here. // diff --git a/iree/compiler/Dialect/Flow/Utils/BUILD b/iree/compiler/Dialect/Flow/Utils/BUILD index 5f75f7f9f2ec..a15ec70ce4ad 100644 --- a/iree/compiler/Dialect/Flow/Utils/BUILD +++ b/iree/compiler/Dialect/Flow/Utils/BUILD @@ -36,6 +36,7 @@ cc_library( "@llvm-project//mlir:LinalgOps", "@llvm-project//mlir:StandardOps", "@llvm-project//mlir:Support", + "@llvm-project//mlir:TosaDialect", "@mlir-hlo//:hlo", ], ) diff --git a/iree/compiler/Dialect/Flow/Utils/CMakeLists.txt b/iree/compiler/Dialect/Flow/Utils/CMakeLists.txt index 76da04c1d92c..0ccb47e7a46e 100644 --- a/iree/compiler/Dialect/Flow/Utils/CMakeLists.txt +++ b/iree/compiler/Dialect/Flow/Utils/CMakeLists.txt @@ -29,6 +29,7 @@ iree_cc_library( MLIRLinalg MLIRStandard MLIRSupport + MLIRTosa iree::compiler::Dialect::Flow::IR iree::compiler::Dialect::Shape::IR tensorflow::mlir_hlo diff --git a/iree/compiler/Dialect/Flow/Utils/DispatchUtils.cpp b/iree/compiler/Dialect/Flow/Utils/DispatchUtils.cpp index 7fc8240757be..a90d9a512a6f 100644 --- a/iree/compiler/Dialect/Flow/Utils/DispatchUtils.cpp +++ b/iree/compiler/Dialect/Flow/Utils/DispatchUtils.cpp @@ -20,6 +20,7 @@ #include "mlir-hlo/Dialect/mhlo/IR/hlo_ops.h" #include "mlir/Dialect/Linalg/IR/LinalgTypes.h" #include "mlir/Dialect/StandardOps/IR/Ops.h" +#include "mlir/Dialect/Tosa/IR/TosaOps.h" #include "mlir/IR/BlockAndValueMapping.h" #include "mlir/IR/Builders.h" @@ -37,7 +38,8 @@ bool isOpOfKnownDialect(Operation *op) { dialectNamespace == linalg::LinalgDialect::getDialectNamespace() || dialectNamespace == mhlo::MhloDialect::getDialectNamespace() || dialectNamespace == mlir::StandardOpsDialect::getDialectNamespace() || - dialectNamespace == ShapeDialect::getDialectNamespace(); + dialectNamespace == ShapeDialect::getDialectNamespace() || + dialectNamespace == tosa::TosaDialect::getDialectNamespace(); } namespace { diff --git a/iree/test/e2e/tosa_ops/BUILD b/iree/test/e2e/tosa_ops/BUILD new file mode 100644 index 000000000000..5eecc23ad5c0 --- /dev/null +++ b/iree/test/e2e/tosa_ops/BUILD @@ -0,0 +1,49 @@ +# Copyright 2019 Google LLC +# +# 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 +# +# https://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. + +# Tests of end-to-end IREE support for individual ops in the TOSA dialect. +# Each test file should have a name matching the corresponding TOSA op and test only the +# functionality of that op (though may make use of other ops where necessary). Tests should be +# written using the IREE Check framework. +# See https://google.github.io/iree/developing-iree/testing-guide#iree-core-end-to-end-tests. + +load("//build_tools/bazel:iree_check_test.bzl", "iree_check_single_backend_test_suite") + +package( + default_visibility = ["//visibility:public"], + features = ["layering_check"], + licenses = ["notice"], # Apache 2.0 +) + +iree_check_single_backend_test_suite( + name = "check_vulkan-spirv_vulkan", + srcs = glob(["*.mlir"]), + driver = "vulkan", + target_backend = "vulkan-spirv", +) + +iree_check_single_backend_test_suite( + name = "check_dylib-llvm-aot_dylib", + srcs = glob(["*.mlir"]), + driver = "dylib", + target_backend = "dylib-llvm-aot", +) + +test_suite( + name = "check", + tests = [ + ":check_dylib-llvm-aot_dylib", + ":check_vulkan-spirv_vulkan", + ], +) diff --git a/iree/test/e2e/tosa_ops/CMakeLists.txt b/iree/test/e2e/tosa_ops/CMakeLists.txt new file mode 100644 index 000000000000..356c083835a8 --- /dev/null +++ b/iree/test/e2e/tosa_ops/CMakeLists.txt @@ -0,0 +1,39 @@ +# Copyright 2020 Google LLC +# +# 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 +# +# https://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. + +iree_add_all_subdirs() + +file(GLOB _GLOB_X_MLIR LIST_DIRECTORIES false RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} CONFIGURE_DEPENDS *.mlir) +iree_check_single_backend_test_suite( + NAME + check_vulkan-spirv_vulkan + SRCS + "${_GLOB_X_MLIR}" + TARGET_BACKEND + "vulkan-spirv" + DRIVER + "vulkan" +) + +file(GLOB _GLOB_X_MLIR LIST_DIRECTORIES false RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} CONFIGURE_DEPENDS *.mlir) +iree_check_single_backend_test_suite( + NAME + check_dylib-llvm-aot_dylib + SRCS + "${_GLOB_X_MLIR}" + TARGET_BACKEND + "dylib-llvm-aot" + DRIVER + "dylib" +) diff --git a/iree/test/e2e/tosa_ops/abs.mlir b/iree/test/e2e/tosa_ops/abs.mlir new file mode 100644 index 000000000000..db84f2a8fada --- /dev/null +++ b/iree/test/e2e/tosa_ops/abs.mlir @@ -0,0 +1,6 @@ +func @tensor_float() attributes { iree.module.export } { + %0 = iree.unfoldable_constant dense<[-1.0, -0.5, 0.0, 1.0]> : tensor<4xf32> + %result = "tosa.abs"(%0) : (tensor<4xf32>) -> tensor<4xf32> + check.expect_almost_eq_const(%result, dense<[1.0, 0.5, 0.0, 1.0]> : tensor<4xf32>) : tensor<4xf32> + return +} diff --git a/iree/test/e2e/tosa_ops/add.mlir b/iree/test/e2e/tosa_ops/add.mlir new file mode 100644 index 000000000000..520220c99693 --- /dev/null +++ b/iree/test/e2e/tosa_ops/add.mlir @@ -0,0 +1,16 @@ +func @tensor_float() attributes { iree.module.export } { + %0 = iree.unfoldable_constant dense<[1.0, 2.0, 3.0, 4.0]> : tensor<4xf32> + %1 = iree.unfoldable_constant dense<[5.0, 6.0, 7.0, 8.0]> : tensor<4xf32> + %result = "tosa.add"(%0, %1) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32> + check.expect_almost_eq_const(%result, dense<[6.0, 8.0, 10.0, 12.0]> : tensor<4xf32>) : tensor<4xf32> + return +} + +func @tensor_int() attributes { iree.module.export } { + %0 = iree.unfoldable_constant dense<[1, 2, 3, 4]> : tensor<4xi32> + %1 = iree.unfoldable_constant dense<[5, 6, 7, 8]> : tensor<4xi32> + %result = "tosa.add"(%0, %1) : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi32> + check.expect_eq_const(%result, dense<[6, 8, 10, 12]> : tensor<4xi32>) : tensor<4xi32> + return +} + diff --git a/iree/test/e2e/tosa_ops/bitwise_and.mlir b/iree/test/e2e/tosa_ops/bitwise_and.mlir new file mode 100644 index 000000000000..3c5453e41a88 --- /dev/null +++ b/iree/test/e2e/tosa_ops/bitwise_and.mlir @@ -0,0 +1,7 @@ +func @tensor() attributes { iree.module.export } { + %0 = iree.unfoldable_constant dense<[0x0, 0x011, 0x101, 0x111]> : tensor<4xi32> + %1 = iree.unfoldable_constant dense<[0x0, 0x010, 0x111, 0x000]> : tensor<4xi32> + %result = "tosa.bitwise_and"(%0, %1) : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi32> + check.expect_eq_const(%result, dense<[0x0, 0x010, 0x101, 0x000]> : tensor<4xi32>) : tensor<4xi32> + return +} diff --git a/iree/test/e2e/tosa_ops/bitwise_or.mlir b/iree/test/e2e/tosa_ops/bitwise_or.mlir new file mode 100644 index 000000000000..a818f2277b32 --- /dev/null +++ b/iree/test/e2e/tosa_ops/bitwise_or.mlir @@ -0,0 +1,7 @@ +func @tensor() attributes { iree.module.export } { + %0 = iree.unfoldable_constant dense<[0x0, 0x11, 0x1101, 0x111]> : tensor<4xi32> + %1 = iree.unfoldable_constant dense<[0x0, 0x10, 0x0111, 0x111]> : tensor<4xi32> + %result = "tosa.bitwise_or"(%0, %1) : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi32> + check.expect_eq_const(%result, dense<[0x0, 0x11, 0x1111, 0x111]> : tensor<4xi32>) : tensor<4xi32> + return +} diff --git a/iree/test/e2e/tosa_ops/bitwise_xor.mlir b/iree/test/e2e/tosa_ops/bitwise_xor.mlir new file mode 100644 index 000000000000..8e9d796d19df --- /dev/null +++ b/iree/test/e2e/tosa_ops/bitwise_xor.mlir @@ -0,0 +1,7 @@ +func @tensor() attributes { iree.module.export } { + %0 = iree.unfoldable_constant dense<[0x0, 0x11, 0x1101, 0x111]> : tensor<4xi32> + %1 = iree.unfoldable_constant dense<[0x0, 0x10, 0x0111, 0x000]> : tensor<4xi32> + %result = "tosa.bitwise_xor"(%0, %1) : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi32> + check.expect_eq_const(%result, dense<[0x0, 0x01, 0x1010, 0x111]> : tensor<4xi32>) : tensor<4xi32> + return +} diff --git a/iree/test/e2e/tosa_ops/logical_left_shift.mlir b/iree/test/e2e/tosa_ops/logical_left_shift.mlir new file mode 100644 index 000000000000..58a7a8a706e6 --- /dev/null +++ b/iree/test/e2e/tosa_ops/logical_left_shift.mlir @@ -0,0 +1,7 @@ +func @tensor() attributes { iree.module.export } { + %0 = iree.unfoldable_constant dense<[5, 3, 2, 7]> : tensor<4xi32> + %1 = iree.unfoldable_constant dense<[0, 1, 2, 3]> : tensor<4xi32> + %result = "tosa.logical_left_shift"(%0, %1) : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi32> + check.expect_eq_const(%result, dense<[5, 6, 8, 56]> : tensor<4xi32>) : tensor<4xi32> + return +} diff --git a/iree/test/e2e/tosa_ops/logical_right_shift.mlir b/iree/test/e2e/tosa_ops/logical_right_shift.mlir new file mode 100644 index 000000000000..12dd54ee8d13 --- /dev/null +++ b/iree/test/e2e/tosa_ops/logical_right_shift.mlir @@ -0,0 +1,7 @@ +func @tensor() attributes { iree.module.export } { + %0 = iree.unfoldable_constant dense<[5, 8, 9, 256]> : tensor<4xi32> + %1 = iree.unfoldable_constant dense<[0, 1, 2, 8]> : tensor<4xi32> + %result = "tosa.logical_right_shift"(%0, %1) : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi32> + check.expect_eq_const(%result, dense<[5, 4, 2, 1]> : tensor<4xi32>) : tensor<4xi32> + return +} diff --git a/iree/test/e2e/tosa_ops/sub.mlir b/iree/test/e2e/tosa_ops/sub.mlir new file mode 100644 index 000000000000..d2af2f51f073 --- /dev/null +++ b/iree/test/e2e/tosa_ops/sub.mlir @@ -0,0 +1,15 @@ +func @tensor_float() attributes { iree.module.export } { + %0 = iree.unfoldable_constant dense<[1.0, 5.0, 3.0, 4.0]> : tensor<4xf32> + %1 = iree.unfoldable_constant dense<[5.0, 1.0, 3.0, 1.5]> : tensor<4xf32> + %result = "tosa.sub"(%0, %1) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32> + check.expect_almost_eq_const(%result, dense<[-4.0, 4.0, 0.0, 2.5]> : tensor<4xf32>) : tensor<4xf32> + return +} + +func @tensor_int() attributes { iree.module.export } { + %0 = iree.unfoldable_constant dense<[1, 5, 3, 4]> : tensor<4xi32> + %1 = iree.unfoldable_constant dense<[5, 1, 3, 1]> : tensor<4xi32> + %result = "tosa.sub"(%0, %1) : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi32> + check.expect_eq_const(%result, dense<[-4, 4, 0, 3]> : tensor<4xi32>) : tensor<4xi32> + return +} diff --git a/iree/test/e2e/tosa_ops/tanh.mlir b/iree/test/e2e/tosa_ops/tanh.mlir new file mode 100644 index 000000000000..e337e51e6697 --- /dev/null +++ b/iree/test/e2e/tosa_ops/tanh.mlir @@ -0,0 +1,6 @@ +func @tensor_float() attributes { iree.module.export } { + %0 = iree.unfoldable_constant dense<[-1.0, 0.0, 0.5, 1.0]> : tensor<4xf32> + %result = "tosa.tanh"(%0) : (tensor<4xf32>) -> tensor<4xf32> + check.expect_almost_eq_const(%result, dense<[-0.761594, 0.0, 0.462117, 0.761594]> : tensor<4xf32>) : tensor<4xf32> + return +} diff --git a/iree/tools/CMakeLists.txt b/iree/tools/CMakeLists.txt index 8af75544733b..318e8b30d92d 100644 --- a/iree/tools/CMakeLists.txt +++ b/iree/tools/CMakeLists.txt @@ -208,6 +208,8 @@ if(${IREE_BUILD_COMPILER}) MLIRShape MLIRStandard MLIRStandardToSPIRV + MLIRTosa + MLIRTosaTransforms MLIRTransforms MLIRVector PUBLIC diff --git a/iree/tools/init_mlir_dialects.h b/iree/tools/init_mlir_dialects.h index 80dbf059e840..f93952713d56 100644 --- a/iree/tools/init_mlir_dialects.h +++ b/iree/tools/init_mlir_dialects.h @@ -31,6 +31,7 @@ #include "mlir/Dialect/Shape/IR/Shape.h" #include "mlir/Dialect/StandardOps/IR/Ops.h" #include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/Dialect/Tosa/IR/TosaOps.h" #include "mlir/Dialect/Vector/VectorOps.h" #include "mlir/IR/Dialect.h" @@ -49,6 +50,7 @@ inline void registerMlirDialects(DialectRegistry ®istry) { StandardOpsDialect, vector::VectorDialect, tensor::TensorDialect, + tosa::TosaDialect, SDBMDialect, shape::ShapeDialect>(); // clang-format on diff --git a/iree/tools/init_mlir_passes.h b/iree/tools/init_mlir_passes.h index 9d56eeca3100..37e7e78f6182 100644 --- a/iree/tools/init_mlir_passes.h +++ b/iree/tools/init_mlir_passes.h @@ -30,6 +30,7 @@ #include "mlir/Dialect/SCF/Passes.h" #include "mlir/Dialect/SPIRV/Transforms/Passes.h" #include "mlir/Dialect/Shape/Transforms/Passes.h" +#include "mlir/Dialect/Tosa/Transforms/Passes.h" #include "mlir/Transforms/Passes.h" namespace mlir { @@ -80,6 +81,9 @@ inline void registerMlirPasses() { registerConvertStandardToSPIRVPass(); registerLegalizeStandardForSPIRVPass(); registerConvertLinalgToSPIRVPass(); + + // TOSA. + registerTosaToLinalgOnTensorsPass(); } } // namespace mlir