Skip to content

Instantly share code, notes, and snippets.

@silvasean
Created March 28, 2022 20:37
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
Star You must be signed in to star a gist
Save silvasean/b9c5f60dfbe3f51cf079bad3c76d095a to your computer and use it in GitHub Desktop.
#map0 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1)>
#map2 = affine_map<(d0, d1) -> (d0, d1)>
#map3 = affine_map<(d0, d1, d2, d3, d4) -> (d2, d3, d4)>
module attributes {torch.debug_module_name = "LayerNormModule"} {
func @forward(%arg0: tensor<2x5x2x2x3xf32>) -> tensor<2x5x2x2x3xf32> {
%cst = arith.constant 0.000000e+00 : f32
%c3_i64 = arith.constant 3 : i64
%c2_i64 = arith.constant 2 : i64
%cst_0 = arith.constant 1.000000e-05 : f64
%cst_1 = arith.constant dense<[[[3.000000e+00, 2.000000e+00, 4.000000e+00], [2.000000e+00, 3.000000e+00, 3.000000e+00]], [[3.000000e+00, 2.000000e+00, 4.000000e+00], [2.000000e+00, 3.000000e+00, 3.000000e+00]]]> : tensor<2x2x3xf32>
%cst_2 = arith.constant dense<[[[5.000000e-01, 4.000000e-01, 3.000000e-01], [2.000000e-01, 4.000000e-01, 3.000000e-01]], [[5.000000e-01, 4.000000e-01, 3.000000e-01], [2.000000e-01, 4.000000e-01, 3.000000e-01]]]> : tensor<2x2x3xf32>
%0 = arith.cmpi eq, %c2_i64, %c2_i64 : i64
cf.assert %0, "mismatching contracting dimension"
cf.assert %0, "mismatching contracting dimension"
cf.assert %0, "mismatching contracting dimension"
cf.assert %0, "mismatching contracting dimension"
cf.assert %0, "mismatching contracting dimension"
cf.assert %0, "mismatching contracting dimension"
%1 = arith.cmpi eq, %c3_i64, %c3_i64 : i64
cf.assert %1, "mismatching contracting dimension"
cf.assert %1, "mismatching contracting dimension"
cf.assert %1, "mismatching contracting dimension"
%2 = arith.muli %c2_i64, %c2_i64 : i64
%3 = arith.muli %2, %c3_i64 : i64
%4 = arith.sitofp %3 : i64 to f32
%5 = linalg.init_tensor [2, 5] : tensor<2x5xf32>
%6 = linalg.fill ins(%cst : f32) outs(%5 : tensor<2x5xf32>) -> tensor<2x5xf32>
%7 = linalg.generic {indexing_maps = [#map0, #map1], iterator_types = ["parallel", "parallel", "reduction", "reduction", "reduction"]} ins(%arg0 : tensor<2x5x2x2x3xf32>) outs(%6 : tensor<2x5xf32>) {
^bb0(%arg1: f32, %arg2: f32):
%15 = arith.addf %arg2, %arg1 : f32
linalg.yield %15 : f32
} -> tensor<2x5xf32>
%8 = linalg.generic {indexing_maps = [#map2, #map2], iterator_types = ["parallel", "parallel"]} ins(%7 : tensor<2x5xf32>) outs(%5 : tensor<2x5xf32>) {
^bb0(%arg1: f32, %arg2: f32):
%15 = arith.divf %arg1, %4 : f32
linalg.yield %15 : f32
} -> tensor<2x5xf32>
%9 = linalg.fill ins(%cst : f32) outs(%5 : tensor<2x5xf32>) -> tensor<2x5xf32>
%10 = linalg.generic {indexing_maps = [#map0, #map1, #map1], iterator_types = ["parallel", "parallel", "reduction", "reduction", "reduction"]} ins(%arg0, %8 : tensor<2x5x2x2x3xf32>, tensor<2x5xf32>) outs(%9 : tensor<2x5xf32>) {
^bb0(%arg1: f32, %arg2: f32, %arg3: f32):
%15 = arith.subf %arg1, %arg2 : f32
%16 = arith.mulf %15, %15 : f32
%17 = arith.addf %arg3, %16 : f32
linalg.yield %17 : f32
} -> tensor<2x5xf32>
%11 = linalg.generic {indexing_maps = [#map2, #map2], iterator_types = ["parallel", "parallel"]} ins(%10 : tensor<2x5xf32>) outs(%5 : tensor<2x5xf32>) {
^bb0(%arg1: f32, %arg2: f32):
%15 = arith.divf %arg1, %4 : f32
linalg.yield %15 : f32
} -> tensor<2x5xf32>
%12 = linalg.generic {indexing_maps = [#map2, #map2], iterator_types = ["parallel", "parallel"]} ins(%11 : tensor<2x5xf32>) outs(%5 : tensor<2x5xf32>) {
^bb0(%arg1: f32, %arg2: f32):
%15 = arith.truncf %cst_0 : f64 to f32
%16 = arith.addf %arg1, %15 : f32
%17 = math.rsqrt %16 : f32
linalg.yield %17 : f32
} -> tensor<2x5xf32>
%13 = linalg.init_tensor [2, 5, 2, 2, 3] : tensor<2x5x2x2x3xf32>
%14 = linalg.generic {indexing_maps = [#map0, #map1, #map1, #map3, #map3, #map0], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%arg0, %8, %12, %cst_1, %cst_2 : tensor<2x5x2x2x3xf32>, tensor<2x5xf32>, tensor<2x5xf32>, tensor<2x2x3xf32>, tensor<2x2x3xf32>) outs(%13 : tensor<2x5x2x2x3xf32>) {
^bb0(%arg1: f32, %arg2: f32, %arg3: f32, %arg4: f32, %arg5: f32, %arg6: f32):
%15 = arith.subf %arg1, %arg2 : f32
%16 = arith.mulf %15, %arg3 : f32
%17 = arith.mulf %16, %arg4 : f32
%18 = arith.addf %17, %arg5 : f32
linalg.yield %18 : f32
} -> tensor<2x5x2x2x3xf32>
return %14 : tensor<2x5x2x2x3xf32>
}
}
#map0 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
#map1 = affine_map<(d0, d1, d2) -> (d0, d1, 0)>
module attributes {torch.debug_module_name = "SoftmaxIntModule"} {
func @forward(%arg0: tensor<?x?x?xf32>) -> tensor<?x?x?xf32> {
%c1 = arith.constant 1 : index
%c0 = arith.constant 0 : index
%c2 = arith.constant 2 : index
%cst = arith.constant 0.000000e+00 : f32
%cst_0 = arith.constant 1.000000e+00 : f64
%cst_1 = arith.constant -3.40282347E+38 : f32
%c0_i64 = arith.constant 0 : i64
%0 = tensor.dim %arg0, %c0 : tensor<?x?x?xf32>
%1 = tensor.dim %arg0, %c1 : tensor<?x?x?xf32>
%2 = linalg.init_tensor [%0, %1, 1] : tensor<?x?x1xi64>
%3 = linalg.fill ins(%c0_i64 : i64) outs(%2 : tensor<?x?x1xi64>) -> tensor<?x?x1xi64>
%4 = linalg.init_tensor [%0, %1, 1] : tensor<?x?x1xf32>
%5 = linalg.fill ins(%cst_1 : f32) outs(%4 : tensor<?x?x1xf32>) -> tensor<?x?x1xf32>
%6:2 = linalg.generic {indexing_maps = [#map0, #map1, #map1], iterator_types = ["parallel", "parallel", "reduction"]} ins(%arg0 : tensor<?x?x?xf32>) outs(%5, %3 : tensor<?x?x1xf32>, tensor<?x?x1xi64>) {
^bb0(%arg1: f32, %arg2: f32, %arg3: i64):
%16 = linalg.index 2 : index
%17 = arith.index_cast %16 : index to i64
%18 = arith.cmpf ogt, %arg1, %arg2 : f32
%19 = arith.select %18, %arg1, %arg2 : f32
%20 = arith.select %18, %17, %arg3 : i64
linalg.yield %19, %20 : f32, i64
} -> (tensor<?x?x1xf32>, tensor<?x?x1xi64>)
%7 = tensor.dim %arg0, %c2 : tensor<?x?x?xf32>
%8 = arith.cmpi eq, %0, %0 : index
cf.assert %8, "mismatched size for broadcast"
%9 = arith.cmpi eq, %1, %1 : index
cf.assert %9, "mismatched size for broadcast"
%10 = linalg.init_tensor [%0, %1, %7] : tensor<?x?x?xf32>
%11 = linalg.generic {indexing_maps = [#map0, #map1, #map0], iterator_types = ["parallel", "parallel", "parallel"]} ins(%arg0, %6#0 : tensor<?x?x?xf32>, tensor<?x?x1xf32>) outs(%10 : tensor<?x?x?xf32>) {
^bb0(%arg1: f32, %arg2: f32, %arg3: f32):
%16 = arith.truncf %cst_0 : f64 to f32
%17 = arith.mulf %arg2, %16 : f32
%18 = arith.subf %arg1, %17 : f32
linalg.yield %18 : f32
} -> tensor<?x?x?xf32>
%12 = linalg.generic {indexing_maps = [#map0, #map0], iterator_types = ["parallel", "parallel", "parallel"]} ins(%11 : tensor<?x?x?xf32>) outs(%10 : tensor<?x?x?xf32>) {
^bb0(%arg1: f32, %arg2: f32):
%16 = math.exp %arg1 : f32
linalg.yield %16 : f32
} -> tensor<?x?x?xf32>
%13 = linalg.fill ins(%cst : f32) outs(%4 : tensor<?x?x1xf32>) -> tensor<?x?x1xf32>
%14 = linalg.generic {indexing_maps = [#map0, #map1], iterator_types = ["parallel", "parallel", "reduction"]} ins(%12 : tensor<?x?x?xf32>) outs(%13 : tensor<?x?x1xf32>) {
^bb0(%arg1: f32, %arg2: f32):
%16 = arith.addf %arg1, %arg2 : f32
linalg.yield %16 : f32
} -> tensor<?x?x1xf32>
cf.assert %8, "mismatched size for broadcast"
cf.assert %9, "mismatched size for broadcast"
%15 = linalg.generic {indexing_maps = [#map0, #map1, #map0], iterator_types = ["parallel", "parallel", "parallel"]} ins(%12, %14 : tensor<?x?x?xf32>, tensor<?x?x1xf32>) outs(%10 : tensor<?x?x?xf32>) {
^bb0(%arg1: f32, %arg2: f32, %arg3: f32):
%16 = arith.divf %arg1, %arg2 : f32
linalg.yield %16 : f32
} -> tensor<?x?x?xf32>
return %15 : tensor<?x?x?xf32>
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment