[GPUToSPIRV] Adding support for vector.constant_mask (#1071)

Leverage upstream pattern for vector.constant_mask lowering to SPIR-V

Author: Garra1980

lib/Conversion/GPUToSPIRV/GPUToSPIRVPass.cpp

@@ -36,6 +36,7 @@
 #include <mlir/Dialect/SPIRV/IR/SPIRVOps.h>
 #include <mlir/Dialect/SPIRV/IR/SPIRVTypes.h>
 #include <mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h>
+#include <mlir/Dialect/Vector/Transforms/LoweringPatterns.h>
 #include <mlir/Dialect/XeGPU/IR/XeGPU.h>
 #include <mlir/IR/BuiltinOps.h>
 #include <mlir/IR/Matchers.h>
@@ -371,6 +372,10 @@ void GPUXToSPIRVPass::runOnOperation() {
     mlir::arith::populateArithToSPIRVPatterns(typeConverter, patterns);
     mlir::populateBuiltinFuncToSPIRVPatterns(typeConverter, patterns);
     mlir::populateVectorToSPIRVPatterns(typeConverter, patterns);
+    // this is for lowering of ConstantMaskOpLowering
+    // but we should also consider replacing VectorMaskConversionPattern
+    // with the upstream one
+    mlir::vector::populateVectorMaskOpLoweringPatterns(patterns);
     mlir::populateMathToSPIRVPatterns(typeConverter, patterns);
     mlir::index::populateIndexToSPIRVPatterns(typeConverter, patterns);
     mlir::populateMemRefToSPIRVPatterns(typeConverter, patterns);
@@ -383,7 +388,6 @@ void GPUXToSPIRVPass::runOnOperation() {
     mlir::populateSCFToSPIRVPatterns(typeConverter, scfToSpirvCtx, patterns);
     mlir::cf::populateControlFlowToSPIRVPatterns(typeConverter, patterns);
     mlir::populateMathToSPIRVPatterns(typeConverter, patterns);
-    // for ub.poison op with vector operand
     imex::populateVectorToSPIRVPatterns(typeConverter, patterns);
 
     if (failed(applyFullConversion(gpuModule, *target, std::move(patterns))))

test/Conversion/GPUToSPIRV/create_mask.mlir

@@ -30,3 +30,36 @@ module attributes {
 // CHECK-NEXT: %[[CAST:.*]] = spirv.SConvert %[[SELECT2]] : i32 to i16
 // CHECK-NEXT: spirv.Bitcast %[[CAST]] : i16 to vector<16xi1>
 // CHECK-NEXT: spirv.Return
+
+// -----
+
+module attributes {
+  gpu.container_module,
+  spirv.target_env = #spirv.target_env<#spirv.vce<v1.0,
+      [Addresses, Float16Buffer, Int64, Int16, Int8, Kernel, Linkage, Vector16, GenericPointer, Groups, Float16, Float64, AtomicFloat32AddEXT, ExpectAssumeKHR],
+      [SPV_EXT_shader_atomic_float_add, SPV_KHR_expect_assume]>, #spirv.resource_limits<>>
+} {
+  gpu.module @kernels {
+    // CHECK-LABEL: spirv.func @constant_mask_0
+    gpu.func @constant_mask_0() kernel attributes {spirv.entry_point_abi = #spirv.entry_point_abi<>} {
+      %0 = vector.constant_mask [0] : vector<16xi1>
+      // CHECK-NEXT: spirv.Constant dense<false> : vector<16xi1>
+      // CHECK-NEXT: spirv.Return
+      gpu.return
+    }
+    // CHECK-LABEL: spirv.func @constant_mask_7
+    gpu.func @constant_mask_7() kernel attributes {spirv.entry_point_abi = #spirv.entry_point_abi<>} {
+      %7 = vector.constant_mask [7] : vector<16xi1>
+      // CHECK-NEXT: spirv.Constant dense<[true, true, true, true, true, true, true, false, false, false, false, false, false, false, false, false]> : vector<16xi1>
+      // CHECK-NEXT: spirv.Return
+      gpu.return
+    }
+    // CHECK-LABEL: spirv.func @constant_mask_16
+    gpu.func @constant_mask_16() kernel attributes {spirv.entry_point_abi = #spirv.entry_point_abi<>} {
+      %16 = vector.constant_mask [16] : vector<16xi1>
+      // CHECK-NEXT: spirv.Constant dense<true> : vector<16xi1>
+      // CHECK-NEXT: spirv.Return
+      gpu.return
+    }
+  }
+}