[mlir][spirv] Add lowering for composite constant.

[denis13]

Summary

Add lowering for constant operation with ranked tensor type to
spv.constant with spv.array type.

[denis13]

Added a missing space in comment.

[merge_guards_bot]

Unit tests: pass. 62009 tests passed, 0 failed and 783 were skipped.

clang-tidy: unknown.

clang-format: fail. Please format your changes with clang-format by running git-clang-format HEAD^ or applying this patch.

Build artifacts: diff.json, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[merge_guards_bot]

Unit tests: pass. 62009 tests passed, 0 failed and 783 were skipped.

clang-tidy: unknown.

clang-format: pass.

Build artifacts: diff.json, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[mravishankar]

mlir/lib/Conversion/StandardToSPIRV/ConvertStandardToSPIRV.cpp

194

Why not use RankedTensorType.

mlir/test/Conversion/StandardToSPIRV/std-to-spirv.mlir

225

Having the constant be a tensor<2x3xi32> while the result is a !spv.array<6 x i32> is weird. To a great extent this is an artifact of the way constants are stored in MLIR, but we should probably linearize the constant to be a tensor<6xi32> to make it less weird.

227

Same as above. We probably need to linearize the constant attribute to a tensor<6xi32>

[denis13]

mlir/test/Conversion/StandardToSPIRV/std-to-spirv.mlir

225

@mravishankar thanks for review, I took those tests from structure-ops.mlir and linearize the type in same way as it works for memref, memref<2x3xf32> -> !spv.ptr<!spv.struct<!spv,array<6xf32>>>. I'll update according your suggestion. Thanks!

[denis13]

Addresses comments.

[denis13]

Update tests.

[merge_guards_bot]

Unit tests: pass. 62028 tests passed, 0 failed and 783 were skipped.

clang-tidy: unknown.

clang-format: pass.

Build artifacts: diff.json, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[merge_guards_bot]

Unit tests: pass. 62028 tests passed, 0 failed and 783 were skipped.

clang-tidy: unknown.

clang-format: pass.

Build artifacts: diff.json, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[denis13]

Remove check for func signature type converter introduced in D72999.

[merge_guards_bot]

Unit tests: pass. 62028 tests passed, 0 failed and 783 were skipped.

clang-tidy: unknown.

clang-format: pass.

Build artifacts: diff.json, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[antiagainst]

mlir/test/Conversion/StandardToSPIRV/std-to-spirv.mlir

225

This does not seem to address Mahesh's comment about linearizing the tensor literal itself? (It should be easy given that the data blob remains the same; we just need to call DenseElementsAttr::reshape() on it.) We'll need a test case for n-D tensors. :)

[denis13]

@antiagainst @mravishankar oops, sorry for the misunderstanding

[denis13]

Linearize DenseElementsAttr.

[merge_guards_bot]

Unit tests: pass. 62028 tests passed, 0 failed and 783 were skipped.

clang-tidy: unknown.

clang-format: pass.

Build artifacts: diff.json, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[antiagainst]

Nice, thanks Denis for adding this!

mlir/lib/Conversion/StandardToSPIRV/ConvertStandardToSPIRV.cpp

 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // Operation conversion
 //===----------------------------------------------------------------------===//
 
 namespace {
 
+/// Convert composite constant operation to SPIR-V dialect.
+// TODO(denis0x0D) : move to DRR.
+class ConstantCompositeOpConversion final : public SPIRVOpLowering<ConstantOp> {
+public:
+  using SPIRVOpLowering<ConstantOp>::SPIRVOpLowering;
+
+  PatternMatchResult
+  matchAndRewrite(ConstantOp constCompositeOp, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override;
+};
+
 /// Convert constant operation with IndexType return to SPIR-V constant
 /// operation. Since IndexType is not used within SPIR-V dialect, this needs
 /// special handling to make sure the result type and the type of the value
 /// attribute are consistent.
 // TODO(ravishankarm) : This should be moved into DRR.
 class ConstantIndexOpConversion final : public SPIRVOpLowering<ConstantOp> {
 public:
   using SPIRVOpLowering<ConstantOp>::SPIRVOpLowering;
static spirv::AccessChainOp getElementPtr(OpBuilder &builder,
   // Add a '0' at the start to index into the struct.
   linearizedIndices.push_back(builder.create<spirv::ConstantOp>(
       loc, indexType, IntegerAttr::get(indexType, 0)));
   linearizedIndices.push_back(ptrLoc);
   return builder.create<spirv::AccessChainOp>(loc, basePtr, linearizedIndices);
 }
 
 //===----------------------------------------------------------------------===//
+// ConstantOp with composite type.
+//===----------------------------------------------------------------------===//
+
+PatternMatchResult ConstantCompositeOpConversion::matchAndRewrite(
+    ConstantOp constCompositeOp, ArrayRef<Value> operands,
+    ConversionPatternRewriter &rewriter) const {
+  auto compositeType =
+      constCompositeOp.getResult().getType().dyn_cast<RankedTensorType>();
+  if (!compositeType) {
+    return matchFailure();
+  }
+  auto spirvCompositeType = typeConverter.convertType(compositeType);
+  if (!spirvCompositeType) {
+    return matchFailure();
+  }
+  rewriter.replaceOpWithNewOp<spirv::ConstantOp>(
+      constCompositeOp, spirvCompositeType, constCompositeOp.value());
+  return matchSuccess();
+}
+
+//===----------------------------------------------------------------------===//
 // ConstantOp with index type.
 //===----------------------------------------------------------------------===//
 
 PatternMatchResult ConstantIndexOpConversion::matchAndRewrite(
     ConstantOp constIndexOp, ArrayRef<Value> operands,
     ConversionPatternRewriter &rewriter) const {
   if (!constIndexOp.getResult().getType().isa<IndexType>()) {
     return matchFailure();

 } // namespace
 
 namespace mlir {
 void populateStandardToSPIRVPatterns(MLIRContext *context,
                                      SPIRVTypeConverter &typeConverter,
                                      OwningRewritePatternList &patterns) {
   // Add patterns that lower operations into SPIR-V dialect.
   populateWithGenerated(context, &patterns);
-  patterns.insert<ConstantIndexOpConversion, CmpFOpConversion, CmpIOpConversion,
+  patterns.insert<ConstantCompositeOpConversion, ConstantIndexOpConversion,
+                  CmpFOpConversion, CmpIOpConversion,
                   IntegerOpConversion<AddIOp, spirv::IAddOp>,
                   IntegerOpConversion<MulIOp, spirv::IMulOp>,
                   IntegerOpConversion<SignedDivIOp, spirv::SDivOp>,
                   IntegerOpConversion<SignedRemIOp, spirv::SModOp>,
                   IntegerOpConversion<SubIOp, spirv::ISubOp>, LoadOpConversion,
                   ReturnOpConversion, SelectOpConversion, StoreOpConversion>(
       context, typeConverter);
 }
 } // namespace mlir

mlir/lib/Conversion/StandardToSPIRV/ConvertStandardToSPIRVPass.cpp

FuncOpConversion::matchAndRewrite(FuncOp funcOp, ArrayRef<Value> operands,
   if (fnType.getNumResults()) {
     return matchFailure();
   }
 
   TypeConverter::SignatureConversion signatureConverter(fnType.getNumInputs());
   {
     for (auto argType : enumerate(funcOp.getType().getInputs())) {
       auto convertedType = typeConverter.convertType(argType.value());
+      if (!convertedType) {
+        return matchFailure();
+      }
       signatureConverter.addInputs(argType.index(), convertedType);
     }
   }
 
   rewriter.updateRootInPlace(funcOp, [&] {
     funcOp.setType(rewriter.getFunctionType(
         signatureConverter.getConvertedTypes(), llvm::None));
     rewriter.applySignatureConversion(&funcOp.getBody(), signatureConverter);

mlir/lib/Dialect/SPIRV/SPIRVLowering.cpp

if (llvm::is_contained(dims, ShapedType::kDynamicSize) ||
         llvm::is_contained(strides, MemRefType::getDynamicStrideOrOffset())) {
       return llvm::None;
     }
     int64_t memrefSize = -1;
     for (auto shape : enumerate(dims)) {
       memrefSize = std::max(memrefSize, shape.value() * strides[shape.index()]);
     }
     return (offset + memrefSize) * elementSize.getValue();
+  } else if (auto tensorType = t.dyn_cast<TensorType>()) {
+    if (!tensorType.hasStaticShape()) {
+      return llvm::None;
+    }
+    auto elementSize = getTypeNumBytes(tensorType.getElementType());
+    if (!elementSize) {
+      return llvm::None;
+    }
+    int64_t size = elementSize.getValue();
+    for (auto shape : tensorType.getShape()) {
+      size *= shape;
+    }
+    return size;
   }
   // TODO: Add size computation for other types.
   return llvm::None;
 }
 
 static Type convertStdType(Type type) {
   // If the type is already valid in SPIR-V, directly return.
   if (spirv::SPIRVDialect::isValidType(type)) {
if (memRefType.hasStaticShape()) {
       auto structType = spirv::StructType::get(arrayType, 0);
       // For now initialize the storage class to StorageBuffer. This will be
       // updated later based on whats passed in w.r.t to the ABI attributes.
       return spirv::PointerType::get(structType,
                                      spirv::StorageClass::StorageBuffer);
     }
   }
 
+  if (auto tensorType = type.dyn_cast<TensorType>()) {
+    // TODO(ravishankarm) : Handle dynamic shapes.
+    if (!tensorType.hasStaticShape()) {
+      return Type();
+    }
+    auto elementType = convertStdType(tensorType.getElementType());
+    if (!elementType) {
+      return Type();
+    }
+    auto elementSize = getTypeNumBytes(elementType);
+    if (!elementSize) {
+      return Type();
+    }
+    auto tensorSize = getTypeNumBytes(tensorType);
+    if (!tensorSize) {
+      return Type();
+    }
+    return spirv::ArrayType::get(elementType,
+                                 tensorSize.getValue() / elementSize.getValue(),
+                                 elementSize.getValue());
+  }
   return Type();
 }
 
 Type SPIRVTypeConverter::convertType(Type type) { return convertStdType(type); }
 
 //===----------------------------------------------------------------------===//
 // Builtin Variables
 //===----------------------------------------------------------------------===//

mlir/test/Conversion/StandardToSPIRV/std-to-spirv.mlir

func @constant() {
   // CHECK: spv.constant 42 : i64
   %1 = constant 42
   // CHECK: spv.constant {{[0-9]*\.[0-9]*e?-?[0-9]*}} : f32
   %2 = constant 0.5 : f32
   // CHECK: spv.constant dense<[2, 3]> : vector<2xi32>
   %3 = constant dense<[2, 3]> : vector<2xi32>
   // CHECK: spv.constant 1 : i32
   %4 = constant 1 : index
+  // CHECK: spv.constant dense<1> : tensor<6xi32> : !spv.array<6 x i32 [4]>
+  %5 = constant dense<1> : tensor<6xi32>
+  // CHECK: spv.constant dense<1.000000e+00> : tensor<6xf32> : !spv.array<6 x f32 [4]>
+  %6 = constant dense<1.0> : tensor<6xf32>
+  // CHECK: spv.constant dense<{{\[}}1, 2, 3, 4, 5, 6]> : tensor<6xi32> : !spv.array<6 x i32 [4]>
+  %7 = constant dense<[1, 2, 3, 4, 5, 6]> : tensor<6xi32>
+  // CHECK: spv.constant dense<{{\[}}1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00, 5.000000e+00, 6.000000e+00]> : tensor<6xf32> : !spv.array<6 x f32 [4]>
+  %8 = constant dense<[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]> : tensor<6xf32>
   return
 }
 
 //===----------------------------------------------------------------------===//
 // std logical binary operations
 //===----------------------------------------------------------------------===//
 
 // CHECK-LABEL: @logical_scalar