-
Notifications
You must be signed in to change notification settings - Fork 15.2k
[mlir][vector][nfc] Add tests + update docs for narrow-type emulation #115460
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from 3 commits
8a9abf6
eef52e2
61e7984
933694b
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change | ||||
|---|---|---|---|---|---|---|
| @@ -1,6 +1,10 @@ | ||||||
| // RUN: mlir-opt --test-emulate-narrow-int="arith-compute-bitwidth=1 memref-load-bitwidth=8" --cse --split-input-file %s | FileCheck %s | ||||||
| // RUN: mlir-opt --test-emulate-narrow-int="arith-compute-bitwidth=1 memref-load-bitwidth=32" --cse --split-input-file %s | FileCheck %s --check-prefix=CHECK32 | ||||||
|
|
||||||
| ///---------------------------------------------------------------------------------------- | ||||||
| /// vector.load | ||||||
| ///---------------------------------------------------------------------------------------- | ||||||
|
|
||||||
| func.func @vector_load_i8(%arg1: index, %arg2: index) -> vector<4xi8> { | ||||||
| %0 = memref.alloc() : memref<3x4xi8> | ||||||
| %1 = vector.load %0[%arg1, %arg2] : memref<3x4xi8>, vector<4xi8> | ||||||
|
|
@@ -82,6 +86,10 @@ func.func @vector_load_i4_dynamic(%arg0 : index, %arg1 : index, %arg2 : index, % | |||||
|
|
||||||
| // ----- | ||||||
|
|
||||||
| ///---------------------------------------------------------------------------------------- | ||||||
| /// vector.transfer_read | ||||||
| ///---------------------------------------------------------------------------------------- | ||||||
|
|
||||||
| func.func @vector_transfer_read_i4(%arg1: index, %arg2: index) -> vector<8xi4> { | ||||||
| %c0 = arith.constant 0 : i4 | ||||||
| %0 = memref.alloc() : memref<3x8xi4> | ||||||
|
|
@@ -111,6 +119,10 @@ func.func @vector_transfer_read_i4(%arg1: index, %arg2: index) -> vector<8xi4> { | |||||
|
|
||||||
| // ----- | ||||||
|
|
||||||
| ///---------------------------------------------------------------------------------------- | ||||||
| /// vector.maskedload | ||||||
| ///---------------------------------------------------------------------------------------- | ||||||
|
|
||||||
| func.func @vector_maskedload_i8(%arg1: index, %arg2: index, %arg3: index, %passthru: vector<4xi8>) -> vector<4xi8> { | ||||||
| %0 = memref.alloc() : memref<3x4xi8> | ||||||
| %mask = vector.create_mask %arg3 : vector<4xi1> | ||||||
|
|
@@ -190,15 +202,15 @@ func.func @vector_maskedload_i4(%arg1: index, %arg2: index, %arg3: index, %passt | |||||
|
|
||||||
| // ----- | ||||||
|
|
||||||
| func.func @vector_cst_maskedload_i8(%arg1: index, %arg2: index, %passthru: vector<4xi8>) -> vector<4xi8> { | ||||||
| func.func @vector_maskedload_i8_constant_mask(%arg1: index, %arg2: index, %passthru: vector<4xi8>) -> vector<4xi8> { | ||||||
| %0 = memref.alloc() : memref<3x4xi8> | ||||||
| %mask = vector.constant_mask [2] : vector<4xi1> | ||||||
| %1 = vector.maskedload %0[%arg1, %arg2], %mask, %passthru : | ||||||
| memref<3x4xi8>, vector<4xi1>, vector<4xi8> into vector<4xi8> | ||||||
| return %1 : vector<4xi8> | ||||||
| } | ||||||
| // Expect no conversions, i8 is supported. | ||||||
| // CHECK: func @vector_cst_maskedload_i8( | ||||||
| // CHECK: func @vector_maskedload_i8_constant_mask( | ||||||
| // CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: index, %[[ARG1:[a-zA-Z0-9]+]]: index, | ||||||
| // CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: vector<4xi8>) | ||||||
| // CHECK-NEXT: %[[ALLOC:.+]] = memref.alloc() : memref<3x4xi8> | ||||||
|
|
@@ -208,7 +220,7 @@ func.func @vector_cst_maskedload_i8(%arg1: index, %arg2: index, %passthru: vecto | |||||
| // CHECK-NEXT: return | ||||||
|
|
||||||
| // CHECK32-DAG: #[[LOAD_IDX_MAP:.+]] = affine_map<()[s0, s1] -> (s0 + s1 floordiv 4)> | ||||||
| // CHECK32: func @vector_cst_maskedload_i8( | ||||||
| // CHECK32: func @vector_maskedload_i8_constant_mask( | ||||||
| // CHECK32-SAME: %[[ARG0:[a-zA-Z0-9]+]]: index, %[[ARG1:[a-zA-Z0-9]+]]: index, | ||||||
| // CHECK32-SAME: %[[ARG3:[a-zA-Z0-9]+]]: vector<4xi8>) | ||||||
| // CHECK32: %[[ALLOC:.+]] = memref.alloc() : memref<3xi32> | ||||||
|
|
@@ -224,7 +236,7 @@ func.func @vector_cst_maskedload_i8(%arg1: index, %arg2: index, %passthru: vecto | |||||
|
|
||||||
| // ----- | ||||||
|
|
||||||
| func.func @vector_cst_maskedload_i4(%arg1: index, %arg2: index, %passthru: vector<8xi4>) -> vector<3x8xi4> { | ||||||
| func.func @vector_maskedload_i4_constant_mask(%arg1: index, %arg2: index, %passthru: vector<8xi4>) -> vector<3x8xi4> { | ||||||
| %0 = memref.alloc() : memref<3x8xi4> | ||||||
| %cst = arith.constant dense<0> : vector<3x8xi4> | ||||||
| %mask = vector.constant_mask [4] : vector<8xi1> | ||||||
|
|
@@ -234,7 +246,7 @@ func.func @vector_cst_maskedload_i4(%arg1: index, %arg2: index, %passthru: vecto | |||||
| return %2 : vector<3x8xi4> | ||||||
| } | ||||||
| // CHECK-DAG: #[[LOAD_IDX_MAP:.+]] = affine_map<()[s0, s1] -> (s0 * 4 + s1 floordiv 2)> | ||||||
| // CHECK: func @vector_cst_maskedload_i4( | ||||||
| // CHECK: func @vector_maskedload_i4_constant_mask( | ||||||
| // CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: index, %[[ARG1:[a-zA-Z0-9]+]]: index, | ||||||
| // CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: vector<8xi4>) | ||||||
| // CHECK: %[[ALLOC:.+]] = memref.alloc() : memref<12xi8> | ||||||
|
|
@@ -248,7 +260,7 @@ func.func @vector_cst_maskedload_i4(%arg1: index, %arg2: index, %passthru: vecto | |||||
| // CHECK: %[[SELECT:.+]] = arith.select %[[ORIG_MASK]], %[[BITCAST]], %[[ARG2]] : vector<8xi1>, vector<8xi4> | ||||||
|
|
||||||
| // CHECK32-DAG: #[[LOAD_IDX_MAP:.+]] = affine_map<()[s0, s1] -> (s0 + s1 floordiv 8)> | ||||||
| // CHECK32: func @vector_cst_maskedload_i4( | ||||||
| // CHECK32: func @vector_maskedload_i4_constant_mask( | ||||||
| // CHECK32-SAME: %[[ARG0:[a-zA-Z0-9]+]]: index, %[[ARG1:[a-zA-Z0-9]+]]: index, | ||||||
| // CHECK32-SAME: %[[ARG2:[a-zA-Z0-9]+]]: vector<8xi4>) | ||||||
| // CHECK32: %[[ALLOC:.+]] = memref.alloc() : memref<3xi32> | ||||||
|
|
@@ -263,6 +275,10 @@ func.func @vector_cst_maskedload_i4(%arg1: index, %arg2: index, %passthru: vecto | |||||
|
|
||||||
| // ----- | ||||||
|
|
||||||
| ///---------------------------------------------------------------------------------------- | ||||||
| /// vector.extract -> vector.masked_load | ||||||
| ///---------------------------------------------------------------------------------------- | ||||||
|
|
||||||
| func.func @vector_extract_maskedload_i4(%arg1: index) -> vector<8x8x16xi4> { | ||||||
| %0 = memref.alloc() : memref<8x8x16xi4> | ||||||
| %c0 = arith.constant 0 : index | ||||||
|
|
@@ -353,6 +369,10 @@ func.func @vector_extract_cst_maskedload_i4() -> vector<8x8x16xi4> { | |||||
|
|
||||||
| // ----- | ||||||
|
|
||||||
| ///---------------------------------------------------------------------------------------- | ||||||
| /// vector.store | ||||||
| ///---------------------------------------------------------------------------------------- | ||||||
|
|
||||||
| func.func @vector_store_i8(%arg0: vector<8xi8>, %arg1: index, %arg2: index) { | ||||||
| %0 = memref.alloc() : memref<4x8xi8> | ||||||
| vector.store %arg0, %0[%arg1, %arg2] :memref<4x8xi8>, vector<8xi8> | ||||||
|
|
@@ -431,6 +451,10 @@ func.func @vector_store_i4_dynamic(%arg0: vector<8xi4>, %arg1: index, %arg2: ind | |||||
|
|
||||||
| // ----- | ||||||
|
|
||||||
| ///---------------------------------------------------------------------------------------- | ||||||
| /// vector.maskedstore | ||||||
| ///---------------------------------------------------------------------------------------- | ||||||
|
|
||||||
| func.func @vector_maskedstore_i8(%arg0: index, %arg1: index, %arg2: index, %value: vector<8xi8>) { | ||||||
| %0 = memref.alloc() : memref<3x8xi8> | ||||||
| %mask = vector.create_mask %arg2 : vector<8xi1> | ||||||
|
|
@@ -469,14 +493,68 @@ func.func @vector_maskedstore_i8(%arg0: index, %arg1: index, %arg2: index, %valu | |||||
|
|
||||||
| // ----- | ||||||
|
|
||||||
| func.func @vector_cst_maskedstore_i8(%arg0: index, %arg1: index, %value: vector<8xi8>) { | ||||||
| func.func @vector_maskedstore_i4( | ||||||
| %idx1: index, | ||||||
| %idx2: index, | ||||||
| %num_elements_to_store: index, | ||||||
| %value: vector<8xi4>) { | ||||||
|
|
||||||
| %0 = memref.alloc() : memref<3x8xi4> | ||||||
| %mask = vector.create_mask %num_elements_to_store : vector<8xi1> | ||||||
| vector.maskedstore %0[%idx1, %idx2], %mask, %value : | ||||||
| memref<3x8xi4>, vector<8xi1>, vector<8xi4> | ||||||
| return | ||||||
| } | ||||||
| // CHECK: #[[$ATTR_10:.+]] = affine_map<()[s0, s1] -> (s0 * 4 + s1 floordiv 2)> | ||||||
| // CHECK: #[[$ATTR_11:.+]] = affine_map<()[s0] -> ((s0 + 1) floordiv 2)> | ||||||
|
|
||||||
| // CHECK-LABEL: func.func @vector_maskedstore_i4( | ||||||
| // CHECK-SAME: %[[IDX_1:[a-zA-Z0-9]+]]: index, | ||||||
| // CHECK-SAME: %[[IDX_2:[a-zA-Z0-9]+]]: index, | ||||||
| // CHECK-SAME: %[[NUM_EL_TO_STORE:[a-zA-Z0-9]+]]: index, | ||||||
| // CHECK-SAME: %[[VAL_TO_STORE:[a-zA-Z0-9]+]]: vector<8xi4>) { | ||||||
| // CHECK: %[[ALLOC:.+]] = memref.alloc() : memref<12xi8> | ||||||
| // CHECK: %[[ORIG_MASK:.+]] = vector.create_mask %[[NUM_EL_TO_STORE]] : vector<8xi1> | ||||||
| // CHECK: %[[LIDX:.+]] = affine.apply #[[$ATTR_10]](){{\[}}%[[IDX_1]], %[[IDX_2]]] | ||||||
|
||||||
| // CHECK: %[[LIDX:.+]] = affine.apply #[[$ATTR_10]](){{\[}}%[[IDX_1]], %[[IDX_2]]] | |
| // CHECK: %[[LIDX:.+]] = affine.apply #[[$ATTR_10]]()[%[[IDX_1]], %[[IDX_2]]] |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Oh, very nice, thanks for pointing this out 🙏🏻
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Note that there is actually nothing fundamental about power of 2 here. It is an implementation detail and just NYI. Maybe might be better to make that clear.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
+1, perhaps we can explicitly say that the non-power-of-two types are not implemented yet in the description. It looks more clear for newcomers for the narrow type emulations.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I will add a TODO at the top - clear indication for anyone brave enough to tackle this :)