[mlir][spirv] Add Vector to SPIR-V conversion pass

[ThomasRaoux]

Summary

Add conversion pass for Vector dialect to SPIR-V dialect and add some simple conversion pattern for vector.broadcast, vector.insert, vector.extract.

[antiagainst]

Nice! Thanks Thomas for starting this! I wanted to do this for ages. :)

mlir/include/mlir/Conversion/Passes.td

389

s/SPIRV/SPIR-V/

mlir/include/mlir/Conversion/VectorToSPIRV/ConvertVectorToSPIRVPass.h

10

SPIR-V

mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp

38

Should we check that the input is just a scalar and not some n-D vector? Vector in SPIR-V is constrained; we can generate invalid SPIR-V ops here.

Similarly for the following two patterns.

86

Use empty lines to separate the function and namespace. (Won't clang-format do that actually?)

mlir/test/Conversion/VectorToSPIRV/simple.mlir

4

CHECK-LABEL

16

CHECK-LABEL

[ThomasRaoux]

Address review comments.

[ThomasRaoux]

Thanks Lei. Please take another look.

mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp

38

Good point. I added a check. There was already a check for the other patterns so I don't think I need anything else.

86

Yeah somehow clang-format doesn't seem to add it.

[antiagainst]

mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp

38

Vectors in SPIR-V can only be <= 4 elements by default. If the source is broadcasting a scalar into a, say, 128-element vector, it will create invalid SPIR-V. I think we also need to check the result vector is okay with spirv::CompositeType::isValid and otherwise return false. We can handle the large vector case later with type conversions.

[ThomasRaoux]

Address review comment

[ThomasRaoux]

mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp

38

I added check that the vector result is of size <= 4. Can the composite type not be valid for this case?
(I did the same for the other patterns)

[antiagainst]

mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp

38

Can the composite type not be valid for this case?

spirv::CompositeType::isValid(VectorType) is a static function that wraps around the logic here and checks whether a general VectorType is allowed in SPIR-V: https://github.com/llvm/llvm-project/blob/master/mlir/lib/Dialect/SPIRV/SPIRVTypes.cpp#L174-L177. It can avoid us duplicating the logic multiple times. :)

[ThomasRaoux]

Use spirv::CompositeType::isValid

mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp

38

My bad I should have looked at it. Makes sense, I replaced the checks with this function.

mlir/include/mlir/Conversion/Passes.h

 #include "mlir/Conversion/SCFToStandard/SCFToStandard.h"
 #include "mlir/Conversion/SPIRVToLLVM/ConvertSPIRVToLLVMPass.h"
 #include "mlir/Conversion/ShapeToStandard/ShapeToStandard.h"
 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h"
 #include "mlir/Conversion/StandardToSPIRV/ConvertStandardToSPIRVPass.h"
 #include "mlir/Conversion/VectorToLLVM/ConvertVectorToLLVM.h"
 #include "mlir/Conversion/VectorToROCDL/VectorToROCDL.h"
 #include "mlir/Conversion/VectorToSCF/VectorToSCF.h"
+#include "mlir/Conversion/VectorToSPIRV/ConvertVectorToSPIRVPass.h"
 
 namespace mlir {
 
 /// Generate the code for registering conversion passes.
 #define GEN_PASS_REGISTRATION
 #include "mlir/Conversion/Passes.h.inc"
 
 } // namespace mlir
 
 #endif // MLIR_CONVERSION_PASSES_H

mlir/include/mlir/Conversion/Passes.td

 
 def ConvertVectorToROCDL : Pass<"convert-vector-to-rocdl", "ModuleOp"> {
   let summary = "Lower the operations from the vector dialect into the ROCDL "
                 "dialect";
   let constructor = "mlir::createConvertVectorToROCDLPass()";
   let dependentDialects = ["ROCDL::ROCDLDialect"];
 }
 
+//===----------------------------------------------------------------------===//
+// VectorToSPIRV
+//===----------------------------------------------------------------------===//
+
+def ConvertVectorToSPIRV : Pass<"convert-vector-to-spirv", "ModuleOp"> {
+  let summary = "Lower the operations from the vector dialect into the SPIR-V "
+                "dialect";
+  let constructor = "mlir::createConvertVectorToSPIRVPass()";
+  let dependentDialects = ["spirv::SPIRVDialect"];
+}
+
 #endif // MLIR_CONVERSION_PASSES

mlir/include/mlir/Conversion/VectorToSPIRV/ConvertVectorToSPIRV.h

+//=- ConvertVectorToSPIRV.h - Vector Ops to SPIR-V dialect patterns - C++ -*-=//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Provides patterns for lowering Vector Ops to SPIR-V dialect.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_INCLUDE_MLIR_CONVERSION_VECTORTOSPIRV_CONVERTVECTORTOSPIRV_H_
+#define MLIR_INCLUDE_MLIR_CONVERSION_VECTORTOSPIRV_CONVERTVECTORTOSPIRV_H_
+
+#include "mlir/Transforms/DialectConversion.h"
+
+namespace mlir {
+class SPIRVTypeConverter;
+
+/// Appends to a pattern list additional patterns for translating Vector Ops to
+/// SPIR-V ops.
+void populateVectorToSPIRVPatterns(MLIRContext *context,
+                                   SPIRVTypeConverter &typeConverter,
+                                   OwningRewritePatternList &patterns);
+
+} // namespace mlir
+
+#endif // MLIR_INCLUDE_MLIR_CONVERSION_VECTORTOSPIRV_CONVERTVECTORTOSPIRV_H_

mlir/include/mlir/Conversion/VectorToSPIRV/ConvertVectorToSPIRVPass.h

+//=- ConvertVectorToSPIRVPass.h - Pass converting Vector to SPIRV -*- C++ -*-=//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Provides a pass to convert Vector ops to SPIR-V ops.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_CONVERSION_VECTORTOSPIRV_CONVERTGPUTOSPIRVPASS_H
+#define MLIR_CONVERSION_VECTORTOSPIRV_CONVERTGPUTOSPIRVPASS_H
+
+#include "mlir/Pass/Pass.h"
+
+namespace mlir {
+
+/// Pass to convert Vector Ops to SPIR-V ops.
+std::unique_ptr<OperationPass<ModuleOp>> createConvertVectorToSPIRVPass();
+
+} // namespace mlir
+
+#endif // MLIR_CONVERSION_VECTORTOSPIRV_CONVERTGPUTOSPIRVPASS_H

mlir/include/mlir/Dialect/SPIRV/SPIRVCompositeOps.td

let arguments = (ins
     SPV_Type:$object,
     SPV_Composite:$composite,
     I32ArrayAttr:$indices
   );
 
   let results = (outs
     SPV_Composite:$result
   );
+
+  let builders = [
+    OpBuilder<[{OpBuilder &builder, OperationState &state, Value object,
+                Value composite, ArrayRef<int32_t> indices}]>
+  ];
 }
 
 #endif // SPIRV_COMPOSITE_OPS

mlir/lib/Conversion/CMakeLists.txt

 add_subdirectory(SCFToStandard)
 add_subdirectory(ShapeToStandard)
 add_subdirectory(SPIRVToLLVM)
 add_subdirectory(StandardToLLVM)
 add_subdirectory(StandardToSPIRV)
 add_subdirectory(VectorToROCDL)
 add_subdirectory(VectorToLLVM)
 add_subdirectory(VectorToSCF)
+add_subdirectory(VectorToSPIRV)

mlir/lib/Conversion/VectorToSPIRV/CMakeLists.txt

+add_mlir_conversion_library(MLIRVectorToSPIRV
+  VectorToSPIRV.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Conversion/VectorToSPIRV
+
+  DEPENDS
+  MLIRConversionPassIncGen
+  intrinsics_gen
+
+  LINK_LIBS PUBLIC
+  MLIRSPIRV
+  MLIRVector
+  MLIRTransforms
+  )

mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp

+//===------- VectorToSPIRV.cpp - Vector to SPIRV lowering passes ----------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements a pass to generate SPIRV operations for Vector
+// operations.
+//
+//===----------------------------------------------------------------------===//
+
+#include "../PassDetail.h"
+#include "mlir/Conversion/VectorToSPIRV/ConvertVectorToSPIRV.h"
+#include "mlir/Conversion/VectorToSPIRV/ConvertVectorToSPIRVPass.h"
+#include "mlir/Dialect/SPIRV/SPIRVDialect.h"
+#include "mlir/Dialect/SPIRV/SPIRVLowering.h"
+#include "mlir/Dialect/SPIRV/SPIRVOps.h"
+#include "mlir/Dialect/SPIRV/SPIRVTypes.h"
+#include "mlir/Dialect/Vector/VectorOps.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+using namespace mlir;
+
+namespace {
+struct VectorBroadcastConvert final
+    : public SPIRVOpLowering<vector::BroadcastOp> {
+  using SPIRVOpLowering<vector::BroadcastOp>::SPIRVOpLowering;
+  LogicalResult
+  matchAndRewrite(vector::BroadcastOp broadcastOp, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    if (broadcastOp.source().getType().isa<VectorType>() ||
+        !spirv::CompositeType::isValid(broadcastOp.getVectorType()))
+      return failure();
+    vector::BroadcastOp::Adaptor adaptor(operands);
+    SmallVector<Value, 4> source(broadcastOp.getVectorType().getNumElements(),
+                                 adaptor.source());
+    Value construct = rewriter.create<spirv::CompositeConstructOp>(
+        broadcastOp.getLoc(), broadcastOp.getVectorType(), source);
+    rewriter.replaceOp(broadcastOp, construct);
+    return success();
+  }
+};
+
+struct VectorExtractOpConvert final
+    : public SPIRVOpLowering<vector::ExtractOp> {
+  using SPIRVOpLowering<vector::ExtractOp>::SPIRVOpLowering;
+  LogicalResult
+  matchAndRewrite(vector::ExtractOp extractOp, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    if (extractOp.getType().isa<VectorType>() ||
+        !spirv::CompositeType::isValid(extractOp.getVectorType()))
+      return failure();
+    vector::ExtractOp::Adaptor adaptor(operands);
+    int32_t id = extractOp.position().begin()->cast<IntegerAttr>().getInt();
+    Value newExtract = rewriter.create<spirv::CompositeExtractOp>(
+        extractOp.getLoc(), adaptor.vector(), id);
+    rewriter.replaceOp(extractOp, newExtract);
+    return success();
+  }
+};
+
+struct VectorInsertOpConvert final : public SPIRVOpLowering<vector::InsertOp> {
+  using SPIRVOpLowering<vector::InsertOp>::SPIRVOpLowering;
+  LogicalResult
+  matchAndRewrite(vector::InsertOp insertOp, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    if (insertOp.getSourceType().isa<VectorType>() ||
+        !spirv::CompositeType::isValid(insertOp.getDestVectorType())) 
+      return failure();
+    vector::InsertOp::Adaptor adaptor(operands);
+    int32_t id = insertOp.position().begin()->cast<IntegerAttr>().getInt();
+    Value newInsert = rewriter.create<spirv::CompositeInsertOp>(
+        insertOp.getLoc(), adaptor.source(), adaptor.dest(), id);
+    rewriter.replaceOp(insertOp, newInsert);
+    return success();
+  }
+};
+} // namespace
+
+void mlir::populateVectorToSPIRVPatterns(MLIRContext *context,
+                                         SPIRVTypeConverter &typeConverter,
+                                         OwningRewritePatternList &patterns) {
+  patterns.insert<VectorBroadcastConvert, VectorExtractOpConvert,
+                  VectorInsertOpConvert>(context, typeConverter);
+}
+
+namespace {
+struct LowerVectorToSPIRVPass
+    : public ConvertVectorToSPIRVBase<LowerVectorToSPIRVPass> {
+  void runOnOperation() override;
+};
+} // namespace
+
+void LowerVectorToSPIRVPass::runOnOperation() {
+  MLIRContext *context = &getContext();
+  ModuleOp module = getOperation();
+
+  auto targetAttr = spirv::lookupTargetEnvOrDefault(module);
+  std::unique_ptr<ConversionTarget> target =
+      spirv::SPIRVConversionTarget::get(targetAttr);
+
+  SPIRVTypeConverter typeConverter(targetAttr);
+  OwningRewritePatternList patterns;
+  populateVectorToSPIRVPatterns(context, typeConverter, patterns);
+
+  target->addLegalOp<ModuleOp, ModuleTerminatorOp>();
+  target->addLegalOp<FuncOp>();
+
+  if (failed(applyFullConversion(module, *target, patterns)))
+    return signalPassFailure();
+}
+
+std::unique_ptr<OperationPass<ModuleOp>>
+mlir::createConvertVectorToSPIRVPass() {
+  return std::make_unique<LowerVectorToSPIRVPass>();
+}

mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

static LogicalResult verify(spirv::CompositeExtractOp compExOp) {
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // spv.CompositeInsert
 //===----------------------------------------------------------------------===//
 
+void spirv::CompositeInsertOp::build(OpBuilder &builder, OperationState &state,
+                                     Value object, Value composite,
+                                     ArrayRef<int32_t> indices) {
+  auto indexAttr = builder.getI32ArrayAttr(indices);
+  build(builder, state, composite.getType(), object, composite, indexAttr);
+}
+
 static ParseResult parseCompositeInsertOp(OpAsmParser &parser,
                                           OperationState &state) {
   SmallVector<OpAsmParser::OperandType, 2> operands;
   Type objectType, compositeType;
   Attribute indicesAttr;
   auto loc = parser.getCurrentLocation();
 
   return failure(

mlir/test/Conversion/VectorToSPIRV/simple.mlir

+// RUN: mlir-opt -split-input-file -convert-vector-to-spirv %s -o - | FileCheck %s
+
+// CHECK-LABEL: broadcast
+//  CHECK-SAME: %[[A:.*]]: f32
+//       CHECK:   spv.CompositeConstruct %[[A]], %[[A]], %[[A]], %[[A]] : vector<4xf32>
+//       CHECK:   spv.CompositeConstruct %[[A]], %[[A]] : vector<2xf32>
+func @broadcast(%arg0 : f32) {
+  %0 = vector.broadcast %arg0 : f32 to vector<4xf32>
+  %1 = vector.broadcast %arg0 : f32 to vector<2xf32>
+  spv.Return
+}
+
+// -----
+
+// CHECK-LABEL: extract_insert
+//  CHECK-SAME: %[[V:.*]]: vector<4xf32>
+//       CHECK:   %[[S:.*]] = spv.CompositeExtract %[[V]][1 : i32] : vector<4xf32>
+//       CHECK:   spv.CompositeInsert %[[S]], %[[V]][0 : i32] : f32 into vector<4xf32>
+func @extract_insert(%arg0 : vector<4xf32>) {
+  %0 = vector.extract %arg0[1] : vector<4xf32>
+  %1 = vector.insert %0, %arg0[0] : f32 into vector<4xf32>
+  spv.Return
+}