Skip to content

Instantly share code, notes, and snippets.

@AmosLewis

AmosLewis/broadcast_to.mlir

Created February 7, 2023 03:26
Show Gist options
  • Save AmosLewis/2890afedc18b40f2860fc9605bf272e0 to your computer and use it in GitHub Desktop.
Save AmosLewis/2890afedc18b40f2860fc9605bf272e0 to your computer and use it in GitHub Desktop.
Download ZIP
Raw
broadcast_to.mlir
func.func @torch.aten.broadcast_to(%arg0: !torch.vtensor<[1,1,1,128],i1>) -> !torch.vtensor<[1,1,128,128],i1> {
%int1 = torch.constant.int 1
%int128 = torch.constant.int 128
%1 = torch.prim.ListConstruct %int1, %int1, %int128, %int128 : (!torch.int, !torch.int, !torch.int, !torch.int) -> !torch.list<int>
%0 = torch.aten.broadcast_to %arg0, %1 : !torch.vtensor<[1,1,1,128],i1>, !torch.list<int> -> !torch.vtensor<[1,1,128,128],i1>
return %0 : !torch.vtensor<[1,1,128,128],i1>
}
@AmosLewis
Copy link
Author 

func.func @torch.aten.broadcast_to(%arg0: !torch.vtensor<[1,1,1,128],i1>) -> !torch.vtensor<[1,1,128,256],i1> {
  %int1 = torch.constant.int 1
  %int128 = torch.constant.int 128
  %int256 = torch.constant.int 256
  %1 = torch.prim.ListConstruct %int1, %int1, %int128, %int256 : (!torch.int, !torch.int, !torch.int, !torch.int) -> !torch.list<int>
  %0 = torch.aten.broadcast_to %arg0, %1 : !torch.vtensor<[1,1,1,128],i1>, !torch.list<int> -> !torch.vtensor<[1,1,128,256],i1>
  return %0 : !torch.vtensor<[1,1,128,256],i1>
}

@AmosLewis
Copy link
Author 

// CHECK-LABEL:   func.func @torch.aten.broadcast_to(
// CHECK-SAME:                                       %[[VAL_0:.*]]: !torch.vtensor<[1,1,1,128],i1>) -> !torch.vtensor<[1,1,128,128],i1> {
// CHECK:           %[[VAL_1:.*]] = torch_c.to_builtin_tensor %[[VAL_0]] : !torch.vtensor<[1,1,1,128],i1> -> tensor<1x1x1x128xi1>
// CHECK:           %[[VAL_2:.*]] = torch.constant.int 1
// CHECK:           %[[VAL_3:.*]] = torch.constant.int 128
// CHECK:           %[[VAL_4:.*]] = torch.prim.ListConstruct %[[VAL_2]], %[[VAL_2]], %[[VAL_3]], %[[VAL_3]] : (!torch.int, !torch.int, !torch.int, !torch.int) -> !torch.list<int>
// CHECK:           %[[VAL_5:.*]] = "tosa.const"() {value = dense<0> : tensor<1x1x128x1xi64>} : () -> tensor<1x1x128x1xi64>
// CHECK:           %[[VAL_6:.*]] = "tosa.add"(%[[VAL_1]], %[[VAL_5]]) : (tensor<1x1x1x128xi1>, tensor<1x1x128x1xi64>) -> tensor<1x1x128x128xi1>
// CHECK:           %[[VAL_7:.*]] = torch_c.from_builtin_tensor %[[VAL_6]] : tensor<1x1x128x128xi1> -> !torch.vtensor<[1,1,128,128],i1>
// CHECK:           return %[[VAL_7]] : !torch.vtensor<[1,1,128,128],i1>
// CHECK:         }

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment