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.

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"],

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",

iree/compiler/Dialect/Flow/Transforms/CMakeLists.txt

@@ -58,6 +58,7 @@ iree_cc_library(
     MLIRStandard
     MLIRSupport
     MLIRTensor
+    MLIRTosaToLinalg
     MLIRTransformUtils
     MLIRTransforms
     iree::base::signature_mangle

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<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.
   //

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",
     ],
 )

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

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 {

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",
+    ],
+)

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"
+)

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
+}

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
+}
+

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
+}

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
+}

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
+}

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
+}

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
+}

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
+}

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
+}

iree/tools/CMakeLists.txt

@@ -208,6 +208,8 @@ if(${IREE_BUILD_COMPILER})
       MLIRShape
       MLIRStandard
       MLIRStandardToSPIRV
+      MLIRTosa
+      MLIRTosaTransforms
       MLIRTransforms
       MLIRVector
     PUBLIC

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 &registry) {
                   StandardOpsDialect,
                   vector::VectorDialect,
                   tensor::TensorDialect,
+                  tosa::TosaDialect,
                   SDBMDialect,
                   shape::ShapeDialect>();
   // clang-format on

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