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
1929 lines
77 KiB
MLIR
1929 lines
77 KiB
MLIR
// RUN: mlir-opt %s -split-input-file -linalg-morph-ops=generic-to-named \
|
|
// RUN: | FileCheck %s --check-prefix=NAMED,ALL
|
|
|
|
// RUN: mlir-opt %s -split-input-file -linalg-morph-ops=generic-to-category \
|
|
// RUN: | FileCheck %s --check-prefix=CATEGORY,ALL
|
|
|
|
#umap = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
|
|
func.func @unary_ops(%A: tensor<?x?x?xf32>, %Out: tensor<?x?x?xf32>) -> tensor<?x?x?xf32> {
|
|
%0 = linalg.generic
|
|
{indexing_maps = [#umap, #umap],
|
|
iterator_types = ["parallel", "parallel","parallel"]}
|
|
ins(%A : tensor<?x?x?xf32>)
|
|
outs(%Out : tensor<?x?x?xf32>) {
|
|
^bb0(%in: f32, %out: f32):
|
|
%v = math.exp %in : f32
|
|
linalg.yield %v : f32
|
|
} -> tensor<?x?x?xf32>
|
|
%1 = linalg.generic
|
|
{indexing_maps = [#umap, #umap], iterator_types = ["parallel", "parallel","parallel"]}
|
|
ins(%0 : tensor<?x?x?xf32>) outs(%Out : tensor<?x?x?xf32>) {
|
|
^bb0(%in: f32, %out: f32):
|
|
%v = math.log %in : f32
|
|
linalg.yield %v : f32
|
|
} -> tensor<?x?x?xf32>
|
|
%2 = linalg.generic
|
|
{indexing_maps = [#umap, #umap], iterator_types = ["parallel", "parallel","parallel"]}
|
|
ins(%1 : tensor<?x?x?xf32>) outs(%Out : tensor<?x?x?xf32>) {
|
|
^bb0(%in: f32, %out: f32):
|
|
%v = math.absf %in : f32
|
|
linalg.yield %v : f32
|
|
} -> tensor<?x?x?xf32>
|
|
%3 = linalg.generic
|
|
{indexing_maps = [#umap, #umap], iterator_types = ["parallel", "parallel","parallel"]}
|
|
ins(%2 : tensor<?x?x?xf32>) outs(%Out : tensor<?x?x?xf32>) {
|
|
^bb0(%in: f32, %out: f32):
|
|
%v = math.ceil %in : f32
|
|
linalg.yield %v : f32
|
|
} -> tensor<?x?x?xf32>
|
|
%4 = linalg.generic
|
|
{indexing_maps = [#umap, #umap], iterator_types = ["parallel", "parallel","parallel"]}
|
|
ins(%3 : tensor<?x?x?xf32>) outs(%Out : tensor<?x?x?xf32>) {
|
|
^bb0(%in: f32, %out: f32):
|
|
%v = math.floor %in : f32
|
|
linalg.yield %v : f32
|
|
} -> tensor<?x?x?xf32>
|
|
%5 = linalg.generic
|
|
{indexing_maps = [#umap, #umap], iterator_types = ["parallel", "parallel","parallel"]}
|
|
ins(%4 : tensor<?x?x?xf32>) outs(%Out : tensor<?x?x?xf32>) {
|
|
^bb0(%in: f32, %out: f32):
|
|
%v = arith.negf %in : f32
|
|
linalg.yield %v : f32
|
|
} -> tensor<?x?x?xf32>
|
|
%cst_1 = arith.constant 1.0 : f32
|
|
%6 = linalg.generic
|
|
{indexing_maps = [#umap, #umap], iterator_types = ["parallel", "parallel","parallel"]}
|
|
ins(%5 : tensor<?x?x?xf32>) outs(%Out : tensor<?x?x?xf32>) {
|
|
^bb0(%in: f32, %out: f32):
|
|
%v = arith.divf %cst_1, %in : f32
|
|
linalg.yield %v : f32
|
|
} -> tensor<?x?x?xf32>
|
|
%7 = linalg.generic
|
|
{indexing_maps = [#umap, #umap], iterator_types = ["parallel", "parallel","parallel"]}
|
|
ins(%6 : tensor<?x?x?xf32>) outs(%Out : tensor<?x?x?xf32>) {
|
|
^bb0(%in: f32, %out: f32):
|
|
%v = math.round %in : f32
|
|
linalg.yield %v : f32
|
|
} -> tensor<?x?x?xf32>
|
|
%8 = linalg.generic
|
|
{indexing_maps = [#umap, #umap], iterator_types = ["parallel", "parallel","parallel"]}
|
|
ins(%7 : tensor<?x?x?xf32>) outs(%Out : tensor<?x?x?xf32>) {
|
|
^bb0(%in: f32, %out: f32):
|
|
%v = math.sqrt %in : f32
|
|
linalg.yield %v : f32
|
|
} -> tensor<?x?x?xf32>
|
|
%9 = linalg.generic
|
|
{indexing_maps = [#umap, #umap], iterator_types = ["parallel", "parallel","parallel"]}
|
|
ins(%8 : tensor<?x?x?xf32>) outs(%Out : tensor<?x?x?xf32>) {
|
|
^bb0(%in: f32, %out: f32):
|
|
%v = math.rsqrt %in : f32
|
|
linalg.yield %v : f32
|
|
} -> tensor<?x?x?xf32>
|
|
%10 = linalg.generic
|
|
{indexing_maps = [#umap, #umap], iterator_types = ["parallel", "parallel","parallel"]}
|
|
ins(%9 : tensor<?x?x?xf32>) outs(%Out : tensor<?x?x?xf32>) {
|
|
^bb0(%in: f32, %out: f32):
|
|
%v = arith.mulf %in, %in : f32
|
|
linalg.yield %v : f32
|
|
} -> tensor<?x?x?xf32>
|
|
%11 = linalg.generic
|
|
{indexing_maps = [#umap, #umap], iterator_types = ["parallel", "parallel","parallel"]}
|
|
ins(%10 : tensor<?x?x?xf32>) outs(%Out : tensor<?x?x?xf32>) {
|
|
^bb0(%in: f32, %out: f32):
|
|
%v = math.tanh %in : f32
|
|
linalg.yield %v : f32
|
|
} -> tensor<?x?x?xf32>
|
|
%12 = linalg.generic
|
|
{indexing_maps = [#umap, #umap], iterator_types = ["parallel", "parallel","parallel"]}
|
|
ins(%11 : tensor<?x?x?xf32>) outs(%Out : tensor<?x?x?xf32>) {
|
|
^bb0(%in: f32, %out: f32):
|
|
%v = math.erf %in : f32
|
|
linalg.yield %v : f32
|
|
} -> tensor<?x?x?xf32>
|
|
return %12 : tensor<?x?x?xf32>
|
|
}
|
|
|
|
// ALL-LABEL: unary_ops
|
|
// ALL-SAME: %[[A:.+]]: tensor<?x?x?xf32>, %[[OUT:.+]]: tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
|
|
|
|
// NAMED-NOT: linalg.generic
|
|
// NAMED: %[[RES0:.+]] = linalg.exp
|
|
// NAMED-SAME: ins(%[[A]] : tensor<?x?x?xf32>)
|
|
// NAMED-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
|
|
// NAMED: %[[RES1:.+]] = linalg.log
|
|
// NAMED-SAME: ins(%[[RES0]] : tensor<?x?x?xf32>)
|
|
// NAMED-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
|
|
// NAMED: %[[RES2:.+]] = linalg.abs
|
|
// NAMED-SAME: ins(%[[RES1]] : tensor<?x?x?xf32>)
|
|
// NAMED-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
|
|
// NAMED: %[[RES3:.+]] = linalg.ceil
|
|
// NAMED-SAME: ins(%[[RES2]] : tensor<?x?x?xf32>)
|
|
// NAMED-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
|
|
// NAMED: %[[RES4:.+]] = linalg.floor
|
|
// NAMED-SAME: ins(%[[RES3]] : tensor<?x?x?xf32>)
|
|
// NAMED-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
|
|
// NAMED: %[[RES5:.+]] = linalg.negf
|
|
// NAMED-SAME: ins(%[[RES4]] : tensor<?x?x?xf32>)
|
|
// NAMED-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
|
|
// NAMED: %[[RES6:.+]] = linalg.reciprocal
|
|
// NAMED-SAME: ins(%[[RES5]] : tensor<?x?x?xf32>)
|
|
// NAMED-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
|
|
// NAMED: %[[RES7:.+]] = linalg.round
|
|
// NAMED-SAME: ins(%[[RES6]] : tensor<?x?x?xf32>)
|
|
// NAMED-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
|
|
// NAMED: %[[RES8:.+]] = linalg.sqrt
|
|
// NAMED-SAME: ins(%[[RES7]] : tensor<?x?x?xf32>)
|
|
// NAMED-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
|
|
// NAMED: %[[RES9:.+]] = linalg.rsqrt
|
|
// NAMED-SAME: ins(%[[RES8]] : tensor<?x?x?xf32>)
|
|
// NAMED-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
|
|
// NAMED: %[[RES10:.+]] = linalg.square
|
|
// NAMED-SAME: ins(%[[RES9]] : tensor<?x?x?xf32>)
|
|
// NAMED-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
|
|
// NAMED: %[[RES11:.+]] = linalg.tanh
|
|
// NAMED-SAME: ins(%[[RES10]] : tensor<?x?x?xf32>)
|
|
// NAMED-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
|
|
// NAMED: %[[RES12:.+]] = linalg.erf
|
|
// NAMED-SAME: ins(%[[RES11]] : tensor<?x?x?xf32>)
|
|
// NAMED-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
|
|
|
|
// CATEGORY: %[[RES0:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<exp>
|
|
// CATEGORY-SAME: ins(%[[A]] : tensor<?x?x?xf32>)
|
|
// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
|
|
// CATEGORY: %[[RES1:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<log>
|
|
// CATEGORY-SAME: ins(%[[RES0]] : tensor<?x?x?xf32>)
|
|
// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
|
|
// CATEGORY: %[[RES2:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<abs>
|
|
// CATEGORY-SAME: ins(%[[RES1]] : tensor<?x?x?xf32>)
|
|
// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
|
|
// CATEGORY: %[[RES3:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<ceil>
|
|
// CATEGORY-SAME: ins(%[[RES2]] : tensor<?x?x?xf32>)
|
|
// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
|
|
// CATEGORY: %[[RES4:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<floor>
|
|
// CATEGORY-SAME: ins(%[[RES3]] : tensor<?x?x?xf32>)
|
|
// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
|
|
// CATEGORY: %[[RES5:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<negf>
|
|
// CATEGORY-SAME: ins(%[[RES4]] : tensor<?x?x?xf32>)
|
|
// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
|
|
// CATEGORY: %[[RES6:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<reciprocal>
|
|
// CATEGORY-SAME: ins(%[[RES5]] : tensor<?x?x?xf32>)
|
|
// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
|
|
// CATEGORY: %[[RES7:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<round>
|
|
// CATEGORY-SAME: ins(%[[RES6]] : tensor<?x?x?xf32>)
|
|
// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
|
|
// CATEGORY: %[[RES8:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<sqrt>
|
|
// CATEGORY-SAME: ins(%[[RES7]] : tensor<?x?x?xf32>)
|
|
// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
|
|
// CATEGORY: %[[RES9:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<rsqrt>
|
|
// CATEGORY-SAME: ins(%[[RES8]] : tensor<?x?x?xf32>)
|
|
// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
|
|
// CATEGORY: %[[RES10:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<square>
|
|
// CATEGORY-SAME: ins(%[[RES9]] : tensor<?x?x?xf32>)
|
|
// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
|
|
// CATEGORY: %[[RES11:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<tanh>
|
|
// CATEGORY-SAME: ins(%[[RES10]] : tensor<?x?x?xf32>)
|
|
// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
|
|
// CATEGORY: %[[RES12:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<erf>
|
|
// CATEGORY-SAME: ins(%[[RES11]] : tensor<?x?x?xf32>)
|
|
// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
|
|
|
|
// -----
|
|
|
|
func.func @unary_ops_non_identity(%A: tensor<?xf32>, %Out: tensor<?x?xf32>) -> tensor<?x?xf32> {
|
|
%0 = linalg.generic
|
|
{indexing_maps = [affine_map<(d0, d1) -> (d1)>, affine_map<(d0, d1) -> (d1, d0)>],
|
|
iterator_types = ["parallel", "parallel"]}
|
|
ins(%A : tensor<?xf32>)
|
|
outs(%Out : tensor<?x?xf32>) {
|
|
^bb0(%in: f32, %out: f32):
|
|
%v = math.exp %in : f32
|
|
linalg.yield %v : f32
|
|
} -> tensor<?x?xf32>
|
|
return %0 : tensor<?x?xf32>
|
|
}
|
|
|
|
|
|
// ALL-DAG: #[[MAP_BC:.+]] = affine_map<(d0, d1) -> (d1)>
|
|
// ALL-DAG: #[[MAP_TP:.+]] = affine_map<(d0, d1) -> (d1, d0)>
|
|
// ALL: unary_ops_non_identity
|
|
// ALL-SAME: %[[A:.+]]: tensor<?xf32>, %[[OUT:.+]]: tensor<?x?xf32>) -> tensor<?x?xf32>
|
|
|
|
// Named ops cannot carry user-defined indexing maps -> expect no change.
|
|
// NAMED-NOT: linalg.exp
|
|
// NAMED: linalg.generic
|
|
|
|
// CATEGORY-NOT: linalg.generic
|
|
// CATEGORY: linalg.elementwise kind=#linalg.elementwise_kind<exp>
|
|
// CATEGORY-SAME: indexing_maps = [#[[MAP_BC]], #[[MAP_TP]]]
|
|
// CATEGORY-SAME: ins(%[[A]] : tensor<?xf32>)
|
|
// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?xf32>) -> tensor<?x?xf32>
|
|
|
|
// -----
|
|
|
|
#map = affine_map<(d0, d1) -> (d0, d1)>
|
|
func.func @binary_ops_int(%A: tensor<?x?xi32>, %B: tensor<?x?xi32>,
|
|
%Out: tensor<?x?xi32>) -> tensor<?x?xi32> {
|
|
%0 = linalg.generic
|
|
{indexing_maps = [#map, #map, #map],
|
|
iterator_types = ["parallel", "parallel"]}
|
|
ins(%A, %B : tensor<?x?xi32>, tensor<?x?xi32>)
|
|
outs(%Out : tensor<?x?xi32>) {
|
|
^bb0(%in: i32, %in_0: i32, %out: i32):
|
|
%v = arith.addi %in, %in_0 : i32
|
|
linalg.yield %v : i32
|
|
} -> tensor<?x?xi32>
|
|
%1 = linalg.generic
|
|
{indexing_maps = [#map, #map, #map],
|
|
iterator_types = ["parallel", "parallel"]}
|
|
ins(%0, %B : tensor<?x?xi32>, tensor<?x?xi32>)
|
|
outs(%Out : tensor<?x?xi32>) {
|
|
^bb0(%in: i32, %in_0: i32, %out: i32):
|
|
%v = arith.subi %in, %in_0 : i32
|
|
linalg.yield %v : i32
|
|
} -> tensor<?x?xi32>
|
|
%2 = linalg.generic
|
|
{indexing_maps = [#map, #map, #map],
|
|
iterator_types = ["parallel", "parallel"]}
|
|
ins(%1, %B : tensor<?x?xi32>, tensor<?x?xi32>)
|
|
outs(%Out : tensor<?x?xi32>) {
|
|
^bb0(%in: i32, %in_0: i32, %out: i32):
|
|
%v = arith.muli %in, %in_0 : i32
|
|
linalg.yield %v : i32
|
|
} -> tensor<?x?xi32>
|
|
%3 = linalg.generic
|
|
{indexing_maps = [#map, #map, #map],
|
|
iterator_types = ["parallel", "parallel"]}
|
|
ins(%2, %B : tensor<?x?xi32>, tensor<?x?xi32>)
|
|
outs(%Out : tensor<?x?xi32>) {
|
|
^bb0(%in: i32, %in_0: i32, %out: i32):
|
|
%v = arith.divsi %in, %in_0 : i32
|
|
linalg.yield %v : i32
|
|
} -> tensor<?x?xi32>
|
|
%4 = linalg.generic
|
|
{indexing_maps = [#map, #map, #map],
|
|
iterator_types = ["parallel", "parallel"]}
|
|
ins(%3, %B : tensor<?x?xi32>, tensor<?x?xi32>)
|
|
outs(%Out : tensor<?x?xi32>) {
|
|
^bb0(%in: i32, %in_0: i32, %out: i32):
|
|
%v = arith.divui %in, %in_0 : i32
|
|
linalg.yield %v : i32
|
|
} -> tensor<?x?xi32>
|
|
%5 = linalg.generic
|
|
{indexing_maps = [#map, #map, #map],
|
|
iterator_types = ["parallel", "parallel"]}
|
|
ins(%4, %B : tensor<?x?xi32>, tensor<?x?xi32>)
|
|
outs(%Out : tensor<?x?xi32>) {
|
|
^bb0(%in: i32, %in_0: i32, %out: i32):
|
|
%v = arith.maxsi %in, %in_0 : i32
|
|
linalg.yield %v : i32
|
|
} -> tensor<?x?xi32>
|
|
%6 = linalg.generic
|
|
{indexing_maps = [#map, #map, #map],
|
|
iterator_types = ["parallel", "parallel"]}
|
|
ins(%5, %B : tensor<?x?xi32>, tensor<?x?xi32>)
|
|
outs(%Out : tensor<?x?xi32>) {
|
|
^bb0(%in: i32, %in_0: i32, %out: i32):
|
|
%v = arith.minsi %in, %in_0 : i32
|
|
linalg.yield %v : i32
|
|
} -> tensor<?x?xi32>
|
|
return %6 : tensor<?x?xi32>
|
|
}
|
|
|
|
// ALL-LABEL: binary_ops_int
|
|
// ALL-SAME: %[[A:.+]]: [[TTY:tensor<\?x\?xi32>]], %[[B:.+]]: [[TTY]],
|
|
// ALL-SAME: %[[OUT:.+]]: [[TTY]]) -> [[TTY]]
|
|
|
|
// NAMED-NOT: linalg.generic
|
|
// NAMED: %[[RES0:.+]] = linalg.add
|
|
// NAMED-SAME: ins(%[[A]], %[[B]] : [[TTY]], [[TTY]])
|
|
// NAMED-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
|
|
// NAMED: %[[RES1:.+]] = linalg.sub
|
|
// NAMED-SAME: ins(%[[RES0]], %[[B]] : [[TTY]], [[TTY]])
|
|
// NAMED-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
|
|
// NAMED: %[[RES2:.+]] = linalg.mul
|
|
// NAMED-SAME: ins(%[[RES1]], %[[B]] : [[TTY]], [[TTY]])
|
|
// NAMED-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
|
|
// NAMED: %[[RES3:.+]] = linalg.div
|
|
// NAMED-SAME: ins(%[[RES2]], %[[B]] : [[TTY]], [[TTY]])
|
|
// NAMED-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
|
|
// NAMED: %[[RES4:.+]] = linalg.div_unsigned
|
|
// NAMED-SAME: ins(%[[RES3]], %[[B]] : [[TTY]], [[TTY]])
|
|
// NAMED-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
|
|
// NAMED: %[[RES5:.+]] = linalg.max
|
|
// NAMED-SAME: ins(%[[RES4]], %[[B]] : [[TTY]], [[TTY]])
|
|
// NAMED-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
|
|
// NAMED: %[[RES6:.+]] = linalg.min
|
|
// NAMED-SAME: ins(%[[RES5]], %[[B]] : [[TTY]], [[TTY]])
|
|
// NAMED-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
|
|
|
|
// CATEGORY-NOT: linalg.generic
|
|
// CATEGORY: %[[RES0:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<add>
|
|
// CATEGORY-SAME: ins(%[[A]], %[[B]] : [[TTY]], [[TTY]])
|
|
// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
|
|
// CATEGORY: %[[RES1:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<sub>
|
|
// CATEGORY-SAME: ins(%[[RES0]], %[[B]] : [[TTY]], [[TTY]])
|
|
// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
|
|
// CATEGORY: %[[RES2:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<mul>
|
|
// CATEGORY-SAME: ins(%[[RES1]], %[[B]] : [[TTY]], [[TTY]])
|
|
// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
|
|
// CATEGORY: %[[RES3:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<div>
|
|
// CATEGORY-SAME: ins(%[[RES2]], %[[B]] : [[TTY]], [[TTY]])
|
|
// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
|
|
// CATEGORY: %[[RES4:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<div_unsigned>
|
|
// CATEGORY-SAME: ins(%[[RES3]], %[[B]] : [[TTY]], [[TTY]])
|
|
// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
|
|
// CATEGORY: %[[RES5:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<max_signed>
|
|
// CATEGORY-SAME: ins(%[[RES4]], %[[B]] : [[TTY]], [[TTY]])
|
|
// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
|
|
// CATEGORY: %[[RES6:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<min_signed>
|
|
// CATEGORY-SAME: ins(%[[RES5]], %[[B]] : [[TTY]], [[TTY]])
|
|
// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
|
|
|
|
// -----
|
|
|
|
#map = affine_map<(d0, d1) -> (d0, d1)>
|
|
func.func @binary_ops_float(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>,
|
|
%Out: tensor<?x?xf32>) -> tensor<?x?xf32> {
|
|
%0 = linalg.generic
|
|
{indexing_maps = [#map, #map, #map],
|
|
iterator_types = ["parallel", "parallel"]}
|
|
ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>)
|
|
outs(%Out : tensor<?x?xf32>) {
|
|
^bb0(%in: f32, %in_0: f32, %out: f32):
|
|
%v = arith.addf %in, %in_0 : f32
|
|
linalg.yield %v : f32
|
|
} -> tensor<?x?xf32>
|
|
%1 = linalg.generic
|
|
{indexing_maps = [#map, #map, #map],
|
|
iterator_types = ["parallel", "parallel"]}
|
|
ins(%0, %B : tensor<?x?xf32>, tensor<?x?xf32>)
|
|
outs(%Out : tensor<?x?xf32>) {
|
|
^bb0(%in: f32, %in_0: f32, %out: f32):
|
|
%v = arith.subf %in, %in_0 : f32
|
|
linalg.yield %v : f32
|
|
} -> tensor<?x?xf32>
|
|
%2 = linalg.generic
|
|
{indexing_maps = [#map, #map, #map],
|
|
iterator_types = ["parallel", "parallel"]}
|
|
ins(%1, %B : tensor<?x?xf32>, tensor<?x?xf32>)
|
|
outs(%Out : tensor<?x?xf32>) {
|
|
^bb0(%in: f32, %in_0: f32, %out: f32):
|
|
%v = arith.mulf %in, %in_0 : f32
|
|
linalg.yield %v : f32
|
|
} -> tensor<?x?xf32>
|
|
%3 = linalg.generic
|
|
{indexing_maps = [#map, #map, #map],
|
|
iterator_types = ["parallel", "parallel"]}
|
|
ins(%2, %B : tensor<?x?xf32>, tensor<?x?xf32>)
|
|
outs(%Out : tensor<?x?xf32>) {
|
|
^bb0(%in: f32, %in_0: f32, %out: f32):
|
|
%v = arith.divf %in, %in_0 : f32
|
|
linalg.yield %v : f32
|
|
} -> tensor<?x?xf32>
|
|
%4 = linalg.generic
|
|
{indexing_maps = [#map, #map, #map],
|
|
iterator_types = ["parallel", "parallel"]}
|
|
ins(%3, %B : tensor<?x?xf32>, tensor<?x?xf32>)
|
|
outs(%Out : tensor<?x?xf32>) {
|
|
^bb0(%in: f32, %in_0: f32, %out: f32):
|
|
%v = arith.maximumf %in, %in_0 : f32
|
|
linalg.yield %v : f32
|
|
} -> tensor<?x?xf32>
|
|
%5 = linalg.generic
|
|
{indexing_maps = [#map, #map, #map],
|
|
iterator_types = ["parallel", "parallel"]}
|
|
ins(%4, %B : tensor<?x?xf32>, tensor<?x?xf32>)
|
|
outs(%Out : tensor<?x?xf32>) {
|
|
^bb0(%in: f32, %in_0: f32, %out: f32):
|
|
%v = arith.minimumf %in, %in_0 : f32
|
|
linalg.yield %v : f32
|
|
} -> tensor<?x?xf32>
|
|
%6 = linalg.generic
|
|
{indexing_maps = [#map, #map, #map],
|
|
iterator_types = ["parallel", "parallel"]}
|
|
ins(%5, %B : tensor<?x?xf32>, tensor<?x?xf32>)
|
|
outs(%Out : tensor<?x?xf32>) {
|
|
^bb0(%in: f32, %in_0: f32, %out: f32):
|
|
%v = math.powf %in, %in_0 : f32
|
|
linalg.yield %v : f32
|
|
} -> tensor<?x?xf32>
|
|
return %6 : tensor<?x?xf32>
|
|
}
|
|
|
|
// ALL-LABEL: binary_ops_float
|
|
// ALL-SAME: %[[A:.+]]: [[TTY:tensor<\?x\?xf32>]], %[[B:.+]]: [[TTY]],
|
|
// ALL-SAME: %[[OUT:.+]]: [[TTY]]) -> [[TTY]]
|
|
|
|
// NAMED-NOT: linalg.generic
|
|
// NAMED: %[[RES0:.+]] = linalg.add
|
|
// NAMED-SAME: ins(%[[A]], %[[B]] : [[TTY]], [[TTY]])
|
|
// NAMED-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
|
|
// NAMED: %[[RES1:.+]] = linalg.sub
|
|
// NAMED-SAME: ins(%[[RES0]], %[[B]] : [[TTY]], [[TTY]])
|
|
// NAMED-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
|
|
// NAMED: %[[RES2:.+]] = linalg.mul
|
|
// NAMED-SAME: ins(%[[RES1]], %[[B]] : [[TTY]], [[TTY]])
|
|
// NAMED-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
|
|
// NAMED: %[[RES3:.+]] = linalg.div
|
|
// NAMED-SAME: ins(%[[RES2]], %[[B]] : [[TTY]], [[TTY]])
|
|
// NAMED-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
|
|
// NAMED: %[[RES4:.+]] = linalg.max
|
|
// NAMED-SAME: ins(%[[RES3]], %[[B]] : [[TTY]], [[TTY]])
|
|
// NAMED-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
|
|
// NAMED: %[[RES5:.+]] = linalg.min
|
|
// NAMED-SAME: ins(%[[RES4]], %[[B]] : [[TTY]], [[TTY]])
|
|
// NAMED-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
|
|
// NAMED: %[[RES6:.+]] = linalg.powf
|
|
// NAMED-SAME: ins(%[[RES5]], %[[B]] : [[TTY]], [[TTY]])
|
|
// NAMED-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
|
|
|
|
// CATEGORY-NOT: linalg.generic
|
|
// CATEGORY: %[[RES0:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<add>
|
|
// CATEGORY-SAME: ins(%[[A]], %[[B]] : [[TTY]], [[TTY]])
|
|
// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
|
|
// CATEGORY: %[[RES1:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<sub>
|
|
// CATEGORY-SAME: ins(%[[RES0]], %[[B]] : [[TTY]], [[TTY]])
|
|
// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
|
|
// CATEGORY: %[[RES2:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<mul>
|
|
// CATEGORY-SAME: ins(%[[RES1]], %[[B]] : [[TTY]], [[TTY]])
|
|
// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
|
|
// CATEGORY: %[[RES3:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<div>
|
|
// CATEGORY-SAME: ins(%[[RES2]], %[[B]] : [[TTY]], [[TTY]])
|
|
// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
|
|
// CATEGORY: %[[RES4:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<max_signed>
|
|
// CATEGORY-SAME: ins(%[[RES3]], %[[B]] : [[TTY]], [[TTY]])
|
|
// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
|
|
// CATEGORY: %[[RES5:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<min_signed>
|
|
// CATEGORY-SAME: ins(%[[RES4]], %[[B]] : [[TTY]], [[TTY]])
|
|
// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
|
|
// CATEGORY: %[[RES6:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<powf>
|
|
// CATEGORY-SAME: ins(%[[RES5]], %[[B]] : [[TTY]], [[TTY]])
|
|
// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
|
|
|
|
// -----
|
|
|
|
#map = affine_map<(d0, d1) -> (d0, d1)>
|
|
func.func @binary_ops_complex(%A: tensor<?x?xcomplex<f32>>,
|
|
%B: tensor<?x?xcomplex<f32>>,
|
|
%Out: tensor<?x?xcomplex<f32>>)
|
|
-> tensor<?x?xcomplex<f32>> {
|
|
%0 = linalg.generic
|
|
{indexing_maps = [#map, #map, #map],
|
|
iterator_types = ["parallel", "parallel"]}
|
|
ins(%A, %B : tensor<?x?xcomplex<f32>>, tensor<?x?xcomplex<f32>>)
|
|
outs(%Out : tensor<?x?xcomplex<f32>>) {
|
|
^bb0(%in: complex<f32>, %in_0: complex<f32>, %out: complex<f32>):
|
|
%v = complex.add %in, %in_0 : complex<f32>
|
|
linalg.yield %v : complex<f32>
|
|
} -> tensor<?x?xcomplex<f32>>
|
|
%1 = linalg.generic
|
|
{indexing_maps = [#map, #map, #map],
|
|
iterator_types = ["parallel", "parallel"]}
|
|
ins(%0, %B : tensor<?x?xcomplex<f32>>, tensor<?x?xcomplex<f32>>)
|
|
outs(%Out : tensor<?x?xcomplex<f32>>) {
|
|
^bb0(%in: complex<f32>, %in_0: complex<f32>, %out: complex<f32>):
|
|
%v = complex.sub %in, %in_0 : complex<f32>
|
|
linalg.yield %v : complex<f32>
|
|
} -> tensor<?x?xcomplex<f32>>
|
|
%2 = linalg.generic
|
|
{indexing_maps = [#map, #map, #map],
|
|
iterator_types = ["parallel", "parallel"]}
|
|
ins(%1, %B : tensor<?x?xcomplex<f32>>, tensor<?x?xcomplex<f32>>)
|
|
outs(%Out : tensor<?x?xcomplex<f32>>) {
|
|
^bb0(%in: complex<f32>, %in_0: complex<f32>, %out: complex<f32>):
|
|
%v = complex.mul %in, %in_0 : complex<f32>
|
|
linalg.yield %v : complex<f32>
|
|
} -> tensor<?x?xcomplex<f32>>
|
|
%3 = linalg.generic
|
|
{indexing_maps = [#map, #map, #map],
|
|
iterator_types = ["parallel", "parallel"]}
|
|
ins(%2, %B : tensor<?x?xcomplex<f32>>, tensor<?x?xcomplex<f32>>)
|
|
outs(%Out : tensor<?x?xcomplex<f32>>) {
|
|
^bb0(%in: complex<f32>, %in_0: complex<f32>, %out: complex<f32>):
|
|
%v = complex.div %in, %in_0 : complex<f32>
|
|
linalg.yield %v : complex<f32>
|
|
} -> tensor<?x?xcomplex<f32>>
|
|
return %3 : tensor<?x?xcomplex<f32>>
|
|
}
|
|
|
|
// ALL-LABEL: binary_ops_complex
|
|
// ALL-SAME: %[[A:.+]]: [[TTY:tensor<\?x\?xcomplex<f32>>]], %[[B:.+]]: [[TTY]],
|
|
// ALL-SAME: %[[OUT:.+]]: [[TTY]]) -> [[TTY]]
|
|
|
|
// NAMED-NOT: linalg.generic
|
|
// NAMED: %[[RES0:.+]] = linalg.add
|
|
// NAMED-SAME: ins(%[[A]], %[[B]] : [[TTY]], [[TTY]])
|
|
// NAMED-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
|
|
// NAMED: %[[RES1:.+]] = linalg.sub
|
|
// NAMED-SAME: ins(%[[RES0]], %[[B]] : [[TTY]], [[TTY]])
|
|
// NAMED-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
|
|
// NAMED: %[[RES2:.+]] = linalg.mul
|
|
// NAMED-SAME: ins(%[[RES1]], %[[B]] : [[TTY]], [[TTY]])
|
|
// NAMED-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
|
|
// NAMED: %[[RES3:.+]] = linalg.div
|
|
// NAMED-SAME: ins(%[[RES2]], %[[B]] : [[TTY]], [[TTY]])
|
|
// NAMED-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
|
|
|
|
// CATEGORY-NOT: linalg.generic
|
|
// CATEGORY: %[[RES0:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<add>
|
|
// CATEGORY-SAME: ins(%[[A]], %[[B]] : [[TTY]], [[TTY]])
|
|
// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
|
|
// CATEGORY: %[[RES1:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<sub>
|
|
// CATEGORY-SAME: ins(%[[RES0]], %[[B]] : [[TTY]], [[TTY]])
|
|
// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
|
|
// CATEGORY: %[[RES2:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<mul>
|
|
// CATEGORY-SAME: ins(%[[RES1]], %[[B]] : [[TTY]], [[TTY]])
|
|
// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
|
|
// CATEGORY: %[[RES3:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<div>
|
|
// CATEGORY-SAME: ins(%[[RES2]], %[[B]] : [[TTY]], [[TTY]])
|
|
// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
|
|
|
|
// -----
|
|
|
|
#map = affine_map<(d0, d1) -> (d0, d1)>
|
|
func.func @binary_ops_bool(%A: tensor<?x?xi1>, %B: tensor<?x?xi1>,
|
|
%Out: tensor<?x?xi1>) -> tensor<?x?xi1> {
|
|
%0 = linalg.generic
|
|
{indexing_maps = [#map, #map, #map],
|
|
iterator_types = ["parallel", "parallel"]}
|
|
ins(%A, %B : tensor<?x?xi1>, tensor<?x?xi1>)
|
|
outs(%Out : tensor<?x?xi1>) {
|
|
^bb0(%in: i1, %in_0: i1, %out: i1):
|
|
%v = arith.ori %in, %in_0 : i1
|
|
linalg.yield %v : i1
|
|
} -> tensor<?x?xi1>
|
|
%1 = linalg.generic
|
|
{indexing_maps = [#map, #map, #map],
|
|
iterator_types = ["parallel", "parallel"]}
|
|
ins(%0, %B : tensor<?x?xi1>, tensor<?x?xi1>)
|
|
outs(%Out : tensor<?x?xi1>) {
|
|
^bb0(%in: i1, %in_0: i1, %out: i1):
|
|
%v = arith.andi %in, %in_0 : i1
|
|
linalg.yield %v : i1
|
|
} -> tensor<?x?xi1>
|
|
return %1 : tensor<?x?xi1>
|
|
}
|
|
|
|
// ALL-LABEL: binary_ops_bool
|
|
// ALL-SAME: %[[A:.+]]: [[TTY:tensor<\?x\?xi1>]], %[[B:.+]]: [[TTY]],
|
|
// ALL-SAME: %[[OUT:.+]]: [[TTY]]) -> [[TTY]]
|
|
|
|
// NAMED-NOT: linalg.generic
|
|
// NAMED: %[[RES0:.+]] = linalg.add
|
|
// NAMED-SAME: ins(%[[A]], %[[B]] : [[TTY]], [[TTY]])
|
|
// NAMED-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
|
|
// NAMED: %[[RES1:.+]] = linalg.mul
|
|
// NAMED-SAME: ins(%[[RES0]], %[[B]] : [[TTY]], [[TTY]])
|
|
// NAMED-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
|
|
|
|
// CATEGORY-NOT: linalg.generic
|
|
// CATEGORY: %[[RES0:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<add>
|
|
// CATEGORY-SAME: ins(%[[A]], %[[B]] : [[TTY]], [[TTY]])
|
|
// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
|
|
// CATEGORY: %[[RES1:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<mul>
|
|
// CATEGORY-SAME: ins(%[[RES0]], %[[B]] : [[TTY]], [[TTY]])
|
|
// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
|
|
|
|
// -----
|
|
|
|
#map = affine_map<(d0, d1) -> (d0, d1)>
|
|
func.func @binary_ops_uint(%A: tensor<?x?xi32>, %B: tensor<?x?xi32>,
|
|
%Out: tensor<?x?xi32>) -> tensor<?x?xi32> {
|
|
%0 = linalg.generic
|
|
{indexing_maps = [#map, #map, #map],
|
|
iterator_types = ["parallel", "parallel"]}
|
|
ins(%A, %B : tensor<?x?xi32>, tensor<?x?xi32>)
|
|
outs(%Out : tensor<?x?xi32>) {
|
|
^bb0(%in: i32, %in_0: i32, %out: i32):
|
|
%v = arith.maxui %in, %in_0 : i32
|
|
linalg.yield %v : i32
|
|
} -> tensor<?x?xi32>
|
|
%1 = linalg.generic
|
|
{indexing_maps = [#map, #map, #map],
|
|
iterator_types = ["parallel", "parallel"]}
|
|
ins(%0, %B : tensor<?x?xi32>, tensor<?x?xi32>)
|
|
outs(%Out : tensor<?x?xi32>) {
|
|
^bb0(%in: i32, %in_0: i32, %out: i32):
|
|
%v = arith.minui %in, %in_0 : i32
|
|
linalg.yield %v : i32
|
|
} -> tensor<?x?xi32>
|
|
return %1 : tensor<?x?xi32>
|
|
}
|
|
|
|
// ALL-LABEL: binary_ops_uint
|
|
// ALL-SAME: %[[A:.+]]: [[TTY:tensor<\?x\?xi32>]], %[[B:.+]]: [[TTY]],
|
|
// ALL-SAME: %[[OUT:.+]]: [[TTY]]) -> [[TTY]]
|
|
|
|
// No named ops yet for unsigned max/min -> expect no change.
|
|
// NAMED-NOT: linalg.{{max|min}}
|
|
// NAMED: linalg.generic
|
|
|
|
// CATEGORY-NOT: linalg.generic
|
|
// CATEGORY: %[[RES0:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<max_unsigned>
|
|
// CATEGORY-SAME: ins(%[[A]], %[[B]] : [[TTY]], [[TTY]])
|
|
// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
|
|
// CATEGORY: %[[RES1:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<min_unsigned>
|
|
// CATEGORY-SAME: ins(%[[RES0]], %[[B]] : [[TTY]], [[TTY]])
|
|
// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
|
|
|
|
// -----
|
|
|
|
func.func @binary_ops_non_identity(%A: tensor<?xf32>, %B: tensor<?x?xf32>,
|
|
%Out: tensor<?x?xf32>) -> tensor<?x?xf32> {
|
|
%0 = linalg.generic
|
|
{indexing_maps = [affine_map<(d0, d1) -> (d1)>, affine_map<(d0, d1) -> (d1, d0)>,
|
|
affine_map<(d0, d1) -> (d0, d1)>],
|
|
iterator_types = ["parallel", "parallel"]}
|
|
ins(%A, %B : tensor<?xf32>, tensor<?x?xf32>)
|
|
outs(%Out : tensor<?x?xf32>) {
|
|
^bb0(%in: f32, %in_0: f32, %out: f32):
|
|
%v = arith.addf %in, %in_0 : f32
|
|
linalg.yield %v : f32
|
|
} -> tensor<?x?xf32>
|
|
return %0 : tensor<?x?xf32>
|
|
}
|
|
|
|
// ALL-DAG: #[[MAP_BC:.+]] = affine_map<(d0, d1) -> (d1)>
|
|
// ALL-DAG: #[[MAP_TP:.+]] = affine_map<(d0, d1) -> (d1, d0)>
|
|
// ALL-DAG: #[[MAP_ID:.+]] = affine_map<(d0, d1) -> (d0, d1)>
|
|
// ALL: binary_ops_non_identity
|
|
// ALL-SAME: %[[A:.+]]: [[TTY1D:tensor<\?xf32>]], %[[B:.+]]: [[TTY:tensor<\?x\?xf32>]],
|
|
// ALL-SAME: %[[OUT:.+]]: [[TTY]]) -> [[TTY]]
|
|
|
|
// Named ops cannot carry user-defined indexing maps -> expect no change.
|
|
// NAMED-NOT: linalg.add
|
|
// NAMED: linalg.generic
|
|
|
|
// CATEGORY-NOT: linalg.generic
|
|
// CATEGORY: linalg.elementwise kind=#linalg.elementwise_kind<add>
|
|
// CATEGORY-SAME: indexing_maps = [#[[MAP_BC]], #[[MAP_TP]], #[[MAP_ID]]]
|
|
// CATEGORY-SAME: ins(%[[A]], %[[B]] : [[TTY1D]], [[TTY]])
|
|
// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
|
|
|
|
// -----
|
|
|
|
#map = affine_map<(d0, d1) -> (d0, d1)>
|
|
#bcast = affine_map<(d0, d1) -> (d0)>
|
|
func.func @binary_ops_swapped(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>,
|
|
%C: tensor<?xf32>, %Out: tensor<?x?xf32>)
|
|
-> tensor<?x?xf32> {
|
|
%0 = linalg.generic
|
|
{indexing_maps = [#map, #map, #map],
|
|
iterator_types = ["parallel", "parallel"]}
|
|
ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>)
|
|
outs(%Out : tensor<?x?xf32>) {
|
|
^bb0(%in: f32, %in_0: f32, %out: f32):
|
|
%v = arith.addf %in_0, %in : f32
|
|
linalg.yield %v : f32
|
|
} -> tensor<?x?xf32>
|
|
%1 = linalg.generic
|
|
{indexing_maps = [#map, #bcast, #map],
|
|
iterator_types = ["parallel", "parallel"]}
|
|
ins(%0, %C : tensor<?x?xf32>, tensor<?xf32>)
|
|
outs(%Out : tensor<?x?xf32>) {
|
|
^bb0(%in: f32, %in_0: f32, %out: f32):
|
|
%v = arith.subf %in_0, %in : f32
|
|
linalg.yield %v : f32
|
|
} -> tensor<?x?xf32>
|
|
return %1 : tensor<?x?xf32>
|
|
}
|
|
|
|
// ALL-DAG: #[[MAP_BC:.+]] = affine_map<(d0, d1) -> (d0)>
|
|
// ALL-DAG: #[[MAP_ID:.+]] = affine_map<(d0, d1) -> (d0, d1)>
|
|
// ALL: binary_ops_swapped
|
|
// ALL-SAME: %[[A:.+]]: [[TTY:tensor<\?x\?xf32>]], %[[B:.+]]: [[TTY]],
|
|
// ALL-SAME: %[[C:.+]]: [[TTY1D:tensor<\?xf32>]],
|
|
// ALL-SAME: %[[OUT:.+]]: [[TTY]]) -> [[TTY]]
|
|
|
|
// NAMED: %[[RES0:.+]] = linalg.add
|
|
// NAMED-SAME: ins(%[[B]], %[[A]] : [[TTY]], [[TTY]])
|
|
// NAMED-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
|
|
// NAMED-NOT: linalg.sub
|
|
// NAMED: linalg.generic
|
|
// NAMED-SAME: ins(%[[RES0]], %[[C]] : [[TTY]], [[TTY1D]])
|
|
|
|
// CATEGORY-NOT: linalg.generic
|
|
// CATEGORY: %[[RES0:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<add>
|
|
// CATEGORY-SAME: ins(%[[B]], %[[A]] : [[TTY]], [[TTY]])
|
|
// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
|
|
// CATEGORY: %[[RES1:.+]] = linalg.elementwise kind=#linalg.elementwise_kind<sub>
|
|
// CATEGORY-SAME: indexing_maps = [#[[MAP_BC]], #[[MAP_ID]], #[[MAP_ID]]]
|
|
// CATEGORY-SAME: ins(%[[C]], %[[RES0]] : [[TTY1D]], [[TTY]])
|
|
// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
|
|
|
|
// -----
|
|
|
|
#map = affine_map<(d0) -> (d0)>
|
|
func.func @unary_op_with_scalar(%A: tensor<?xi32>, %Out: tensor<?xi32>)
|
|
-> tensor<?xi32> {
|
|
%cst = arith.constant 123 : i32
|
|
%0 = linalg.generic
|
|
{indexing_maps = [#map, #map],
|
|
iterator_types = ["parallel"]}
|
|
ins(%A : tensor<?xi32>)
|
|
outs(%Out : tensor<?xi32>) {
|
|
^bb0(%in: i32, %out: i32):
|
|
%v = arith.addi %cst, %in : i32
|
|
linalg.yield %v : i32
|
|
} -> tensor<?xi32>
|
|
return %0 : tensor<?xi32>
|
|
}
|
|
|
|
// CATEGORY-DAG: #[[MAP_ID:.+]] = affine_map<(d0) -> (d0)>
|
|
// CATEGORY-DAG: #[[MAP_BC:.+]] = affine_map<(d0) -> ()>
|
|
// ALL: unary_op_with_scalar
|
|
// CATEGORY-SAME: %[[A:.+]]: [[TTY:tensor<\?xi32>]],
|
|
// CATEGORY-SAME: %[[OUT:.+]]: [[TTY]]) -> [[TTY]]
|
|
|
|
// Named ops cannot broadcast from a scalar operand -> expect no change.
|
|
// NAMED-NOT: linalg.add
|
|
// NAMED: linalg.generic
|
|
|
|
// CATEGORY-NOT: linalg.generic
|
|
// CATEGORY: %[[CST:.+]] = arith.constant 123 : i32
|
|
// CATEGORY: linalg.elementwise kind=#linalg.elementwise_kind<add>
|
|
// CATEGORY-SAME: indexing_maps = [#[[MAP_BC]], #[[MAP_ID]], #[[MAP_ID]]]
|
|
// CATEGORY-SAME: ins(%[[CST]], %[[A]] : i32, [[TTY]])
|
|
// CATEGORY-SAME: outs(%[[OUT]] : [[TTY]]) -> [[TTY]]
|
|
|
|
// -----
|
|
|
|
#map = affine_map<(d0) -> (d0)>
|
|
func.func @negative_unary_op_using_block_arg_twice(%A: tensor<?xi32>,
|
|
%Out: tensor<?xi32>)
|
|
-> tensor<?xi32> {
|
|
%0 = linalg.generic
|
|
{indexing_maps = [#map, #map],
|
|
iterator_types = ["parallel"]}
|
|
ins(%A : tensor<?xi32>)
|
|
outs(%Out : tensor<?xi32>) {
|
|
^bb0(%in: i32, %out: i32):
|
|
%v = arith.addi %in, %in : i32
|
|
linalg.yield %v : i32
|
|
} -> tensor<?xi32>
|
|
return %0 : tensor<?xi32>
|
|
}
|
|
|
|
// ALL-LABEL: negative_unary_op_using_block_arg_twice
|
|
|
|
// Named ops cannot broadcast from a scalar operand -> expect no change.
|
|
// NAMED-NOT: linalg.add
|
|
|
|
// There is no scalar operand to hoist -> expect no change.
|
|
// CATEGORY-NOT: linalg.elementwise kind=#linalg.elementwise_kind<add>
|
|
|
|
// ALL: linalg.generic
|
|
|
|
// -----
|
|
|
|
///----------------------------------------------------------------------------------------
|
|
/// Tests for linalg.matmul
|
|
///----------------------------------------------------------------------------------------
|
|
|
|
#map = affine_map<(d0, d1, d2) -> (d0, d2)>
|
|
#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
|
|
#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
|
|
func.func @op_matmul(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>,
|
|
%Out: tensor<?x?xf32>) -> tensor<?x?xf32> {
|
|
%0 = linalg.generic
|
|
{indexing_maps = [#map, #map1, #map2],
|
|
iterator_types = ["parallel", "parallel", "reduction"]}
|
|
ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>)
|
|
outs(%Out : tensor<?x?xf32>) {
|
|
^bb0(%in: f32, %in_0: f32, %out: f32):
|
|
%1 = arith.mulf %in, %in_0 : f32
|
|
%2 = arith.addf %out, %1 : f32
|
|
linalg.yield %2 : f32
|
|
} -> tensor<?x?xf32>
|
|
return %0 : tensor<?x?xf32>
|
|
}
|
|
|
|
// CATEGORY-DAG: #[[$MAP_A:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>
|
|
// CATEGORY-DAG: #[[$MAP_B:.+]] = affine_map<(d0, d1, d2) -> (d2, d1)>
|
|
// CATEGORY-DAG: #[[$MAP_C:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
|
|
|
|
// ALL-LABEL: op_matmul
|
|
// ALL-SAME: %[[A:.+]]: tensor<?x?xf32>, %[[B:.+]]: tensor<?x?xf32>,
|
|
// ALL-SAME: %[[OUT:.+]]: tensor<?x?xf32>) -> tensor<?x?xf32>
|
|
|
|
// NAMED-NOT: linalg.generic
|
|
// NAMED: linalg.matmul
|
|
// NAMED-SAME: ins(%[[A]], %[[B]] : tensor<?x?xf32>, tensor<?x?xf32>)
|
|
// NAMED-SAME: outs(%[[OUT]] : tensor<?x?xf32>) -> tensor<?x?xf32>
|
|
|
|
// CATEGORY-NOT: linalg.generic
|
|
// CATEGORY: linalg.contract
|
|
// CATEGORY-SAME: indexing_maps = {{\[}}#[[$MAP_A]], #[[$MAP_B]], #[[$MAP_C]]{{\]}}
|
|
// CATEGORY-SAME: ins(%[[A]], %[[B]] : tensor<?x?xf32>, tensor<?x?xf32>)
|
|
// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?xf32>) -> tensor<?x?xf32>
|
|
|
|
// Cast-auditing tests: ensure we only specialize when the cast semantics can
|
|
// be expressed by linalg.matmul, and use the cast attribute when needed.
|
|
|
|
// Check matmul with unsigned cast is correctly raised back to named op.
|
|
func.func @op_matmul_unsigned_cast(%A: tensor<16x8xi16>, %B: tensor<8x32xi32>,
|
|
%Out: tensor<16x32xi32>) -> tensor<16x32xi32> {
|
|
%0 = linalg.generic
|
|
{indexing_maps = [#map, #map1, #map2],
|
|
iterator_types = ["parallel", "parallel", "reduction"]}
|
|
ins(%A, %B : tensor<16x8xi16>, tensor<8x32xi32>)
|
|
outs(%Out : tensor<16x32xi32>) {
|
|
^bb0(%in: i16, %in_0: i32, %out: i32):
|
|
%1 = arith.extui %in : i16 to i32
|
|
%3 = arith.muli %1, %in_0 : i32
|
|
%4 = arith.addi %out, %3 : i32
|
|
linalg.yield %4 : i32
|
|
} -> tensor<16x32xi32>
|
|
return %0 : tensor<16x32xi32>
|
|
}
|
|
|
|
// ALL-LABEL: op_matmul_unsigned_cast
|
|
|
|
// NAMED-NOT: linalg.generic
|
|
// NAMED: linalg.matmul {cast = #linalg.type_fn<cast_unsigned>}
|
|
|
|
// CATEGORY-NOT: linalg.generic
|
|
// CATEGORY: linalg.contract
|
|
// CATEGORY-SAME: indexing_maps = {{\[}}#[[$MAP_A]], #[[$MAP_B]], #[[$MAP_C]]{{\]}}
|
|
// CATEGORY-SAME: {cast = #linalg.type_fn<cast_unsigned>}
|
|
|
|
// Ensures truncation rounding is tolerated with unsigned cases.
|
|
// Note: We only consider casts as conflicting if they have different
|
|
// signedness behaviours, and then we do not specialize if they do
|
|
// conflict. Since this is not such a case, we do not block specialization.
|
|
// Also the roundtrip lowering back to linalg.generic for such an op is
|
|
// expected to produce the same thing again, so we are not loosing
|
|
// information here.
|
|
func.func @op_matmul_unsigned_cast_and_truncate(%A: tensor<16x8xi16>, %B: tensor<8x32xi64>,
|
|
%Out: tensor<16x32xi32>) -> tensor<16x32xi32> {
|
|
%0 = linalg.generic
|
|
{indexing_maps = [#map, #map1, #map2],
|
|
iterator_types = ["parallel", "parallel", "reduction"]}
|
|
ins(%A, %B : tensor<16x8xi16>, tensor<8x32xi64>)
|
|
outs(%Out : tensor<16x32xi32>) {
|
|
^bb0(%in: i16, %in_0: i64, %out: i32):
|
|
%1 = arith.extui %in : i16 to i32
|
|
%2 = arith.trunci %in_0 : i64 to i32
|
|
%3 = arith.muli %1, %2 : i32
|
|
%4 = arith.addi %out, %3 : i32
|
|
linalg.yield %4 : i32
|
|
} -> tensor<16x32xi32>
|
|
return %0 : tensor<16x32xi32>
|
|
}
|
|
|
|
// ALL-LABEL: op_matmul_unsigned_cast_and_truncate
|
|
|
|
// NAMED-NOT: linalg.generic
|
|
// NAMED: linalg.matmul {cast = #linalg.type_fn<cast_unsigned>}
|
|
|
|
// CATEGORY-NOT: linalg.generic
|
|
// CATEGORY: linalg.contract
|
|
// CATEGORY-SAME: indexing_maps = {{\[}}#[[$MAP_A]], #[[$MAP_B]], #[[$MAP_C]]{{\]}}
|
|
// CATEGORY-SAME: {cast = #linalg.type_fn<cast_unsigned>}
|
|
|
|
// Signed casts are the default, no cast attribute is required.
|
|
func.func @op_matmul_signed_cast(%A: tensor<16x8xi16>, %B: tensor<8x32xi16>,
|
|
%Out: tensor<16x32xi32>) -> tensor<16x32xi32> {
|
|
%0 = linalg.generic
|
|
{indexing_maps = [#map, #map1, #map2],
|
|
iterator_types = ["parallel", "parallel", "reduction"]}
|
|
ins(%A, %B : tensor<16x8xi16>, tensor<8x32xi16>)
|
|
outs(%Out : tensor<16x32xi32>) {
|
|
^bb0(%in: i16, %in_0: i16, %out: i32):
|
|
%1 = arith.extsi %in : i16 to i32
|
|
%2 = arith.extsi %in_0 : i16 to i32
|
|
%3 = arith.muli %1, %2 : i32
|
|
%4 = arith.addi %out, %3 : i32
|
|
linalg.yield %4 : i32
|
|
} -> tensor<16x32xi32>
|
|
return %0 : tensor<16x32xi32>
|
|
}
|
|
|
|
// ALL-LABEL: op_matmul_signed_cast
|
|
|
|
// NAMED-NOT: linalg.generic
|
|
// NAMED-NOT: linalg.matmul {cast = #linalg.type_fn<cast_unsigned>}
|
|
// NAMED: linalg.matmul
|
|
|
|
// CATEGORY-NOT: linalg.generic
|
|
// CATEGORY: %[[RES:.+]] = linalg.contract
|
|
// CATEGORY-SAME: indexing_maps = {{\[}}#[[$MAP_A]], #[[$MAP_B]], #[[$MAP_C]]{{\]}}
|
|
// CATEGORY-NOT: {cast =
|
|
// CATEGORY-SAME: ins
|
|
// CATEGORY: return %[[RES]]
|
|
|
|
// Mixed signed/unsigned inputs cannot be encoded with a single cast attribute.
|
|
func.func @negative_op_matmul_mixed_cast(%A: tensor<16x8xi16>, %B: tensor<8x32xi16>,
|
|
%Out: tensor<16x32xi32>) -> tensor<16x32xi32> {
|
|
%0 = linalg.generic
|
|
{indexing_maps = [#map, #map1, #map2],
|
|
iterator_types = ["parallel", "parallel", "reduction"]}
|
|
ins(%A, %B : tensor<16x8xi16>, tensor<8x32xi16>)
|
|
outs(%Out : tensor<16x32xi32>) {
|
|
^bb0(%in: i16, %in_0: i16, %out: i32):
|
|
%1 = arith.extui %in : i16 to i32
|
|
%2 = arith.extsi %in_0 : i16 to i32
|
|
%3 = arith.muli %1, %2 : i32
|
|
%4 = arith.addi %out, %3 : i32
|
|
linalg.yield %4 : i32
|
|
} -> tensor<16x32xi32>
|
|
return %0 : tensor<16x32xi32>
|
|
}
|
|
|
|
// ALL-LABEL: negative_op_matmul_mixed_cast
|
|
|
|
// NAMED: linalg.generic
|
|
// NAMED-NOT: linalg.matmul
|
|
|
|
// CATEGORY: linalg.generic
|
|
// CATEGORY-NOT: linalg.contract
|
|
|
|
// Output-side casts are not representable by the named matmul ops.
|
|
func.func @negative_op_matmul_output_cast(%A: tensor<16x8xi32>, %B: tensor<8x32xi32>,
|
|
%Out: tensor<16x32xi64>) -> tensor<16x32xi64> {
|
|
%0 = linalg.generic
|
|
{indexing_maps = [#map, #map1, #map2],
|
|
iterator_types = ["parallel", "parallel", "reduction"]}
|
|
ins(%A, %B : tensor<16x8xi32>, tensor<8x32xi32>)
|
|
outs(%Out : tensor<16x32xi64>) {
|
|
^bb0(%in: i32, %in_0: i32, %out: i64):
|
|
%3 = arith.trunci %out : i64 to i32
|
|
%4 = arith.muli %in, %in_0 : i32
|
|
%5 = arith.addi %3, %4 : i32
|
|
%6 = arith.extsi %5 : i32 to i64
|
|
linalg.yield %6 : i64
|
|
} -> tensor<16x32xi64>
|
|
return %0 : tensor<16x32xi64>
|
|
}
|
|
|
|
// ALL-LABEL: negative_op_matmul_output_cast
|
|
|
|
// NAMED: linalg.generic
|
|
// NAMED-NOT: linalg.matmul
|
|
|
|
// CATEGORY: linalg.generic
|
|
// CATEGORY-NOT: linalg.contract
|
|
|
|
// Bitcasts are not modeled by the cast attribute, but should not block
|
|
// specialization.
|
|
// NOTE: Bitcasts are not preserved by the matmul named op during
|
|
// roundtrip, so this is potentially loosing information here.
|
|
// See #177593 for more details.
|
|
func.func @op_matmul_bitcast_int_to_float(%A: tensor<16x8xi32>,
|
|
%B: tensor<8x32xi32>,
|
|
%Out: tensor<16x32xf32>) -> tensor<16x32xf32> {
|
|
%0 = linalg.generic
|
|
{indexing_maps = [#map, #map1, #map2],
|
|
iterator_types = ["parallel", "parallel", "reduction"]}
|
|
ins(%A, %B : tensor<16x8xi32>, tensor<8x32xi32>)
|
|
outs(%Out : tensor<16x32xf32>) {
|
|
^bb0(%in: i32, %in_0: i32, %out: f32):
|
|
%1 = arith.bitcast %in : i32 to f32
|
|
%2 = arith.bitcast %in_0 : i32 to f32
|
|
%3 = arith.mulf %1, %2 : f32
|
|
%4 = arith.addf %out, %3 : f32
|
|
linalg.yield %4 : f32
|
|
} -> tensor<16x32xf32>
|
|
return %0 : tensor<16x32xf32>
|
|
}
|
|
|
|
// ALL-LABEL: op_matmul_bitcast_int_to_float
|
|
|
|
// NAMED-NOT: linalg.generic
|
|
// NAMED: linalg.matmul
|
|
|
|
// CATEGORY-NOT: linalg.generic
|
|
// CATEGORY: linalg.contract
|
|
|
|
// Signed float casts only use sitofp, which defaults to signed semantics.
|
|
func.func @op_matmul_signed_cast_float(%A: tensor<16x8xi16>, %B: tensor<8x32xi16>,
|
|
%Out: tensor<16x32xf32>) -> tensor<16x32xf32> {
|
|
%0 = linalg.generic
|
|
{indexing_maps = [#map, #map1, #map2],
|
|
iterator_types = ["parallel", "parallel", "reduction"]}
|
|
ins(%A, %B : tensor<16x8xi16>, tensor<8x32xi16>)
|
|
outs(%Out : tensor<16x32xf32>) {
|
|
^bb0(%in: i16, %in_0: i16, %out: f32):
|
|
%1 = arith.sitofp %in : i16 to f32
|
|
%2 = arith.sitofp %in_0 : i16 to f32
|
|
%3 = arith.mulf %1, %2 : f32
|
|
%4 = arith.addf %out, %3 : f32
|
|
linalg.yield %4 : f32
|
|
} -> tensor<16x32xf32>
|
|
return %0 : tensor<16x32xf32>
|
|
}
|
|
|
|
// ALL-LABEL: op_matmul_signed_cast_float
|
|
|
|
// NAMED-NOT: linalg.generic
|
|
// NAMED-NOT: linalg.matmul {cast = #linalg.type_fn<cast_unsigned>}
|
|
// NAMED: linalg.matmul
|
|
|
|
// CATEGORY-NOT: linalg.generic
|
|
// CATEGORY-NOT: linalg.contract{{.*}}{cast =
|
|
// CATEGORY: %[[RES:.+]] = linalg.contract
|
|
// CATEGORY: return %[[RES]]
|
|
|
|
// Unsigned float casts are expressed via uitofp and use the unsigned cast attr.
|
|
func.func @op_matmul_unsigned_cast_float(%A: tensor<16x8xi16>, %B: tensor<8x32xi16>,
|
|
%Out: tensor<16x32xf32>) -> tensor<16x32xf32> {
|
|
%0 = linalg.generic
|
|
{indexing_maps = [#map, #map1, #map2],
|
|
iterator_types = ["parallel", "parallel", "reduction"]}
|
|
ins(%A, %B : tensor<16x8xi16>, tensor<8x32xi16>)
|
|
outs(%Out : tensor<16x32xf32>) {
|
|
^bb0(%in: i16, %in_0: i16, %out: f32):
|
|
%1 = arith.uitofp %in : i16 to f32
|
|
%2 = arith.uitofp %in_0 : i16 to f32
|
|
%3 = arith.mulf %1, %2 : f32
|
|
%4 = arith.addf %out, %3 : f32
|
|
linalg.yield %4 : f32
|
|
} -> tensor<16x32xf32>
|
|
return %0 : tensor<16x32xf32>
|
|
}
|
|
|
|
// ALL-LABEL: op_matmul_unsigned_cast_float
|
|
|
|
// NAMED-NOT: linalg.generic
|
|
// NAMED: linalg.matmul {cast = #linalg.type_fn<cast_unsigned>}
|
|
|
|
// CATEGORY-NOT: linalg.generic
|
|
// CATEGORY: linalg.contract{{.*}}{cast = #linalg.type_fn<cast_unsigned>}
|
|
|
|
// -----
|
|
|
|
///----------------------------------------------------------------------------------------
|
|
/// Tests for linalg.batch_matmul
|
|
///----------------------------------------------------------------------------------------
|
|
|
|
#map = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
|
|
#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
|
|
#map2 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
|
|
func.func @op_batch_matmul(%A: tensor<2x16x8xf32>, %B: tensor<2x8x16xf32>,
|
|
%Out: tensor<2x16x16xf32>) -> tensor<2x16x16xf32> {
|
|
%0 = linalg.generic
|
|
{indexing_maps = [#map, #map1, #map2],
|
|
iterator_types = ["parallel", "parallel", "parallel", "reduction"]}
|
|
ins(%A, %B : tensor<2x16x8xf32>, tensor<2x8x16xf32>)
|
|
outs(%Out : tensor<2x16x16xf32>) {
|
|
^bb0(%in: f32, %in_0: f32, %out: f32):
|
|
%1 = arith.mulf %in, %in_0 : f32
|
|
%2 = arith.addf %out, %1 : f32
|
|
linalg.yield %2 : f32
|
|
} -> tensor<2x16x16xf32>
|
|
return %0 : tensor<2x16x16xf32>
|
|
}
|
|
|
|
// CATEGORY-DAG: #[[$MAP_A:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
|
|
// CATEGORY-DAG: #[[$MAP_B:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
|
|
// CATEGORY-DAG: #[[$MAP_C:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
|
|
|
|
// ALL-LABEL: op_batch_matmul
|
|
// ALL-SAME: %[[A:.+]]: tensor<2x16x8xf32>, %[[B:.+]]: tensor<2x8x16xf32>, %[[OUT:.+]]: tensor<2x16x16xf32>) -> tensor<2x16x16xf32>
|
|
|
|
// NAMED-NOT: linalg.generic
|
|
// NAMED: linalg.batch_matmul ins(%[[A]], %[[B]] : tensor<2x16x8xf32>, tensor<2x8x16xf32>) outs(%[[OUT]] : tensor<2x16x16xf32>) -> tensor<2x16x16xf32>
|
|
|
|
// CATEGORY-NOT: linalg.generic
|
|
// CATEGORY: linalg.contract indexing_maps = {{\[}}#[[$MAP_A]], #[[$MAP_B]], #[[$MAP_C]]{{\]}}
|
|
// CATEGORY-SAME: ins(%[[A]], %[[B]] : tensor<2x16x8xf32>, tensor<2x8x16xf32>)
|
|
// CATEGORY-SAME: outs(%[[OUT]] : tensor<2x16x16xf32>) -> tensor<2x16x16xf32>
|
|
|
|
// Ensure that the unsigned cast path for cast detection is exercised for
|
|
// batch_matmul as well.
|
|
func.func @op_batch_matmul_unsigned_cast(%A: tensor<2x16x8xi16>,
|
|
%B: tensor<2x8x16xi64>,
|
|
%Out: tensor<2x16x16xi32>) -> tensor<2x16x16xi32> {
|
|
%0 = linalg.generic
|
|
{indexing_maps = [#map, #map1, #map2],
|
|
iterator_types = ["parallel", "parallel", "parallel", "reduction"]}
|
|
ins(%A, %B : tensor<2x16x8xi16>, tensor<2x8x16xi64>)
|
|
outs(%Out : tensor<2x16x16xi32>) {
|
|
^bb0(%in: i16, %in_0: i64, %out: i32):
|
|
%1 = arith.extui %in : i16 to i32
|
|
%2 = arith.trunci %in_0 : i64 to i32
|
|
%3 = arith.muli %1, %2 : i32
|
|
%4 = arith.addi %out, %3 : i32
|
|
linalg.yield %4 : i32
|
|
} -> tensor<2x16x16xi32>
|
|
return %0 : tensor<2x16x16xi32>
|
|
}
|
|
|
|
// ALL-LABEL: op_batch_matmul_unsigned_cast
|
|
|
|
// NAMED-NOT: linalg.generic
|
|
// NAMED: linalg.batch_matmul {cast = #linalg.type_fn<cast_unsigned>}
|
|
|
|
// CATEGORY-NOT: linalg.generic
|
|
// CATEGORY: linalg.contract indexing_maps = {{\[}}#[[$MAP_A]], #[[$MAP_B]], #[[$MAP_C]]{{\]}}
|
|
// CATEGORY-SAME: {cast = #linalg.type_fn<cast_unsigned>}
|
|
|
|
// -----
|
|
|
|
// A multi-reduction contraction.
|
|
#mapA = affine_map<(m, n, k1, k2) -> (m, k1, k2)>
|
|
#mapB = affine_map<(m, n, k1, k2) -> (k2, k1, n)>
|
|
#mapC = affine_map<(m, n, k1, k2) -> (m, n)>
|
|
func.func @op_multi_reduction(%A: tensor<10x20x30xf32>,
|
|
%B: tensor<30x20x40xf32>,
|
|
%C: tensor<10x40xf32>) -> tensor<10x40xf32> {
|
|
%0 = linalg.generic
|
|
{indexing_maps = [#mapA, #mapB, #mapC],
|
|
iterator_types = ["parallel", "parallel", "reduction", "reduction"]}
|
|
ins(%A, %B : tensor<10x20x30xf32>, tensor<30x20x40xf32>)
|
|
outs(%C : tensor<10x40xf32>) {
|
|
^bb0(%a: f32, %b: f32, %c: f32):
|
|
%1 = arith.mulf %a, %b : f32
|
|
%2 = arith.addf %c, %1 : f32
|
|
linalg.yield %2 : f32
|
|
} -> tensor<10x40xf32>
|
|
return %0 : tensor<10x40xf32>
|
|
}
|
|
|
|
// ALL-LABEL: op_multi_reduction
|
|
|
|
// Cannot be lifted to named matrix multiply.
|
|
// NAMED: linalg.generic
|
|
|
|
// CATEGORY-NOT: linalg.generic
|
|
// CATEGORY: linalg.contract
|
|
|
|
// -----
|
|
|
|
// Batch dim not in identity position: batch dim d0 appears at result
|
|
// position 1 in A's map instead of position 0.
|
|
#mapBni0 = affine_map<(d0, d1, d2, d3) -> (d1, d0, d3)>
|
|
#mapBni1 = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
|
|
#mapBni2 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
|
|
func.func @batch_matmul_non_identity_batch(%A: tensor<4x2x8xf32>, %B: tensor<2x8x16xf32>,
|
|
%Out: tensor<2x4x16xf32>) -> tensor<2x4x16xf32> {
|
|
%0 = linalg.generic
|
|
{indexing_maps = [#mapBni0, #mapBni1, #mapBni2],
|
|
iterator_types = ["parallel", "parallel", "parallel", "reduction"]}
|
|
ins(%A, %B : tensor<4x2x8xf32>, tensor<2x8x16xf32>)
|
|
outs(%Out : tensor<2x4x16xf32>) {
|
|
^bb0(%in: f32, %in_0: f32, %out: f32):
|
|
%1 = arith.mulf %in, %in_0 : f32
|
|
%2 = arith.addf %out, %1 : f32
|
|
linalg.yield %2 : f32
|
|
} -> tensor<2x4x16xf32>
|
|
return %0 : tensor<2x4x16xf32>
|
|
}
|
|
|
|
// ALL-LABEL: batch_matmul_non_identity_batch
|
|
|
|
// Cannot be lifted to named matrix multiply.
|
|
// NAMED: linalg.generic
|
|
|
|
// CATEGORY-NOT: linalg.generic
|
|
// CATEGORY: linalg.contract
|
|
|
|
// -----
|
|
|
|
// TODO: named matvec
|
|
#map = affine_map<(d0, d1) -> (d0, d1)>
|
|
#map1 = affine_map<(d0, d1) -> (d1)>
|
|
#map2 = affine_map<(d0, d1) -> (d0)>
|
|
func.func @op_matvec(%A: tensor<?x?xf32>, %B: tensor<?xf32>, %Out: tensor<?xf32>) -> tensor<?xf32> {
|
|
%0 = linalg.generic
|
|
{indexing_maps = [#map, #map1, #map2],
|
|
iterator_types = ["parallel", "reduction"]}
|
|
ins(%A, %B : tensor<?x?xf32>, tensor<?xf32>)
|
|
outs(%Out : tensor<?xf32>) {
|
|
^bb0(%in: f32, %in_0: f32, %out: f32):
|
|
%1 = arith.mulf %in, %in_0 : f32
|
|
%2 = arith.addf %out, %1 : f32
|
|
linalg.yield %2 : f32
|
|
} -> tensor<?xf32>
|
|
return %0 : tensor<?xf32>
|
|
}
|
|
|
|
// ALL-LABEL: op_matvec
|
|
|
|
// NAMED: linalg.generic
|
|
|
|
// CATEGORY-NOT: linalg.generic
|
|
// CATEGORY: linalg.contract
|
|
|
|
// -----
|
|
|
|
// Matmul transpose A: A is accessed as (k, m) instead of (m, k)
|
|
#map_ta = affine_map<(d0, d1, d2) -> (d2, d0)>
|
|
#map_b = affine_map<(d0, d1, d2) -> (d2, d1)>
|
|
#map_c = affine_map<(d0, d1, d2) -> (d0, d1)>
|
|
func.func @op_matmul_transpose_a(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>,
|
|
%Out: tensor<?x?xf32>) -> tensor<?x?xf32> {
|
|
%0 = linalg.generic
|
|
{indexing_maps = [#map_ta, #map_b, #map_c],
|
|
iterator_types = ["parallel", "parallel", "reduction"]}
|
|
ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>) outs(%Out : tensor<?x?xf32>) {
|
|
^bb0(%in: f32, %in_0: f32, %out: f32):
|
|
%1 = arith.mulf %in, %in_0 : f32
|
|
%2 = arith.addf %out, %1 : f32
|
|
linalg.yield %2 : f32
|
|
} -> tensor<?x?xf32>
|
|
return %0 : tensor<?x?xf32>
|
|
}
|
|
|
|
// ALL-DAG: #[[$MAP_TA:.+]] = affine_map<(d0, d1, d2) -> (d2, d0)>
|
|
// ALL-DAG: #[[$MAP_B:.+]] = affine_map<(d0, d1, d2) -> (d2, d1)>
|
|
// ALL-DAG: #[[$MAP_C:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
|
|
|
|
// ALL-LABEL: op_matmul_transpose_a
|
|
// ALL-SAME: %[[A:.+]]: tensor<?x?xf32>, %[[B:.+]]: tensor<?x?xf32>, %[[OUT:.+]]: tensor<?x?xf32>
|
|
|
|
// NAMED-NOT: linalg.generic
|
|
// NAMED: linalg.matmul
|
|
// NAMED-SAME: indexing_maps = [#[[$MAP_TA]], #[[$MAP_B]], #[[$MAP_C]]]
|
|
// NAMED-SAME: ins(%[[A]], %[[B]] : tensor<?x?xf32>, tensor<?x?xf32>)
|
|
// NAMED-SAME: outs(%[[OUT]] : tensor<?x?xf32>) -> tensor<?x?xf32>
|
|
|
|
// CATEGORY-NOT: linalg.generic
|
|
// CATEGORY: linalg.contract
|
|
// CATEGORY-SAME: indexing_maps = [#[[$MAP_TA]], #[[$MAP_B]], #[[$MAP_C]]]
|
|
|
|
// -----
|
|
|
|
// Matmul transpose B: B is accessed as (n, k) instead of (k, n)
|
|
#map_a = affine_map<(d0, d1, d2) -> (d0, d2)>
|
|
#map_tb = affine_map<(d0, d1, d2) -> (d1, d2)>
|
|
#map_c = affine_map<(d0, d1, d2) -> (d0, d1)>
|
|
func.func @op_matmul_transpose_b(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>,
|
|
%Out: tensor<?x?xf32>) -> tensor<?x?xf32> {
|
|
%0 = linalg.generic
|
|
{indexing_maps = [#map_a, #map_tb, #map_c],
|
|
iterator_types = ["parallel", "parallel", "reduction"]}
|
|
ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>)
|
|
outs(%Out : tensor<?x?xf32>) {
|
|
^bb0(%in: f32, %in_0: f32, %out: f32):
|
|
%1 = arith.mulf %in, %in_0 : f32
|
|
%2 = arith.addf %out, %1 : f32
|
|
linalg.yield %2 : f32
|
|
} -> tensor<?x?xf32>
|
|
return %0 : tensor<?x?xf32>
|
|
}
|
|
|
|
// ALL-DAG: #[[$MAP_A:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>
|
|
// ALL-DAG: #[[$MAP_TB:.+]] = affine_map<(d0, d1, d2) -> (d1, d2)>
|
|
// ALL-DAG: #[[$MAP_C:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
|
|
|
|
// ALL-LABEL: op_matmul_transpose_b
|
|
// ALL-SAME: %[[A:.+]]: tensor<?x?xf32>, %[[B:.+]]: tensor<?x?xf32>, %[[OUT:.+]]: tensor<?x?xf32>
|
|
|
|
// NAMED-NOT: linalg.generic
|
|
// NAMED: linalg.matmul
|
|
// NAMED-SAME: indexing_maps = [#[[$MAP_A]], #[[$MAP_TB]], #[[$MAP_C]]]
|
|
// NAMED-SAME: ins(%[[A]], %[[B]] : tensor<?x?xf32>, tensor<?x?xf32>)
|
|
// NAMED-SAME: outs(%[[OUT]] : tensor<?x?xf32>) -> tensor<?x?xf32>
|
|
|
|
// CATEGORY-NOT: linalg.generic
|
|
// CATEGORY: linalg.contract
|
|
// CATEGORY-SAME: indexing_maps = [#[[$MAP_A]], #[[$MAP_TB]], #[[$MAP_C]]]
|
|
// CATEGORY-SAME: ins(%[[A]], %[[B]] : tensor<?x?xf32>, tensor<?x?xf32>)
|
|
// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?xf32>) -> tensor<?x?xf32>
|
|
|
|
// -----
|
|
|
|
// Batch matmul transpose A: A is accessed as (b, k, m) instead of (b, m, k)
|
|
#map_ta = affine_map<(d0, d1, d2, d3) -> (d0, d3, d1)>
|
|
#map_b = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
|
|
#map_c = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
|
|
func.func @op_batch_matmul_transpose_a(%A: tensor<2x8x4xf32>, %B: tensor<2x8x16xf32>,
|
|
%Out: tensor<2x4x16xf32>) -> tensor<2x4x16xf32> {
|
|
%0 = linalg.generic
|
|
{indexing_maps = [#map_ta, #map_b, #map_c],
|
|
iterator_types = ["parallel", "parallel", "parallel", "reduction"]}
|
|
ins(%A, %B : tensor<2x8x4xf32>, tensor<2x8x16xf32>)
|
|
outs(%Out : tensor<2x4x16xf32>) {
|
|
^bb0(%in: f32, %in_0: f32, %out: f32):
|
|
%1 = arith.mulf %in, %in_0 : f32
|
|
%2 = arith.addf %out, %1 : f32
|
|
linalg.yield %2 : f32
|
|
} -> tensor<2x4x16xf32>
|
|
return %0 : tensor<2x4x16xf32>
|
|
}
|
|
|
|
// ALL-DAG: #[[$MAP_TA:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d1)>
|
|
// ALL-DAG: #[[$MAP_B:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
|
|
// ALL-DAG: #[[$MAP_C:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
|
|
|
|
// ALL-LABEL: op_batch_matmul_transpose_a
|
|
// ALL-SAME: %[[A:.+]]: tensor<2x8x4xf32>, %[[B:.+]]: tensor<2x8x16xf32>, %[[OUT:.+]]: tensor<2x4x16xf32>
|
|
|
|
// NAMED-NOT: linalg.generic
|
|
// NAMED: linalg.batch_matmul
|
|
// NAMED-SAME: indexing_maps = [#[[$MAP_TA]], #[[$MAP_B]], #[[$MAP_C]]]
|
|
// NAMED-SAME: ins(%[[A]], %[[B]] : tensor<2x8x4xf32>, tensor<2x8x16xf32>)
|
|
// NAMED-SAME: outs(%[[OUT]] : tensor<2x4x16xf32>) -> tensor<2x4x16xf32>
|
|
|
|
// CATEGORY-NOT: linalg.generic
|
|
// CATEGORY: linalg.contract
|
|
// CATEGORY-SAME: indexing_maps = [#[[$MAP_TA]], #[[$MAP_B]], #[[$MAP_C]]]
|
|
// CATEGORY-SAME: ins(%[[A]], %[[B]] : tensor<2x8x4xf32>, tensor<2x8x16xf32>)
|
|
// CATEGORY-SAME: outs(%[[OUT]] : tensor<2x4x16xf32>) -> tensor<2x4x16xf32>
|
|
|
|
// -----
|
|
|
|
// Batch matmul transpose B: B is accessed as (b, n, k) instead of (b, k, n)
|
|
#map_a = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
|
|
#map_tb = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>
|
|
#map_c = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
|
|
func.func @op_batch_matmul_transpose_b(%A: tensor<2x4x8xf32>, %B: tensor<2x16x8xf32>,
|
|
%Out: tensor<2x4x16xf32>) -> tensor<2x4x16xf32> {
|
|
%0 = linalg.generic
|
|
{indexing_maps = [#map_a, #map_tb, #map_c],
|
|
iterator_types = ["parallel", "parallel", "parallel", "reduction"]}
|
|
ins(%A, %B : tensor<2x4x8xf32>, tensor<2x16x8xf32>)
|
|
outs(%Out : tensor<2x4x16xf32>) {
|
|
^bb0(%in: f32, %in_0: f32, %out: f32):
|
|
%1 = arith.mulf %in, %in_0 : f32
|
|
%2 = arith.addf %out, %1 : f32
|
|
linalg.yield %2 : f32
|
|
} -> tensor<2x4x16xf32>
|
|
return %0 : tensor<2x4x16xf32>
|
|
}
|
|
|
|
// ALL-DAG: #[[$MAP_A:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
|
|
// ALL-DAG: #[[$MAP_TB:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>
|
|
// ALL-DAG: #[[$MAP_C:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
|
|
|
|
// ALL-LABEL: op_batch_matmul_transpose_b
|
|
// ALL-SAME: %[[A:.+]]: tensor<2x4x8xf32>, %[[B:.+]]: tensor<2x16x8xf32>, %[[OUT:.+]]: tensor<2x4x16xf32>
|
|
|
|
// NAMED-NOT: linalg.generic
|
|
// NAMED: linalg.batch_matmul
|
|
// NAMED-SAME: indexing_maps = [#[[$MAP_A]], #[[$MAP_TB]], #[[$MAP_C]]]
|
|
// NAMED-SAME: ins(%[[A]], %[[B]] : tensor<2x4x8xf32>, tensor<2x16x8xf32>)
|
|
// NAMED-SAME: outs(%[[OUT]] : tensor<2x4x16xf32>) -> tensor<2x4x16xf32>
|
|
|
|
// CATEGORY-NOT: linalg.generic
|
|
// CATEGORY: linalg.contract
|
|
// CATEGORY-SAME: indexing_maps = [#[[$MAP_A]], #[[$MAP_TB]], #[[$MAP_C]]]
|
|
// CATEGORY-SAME: ins(%[[A]], %[[B]] : tensor<2x4x8xf32>, tensor<2x16x8xf32>)
|
|
// CATEGORY-SAME: outs(%[[OUT]] : tensor<2x4x16xf32>) -> tensor<2x4x16xf32>
|
|
|
|
// -----
|
|
|
|
// Both A and B transposed.
|
|
#map_ta = affine_map<(d0, d1, d2) -> (d2, d0)>
|
|
#map_tb = affine_map<(d0, d1, d2) -> (d1, d2)>
|
|
#map_c = affine_map<(d0, d1, d2) -> (d0, d1)>
|
|
func.func @op_matmul_transpose_a_and_b(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>,
|
|
%Out: tensor<?x?xf32>) -> tensor<?x?xf32> {
|
|
%0 = linalg.generic
|
|
{indexing_maps = [#map_ta, #map_tb, #map_c],
|
|
iterator_types = ["parallel", "parallel", "reduction"]}
|
|
ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>)
|
|
outs(%Out : tensor<?x?xf32>) {
|
|
^bb0(%in: f32, %in_0: f32, %out: f32):
|
|
%1 = arith.mulf %in, %in_0 : f32
|
|
%2 = arith.addf %out, %1 : f32
|
|
linalg.yield %2 : f32
|
|
} -> tensor<?x?xf32>
|
|
return %0 : tensor<?x?xf32>
|
|
}
|
|
|
|
// ALL-DAG: #[[$MAP_TA:.+]] = affine_map<(d0, d1, d2) -> (d2, d0)>
|
|
// ALL-DAG: #[[$MAP_TB:.+]] = affine_map<(d0, d1, d2) -> (d1, d2)>
|
|
// ALL-DAG: #[[$MAP_C:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
|
|
|
|
// ALL-LABEL: op_matmul_transpose_a_and_b
|
|
// ALL-SAME: %[[A:.+]]: tensor<?x?xf32>, %[[B:.+]]: tensor<?x?xf32>, %[[OUT:.+]]: tensor<?x?xf32>
|
|
|
|
// NAMED-NOT: linalg.generic
|
|
// NAMED: linalg.matmul
|
|
// NAMED-SAME: indexing_maps = [#[[$MAP_TA]], #[[$MAP_TB]], #[[$MAP_C]]]
|
|
// NAMED-SAME: ins(%[[A]], %[[B]] : tensor<?x?xf32>, tensor<?x?xf32>)
|
|
// NAMED-SAME: outs(%[[OUT]] : tensor<?x?xf32>) -> tensor<?x?xf32>
|
|
|
|
// CATEGORY-NOT: linalg.generic
|
|
// CATEGORY: linalg.contract
|
|
// CATEGORY-SAME: indexing_maps = [#[[$MAP_TA]], #[[$MAP_TB]], #[[$MAP_C]]]
|
|
// CATEGORY-SAME: ins(%[[A]], %[[B]] : tensor<?x?xf32>, tensor<?x?xf32>)
|
|
// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?xf32>) -> tensor<?x?xf32>
|
|
|
|
// -----
|
|
|
|
// Output transposed: C is accessed as (n, m) instead of (m, n).
|
|
#map_a = affine_map<(d0, d1, d2) -> (d0, d2)>
|
|
#map_b = affine_map<(d0, d1, d2) -> (d2, d1)>
|
|
#map_tc = affine_map<(d0, d1, d2) -> (d1, d0)>
|
|
func.func @op_matmul_transposed_output(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>,
|
|
%Out: tensor<?x?xf32>) -> tensor<?x?xf32> {
|
|
%0 = linalg.generic
|
|
{indexing_maps = [#map_a, #map_b, #map_tc],
|
|
iterator_types = ["parallel", "parallel", "reduction"]}
|
|
ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>)
|
|
outs(%Out : tensor<?x?xf32>) {
|
|
^bb0(%in: f32, %in_0: f32, %out: f32):
|
|
%1 = arith.mulf %in, %in_0 : f32
|
|
%2 = arith.addf %out, %1 : f32
|
|
linalg.yield %2 : f32
|
|
} -> tensor<?x?xf32>
|
|
return %0 : tensor<?x?xf32>
|
|
}
|
|
|
|
// ALL-DAG: #[[$MAP_A:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>
|
|
// ALL-DAG: #[[$MAP_B:.+]] = affine_map<(d0, d1, d2) -> (d2, d1)>
|
|
// ALL-DAG: #[[$MAP_TC:.+]] = affine_map<(d0, d1, d2) -> (d1, d0)>
|
|
|
|
// ALL-LABEL: op_matmul_transposed_output
|
|
// ALL-SAME: %[[A:.+]]: tensor<?x?xf32>, %[[B:.+]]: tensor<?x?xf32>, %[[OUT:.+]]: tensor<?x?xf32>
|
|
|
|
// NAMED-NOT: linalg.generic
|
|
// NAMED: linalg.matmul
|
|
// NAMED-SAME: indexing_maps = [#[[$MAP_A]], #[[$MAP_B]], #[[$MAP_TC]]]
|
|
// NAMED-SAME: ins(%[[A]], %[[B]] : tensor<?x?xf32>, tensor<?x?xf32>)
|
|
// NAMED-SAME: outs(%[[OUT]] : tensor<?x?xf32>) -> tensor<?x?xf32>
|
|
|
|
// CATEGORY-NOT: linalg.generic
|
|
// CATEGORY: linalg.contract
|
|
// CATEGORY-SAME: indexing_maps = [#[[$MAP_A]], #[[$MAP_B]], #[[$MAP_TC]]]
|
|
// CATEGORY-SAME: ins(%[[A]], %[[B]] : tensor<?x?xf32>, tensor<?x?xf32>)
|
|
// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?xf32>) -> tensor<?x?xf32>
|
|
|
|
// -----
|
|
|
|
// Matmul with non-canonical loop ordering.
|
|
#map_nc_a = affine_map<(m, k, n) -> (m, k)>
|
|
#map_nc_b = affine_map<(m, k, n) -> (k, n)>
|
|
#map_nc_c = affine_map<(m, k, n) -> (m, n)>
|
|
func.func @op_matmul_non_canonical_loops(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>,
|
|
%Out: tensor<?x?xf32>) -> tensor<?x?xf32> {
|
|
%0 = linalg.generic
|
|
{indexing_maps = [#map_nc_a, #map_nc_b, #map_nc_c],
|
|
iterator_types = ["parallel", "reduction", "parallel"]}
|
|
ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>)
|
|
outs(%Out : tensor<?x?xf32>) {
|
|
^bb0(%in: f32, %in_0: f32, %out: f32):
|
|
%1 = arith.mulf %in, %in_0 : f32
|
|
%2 = arith.addf %out, %1 : f32
|
|
linalg.yield %2 : f32
|
|
} -> tensor<?x?xf32>
|
|
return %0 : tensor<?x?xf32>
|
|
}
|
|
|
|
// CATEGORY-DAG: #[[$MAP_A:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
|
|
// CATEGORY-DAG: #[[$MAP_B:.+]] = affine_map<(d0, d1, d2) -> (d1, d2)>
|
|
// CATEGORY-DAG: #[[$MAP_C:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>
|
|
|
|
// ALL-LABEL: op_matmul_non_canonical_loops
|
|
// ALL-SAME: %[[A:.+]]: tensor<?x?xf32>, %[[B:.+]]: tensor<?x?xf32>, %[[OUT:.+]]: tensor<?x?xf32>
|
|
|
|
// NAMED-NOT: linalg.generic
|
|
// NAMED: linalg.matmul
|
|
// NAMED-SAME: ins(%[[A]], %[[B]] : tensor<?x?xf32>, tensor<?x?xf32>)
|
|
// NAMED-SAME: outs(%[[OUT]] : tensor<?x?xf32>) -> tensor<?x?xf32>
|
|
|
|
// CATEGORY-NOT: linalg.generic
|
|
// CATEGORY: linalg.contract
|
|
// CATEGORY-SAME: indexing_maps = [#[[$MAP_A]], #[[$MAP_B]], #[[$MAP_C]]]
|
|
// CATEGORY-SAME: ins(%[[A]], %[[B]] : tensor<?x?xf32>, tensor<?x?xf32>)
|
|
// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?xf32>) -> tensor<?x?xf32>
|
|
// -----
|
|
|
|
// Batch matmul with non-canonical loop ordering.
|
|
#map_bnc_a = affine_map<(batch, m, k, n) -> (batch, m, k)>
|
|
#map_bnc_b = affine_map<(batch, m, k, n) -> (batch, k, n)>
|
|
#map_bnc_c = affine_map<(batch, m, k, n) -> (batch, m, n)>
|
|
func.func @op_batch_matmul_non_canonical_loops(%A: tensor<2x16x8xf32>, %B: tensor<2x8x16xf32>,
|
|
%Out: tensor<2x16x16xf32>) -> tensor<2x16x16xf32> {
|
|
%0 = linalg.generic
|
|
{indexing_maps = [#map_bnc_a, #map_bnc_b, #map_bnc_c],
|
|
iterator_types = ["parallel", "parallel", "reduction", "parallel"]}
|
|
ins(%A, %B : tensor<2x16x8xf32>, tensor<2x8x16xf32>)
|
|
outs(%Out : tensor<2x16x16xf32>) {
|
|
^bb0(%in: f32, %in_0: f32, %out: f32):
|
|
%1 = arith.mulf %in, %in_0 : f32
|
|
%2 = arith.addf %out, %1 : f32
|
|
linalg.yield %2 : f32
|
|
} -> tensor<2x16x16xf32>
|
|
return %0 : tensor<2x16x16xf32>
|
|
}
|
|
|
|
// CATEGORY-DAG: #[[$MAP_A:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
|
|
// CATEGORY-DAG: #[[$MAP_B:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>
|
|
// CATEGORY-DAG: #[[$MAP_C:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
|
|
|
|
// ALL-LABEL: op_batch_matmul_non_canonical_loops
|
|
// ALL-SAME: %[[A:.+]]: tensor<2x16x8xf32>, %[[B:.+]]: tensor<2x8x16xf32>, %[[OUT:.+]]: tensor<2x16x16xf32>
|
|
|
|
// NAMED-NOT: linalg.generic
|
|
// NAMED: linalg.batch_matmul
|
|
// NAMED-SAME: ins(%[[A]], %[[B]] : tensor<2x16x8xf32>, tensor<2x8x16xf32>)
|
|
// NAMED-SAME: outs(%[[OUT]] : tensor<2x16x16xf32>) -> tensor<2x16x16xf32>
|
|
|
|
// CATEGORY-NOT: linalg.generic
|
|
// CATEGORY: linalg.contract
|
|
// CATEGORY-SAME: indexing_maps = [#[[$MAP_A]], #[[$MAP_B]], #[[$MAP_C]]]
|
|
// CATEGORY-SAME: ins(%[[A]], %[[B]] : tensor<2x16x8xf32>, tensor<2x8x16xf32>)
|
|
// CATEGORY-SAME: outs(%[[OUT]] : tensor<2x16x16xf32>) -> tensor<2x16x16xf32>
|
|
|
|
// -----
|
|
|
|
// Matmul with non-canonical loop ordering (d0=m, d1=k, d2=n) and B transposed.
|
|
#map_nc_tb_a = affine_map<(d0, d1, d2) -> (d0, d1)>
|
|
#map_nc_tb_b = affine_map<(d0, d1, d2) -> (d2, d1)>
|
|
#map_nc_tb_c = affine_map<(d0, d1, d2) -> (d0, d2)>
|
|
func.func @op_matmul_non_canonical_transpose_b(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>,
|
|
%Out: tensor<?x?xf32>) -> tensor<?x?xf32> {
|
|
%0 = linalg.generic
|
|
{indexing_maps = [#map_nc_tb_a, #map_nc_tb_b, #map_nc_tb_c],
|
|
iterator_types = ["parallel", "reduction", "parallel"]}
|
|
ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>) outs(%Out : tensor<?x?xf32>) {
|
|
^bb0(%in: f32, %in_0: f32, %out: f32):
|
|
%1 = arith.mulf %in, %in_0 : f32
|
|
%2 = arith.addf %out, %1 : f32
|
|
linalg.yield %2 : f32
|
|
} -> tensor<?x?xf32>
|
|
return %0 : tensor<?x?xf32>
|
|
}
|
|
|
|
// NAMED-DAG: #[[$MAP_A:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>
|
|
// NAMED-DAG: #[[$MAP_TB:.+]] = affine_map<(d0, d1, d2) -> (d1, d2)>
|
|
// NAMED-DAG: #[[$MAP_C:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
|
|
|
|
// CATEGORY-DAG: #[[$MAP_A:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
|
|
// CATEGORY-DAG: #[[$MAP_TB:.+]] = affine_map<(d0, d1, d2) -> (d2, d1)>
|
|
// CATEGORY-DAG: #[[$MAP_C:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>
|
|
|
|
// ALL-LABEL: op_matmul_non_canonical_transpose_b
|
|
// ALL-SAME: %[[A:.+]]: tensor<?x?xf32>, %[[B:.+]]: tensor<?x?xf32>, %[[OUT:.+]]: tensor<?x?xf32>
|
|
|
|
// NAMED-NOT: linalg.generic
|
|
// NAMED: linalg.matmul
|
|
// NAMED-SAME: indexing_maps = [#[[$MAP_A]], #[[$MAP_TB]], #[[$MAP_C]]]
|
|
// NAMED-SAME: ins(%[[A]], %[[B]] : tensor<?x?xf32>, tensor<?x?xf32>)
|
|
// NAMED-SAME: outs(%[[OUT]] : tensor<?x?xf32>) -> tensor<?x?xf32>
|
|
|
|
// CATEGORY-NOT: linalg.generic
|
|
// CATEGORY: linalg.contract
|
|
// CATEGORY-SAME: indexing_maps = [#[[$MAP_A]], #[[$MAP_TB]], #[[$MAP_C]]]
|
|
// CATEGORY-SAME: ins(%[[A]], %[[B]] : tensor<?x?xf32>, tensor<?x?xf32>)
|
|
// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?xf32>) -> tensor<?x?xf32>
|
|
|
|
// -----
|
|
|
|
// Batch matmul with non-canonical loop ordering (d0=batch, d1=m, d2=k, d3=n)
|
|
// and B Transposed.
|
|
#map_bnc_tb_a = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
|
|
#map_bnc_tb_b = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
|
|
#map_bnc_tb_c = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
|
|
func.func @op_batch_matmul_non_canonical_transpose_b(%A: tensor<2x16x8xf32>, %B: tensor<2x16x8xf32>,
|
|
%Out: tensor<2x16x16xf32>) -> tensor<2x16x16xf32> {
|
|
%0 = linalg.generic
|
|
{indexing_maps = [#map_bnc_tb_a, #map_bnc_tb_b, #map_bnc_tb_c],
|
|
iterator_types = ["parallel", "parallel", "reduction", "parallel"]}
|
|
ins(%A, %B : tensor<2x16x8xf32>, tensor<2x16x8xf32>)
|
|
outs(%Out : tensor<2x16x16xf32>) {
|
|
^bb0(%in: f32, %in_0: f32, %out: f32):
|
|
%1 = arith.mulf %in, %in_0 : f32
|
|
%2 = arith.addf %out, %1 : f32
|
|
linalg.yield %2 : f32
|
|
} -> tensor<2x16x16xf32>
|
|
return %0 : tensor<2x16x16xf32>
|
|
}
|
|
|
|
// NAMED-DAG: #[[$MAP_A:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
|
|
// NAMED-DAG: #[[$MAP_TB:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>
|
|
// NAMED-DAG: #[[$MAP_C:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
|
|
|
|
// CATEGORY-DAG: #[[$MAP_A:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
|
|
// CATEGORY-DAG: #[[$MAP_TB:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
|
|
// CATEGORY-DAG: #[[$MAP_C:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
|
|
|
|
// ALL-LABEL: op_batch_matmul_non_canonical_transpose_b
|
|
// ALL-SAME: %[[A:.+]]: tensor<2x16x8xf32>, %[[B:.+]]: tensor<2x16x8xf32>, %[[OUT:.+]]: tensor<2x16x16xf32>
|
|
|
|
// NAMED-NOT: linalg.generic
|
|
// NAMED: linalg.batch_matmul
|
|
// NAMED-SAME: indexing_maps = [#[[$MAP_A]], #[[$MAP_TB]], #[[$MAP_C]]]
|
|
// NAMED-SAME: ins(%[[A]], %[[B]] : tensor<2x16x8xf32>, tensor<2x16x8xf32>)
|
|
// NAMED-SAME: outs(%[[OUT]] : tensor<2x16x16xf32>) -> tensor<2x16x16xf32>
|
|
|
|
// CATEGORY-NOT: linalg.generic
|
|
// CATEGORY: linalg.contract
|
|
// CATEGORY-SAME: indexing_maps = [#[[$MAP_A]], #[[$MAP_TB]], #[[$MAP_C]]]
|
|
// CATEGORY-SAME: ins(%[[A]], %[[B]] : tensor<2x16x8xf32>, tensor<2x16x8xf32>)
|
|
// CATEGORY-SAME: outs(%[[OUT]] : tensor<2x16x16xf32>) -> tensor<2x16x16xf32>
|
|
|
|
// -----
|
|
|
|
// Matmul with fully permuted loop ordering.
|
|
#map_fs_a = affine_map<(d0, d1, d2) -> (d1, d0)>
|
|
#map_fs_b = affine_map<(d0, d1, d2) -> (d0, d2)>
|
|
#map_fs_c = affine_map<(d0, d1, d2) -> (d1, d2)>
|
|
func.func @op_matmul_fully_shuffled_loops(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>,
|
|
%Out: tensor<?x?xf32>) -> tensor<?x?xf32> {
|
|
%0 = linalg.generic
|
|
{indexing_maps = [#map_fs_a, #map_fs_b, #map_fs_c],
|
|
iterator_types = ["reduction", "parallel", "parallel"]}
|
|
ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>)
|
|
outs(%Out : tensor<?x?xf32>) {
|
|
^bb0(%in: f32, %in_0: f32, %out: f32):
|
|
%1 = arith.mulf %in, %in_0 : f32
|
|
%2 = arith.addf %out, %1 : f32
|
|
linalg.yield %2 : f32
|
|
} -> tensor<?x?xf32>
|
|
return %0 : tensor<?x?xf32>
|
|
}
|
|
|
|
// CATEGORY-DAG: #[[$MAP_A:.+]] = affine_map<(d0, d1, d2) -> (d1, d0)>
|
|
// CATEGORY-DAG: #[[$MAP_B:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>
|
|
// CATEGORY-DAG: #[[$MAP_C:.+]] = affine_map<(d0, d1, d2) -> (d1, d2)>
|
|
|
|
// ALL-LABEL: op_matmul_fully_shuffled_loops
|
|
// ALL-SAME: %[[A:.+]]: tensor<?x?xf32>, %[[B:.+]]: tensor<?x?xf32>, %[[OUT:.+]]: tensor<?x?xf32>
|
|
|
|
// NAMED-NOT: linalg.generic
|
|
// NAMED: linalg.matmul
|
|
// NAMED-SAME: ins(%[[A]], %[[B]] : tensor<?x?xf32>, tensor<?x?xf32>)
|
|
// NAMED-SAME: outs(%[[OUT]] : tensor<?x?xf32>) -> tensor<?x?xf32>
|
|
|
|
// CATEGORY-NOT: linalg.generic
|
|
// CATEGORY: linalg.contract
|
|
// CATEGORY-SAME: indexing_maps = [#[[$MAP_A]], #[[$MAP_TB]], #[[$MAP_C]]]
|
|
// CATEGORY-SAME: ins(%[[A]], %[[B]] : tensor<?x?xf32>, tensor<?x?xf32>)
|
|
// CATEGORY-SAME: outs(%[[OUT]] : tensor<?x?xf32>) -> tensor<?x?xf32>
|
|
|
|
// -----
|
|
|
|
// TODO: this could also be specialized to a named matmul.
|
|
#map_bcast_a = affine_map<(d0, d1, d2) -> (d2)>
|
|
#map_bcast_b = affine_map<(d0, d1, d2) -> (d2, d1)>
|
|
#map_bcast_c = affine_map<(d0, d1, d2) -> (d0, d1)>
|
|
func.func @op_matmul_broadcast_a(%A: tensor<?xf32>, %B: tensor<?x?xf32>,
|
|
%Out: tensor<?x?xf32>) -> tensor<?x?xf32> {
|
|
%0 = linalg.generic
|
|
{indexing_maps = [#map_bcast_a, #map_bcast_b, #map_bcast_c],
|
|
iterator_types = ["parallel", "parallel", "reduction"]}
|
|
ins(%A, %B : tensor<?xf32>, tensor<?x?xf32>)
|
|
outs(%Out : tensor<?x?xf32>) {
|
|
^bb0(%in: f32, %in_0: f32, %out: f32):
|
|
%1 = arith.mulf %in, %in_0 : f32
|
|
%2 = arith.addf %out, %1 : f32
|
|
linalg.yield %2 : f32
|
|
} -> tensor<?x?xf32>
|
|
return %0 : tensor<?x?xf32>
|
|
}
|
|
|
|
// ALL-LABEL: op_matmul_broadcast_a
|
|
|
|
// NAMED: linalg.generic
|
|
// NAMED-NOT: linalg.matmul
|
|
|
|
// CATEGORY-NOT: linalg.generic
|
|
// CATEGORY: linalg.contract
|
|
|
|
// -----
|
|
|
|
// TODO: this could also be specialized to a named batch_matmul.
|
|
#map_bbcast_a = affine_map<(d0, d1, d2, d3) -> (d1, d3)>
|
|
#map_bbcast_b = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
|
|
#map_bbcast_c = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
|
|
func.func @op_batch_matmul_broadcast_a(%A: tensor<16x8xf32>, %B: tensor<2x8x16xf32>,
|
|
%Out: tensor<2x16x16xf32>) -> tensor<2x16x16xf32> {
|
|
%0 = linalg.generic
|
|
{indexing_maps = [#map_bbcast_a, #map_bbcast_b, #map_bbcast_c],
|
|
iterator_types = ["parallel", "parallel", "parallel", "reduction"]}
|
|
ins(%A, %B : tensor<16x8xf32>, tensor<2x8x16xf32>)
|
|
outs(%Out : tensor<2x16x16xf32>) {
|
|
^bb0(%in: f32, %in_0: f32, %out: f32):
|
|
%1 = arith.mulf %in, %in_0 : f32
|
|
%2 = arith.addf %out, %1 : f32
|
|
linalg.yield %2 : f32
|
|
} -> tensor<2x16x16xf32>
|
|
return %0 : tensor<2x16x16xf32>
|
|
}
|
|
|
|
// ALL-LABEL: op_batch_matmul_broadcast_a
|
|
|
|
// NAMED: linalg.generic
|
|
// NAMED-NOT: linalg.batch_matmul
|
|
|
|
// CATEGORY-NOT: linalg.generic
|
|
// CATEGORY: linalg.contract
|
|
|
|
// -----
|
|
|
|
// TODO: this could also be specialized to a named batch_matmul.
|
|
#map_bbcast2_a = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
|
|
#map_bbcast2_b = affine_map<(d0, d1, d2, d3) -> (d3)>
|
|
#map_bbcast2_c = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
|
|
func.func @op_batch_matmul_broadcast_b(%A: tensor<2x16x8xf32>, %B: tensor<8xf32>,
|
|
%Out: tensor<2x16x16xf32>) -> tensor<2x16x16xf32> {
|
|
%0 = linalg.generic
|
|
{indexing_maps = [#map_bbcast2_a, #map_bbcast2_b, #map_bbcast2_c],
|
|
iterator_types = ["parallel", "parallel", "parallel", "reduction"]}
|
|
ins(%A, %B : tensor<2x16x8xf32>, tensor<8xf32>)
|
|
outs(%Out : tensor<2x16x16xf32>) {
|
|
^bb0(%in: f32, %in_0: f32, %out: f32):
|
|
%1 = arith.mulf %in, %in_0 : f32
|
|
%2 = arith.addf %out, %1 : f32
|
|
linalg.yield %2 : f32
|
|
} -> tensor<2x16x16xf32>
|
|
return %0 : tensor<2x16x16xf32>
|
|
}
|
|
|
|
// ALL-LABEL: op_batch_matmul_broadcast_b
|
|
|
|
// NAMED: linalg.generic
|
|
// NAMED-NOT: linalg.batch_matmul
|
|
|
|
// CATEGORY-NOT: linalg.generic
|
|
// CATEGORY: linalg.contract
|
|
|
|
// -----
|
|
|
|
///----------------------------------------------------------------------------------------
|
|
/// Tests for linalg.mmt4d
|
|
///----------------------------------------------------------------------------------------
|
|
|
|
#mapA = affine_map<(m, n, k, m0, n0, k0) -> (m, k, m0, k0)>
|
|
#mapB = affine_map<(m, n, k, m0, n0, k0) -> (n, k, n0, k0)>
|
|
#mapC = affine_map<(m, n, k, m0, n0, k0) -> (m, n, m0, n0)>
|
|
func.func @op_mmt4d(%A: tensor<?x?x?x?xf32>, %B: tensor<?x?x?x?xf32>,
|
|
%C: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> {
|
|
%0 = linalg.generic
|
|
{indexing_maps = [#mapA, #mapB, #mapC],
|
|
iterator_types = ["parallel", "parallel", "reduction",
|
|
"parallel", "parallel", "reduction"]}
|
|
ins(%A, %B : tensor<?x?x?x?xf32>, tensor<?x?x?x?xf32>)
|
|
outs(%C : tensor<?x?x?x?xf32>) {
|
|
^bb0(%in: f32, %in_0: f32, %out: f32):
|
|
%1 = arith.mulf %in, %in_0 : f32
|
|
%2 = arith.addf %out, %1 : f32
|
|
linalg.yield %2 : f32
|
|
} -> tensor<?x?x?x?xf32>
|
|
return %0 : tensor<?x?x?x?xf32>
|
|
}
|
|
|
|
// ALL-LABEL: op_mmt4d
|
|
|
|
// NAMED-NOT: linalg.generic
|
|
// NAMED: linalg.mmt4d
|
|
|
|
// CATEGORY-NOT: linalg.generic
|
|
// CATEGORY: linalg.contract
|
|
|
|
// MMT4D transpose A inner and outer:
|
|
// A is accessed as (k, m, k0, m0) instead of (m, k, m0, k0)
|
|
#map_tA = affine_map<(m, n, k, m0, n0, k0) -> (k, m, k0, m0)>
|
|
func.func @op_mmt4d_transpose_a(%A: tensor<?x?x?x?xf32>, %B: tensor<?x?x?x?xf32>,
|
|
%C: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> {
|
|
%0 = linalg.generic
|
|
{indexing_maps = [#map_tA, #mapB, #mapC],
|
|
iterator_types = ["parallel", "parallel", "reduction",
|
|
"parallel", "parallel", "reduction"]}
|
|
ins(%A, %B : tensor<?x?x?x?xf32>, tensor<?x?x?x?xf32>)
|
|
outs(%C : tensor<?x?x?x?xf32>) {
|
|
^bb0(%in: f32, %in_0: f32, %out: f32):
|
|
%1 = arith.mulf %in, %in_0 : f32
|
|
%2 = arith.addf %out, %1 : f32
|
|
linalg.yield %2 : f32
|
|
} -> tensor<?x?x?x?xf32>
|
|
return %0 : tensor<?x?x?x?xf32>
|
|
}
|
|
|
|
// ALL-LABEL: op_mmt4d_transpose_a
|
|
|
|
// NAMED-NOT: linalg.generic
|
|
// NAMED: linalg.mmt4d
|
|
|
|
// CATEGORY-NOT: linalg.generic
|
|
// CATEGORY: linalg.contract
|
|
|
|
// MMT4D transpose B inner and outer:
|
|
// B is accessed as (k, n, k0, n0) instead of (n, k, n0, k0)
|
|
#map_tB = affine_map<(m, n, k, m0, n0, k0) -> (k, n, k0, n0)>
|
|
func.func @op_mmt4d_transpose_b(%A: tensor<?x?x?x?xf32>, %B: tensor<?x?x?x?xf32>,
|
|
%C: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> {
|
|
%0 = linalg.generic
|
|
{indexing_maps = [#mapA, #map_tB, #mapC],
|
|
iterator_types = ["parallel", "parallel", "reduction",
|
|
"parallel", "parallel", "reduction"]}
|
|
ins(%A, %B : tensor<?x?x?x?xf32>, tensor<?x?x?x?xf32>)
|
|
outs(%C : tensor<?x?x?x?xf32>) {
|
|
^bb0(%in: f32, %in_0: f32, %out: f32):
|
|
%1 = arith.mulf %in, %in_0 : f32
|
|
%2 = arith.addf %out, %1 : f32
|
|
linalg.yield %2 : f32
|
|
} -> tensor<?x?x?x?xf32>
|
|
return %0 : tensor<?x?x?x?xf32>
|
|
}
|
|
|
|
// ALL-LABEL: op_mmt4d_transpose_b
|
|
|
|
// NAMED-NOT: linalg.generic
|
|
// NAMED: linalg.mmt4d
|
|
|
|
// CATEGORY-NOT: linalg.generic
|
|
// CATEGORY: linalg.contract
|
|
|
|
// MMT4D transpose both A and B inner and outer:
|
|
func.func @op_mmt4d_transpose_a_and_b(
|
|
%A: tensor<?x?x?x?xf32>, %B: tensor<?x?x?x?xf32>,
|
|
%C: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> {
|
|
%0 = linalg.generic
|
|
{indexing_maps = [#map_tA, #map_tB, #mapC],
|
|
iterator_types = ["parallel", "parallel", "reduction",
|
|
"parallel", "parallel", "reduction"]}
|
|
ins(%A, %B : tensor<?x?x?x?xf32>, tensor<?x?x?x?xf32>)
|
|
outs(%C : tensor<?x?x?x?xf32>) {
|
|
^bb0(%in: f32, %in_0: f32, %out: f32):
|
|
%1 = arith.mulf %in, %in_0 : f32
|
|
%2 = arith.addf %out, %1 : f32
|
|
linalg.yield %2 : f32
|
|
} -> tensor<?x?x?x?xf32>
|
|
return %0 : tensor<?x?x?x?xf32>
|
|
}
|
|
|
|
// ALL-LABEL: op_mmt4d_transpose_a_and_b
|
|
|
|
// NAMED-NOT: linalg.generic
|
|
// NAMED: linalg.mmt4d
|
|
|
|
// CATEGORY-NOT: linalg.generic
|
|
// CATEGORY: linalg.contract
|
|
|
|
// MMT4D transpose C inner and outer:
|
|
// C is accessed as (n, m, n0, m0) instead of (m, n, m0, n0)
|
|
#map_tC = affine_map<(m, n, k, m0, n0, k0) -> (n, m, n0, m0)>
|
|
func.func @op_mmt4d_transpose_c(%A: tensor<?x?x?x?xf32>, %B: tensor<?x?x?x?xf32>,
|
|
%C: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> {
|
|
%0 = linalg.generic
|
|
{indexing_maps = [#mapA, #mapB, #map_tC],
|
|
iterator_types = ["parallel", "parallel", "reduction",
|
|
"parallel", "parallel", "reduction"]}
|
|
ins(%A, %B : tensor<?x?x?x?xf32>, tensor<?x?x?x?xf32>)
|
|
outs(%C : tensor<?x?x?x?xf32>) {
|
|
^bb0(%in: f32, %in_0: f32, %out: f32):
|
|
%1 = arith.mulf %in, %in_0 : f32
|
|
%2 = arith.addf %out, %1 : f32
|
|
linalg.yield %2 : f32
|
|
} -> tensor<?x?x?x?xf32>
|
|
return %0 : tensor<?x?x?x?xf32>
|
|
}
|
|
|
|
// ALL-LABEL: op_mmt4d_transpose_c
|
|
|
|
// NAMED-NOT: linalg.generic
|
|
// NAMED: linalg.mmt4d
|
|
|
|
// CATEGORY-NOT: linalg.generic
|
|
// CATEGORY: linalg.contract
|
|
|
|
// MMT4D transpose C inner only:
|
|
// C is accessed as (m, n, n0, m0) instead of (m, n, m0, n0)
|
|
#map_tC_inner = affine_map<(m, n, k, m0, n0, k0) -> (m, n, n0, m0)>
|
|
func.func @op_mmt4d_transpose_c_inner(%A: tensor<?x?x?x?xf32>, %B: tensor<?x?x?x?xf32>,
|
|
%C: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> {
|
|
%0 = linalg.generic
|
|
{indexing_maps = [#mapA, #mapB, #map_tC_inner],
|
|
iterator_types = ["parallel", "parallel", "reduction",
|
|
"parallel", "parallel", "reduction"]}
|
|
ins(%A, %B : tensor<?x?x?x?xf32>, tensor<?x?x?x?xf32>)
|
|
outs(%C : tensor<?x?x?x?xf32>) {
|
|
^bb0(%in: f32, %in_0: f32, %out: f32):
|
|
%1 = arith.mulf %in, %in_0 : f32
|
|
%2 = arith.addf %out, %1 : f32
|
|
linalg.yield %2 : f32
|
|
} -> tensor<?x?x?x?xf32>
|
|
return %0 : tensor<?x?x?x?xf32>
|
|
}
|
|
|
|
// ALL-LABEL: op_mmt4d_transpose_c_inner
|
|
|
|
// NAMED-NOT: linalg.generic
|
|
// NAMED: linalg.mmt4d
|
|
|
|
// CATEGORY-NOT: linalg.generic
|
|
// CATEGORY: linalg.contract
|
|
|
|
// Negative MMT4D:
|
|
// A can only be accessed as inner transpose or outer transpose of (m, k, m0, k0)
|
|
#mapA_negative = affine_map<(m, n, k, m0, n0, k0) -> (n, k, n0, k0)>
|
|
func.func @negative_op_mmt4d(%A: tensor<?x?x?x?xf32>, %B: tensor<?x?x?x?xf32>,
|
|
%C: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> {
|
|
%0 = linalg.generic
|
|
{indexing_maps = [#mapA_negative, #mapB, #mapC],
|
|
iterator_types = ["parallel", "parallel", "reduction",
|
|
"parallel", "parallel", "reduction"]}
|
|
ins(%A, %B : tensor<?x?x?x?xf32>, tensor<?x?x?x?xf32>)
|
|
outs(%C : tensor<?x?x?x?xf32>) {
|
|
^bb0(%in: f32, %in_0: f32, %out: f32):
|
|
%1 = arith.mulf %in, %in_0 : f32
|
|
%2 = arith.addf %out, %1 : f32
|
|
linalg.yield %2 : f32
|
|
} -> tensor<?x?x?x?xf32>
|
|
return %0 : tensor<?x?x?x?xf32>
|
|
}
|
|
|
|
// ALL-LABEL: negative_op_mmt4d
|
|
|
|
// NAMED-NOT: linalg.mmt4d
|
|
// NAMED: linalg.generic
|
|
|
|
// CATEGORY-NOT: linalg.generic
|
|
// CATEGORY: linalg.contract
|