[MLIR][SPIRVToLLVM] Conversion for composite extract and insert

[georgemitenkov]

Summary

A pattern to convert spv.CompositeInsert and spv.CompositeExtract.
In LLVM, there are 2 ops that correspond to each instruction depending
on the container type. If the container type is a vector type, then the result of
conversion is llvm.insertelement or llvm.extractelement. If the container
type is an aggregate type (i.e. struct, array), the result of conversion
is llvm.insertvalue or llvm.extractvalue.

[mravishankar]

mlir/lib/Conversion/SPIRVToLLVM/ConvertSPIRVToLLVM.cpp

242

Can we split this? THere doesnt seem to much overlap between extract and insert and separate methods might better.

251

If you do a dyn_cast, the value can be null. If you are not checking for the null value, you can just use cast<..> here.

260

Please split these as well. Actually I really dont see a need for these helper functions. You could just do this at the call site itself.

267

Same as above. Just use cast<..> if you are not checking for null value.

563

Similar to other comments above. Might just be cleaner to split this into two patterns. One for insert and other for extract. There is not much to share here.

[georgemitenkov]

Splitted the pattern into two and removed unnecessary helper functions.

mlir/lib/Conversion/SPIRVToLLVM/ConvertSPIRVToLLVM.cpp

static LogicalResult replaceWithLoadOrStore(Operation *op,
   rewriter.replaceOpWithNewOp<LLVM::StoreOp>(storeOp, storeOp.value(),
                                              storeOp.ptr(), alignment,
                                              isVolatile, isNonTemporal);
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // Type conversion
 //===----------------------------------------------------------------------===//
 
 /// Converts SPIR-V array type to LLVM array. Natural stride (according to
 /// `VulkanLayoutUtils`) is also mapped to LLVM array. This has to be respected
 /// when converting ops that manipulate array types.
 static Optional<Type> convertArrayType(spirv::ArrayType type,
                                        TypeConverter &converter) {
   unsigned stride = type.getArrayStride();
   Type elementType = type.getElementType();
   auto sizeInBytes = elementType.cast<spirv::SPIRVType>().getSizeInBytes();
   if (stride != 0 &&
       !(sizeInBytes.hasValue() && sizeInBytes.getValue() == stride))
     return llvm::None;
 
   auto llvmElementType =
       converter.convertType(elementType).cast<LLVM::LLVMType>();
   unsigned numElements = type.getNumElements();
   return LLVM::LLVMType::getArrayTy(llvmElementType, numElements);
 }
 
 /// Converts SPIR-V pointer type to LLVM pointer. Pointer's storage class is not
 /// modelled at the moment.
 static Type convertPointerType(spirv::PointerType type,
                                TypeConverter &converter) {
   auto pointeeType =
       converter.convertType(type.getPointeeType()).cast<LLVM::LLVMType>();
   return pointeeType.getPointerTo();
 }
 
 /// Converts SPIR-V runtime array to LLVM array. Since LLVM allows indexing over
 /// the bounds, the runtime array is converted to a 0-sized LLVM array. There is
 /// no modelling of array stride at the moment.
 static Optional<Type> convertRuntimeArrayType(spirv::RuntimeArrayType type,
                                               TypeConverter &converter) {
matchAndRewrite(spirv::BranchConditionalOp op, ArrayRef<Value> operands,
     rewriter.replaceOpWithNewOp<LLVM::CondBrOp>(
         op, op.condition(), op.getTrueBlockArguments(),
         op.getFalseBlockArguments(), branchWeights, op.getTrueBlock(),
         op.getFalseBlock());
     return success();
   }
 };
 
+/// Converts `spv.CompositeExtract` to `llvm.extractvalue` if the container type
+/// is an aggregate type (struct or array). Otherwise, converts to
+/// `llvm.extractelement` that operates on vectors.
+class CompositeExtractPattern
+    : public SPIRVToLLVMConversion<spirv::CompositeExtractOp> {
+public:
+  using SPIRVToLLVMConversion<spirv::CompositeExtractOp>::SPIRVToLLVMConversion;
+
+  LogicalResult
+  matchAndRewrite(spirv::CompositeExtractOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto dstType = this->typeConverter.convertType(op.getType());
+    if (!dstType)
+      return failure();
+
+    Type containerType = op.composite().getType();
+    if (containerType.isa<VectorType>()) {
+      Location loc = op.getLoc();
+      IntegerAttr value = op.indices()[0].cast<IntegerAttr>();
+      Value index = createI32ConstantOf(loc, rewriter, value.getInt());
+      rewriter.replaceOpWithNewOp<LLVM::ExtractElementOp>(
+          op, dstType, op.composite(), index);
+      return success();
+    }
+    rewriter.replaceOpWithNewOp<LLVM::ExtractValueOp>(
+        op, dstType, op.composite(), op.indices());
+    return success();
+  }
+};
+
+/// Converts `spv.CompositeInsert` to `llvm.insertvalue` if the container type
+/// is an aggregate type (struct or array). Otherwise, converts to
+/// `llvm.insertelement` that operates on vectors.
+class CompositeInsertPattern
+    : public SPIRVToLLVMConversion<spirv::CompositeInsertOp> {
+public:
+  using SPIRVToLLVMConversion<spirv::CompositeInsertOp>::SPIRVToLLVMConversion;
+
+  LogicalResult
+  matchAndRewrite(spirv::CompositeInsertOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto dstType = this->typeConverter.convertType(op.getType());
+    if (!dstType)
+      return failure();
+
+    Type containerType = op.composite().getType();
+    if (containerType.isa<VectorType>()) {
+      Location loc = op.getLoc();
+      IntegerAttr value = op.indices()[0].cast<IntegerAttr>();
+      Value index = createI32ConstantOf(loc, rewriter, value.getInt());
+      rewriter.replaceOpWithNewOp<LLVM::InsertElementOp>(
+          op, dstType, op.composite(), op.object(), index);
+      return success();
+    }
+    rewriter.replaceOpWithNewOp<LLVM::InsertValueOp>(
+        op, dstType, op.composite(), op.object(), op.indices());
+    return success();
+  }
+};
+
 /// Converts SPIR-V operations that have straightforward LLVM equivalent
 /// into LLVM dialect operations.
 template <typename SPIRVOp, typename LLVMOp>
 class DirectConversionPattern : public SPIRVToLLVMConversion<SPIRVOp> {
 public:
   using SPIRVToLLVMConversion<SPIRVOp>::SPIRVToLLVMConversion;
 
   LogicalResult
patterns.insert<
       NotPattern<spirv::LogicalNotOp>,
 
       // Memory ops
       AccessChainPattern, AddressOfPattern, GlobalVariablePattern,
       LoadStorePattern<spirv::LoadOp>, LoadStorePattern<spirv::StoreOp>,
       VariablePattern,
 
       // Miscellaneous ops
+      CompositeExtractPattern, CompositeInsertPattern,
       DirectConversionPattern<spirv::SelectOp, LLVM::SelectOp>,
       DirectConversionPattern<spirv::UndefOp, LLVM::UndefOp>,
 
       // Shift ops
       ShiftPattern<spirv::ShiftRightArithmeticOp, LLVM::AShrOp>,
       ShiftPattern<spirv::ShiftRightLogicalOp, LLVM::LShrOp>,
       ShiftPattern<spirv::ShiftLeftLogicalOp, LLVM::ShlOp>,
 

mlir/test/Conversion/SPIRVToLLVM/misc-ops-to-llvm.mlir

 // RUN: mlir-opt -convert-spirv-to-llvm %s | FileCheck %s
 
 //===----------------------------------------------------------------------===//
+// spv.CompositeExtract
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: @composite_extract_array
+spv.func @composite_extract_array(%arg: !spv.array<4x!spv.array<4xf32>>) "None" {
+  // CHECK: llvm.extractvalue %{{.*}}[1 : i32, 3 : i32] : !llvm.array<4 x array<4 x float>>
+  %0 = spv.CompositeExtract %arg[1 : i32, 3 : i32] : !spv.array<4x!spv.array<4xf32>>
+  spv.Return
+}
+
+// CHECK-LABEL: @composite_extract_vector
+spv.func @composite_extract_vector(%arg: vector<3xf32>) "None" {
+  // CHECK: %[[ZERO:.*]] = llvm.mlir.constant(0 : i32) : !llvm.i32
+  // CHECK: llvm.extractelement %{{.*}}[%[[ZERO]] : !llvm.i32] : !llvm.vec<3 x float>
+  %0 = spv.CompositeExtract %arg[0 : i32] : vector<3xf32>
+  spv.Return
+}
+
+//===----------------------------------------------------------------------===//
+// spv.CompositeInsert
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: @composite_insert_struct
+spv.func @composite_insert_struct(%arg0: i32, %arg1: !spv.struct<f32, !spv.array<4xi32>>) "None" {
+  // CHECK: llvm.insertvalue %{{.*}}, %{{.*}}[1 : i32, 3 : i32] : !llvm.struct<packed (float, array<4 x i32>)>
+  %0 = spv.CompositeInsert %arg0, %arg1[1 : i32, 3 : i32] : i32 into !spv.struct<f32, !spv.array<4xi32>>
+  spv.Return
+}
+
+// CHECK-LABEL: @composite_insert_vector
+spv.func @composite_insert_vector(%arg0: vector<3xf32>, %arg1: f32) "None" {
+  // CHECK: %[[ONE:.*]] = llvm.mlir.constant(1 : i32) : !llvm.i32
+  // CHECK: llvm.insertelement %{{.*}}, %{{.*}}[%[[ONE]] : !llvm.i32] : !llvm.vec<3 x float>
+  %0 = spv.CompositeInsert %arg1, %arg0[1 : i32] : f32 into vector<3xf32>
+  spv.Return
+}
+
+//===----------------------------------------------------------------------===//
 // spv.Select
 //===----------------------------------------------------------------------===//
 
 // CHECK-LABEL: @select_scalar
 spv.func @select_scalar(%arg0: i1, %arg1: vector<3xi32>, %arg2: f32) "None" {
   // CHECK: llvm.select %{{.*}}, %{{.*}}, %{{.*}} : !llvm.i1, !llvm.vec<3 x i32>
   %0 = spv.Select %arg0, %arg1, %arg1 : i1, vector<3xi32>
   // CHECK: llvm.select %{{.*}}, %{{.*}}, %{{.*}} : !llvm.i1, !llvm.float