From 4396b53e61897b27b42cf8dc19cc325bc0f22322 Mon Sep 17 00:00:00 2001 From: "Schlimbach, Frank" Date: Fri, 15 Nov 2024 18:02:04 +0100 Subject: [PATCH] fixing ndarray tests --- .../NDArrayToLinalg/NDArrayFusion.mlir | 97 -------- .../Extensions/sharding_propagation.mlir | 22 +- .../NDArray/IR/NDArrayCanonicalize.mlir | 222 +++++++----------- test/Dialect/NDArray/IR/NDArrayOps.mlir | 194 ++++++--------- .../NDArray/Transforms/AddGPURegions.mlir | 74 +++--- 5 files changed, 202 insertions(+), 407 deletions(-) delete mode 100644 test/Conversion/NDArrayToLinalg/NDArrayFusion.mlir diff --git a/test/Conversion/NDArrayToLinalg/NDArrayFusion.mlir b/test/Conversion/NDArrayToLinalg/NDArrayFusion.mlir deleted file mode 100644 index 6c4c50036..000000000 --- a/test/Conversion/NDArrayToLinalg/NDArrayFusion.mlir +++ /dev/null @@ -1,97 +0,0 @@ -// RUN: imex-opt --split-input-file --pass-pipeline="builtin.module(convert-ndarray-to-linalg,func.func(tosa-to-linalg,canonicalize,linalg-fuse-elementwise-ops))" %s -verify-diagnostics -o -| FileCheck %s - -func.func @test_binop_fusion_arith(%arg0: !ndarray.ndarray, %arg1: !ndarray.ndarray) -> !ndarray.ndarray { - %0 = ndarray.ewbin %arg0, %arg1 {op = 0 : i32} : (!ndarray.ndarray, !ndarray.ndarray) -> !ndarray.ndarray - %1 = ndarray.ewbin %0, %arg0 {op = 21 : i32} : (!ndarray.ndarray, !ndarray.ndarray) -> !ndarray.ndarray - %2 = ndarray.ewbin %arg0, %1 {op = 0 : i32} : (!ndarray.ndarray, !ndarray.ndarray) -> !ndarray.ndarray - %3 = ndarray.ewbin %arg1, %2 {op = 21 : i32} : (!ndarray.ndarray, !ndarray.ndarray) -> !ndarray.ndarray - return %3 : !ndarray.ndarray -} -// CHECK-LABEL: @test_binop_fusion_arith -// CHECK-NEXT: arith.constant -// CHECK: memref.dim -// CHECK-NEXT: tensor.empty -// CHECK-NEXT: linalg.generic -// CHECK-NEXT: bb -// CHECK-NEXT: arith.addi -// CHECK-NEXT: arith.muli -// CHECK-NEXT: arith.addi -// CHECK-NEXT: arith.muli -// CHECK: return %{{[0-9]+}} : memref> - -// NOTE tosa ewbinop with dynamic shapes are broadcast-aware and thus do not fuse -func.func @test_binop_fusion_tosa(%arg0: !ndarray.ndarray<5xi64>, %arg1: !ndarray.ndarray<5xi64>) -> !ndarray.ndarray<5xi64> { - %0 = ndarray.ewbin %arg0, %arg1 {op = 4 : i32} : (!ndarray.ndarray<5xi64>, !ndarray.ndarray<5xi64>) -> !ndarray.ndarray<5xi64> - %1 = ndarray.ewbin %0, %arg0 {op = 2 : i32} : (!ndarray.ndarray<5xi64>, !ndarray.ndarray<5xi64>) -> !ndarray.ndarray<5xi64> - %2 = ndarray.ewbin %arg0, %1 {op = 4 : i32} : (!ndarray.ndarray<5xi64>, !ndarray.ndarray<5xi64>) -> !ndarray.ndarray<5xi64> - %3 = ndarray.ewbin %arg1, %2 {op = 2 : i32} : (!ndarray.ndarray<5xi64>, !ndarray.ndarray<5xi64>) -> !ndarray.ndarray<5xi64> - return %3 : !ndarray.ndarray<5xi64> -} -// CHECK-LABEL: @test_binop_fusion_tosa -// CHECK: tensor.empty -// CHECK-NEXT: linalg.generic -// CHECK-NEXT: bb -// CHECK-NEXT: arith.ori -// CHECK-NEXT: arith.andi -// CHECK-NEXT: arith.ori -// CHECK-NEXT: arith.andi -// CHECK: return %{{[0-9]+}} : memref<5xi64, strided<[?], offset: ?>> - -func.func @test_binop_fusion_bcast(%arg0: !ndarray.ndarray, %arg1: !ndarray.ndarray) -> !ndarray.ndarray { - %0 = ndarray.ewbin %arg0, %arg1 {op = 0 : i32} : (!ndarray.ndarray, !ndarray.ndarray) -> !ndarray.ndarray - %1 = ndarray.ewbin %0, %arg0 {op = 21 : i32} : (!ndarray.ndarray, !ndarray.ndarray) -> !ndarray.ndarray - %2 = ndarray.ewbin %arg0, %1 {op = 0 : i32} : (!ndarray.ndarray, !ndarray.ndarray) -> !ndarray.ndarray - %3 = ndarray.ewbin %arg1, %2 {op = 21 : i32} : (!ndarray.ndarray, !ndarray.ndarray) -> !ndarray.ndarray - return %3 : !ndarray.ndarray -} -// CHECK-LABEL: @test_binop_fusion_bcast -// CHECK-NEXT: arith.constant -// CHECK-NEXT: arith.constant -// CHECK: memref.dim -// CHECK-NEXT: memref.dim -// CHECK-NEXT: tensor.empty -// CHECK-NEXT: linalg.generic -// CHECK-NEXT: bb -// CHECK-NEXT: arith.addi -// CHECK-NEXT: arith.muli -// CHECK-NEXT: arith.addi -// CHECK-NEXT: arith.muli -// CHECK: return %{{[0-9]+}} : memref - -func.func @test_binop_fusion_bcast2(%arg0: !ndarray.ndarray, %arg1: !ndarray.ndarray<1xi64>) -> !ndarray.ndarray { - %0 = ndarray.ewbin %arg0, %arg1 {op = 0 : i32} : (!ndarray.ndarray, !ndarray.ndarray<1xi64>) -> !ndarray.ndarray - %1 = ndarray.ewbin %0, %arg0 {op = 21 : i32} : (!ndarray.ndarray, !ndarray.ndarray) -> !ndarray.ndarray - %2 = ndarray.ewbin %arg0, %1 {op = 0 : i32} : (!ndarray.ndarray, !ndarray.ndarray) -> !ndarray.ndarray - %3 = ndarray.ewbin %arg1, %2 {op = 21 : i32} : (!ndarray.ndarray<1xi64>, !ndarray.ndarray) -> !ndarray.ndarray - return %3 : !ndarray.ndarray -} -// CHECK-LABEL: @test_binop_fusion_bcast2 -// CHECK-NEXT: arith.constant -// CHECK: memref.dim -// CHECK-NEXT: tensor.empty -// CHECK-NEXT: linalg.generic -// CHECK-NEXT: bb -// CHECK-NEXT: arith.addi -// CHECK-NEXT: arith.muli -// CHECK-NEXT: arith.addi -// CHECK-NEXT: arith.muli -// CHECK: return %{{[0-9]+}} : memref> - -func.func @test_binop_fusion_bcast3(%arg0: !ndarray.ndarray, %arg1: !ndarray.ndarray<1x?xi64>) -> !ndarray.ndarray { - %0 = ndarray.ewbin %arg0, %arg1 {op = 0 : i32} : (!ndarray.ndarray, !ndarray.ndarray<1x?xi64>) -> !ndarray.ndarray - %1 = ndarray.ewbin %0, %arg1 {op = 21 : i32} : (!ndarray.ndarray, !ndarray.ndarray<1x?xi64>) -> !ndarray.ndarray - %2 = ndarray.ewbin %arg1, %1 {op = 0 : i32} : (!ndarray.ndarray<1x?xi64>, !ndarray.ndarray) -> !ndarray.ndarray - return %2 : !ndarray.ndarray -} -// CHECK-LABEL: @test_binop_fusion_bcast3 -// CHECK-NEXT: arith.constant -// CHECK-NEXT: arith.constant -// CHECK: memref.dim -// CHECK-NEXT: memref.dim -// CHECK-NEXT: tensor.empty -// CHECK-NEXT: linalg.generic -// CHECK-NEXT: bb -// CHECK-NEXT: arith.addi -// CHECK-NEXT: arith.muli -// CHECK-NEXT: arith.addi -// CHECK: return %{{[0-9]+}} : memref> diff --git a/test/Dialect/NDArray/Extensions/sharding_propagation.mlir b/test/Dialect/NDArray/Extensions/sharding_propagation.mlir index a57f3992d..156d035d3 100644 --- a/test/Dialect/NDArray/Extensions/sharding_propagation.mlir +++ b/test/Dialect/NDArray/Extensions/sharding_propagation.mlir @@ -9,7 +9,7 @@ func.func @test_shard_propagate_subview_balanced(%arg0: tensor<1024x1024xi64>) - %s = mesh.sharding @mesh4 split_axes = [[0]] : !mesh.sharding // CHECK: mesh.shard %arg0 to [[S]] : tensor<1024x1024xi64> %0 = mesh.shard %arg0 to %s : tensor<1024x1024xi64> - // CHECK: mesh.sharding @mesh4 split_axes = {{\[\[}}0]] sharded_dims_offsets = [1, 1, 1, 1] : !mesh.sharding + // CHECK: mesh.sharding @mesh4 split_axes = {{\[\[}}0]] sharded_dims_offsets = [0, 1, 2, 3, 4] : !mesh.sharding %1 = ndarray.subview %0[1, 0][4, 3][256, 1] : tensor<1024x1024xi64> to tensor<4x3xi64> return %1 : tensor<4x3xi64> } @@ -20,7 +20,7 @@ func.func @test_shard_propagate_subview_leading(%arg0: tensor<1024x1024xi64>) -> %s = mesh.sharding @mesh4 split_axes = [[0]] : !mesh.sharding // CHECK: mesh.shard %arg0 to [[S]] : tensor<1024x1024xi64> %0 = mesh.shard %arg0 to %s : tensor<1024x1024xi64> - // CHECK: mesh.sharding @mesh4 split_axes = {{\[\[}}0]] sharded_dims_offsets = [3, 0, 0, 0] : !mesh.sharding + // CHECK: mesh.sharding @mesh4 split_axes = {{\[\[}}0]] sharded_dims_offsets = [0, 3, 3, 3, 3] : !mesh.sharding %1 = ndarray.subview %0[0, 0][3, 3][3, 1] : tensor<1024x1024xi64> to tensor<3x3xi64> return %1 : tensor<3x3xi64> } @@ -31,7 +31,7 @@ func.func @test_shard_propagate_subview_mid(%arg0: tensor<1024x1024xi64>) -> ten %s = mesh.sharding @mesh4 split_axes = [[0]] : !mesh.sharding // CHECK: mesh.shard %arg0 to [[S]] : tensor<1024x1024xi64> %0 = mesh.shard %arg0 to %s : tensor<1024x1024xi64> - // CHECK: mesh.sharding @mesh4 split_axes = {{\[\[}}0]] sharded_dims_offsets = [0, 1, 2, 0] : !mesh.sharding + // CHECK: mesh.sharding @mesh4 split_axes = {{\[\[}}0]] sharded_dims_offsets = [0, 0, 1, 3, 3] : !mesh.sharding %1 = ndarray.subview %0[511, 0][3, 3][1, 1] : tensor<1024x1024xi64> to tensor<3x3xi64> return %1 : tensor<3x3xi64> } @@ -42,7 +42,7 @@ func.func @test_shard_propagate_subview_trailing(%arg0: tensor<1024x1024xi64>) - %s = mesh.sharding @mesh4 split_axes = [[0]] : !mesh.sharding // CHECK: mesh.shard %arg0 to [[S]] : tensor<1024x1024xi64> %0 = mesh.shard %arg0 to %s : tensor<1024x1024xi64> - // CHECK: mesh.sharding @mesh4 split_axes = {{\[\[}}0]] sharded_dims_offsets = [0, 0, 0, 3] : !mesh.sharding + // CHECK: mesh.sharding @mesh4 split_axes = {{\[\[}}0]] sharded_dims_offsets = [0, 0, 0, 0, 3] : !mesh.sharding %1 = ndarray.subview %0[1000, 0][3, 3][1, 1] : tensor<1024x1024xi64> to tensor<3x3xi64> return %1 : tensor<3x3xi64> } @@ -53,7 +53,7 @@ func.func @test_shard_propagate_subview_gap(%arg0: tensor<1024x1024xi64>) -> ten %s = mesh.sharding @mesh4 split_axes = [[0]] : !mesh.sharding // CHECK: mesh.shard %arg0 to [[S]] : tensor<1024x1024xi64> %0 = mesh.shard %arg0 to %s : tensor<1024x1024xi64> - // CHECK: mesh.sharding @mesh4 split_axes = {{\[\[}}0]] sharded_dims_offsets = [1, 0, 1, 1] : !mesh.sharding + // CHECK: mesh.sharding @mesh4 split_axes = {{\[\[}}0]] sharded_dims_offsets = [0, 1, 1, 2, 3] : !mesh.sharding %1 = ndarray.subview %0[255, 0][3, 3][257, 1] : tensor<1024x1024xi64> to tensor<3x3xi64> return %1 : tensor<3x3xi64> } @@ -62,9 +62,9 @@ func.func @test_shard_propagate_subview_gap(%arg0: tensor<1024x1024xi64>) -> ten func.func @test_shard_propagate_insert_slice(%arg0: tensor<1024x1024xi64>, %arg1: tensor<3x3xi64>) { %s = mesh.sharding @mesh4 split_axes = [[0]] : !mesh.sharding %0 = mesh.shard %arg0 to %s : tensor<1024x1024xi64> - // CHECK: %[[sharding_2:.*]] = mesh.sharding @mesh4 split_axes = {{\[\[}}0]] sharded_dims_offsets = [3, 0, 0, 0] : !mesh.sharding + // CHECK: %[[sharding_2:.*]] = mesh.sharding @mesh4 split_axes = {{\[\[}}0]] sharded_dims_offsets = [0, 3, 3, 3, 3] : !mesh.sharding // CHECK: %[[sharding_annotated_1:.*]] = mesh.shard %arg1 to %[[sharding_2]] annotate_for_users : tensor<3x3xi64> - // CHECK-NEXT: ndarray.insert_slice %[[sharding_annotated_1]] into + // CHECK: ndarray.insert_slice %[[sharding_annotated_1]] into ndarray.insert_slice %arg1 into %0[0, 0][3, 3][1, 1] : tensor<3x3xi64> into tensor<1024x1024xi64> return } @@ -75,9 +75,9 @@ mesh.mesh @mesh4x4(shape = 4x4) func.func @test_shard_propagate_insert_slice_2d(%arg0: tensor<1024x1024xi64>, %arg1: tensor<3x3xi64>) { %s = mesh.sharding @mesh4x4 split_axes = [[0], [1]] : !mesh.sharding %0 = mesh.shard %arg0 to %s : tensor<1024x1024xi64> - // CHECK: %[[sharding_2:.*]] = mesh.sharding @mesh4x4 split_axes = {{\[\[}}0], [1]] sharded_dims_offsets = [3, 0, 0, 0, 1, 0, 1, 1] : !mesh.sharding + // CHECK: %[[sharding_2:.*]] = mesh.sharding @mesh4x4 split_axes = {{\[\[}}0], [1]] sharded_dims_offsets = [0, 3, 3, 3, 3, 0, 1, 1, 2, 3] : !mesh.sharding // CHECK: %[[sharding_annotated_1:.*]] = mesh.shard %arg1 to %[[sharding_2]] annotate_for_users : tensor<3x3xi64> - // CHECK-NEXT: ndarray.insert_slice %[[sharding_annotated_1]] into + // CHECK: ndarray.insert_slice %[[sharding_annotated_1]] into ndarray.insert_slice %arg1 into %0[0, 255][3, 3][1, 257] : tensor<3x3xi64> into tensor<1024x1024xi64> return } @@ -86,9 +86,9 @@ func.func @test_shard_propagate_insert_slice_2d(%arg0: tensor<1024x1024xi64>, %a func.func @test_shard_propagate_insert_slice_2d_2(%arg0: tensor<1024x1024xi64>, %arg1: tensor<600x3xi64>) { %s = mesh.sharding @mesh4x4 split_axes = [[0], [1]] : !mesh.sharding %0 = mesh.shard %arg0 to %s : tensor<1024x1024xi64> - // CHECK: %[[sharding_2:.*]] = mesh.sharding @mesh4x4 split_axes = {{\[\[}}0], [1]] sharded_dims_offsets = [156, 256, 188, 0, 1, 0, 1, 1] : !mesh.sharding + // CHECK: %[[sharding_2:.*]] = mesh.sharding @mesh4x4 split_axes = {{\[\[}}0], [1]] sharded_dims_offsets = [0, 156, 412, 600, 600, 0, 1, 1, 2, 3] : !mesh.sharding // CHECK: %[[sharding_annotated_1:.*]] = mesh.shard %arg1 to %[[sharding_2]] annotate_for_users : tensor<600x3xi64> - // CHECK-NEXT: ndarray.insert_slice %[[sharding_annotated_1]] into + // CHECK: ndarray.insert_slice %[[sharding_annotated_1]] into ndarray.insert_slice %arg1 into %0[100, 255][600, 3][1, 257] : tensor<600x3xi64> into tensor<1024x1024xi64> return } diff --git a/test/Dialect/NDArray/IR/NDArrayCanonicalize.mlir b/test/Dialect/NDArray/IR/NDArrayCanonicalize.mlir index 5be5bc4b0..95697d344 100644 --- a/test/Dialect/NDArray/IR/NDArrayCanonicalize.mlir +++ b/test/Dialect/NDArray/IR/NDArrayCanonicalize.mlir @@ -1,93 +1,72 @@ // RUN: imex-opt %s -canonicalize | FileCheck %s -func.func @test_subview(%arg0: !ndarray.ndarray) -> !ndarray.ndarray { +func.func @test_subview(%arg0: tensor) -> tensor { %c0 = arith.constant 0 : index %c3 = arith.constant 3 : index - %0 = ndarray.subview %arg0[%c0][%c3][%c3] : !ndarray.ndarray to !ndarray.ndarray - return %0 : !ndarray.ndarray + %0 = ndarray.subview %arg0[%c0][%c3][%c3] : tensor to tensor + return %0 : tensor } // CHECK-LABEL: @test_subview -// CHECK-NEXT: [[C0:%.*]] = ndarray.subview %arg0[0] [3] [3] : !ndarray.ndarray to !ndarray.ndarray<3xi64> -// CHECK-NEXT: [[C1:%.*]] = ndarray.cast %0 : !ndarray.ndarray<3xi64> to !ndarray.ndarray +// CHECK-NEXT: [[C0:%.*]] = ndarray.subview %arg0[0] [3] [3] : tensor to tensor<3xi64> +// CHECK-NEXT: [[C1:%.*]] = tensor.cast %0 : tensor<3xi64> to tensor -func.func @test_subview_cast(%arg0: !ndarray.ndarray<5xi64>) -> !ndarray.ndarray<3xi64> { - %c1 = arith.constant 1 : index - %c3 = arith.constant 3 : index - %0 = ndarray.subview %arg0[%c1] [%c3] [%c1] : !ndarray.ndarray<5xi64> to !ndarray.ndarray - %1 = ndarray.cast %0 : !ndarray.ndarray to !ndarray.ndarray<3xi64> - return %1 : !ndarray.ndarray<3xi64> -} -// CHECK-LABEL: @test_subview_cast -// CHECK-NEXT: [[C0:%.*]] = ndarray.subview %arg0[1] [3] [1] : !ndarray.ndarray<5xi64> to !ndarray.ndarray<3xi64> - -func.func @test_subview_cast2(%arg0: !ndarray.ndarray<5xi64>) -> !ndarray.ndarray<3xi64> { - %0 = ndarray.subview %arg0[0] [3] [1] : !ndarray.ndarray<5xi64> to !ndarray.ndarray<3xi64> - %1 = ndarray.cast %0 : !ndarray.ndarray<3xi64> to !ndarray.ndarray - %2 = ndarray.subview %1[0] [3] [1] : !ndarray.ndarray to !ndarray.ndarray<3xi64> - return %2 : !ndarray.ndarray<3xi64> -} -// CHECK-LABEL: @test_subview_cast2 -// CHECK-NEXT: [[C0:%.*]] = ndarray.subview %arg0[0] [3] [1] : !ndarray.ndarray<5xi64> to !ndarray.ndarray<3xi64> -// CHECK-NEXT: [[C1:%.*]] = ndarray.subview %0[0] [3] [1] : !ndarray.ndarray<3xi64> to !ndarray.ndarray<3xi64> -// CHECK-NEXT: return [[C1:%.*]] : !ndarray.ndarray<3xi64> - -func.func @test_extract_slice(%arg0: !ndarray.ndarray) -> !ndarray.ndarray { +func.func @test_extract_slice(%arg0: tensor) -> tensor { %c0 = arith.constant 0 : index %c3 = arith.constant 3 : index - %0 = ndarray.extract_slice %arg0[%c0][%c3][%c3] : !ndarray.ndarray to !ndarray.ndarray - return %0 : !ndarray.ndarray + %0 = ndarray.extract_slice %arg0[%c0][%c3][%c3] : tensor to tensor + return %0 : tensor } // CHECK-LABEL: @test_extract_slice -// CHECK-NEXT: [[C0:%.*]] = ndarray.extract_slice %arg0[0] [3] [3] : !ndarray.ndarray to !ndarray.ndarray<3xi64> -// CHECK-NEXT: [[C1:%.*]] = ndarray.cast %0 : !ndarray.ndarray<3xi64> to !ndarray.ndarray +// CHECK-NEXT: [[C0:%.*]] = ndarray.extract_slice %arg0[0] [3] [3] : tensor to tensor<3xi64> +// CHECK-NEXT: [[C1:%.*]] = tensor.cast %0 : tensor<3xi64> to tensor -func.func @test_extract_slice_cast(%arg0: !ndarray.ndarray<5xi64>) -> !ndarray.ndarray<3xi64> { +func.func @test_extract_slice_cast(%arg0: tensor<5xi64>) -> tensor<3xi64> { %c1 = arith.constant 1 : index %c3 = arith.constant 3 : index - %0 = ndarray.extract_slice %arg0[%c1] [%c3] [%c1] : !ndarray.ndarray<5xi64> to !ndarray.ndarray - %1 = ndarray.cast %0 : !ndarray.ndarray to !ndarray.ndarray<3xi64> - return %1 : !ndarray.ndarray<3xi64> + %0 = ndarray.extract_slice %arg0[%c1] [%c3] [%c1] : tensor<5xi64> to tensor + %1 = tensor.cast %0 : tensor to tensor<3xi64> + return %1 : tensor<3xi64> } // CHECK-LABEL: @test_extract_slice_cast -// CHECK-NEXT: [[C0:%.*]] = ndarray.extract_slice %arg0[1] [3] [1] : !ndarray.ndarray<5xi64> to !ndarray.ndarray<3xi64> +// CHECK-NEXT: [[C0:%.*]] = ndarray.extract_slice %arg0[1] [3] [1] : tensor<5xi64> to tensor<3xi64> -func.func @test_extract_slice_cast2(%arg0: !ndarray.ndarray<5xi64>) -> !ndarray.ndarray<3xi64> { - %0 = ndarray.extract_slice %arg0[0] [3] [1] : !ndarray.ndarray<5xi64> to !ndarray.ndarray<3xi64> - %1 = ndarray.cast %0 : !ndarray.ndarray<3xi64> to !ndarray.ndarray - %2 = ndarray.extract_slice %1[0] [3] [1] : !ndarray.ndarray to !ndarray.ndarray<3xi64> - return %2 : !ndarray.ndarray<3xi64> +func.func @test_extract_slice_cast2(%arg0: tensor<5xi64>) -> tensor<3xi64> { + %0 = ndarray.extract_slice %arg0[0] [3] [1] : tensor<5xi64> to tensor<3xi64> + %1 = tensor.cast %0 : tensor<3xi64> to tensor + %2 = ndarray.extract_slice %1[0] [3] [1] : tensor to tensor<3xi64> + return %2 : tensor<3xi64> } // CHECK-LABEL: @test_extract_slice_cast2 -// CHECK-NEXT: [[C0:%.*]] = ndarray.extract_slice %arg0[0] [3] [1] : !ndarray.ndarray<5xi64> to !ndarray.ndarray<3xi64> -// CHECK-NEXT: [[C1:%.*]] = ndarray.extract_slice %0[0] [3] [1] : !ndarray.ndarray<3xi64> to !ndarray.ndarray<3xi64> -// CHECK-NEXT: return [[C1:%.*]] : !ndarray.ndarray<3xi64> +// CHECK-NEXT: [[C0:%.*]] = ndarray.extract_slice %arg0[0] [3] [1] : tensor<5xi64> to tensor<3xi64> +// CHECK-NEXT: [[C1:%.*]] = ndarray.extract_slice %0[0] [3] [1] : tensor<3xi64> to tensor<3xi64> +// CHECK-NEXT: return [[C1:%.*]] : tensor<3xi64> -func.func @test_extract_immutable_insert_slice(%arg0: !ndarray.ndarray<16x16xi64>) -> !ndarray.ndarray<3x3xi64> { +func.func @test_extract_immutable_insert_slice(%arg0: tensor<16x16xi64>) -> tensor<3x3xi64> { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index %c3 = arith.constant 3 : index - %0 = ndarray.create %c3, %c3 {dtype = 2 : i8} : (index, index) -> !ndarray.ndarray<3x3xi64> - %1 = ndarray.create %c3, %c3 {dtype = 2 : i8} : (index, index) -> !ndarray.ndarray<3x3xi64> - %2 = ndarray.immutable_insert_slice %0 into %arg0[%c0, %c0] [%c3, %c3] [%c1, %c1] : !ndarray.ndarray<3x3xi64> into !ndarray.ndarray<16x16xi64> - %3 = ndarray.immutable_insert_slice %1 into %2[%c0, %c3] [%c3, %c3] [%c1, %c1] : !ndarray.ndarray<3x3xi64> into !ndarray.ndarray<16x16xi64> - %4 = ndarray.extract_slice %3[0, 0] [3, 3] [1, 1] : !ndarray.ndarray<16x16xi64> to !ndarray.ndarray<3x3xi64> - return %4 : !ndarray.ndarray<3x3xi64> + %0 = ndarray.create %c3, %c3 {dtype = 2 : i8} : (index, index) -> tensor<3x3xi64> + %1 = ndarray.create %c3, %c3 {dtype = 2 : i8} : (index, index) -> tensor<3x3xi64> + %2 = ndarray.immutable_insert_slice %0 into %arg0[%c0, %c0] [%c3, %c3] [%c1, %c1] : tensor<3x3xi64> into tensor<16x16xi64> + %3 = ndarray.immutable_insert_slice %1 into %2[%c0, %c3] [%c3, %c3] [%c1, %c1] : tensor<3x3xi64> into tensor<16x16xi64> + %4 = ndarray.extract_slice %3[0, 0] [3, 3] [1, 1] : tensor<16x16xi64> to tensor<3x3xi64> + return %4 : tensor<3x3xi64> } // CHECK-LABEL: func.func @test_extract_immutable_insert_slice // CHECK: arith.constant // CHECK-NEXT: [[V0:%.*]] = ndarray.create -// CHECK-NEXT: return [[V0]] : !ndarray.ndarray<3x3xi64> +// CHECK-NEXT: return [[V0]] : tensor<3x3xi64> -func.func @test_extract_immutable_insert_slice_overwrite(%arg0: !ndarray.ndarray<16x16xi64>) -> !ndarray.ndarray<3x3xi64> { +func.func @test_extract_immutable_insert_slice_overwrite(%arg0: tensor<16x16xi64>) -> tensor<3x3xi64> { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index %c3 = arith.constant 3 : index - %0 = ndarray.create %c3, %c3 {dtype = 2 : i8} : (index, index) -> !ndarray.ndarray<3x3xi64> - %1 = ndarray.create %c3, %c3 {dtype = 2 : i8} : (index, index) -> !ndarray.ndarray<3x3xi64> - %2 = ndarray.immutable_insert_slice %0 into %arg0[%c0, %c0] [%c3, %c3] [%c1, %c1] : !ndarray.ndarray<3x3xi64> into !ndarray.ndarray<16x16xi64> - %3 = ndarray.immutable_insert_slice %1 into %2[%c0, %c1] [%c3, %c3] [%c1, %c1] : !ndarray.ndarray<3x3xi64> into !ndarray.ndarray<16x16xi64> - %4 = ndarray.extract_slice %3[0, 0] [3, 3] [1, 1] : !ndarray.ndarray<16x16xi64> to !ndarray.ndarray<3x3xi64> - return %4 : !ndarray.ndarray<3x3xi64> + %0 = ndarray.create %c3, %c3 {dtype = 2 : i8} : (index, index) -> tensor<3x3xi64> + %1 = ndarray.create %c3, %c3 {dtype = 2 : i8} : (index, index) -> tensor<3x3xi64> + %2 = ndarray.immutable_insert_slice %0 into %arg0[%c0, %c0] [%c3, %c3] [%c1, %c1] : tensor<3x3xi64> into tensor<16x16xi64> + %3 = ndarray.immutable_insert_slice %1 into %2[%c0, %c1] [%c3, %c3] [%c1, %c1] : tensor<3x3xi64> into tensor<16x16xi64> + %4 = ndarray.extract_slice %3[0, 0] [3, 3] [1, 1] : tensor<16x16xi64> to tensor<3x3xi64> + return %4 : tensor<3x3xi64> } // CHECK-LABEL: func.func @test_extract_immutable_insert_slice_overwrite // CHECK-NEXT: arith.constant @@ -96,126 +75,93 @@ func.func @test_extract_immutable_insert_slice_overwrite(%arg0: !ndarray.ndarray // CHECK-NEXT: ndarray.immutable_insert_slice // CHECK-NEXT: ndarray.immutable_insert_slice // CHECK-NEXT: [[V0:%.*]] = ndarray.extract_slice -// CHECK-NEXT: return [[V0]] : !ndarray.ndarray<3x3xi64> +// CHECK-NEXT: return [[V0]] : tensor<3x3xi64> -func.func @test_extract_immutable_insert_slice_strided(%arg0: !ndarray.ndarray<16xi64>) -> !ndarray.ndarray<3xi64> { +func.func @test_extract_immutable_insert_slice_strided(%arg0: tensor<16xi64>) -> tensor<3xi64> { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index %c3 = arith.constant 3 : index - %0 = ndarray.create %c3 {dtype = 2 : i8} : (index) -> !ndarray.ndarray<3xi64> - %2 = ndarray.immutable_insert_slice %0 into %arg0[%c0] [%c3] [%c3] : !ndarray.ndarray<3xi64> into !ndarray.ndarray<16xi64> - %4 = ndarray.extract_slice %2[0] [3] [%c1] : !ndarray.ndarray<16xi64> to !ndarray.ndarray<3xi64> - return %4 : !ndarray.ndarray<3xi64> + %0 = ndarray.create %c3 {dtype = 2 : i8} : (index) -> tensor<3xi64> + %2 = ndarray.immutable_insert_slice %0 into %arg0[%c0] [%c3] [%c3] : tensor<3xi64> into tensor<16xi64> + %4 = ndarray.extract_slice %2[0] [3] [%c1] : tensor<16xi64> to tensor<3xi64> + return %4 : tensor<3xi64> } // CHECK-LABEL: func.func @test_extract_immutable_insert_slice_strided // CHECK-NEXT: arith.constant // CHECK-NEXT: ndarray.create // CHECK-NEXT: ndarray.immutable_insert_slice // CHECK-NEXT: [[V0:%.*]] = ndarray.extract_slice -// CHECK-NEXT: return [[V0]] : !ndarray.ndarray<3xi64> +// CHECK-NEXT: return [[V0]] : tensor<3xi64> -func.func @test_extract_immutable_insert_slice_strided2(%arg0: !ndarray.ndarray<16xi64>) -> !ndarray.ndarray<3xi64> { +func.func @test_extract_immutable_insert_slice_strided2(%arg0: tensor<16xi64>) -> tensor<3xi64> { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index %c3 = arith.constant 3 : index - %0 = ndarray.create %c3 {dtype = 2 : i8} : (index) -> !ndarray.ndarray<3xi64> - %2 = ndarray.immutable_insert_slice %0 into %arg0[%c0] [%c3] [%c3] : !ndarray.ndarray<3xi64> into !ndarray.ndarray<16xi64> - %4 = ndarray.extract_slice %2[0] [3] [%c3] : !ndarray.ndarray<16xi64> to !ndarray.ndarray<3xi64> - return %4 : !ndarray.ndarray<3xi64> + %0 = ndarray.create %c3 {dtype = 2 : i8} : (index) -> tensor<3xi64> + %2 = ndarray.immutable_insert_slice %0 into %arg0[%c0] [%c3] [%c3] : tensor<3xi64> into tensor<16xi64> + %4 = ndarray.extract_slice %2[0] [3] [%c3] : tensor<16xi64> to tensor<3xi64> + return %4 : tensor<3xi64> } // CHECK-LABEL: func.func @test_extract_immutable_insert_slice_strided2 // CHECK-NEXT: arith.constant // CHECK-NEXT: [[V0:%.*]] = ndarray.create -// CHECK-NEXT: return [[V0]] : !ndarray.ndarray<3xi64> +// CHECK-NEXT: return [[V0]] : tensor<3xi64> -func.func @test_insert_slice(%arg0: !ndarray.ndarray, %arg1: !ndarray.ndarray) -> !ndarray.ndarray { +func.func @test_insert_slice(%arg0: tensor, %arg1: tensor) -> tensor { %c0 = arith.constant 0 : index %c3 = arith.constant 3 : index - ndarray.insert_slice %arg1 into %arg0[%c0] [%c3] [%c3] : !ndarray.ndarray into !ndarray.ndarray - return %arg0 : !ndarray.ndarray + ndarray.insert_slice %arg1 into %arg0[%c0] [%c3] [%c3] : tensor into tensor + return %arg0 : tensor } // CHECK-LABEL: @test_insert_slice -// CHECK-NEXT: [[C0:%.*]] = ndarray.cast %arg1 : !ndarray.ndarray to !ndarray.ndarray<3xi64> -// CHECK-NEXT: ndarray.insert_slice %0 into %arg0[0] [3] [3] : !ndarray.ndarray<3xi64> into !ndarray.ndarray +// CHECK-NEXT: [[C0:%.*]] = tensor.cast %arg1 : tensor to tensor<3xi64> +// CHECK-NEXT: ndarray.insert_slice [[C0]] into %arg0[0] [3] [3] : tensor<3xi64> into tensor -func.func @test_insert_slice_cast(%arg0: !ndarray.ndarray<5xi64>, %arg1: !ndarray.ndarray<3xi64>) -> !ndarray.ndarray<5xi64> { +func.func @test_insert_slice_cast(%arg0: tensor<5xi64>, %arg1: tensor<3xi64>) -> tensor<5xi64> { %c0 = arith.constant 0 : index %c3 = arith.constant 3 : index - %1 = ndarray.cast %arg0 : !ndarray.ndarray<5xi64> to !ndarray.ndarray - %2 = ndarray.cast %arg1 : !ndarray.ndarray<3xi64> to !ndarray.ndarray - ndarray.insert_slice %2 into %1[%c0] [%c3] [%c3] : !ndarray.ndarray into !ndarray.ndarray - return %arg0 : !ndarray.ndarray<5xi64> + %1 = tensor.cast %arg0 : tensor<5xi64> to tensor + %2 = tensor.cast %arg1 : tensor<3xi64> to tensor + ndarray.insert_slice %2 into %1[%c0] [%c3] [%c3] : tensor into tensor + return %arg0 : tensor<5xi64> } // CHECK-LABEL: @test_insert_slice_cast -// CHECK-NEXT: ndarray.insert_slice %arg1 into %arg0[0] [3] [3] : !ndarray.ndarray<3xi64> into !ndarray.ndarray<5xi64> -// CHECK-NEXT: return %arg0 : !ndarray.ndarray<5xi64> +// CHECK-NEXT: ndarray.insert_slice %arg1 into %arg0[0] [3] [3] : tensor<3xi64> into tensor<5xi64> +// CHECK-NEXT: return %arg0 : tensor<5xi64> -func.func @test_immutable_insert_slice(%arg0: !ndarray.ndarray, %arg1: !ndarray.ndarray) -> !ndarray.ndarray { +func.func @test_immutable_insert_slice(%arg0: tensor, %arg1: tensor) -> tensor { %c0 = arith.constant 0 : index %c3 = arith.constant 3 : index - %0 = ndarray.immutable_insert_slice %arg1 into %arg0[%c0] [%c3] [%c3] : !ndarray.ndarray into !ndarray.ndarray - return %0 : !ndarray.ndarray + %0 = ndarray.immutable_insert_slice %arg1 into %arg0[%c0] [%c3] [%c3] : tensor into tensor + return %0 : tensor } // CHECK-LABEL: @test_immutable_insert_slice -// CHECK-NEXT: [[C0:%.*]] = ndarray.cast %arg1 : !ndarray.ndarray to !ndarray.ndarray<3xi64> -// CHECK-NEXT: [[V0:%.*]] = ndarray.immutable_insert_slice %0 into %arg0 [0] [3] [3] : !ndarray.ndarray<3xi64> into !ndarray.ndarray -// CHECK-NEXT: return [[V0:%.*]] : !ndarray.ndarray +// CHECK-NEXT: [[C0:%.*]] = tensor.cast %arg1 : tensor to tensor<3xi64> +// CHECK-NEXT: [[V0:%.*]] = ndarray.immutable_insert_slice [[C0]] into %arg0 [0] [3] [3] : tensor<3xi64> into tensor +// CHECK-NEXT: return [[V0:%.*]] : tensor -func.func @test_immutable_insert_slice_cast(%arg0: !ndarray.ndarray<5xi64>, %arg1: !ndarray.ndarray<3xi64>) -> !ndarray.ndarray<5xi64> { +func.func @test_immutable_insert_slice_cast(%arg0: tensor<5xi64>, %arg1: tensor<3xi64>) -> tensor<5xi64> { %c0 = arith.constant 0 : index %c3 = arith.constant 3 : index - %1 = ndarray.cast %arg0 : !ndarray.ndarray<5xi64> to !ndarray.ndarray - %2 = ndarray.cast %arg1 : !ndarray.ndarray<3xi64> to !ndarray.ndarray - %3 = ndarray.immutable_insert_slice %2 into %1[%c0] [%c3] [%c3] : !ndarray.ndarray into !ndarray.ndarray - %4 = ndarray.cast %3 : !ndarray.ndarray to !ndarray.ndarray<5xi64> - return %4 : !ndarray.ndarray<5xi64> + %1 = tensor.cast %arg0 : tensor<5xi64> to tensor + %2 = tensor.cast %arg1 : tensor<3xi64> to tensor + %3 = ndarray.immutable_insert_slice %2 into %1[%c0] [%c3] [%c3] : tensor into tensor + %4 = tensor.cast %3 : tensor to tensor<5xi64> + return %4 : tensor<5xi64> } // CHECK-LABEL: @test_immutable_insert_slice_cast -// CHECK-NEXT: [[V0:%.*]] = ndarray.immutable_insert_slice %arg1 into %arg0 [0] [3] [3] : !ndarray.ndarray<3xi64> into !ndarray.ndarray<5xi64> -// CHECK-NEXT: return [[V0:%.*]] : !ndarray.ndarray<5xi64> +// CHECK-NEXT: [[V0:%.*]] = ndarray.immutable_insert_slice %arg1 into %arg0 [0] [3] [3] : tensor<3xi64> into tensor<5xi64> +// CHECK-NEXT: return [[V0:%.*]] : tensor<5xi64> -func.func @test_dim(%arg0: !ndarray.ndarray<5xi64>) -> index { - %c0 = arith.constant 0 : index - %1 = ndarray.dim %arg0 %c0 : !ndarray.ndarray<5xi64> -> index - return %1 : index -} -// CHECK-LABEL: @test_dim -// CHECK-NEXT: [[C0:%.*]] = arith.constant 5 : index -// CHECK-NEXT: return [[C0:%.*]] : index - -func.func @test_linspace() -> !ndarray.ndarray { +func.func @test_linspace() -> tensor { %c4 = arith.constant 4 : index %cst_1 = arith.constant 4.000000e+00 : f64 %cst_2 = arith.constant 8.000000e+00 : f64 - %0 = ndarray.linspace %cst_1 %cst_2 %c4 false : (f64, f64, index) -> !ndarray.ndarray - return %0 : !ndarray.ndarray + %0 = ndarray.linspace %cst_1 %cst_2 %c4 false : (f64, f64, index) -> tensor + return %0 : tensor } -// CHECK-LABEL: func.func @test_linspace() -> !ndarray.ndarray { +// CHECK-LABEL: func.func @test_linspace() -> tensor { // CHECK: ndarray.linspace -// CHECK-SAME: !ndarray.ndarray<4xi64> -// CHECK-NEXT: ndarray.cast -// CHECK-SAME: !ndarray.ndarray - -func.func @test_permute_dims_cast(%arg0: !ndarray.ndarray<3x4x5xi64>) -> !ndarray.ndarray { - %0 = ndarray.permute_dims %arg0 [2, 1, 0] : !ndarray.ndarray<3x4x5xi64> -> !ndarray.ndarray - return %0 : !ndarray.ndarray -} -// CHECK-LABEL: func.func @test_permute_dims_cast -// CHECK: ndarray.permute_dims -// CHECK-SAME: !ndarray.ndarray<3x4x5xi64> -> !ndarray.ndarray<5x4x3xi64> -// CHECK-NEXT: ndarray.cast - -func.func @test_permute_dims_identity1(%arg0: !ndarray.ndarray<3x4x5xi64>) -> !ndarray.ndarray { - %0 = ndarray.permute_dims %arg0 [0, 1, 2] : !ndarray.ndarray<3x4x5xi64> -> !ndarray.ndarray - return %0 : !ndarray.ndarray -} -// CHECK-LABEL: func.func @test_permute_dims_identity1 -// CHECK: ndarray.cast %arg0 : !ndarray.ndarray<3x4x5xi64> to !ndarray.ndarray -// CHECK-NOT: ndarray.permute_dims - -func.func @test_permute_dims_identity2(%arg0: !ndarray.ndarray<3x4x5xi64>) -> !ndarray.ndarray<3x4x5xi64> { - %0 = ndarray.permute_dims %arg0 [0, 1, 2] : !ndarray.ndarray<3x4x5xi64> -> !ndarray.ndarray<3x4x5xi64> - return %0 : !ndarray.ndarray<3x4x5xi64> -} -// CHECK-LABEL: func.func @test_permute_dims_identity2 -// CHECK-NOT: ndarray.permute_dims +// CHECK-SAME: tensor<4xi64> +// CHECK-NEXT: tensor.cast +// CHECK-SAME: tensor diff --git a/test/Dialect/NDArray/IR/NDArrayOps.mlir b/test/Dialect/NDArray/IR/NDArrayOps.mlir index 37512425d..3aff3dd41 100644 --- a/test/Dialect/NDArray/IR/NDArrayOps.mlir +++ b/test/Dialect/NDArray/IR/NDArrayOps.mlir @@ -1,213 +1,159 @@ -// RUN: imex-opt %s | sed s/true\>/1\>/g | FileCheck %s +// RUN: imex-opt %s | FileCheck %s // Verify the printed output can be parsed. -// RUN: imex-opt %s | sed s/true\>/1\>/g | imex-opt | FileCheck %s -// RUN: imex-opt -mlir-print-op-generic %s | sed s/true\>/1\>/g | imex-opt | FileCheck %s +// RUN: imex-opt %s | imex-opt | FileCheck %s +// RUN: imex-opt -mlir-print-op-generic %s | imex-opt | FileCheck %s // FIXME sed above, for using 1 instead of true // ----- -func.func @test_subview(%arg0: !ndarray.ndarray) -> !ndarray.ndarray { +func.func @test_subview(%arg0: tensor) -> tensor { %c0 = arith.constant 0 : index %c3 = arith.constant 3 : index - %0 = ndarray.subview %arg0[%c0][%c3][%c3] : !ndarray.ndarray to !ndarray.ndarray - return %0 : !ndarray.ndarray + %0 = ndarray.subview %arg0[%c0][%c3][%c3] : tensor to tensor + return %0 : tensor } // CHECK-LABEL: @test_subview // CHECK-NEXT: [[C0:%.*]] = arith.constant // CHECK-NEXT: [[C1:%.*]] = arith.constant -// CHECK-NEXT: [[V0:%.*]] = ndarray.subview %arg0[[[C0]]] [[[C1]]] [[[C1]]] : !ndarray.ndarray to !ndarray.ndarray -// CHECK-NEXT: return [[V0:%.*]] : !ndarray.ndarray +// CHECK-NEXT: [[V0:%.*]] = ndarray.subview %arg0[[[C0]]] [[[C1]]] [[[C1]]] : tensor to tensor +// CHECK-NEXT: return [[V0:%.*]] : tensor // ----- -func.func @test_subview_const(%arg0: !ndarray.ndarray) -> !ndarray.ndarray<3xi64> { - %0 = ndarray.subview %arg0[0][3][3] : !ndarray.ndarray to !ndarray.ndarray<3xi64> - return %0 : !ndarray.ndarray<3xi64> +func.func @test_subview_const(%arg0: tensor) -> tensor<3xi64> { + %0 = ndarray.subview %arg0[0][3][3] : tensor to tensor<3xi64> + return %0 : tensor<3xi64> } // CHECK-LABEL: @test_subview_const -// CHECK-NEXT: [[V0:%.*]] = ndarray.subview %arg0[0] [3] [3] : !ndarray.ndarray to !ndarray.ndarray<3xi64> -// CHECK-NEXT: return [[V0:%.*]] : !ndarray.ndarray<3xi64> +// CHECK-NEXT: [[V0:%.*]] = ndarray.subview %arg0[0] [3] [3] : tensor to tensor<3xi64> +// CHECK-NEXT: return [[V0:%.*]] : tensor<3xi64> // ----- -func.func @test_extract_slice(%arg0: !ndarray.ndarray) -> !ndarray.ndarray { +func.func @test_extract_slice(%arg0: tensor) -> tensor { %c0 = arith.constant 0 : index %c3 = arith.constant 3 : index - %0 = ndarray.extract_slice %arg0[%c0][%c3][%c3] : !ndarray.ndarray to !ndarray.ndarray - return %0 : !ndarray.ndarray + %0 = ndarray.extract_slice %arg0[%c0][%c3][%c3] : tensor to tensor + return %0 : tensor } // CHECK-LABEL: @test_extract_slice // CHECK-NEXT: [[C0:%.*]] = arith.constant // CHECK-NEXT: [[C1:%.*]] = arith.constant -// CHECK-NEXT: [[V0:%.*]] = ndarray.extract_slice %arg0[[[C0]]] [[[C1]]] [[[C1]]] : !ndarray.ndarray to !ndarray.ndarray -// CHECK-NEXT: return [[V0:%.*]] : !ndarray.ndarray +// CHECK-NEXT: [[V0:%.*]] = ndarray.extract_slice %arg0[[[C0]]] [[[C1]]] [[[C1]]] : tensor to tensor +// CHECK-NEXT: return [[V0:%.*]] : tensor // ----- -func.func @test_extract_slice_const(%arg0: !ndarray.ndarray) -> !ndarray.ndarray<3xi64> { - %0 = ndarray.extract_slice %arg0[0][3][3] : !ndarray.ndarray to !ndarray.ndarray<3xi64> - return %0 : !ndarray.ndarray<3xi64> +func.func @test_extract_slice_const(%arg0: tensor) -> tensor<3xi64> { + %0 = ndarray.extract_slice %arg0[0][3][3] : tensor to tensor<3xi64> + return %0 : tensor<3xi64> } // CHECK-LABEL: @test_extract_slice_const -// CHECK-NEXT: [[V0:%.*]] = ndarray.extract_slice %arg0[0] [3] [3] : !ndarray.ndarray to !ndarray.ndarray<3xi64> -// CHECK-NEXT: return [[V0:%.*]] : !ndarray.ndarray<3xi64> +// CHECK-NEXT: [[V0:%.*]] = ndarray.extract_slice %arg0[0] [3] [3] : tensor to tensor<3xi64> +// CHECK-NEXT: return [[V0:%.*]] : tensor<3xi64> // ----- -func.func @test_insert_slice(%arg0: !ndarray.ndarray, %arg1: !ndarray.ndarray) -> !ndarray.ndarray { +func.func @test_insert_slice(%arg0: tensor, %arg1: tensor) -> tensor { %c0 = arith.constant 0 : index %c3 = arith.constant 3 : index - ndarray.insert_slice %arg1 into %arg0[%c0] [%c3] [%c3] : !ndarray.ndarray into !ndarray.ndarray - return %arg0 : !ndarray.ndarray + ndarray.insert_slice %arg1 into %arg0[%c0] [%c3] [%c3] : tensor into tensor + return %arg0 : tensor } // CHECK-LABEL: @test_insert_slice // CHECK-NEXT: [[C0:%.*]] = arith.constant // CHECK-NEXT: [[C1:%.*]] = arith.constant -// CHECK-NEXT: ndarray.insert_slice %arg1 into %arg0[[[C0]]] [[[C1]]] [[[C1]]] : !ndarray.ndarray into !ndarray.ndarray +// CHECK-NEXT: ndarray.insert_slice %arg1 into %arg0[[[C0]]] [[[C1]]] [[[C1]]] : tensor into tensor // ----- -func.func @test_insert_slice_const(%arg0: !ndarray.ndarray, %arg1: !ndarray.ndarray<3xi64>) -> !ndarray.ndarray { - ndarray.insert_slice %arg1 into %arg0[0] [3] [3] : !ndarray.ndarray<3xi64> into !ndarray.ndarray - return %arg0 : !ndarray.ndarray +func.func @test_insert_slice_const(%arg0: tensor, %arg1: tensor<3xi64>) -> tensor { + ndarray.insert_slice %arg1 into %arg0[0] [3] [3] : tensor<3xi64> into tensor + return %arg0 : tensor } // CHECK-LABEL: @test_insert_slice_const -// CHECK-NEXT: ndarray.insert_slice %arg1 into %arg0[0] [3] [3] : !ndarray.ndarray<3xi64> into !ndarray.ndarray +// CHECK-NEXT: ndarray.insert_slice %arg1 into %arg0[0] [3] [3] : tensor<3xi64> into tensor // ----- -func.func @test_insert_slice_scalar(%arg0: !ndarray.ndarray, %arg1: !ndarray.ndarray) -> !ndarray.ndarray { +func.func @test_insert_slice_scalar(%arg0: tensor, %arg1: tensor) -> tensor { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index %c3 = arith.constant 3 : index - ndarray.insert_slice %arg1 into %arg0[%c0] [%c1] [%c3] : !ndarray.ndarray into !ndarray.ndarray - return %arg0 : !ndarray.ndarray + ndarray.insert_slice %arg1 into %arg0[%c0] [%c1] [%c3] : tensor into tensor + return %arg0 : tensor } // CHECK-LABEL: @test_insert_slice_scalar // CHECK-NEXT: [[C0:%.*]] = arith.constant // CHECK-NEXT: [[C1:%.*]] = arith.constant // CHECK-NEXT: [[C3:%.*]] = arith.constant -// CHECK-NEXT: ndarray.insert_slice %arg1 into %arg0[[[C0]]] [[[C1]]] [[[C3]]] : !ndarray.ndarray into !ndarray.ndarray +// CHECK-NEXT: ndarray.insert_slice %arg1 into %arg0[[[C0]]] [[[C1]]] [[[C3]]] : tensor into tensor // ----- -func.func @test_immutable_insert_slice(%arg0: !ndarray.ndarray, %arg1: !ndarray.ndarray) -> !ndarray.ndarray { +func.func @test_immutable_insert_slice(%arg0: tensor, %arg1: tensor) -> tensor { %c0 = arith.constant 0 : index %c3 = arith.constant 3 : index - %0 = ndarray.immutable_insert_slice %arg1 into %arg0[%c0] [%c3] [%c3] : !ndarray.ndarray into !ndarray.ndarray - return %arg0 : !ndarray.ndarray + %0 = ndarray.immutable_insert_slice %arg1 into %arg0[%c0] [%c3] [%c3] : tensor into tensor + return %arg0 : tensor } // CHECK-LABEL: @test_immutable_insert_slice // CHECK-NEXT: [[C0:%.*]] = arith.constant // CHECK-NEXT: [[C1:%.*]] = arith.constant -// CHECK-NEXT: [[V0:%.*]] = ndarray.immutable_insert_slice %arg1 into %arg0 [[[C0]]] [[[C1]]] [[[C1]]] : !ndarray.ndarray into !ndarray.ndarray -// CHECK-NEXT: return [[V0:%.*]] : !ndarray.ndarray +// CHECK-NEXT: [[V0:%.*]] = ndarray.immutable_insert_slice %arg1 into %arg0 [[[C0]]] [[[C1]]] [[[C1]]] : tensor into tensor +// CHECK-NEXT: return [[V0:%.*]] : tensor // ----- -func.func @test_immutable_insert_slice_const(%arg0: !ndarray.ndarray, %arg1: !ndarray.ndarray<3xi64>) -> !ndarray.ndarray { - %0 = ndarray.immutable_insert_slice %arg1 into %arg0[0] [3] [3] : !ndarray.ndarray<3xi64> into !ndarray.ndarray - return %0 : !ndarray.ndarray +func.func @test_immutable_insert_slice_const(%arg0: tensor, %arg1: tensor<3xi64>) -> tensor { + %0 = ndarray.immutable_insert_slice %arg1 into %arg0[0] [3] [3] : tensor<3xi64> into tensor + return %0 : tensor } // CHECK-LABEL: @test_immutable_insert_slice_const -// CHECK-NEXT: [[V0:%.*]] = ndarray.immutable_insert_slice %arg1 into %arg0 [0] [3] [3] : !ndarray.ndarray<3xi64> into !ndarray.ndarray -// CHECK-NEXT: return [[V0:%.*]] : !ndarray.ndarray +// CHECK-NEXT: [[V0:%.*]] = ndarray.immutable_insert_slice %arg1 into %arg0 [0] [3] [3] : tensor<3xi64> into tensor +// CHECK-NEXT: return [[V0:%.*]] : tensor // ----- -func.func @test_linspace(%arg0: si64, %arg1: si64, %arg2: si64) -> !ndarray.ndarray { - %0 = ndarray.linspace %arg0 %arg1 %arg2 false : (si64, si64, si64) -> !ndarray.ndarray - return %0 : !ndarray.ndarray +func.func @test_linspace(%arg0: si64, %arg1: si64, %arg2: si64) -> tensor { + %0 = ndarray.linspace %arg0 %arg1 %arg2 false : (si64, si64, si64) -> tensor + return %0 : tensor } // CHECK-LABEL: @test_linspace -// CHECK-NEXT: ndarray.linspace %arg0 %arg1 %arg2 false : (si64, si64, si64) -> !ndarray.ndarray +// CHECK-NEXT: ndarray.linspace %arg0 %arg1 %arg2 false : (si64, si64, si64) -> tensor -func.func @test_create(%arg0: index, %arg1: index, %arg2: index, %arg3: i64) -> !ndarray.ndarray { - %0 = ndarray.create %arg0, %arg1, %arg2 {dtype = 0 : i8} : (index, index, index) -> !ndarray.ndarray - return %0 : !ndarray.ndarray +func.func @test_create(%arg0: index, %arg1: index, %arg2: index, %arg3: i64) -> tensor { + %0 = ndarray.create %arg0, %arg1, %arg2 {dtype = 0 : i8} : (index, index, index) -> tensor + return %0 : tensor } // CHECK-LABEL: @test_create -// CHECK: %arg0, %arg1, %arg2 {dtype = 0 : i8} : (index, index, index) -> !ndarray.ndarray +// CHECK: %arg0, %arg1, %arg2 {dtype = 0 : i8} : (index, index, index) -> tensor -func.func @test_create2(%arg0: index, %arg1: index, %arg2: index, %arg3: i64) -> !ndarray.ndarray { - %0 = ndarray.create %arg0, %arg1, %arg2 value %arg3 {environment = 3 : i64, team = 3 : i64, dtype = 2 : i8} : (index, index, index, i64) -> !ndarray.ndarray - return %0 : !ndarray.ndarray +func.func @test_create2(%arg0: index, %arg1: index, %arg2: index, %arg3: i64) -> tensor { + %0 = ndarray.create %arg0, %arg1, %arg2 value %arg3 {environment = 3 : i64, team = 3 : i64, dtype = 2 : i8} : (index, index, index, i64) -> tensor + return %0 : tensor } // CHECK-LABEL: @test_create2 -// CHECK: ndarray.create %arg0, %arg1, %arg2 value %arg3 {dtype = 2 : i8, environment = 3 : i64, team = 3 : i64} : (index, index, index, i64) -> !ndarray.ndarray +// CHECK: ndarray.create %arg0, %arg1, %arg2 value %arg3 {dtype = 2 : i8, environment = 3 : i64, team = 3 : i64} : (index, index, index, i64) -> tensor // ----- -func.func @test_reshape(%arg0: index) -> !ndarray.ndarray { - %0 = ndarray.create %arg0 {dtype = 2 : i8} : (index) -> !ndarray.ndarray +func.func @test_reshape(%arg0: index) -> tensor { + %0 = ndarray.create %arg0 {dtype = 2 : i8} : (index) -> tensor %c0 = arith.constant 0 : index %c3 = arith.constant 3 : index - %1 = "ndarray.reshape"(%0, %c0, %c3) : (!ndarray.ndarray, index, index) -> !ndarray.ndarray - return %1 : !ndarray.ndarray + %1 = "ndarray.reshape"(%0, %c0, %c3) : (tensor, index, index) -> tensor + return %1 : tensor } // CHECK-LABEL: @test_reshape // CHECK: ndarray.create // CHECK: ndarray.reshape -// CHECK-SAME: -> !ndarray.ndarray +// CHECK-SAME: -> tensor // ----- -func.func @test_ewbin(%arg0: !ndarray.ndarray) -> !ndarray.ndarray { - %0 = ndarray.ewbin %arg0, %arg0 {op = 0 : i32} : (!ndarray.ndarray, !ndarray.ndarray) -> !ndarray.ndarray - return %0 : !ndarray.ndarray -} -// CHECK-LABEL: @test_ewbin -// CHECK-NEXT: ndarray.ewbin %arg0, %arg0 {op = 0 : i32} : (!ndarray.ndarray, !ndarray.ndarray) -> !ndarray.ndarray - -// ----- -func.func @test_reduction(%arg0: !ndarray.ndarray) -> si64 { - %0 = ndarray.reduction %arg0 {op = 4 : i32} : !ndarray.ndarray -> !ndarray.ndarray - %1 = builtin.unrealized_conversion_cast %0 : !ndarray.ndarray to si64 - return %1 : si64 -} -// CHECK-LABEL: @test_reduction -// CHECK-NEXT: ndarray.reduction %arg0 {op = 4 : i32} : !ndarray.ndarray -> !ndarray.ndarray - -// ----- -func.func @test_dim(%arg0: !ndarray.ndarray) -> index { - %c0 = arith.constant 0 : index - %1 = ndarray.dim %arg0 %c0 : !ndarray.ndarray -> index - return %1 : index -} -// CHECK-LABEL: func.func @test_dim -// CHECK: [[V0:%.*]] = ndarray.dim -// CHECK-NEXT: return [[V0]] : index - -// ----- -func.func @test_cast(%arg0: !ndarray.ndarray<5xi64>) -> !ndarray.ndarray { - %0 = ndarray.cast %arg0 : !ndarray.ndarray<5xi64> to !ndarray.ndarray - return %0 : !ndarray.ndarray -} -// CHECK-LABEL: func.func @test_cast -// CHECK: [[V0:%.*]] = ndarray.cast -// CHECK-NEXT: return [[V0]] : !ndarray.ndarray - -// ----- -func.func @test_copy(%arg0: !ndarray.ndarray<5xi64>) -> !ndarray.ndarray<5xi64> { - %0 = ndarray.copy %arg0 : !ndarray.ndarray<5xi64> -> !ndarray.ndarray<5xi64> - return %0 : !ndarray.ndarray<5xi64> +func.func @test_copy(%arg0: tensor<5xi64>) -> tensor<5xi64> { + %0 = ndarray.copy %arg0 : tensor<5xi64> -> tensor<5xi64> + return %0 : tensor<5xi64> } // CHECK-LABEL: func.func @test_copy // CHECK: [[V0:%.*]] = ndarray.copy -// CHECK-NEXT: return [[V0]] : !ndarray.ndarray<5xi64> +// CHECK-NEXT: return [[V0]] : tensor<5xi64> // ----- -func.func @test_castelem(%arg0: !ndarray.ndarray<5xi64>) -> !ndarray.ndarray<5xi32> { - %0 = ndarray.cast_elemtype %arg0 : !ndarray.ndarray<5xi64> to !ndarray.ndarray<5xi32> - return %0 : !ndarray.ndarray<5xi32> +func.func @test_castelem(%arg0: tensor<5xi64>) -> tensor<5xi32> { + %0 = ndarray.cast_elemtype %arg0 : tensor<5xi64> to tensor<5xi32> + return %0 : tensor<5xi32> } // CHECK-LABEL: func.func @test_castelem // CHECK: [[V0:%.*]] = ndarray.cast_elemtype -// CHECK-NEXT: return [[V0]] : !ndarray.ndarray<5xi32> - -// ----- -func.func @test_from_memref(%arg0: memref>) -> !ndarray.ndarray { - %0 = ndarray.from_memref %arg0 : memref> -> !ndarray.ndarray - return %0 : !ndarray.ndarray -} -// CHECK-LABEL: func.func @test_from_memref -// CHECK: [[V0:%.*]] = ndarray.from_memref -// CHECK-NEXT: return [[V0]] : !ndarray.ndarray - -// ----- -func.func @test_permute_dims(%arg0: !ndarray.ndarray) -> !ndarray.ndarray { - %0 = ndarray.permute_dims %arg0 [0, 1, 2] : !ndarray.ndarray -> !ndarray.ndarray - return %0 : !ndarray.ndarray -} -// CHECK-LABEL: func.func @test_permute_dims -// CHECK: [[V0:%.*]] = ndarray.permute_dims -// CHECK-NEXT: return [[V0]] : !ndarray.ndarray +// CHECK-NEXT: return [[V0]] : tensor<5xi32> diff --git a/test/Dialect/NDArray/Transforms/AddGPURegions.mlir b/test/Dialect/NDArray/Transforms/AddGPURegions.mlir index 0db09a635..6a00cd4ce 100644 --- a/test/Dialect/NDArray/Transforms/AddGPURegions.mlir +++ b/test/Dialect/NDArray/Transforms/AddGPURegions.mlir @@ -7,76 +7,76 @@ func.func @test_region(%arg0: i64, %arg1: i64, %arg2: i64) -> i64 { %c22 = arith.constant 22 : index %v = arith.constant 55 : i64 %s = arith.index_cast %arg0 : i64 to index - %0 = ndarray.linspace %arg0 %arg1 %c33 false {device = "XeGPU", team = 1 : i64} : (i64, i64, i64) -> !ndarray.ndarray<33xi64, #region.gpu_env> - %1 = ndarray.create %c22 value %v {dtype = 2 : i8, device = "XeGPU", team = 1 : i64} : (index, i64) -> !ndarray.ndarray> - %10 = ndarray.subview %0[%c0][22][%c3] : !ndarray.ndarray<33xi64, #region.gpu_env> to !ndarray.ndarray> - %20 = ndarray.ewbin %10, %1 {op = 0 : i32} : (!ndarray.ndarray>, !ndarray.ndarray>) -> !ndarray.ndarray> - %21 = ndarray.reduction %20 {op = 4 : i32} : !ndarray.ndarray> -> !ndarray.ndarray> - %30 = builtin.unrealized_conversion_cast %21 : !ndarray.ndarray> to i64 - ndarray.delete %0 : !ndarray.ndarray<33xi64, #region.gpu_env> - ndarray.delete %1 : !ndarray.ndarray> + %0 = ndarray.linspace %arg0 %arg1 %c33 false {device = "XeGPU", team = 1 : i64} : (i64, i64, i64) -> tensor<33xi64, #region.gpu_env> + %1 = ndarray.create %c22 value %v {dtype = 2 : i8, device = "XeGPU", team = 1 : i64} : (index, i64) -> tensor> + %10 = ndarray.subview %0[%c0][22][%c3] : tensor<33xi64, #region.gpu_env> to tensor> + %20 = ndarray.ewbin %10, %1 {op = 0 : i32} : (tensor>, tensor>) -> tensor> + %21 = ndarray.reduction %20 {op = 4 : i32} : tensor> -> tensor> + %30 = builtin.unrealized_conversion_cast %21 : tensor> to i64 + ndarray.delete %0 : tensor<33xi64, #region.gpu_env> + ndarray.delete %1 : tensor> return %30 : i64 } // CHECK-LABEL: func.func @test_region -// CHECK: [[V0:%.*]] = region.env_region #region.gpu_env -> !ndarray.ndarray<33xi64, #region.gpu_env> { +// CHECK: [[V0:%.*]] = region.env_region #region.gpu_env -> tensor<33xi64, #region.gpu_env> { // CHECK-NEXT: ndarray.linspace // CHECK-NEXT: region.env_region_yield -// CHECK: [[V1:%.*]] = region.env_region #region.gpu_env -> !ndarray.ndarray> { +// CHECK: [[V1:%.*]] = region.env_region #region.gpu_env -> tensor> { // CHECK-NEXT: ndarray.create // CHECK-NEXT: region.env_region_yield -// CHECK: [[V2:%.*]] = region.env_region #region.gpu_env -> !ndarray.ndarray> { +// CHECK: [[V2:%.*]] = region.env_region #region.gpu_env -> tensor> { // CHECK-NEXT: ndarray.subview [[V0]] // CHECK-NEXT: region.env_region_yield -// CHECK: [[V3:%.*]] = region.env_region #region.gpu_env -> !ndarray.ndarray> { +// CHECK: [[V3:%.*]] = region.env_region #region.gpu_env -> tensor> { // CHECK-NEXT: ndarray.ewbin [[V2]], [[V1]] -// CHECK: [[V4:%.*]] = region.env_region #region.gpu_env -> !ndarray.ndarray> { +// CHECK: [[V4:%.*]] = region.env_region #region.gpu_env -> tensor> { // CHECK-NEXT: ndarray.reduction [[V3]] // CHECK-NEXT: region.env_region_yield // CHECK-NEXT: } // CHECK-NEXT: [[V5:%.*]] = builtin.unrealized_conversion_cast // CHECK: region.env_region #region.gpu_env { -// CHECK-NEXT: ndarray.delete [[V0]] : !ndarray.ndarray<33xi64, #region.gpu_env> +// CHECK-NEXT: ndarray.delete [[V0]] : tensor<33xi64, #region.gpu_env> // CHECK-NEXT: } // CHECK-NEXT: region.env_region #region.gpu_env { -// CHECK-NEXT: ndarray.delete [[V1]] : !ndarray.ndarray> +// CHECK-NEXT: ndarray.delete [[V1]] : tensor> // CHECK-NEXT: } // CHECK-NEXT: return [[V5]] // ----- -func.func @test_copy() -> !ndarray.ndarray<33xi64> { +func.func @test_copy() -> tensor<33xi64> { %c0 = arith.constant 0 : i64 %c3 = arith.constant 3 : i64 %c33 = arith.constant 33 : i64 - %0 = ndarray.linspace %c0 %c3 %c33 false {device = "XeGPU", team = 1 : i64} : (i64, i64, i64) -> !ndarray.ndarray<33xi64, #region.gpu_env> - %1 = ndarray.copy %0 : !ndarray.ndarray<33xi64, #region.gpu_env> -> !ndarray.ndarray<33xi64> - %2 = ndarray.copy %1 : !ndarray.ndarray<33xi64> -> !ndarray.ndarray<33xi64, #region.gpu_env> - %3 = ndarray.copy %2 : !ndarray.ndarray<33xi64, #region.gpu_env> -> !ndarray.ndarray<33xi64, #region.gpu_env> - %4 = ndarray.copy %3 : !ndarray.ndarray<33xi64, #region.gpu_env> -> !ndarray.ndarray<33xi64> - %5 = ndarray.copy %4 : !ndarray.ndarray<33xi64> -> !ndarray.ndarray<33xi64> - return %5 : !ndarray.ndarray<33xi64> + %0 = ndarray.linspace %c0 %c3 %c33 false {device = "XeGPU", team = 1 : i64} : (i64, i64, i64) -> tensor<33xi64, #region.gpu_env> + %1 = ndarray.copy %0 : tensor<33xi64, #region.gpu_env> -> tensor<33xi64> + %2 = ndarray.copy %1 : tensor<33xi64> -> tensor<33xi64, #region.gpu_env> + %3 = ndarray.copy %2 : tensor<33xi64, #region.gpu_env> -> tensor<33xi64, #region.gpu_env> + %4 = ndarray.copy %3 : tensor<33xi64, #region.gpu_env> -> tensor<33xi64> + %5 = ndarray.copy %4 : tensor<33xi64> -> tensor<33xi64> + return %5 : tensor<33xi64> } -// CHECK-LABEL: func.func @test_copy() -> !ndarray.ndarray<33xi64> { -// CHECK: region.env_region #region.gpu_env -> !ndarray.ndarray<33xi64, #region.gpu_env> { +// CHECK-LABEL: func.func @test_copy() -> tensor<33xi64> { +// CHECK: region.env_region #region.gpu_env -> tensor<33xi64, #region.gpu_env> { // CHECK: ndarray.linspace -// CHECK-SAME: -> !ndarray.ndarray<33xi64, #region.gpu_env> +// CHECK-SAME: -> tensor<33xi64, #region.gpu_env> // CHECK: region.env_region_yield -// CHECK-SAME: !ndarray.ndarray<33xi64, #region.gpu_env> +// CHECK-SAME: tensor<33xi64, #region.gpu_env> // CHECK: ndarray.copy -// CHECK-SAME: !ndarray.ndarray<33xi64, #region.gpu_env> -> !ndarray.ndarray<33xi64> -// CHECK: region.env_region #region.gpu_env -> !ndarray.ndarray<33xi64, #region.gpu_env> { +// CHECK-SAME: tensor<33xi64, #region.gpu_env> -> tensor<33xi64> +// CHECK: region.env_region #region.gpu_env -> tensor<33xi64, #region.gpu_env> { // CHECK: ndarray.copy -// CHECK-SAME: !ndarray.ndarray<33xi64> -> !ndarray.ndarray<33xi64, #region.gpu_env> +// CHECK-SAME: tensor<33xi64> -> tensor<33xi64, #region.gpu_env> // CHECK: region.env_region_yield -// CHECK-SAME: !ndarray.ndarray<33xi64, #region.gpu_env> -// CHECK: region.env_region #region.gpu_env -> !ndarray.ndarray<33xi64, #region.gpu_env> { +// CHECK-SAME: tensor<33xi64, #region.gpu_env> +// CHECK: region.env_region #region.gpu_env -> tensor<33xi64, #region.gpu_env> { // CHECK: ndarray.copy -// CHECK-SAME: !ndarray.ndarray<33xi64, #region.gpu_env> -> !ndarray.ndarray<33xi64, #region.gpu_env> +// CHECK-SAME: tensor<33xi64, #region.gpu_env> -> tensor<33xi64, #region.gpu_env> // CHECK: region.env_region_yield -// CHECK-SAME: !ndarray.ndarray<33xi64, #region.gpu_env> +// CHECK-SAME: tensor<33xi64, #region.gpu_env> // CHECK: ndarray.copy -// CHECK-SAME: !ndarray.ndarray<33xi64, #region.gpu_env> -> !ndarray.ndarray<33xi64> +// CHECK-SAME: tensor<33xi64, #region.gpu_env> -> tensor<33xi64> // CHECK: ndarray.copy -// CHECK-SAME: !ndarray.ndarray<33xi64> -> !ndarray.ndarray<33xi64> +// CHECK-SAME: tensor<33xi64> -> tensor<33xi64> // CHECK: return -// CHECK-SAME: !ndarray.ndarray<33xi64> +// CHECK-SAME: tensor<33xi64>