Searched refs:tiled_loop (Results 1 – 7 of 7) sorted by relevance
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/tfrt/jit/transforms/ |
| H A D | tf_jitrt_fuse_fill_into_tiled_reduction.cc | 62 SmallVector<OpFoldResult> GetParallelDimStep(LoopOp tiled_loop) { in GetParallelDimStep() argument
63 assert(tiled_loop.getNumLoops() == 2 && "Expected a 2D loop"); in GetParallelDimStep()
64 Value step = tiled_loop.isParallelDimension(0) ? tiled_loop.step().front() in GetParallelDimStep()
65 : tiled_loop.step().back(); in GetParallelDimStep()
112 LoopOp tiled_loop) const { in CloneAndAppendInitTensorToTiledLoop()
119 init.getLoc(), GetParallelDimStep(tiled_loop), in CloneAndAppendInitTensorToTiledLoop()
122 rewriter.updateRootInPlace(tiled_loop, [&]() { in CloneAndAppendInitTensorToTiledLoop()
124 &tiled_loop.appendOutputOperand(rewriter, init_clone); in CloneAndAppendInitTensorToTiledLoop()
126 return tiled_loop.getTiedBlockArgument(*init_clone_output_operand); in CloneAndAppendInitTensorToTiledLoop()
265 LoopOp tiled_loop) const { in CreateLoopWithUpdatedResults()
[all …]
|
| H A D | tf_jitrt_tile_transpose.cc | 103 auto tiled_loop = in matchAndRewrite() local
105 if (!tiled_loop) return failure(); in matchAndRewrite()
107 tiled_loop->walk( in matchAndRewrite()
110 rewriter.replaceOp(linalg_op, tiled_loop->getResults()); in matchAndRewrite()
|
| H A D | tf_jitrt_tile_cwise.cc | 67 LoopOp tiled_loop = in matchAndRewrite() local
69 if (!tiled_loop) return failure(); in matchAndRewrite()
71 tiled_loop->walk( in matchAndRewrite()
74 rewriter.replaceOp(linalg_op, tiled_loop->getResults()); in matchAndRewrite()
|
| H A D | tf_jitrt_passes.td | 123 let summary = "Fuse `linalg.fill` into `linalg.tiled_loop` with a reduction.";
132 `linalg.tiled_loop`.
214 let summary = "Vectorize tiled Linalg ops inside `tiled_loop`";
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/mlir_hlo/include/mlir-hlo/Dialect/gml_st/transforms/ |
| H A D | passes.td | 56 let summary = "Pass to bufferize linalg.tiled_loop with the ops inside it.";
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/mlir_hlo/include/mlir-hlo/Dialect/gml_st/ |
| H A D | README.md | 74 %sum = linalg.tiled_loop (%i, %j) = (%c0, %c0) to (%c80, %c60) step (%c4, %c4)
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/mlir_hlo/tests/Dialect/gml_st/ |
| H A D | ops.mlir | 120 func.func @tiled_loop(%lhs: tensor<24x64xi8>, %rhs: tensor<24x64xi8>,
151 // CHECK-LABEL: func @tiled_loop
|