23 lines · plain
1// RUN: mlir-opt -test-grid-simplifications %s | FileCheck %s2 3shard.grid @grid0(shape = 4x?x2)4shard.grid @grid1(shape = 2x3)5 6// CHECK-LABEL: func.func @grid_shape_op_folding7func.func @grid_shape_op_folding() -> (index, index) {8 // CHECK: %[[AXIS_2_SIZE:.*]] = arith.constant 2 : index9 // CHECK: %[[AXIS_1_SIZE:.*]] = shard.grid_shape @grid0 axes = [1] : index10 %0:2 = shard.grid_shape @grid0 axes = [2, 1] : index, index11 // CHECK: return %[[AXIS_2_SIZE]], %[[AXIS_1_SIZE]]12 return %0#0, %0#1 : index, index13}14 15// CHECK-LABEL: func.func @grid_shape_op_folding_all_axes_static_grid16func.func @grid_shape_op_folding_all_axes_static_grid() -> (index, index) {17 // CHECK: %[[AXIS_0_SIZE:.*]] = arith.constant 2 : index18 // CHECK: %[[AXIS_1_SIZE:.*]] = arith.constant 3 : index19 %0:2 = shard.grid_shape @grid1 : index, index20 // CHECK: return %[[AXIS_0_SIZE]], %[[AXIS_1_SIZE]]21 return %0#0, %0#1 : index, index22}23