18 lines · plain
1// RUN: mlir-opt \2// RUN: --pass-pipeline="builtin.module(func.func(shard-partition))" \3// RUN: %s | FileCheck %s4 5shard.grid @grid4x4(shape = 4x4)6 7// CHECK-LABEL: func @test_partition_constant8// CHECK-NEXT: [[vcst:%.*]] = arith.constant dense<0.000000e+00> :9// tensor<256x1024xf32> CHECK-NEXT: [[vc434_i32:%.*]] = arith.constant 434 :10// i32 CHECK-NEXT: return [[vcst]] : tensor<256x1024xf32>11func.func @test_partition_constant() ->(tensor<1024x1024xf32>)attributes{llvm.emit_c_interface} {12 %cst = arith.constant dense<0.000000e+00> : tensor<1024x1024xf32>13 %sharding_1 = shard.sharding @grid4x4 split_axes = [[0]] : !shard.sharding14 %sharded_1 = shard.shard %cst to %sharding_1 : tensor<1024x1024xf32>15 %ci = arith.constant 434 : i3216 return %sharded_1 : tensor<1024x1024xf32>17}18