Lower ONNXConstantOfShapeOp to Krnl dialect (#296)

* Lower ONNXConstantOfShapeOp to Krnl dialect

* Change a variable name

* Add comments to lit tests

Co-authored-by: Alexandre Eichenberger <alexe@us.ibm.com>

src/Conversion/ONNXToKrnl/CMakeLists.txt

@@ -20,6 +20,7 @@ add_library(OMONNXToKrnl
         Tensor/Squeeze.cpp
         Tensor/Unsqueeze.cpp
         Tensor/Constant.cpp
+        Tensor/ConstantOfShape.cpp
         Tensor/Concat.cpp
         Tensor/Split.cpp
         Tensor/Gather.cpp

src/Conversion/ONNXToKrnl/ConvertONNXToKrnl.cpp

@@ -98,6 +98,7 @@ void FrontendToKrnlLoweringPass::runOnOperation() {
   populateLoweringONNXTransposeOpPattern(patterns, &getContext());
   populateLoweringONNXGatherOpPattern(patterns, &getContext());
   populateLoweringONNXIdentityOpPattern(patterns, &getContext());
+  populateLoweringONNXConstantOfShapeOpPattern(patterns, &getContext());
   populateLoweringONNXConstantOpPattern(patterns, &getContext());
   populateLoweringONNXConcatOpPattern(patterns, &getContext());
   populateLoweringONNXSqueezeOpPattern(patterns, &getContext());

src/Conversion/ONNXToKrnl/ONNXToKrnlCommon.hpp

@@ -238,6 +238,9 @@ void populateLoweringONNXReshapeOpPattern(
 void populateLoweringONNXIdentityOpPattern(
     OwningRewritePatternList &patterns, MLIRContext *ctx);
 
+void populateLoweringONNXConstantOfShapeOpPattern(
+    OwningRewritePatternList &patterns, MLIRContext *ctx);
+
 void populateLoweringONNXConstantOpPattern(
     OwningRewritePatternList &patterns, MLIRContext *ctx);
 

src/Conversion/ONNXToKrnl/Tensor/ConstantOfShape.cpp

@@ -0,0 +1,100 @@
+//===------------ ConstantOfShape.cpp - Lowering ConstantOfShape Op -------===//
+//
+// Copyright 2019 The IBM Research Authors.
+//
+// =============================================================================
+//
+// This file lowers the ONNX ConstantOfShape Operator to Krnl dialect.
+//
+//===----------------------------------------------------------------------===//
+
+#include "src/Conversion/ONNXToKrnl/ONNXToKrnlCommon.hpp"
+
+using namespace mlir;
+
+struct ONNXConstantOfShapeOpLowering : public ConversionPattern {
+  ONNXConstantOfShapeOpLowering(MLIRContext *ctx)
+      : ConversionPattern(
+            mlir::ONNXConstantOfShapeOp::getOperationName(), 1, ctx) {}
+
+  LogicalResult matchAndRewrite(Operation *op, ArrayRef<Value> operands,
+      ConversionPatternRewriter &rewriter) const final {
+    auto loc = op->getLoc();
+    ONNXConstantOfShapeOpAdaptor operandAdaptor(operands);
+
+    auto valueAttr = llvm::cast<ONNXConstantOfShapeOp>(op)
+                         .value()
+                         .getValue()
+                         .cast<DenseElementsAttr>();
+
+    auto memRefType = convertToMemRefType(*op->result_type_begin());
+    auto elementType = memRefType.getElementType();
+    size_t rank = memRefType.cast<ShapedType>().getRank();
+
+    // Allocate memory for the output.
+    Value alloc;
+    bool insertDealloc = checkInsertDealloc(op);
+    if (hasAllConstantDimensions(memRefType))
+      alloc = insertAllocAndDealloc(memRefType, loc, rewriter, insertDealloc);
+    else {
+      SmallVector<Value, 2> allocOperands;
+      // Load dimensions from the input.
+      for (decltype(rank) i = 0; i < rank; ++i) {
+        auto index = emitConstantOp(rewriter, loc, rewriter.getIndexType(), i);
+        auto dim =
+            rewriter.create<AffineLoadOp>(loc, operandAdaptor.input(), index);
+        auto dimIndex =
+            rewriter.create<IndexCastOp>(loc, rewriter.getIndexType(), dim);
+        allocOperands.emplace_back(dimIndex);
+      }
+      // Allocate memory.
+      alloc = rewriter.create<AllocOp>(loc, memRefType, allocOperands);
+      // Insert deallocation if needed.
+      if (insertDealloc) {
+        Block *parentBlock = alloc.getDefiningOp()->getBlock();
+        DeallocOp dealloc = rewriter.create<DeallocOp>(loc, alloc);
+        dealloc.getOperation()->moveBefore(&parentBlock->back());
+      }
+    }
+
+    // Get the constant value from the attribute 'value'.
+    Value constantVal;
+    if (elementType.isa<IntegerType>()) {
+      auto valueIt = valueAttr.getValues<IntegerAttr>().begin();
+      auto valueInt = (*valueIt++).cast<IntegerAttr>().getInt();
+      constantVal = emitConstantOp(rewriter, loc, elementType, valueInt);
+    } else if (elementType.isa<FloatType>()) {
+      auto valueIt = valueAttr.getValues<FloatAttr>().begin();
+      auto valueFloat = (*valueIt++).cast<FloatAttr>().getValueAsDouble();
+      constantVal = emitConstantOp(rewriter, loc, elementType, valueFloat);
+    } else {
+      llvm_unreachable("unsupported element type");
+    }
+
+    SmallVector<Value, 4> loopIVs;
+    // Create a Krnl iterate if the output is not a scalar tensor.
+    if (!hasAllScalarValues({alloc})) {
+      BuildKrnlLoop loops(rewriter, loc, rank);
+      loops.createDefineAndIterateOp(alloc);
+      Block *iterationBlock = loops.getIterateBlock();
+      // Get IVs.
+      for (auto arg : iterationBlock->getArguments())
+        loopIVs.push_back(arg);
+      // Insert instructions inside the KernelIterateOp body.
+      rewriter.setInsertionPointToStart(iterationBlock);
+    }
+
+    // Store the constant value to the output.
+    rewriter.create<AffineStoreOp>(loc, constantVal, alloc, loopIVs);
+
+    // Replace this operation with the generated alloc.
+    rewriter.replaceOp(op, alloc);
+
+    return success();
+  }
+};
+
+void populateLoweringONNXConstantOfShapeOpPattern(
+    OwningRewritePatternList &patterns, MLIRContext *ctx) {
+  patterns.insert<ONNXConstantOfShapeOpLowering>(ctx);
+}

test/backend/test.py

@@ -420,6 +420,18 @@ test_to_enable = [
     "test_split_variable_parts_2d_cpu",
     "test_split_variable_parts_default_axis_cpu",
 
+    # ConstantOfShape
+    "test_constantofshape_float_ones_cpu",
+    # Error:
+    #    Items are not equal:
+    #     ACTUAL: dtype('int32')
+    #     DESIRED: dtype('uint8')
+    # In this test, 'int32' was specified for value attribute as in
+    # onnx/onnx/backend/test/case/node/constantofshape.py
+    # and onnx-mlir correctly imported and converted the model.
+    # It is unknown why 'uint8' came from.
+    #"test_constantofshape_int_zeros_cpu",
+
     # Model
     "test_resnet50_cpu",
     "test_vgg19_cpu",

test/mlir/onnx/onnx_lowering.mlir

@@ -2170,3 +2170,83 @@ func @test_gather_axis1(%arg0 : tensor<3x3xf32>) -> tensor<1x3x2xf32> {
   // CHECK: [[DATA:%.+]] = load %arg0{{.}}[[ARG1]], [[AFFINE2]]{{.}} : memref<3x3xf32>
   // CHECK: affine.store [[DATA]], [[ALLOC]]{{.}}[[ARG1]], [[ARG2]], [[ARG3]]{{.}} : memref<1x3x2xf32>
 }
+
+// -----
+
+// Check the lowering of ConstantOfShape when:
+//   - No value attribute.
+//   - The input is an empty tensor.
+// Expected emitted code:
+//   - No need a Krnl iterate.
+//   - The output is a scalar tensor.
+func @test_constant_of_shape_empty_tensor(%arg0 : tensor<0xi64>) -> tensor<*xf32> {
+  %0 = "onnx.ConstantOfShape"(%arg0) : (tensor<0xi64>) -> tensor<*xf32>
+  "std.return"(%0) : (tensor<*xf32>) -> ()
+
+  // CHECK-LABEL: test_constant_of_shape_empty_tensor
+  // CHECK: [[RES:%.+]] = alloc() : memref<f32>
+  // CHECK: [[CST_VALUE:%.+]] = constant 0.000000e+00 : f32
+  // CHECK: affine.store [[CST_VALUE]], [[RES]][] : memref<f32>
+  // CHECK: return [[RES]] : memref<f32>
+}
+
+// -----
+
+// Check the lowering of ConstantOfShape when:
+//   - The input is not a constant tensor.
+// Expected emitted code:
+//   - Emit code to compute output dimensions from the input's dimensions.
+//   - Krnl iterates are used to set values to the output.
+func @test_constant_of_shape_dynamic_dims(%arg0 : tensor<3xi64>) -> tensor<*xf32> {
+  %0 = "onnx.ConstantOfShape"(%arg0) {value = dense<[1.0]> : tensor<1xf32>} : (tensor<3xi64>) -> tensor<*xf32>
+  "std.return"(%0) : (tensor<*xf32>) -> ()
+
+  // CHECK-LABEL: test_constant_of_shape_dynamic_dims
+  // CHECK: [[CST0:%.+]] = constant 0 : index
+  // CHECK: [[LOAD_DIM_0:%.+]] = affine.load %arg0{{\[}}[[CST0]]{{\]}} : memref<3xi64>
+  // CHECK: [[DIM_0:%.+]] = index_cast [[LOAD_DIM_0]] : i64 to index
+  // CHECK: [[CST1:%.+]] = constant 1 : index
+  // CHECK: [[LOAD_DIM_1:%.+]] = affine.load %arg0{{\[}}[[CST1]]{{\]}} : memref<3xi64>
+  // CHECK: [[DIM_1:%.+]] = index_cast [[LOAD_DIM_1]] : i64 to index
+  // CHECK: [[CST2:%.+]] = constant 2 : index
+  // CHECK: [[LOAD_DIM_2:%.+]] = affine.load %arg0{{\[}}[[CST2]]{{\]}} : memref<3xi64>
+  // CHECK: [[DIM_2:%.+]] = index_cast [[LOAD_DIM_2]] : i64 to index
+  // CHECK: [[RES:%.+]] = alloc([[DIM_0]], [[DIM_1]], [[DIM_2]]) : memref<?x?x?xf32>
+
+  // CHECK: [[CST_VALUE:%.+]] = constant 1.000000e+00 : f32
+  // CHECK: [[LOOP_DEF:%.+]]:3 = krnl.define_loops 3
+  // CHECK: [[CST00:%.+]] = constant 0 : index
+  // CHECK: [[RES_DIM_0:%.+]] = dim [[RES]], [[CST00]] : memref<?x?x?xf32>
+  // CHECK: [[CST11:%.+]] = constant 1 : index
+  // CHECK: [[RES_DIM_1:%.+]] = dim [[RES]], [[CST11]] : memref<?x?x?xf32>
+  // CHECK: [[CST22:%.+]] = constant 2 : index
+  // CHECK: [[RES_DIM_2:%.+]] = dim [[RES]], [[CST22]] : memref<?x?x?xf32>
+  // CHECK: krnl.iterate([[LOOP_DEF]]#0, [[LOOP_DEF]]#1, [[LOOP_DEF]]#2) with ([[LOOP_DEF]]#0 -> %arg1 = 0 to [[RES_DIM_0]], [[LOOP_DEF]]#1 -> %arg2 = 0 to [[RES_DIM_1]], [[LOOP_DEF]]#2 -> %arg3 = 0 to [[RES_DIM_2]]) {
+  // CHECK:   affine.store [[CST_VALUE]], [[RES]][%arg1, %arg2, %arg3] : memref<?x?x?xf32>
+  // CHECK: }
+  // CHECK: return [[RES]] : memref<?x?x?xf32>
+}
+
+// -----
+
+// Check the lowering of ConstantOfShape when:
+//   - The input is a constant tensor.
+// Expected emitted code:
+//   - Output dimensions are computed during compilation time.
+//   - Krnl iterates are used to set values to the output.
+func @test_constant_of_shape_static_dims() -> tensor<*xf32> {
+  %0 = "onnx.Constant"() {value = dense<[3, 4, 5]> : tensor<3xi64> } : () -> tensor<3xi64>
+  %1 = "onnx.ConstantOfShape"(%0) {value = dense<[1.0]> : tensor<1xf32>} : (tensor<3xi64>) -> tensor<*xf32>
+  "std.return"(%1) : (tensor<*xf32>) -> ()
+
+  // CHECK-LABEL: test_constant_of_shape_static_dims
+  // CHECK: [[RES:%.+]] = alloc() : memref<3x4x5xf32>
+  // CHECK: [[GLOBAL_CST:%.+]] = "krnl.global"() {name = "constant_0", shape = [3], value = dense<[3, 4, 5]> : tensor<3xi64>} : () -> memref<3xi64>
+  // CHECK: [[CST_VALUE:%.+]] = constant 1.000000e+00 : f32
+  // CHECK: [[LOOP_DEF:%.+]]:3 = krnl.define_loops 3
+  // CHECK: krnl.iterate([[LOOP_DEF]]#0, [[LOOP_DEF]]#1, [[LOOP_DEF]]#2) with ([[LOOP_DEF]]#0 -> %arg0 = 0 to 3, [[LOOP_DEF]]#1 -> %arg1 = 0 to 4, [[LOOP_DEF]]#2 -> %arg2 = 0 to 5) {
+  // CHECK:   affine.store [[CST_VALUE]], [[RES]][%arg0, %arg1, %arg2] : memref<3x4x5xf32>
+  // CHECK: }
+  // CHECK: return [[RES]] : memref<3x4x5xf32>
+}
+