Files
TheRock/compiler/amd-llvm/mlir/test/Dialect/ArmSVE/legalize-transfer-read.mlir
T
derek.holloway 7502bf64a4
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
Flatten subprojects into a single base monorepo
2026-09-02 16:34:47 -07:00

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>
}