You can not select more than 25 topics Topics must start with a chinese character,a letter or number, can include dashes ('-') and can be up to 35 characters long.

README.md 3.8 kB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899
  1. ## Build
  2. Same as [MegEngine](/README.md) except passing the additional flag "-DMGE_BUILD_IMPERATIVE_RT=ON" to cmake configure command.
  3. ## Test
  4. 1. Make sure `make develop` is executed
  5. 2. Setup `PYTHONPATH`
  6. ```bash
  7. export PYTHONPATH="$(git rev-parse --show-toplevel)/imperative/python"
  8. ```
  9. 3. Run `pytest` (pip install as needed)
  10. ```bash
  11. cd $(git rev-parse --show-toplevel)/imperative/python/test
  12. pytest
  13. ```
  14. ## Concepts
  15. ### Op and Tensor-like
  16. An op is a subclass of `OpBase` representing some operation, for example
  17. * `Elemwise`
  18. * `Reduce`
  19. Op can be parametrized. For example, `Elemwise` has a single parameter `mode`, which is required by its constructor.
  20. A tensor-like is a subclass of `TensorBase` that defines how ops should apply on it, for example
  21. * `RawTensor` launch kernel associated with op
  22. * `Tracer` record information for autodiff
  23. Op instances are callable with signature `(*args: TensorBase) -> Tuple[TensorBase]`. It will invoke the correct implementation for that specific op and tensor-like, e.g. launch kernel if `args` is `RawTensor`, record information for autodiff if `args` is `Tracer`.
  24. ### The `Const` Op
  25. The `Const` op is a special op that is used to convert literal to tensor-likes. Although it does not really use any input, at least one should be provided, otherwise it can't know which specific tensor-like to return.
  26. ### Tensor Wrapper
  27. Tensor-likes have a dataflow semantic, thus immutable. `TensorWrapper` provide a mutable layer on top of tensor-likes by replacing the wrapped tensor-like on demand.
  28. ## How to Wrap a MegBrain Op
  29. 1. Define the op
  30. Most ops have been automatically generated in `ops.builtin` using `.oprdecl` files (take a look at [basic_arith.oprdecl](/src/opr/impl/basic_arith.oprdecl)). If your op is already there, skip to next step.
  31. For other ops, this work can still be done automatically with the help of an Python op serializer that matches MegBrain's own.
  32. Before proceeding, if you are unfamiliar with MegBrain's serializer, here is a brief introduction. Each MegBrain op has a registered name, which is found at `MGB_SEREG_OPR(this_is_the_name, ...)` in some `.sereg.h` file. The default serializer simply write the memory of struct returned by `opr->param()`.
  33. You can create a serializer by subclassing `ops._internal.helper.OpDef` as follows
  34. ```python
  35. class WhateverDef(OpDef): # must end with "Def"
  36. name = 'Whatever' # name in MegBrain serialization registry
  37. param_names = ('param',) # Does not have to be 'param', but it is a good practice to mirror
  38. # C++ name, which is usually param(). It can also contain more
  39. # than one element, for example if the C++ serializer writes
  40. # `opr->param1()` followed by `opr->param2()`, you should use
  41. # ('param1', 'param2') instead.
  42. class Param:
  43. def serialize(self):
  44. c_struct_memory = bytes(...) # memory of a C++ `Param` struct
  45. return b'\x00'*4 + c_struct_memory # remember to add 4 leading bytes
  46. def __init__(self):
  47. self.param = self.Param(...) # must assign to attribute(s) specified in param_names
  48. ```
  49. A concrete example can be found at `ops._internal.misc_ops.DimshuffleDef`.
  50. Lastly, make sure it is imported in `ops._internal.all_ops` and a corresponding op will show up in `ops.builtin`
  51. 2. Define a convenience function
  52. Use `functional` as a reference.
  53. Tips:
  54. * an op instance has to be constructed before applying it
  55. `op = WhateverOp(param=...)`
  56. * apply an op by calling the op instance
  57. `outputs = op(*inputs)`
  58. * op always return a tuple
  59. `result = outputs`
  60. * input can be any tensor-like