Op fusion experiment

src/Accelerators/NNPA/Transform/ZHigh/Stickify/Stickify.cpp

@@ -12,6 +12,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include <errno.h>
 #include <fenv.h>
 #include <stdarg.h>
 #include <stdbool.h>

src/Conversion/ONNXToKrnl/CMakeLists.txt

@@ -39,6 +39,7 @@ add_onnx_mlir_library(OMONNXToKrnl
   Tensor/ArgMinMax.cpp
   Tensor/Compress.cpp  
   Tensor/Concat.cpp
+  Tensor/ConcatShapeTranspose.cpp
   Tensor/Constant.cpp
   Tensor/ConstantOfShape.cpp
   Tensor/DepthToSpace.cpp

src/Conversion/ONNXToKrnl/ConvertONNXToKrnl.cpp

@@ -214,6 +214,8 @@ void populateONNXToKrnlConversionPattern(RewritePatternSet &patterns,
   populateLoweringONNXConstantOfShapeOpPattern(patterns, typeConverter, ctx);
   populateLoweringONNXConstantOpPattern(patterns, typeConverter, ctx);
   populateLoweringONNXConcatOpPattern(patterns, typeConverter, ctx);
+  populateLoweringONNXConcatShapeTransposeOpPattern(
+      patterns, typeConverter, ctx);
   populateLoweringONNXDepthToSpaceOpPattern(patterns, typeConverter, ctx);
   populateLoweringONNXScatterElementsOpPattern(patterns, typeConverter, ctx);
   populateLoweringONNXScatterNDOpPattern(patterns, typeConverter, ctx);

src/Conversion/ONNXToKrnl/ONNXToKrnlCommon.hpp

@@ -357,6 +357,8 @@ void populateLoweringONNXConstantOpPattern(
     mlir::RewritePatternSet &, mlir::TypeConverter &, mlir::MLIRContext *);
 void populateLoweringONNXConcatOpPattern(
     mlir::RewritePatternSet &, mlir::TypeConverter &, mlir::MLIRContext *);
+void populateLoweringONNXConcatShapeTransposeOpPattern(
+    mlir::RewritePatternSet &, mlir::TypeConverter &, mlir::MLIRContext *);
 void populateLoweringONNXDepthToSpaceOpPattern(
     mlir::RewritePatternSet &, mlir::TypeConverter &, mlir::MLIRContext *);
 void populateLoweringONNXSpaceToDepthOpPattern(

src/Conversion/ONNXToKrnl/Tensor/ConcatShapeTranspose.cpp

@@ -0,0 +1,183 @@
+
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+//===-------------- Concat.cpp - Lowering ConcatShapeTranspose Op ---------===//
+//
+// Copyright 2022 The IBM Research Authors.
+//
+// =============================================================================
+//
+// This file lowers the ONNX Concat OperatorShapeTranspose to Krnl dialect.
+//
+//===----------------------------------------------------------------------===//
+
+#include "src/Conversion/ONNXToKrnl/ONNXToKrnlCommon.hpp"
+#include "src/Dialect/Krnl/KrnlHelper.hpp"
+#include "src/Dialect/ONNX/ONNXOps/ShapeHelper.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+
+struct ONNXConcatShapeTransposeOpLowering : public ConversionPattern {
+  ONNXConcatShapeTransposeOpLowering(
+      TypeConverter &typeConverter, MLIRContext *ctx)
+      : ConversionPattern(typeConverter,
+            mlir::ONNXConcatShapeTransposeOp::getOperationName(), 1, ctx) {}
+
+  LogicalResult matchAndRewrite(Operation *op, ArrayRef<Value> operands,
+      ConversionPatternRewriter &rewriter) const final {
+    Location loc = op->getLoc();
+    ONNXConcatShapeTransposeOp concatShapeTransposeOp =
+        llvm::cast<ONNXConcatShapeTransposeOp>(op);
+    ONNXConcatShapeTransposeOpAdaptor operandAdaptor(
+        operands, op->getAttrDictionary());
+    ONNXConcatShapeTransposeOpShapeHelper shapeHelper(&concatShapeTransposeOp,
+        &rewriter, krnl::getDenseElementAttributeFromKrnlValue,
+        krnl::loadDenseElementArrayValueAtIndex);
+    auto shapecomputed = shapeHelper.computeShape(operandAdaptor);
+    (void)shapecomputed;
+    assert(succeeded(shapecomputed) && "Could not compute output shape");
+
+    MultiDialectBuilder<KrnlBuilder, MathBuilder> create(rewriter, loc);
+
+    // Compute concat output shape.
+    unsigned numInputs = op->getNumOperands();
+    Value firstInput = operandAdaptor.inputs().front();
+    ArrayRef<int64_t> commonShape =
+        firstInput.getType().cast<ShapedType>().getShape();
+    // firstInput.getType().cast<ShapedType>().getElementType();
+    uint64_t rank = commonShape.size();
+    int64_t axis = operandAdaptor.axis();
+
+    // Negative axis means values are counted from the opposite side.
+    // TOFIX should be in normalization pass
+    if (axis < 0)
+      axis += rank;
+
+    IndexExprScope IEScope(&rewriter, loc);
+    DimsExpr outputConcatDims(rank);
+    MemRefBoundsIndexCapture firstInputBounds(operandAdaptor.inputs()[0]);
+    for (unsigned dim = 0; dim < rank; dim++) {
+      outputConcatDims[dim] = firstInputBounds.getDim(dim);
+    }
+    IndexExpr cumulativeAxisSize = DimIndexExpr(firstInputBounds.getDim(axis));
+
+    // Handle the rest of input
+    for (unsigned i = 1; i < numInputs; ++i) {
+      Value currentInput = operandAdaptor.inputs()[i];
+      MemRefBoundsIndexCapture currInputBounds(currentInput);
+      for (unsigned dim = 0; dim < rank; dim++) {
+        if (dim == axis) {
+          DimIndexExpr currentSize(currInputBounds.getDim(axis));
+          cumulativeAxisSize = cumulativeAxisSize + currentSize;
+        } else {
+          if (currInputBounds.getDim(dim).isLiteral()) {
+            // The size of current dimension of current input  is a constant
+            outputConcatDims[dim] = currInputBounds.getDim(dim);
+          }
+        }
+      }
+    }
+    outputConcatDims[axis] = cumulativeAxisSize;
+
+    // Shape for Shape
+    uint64_t start = operandAdaptor.start();
+    uint64_t end = rank;
+    if (operandAdaptor.end().has_value()) {
+      end = operandAdaptor.end().value();
+    }
+    // Handle negative
+    if (start < 0)
+      start += rank;
+
+    if (end < 0)
+      end += rank;
+
+    Type outputShapeType = op->getResultTypes()[0];
+
+    // Alloc and set value for ShapeOp output
+    auto convertedShapeType =
+        typeConverter->convertType(outputShapeType).cast<MemRefType>();
+    Value shapeAlloc = insertAllocAndDeallocSimple(
+        rewriter, op, convertedShapeType, loc, shapeHelper.dimsForOutput());
+    Type elementType = convertedShapeType.getElementType();
+    for (uint64_t i = start; i < end; i++) {
+      Value intVal =
+          create.math.cast(elementType, outputConcatDims[i].getValue());
+      create.krnl.store(
+          intVal, shapeAlloc, create.math.constantIndex(i - start));
+    }
+
+    // Convert the output type to MemRefType.
+    DimsExpr outputTransposeDims = shapeHelper.dimsForOutput(1);
+    ArrayAttr permAttr = operandAdaptor.permAttr();
+    Type t = op->getResultTypes()[1];
+    auto outputTransposeType = typeConverter->convertType(t).cast<MemRefType>();
+    Value alloc = insertAllocAndDeallocSimple(
+        rewriter, op, outputTransposeType, loc, outputTransposeDims);
+
+    // Creates loops, one for each input.
+    // Since the each input should have same size for each dimension(except
+    // axis), we will try to make the loop upper bound the same for futher
+    // optimization. Difference may come from constant vs. dynamic, or dynamic
+    // dim of different inputs.
+    KrnlBuilder createKrnl(rewriter, loc);
+    SmallVector<IndexExpr, 4> commonUB = outputConcatDims;
+    IndexExpr accumulatedOffset = LiteralIndexExpr(0);
+    for (unsigned int i = 0; i < numInputs; ++i) {
+      // Since the acculatedOffsetValue will be used in a nested IndexExprScope,
+      // we get the Value of this IndexExpr and pass it as a symbol
+      Value accumulatedOffsetValue = accumulatedOffset.getValue();
+      OpBuilder::InsertionGuard insertGuard(rewriter);
+      // Create loop.
+      ValueRange loopDef = createKrnl.defineLoops(rank);
+      SmallVector<IndexExpr, 4> lbs(rank, LiteralIndexExpr(0));
+      MemRefBoundsIndexCapture bounds(operands[i]);
+      SmallVector<IndexExpr, 4> ubs;
+      bounds.getDimList(ubs);
+      // For each input, only the dimension 'axis' is different
+      commonUB[axis] = ubs[axis];
+      createKrnl.iterateIE(loopDef, loopDef, lbs, commonUB,
+          [&](KrnlBuilder &createKrnl, ValueRange loopInd) {
+            // Indices for the read and write.
+            SmallVector<Value, 4> readIndices, writeIndices;
+            for (unsigned int r = 0; r < rank; ++r) {
+              if (r != axis || i == 0)
+                writeIndices.emplace_back(loopInd[r]);
+              else {
+                IndexExprScope IEScope(&rewriter, loc);
+                IndexExpr writeOffset = DimIndexExpr(loopInd[r]);
+                IndexExpr accumulatedOffsetIE =
+                    SymbolIndexExpr(accumulatedOffsetValue);
+                writeOffset = writeOffset + accumulatedOffsetIE;
+                writeIndices.emplace_back(writeOffset.getValue());
+              }
+            }
+            // Insert copy.
+            Value loadData = createKrnl.load(operands[i], loopInd);
+            SmallVector<Value, 4> transposedWriteIndices(rank);
+            for (uint64_t i = 0; i < rank; i++) {
+              transposedWriteIndices[i] =
+                  writeIndices[ArrayAttrIntVal(permAttr, i)];
+            }
+            createKrnl.store(loadData, alloc, transposedWriteIndices);
+          });
+      MemRefBoundsIndexCapture operandJBounds(operands[i]);
+      accumulatedOffset = accumulatedOffset + operandJBounds.getDim(axis);
+    }
+    SmallVector<Value, 2> outputs;
+    rewriter.replaceOp(op, {shapeAlloc, alloc});
+    return success();
+  }
+};
+
+void populateLoweringONNXConcatShapeTransposeOpPattern(
+    RewritePatternSet &patterns, TypeConverter &typeConverter,
+    MLIRContext *ctx) {
+  patterns.insert<ONNXConcatShapeTransposeOpLowering>(typeConverter, ctx);
+}
+
+} // namespace onnx_mlir

src/Dialect/ONNX/AdditionalONNXOps.td

@@ -380,3 +380,24 @@ def ONNXMaxPoolSingleOutOp: ONNX_Op<"MaxPoolSingleOut",
     static std::vector<int> getTypeMap() { return {20}; }
   }];
 }
+
+//===----------------------------------------------------------------------===//
+// ConcatShapeTransposeOp
+def ONNXConcatShapeTransposeOp: ONNX_Op<"ConcatShapeTranspose", [Pure,
+    DeclareOpInterfaceMethods<ShapeInferenceOpInterface>]> {
+  let summary = "ONNX merged operation";
+  let description = [{
+    Merge the following sequence of ops into one op
+    v1 = onnx.concat
+    v2 = onnx.shape(v1)
+    v3 = onnx.transpose(v1)
+  }];
+  let arguments = (ins Variadic<AnyTypeOf<[TensorOf<[UI8]>, TensorOf<[UI16]>, TensorOf<[UI32]>, TensorOf<[UI64]>, TensorOf<[I8]>, TensorOf<[I16]>, TensorOf<[I32]>, TensorOf<[I64]>, TensorOf<[BF16]>, TensorOf<[F16]>, TensorOf<[F32]>, TensorOf<[F64]>, TensorOf<[StringType]>, TensorOf<[I1]>, TensorOf<[Complex<F32>]>, TensorOf<[Complex<F64>]>]>>:$inputs,
+  SI64Attr:$axis,
+  OptionalAttr<SI64Attr>:$end,
+  DefaultValuedAttr<SI64Attr, "0">:$start,
+  OptionalAttr<I64ArrayAttr>:$perm);
+  let results = (outs TensorOf<[I64]>:$shape,
+     AnyTypeOf<[TensorOf<[UI8]>, TensorOf<[UI16]>, TensorOf<[UI32]>, TensorOf<[UI64]>, TensorOf<[I8]>, TensorOf<[I16]>, TensorOf<[I32]>, TensorOf<[I64]>, TensorOf<[BF16]>, TensorOf<[F16]>, TensorOf<[F32]>, TensorOf<[F64]>, TensorOf<[StringType]>, TensorOf<[I1]>, TensorOf<[Complex<F32>]>, TensorOf<[Complex<F64>]>]>:$transposed);
+}
+

src/Dialect/ONNX/CMakeLists.txt

@@ -24,6 +24,7 @@ add_onnx_mlir_library(OMONNXOps
 
   # Support for shape inference and verifiers
   ONNXOps/Additional/CallOp.cpp
+  ONNXOps/Additional/ConcatShapeTranspose.cpp
   ONNXOps/Additional/Custom.cpp
   ONNXOps/Additional/Dim.cpp
   ONNXOps/Additional/EntryPoint.cpp