CodeQL / Analyze (actions) (push) Canceled after 0s
CodeQL / Analyze (python) (push) Canceled after 0s
pre-commit / pre-commit (push) Canceled after 0s
Multi-Arch CI / setup (push) Canceled after 0s
Publish build_manylinux_x86_64 images / publish_build_manylinux_x86_64 (push) Canceled after 0s
Setup (Multi-Arch) / setup (push) Canceled after 0s
Multi-Arch CI / Linux::${{ fromJSON(needs.setup.outputs.linux_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubi10 image / publish_no_rocm_image_ubi10 (push) Canceled after 0s
Multi-Arch CI / Windows::${{ fromJSON(needs.setup.outputs.windows_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Multi-Arch CI / Manifest Diff (push) Canceled after 0s
TheRock Manifest Diff Report / Generate Manifest Diff Report (push) Canceled after 0s
Multi-Arch CI / CI Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_ocl_rt images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_rocgdb images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Unit Tests / Unit Tests :: ubuntu-24.04 (push) Canceled after 0s
Unit Tests / Unit Tests :: windows-2022 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_media (push) Canceled after 0s
Unit Tests / Unit Tests Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_openmpi (push) Canceled after 0s
258 lines
8.7 KiB
MLIR
258 lines
8.7 KiB
MLIR
// RUN: mlir-opt --arm-sve-legalize-vector-storage --split-input-file %s | FileCheck %s
|
|
|
|
|
|
// Test the `LegalizeTransferRead` pattern
|
|
// (mlir/lib/Dialect/ArmSVE/Transforms/LegalizeVectorStorage.cpp)
|
|
|
|
// -----
|
|
|
|
// This is the base case, unremarkable in any way, except that it's our main
|
|
// motivating example and use case.
|
|
|
|
// CHECK-LABEL: @base_case
|
|
// CHECK-SAME: %[[I:.+]]: index, %[[J:.+]]: index, %[[M:.+]]:
|
|
// CHECK-DAG: %[[PAD:.+]] = arith.constant 123 : i8
|
|
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
|
|
// CHECK: %[[COLLAPSE:.+]] = memref.collapse_shape %[[M]]
|
|
// CHECK-SAME{LITERAL}: [[0], [1], [2, 3]]
|
|
// CHECK-SAME: : memref<?x?x?x8xi8> into memref<?x?x?xi8>
|
|
// CHECK-NEXT: %[[T0:.+]] = vector.transfer_read %[[COLLAPSE]][%[[I]], %[[J]], %[[C0]]], %[[PAD]] {in_bounds = [true]}
|
|
// CHECK-SAME: : memref<?x?x?xi8>, vector<[32]xi8>
|
|
// CHECK-NEXT: %[[T1:.+]] = vector.shape_cast %[[T0]] : vector<[32]xi8> to vector<[4]x8xi8>
|
|
// CHECK-NEXT: return %[[T1]] : vector<[4]x8xi8>
|
|
|
|
func.func @base_case(%i : index, %j : index, %M : memref<?x?x?x8xi8>) -> vector<[4]x8xi8> {
|
|
%c0 = arith.constant 0 : index
|
|
%pad = arith.constant 123 : i8
|
|
|
|
%A = vector.transfer_read %M[%i, %j, %c0, %c0], %pad {in_bounds = [true, true]} : memref<?x?x?x8xi8>, vector<[4]x8xi8>
|
|
|
|
return %A : vector<[4]x8xi8>
|
|
}
|
|
|
|
// -----
|
|
|
|
// Test the case where the scalable dimension is not the second-to-last.
|
|
|
|
// CHECK-LABEL: @with_3d_vector
|
|
// CHECK-SAME: %[[I:.+]]: index, %[[J:.+]]: index, %[[M:.+]]:
|
|
// CHECK-DAG: %[[PAD:.+]] = arith.constant 123 : i8
|
|
// CHECK-DAG: %[[COLLAPSED:.+]] = memref.collapse_shape %[[M]]
|
|
// CHECK-SAME{LITERAL}: [[0], [1, 2, 3]]
|
|
// CHECK-SAME: : memref<?x?x2x8xi8> into memref<?x?xi8>
|
|
// CHECK-NEXT: %[[T0:.+]] = vector.transfer_read %[[COLLAPSED]][%[[I]], %[[J]]], %[[PAD]] {in_bounds = [true]}
|
|
// CHECK-SAME: : memref<?x?xi8>, vector<[64]xi8>
|
|
// CHECK-NEXT: %[[T1:.+]] = vector.shape_cast %[[T0]] : vector<[64]xi8> to vector<[4]x2x8xi8>
|
|
// CHECK-NEXT: return %[[T1]] : vector<[4]x2x8xi8>
|
|
|
|
func.func @with_3d_vector(%i : index, %j : index, %M : memref<?x?x2x8xi8>) -> vector<[4]x2x8xi8> {
|
|
%c0 = arith.constant 0 : index
|
|
%pad = arith.constant 123 : i8
|
|
|
|
%A = vector.transfer_read %M[%i, %j, %c0, %c0], %pad {in_bounds = [true, true, true]} : memref<?x?x2x8xi8>, vector<[4]x2x8xi8>
|
|
|
|
return %A : vector<[4]x2x8xi8>
|
|
}
|
|
|
|
// -----
|
|
|
|
// Test the case when the vector is already LLVM-legal (fixed).
|
|
|
|
// CHECK-LABEL: @negative_vector_legal_fixed
|
|
// CHECK-NOT: memref.collapse
|
|
|
|
func.func @negative_vector_legal_fixed(%i : index, %j : index, %M : memref<?x?x?x8xi8>) -> vector<8x8xi8> {
|
|
%c0 = arith.constant 0 : index
|
|
%pad = arith.constant 123 : i8
|
|
|
|
%A = vector.transfer_read %M[%i, %j, %c0, %c0], %pad {in_bounds = [true, true]} : memref<?x?x?x8xi8>, vector<8x8xi8>
|
|
|
|
return %A : vector<8x8xi8>
|
|
}
|
|
|
|
// -----
|
|
|
|
// Test the case when the vector is already LLVM-legal (single-dimension scalable).
|
|
|
|
// CHECK-LABEL: @negative_vector_legal_1d_scalable
|
|
// CHECK-NOT: memref.collapse
|
|
|
|
func.func @negative_vector_legal_1d_scalable(%i : index, %j : index, %M : memref<?x?x?x8xi8>) -> vector<[8]xi8> {
|
|
%c0 = arith.constant 0 : index
|
|
%pad = arith.constant 123 : i8
|
|
|
|
%A = vector.transfer_read %M[%i, %j, %c0, %c0], %pad {in_bounds = [true]} : memref<?x?x?x8xi8>, vector<[8]xi8>
|
|
|
|
return %A : vector<[8]xi8>
|
|
}
|
|
|
|
// -----
|
|
|
|
// Test the case when the vector is already LLVM-legal (single trailing
|
|
// scalable dimension).
|
|
|
|
// CHECK-LABEL: @negative_vector_legal_trailing_scalable_dim
|
|
// CHECK-NOT: memref.collapse
|
|
|
|
func.func @negative_vector_legal_trailing_scalable_dim(%i : index, %j : index, %M : memref<?x?x?x8xi8>) -> vector<8x[8]xi8> {
|
|
%c0 = arith.constant 0 : index
|
|
%pad = arith.constant 123 : i8
|
|
|
|
%A = vector.transfer_read %M[%i, %j, %c0, %c0], %pad {in_bounds = [true, true]} : memref<?x?x?x8xi8>, vector<8x[8]xi8>
|
|
|
|
return %A : vector<8x[8]xi8>
|
|
}
|
|
|
|
// -----
|
|
|
|
// Test the case of unsupported vector type (more than one scalable dimension)
|
|
|
|
// CHECK-LABEL: @negative_vector_type_two_scalable_dims
|
|
// CHECK-NOT: memref.collapse
|
|
|
|
func.func @negative_vector_type_two_scalable_dims(%i : index, %j : index, %M : memref<?x?x?x8xi8>) -> vector<[8]x[8]x8xi8> {
|
|
%c0 = arith.constant 0 : index
|
|
%pad = arith.constant 123 : i8
|
|
|
|
%A = vector.transfer_read %M[%i, %j, %c0, %c0], %pad {in_bounds = [true, true, true]} : memref<?x?x?x8xi8>, vector<[8]x[8]x8xi8>
|
|
|
|
return %A : vector<[8]x[8]x8xi8>
|
|
}
|
|
|
|
// -----
|
|
|
|
// Test the case of reading from a tensor - not supported, since the
|
|
// transform reasons about memory layouts.
|
|
|
|
// CHECK-LABEL: @negative_tensor_transfer
|
|
// CHECK-NOT: memref.collapse
|
|
|
|
func.func @negative_tensor_transfer(%i : index, %j : index, %M : tensor<?x?x?x8xi8>) -> vector<[4]x8xi8> {
|
|
%c0 = arith.constant 0 : index
|
|
%pad = arith.constant 123 : i8
|
|
|
|
%A = vector.transfer_read %M[%i, %j, %c0, %c0], %pad {in_bounds = [true, true]} : tensor<?x?x?x8xi8>, vector<[4]x8xi8>
|
|
|
|
return %A : vector<[4]x8xi8>
|
|
}
|
|
|
|
// -----
|
|
|
|
// Test the case when the transfer is discontiguous because the memref
|
|
// is discontiguous.
|
|
// There are other ways to make a memref discontiguous. The transformation
|
|
// is not concerned with the particular reason a memref is discontiguous, but
|
|
// only with the fact. Therefore there are no variations with the memref made
|
|
// discontiguous by some other mechanism.
|
|
|
|
// CHECK-LABEL: @negative_discontig_mem
|
|
// CHECK-NOT: memref.collapse
|
|
|
|
#strides = strided<[?, ?, 16, 1]>
|
|
|
|
func.func @negative_discontig_mem(%i : index, %j : index, %M : memref<?x?x?x8xi8, #strides>) -> vector<[4]x8xi8> {
|
|
%c0 = arith.constant 0 : index
|
|
%pad = arith.constant 123 : i8
|
|
|
|
%A = vector.transfer_read %M[%i, %j, %c0, %c0], %pad {in_bounds = [true, true]} : memref<?x?x?x8xi8, #strides>, vector<[4]x8xi8>
|
|
|
|
return %A : vector<[4]x8xi8>
|
|
}
|
|
|
|
// -----
|
|
|
|
// Test the case when the transformation is not applied because of
|
|
// a non-trivial permutation map (broadcast).
|
|
|
|
// CHECK-LABEL: @negative_broadcast
|
|
// CHECK-NOT: memref.collapse
|
|
|
|
#perm = affine_map<(i, j, k, p) -> (k, 0)>
|
|
|
|
func.func @negative_broadcast(%i : index, %j : index, %M : memref<?x?x?x8xi8>) -> vector<[4]x8xi8> {
|
|
%c0 = arith.constant 0 : index
|
|
%pad = arith.constant 123 : i8
|
|
|
|
%A = vector.transfer_read %M[%i, %j, %c0, %c0], %pad {permutation_map = #perm, in_bounds = [true, true] } : memref<?x?x?x8xi8>, vector<[4]x8xi8>
|
|
|
|
return %A : vector<[4]x8xi8>
|
|
}
|
|
|
|
// -----
|
|
|
|
// Test the case of a masked read - not supported right now.
|
|
// (see mlir/lib/Dialect/ArmSVE/Transforms/LegalizeVectorStorage.cpp)
|
|
|
|
// CHECK-LABEL: @negative_masked
|
|
// CHECK-NOT: memref.collapse
|
|
|
|
func.func @negative_masked(
|
|
%i : index, %j : index,
|
|
%M : memref<?x?x?x8xi8>, %mask : vector<[4]x8xi1>) -> vector<[4]x8xi8> {
|
|
|
|
%c0 = arith.constant 0 : index
|
|
%pad = arith.constant 123 : i8
|
|
|
|
%A = vector.mask %mask {
|
|
vector.transfer_read %M[%i, %j, %c0, %c0], %pad {in_bounds = [true, true] } : memref<?x?x?x8xi8>, vector<[4]x8xi8>
|
|
} : vector<[4]x8xi1> -> vector<[4]x8xi8>
|
|
|
|
return %A : vector<[4]x8xi8>
|
|
}
|
|
|
|
// -----
|
|
|
|
// Test case with a mask operand - not supported right now.
|
|
// (see mlir/lib/Dialect/ArmSVE/Transforms/LegalizeVectorStorage.cpp)
|
|
|
|
// CHECK-LABEL: @negative_with_mask
|
|
// CHECK-NOT: memref.collapse
|
|
|
|
func.func @negative_with_mask(
|
|
%i : index, %j : index,
|
|
%M : memref<?x?x?x8xi8>, %mask : vector<[4]x8xi1>) -> vector<[4]x8xi8> {
|
|
|
|
%c0 = arith.constant 0 : index
|
|
%pad = arith.constant 123 : i8
|
|
|
|
%A = vector.transfer_read %M[%i, %j, %c0, %c0], %pad, %mask {in_bounds = [true, true] } : memref<?x?x?x8xi8>, vector<[4]x8xi8>
|
|
|
|
return %A : vector<[4]x8xi8>
|
|
}
|
|
|
|
// -----
|
|
|
|
// Test the case when the dimensions to collapse (excluding the scalable one)
|
|
// of the vector and the memref do not match (static non matching dimension).
|
|
|
|
// CHECK-LABEL: @negative_non_matching_dim_static
|
|
// CHECK-NOT: memref.collapse
|
|
|
|
func.func @negative_non_matching_dim_static(%i : index, %j : index, %M : memref<?x?x?x8xi8>) -> vector<[4]x4xi8> {
|
|
|
|
%c0 = arith.constant 0 : index
|
|
%pad = arith.constant 123 : i8
|
|
|
|
%A = vector.transfer_read %M[%i, %j, %c0, %c0], %pad {in_bounds = [true, true] } : memref<?x?x?x8xi8>, vector<[4]x4xi8>
|
|
|
|
return %A : vector<[4]x4xi8>
|
|
}
|
|
|
|
// -----
|
|
|
|
// Test the case when the dimensions to collapse (excluding the scalable one)
|
|
// of the vector and the memref do not match (dynamic non matching dimension).
|
|
|
|
// CHECK-LABEL: @negative_non_matching_dim_dynamic
|
|
// CHECK-NOT: memref.collapse
|
|
|
|
func.func @negative_non_matching_dim_dynamic(%i : index, %j : index, %M : memref<?x?x?x?xi8>) -> vector<[4]x4xi8> {
|
|
|
|
%c0 = arith.constant 0 : index
|
|
%pad = arith.constant 123 : i8
|
|
|
|
%A = vector.transfer_read %M[%i, %j, %c0, %c0], %pad {in_bounds = [true, true] } : memref<?x?x?x?xi8>, vector<[4]x4xi8>
|
|
|
|
return %A : vector<[4]x4xi8>
|
|
}
|