[mlir][SPIRVToLLVM] Add shufflevector conversion for static index

[Weiwei-2021]

Summary

Add the shufflevector conversion. It only handles the static, i.e., IntegerAttr, index.

I am not sure whether it is ok to create multiple constantOps with the same value in llvm dialect, therefore I add a utility to find existing constant Op for a specific value. Correct me if we don't need to do it.

Co-authored: Xinyi Liu <[e-mail]>

[antiagainst]

mlir/lib/Conversion/SPIRVToLLVM/SPIRVToLLVM.cpp

1376

Can use OpAdaptor instead of ArrayRef<Value> so that you can access vector1() and vector2() using methods.

1380

You should use vector1 and vector2 from operands. Dialect conversion works by going through each op in a block sequentially and covert them one by one. The ops generating vector1 and vector2 might be converted previously (and with the converted vector1 and vector2 stored in operands) so we need to use the converted version for them.

1387

Can return early here. We typically prefer early return in MLIR.

1390

This does not need to be converted again?

1393

This map should not be needed. You can create duplicated constants; that's quite common in MLIR.

1396

For unsupported cases, we should fail the pattern match. Assertion should not be used here.

[Weiwei-2021]

mlir/lib/Conversion/SPIRVToLLVM/SPIRVToLLVM.cpp

1376

Cool. Do not know that the OpAdaptor can be used in this way. Thanks.

1393

Got it. Thanks.

[antiagainst]

Nice, just two small issues. Feel free to land after fixing them!

mlir/lib/Conversion/SPIRVToLLVM/SPIRVToLLVM.cpp

1388

Given we have an early return in the above, we can remove the { .. } for the else and save one level of indent. :) That's the typical coding style in MLIR.

1411

Please use the input op's location. Unknown location will make print/debug hard.

[Weiwei-2021]

Thank you for your comment. Fixed and will land soon!

mlir/lib/Conversion/SPIRVToLLVM/SPIRVToLLVM.cpp

1411

good to know it. I see other places are using unknownLoc, thought unknown is expected for new Ops like constant.

mlir/lib/Conversion/SPIRVToLLVM/SPIRVToLLVM.cpp

matchAndRewrite(spirv::ModuleOp spvModuleOp, OpAdaptor adaptor,
 
     // Remove the terminator block that was automatically added by builder
     rewriter.eraseBlock(&newModuleOp.getBodyRegion().back());
     rewriter.eraseOp(spvModuleOp);
     return success();
   }
 };
 
+//===----------------------------------------------------------------------===//
+// VectorShuffleOp conversion
+//===----------------------------------------------------------------------===//
+
+class VectorShufflePattern
+    : public SPIRVToLLVMConversion<spirv::VectorShuffleOp> {
+public:
+  using SPIRVToLLVMConversion<spirv::VectorShuffleOp>::SPIRVToLLVMConversion;
+  LogicalResult
+  matchAndRewrite(spirv::VectorShuffleOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    Location loc = op.getLoc();
+    auto components = adaptor.components();
+    auto vector1 = adaptor.vector1();
+    auto vector2 = adaptor.vector2();
+    int vector1Size = vector1.getType().cast<VectorType>().getNumElements();
+    int vector2Size = vector2.getType().cast<VectorType>().getNumElements();
+    if (vector1Size == vector2Size) {
+      rewriter.replaceOpWithNewOp<LLVM::ShuffleVectorOp>(op, vector1, vector2,
+                                                         components);
+      return success();
+    }
+
+    auto dstType = typeConverter.convertType(op.getType());
+    auto scalarType = dstType.cast<VectorType>().getElementType();
+    auto componentsArray = components.getValue();
+    auto context = rewriter.getContext();
+    auto llvmI32Type = IntegerType::get(context, 32);
+    Value targetOp = rewriter.create<LLVM::UndefOp>(loc, dstType);
+    for (unsigned i = 0; i < componentsArray.size(); i++) {
+      if (componentsArray[i].isa<IntegerAttr>())
+        op.emitError("unable to support non-constant component");
+
+      int indexVal = componentsArray[i].cast<IntegerAttr>().getInt();
+      if (indexVal == -1)
+        continue;
+
+      int offsetVal = 0;
+      Value baseVector = vector1;
+      if (indexVal >= vector1Size) {
+        offsetVal = vector1Size;
+        baseVector = vector2;
+      }
+
+      Value dstIndex = rewriter.create<LLVM::ConstantOp>(
+          loc, llvmI32Type, rewriter.getIntegerAttr(rewriter.getI32Type(), i));
+      Value index = rewriter.create<LLVM::ConstantOp>(
+          loc, llvmI32Type,
+          rewriter.getIntegerAttr(rewriter.getI32Type(), indexVal - offsetVal));
+
+      auto extractOp = rewriter.create<LLVM::ExtractElementOp>(
+          loc, scalarType, baseVector, index);
+      targetOp = rewriter.create<LLVM::InsertElementOp>(loc, dstType, targetOp,
+                                                        extractOp, dstIndex);
+    }
+    rewriter.replaceOp(op, targetOp);
+    return success();
+  }
+};
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // Pattern population
 //===----------------------------------------------------------------------===//
 
 void mlir::populateSPIRVToLLVMTypeConversion(LLVMTypeConverter &typeConverter) {
   typeConverter.addConversion([&](spirv::ArrayType type) {
patterns.add<
       AccessChainPattern, AddressOfPattern, GlobalVariablePattern,
       LoadStorePattern<spirv::LoadOp>, LoadStorePattern<spirv::StoreOp>,
       VariablePattern,
 
       // Miscellaneous ops
       CompositeExtractPattern, CompositeInsertPattern,
       DirectConversionPattern<spirv::SelectOp, LLVM::SelectOp>,
       DirectConversionPattern<spirv::UndefOp, LLVM::UndefOp>,
+      VectorShufflePattern,
 
       // Shift ops
       ShiftPattern<spirv::ShiftRightArithmeticOp, LLVM::AShrOp>,
       ShiftPattern<spirv::ShiftRightLogicalOp, LLVM::LShrOp>,
       ShiftPattern<spirv::ShiftLeftLogicalOp, LLVM::ShlOp>,
 
       // Return ops
       ReturnPattern, ReturnValuePattern>(patterns.getContext(), typeConverter);

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

 // CHECK-LABEL: @select_vector
 spv.func @select_vector(%arg0: vector<2xi1>, %arg1: vector<2xi32>) "None" {
   // CHECK: llvm.select %{{.*}}, %{{.*}}, %{{.*}} : vector<2xi1>, vector<2xi32>
   %0 = spv.Select %arg0, %arg1, %arg1 : vector<2xi1>, vector<2xi32>
   spv.Return
 }
 
 //===----------------------------------------------------------------------===//
+// spv.VectorShuffle
+//===----------------------------------------------------------------------===//
+
+spv.func @vector_shuffle_same_size(%vector1: vector<2xf32>, %vector2: vector<2xf32>) -> vector<3xf32> "None" {
+  //      CHECK: %[[res:.*]] = llvm.shufflevector {{.*}} [0 : i32, 2 : i32, -1 : i32] : vector<2xf32>, vector<2xf32>
+  // CHECK-NEXT: return %[[res]] : vector<3xf32>
+  %0 = spv.VectorShuffle [0: i32, 2: i32, 0xffffffff: i32] %vector1: vector<2xf32>, %vector2: vector<2xf32> -> vector<3xf32>
+  spv.ReturnValue %0: vector<3xf32>
+}
+
+spv.func @vector_shuffle_different_size(%vector1: vector<3xf32>, %vector2: vector<2xf32>) -> vector<3xf32> "None" {
+  //      CHECK: %[[UNDEF:.*]] = llvm.mlir.undef : vector<3xf32>
+  // CHECK-NEXT: %[[C0_0:.*]] = llvm.mlir.constant(0 : i32) : i32
+  // CHECK-NEXT: %[[C0_1:.*]] = llvm.mlir.constant(0 : i32) : i32
+  // CHECK-NEXT: %[[EXT0:.*]] = llvm.extractelement %arg0[%[[C0_1]] : i32] : vector<3xf32>
+  // CHECK-NEXT: %[[INSERT0:.*]] = llvm.insertelement %[[EXT0]], %[[UNDEF]][%[[C0_0]] : i32] : vector<3xf32>
+  // CHECK-NEXT: %[[C1_0:.*]] = llvm.mlir.constant(1 : i32) : i32
+  // CHECK-NEXT: %[[C1_1:.*]] = llvm.mlir.constant(1 : i32) : i32
+  // CHECK-NEXT: %[[EXT1:.*]] = llvm.extractelement {{.*}}[%[[C1_1]] : i32] : vector<2xf32>
+  // CHECK-NEXT: %[[RES:.*]] = llvm.insertelement %[[EXT1]], %[[INSERT0]][%[[C1_0]] : i32] : vector<3xf32>
+  // CHECK-NEXT: llvm.return %[[RES]] : vector<3xf32>
+  %0 = spv.VectorShuffle [0: i32, 4: i32, 0xffffffff: i32] %vector1: vector<3xf32>, %vector2: vector<2xf32> -> vector<3xf32>
+  spv.ReturnValue %0: vector<3xf32>
+}
+
+//===----------------------------------------------------------------------===//
 // spv.EntryPoint and spv.ExecutionMode
 //===----------------------------------------------------------------------===//
 
 //      CHECK: module {
 // CHECK-NEXT:   llvm.mlir.global external constant @{{.*}}() : !llvm.struct<(i32)> {
 // CHECK-NEXT:     %[[UNDEF:.*]] = llvm.mlir.undef : !llvm.struct<(i32)>
 // CHECK-NEXT:     %[[VAL:.*]] = llvm.mlir.constant(31 : i32) : i32
 // CHECK-NEXT:     %[[RET:.*]] = llvm.insertvalue %[[VAL]], %[[UNDEF]][0 : i32] : !llvm.struct<(i32)>