Changeset View
Changeset View
Standalone View
Standalone View
mlir/test/Dialect/Shape/canonicalize.mlir
Show First 20 Lines • Show All 1,313 Lines • ▼ Show 20 Lines | |||||
// CHECK-LABEL: @cast_extent_tensor | // CHECK-LABEL: @cast_extent_tensor | ||||
func @cast_extent_tensor(%arg : tensor<*xf32>) -> tensor<3xindex> { | func @cast_extent_tensor(%arg : tensor<*xf32>) -> tensor<3xindex> { | ||||
// CHECK: tensor.cast %{{.*}} : tensor<?xindex> to tensor<3xindex> | // CHECK: tensor.cast %{{.*}} : tensor<?xindex> to tensor<3xindex> | ||||
%0 = shape.shape_of %arg : tensor<*xf32> -> tensor<?xindex> | %0 = shape.shape_of %arg : tensor<*xf32> -> tensor<?xindex> | ||||
%1 = tensor.cast %0 : tensor<?xindex> to tensor<3xindex> | %1 = tensor.cast %0 : tensor<?xindex> to tensor<3xindex> | ||||
return %1 : tensor<3xindex> | return %1 : tensor<3xindex> | ||||
} | } | ||||
// ---- | // ----- | ||||
// CHECK-LABEL: max_same_arg | // CHECK-LABEL: max_same_arg | ||||
// CHECK-SAME: (%[[SHAPE:.*]]: !shape.shape) | // CHECK-SAME: (%[[SHAPE:.*]]: !shape.shape) | ||||
func @max_same_arg(%a: !shape.shape) -> !shape.shape { | func @max_same_arg(%a: !shape.shape) -> !shape.shape { | ||||
%1 = shape.max %a, %a : !shape.shape, !shape.shape -> !shape.shape | %1 = shape.max %a, %a : !shape.shape, !shape.shape -> !shape.shape | ||||
// CHECK: return %[[SHAPE]] | // CHECK: return %[[SHAPE]] | ||||
return %1 : !shape.shape | return %1 : !shape.shape | ||||
} | } | ||||
// ---- | // ----- | ||||
// CHECK-LABEL: min_same_arg | // CHECK-LABEL: min_same_arg | ||||
// CHECK-SAME: (%[[SHAPE:.*]]: !shape.shape) | // CHECK-SAME: (%[[SHAPE:.*]]: !shape.shape) | ||||
func @min_same_arg(%a: !shape.shape) -> !shape.shape { | func @min_same_arg(%a: !shape.shape) -> !shape.shape { | ||||
%1 = shape.min %a, %a : !shape.shape, !shape.shape -> !shape.shape | %1 = shape.min %a, %a : !shape.shape, !shape.shape -> !shape.shape | ||||
// CHECK: return %[[SHAPE]] | // CHECK: return %[[SHAPE]] | ||||
return %1 : !shape.shape | return %1 : !shape.shape | ||||
} | } | ||||
// ---- | // ----- | ||||
// CHECK-LABEL: @cstr_broadcastable_folding | // CHECK-LABEL: @cstr_broadcastable_folding | ||||
func @cstr_broadcastable_folding(%arg : tensor<?x4xf32>) { | func @cstr_broadcastable_folding(%arg : tensor<?x4xf32>) { | ||||
// CHECK: const_witness true | // CHECK: const_witness true | ||||
%0 = shape.shape_of %arg : tensor<?x4xf32> -> tensor<2xindex> | %0 = shape.shape_of %arg : tensor<?x4xf32> -> tensor<2xindex> | ||||
%1 = shape.const_shape [4] : tensor<1xindex> | %1 = shape.const_shape [4] : tensor<1xindex> | ||||
%2 = shape.cstr_broadcastable %0, %1: tensor<2xindex>, tensor<1xindex> | %2 = shape.cstr_broadcastable %0, %1: tensor<2xindex>, tensor<1xindex> | ||||
"use"(%2) : (!shape.witness) -> () | "use"(%2) : (!shape.witness) -> () | ||||
Show All 34 Lines |