9a8437f504
Rename * Op mesh.cluster -> mesh.mesh * Op mesh.cluster_shape -> mesh.mesh_shape * variables and attributes. The name `mesh` is more specific to what it really represents. It is a mesh of devices. The name `cluster` implies a broader posibility of device configurations. When just the word `mesh` is used the meaning can often be inferred from the context whether it refers to the mesh dialect or a device mesh. The full name can be used when needed.
23 lines
920 B
MLIR
23 lines
920 B
MLIR
// RUN: mlir-opt -test-mesh-simplifications %s | FileCheck %s
|
|
|
|
mesh.mesh @mesh0(shape = 4x?x2)
|
|
mesh.mesh @mesh1(shape = 2x3)
|
|
|
|
// CHECK-LABEL: func.func @mesh_shape_op_folding
|
|
func.func @mesh_shape_op_folding() -> (index, index) {
|
|
// CHECK: %[[AXIS_2_SIZE:.*]] = arith.constant 2 : index
|
|
// CHECK: %[[AXIS_1_SIZE:.*]] = mesh.mesh_shape @mesh0 axes = [1] : index
|
|
%0:2 = mesh.mesh_shape @mesh0 axes = [2, 1] : index, index
|
|
// CHECK: return %[[AXIS_2_SIZE]], %[[AXIS_1_SIZE]]
|
|
return %0#0, %0#1 : index, index
|
|
}
|
|
|
|
// CHECK-LABEL: func.func @mesh_shape_op_folding_all_axes_static_mesh
|
|
func.func @mesh_shape_op_folding_all_axes_static_mesh() -> (index, index) {
|
|
// CHECK: %[[AXIS_0_SIZE:.*]] = arith.constant 2 : index
|
|
// CHECK: %[[AXIS_1_SIZE:.*]] = arith.constant 3 : index
|
|
%0:2 = mesh.mesh_shape @mesh1 : index, index
|
|
// CHECK: return %[[AXIS_0_SIZE]], %[[AXIS_1_SIZE]]
|
|
return %0#0, %0#1 : index, index
|
|
}
|