Skip to content

Instantly share code, notes, and snippets.

@AmosLewis

AmosLewis/where.mlir

Created October 6, 2022 17:36
Show Gist options
  • Save AmosLewis/32847885f8b3ff27b7ef6564154fec59 to your computer and use it in GitHub Desktop.
Save AmosLewis/32847885f8b3ff27b7ef6564154fec59 to your computer and use it in GitHub Desktop.
Download ZIP
wheremlirrun
Raw
where.mlir
func.func @torch.aten.where.self(%arg0: !torch.vtensor<[1,1,5,5],i1>, %arg1: !torch.vtensor<[1,12,5,5],f32>, %arg2: !torch.vtensor<[],f32>) -> !torch.vtensor<[1,12,5,5],f32> {
%0 = torch.aten.where.self %arg0, %arg1, %arg2 : !torch.vtensor<[1,1,5,5],i1>, !torch.vtensor<[1,12,5,5],f32>, !torch.vtensor<[],f32> -> !torch.vtensor<[1,12,5,5],f32>
return %0 : !torch.vtensor<[1,12,5,5],f32>
}
@AmosLewis
Copy link
Author 

(mlir_venv) nod% torch-mlir-opt -pass-pipeline='torch-backend-to-tosa-backend-pipeline' /tmp/where.mlir  | externals/llvm-project/mlir/utils/generate-test-checks.py
// NOTE: Assertions have been autogenerated by utils/generate-test-checks.py

// The script is designed to make adding checks to
// a test case fast, it is *not* designed to be authoritative
// about what constitutes a good test! The CHECK should be
// minimized and named to reflect the test intent.



// CHECK-LABEL:   func.func @torch.aten.where.self(
// CHECK-SAME:                                     %[[VAL_0:.*]]: tensor<1x1x5x5xi1>,
// CHECK-SAME:                                     %[[VAL_1:.*]]: tensor<1x12x5x5xf32>,
// CHECK-SAME:                                     %[[VAL_2:.*]]: tensor<f32>) -> tensor<1x12x5x5xf32> {
// CHECK:           %[[VAL_3:.*]] = "tosa.select"(%[[VAL_0]], %[[VAL_1]], %[[VAL_2]]) : (tensor<1x1x5x5xi1>, tensor<1x12x5x5xf32>, tensor<f32>) -> tensor<1x12x5x5xf32>
// CHECK:           return %[[VAL_3]] : tensor<1x12x5x5xf32>
// CHECK:         }

(mlir_venv) nod%

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