[Mlir-commits] [mlir] b81bedf - [MLIR][SPIRVToLLVM] Conversion for composite extract and insert

[George Mitenkov]

Author: George Mitenkov
Date: 2020-10-06T11:46:25+03:00
New Revision: b81bedf7146ad5a163e9f1d7283c83ffa4e2043f

URL: https://github.com/llvm/llvm-project/commit/b81bedf7146ad5a163e9f1d7283c83ffa4e2043f
DIFF: https://github.com/llvm/llvm-project/commit/b81bedf7146ad5a163e9f1d7283c83ffa4e2043f.diff

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

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`.

Reviewed By: mravishankar

Differential Revision: https://reviews.llvm.org/D88205

Added: 
    

Modified: 
    mlir/lib/Conversion/SPIRVToLLVM/ConvertSPIRVToLLVM.cpp
    mlir/test/Conversion/SPIRVToLLVM/misc-ops-to-llvm.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Conversion/SPIRVToLLVM/ConvertSPIRVToLLVM.cpp b/mlir/lib/Conversion/SPIRVToLLVM/ConvertSPIRVToLLVM.cpp
index 2c11116d295f..6a5ea1bb06c9 100644
--- a/mlir/lib/Conversion/SPIRVToLLVM/ConvertSPIRVToLLVM.cpp
+++ b/mlir/lib/Conversion/SPIRVToLLVM/ConvertSPIRVToLLVM.cpp
@@ -555,6 +555,66 @@ class BranchConditionalConversionPattern
   }
 };
 
+/// 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>
@@ -1360,6 +1420,7 @@ void mlir::populateSPIRVToLLVMConversionPatterns(
       VariablePattern,
 
       // Miscellaneous ops
+      CompositeExtractPattern, CompositeInsertPattern,
       DirectConversionPattern<spirv::SelectOp, LLVM::SelectOp>,
       DirectConversionPattern<spirv::UndefOp, LLVM::UndefOp>,
 

diff  --git a/mlir/test/Conversion/SPIRVToLLVM/misc-ops-to-llvm.mlir b/mlir/test/Conversion/SPIRVToLLVM/misc-ops-to-llvm.mlir
index e114cea6c96f..700c991463ce 100644
--- a/mlir/test/Conversion/SPIRVToLLVM/misc-ops-to-llvm.mlir
+++ b/mlir/test/Conversion/SPIRVToLLVM/misc-ops-to-llvm.mlir
@@ -1,5 +1,43 @@
 // 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
 //===----------------------------------------------------------------------===//