[MLIR][TOSA] Added Tosa to Standard/SCF Lowerings (const, if, while)

[rsuderman]

Summary

Includes a lowering for tosa.const, tosa.if, and tosa.while to Standard/SCF dialects. TosaToStandard is
used for constant lowerings and TosaToSCF handles the if/while ops.

[rsuderman]

Fixed linter errors.

[silvas]

test cases?

[rsuderman]

Fixed upload

[silvas]

mlir/lib/Conversion/TosaToStandard/TosaToStandardPass.cpp

48

it feels like these could be in TosaToSCF and convert to scf.if/scf.while. Would be easier to write them as patterns too.

[rriddle]

mlir/lib/Conversion/TosaToStandard/TosaToStandard.cpp

77

Remove all of the mlir:: in this commit.

[rriddle]

mlir/lib/Conversion/TosaToStandard/TosaToStandardPass.cpp

48

+1

[rsuderman]

Moved control flow lowerings to a TosaToSCF set of passes.

[rsuderman]

Fixed arc lint error.

[rsuderman]

clang-tidy error

[silvas]

mlir/include/mlir/Conversion/TosaToStandard/TosaToStandard.h

24

decl not needed?

mlir/lib/Conversion/TosaToStandard/TosaToStandard.cpp

3

nit: formatting

mlir/lib/Conversion/TosaToStandard/TosaToStandardPass.cpp

32

specify in td file.

39

you should only need ConstantOp legal.

46

use partial conversion so you don't need markUnknownOpDynamicallyLegal

mlir/test/Conversion/TosaToSCF/tosa-to-scf.mlir

39 ↗

(On Diff #326459)

nit: indent

[rriddle]

mlir/lib/Conversion/TosaToSCF/TosaToSCF.cpp

26 ↗

(On Diff #326459)

Functions should be marked static and at the top-level, only classes should be in anonymous namespaces.

30 ↗

(On Diff #326459)

nit: Spell out auto here.

31 ↗

(On Diff #326459)

Remove trivial braces.

42 ↗

(On Diff #326459)

This seems sketchy, how does this work with rollbacks? Why not rewriter.erase?

66 ↗

(On Diff #326459)

Same comment here.

mlir/lib/Conversion/TosaToSCF/TosaToSCFPass.cpp

29 ↗

(On Diff #326459)

Why does this need to be a FuncOp pass?

mlir/lib/Conversion/TosaToStandard/TosaToStandardPass.cpp

36

Same comment, why is this a FuncOp pass?

[rsuderman]

silvas@ and rriddle@ comments

mlir/include/mlir/Conversion/Passes.h

 #include "mlir/Conversion/SCFToOpenMP/SCFToOpenMP.h"
 #include "mlir/Conversion/SCFToSPIRV/SCFToSPIRVPass.h"
 #include "mlir/Conversion/SCFToStandard/SCFToStandard.h"
 #include "mlir/Conversion/SPIRVToLLVM/SPIRVToLLVMPass.h"
 #include "mlir/Conversion/ShapeToStandard/ShapeToStandard.h"
 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h"
 #include "mlir/Conversion/StandardToSPIRV/StandardToSPIRVPass.h"
 #include "mlir/Conversion/TosaToLinalg/TosaToLinalg.h"
+#include "mlir/Conversion/TosaToStandard/TosaToStandard.h"
 #include "mlir/Conversion/VectorToLLVM/ConvertVectorToLLVM.h"
 #include "mlir/Conversion/VectorToROCDL/VectorToROCDL.h"
 #include "mlir/Conversion/VectorToSCF/VectorToSCF.h"
 #include "mlir/Conversion/VectorToSPIRV/VectorToSPIRVPass.h"
 
 namespace mlir {
 
 /// Generate the code for registering conversion passes.
 #define GEN_PASS_REGISTRATION
 #include "mlir/Conversion/Passes.h.inc"
 
 } // namespace mlir
 
 #endif // MLIR_CONVERSION_PASSES_H

mlir/include/mlir/Conversion/Passes.td

let description = [{
     Pass that converts TOSA operations to the equivalent operations using the
     tensor operations in LinAlg.
   }];
 
   let constructor = "tosa::createTosaToLinalgOnTensors()";
 }
 
 //===----------------------------------------------------------------------===//
+// TosaToStandard
+//===----------------------------------------------------------------------===//
+
+def TosaToStandard : FunctionPass<"tosa-to-standard"> {
+  let summary = "Lower TOSA to the Standard dialect";
+  let description = [{
+    Pass that converts TOSA operations to the equivalent operations using the
+    operations in the Standard dialect.
+  }];
+
+  let constructor = "tosa::createTosaToStandard()";
+}
+
+//===----------------------------------------------------------------------===//
 // VectorToSCF
 //===----------------------------------------------------------------------===//
 
 def ConvertVectorToSCF : FunctionPass<"convert-vector-to-scf"> {
   let summary = "Lower the operations from the vector dialect into the SCF "
                 "dialect";
   let constructor = "mlir::createConvertVectorToSCFPass()";
   let dependentDialects = ["AffineDialect", "scf::SCFDialect"];

mlir/include/mlir/Conversion/TosaToStandard/TosaToStandard.h

+//===-- TosaToStandard.h - TOSA optimization pass declarations --*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file declares the passes for the TOSA to Standard Dialect conversion.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_CONVERSION_TOSATOSTANDARD_TOSATOSTANDARD_H
+#define MLIR_CONVERSION_TOSATOSTANDARD_TOSATOSTANDARD_H
+
+#include "mlir/Pass/Pass.h"
+
+namespace mlir {
+namespace tosa {
+
+std::unique_ptr<Pass> createTosaToStandard();
+
+LogicalResult legalizeTosaControlFlow(Region &region);
+
+void populateTosaToStandardConversionPatterns(
+    MLIRContext *context, OwningRewritePatternList *patterns);
+
+/// Populates passes to convert from TOSA to Standard.
+void addTosaToStandardPasses(OpPassManager &pm);
+
+} // namespace tosa
+} // namespace mlir
+
+#endif // MLIR_CONVERSION_TOSATOSTANDARD_TOSATOSTANDARD_H

mlir/lib/Conversion/CMakeLists.txt

 add_subdirectory(SCFToOpenMP)
 add_subdirectory(SCFToSPIRV)
 add_subdirectory(SCFToStandard)
 add_subdirectory(ShapeToStandard)
 add_subdirectory(SPIRVToLLVM)
 add_subdirectory(StandardToLLVM)
 add_subdirectory(StandardToSPIRV)
 add_subdirectory(TosaToLinalg)
+add_subdirectory(TosaToStandard)
 add_subdirectory(ArmSVEToLLVM)
 add_subdirectory(VectorToROCDL)
 add_subdirectory(VectorToLLVM)
 add_subdirectory(VectorToSCF)
 add_subdirectory(VectorToSPIRV)

mlir/lib/Conversion/TosaToStandard/CMakeLists.txt

+add_mlir_conversion_library(MLIRTosaToStandard
+  TosaToStandard.cpp
+  TosaToStandardPass.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Tosa
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/IR
+
+  DEPENDS
+  MLIRConversionPassIncGen
+
+  LINK_LIBS PUBLIC
+  MLIRIR
+  MLIRStandard
+  MLIRPass
+  MLIRTensor
+  MLIRTosa
+  MLIRTosaTransforms
+  MLIRSupport
+  )

mlir/lib/Conversion/TosaToStandard/TosaToStandard.cpp

+//===- TosaToStandard.cpp - Lowering Tosa to Standard Dialect
+//-----------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// These rewriters lower from the Tosa to the Standard dialect.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Conversion/TosaToStandard/TosaToStandard.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Tosa/IR/TosaOps.h"
+#include "mlir/IR/BlockAndValueMapping.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+using namespace mlir;
+using namespace tosa;
+
+namespace {
+
+class ConstOpConverter : public OpRewritePattern<mlir::tosa::ConstOp> {
+public:
+  using OpRewritePattern<mlir::tosa::ConstOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(mlir::tosa::ConstOp op,
+                                PatternRewriter &rewriter) const final {
+    rewriter.replaceOpWithNewOp<mlir::ConstantOp>(op, op.value());
+    return success();
+  }
+};
+
+// Replaces all tosa.yield ops with a branch to the correct target. This is
+// used for processing each branch of tosa.if
+LogicalResult replaceYieldOps(Region *region, Block *targetBlock, Location loc,
+                              const BlockAndValueMapping &mapper,
+                              OpBuilder *builder) {
+  for (auto &oldBlock : region->getBlocks()) {
+    Block *block = mapper.lookup(&oldBlock);
+    auto returnOp = dyn_cast<tosa::YieldOp>(block->getTerminator());
+    if (!returnOp)
+      continue;
+    builder->setInsertionPointToEnd(block);
+    builder->create<mlir::BranchOp>(loc, targetBlock, returnOp.getOperands());
+    returnOp.erase();
+  }
+
+  return success();
+}
+
+LogicalResult legalizeTosaIf(IfOp ifOp) {
+  Operation *operation = ifOp.getOperation();
+  mlir::OpBuilder builder(ifOp);
+  auto *origBlock = operation->getBlock();
+  auto *tailBlock = origBlock->splitBlock(operation);
+  auto loc = ifOp.getLoc();
+
+  // Initially we clone the then/else blocks at the current point in the
+  // region so we can branch to them in the conditional check.
+  BlockAndValueMapping mapper;
+  ifOp.then_branch().cloneInto(origBlock->getParent(),
+                               Region::iterator(tailBlock), mapper);
+  ifOp.else_branch().cloneInto(origBlock->getParent(),
+                               Region::iterator(tailBlock), mapper);
+
+  Block *thenBranch = mapper.lookup(&ifOp.then_branch().front());
+  Block *elseBranch = mapper.lookup(&ifOp.else_branch().front());
+
+  // Move to the conditional section, extract the value from the condition
+  // tensor, and execute the conditional branch to the then/else cases.
+  builder.setInsertionPointToEnd(origBlock);
+  auto condValue = builder.create<mlir::tensor::ExtractOp>(loc, ifOp.cond());
+  builder.create<mlir::CondBranchOp>(loc, condValue, thenBranch, ifOp.inputs(),
+                                     elseBranch, ifOp.inputs());
+
+  // Yields need to be replaced with branches to the tile block.
+  if (failed(replaceYieldOps(&ifOp.then_branch(), tailBlock, loc, mapper,
+                             &builder)))
+    return failure();
+  if (failed(replaceYieldOps(&ifOp.else_branch(), tailBlock, loc, mapper,
+                             &builder)))
+    return failure();
+
+  tailBlock->addArguments(ifOp.getResultTypes());
+
+  for (auto it : llvm::enumerate(ifOp.getResults())) {
+    it.value().replaceAllUsesWith(tailBlock->getArgument(it.index()));
+  }
+
+  operation->erase();
+  return success();
+}
+
+LogicalResult legalizeTosaWhile(WhileOp whileOp) {
+  auto *operation = whileOp.getOperation();
+  mlir::OpBuilder builder(whileOp);
+  auto loc = whileOp.getLoc();
+
+  auto *origBlock = operation->getBlock();
+  auto *tailBlock = origBlock->splitBlock(operation);
+
+  // Copy the condition and body checks into their separate sections,
+  // lookup the newly created copies.
+  BlockAndValueMapping mapper;
+  whileOp.cond().cloneInto(origBlock->getParent(), Region::iterator(tailBlock),
+                           mapper);
+  whileOp.body().cloneInto(origBlock->getParent(), Region::iterator(tailBlock),
+                           mapper);
+
+  auto *condBlock = mapper.lookup(&whileOp.cond().front());
+  auto *bodyBlock = mapper.lookup(&whileOp.body().front());
+
+  // Move back to before the cond/body blocks and branch to the conditiono op.
+  builder.setInsertionPointToEnd(origBlock);
+  builder.create<mlir::BranchOp>(loc, condBlock, whileOp.getOperands());
+
+  // Update the inlined condition blocks by replacing any yield ops with an
+  // tensor.extract and conditional branch to the body block or tail.
+  builder.setInsertionPointToStart(condBlock);
+  for (auto &block : whileOp.cond()) {
+    auto *newBlock = mapper.lookup(&block);
+
+    auto yieldOp = dyn_cast<tosa::YieldOp>(newBlock->getTerminator());
+    if (!yieldOp)
+      continue;
+    builder.setInsertionPointToEnd(newBlock);
+
+    auto yieldValue = yieldOp.getOperand(0);
+    auto condValue = builder.create<mlir::tensor::ExtractOp>(loc, yieldValue);
+
+    // Pass the args passed to the cond block as the arguments for each branch
+    // condition.
+    llvm::SmallVector<Value, 4> successorArgs(condBlock->args_begin(),
+                                              condBlock->args_end());
+    builder.create<mlir::CondBranchOp>(loc, condValue, bodyBlock, successorArgs,
+                                       tailBlock, successorArgs);
+    yieldOp.erase();
+  }
+
+  // Replaces all yield s in the while op body with branches to the condition.
+  for (auto &block : whileOp.body()) {
+    auto *newBlock = mapper.lookup(&block);
+    auto yieldOp = dyn_cast<tosa::YieldOp>(newBlock->getTerminator());
+    if (!yieldOp)
+      continue;
+    builder.setInsertionPointToEnd(newBlock);
+    builder.create<mlir::BranchOp>(loc, condBlock, yieldOp.getOperands());
+    yieldOp.erase();
+  }
+
+  // Erase the original while loop.
+  tailBlock->addArguments(whileOp.getResultTypes());
+  for (auto it : llvm::enumerate(whileOp.getResults())) {
+    it.value().replaceAllUsesWith(tailBlock->getArgument(it.index()));
+  }
+  operation->erase();
+
+  return success();
+}
+
+} // namespace
+
+void mlir::tosa::populateTosaToStandardConversionPatterns(
+    MLIRContext *context, OwningRewritePatternList *patterns) {
+  patterns->insert<ConstOpConverter>(context);
+}
+
+LogicalResult mlir::tosa::legalizeTosaControlFlow(Region &region) {
+  // Both While and If legalization involves inling referenced blocks. As a
+  // result, this may require multiple passes with inline blocks.
+  auto whileOps = llvm::to_vector<5>(region.getOps<WhileOp>());
+  for (auto whileOp : whileOps) {
+    if (failed(legalizeTosaWhile(whileOp)))
+      return failure();
+  }
+
+  auto ifOps = llvm::to_vector<5>(region.getOps<IfOp>());
+  for (auto ifOp : ifOps) {
+    if (failed(legalizeTosaIf(ifOp)))
+      return failure();
+  }
+
+  return success();
+}

mlir/lib/Conversion/TosaToStandard/TosaToStandardPass.cpp

+//===- TosaToStandardPass.cpp - Lowering Tosa to Linalg Dialect -----------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This transformation pass legalizes Tosa operations to the Standard dialect.
+//
+//===----------------------------------------------------------------------===//
+
+#include "../PassDetail.h"
+#include "mlir/Conversion/TosaToStandard/TosaToStandard.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/Tosa/Transforms/PassDetail.h"
+#include "mlir/Dialect/Tosa/Transforms/Passes.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/PassManager.h"
+#include "mlir/Transforms/DialectConversion.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+using namespace mlir;
+using namespace tosa;
+
+namespace {
+struct TosaToStandard : public TosaToStandardBase<TosaToStandard> {
+public:
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<StandardOpsDialect, tensor::TensorDialect>();
+  }
+
+  void runOnFunction() override {
+    OwningRewritePatternList patterns;
+    ConversionTarget target(getContext());
+    target.addLegalDialect<tensor::TensorDialect, StandardOpsDialect>();
+    target.addIllegalOp<tosa::ConstOp>();
+    target.markUnknownOpDynamicallyLegal([](Operation *) { return true; });
+
+    FuncOp func = getFunction();
+    mlir::tosa::populateTosaToStandardConversionPatterns(func.getContext(),
+                                                         &patterns);
+    if (failed(applyFullConversion(func, target, std::move(patterns))))
+      signalPassFailure();
+
+    if (failed(legalizeTosaControlFlow(func.body())))
+      return signalPassFailure();
+  }
+};
+} // namespace
+
+std::unique_ptr<Pass> mlir::tosa::createTosaToStandard() {
+  return std::make_unique<TosaToStandard>();
+}
+
+void mlir::tosa::addTosaToStandardPasses(OpPassManager &pm) {
+  pm.addNestedPass<FuncOp>(createTosaToStandard());
+}

mlir/test/Conversion/TosaToStandard/tosa-to-standard.mlir

+// RUN: mlir-opt --split-input-file --tosa-to-standard %s -verify-diagnostics -o -| FileCheck %s
+
+
+// CHECK-LABEL: func @const_test
+func @const_test() -> (tensor<i32>) {
+  // CHECK: [[C3:%.+]] = constant dense<3> : tensor<i32>
+  %0 = "tosa.const"() {value = dense<3> : tensor<i32>} : () -> tensor<i32>
+
+  // CHECK: return [[C3]]
+  return %0 : tensor<i32>
+}
+
+// ----
+
+// CHECK-LABEL: func @while_test
+// CHECK-SAME: ([[ARG0:%.+]]: tensor<i32>)
+func @while_test(%arg0 : tensor<i32>) -> (tensor<i32>) {
+  // CHECK-NOT: "tosa.while_loop"
+  // CHECK: br [[B1:\^.+]]([[ARG0]] : tensor<i32>)
+  %1 = "tosa.while_loop"(%arg0) ( {
+  // CHECK: [[B1]]([[ARG1:%.+]]: tensor<i32>)
+  ^bb0(%arg2: tensor<i32>):  // no predecessors
+    // CHECK: [[C3:%.+]] = constant dense<3> : tensor<i32>
+    // CHECK: [[GT:%.+]] = "tosa.greater_equal"([[C3]], [[ARG1]])
+    // CHECK: [[EX:%.+]] = tensor.extract [[GT]][]
+    // CHECK: cond_br [[EX]], [[B2:\^.+]]([[ARG1]] : tensor<i32>), [[B3:\^.+]]([[ARG1]] : tensor<i32>)
+    %2 = "tosa.const"() {value = dense<3> : tensor<i32>} : () -> tensor<i32>
+    %3 = "tosa.greater_equal"(%2, %arg2) : (tensor<i32>, tensor<i32>) -> tensor<i1>
+    "tosa.yield"(%3) : (tensor<i1>) -> ()
+  },  {
+  // CHECK: [[B2]]([[ARG2:%.+]]: tensor<i32>)
+  ^bb0(%arg2: tensor<i32>):  // no predecessors
+    // CHECK: [[C1:%.+]] = constant dense<1> : tensor<i32>
+    // CHECK: [[ADD:%.+]] = "tosa.add"([[ARG2]], [[C1:%.+]])
+    // CHECK: br [[B1]]([[ADD]] : tensor<i32>)
+    %2 = "tosa.const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
+    %3 = "tosa.add"(%arg2, %2) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+    "tosa.yield"(%3) : (tensor<i32>) -> ()
+  }) : (tensor<i32>) -> (tensor<i32>)
+  // CHECK: [[B3]]([[RESULT:%.+]]: tensor<i32>)
+  // CHECK:   return [[RESULT]]
+  return %1 : tensor<i32>
+}
+
+// ----
+
+// CHECK-LABEL: func @if_test
+// CHECK-SAME: ([[ARG0:%.+]]: tensor<f32>, [[ARG1:%.+]]: tensor<f32>, [[ARG2:%.+]]: tensor<i1>)
+func @if_test(%arg0 : tensor<f32>, %arg1 : tensor<f32>, %arg2 : tensor<i1>) -> (tensor<f32>) {
+  // CHECK: [[EX:%.+]] = tensor.extract [[ARG2]][] : tensor<i1>
+  // CHECK: cond_br [[EX]], [[B1:\^.+]]([[ARG0]], [[ARG1]] : tensor<f32>, tensor<f32>), [[B2:\^.+]]([[ARG0]], [[ARG1]] : tensor<f32>, tensor<f32>)
+  %0 = "tosa.cond_if"(%arg2, %arg0, %arg1) ({
+  // CHECK: [[B1]]([[ARG3:%.+]]: tensor<f32>, [[ARG4:%.+]]: tensor<f32>)
+  ^bb1(%arg3 : tensor<f32>, %arg4 : tensor<f32>):
+    // CHECK:   br [[B3:\^.+]]([[ARG3]] : tensor<f32>)
+    "tosa.yield"(%arg3) : (tensor<f32>) -> ()
+  }, {
+  // CHECK: [[B2]]([[ARG5:%.+]]: tensor<f32>, [[ARG6:%.+]]: tensor<f32>):
+  ^bb1(%arg5 : tensor<f32>, %arg6 : tensor<f32>):
+    // CHECK:   br [[B3]]([[ARG6]] : tensor<f32>)
+    "tosa.yield"(%arg6) : (tensor<f32>) -> ()
+  }) : (tensor<i1>, tensor<f32>, tensor<f32>) -> (tensor<f32>)
+
+  // CHECK: [[B3]]([[ARG7:%.+]]: tensor<f32>):
+  // CHECK:   return [[ARG7]] : tensor<f32>
+  return %0 : tensor<f32>
+}