[Mlir-commits] [mlir] [mlir][spirv] Add integration test for `vector.interleave` and `vector.shuffle` (PR #93595)

Angel Zhang llvmlistbot at llvm.org
Tue May 28 18:57:24 PDT 2024


https://github.com/angelz913 updated https://github.com/llvm/llvm-project/pull/93595

>From af1b3391651a1ac01589dc1c1e15f208f7ef201f Mon Sep 17 00:00:00 2001
From: Angel Zhang <angel.zhang at amd.com>
Date: Thu, 23 May 2024 20:08:55 +0000
Subject: [PATCH 1/4] Add integration test for vector.interleave

---
 .../mlir-vulkan-runner/vector-interleave.mlir | 45 +++++++++++++++++++
 1 file changed, 45 insertions(+)
 create mode 100644 mlir/test/mlir-vulkan-runner/vector-interleave.mlir

diff --git a/mlir/test/mlir-vulkan-runner/vector-interleave.mlir b/mlir/test/mlir-vulkan-runner/vector-interleave.mlir
new file mode 100644
index 0000000000000..08f9c5a0bf164
--- /dev/null
+++ b/mlir/test/mlir-vulkan-runner/vector-interleave.mlir
@@ -0,0 +1,45 @@
+// RUN: mlir-vulkan-runner %s \
+// RUN:  --shared-libs=%vulkan-runtime-wrappers,%mlir_runner_utils \
+// RUN:  --entry-point-result=void | FileCheck %s
+
+// CHECK: [0, 2, 1, 3]
+module attributes {
+  gpu.container_module,
+  spirv.target_env = #spirv.target_env<
+    #spirv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>, #spirv.resource_limits<>>
+} {
+  gpu.module @kernels {
+    gpu.func @kernel_vector_interleave(%arg0 : vector<2xi32>, %arg1 : vector<2xi32>, %arg2 : memref<4xi32>)
+      kernel attributes { spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [1, 1, 1]>} {
+      %c0 = arith.constant 0 : index
+      %result = vector.interleave %arg0, %arg1 : vector<2xi32>
+      vector.store %result, %arg2[%c0] : memref<4xi32>, vector<4xi32>
+      gpu.return
+    }
+  }
+
+  func.func @main() {
+    // Allocate 3 buffers.
+    %buf0 = arith.constant dense<[0, 1]> : vector<2xi32>
+    %buf1 = arith.constant dense<[2, 3]> : vector<2xi32>
+    %buf2 = memref.alloc() : memref<4xi32>
+    
+    %idx0 = arith.constant 0 : index
+    %idx1 = arith.constant 1 : index
+    %idx4 = arith.constant 4 : index
+
+    // Initialize output buffer.
+    %value0 = arith.constant 0 : i32
+    %buf3 = memref.cast %buf2 : memref<4xi32> to memref<?xi32>
+    call @fillResource1DInt(%buf3, %value0) : (memref<?xi32>, i32) -> ()
+
+    gpu.launch_func @kernels::@kernel_vector_interleave
+        blocks in (%idx4, %idx1, %idx1) threads in (%idx1, %idx1, %idx1)
+        args(%buf0 : vector<2xi32>, %buf1 : vector<2xi32>, %buf2 : memref<4xi32>)
+    %buf4 = memref.cast %buf3 : memref<?xi32> to memref<*xi32>
+    call @printMemrefI32(%buf4) : (memref<*xi32>) -> ()
+    return
+  }
+  func.func private @fillResource1DInt(%0 : memref<?xi32>, %1 : i32)
+  func.func private @printMemrefI32(%ptr : memref<*xi32>)
+}

>From 7b52bd780eacb74854fde327d4b795f330ce806f Mon Sep 17 00:00:00 2001
From: Angel Zhang <angel.zhang at amd.com>
Date: Tue, 28 May 2024 18:36:30 +0000
Subject: [PATCH 2/4] Add VectorToSPIRV patterns to GPUToSPIRVPass, and fix
 errors in e2e test

---
 .../Conversion/GPUToSPIRV/GPUToSPIRVPass.cpp   |  2 ++
 .../mlir-vulkan-runner/vector-interleave.mlir  | 18 +++++++++++++-----
 2 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRVPass.cpp b/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRVPass.cpp
index 1d1db913e3df2..2677d4e24be2c 100644
--- a/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRVPass.cpp
+++ b/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRVPass.cpp
@@ -16,6 +16,7 @@
 #include "mlir/Conversion/ArithToSPIRV/ArithToSPIRV.h"
 #include "mlir/Conversion/FuncToSPIRV/FuncToSPIRV.h"
 #include "mlir/Conversion/GPUToSPIRV/GPUToSPIRV.h"
+#include "mlir/Conversion/VectorToSPIRV/VectorToSPIRV.h"
 #include "mlir/Conversion/MemRefToSPIRV/MemRefToSPIRV.h"
 #include "mlir/Conversion/SCFToSPIRV/SCFToSPIRV.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
@@ -132,6 +133,7 @@ void GPUToSPIRVPass::runOnOperation() {
     mlir::arith::populateArithToSPIRVPatterns(typeConverter, patterns);
     populateMemRefToSPIRVPatterns(typeConverter, patterns);
     populateFuncToSPIRVPatterns(typeConverter, patterns);
+    populateVectorToSPIRVPatterns(typeConverter, patterns);
 
     if (failed(applyFullConversion(gpuModule, *target, std::move(patterns))))
       return signalPassFailure();
diff --git a/mlir/test/mlir-vulkan-runner/vector-interleave.mlir b/mlir/test/mlir-vulkan-runner/vector-interleave.mlir
index 08f9c5a0bf164..d760c1631c473 100644
--- a/mlir/test/mlir-vulkan-runner/vector-interleave.mlir
+++ b/mlir/test/mlir-vulkan-runner/vector-interleave.mlir
@@ -9,10 +9,12 @@ module attributes {
     #spirv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>, #spirv.resource_limits<>>
 } {
   gpu.module @kernels {
-    gpu.func @kernel_vector_interleave(%arg0 : vector<2xi32>, %arg1 : vector<2xi32>, %arg2 : memref<4xi32>)
+    gpu.func @kernel_vector_interleave(%arg0 : memref<2xi32>, %arg1 : memref<2xi32>, %arg2 : memref<4xi32>)
       kernel attributes { spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [1, 1, 1]>} {
       %c0 = arith.constant 0 : index
-      %result = vector.interleave %arg0, %arg1 : vector<2xi32>
+      %vec0 = vector.load %arg0[%c0] : memref<2xi32>, vector<2xi32>
+      %vec1 = vector.load %arg1[%c0] : memref<2xi32>, vector<2xi32>
+      %result = vector.interleave %vec0, %vec1 : vector<2xi32> -> vector<4xi32>
       vector.store %result, %arg2[%c0] : memref<4xi32>, vector<4xi32>
       gpu.return
     }
@@ -20,14 +22,20 @@ module attributes {
 
   func.func @main() {
     // Allocate 3 buffers.
-    %buf0 = arith.constant dense<[0, 1]> : vector<2xi32>
-    %buf1 = arith.constant dense<[2, 3]> : vector<2xi32>
+    %buf0 = memref.alloc() : memref<2xi32>
+    %buf1 = memref.alloc() : memref<2xi32>
     %buf2 = memref.alloc() : memref<4xi32>
     
     %idx0 = arith.constant 0 : index
     %idx1 = arith.constant 1 : index
     %idx4 = arith.constant 4 : index
 
+    // Initialize input buffer
+    %buf0_vals = arith.constant dense<[0, 1]> : vector<2xi32>
+    %buf1_vals = arith.constant dense<[2, 3]> : vector<2xi32>
+    vector.store %buf0_vals, %buf0[%idx0] : memref<2xi32>, vector<2xi32>
+    vector.store %buf1_vals, %buf1[%idx0] : memref<2xi32>, vector<2xi32>
+
     // Initialize output buffer.
     %value0 = arith.constant 0 : i32
     %buf3 = memref.cast %buf2 : memref<4xi32> to memref<?xi32>
@@ -35,7 +43,7 @@ module attributes {
 
     gpu.launch_func @kernels::@kernel_vector_interleave
         blocks in (%idx4, %idx1, %idx1) threads in (%idx1, %idx1, %idx1)
-        args(%buf0 : vector<2xi32>, %buf1 : vector<2xi32>, %buf2 : memref<4xi32>)
+        args(%buf0 : memref<2xi32>, %buf1 : memref<2xi32>, %buf2 : memref<4xi32>)
     %buf4 = memref.cast %buf3 : memref<?xi32> to memref<*xi32>
     call @printMemrefI32(%buf4) : (memref<*xi32>) -> ()
     return

>From 11cea8b5e27f60858a68ab52ac4a36bb3b3f803f Mon Sep 17 00:00:00 2001
From: Angel Zhang <angel.zhang at amd.com>
Date: Tue, 28 May 2024 18:52:28 +0000
Subject: [PATCH 3/4] Add integration test for vector.shuffle

---
 .../mlir-vulkan-runner/vector-shuffle.mlir    | 53 +++++++++++++++++++
 1 file changed, 53 insertions(+)
 create mode 100644 mlir/test/mlir-vulkan-runner/vector-shuffle.mlir

diff --git a/mlir/test/mlir-vulkan-runner/vector-shuffle.mlir b/mlir/test/mlir-vulkan-runner/vector-shuffle.mlir
new file mode 100644
index 0000000000000..e29e054ccd46b
--- /dev/null
+++ b/mlir/test/mlir-vulkan-runner/vector-shuffle.mlir
@@ -0,0 +1,53 @@
+// RUN: mlir-vulkan-runner %s \
+// RUN:  --shared-libs=%vulkan-runtime-wrappers,%mlir_runner_utils \
+// RUN:  --entry-point-result=void | FileCheck %s
+
+// CHECK: [2, 1, 3]
+module attributes {
+  gpu.container_module,
+  spirv.target_env = #spirv.target_env<
+    #spirv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>, #spirv.resource_limits<>>
+} {
+  gpu.module @kernels {
+    gpu.func @kernel_vector_shuffle(%arg0 : memref<2xi32>, %arg1 : memref<2xi32>, %arg2 : memref<3xi32>)
+      kernel attributes { spirv.entry_point_abi = #spirv.entry_point_abi<workgroup_size = [1, 1, 1]>} {
+      %c0 = arith.constant 0 : index
+      %vec0 = vector.load %arg0[%c0] : memref<2xi32>, vector<2xi32>
+      %vec1 = vector.load %arg1[%c0] : memref<2xi32>, vector<2xi32>
+      %result = vector.shuffle %vec0, %vec1[2, 1, 3] : vector<2xi32>, vector<2xi32>
+      vector.store %result, %arg2[%c0] : memref<3xi32>, vector<3xi32>
+      gpu.return
+    }
+  }
+
+  func.func @main() {
+    // Allocate 3 buffers.
+    %buf0 = memref.alloc() : memref<2xi32>
+    %buf1 = memref.alloc() : memref<2xi32>
+    %buf2 = memref.alloc() : memref<3xi32>
+    
+    %idx0 = arith.constant 0 : index
+    %idx1 = arith.constant 1 : index
+    %idx4 = arith.constant 4 : index
+
+    // Initialize input buffer
+    %buf0_vals = arith.constant dense<[0, 1]> : vector<2xi32>
+    %buf1_vals = arith.constant dense<[2, 3]> : vector<2xi32>
+    vector.store %buf0_vals, %buf0[%idx0] : memref<2xi32>, vector<2xi32>
+    vector.store %buf1_vals, %buf1[%idx0] : memref<2xi32>, vector<2xi32>
+
+    // Initialize output buffer.
+    %value0 = arith.constant 0 : i32
+    %buf3 = memref.cast %buf2 : memref<3xi32> to memref<?xi32>
+    call @fillResource1DInt(%buf3, %value0) : (memref<?xi32>, i32) -> ()
+
+    gpu.launch_func @kernels::@kernel_vector_shuffle
+        blocks in (%idx4, %idx1, %idx1) threads in (%idx1, %idx1, %idx1)
+        args(%buf0 : memref<2xi32>, %buf1 : memref<2xi32>, %buf2 : memref<3xi32>)
+    %buf4 = memref.cast %buf3 : memref<?xi32> to memref<*xi32>
+    call @printMemrefI32(%buf4) : (memref<*xi32>) -> ()
+    return
+  }
+  func.func private @fillResource1DInt(%0 : memref<?xi32>, %1 : i32)
+  func.func private @printMemrefI32(%ptr : memref<*xi32>)
+}

>From b2f029ef35fad1b9cef0affce7746d3c78fb41b5 Mon Sep 17 00:00:00 2001
From: Angel Zhang <anzhouzhang913 at gmail.com>
Date: Tue, 28 May 2024 21:57:16 -0400
Subject: [PATCH 4/4] Change comment

Co-authored-by: Jakub Kuderski <kubakuderski at gmail.com>
---
 mlir/test/mlir-vulkan-runner/vector-interleave.mlir | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/test/mlir-vulkan-runner/vector-interleave.mlir b/mlir/test/mlir-vulkan-runner/vector-interleave.mlir
index d760c1631c473..2f5c319e2f5c5 100644
--- a/mlir/test/mlir-vulkan-runner/vector-interleave.mlir
+++ b/mlir/test/mlir-vulkan-runner/vector-interleave.mlir
@@ -30,7 +30,7 @@ module attributes {
     %idx1 = arith.constant 1 : index
     %idx4 = arith.constant 4 : index
 
-    // Initialize input buffer
+    // Initialize input buffer.
     %buf0_vals = arith.constant dense<[0, 1]> : vector<2xi32>
     %buf1_vals = arith.constant dense<[2, 3]> : vector<2xi32>
     vector.store %buf0_vals, %buf0[%idx0] : memref<2xi32>, vector<2xi32>



More information about the Mlir-commits mailing list