diff --git a/python/xmos_ai_tools/xinterpreters/src/dll_interpreter.cc b/python/xmos_ai_tools/xinterpreters/src/dll_interpreter.cc index 2f76ecea1..9aa571fde 100644 --- a/python/xmos_ai_tools/xinterpreters/src/dll_interpreter.cc +++ b/python/xmos_ai_tools/xinterpreters/src/dll_interpreter.cc @@ -90,6 +90,11 @@ DLLEXPORT inference_engine *new_interpreter(size_t max_arena_size, size_t extern resolver->AddIf(); resolver->AddWhile(); resolver->AddCallOnce(); + resolver->AddAbs(); + resolver->AddLog(); + resolver->AddSqrt(); + resolver->AddDiv(); + resolver->AddBatchMatMul(); tflite_micro::ops::micro::xcore::RegisterXCOps(resolver); add_lib_vision_ops(resolver); diff --git a/xformer/IR/XCoreOps.td b/xformer/IR/XCoreOps.td index d36d60cb2..ba3d60c2d 100644 --- a/xformer/IR/XCoreOps.td +++ b/xformer/IR/XCoreOps.td @@ -527,4 +527,17 @@ def XC_Bsign8Op : XC_Op<"bsign_8", [Pure]> { let results = (outs TensorOf<[I32]>:$output); } +def XC_BatchMatMulOp : XC_Op<"batch_matmul", [Pure]> { + let summary = "Batch MatMul op"; + + let description = [{Batch MatMul op.}]; + + let arguments = (ins TensorOf<[QI8]>:$input1, TensorOf<[QI8]>:$input2, + I32ArrayAttr:$compute_shape, + F32Attr:$lhs_zero_point, F32Attr:$rhs_zero_point, F32Attr:$out_zero_point, + F32Attr:$scale); + + let results = (outs TensorOf<[QI8]>:$output); +} + #endif // XCORE_DIALECT diff --git a/xformer/Transforms/Passes.cpp b/xformer/Transforms/Passes.cpp index d791b4751..4c74b48c7 100644 --- a/xformer/Transforms/Passes.cpp +++ b/xformer/Transforms/Passes.cpp @@ -48,6 +48,7 @@ void buildXCoreRemainingPassPipeline(OpPassManager &pm) { pm.addPass(createReplaceMaxPoolPass()); pm.addPass(createReplaceMulPass()); pm.addPass(createReplaceMeanPass()); + pm.addPass(createReplaceBatchMatMulPass()); pm.addPass(createReplaceSumPass()); pm.addPass(createReplaceTransposeConvPass()); pm.addPass(createReplaceConv2DPass()); diff --git a/xformer/Transforms/Passes.h b/xformer/Transforms/Passes.h index 8ed638f50..58815521b 100644 --- a/xformer/Transforms/Passes.h +++ b/xformer/Transforms/Passes.h @@ -37,6 +37,7 @@ std::unique_ptr> createRemoveDynamicShapePass(); std::unique_ptr> createReplaceAddSubPass(); std::unique_ptr> createReplaceMulPass(); std::unique_ptr> createReplaceMeanPass(); +std::unique_ptr> createReplaceBatchMatMulPass(); std::unique_ptr> createReplaceSumPass(); std::unique_ptr> createReplaceMaxPoolPass(); std::unique_ptr> createReplaceStridedSlicePass(); diff --git a/xformer/Transforms/ReplaceBatchMatMul.cpp b/xformer/Transforms/ReplaceBatchMatMul.cpp new file mode 100644 index 000000000..98953ad70 --- /dev/null +++ b/xformer/Transforms/ReplaceBatchMatMul.cpp @@ -0,0 +1,172 @@ +#include "IR/XCoreOps.h" +#include "Utils/Util.h" +#include "lib_nn/api/AggregateFn.hpp" +#include "lib_nn/api/MemCpyFn.hpp" +#include "lib_nn/api/OutputTransformFn.hpp" +extern "C" { +#include "lib_nn/api/nn_layers.h" +} +#include "mlir/IR/TypeUtilities.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" +#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h" +#include "tensorflow/compiler/mlir/lite/utils/validators.h" + +namespace mlir::xcore { + +namespace { +// Replace TFL Batch MatMul with Batch MatMul for XCore. +struct ReplaceBatchMatMul + : public PassWrapper> { + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ReplaceBatchMatMul) + + void getDependentDialects(DialectRegistry ®istry) const final { + registry.insert(); + } + StringRef getArgument() const final { return "xcore-replace-batch-matmul"; } + StringRef getDescription() const final { + return "Replace TFL Batch MatMul with Batch MatMul for XCore."; + } + void runOnOperation() override; +}; + +struct ReplaceBatchMatMulPattern : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(TFL::BatchMatMulOp batchMatMulOp, + PatternRewriter &rewriter) const override { + + auto X = batchMatMulOp.getX(); + auto Y = batchMatMulOp.getY(); + auto output = batchMatMulOp.getOutput(); + bool memcpyOnX = true; + bool reorderOnY = true; + + auto XType = X.getType().cast(); + auto YType = Y.getType().cast(); + auto outputType = output.getType().cast(); + + // Check if input and output are int8. + bool isInt8 = utils::isNBitSignedQType<8>(XType.getElementType()) && + utils::isNBitSignedQType<8>(YType.getElementType()) && + utils::isNBitSignedQType<8>(outputType.getElementType()); + + if (!isInt8) { + return failure(); + } + + auto XShape = XType.getShape(); + if (XShape.size() != 3) { + return mlir::failure(); // We don't support batch matmul with rank != 3 + } + + auto YShape = YType.getShape(); + if (YShape.size() != 3) { + return mlir::failure(); // We don't support batch matmul with rank != 3 + } + + if (XShape[0] != YShape[0] || XShape[2] != YShape[1]) { + return mlir::failure(); // X Y shape not match + } + + int64_t B = XShape[0]; + int64_t N = XShape[1]; + int64_t M = YShape[1]; + int64_t K = YShape[2]; + + auto XConstOp = dyn_cast_or_null(X.getDefiningOp()); + auto YConstOp = dyn_cast_or_null(Y.getDefiningOp()); + + if (XConstOp) { + llvm::outs() << "const op on x\n"; + } + + if (YConstOp) { + llvm::outs() << "const op on y\n"; + } + + if (XConstOp && YConstOp) { + // Shouldn't be both const + return failure(); + } + + bool adjX = batchMatMulOp.getAdjX(); + bool adjY = batchMatMulOp.getAdjY(); + + // Here we take the chance to reorder X or Y if they are const + if (YConstOp) { + auto valueAttr = YConstOp.getValue().cast(); + std::vector dataVec = std::vector( + valueAttr.getValues().begin(), valueAttr.getValues().end()); + llvm::outs() << "shape " << B << " " << M << " " << K << " total: " << dataVec.size() << "\n"; + llvm::outs() << "new shape " << B << " " << K << " " << M << "\n"; + + if (adjY == false) { + // Reorder to column-major (B, K, M), adjY == true means it's already column-major + std::vector oldVec = dataVec; + dataVec = std::vector(B*M*K); + for (int64_t b = 0; b < B; ++b) { + for (int64_t m = 0; m < M; ++m) { + for (int64_t k = 0; k < K; ++k) { + dataVec[b * (K * M) + k * M + m] = oldVec[b * (K * M) + m * K + k]; + } + } + } + } + + // Create new constant attribute with shape [B, K, M] + TypeAttr newTypeAttr = TypeAttr::get( + RankedTensorType::get( + {B, K, M}, YConstOp.getQtype().cast().getElementType())); + DenseElementsAttr newValueAttr = DenseElementsAttr::get( + RankedTensorType::get( + {B, K, M}, rewriter.getIntegerType(8)), dataVec); + Y = rewriter.create( + batchMatMulOp.getLoc(), newTypeAttr, newValueAttr); + reorderOnY = false; + } + + auto xQType = utils::getQType(X); + auto yQType = utils::getQType(Y); + auto outputQType = utils::getQType(output); + + float xZeroPoint = static_cast(xQType.getZeroPoint()); + float yZeroPoint = static_cast(yQType.getZeroPoint()); + float outZeroPoint = static_cast(outputQType.getZeroPoint()); + float xScale = static_cast(xQType.getScale()); + float yScale = static_cast(yQType.getScale()); + float outScale = static_cast(outputQType.getScale()); + + int32_t computeShape[4] = {B, N, M, K}; // batch, lhs row size, channel_size, rhs col size + + auto xcBatchMatMulOp = rewriter.create( + batchMatMulOp.getLoc(), batchMatMulOp.getType(), X, Y, + rewriter.getI32ArrayAttr(computeShape), + rewriter.getF32FloatAttr(xZeroPoint), + rewriter.getF32FloatAttr(yZeroPoint), + rewriter.getF32FloatAttr(outZeroPoint), + rewriter.getF32FloatAttr(xScale*yScale/outScale) + ); + rewriter.replaceOp(batchMatMulOp, xcBatchMatMulOp.getOutput()); + + return success(); + } +}; + +void ReplaceBatchMatMul::runOnOperation() { + auto *ctx = &getContext(); + func::FuncOp func = getOperation(); + RewritePatternSet patterns(ctx); + patterns.insert(ctx); + (void)applyPatternsAndFoldGreedily(func, std::move(patterns)); +} +} // namespace + +// Creates an instance of the ReplaceBatchMatMul pass. +std::unique_ptr> createReplaceBatchMatMulPass() { + return std::make_unique(); +} + +static PassRegistration pass; + +} // namespace mlir::xcore diff --git a/xformer/Transforms/TranslateToCustomOp.cpp b/xformer/Transforms/TranslateToCustomOp.cpp index 23756e020..caa227578 100644 --- a/xformer/Transforms/TranslateToCustomOp.cpp +++ b/xformer/Transforms/TranslateToCustomOp.cpp @@ -281,6 +281,24 @@ std::vector MaxPool2DOp::buildCustomOptions() { return fbb.GetBuffer(); } +std::vector BatchMatMulOp::buildCustomOptions() { + flexbuffers::Builder fbb; + auto rootMap = fbb.StartMap(); + auto computeShapeVec = fbb.StartVector("compute_shape"); + auto computeShape = getComputeShape().cast(); + for (int i = 0; i < 4; ++i) { + fbb.Int(computeShape[i].cast().getInt()); + } + fbb.EndVector(computeShapeVec, false, false); + fbb.IndirectFloat("lhs_zp", getLhsZeroPoint().convertToFloat()); + fbb.IndirectFloat("rhs_zp", getRhsZeroPoint().convertToFloat()); + fbb.IndirectFloat("out_zp", getOutZeroPoint().convertToFloat()); + fbb.IndirectFloat("scale", getScale().convertToFloat()); + fbb.EndMap(rootMap); + fbb.Finish(); + return fbb.GetBuffer(); +} + namespace { /// This pass translates XCore ops to TFLite custom ops. struct TranslateToCustomOp @@ -330,6 +348,7 @@ void TranslateToCustomOp::runOnOperation() { patterns.insert>(ctx); patterns.insert>(ctx); patterns.insert>(ctx); + patterns.insert>(ctx); patterns.insert>(ctx); patterns.insert>(ctx); patterns.insert>(ctx); diff --git a/xformer/lib_tflite_micro.BUILD b/xformer/lib_tflite_micro.BUILD index 1c327eb7f..2e9fde6a1 100644 --- a/xformer/lib_tflite_micro.BUILD +++ b/xformer/lib_tflite_micro.BUILD @@ -55,5 +55,6 @@ filegroup( "lib_tflite_micro/src/tflite-xcore-kernels/xcore_beta_transposeconvf32.cc", "lib_tflite_micro/src/tflite-xcore-kernels/xcore_beta_fcf32.cc", "lib_tflite_micro/src/tflite-xcore-kernels/conv2d_float.c", + "lib_tflite_micro/src/tflite-xcore-kernels/xcore_batch_matmul.cc", ], ) diff --git a/xformer/lib_tflmc.BUILD b/xformer/lib_tflmc.BUILD index 3a6f90370..e55d77493 100644 --- a/xformer/lib_tflmc.BUILD +++ b/xformer/lib_tflmc.BUILD @@ -67,6 +67,7 @@ filegroup( "@tflite_micro//tensorflow/lite/micro/kernels:slice.cc", "@tflite_micro//tensorflow/lite/micro/kernels:arg_min_max.cc", "@tflite_micro//tensorflow/lite/micro/kernels:batch_to_space_nd.cc", + "@tflite_micro//tensorflow/lite/micro/kernels:batch_matmul.cc", "@tflite_micro//tensorflow/lite/micro/kernels:broadcast_args.cc", "@tflite_micro//tensorflow/lite/micro/kernels:broadcast_to.cc", "@tflite_micro//tensorflow/lite/micro/kernels:cast.cc",