TosaToSPIRVTosaPass.cpp

Go to the documentation of this file.

//===- TosaToSPIRVTosaPass.cpp - Lower TOSA to SPIR-V Graph/TOSA ----------===//
//
// 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 pass lowers TOSA IR to the SPIR-V Graph/TOSA representation.
//
//===----------------------------------------------------------------------===//
 
#include "mlir/Conversion/TosaToSPIRVTosa/TosaToSPIRVTosa.h"
 
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
#include "mlir/Dialect/SPIRV/IR/TargetAndABI.h"
#include "mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h"
#include "mlir/Dialect/Tosa/IR/TosaOps.h"
#include "mlir/IR/TypeUtilities.h"
#include "mlir/Transforms/DialectConversion.h"
 
#include <algorithm>
 
namespace mlir {
#define GEN_PASS_DEF_TOSATOSPIRVTOSA
#include "mlir/Conversion/Passes.h.inc"
 
namespace tosa {
 
spirv::VerCapExtAttr getDefaultVerCapExtAttr(MLIRContext *context) {
  return spirv::VerCapExtAttr::get(
      spirv::Version::V_1_5,
      {
          spirv::Capability::VulkanMemoryModel,
          spirv::Capability::Shader,
          spirv::Capability::Int8,
          spirv::Capability::Int16,
          spirv::Capability::Int64,
          spirv::Capability::Float16,
          spirv::Capability::BFloat16TypeKHR,
          spirv::Capability::Float8EXT,
          spirv::Capability::TensorsARM,
          spirv::Capability::GraphARM,
          spirv::Capability::ReplicatedCompositesEXT,
      },
      {
          spirv::Extension::SPV_ARM_tensors,
          spirv::Extension::SPV_ARM_graph,
          spirv::Extension::SPV_KHR_vulkan_memory_model,
          spirv::Extension::SPV_EXT_replicated_composites,
          spirv::Extension::SPV_KHR_bfloat16,
          spirv::Extension::SPV_EXT_float8,
      },
      context);
}
 
spirv::TargetEnvAttr constructTargetEnvAttrWithCapExtDefaults(
    MLIRContext *context, spirv::ResourceLimitsAttr limits,
    spirv::ClientAPI clientAPI, spirv::Vendor vendorID,
    spirv::DeviceType deviceType, uint32_t deviceID) {
  if (!limits)
    limits = spirv::getDefaultResourceLimits(context);
 
  return spirv::TargetEnvAttr::get(getDefaultVerCapExtAttr(context), limits,
                                   clientAPI, vendorID, deviceType, deviceID);
}
 
namespace {
 
LogicalResult verifyGraphTargetEnv(Operation *op,
                                   spirv::TargetEnvAttr targetAttr) {
  spirv::TargetEnv targetEnv(targetAttr);
  if (targetEnv.allows(spirv::Capability::GraphARM) &&
      targetEnv.allows(spirv::Extension::SPV_ARM_graph) &&
      targetEnv.allows(spirv::Extension::SPV_ARM_tensors)) {
    return success();
  }
 
  return op->emitOpError()
         << "requires GraphARM capability and SPV_ARM_graph/SPV_ARM_tensors "
            "extensions in spirv.target_env";
}
 
LogicalResult verifyNoUnsupportedFuncOps(Operation *op) {
  WalkResult result = op->walk([](Operation *op) -> WalkResult {
    if (isa<func::CallOp, func::CallIndirectOp>(op)) {
      op->emitOpError()
          << "is not supported in TOSA to SPIR-V Graph conversion; inline "
             "calls before running this pass";
      return WalkResult::interrupt();
    }
    if (auto funcOp = dyn_cast<func::FuncOp>(op)) {
      if (funcOp->getParentOfType<func::FuncOp>()) {
        funcOp.emitOpError()
            << "nesting is not supported in TOSA to SPIR-V Graph conversion";
        return WalkResult::interrupt();
      }
    }
    return WalkResult::advance();
  });
 
  return failure(result.wasInterrupted());
}
 
struct TosaToSPIRVTosa final : impl::TosaToSPIRVTosaBase<TosaToSPIRVTosa> {
  void runOnOperation() override {
    MLIRContext *context = &getContext();
    RewritePatternSet patterns(context);
    Operation *op = getOperation();
 
    spirv::TargetEnvAttr targetAttr = spirv::lookupTargetEnv(op);
    if (!targetAttr) {
      targetAttr = constructTargetEnvAttrWithCapExtDefaults(context);
    }
 
    if (failed(verifyGraphTargetEnv(op, targetAttr)) ||
        failed(verifyNoUnsupportedFuncOps(op))) {
      signalPassFailure();
      return;
    }
 
    std::unique_ptr<ConversionTarget> target =
        SPIRVConversionTarget::get(targetAttr);
 
    target->addIllegalDialect<tosa::TosaDialect>();
    target->addIllegalOp<func::CallOp, func::CallIndirectOp>();
 
    SPIRVTypeConverter typeConverter(targetAttr);
    typeConverter.addConversion([this](IntegerType integerType) {
      return this->convertIntegerType(integerType);
    });
    typeConverter.addConversion([this](TensorType tensorType) {
      return this->convertTensorType(tensorType);
    });
    typeConverter.addConversion([this](tosa::shapeType shapeType) {
      return this->convertShapeType(shapeType);
    });
 
    populateTosaToSPIRVTosaConversionPatterns(typeConverter, patterns,
                                              targetAttr);
    populateTosaToSPIRVTosaOpsConversionPatterns(typeConverter, patterns);
 
    FrozenRewritePatternSet frozenPatterns(std::move(patterns));
 
    if (failed(applyPartialConversion(op, *target, frozenPatterns))) {
      signalPassFailure();
    }
  }
 
private:
  IntegerType convertIntegerType(IntegerType integerType) {
    if (integerType.getWidth() == 48) {
      return IntegerType::get(&getContext(), 64, integerType.getSignedness());
    }
 
    if (integerType.getWidth() == 4) {
      return IntegerType::get(&getContext(), 8, integerType.getSignedness());
    }
 
    return integerType;
  }
 
  std::optional<SmallVector<int64_t>> convertShape(ArrayRef<int64_t> shape) {
    // Scalar ARM tensors are not supported, so convert them to
    // tensors with shape [1].
    if (shape.empty())
      return SmallVector<int64_t>({1});
 
    if (llvm::is_contained(shape, 0))
      return std::nullopt;
 
    bool isPartiallyDynamic =
        llvm::is_contained(shape, ShapedType::kDynamic) &&
        llvm::any_of(shape, [](int64_t dim) { return dim > 0; });
    // Partially shaped ARM tensors are not supported, so convert them to
    // unshaped tensors.
    if (isPartiallyDynamic)
      return SmallVector<int64_t>(shape.size(), ShapedType::kDynamic);
    return SmallVector<int64_t>(shape);
  }
 
  std::optional<spirv::TensorArmType> convertTensorType(TensorType tensorType) {
    Type elementType = getElementTypeOrSelf(tensorType);
    if (elementType.isIndex())
      elementType = IntegerType::get(&getContext(), 32);
    if (auto integerType = dyn_cast<IntegerType>(elementType))
      elementType = convertIntegerType(integerType);
 
    SmallVector<int64_t> shape;
    if (tensorType.hasRank()) {
      std::optional<SmallVector<int64_t>> convertedShape =
          convertShape(tensorType.getShape());
      if (!convertedShape)
        return std::nullopt;
      shape = std::move(*convertedShape);
    }
 
    return spirv::TensorArmType::get(shape, elementType);
  }
 
  spirv::TensorArmType convertShapeType(tosa::shapeType shapeType) {
    const int64_t rank = std::max(shapeType.getRank(), 1);
    return spirv::TensorArmType::get({rank},
                                     IntegerType::get(&getContext(), 32));
  }
};
} // namespace
 
std::unique_ptr<Pass> createTosaToSPIRVTosa() {
  return std::make_unique<TosaToSPIRVTosa>();
}
 
} // namespace tosa
} // namespace mlir