Integrate TOSA with the LLVM and SPIRV backends.
Hooks up TOSA with the LLVM and SPIRV backends. This includes a set of tests for
supported TOSA lowerings.
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 a5333f9..e192a3d 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 65e408d..c32aba1 100644
--- a/iree/compiler/Dialect/Flow/Transforms/BUILD
+++ b/iree/compiler/Dialect/Flow/Transforms/BUILD
@@ -79,6 +79,7 @@
"@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 55a4825..4c2f4ac 100644
--- a/iree/compiler/Dialect/Flow/Transforms/CMakeLists.txt
+++ b/iree/compiler/Dialect/Flow/Transforms/CMakeLists.txt
@@ -58,6 +58,7 @@
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 231e1f7..5669b70 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 @@
passManager.addNestedPass<FuncOp>(mhlo::createLegalizeControlFlowPass());
passManager.addNestedPass<FuncOp>(createHLOPreprocessingPass());
+ // Convert TOSA ops to Linalg-on-tensor ops.
+ passManager.addNestedPass<FuncOp>(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 5f75f7f..a15ec70 100644
--- a/iree/compiler/Dialect/Flow/Utils/BUILD
+++ b/iree/compiler/Dialect/Flow/Utils/BUILD
@@ -36,6 +36,7 @@
"@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 76da04c..0ccb47e 100644
--- a/iree/compiler/Dialect/Flow/Utils/CMakeLists.txt
+++ b/iree/compiler/Dialect/Flow/Utils/CMakeLists.txt
@@ -29,6 +29,7 @@
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 7fc8240..a90d9a5 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 @@
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 0000000..5eecc23
--- /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 0000000..356c083
--- /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 0000000..db84f2a
--- /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 0000000..520220c
--- /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 0000000..3c5453e
--- /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 0000000..a818f22
--- /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 0000000..8e9d796
--- /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 0000000..58a7a8a
--- /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 0000000..12dd54e
--- /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 0000000..d2af2f5
--- /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 0000000..e337e51
--- /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 8af7554..318e8b3 100644
--- a/iree/tools/CMakeLists.txt
+++ b/iree/tools/CMakeLists.txt
@@ -208,6 +208,8 @@
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 80dbf05..f939527 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 @@
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 9d56eec..37e7e78 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 @@
registerConvertStandardToSPIRVPass();
registerLegalizeStandardForSPIRVPass();
registerConvertLinalgToSPIRVPass();
+
+ // TOSA.
+ registerTosaToLinalgOnTensorsPass();
}
} // namespace mlir
